Tunable Haar-Moment Mixing Regularizer

Integrate the regularizer after the recurrent or state-space transition, rather than on raw inputs. For each training sequence, collect a contiguous hidden-state block H=(h_1,...,h_L) from the RNN or SSM. Either RMS-normalize H or use \(\tilde h_t=h_t/(\|h_t\|_2+\epsilon)\) if the penalty should be scale invariant. Sample q Gaussian probes u_i and normalize each to unit norm. Precompute J normalized Gaussian reference vectors v_j in the hidden dimension d; these approximate Haar-random directions on the real unit sphere. Compute reference moments \(m_r^{ref}(u_i)=J^{-1}\sum_j(u_i^\top v_j)^r\), empirical moments \(\hat m_r(u_i)=L^{-1}\sum_t(u_i^\top\tilde h_t)^r\), and add \(\lambda q^{-1}\sum_i\sum_{r=1}^k(\hat m_r(u_i)-m_r^{ref}(u_i))^2\) to the task loss. Pseudocode: H=encoder(x); H=H/(norm(H,dim=-1)+eps); U=normalize(randn(q,d)); V=normalize(randn(J,d)); R=0; for r in 1..k: A=(H @ U.T)r; B=(V @ U.T)r; R += mean((A.mean(dim=time)-B.mean(dim=reference))**2); loss=task_loss+lambda*R; backward(loss). Start with k=2, then test k=3 and k=4; for a symmetric reference, odd target moments are zero and can be omitted in a first ablation. The paper provides the moment-matching principle, while probe and reference averages are estimated empirically. Use a small GRU and an S4-like SSM on sequential MNIST, permuted MNIST, and a synthetic long-range copy task, comparing against the unregularized model and RMSNorm plus weight decay. Log validation accuracy, hidden covariance condition number, gradient-norm variance, loss spikes, and long-horizon memory accuracy. Success is fewer unstable runs or lower validation loss at equal FLOPs without harming copy accuracy; failure is hidden-state collapse or increasing degradation as k and lambda grow.