Subgaussian orthogonal feature basis

Integrate the method after a transformer MLP, attention output, or residual stream before the next linear projection. For a minibatch feature matrix X of shape B×d, maintain an exponential-moving-average covariance or use a detached minibatch covariance; compute mean-centered Z=X−mean(X), Σ̂=(ZᵀZ/B)+εI, and whiten with a symmetric inverse square root obtained by eigendecomposition. Insert a trainable orthogonal matrix Q, parameterized as a product of Householder reflections or retracted after each optimizer step by Q←qf(Q). Set Y=ZΣ̂−1/2Q. For every channel j, estimate its empirical ψ2-norm by binary-searching t so that B−1∑i exp((Yij/t)^2)≤2; clamp squared values to 25 for numerical stability. Compute the empirical L2 denominator and the smooth maximum loss shown in the formula, then add λLsubG to the task loss. Backpropagate through Q and the feature transform, but optionally stop gradients through the covariance inverse square root during the first implementation. Pseudocode: Z=X-X.mean(0); S=ema_cov(Z); A=inv_sqrt(S+eps*I); Y=Z@A@Q; for j: psi[j]=binary_search_t(mean(exp(clamp(Y[:,j]^2/t^2,max=25)))<=2); r[j]=psi[j]/sqrt(mean(Y[:,j]^2)+eps); loss=task+lambda*tau*logmeanexp(r/tau). The mathematical quantities are the Orlicz threshold and the ratio; covariance estimation, clipping, and temperature τ are engineering choices. First test on a 6-layer ViT-Tiny or small GPT-style transformer on CIFAR-100 or WikiText-2, comparing LayerNorm alone against LayerNorm plus this penalty with matched FLOPs. Measure validation loss or accuracy, gradient-norm variance, the 99.9th percentile of channel activations, and the maximum held-out ψ2/L2 ratio. Success means lower tail ratios and fewer gradient spikes without slower training or degraded accuracy; a useful target is at least 20% lower extreme activation magnitude or gradient variance at equal validation quality.