Complete MLSI Heat Regularization for Matrix Attention

(1) Integration point: modify a multi-head attention or MoE router after it produces matrix logits Z with shape [batch, sequence, heads, heads] or [batch, token, experts, experts]. Interpret the token or expert index as a periodic coordinate θ. Convert each matrix logit into a positive definite density matrix using S=softplus(Z)+εI and ρ=S/tr(S), with ε=10^−4. Apply heat smoothing along the sequence coordinate before scalarizing ρ into attention or routing scores. (2) Pseudocode: S = softplus(Z) + eps*I; rho = S / trace(S); rho_hat = FFT(rho, dim=theta); for m in frequencies: rho_hat[m] *= exp(-tau*m*m); rho = real(IFFT(rho_hat)); rho = project_to_PD_and_trace_one(rho); scores = log(diag(rho)+eps). The multiplier exp(-tau*m*m) is the exact Fourier action of P_t=e^{t∂_θ²}. Add lambda * I_BKM(rho) to the loss, estimating ∂θρ by periodic central differences and evaluating Ω_ρ^−1 by eigendecomposition with eigenvalues clipped below ε. (3) The paper supplies the complete MLSI and predicted contraction e^−2t; FFT discretization, ε, λ, τ, and scalarization are empirical choices. During training estimate the ratio I_BKM(P_tρ)/I_BKM(ρ) and compare it with e^−2t. (4) First cheap experiment: train a 2-layer, 128-dimensional, 4-head Transformer on WikiText-2 or a synthetic periodic sequence task. Compare no regularizer, ordinary attention-entropy regularization, and this method at matched FLOPs. Sweep τ in {0.01,0.05,0.1} and λ. Record validation loss, high-frequency attention power, gradient-norm variance, head-collapse rate, and wall-clock overhead. Success means lower validation loss or fewer training spikes at equal compute, with measured BKM information decreasing near the predicted e^−2t rate and overhead below 10%.