Zero-Temperature Block-Triangular Routing

Integrate this at the attention-probability or MoE-router tensor, after computing ordinary logits and before multiplying by values or dispatching tokens. Partition the \(m\) tokens, experts, or channels into \(K\) ordered groups; assign each group a learned energy \(E_g\), and define \(\gamma_g=\exp(-\beta E_g)/\sum_h\exp(-\beta E_h)\). Build a nonnegative unconstrained matrix \(Q\) from router logits. Parameterize the final matrix with an upper-block mask so entries from a later group into an earlier group are zero, then project it toward the Gibbs-preserving equality \(P\gamma=\gamma\). For an MVP, initialize \(P=\operatorname{softmax}(Q,\mathrm{dim}=1)\), alternately normalize rows and apply a multiplicative correction based on the ratio \(\gamma/(P\gamma)\) for 5--10 iterations, followed by row renormalization. Use \(\beta(t)=\beta_{\min}+ (\beta_{\max}-\beta_{\min})t/T\), choosing \(\beta_{\max}\) so the smallest group weight is approximately \(10^{-3}\). Compute \(Y=PV\) for attention or dispatch tokens using \(P\) for MoE. The paper supplies the Gibbs weights, invariance constraint, critical-temperature intuition, and extremal block structure; group energies and logits are learned empirically. Start with a 6-layer, 8-head Transformer on CIFAR-10 patch classification or WikiText-2, using three token or expert groups. Compare dense softmax attention or routing at equal parameter count and batch size. Record validation accuracy, routing entropy, \(\|P\gamma-\gamma\|_1\), lower-block mass, active-edge count, and wall-clock/FLOPs. Success means equal-or-better accuracy with reduced cross-group traffic and no loss spike during annealing.