Rayleigh–Jeans Condensation Routing

Integrate this at the attention-weight or MoE-router probability computation, initially for a small experimental branch rather than replacing every softmax. For attention, compute similarities \(s_{ij}=q_i^\top k_j/\sqrt d\) for query \(i\), set \(E_{ij}=c_i-s_{ij}\), and choose \(c_i=\max_j s_{ij}+\epsilon\) with \(\epsilon>0\). Choose a target mean energy \(\bar E_i\), initially by matching the mean energy of the baseline softmax attention at temperature 1, or use a learned scalar constrained to lie between the minimum and maximum candidate energies. Solve for \(\mu_i<\min_jE_{ij}\) by bisection on the scalar equation \(g(\mu)=\sum_jE_{ij}/(E_{ij}-\mu)\,/\,\sum_j1/(E_{ij}-\mu)-\bar E_i=0\). Then compute \(T_i=1/\sum_j1/(E_{ij}-\mu_i)\) and probabilities \(p_{ij}=T_i/(E_{ij}-\mu_i)\). In practice use \(\mu_i=\min_jE_{ij}-\delta\), with \(\delta\) bounded below, if the root is numerically close to the pole. For MoE, replace \(j\) by experts, use router scores as \(s_j\), and select top-k experts from \(p_j\), optionally renormalizing the selected probabilities. The mathematical quantities are the inverse-gap probabilities, normalization, mean-energy root, and pole-distance constraint; the target \(\bar E\), offset \(c\), and minimum gap \(\delta\) are engineering hyperparameters or learned values. The first cheap experiment should compare softmax, sparsemax, and this Rayleigh–Jeans router in a 6-layer Transformer on a small language-model corpus, measuring validation loss, attention entropy, effective support, and wall-clock cost. A successful signal is lower effective support at equal validation loss, or lower validation loss at equal FLOPs, without unstable gradients when the best score changes between tokens.