Yang–Baxter Pairwise Router
Implementation & benchmark of arXiv:2607.18959 — Integrable multi-species SSEP with reactive particle species
Source paper: Integrable multi-species SSEP with reactive particle species arXiv:2607.18959 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace unconstrained pairwise token-routing interactions with a structured two-token router derived from an involutive set-theoretical Yang–Baxter solution. The pair operator is a convex interpolation between identity and a permutation of discrete routing states, so it cannot amplify probability mass or logits when applied to routing distributions. The Yang–Baxter relation provides a falsifiable test for whether three-token routing updates are insensitive to the two admissible overlapping-update orderings.
Formulas
Mathematical statement
The paper starts with a finite set of species X and an operator \(\check r:V\otimes V\to V\otimes V\), where V has basis \(\{|i\rangle:i\in X\}\), acting as \(\check r|i,j\rangle=|g_i(j),f_j(i)\rangle\). The maps \(f_j,g_i:X\to X\) are bijections, \(\check r\) satisfies the braid/Yang–Baxter equation \(\check r_{12}\check r_{23}\check r_{12}=\check r_{23}\check r_{12}\check r_{23}\), and it is involutive: \(\check r^2=\mathrm{Id}\). Baxterisation gives \(\check r(u)=\frac{u\check r+\mathrm{Id}_{V\otimes V}}{u+1}\), with \(u\ge0\). If \(\check r\) is represented by a permutation matrix P, then \(R(u)=(uP+I)/(u+1)=(1-a)I+aP\), where \(a=u/(u+1)\in[0,1)\). Thus R(u) is a convex combination of identity and a permutation, preserves nonnegative probability vectors and total mass, and has no amplification in \(\ell_1\). The parameter-dependent consistency relation is \(\check r_{12}(u)\check r_{23}(u+v)\check r_{12}(v)=\check r_{23}(v)\check r_{12}(u+v)\check r_{23}(u)\).
Implementation notes
Use this as an experimental replacement for a small part of MoE routing rather than as a full Transformer replacement. Let each token have a categorical expert distribution \(p_t\in\Delta^{N-1}\), and form pair distributions \(q_{t,t+1}=p_t\otimes p_{t+1}\in\mathbb R^{N^2}\). Choose a concrete involutive set-theoretical solution by defining explicit bijections \(f_j,g_i\) on \(X=\{1,\ldots,N\}\), enumerate the map \((i,j)\mapsto(g_i(j),f_j(i))\), and store its permutation matrix \(P\in\mathbb R^{N^2\times N^2}\). For each adjacent pair, compute \(u_t=\mathrm{softplus}(s_t)\), set \(a_t=u_t/(1+u_t)\), and update \(q'_{t,t+1}=((u_tP+I)/(u_t+1))q_{t,t+1}\). Marginalize \(q'\) over the second and first expert indices to obtain revised distributions \(\tilde p_t,\tilde p_{t+1}\), then apply \(p_t\leftarrow(1-\rho)p_t+\rho\tilde p_t\) with \(\rho\le1\). During training, use straight-through top-1 or top-2 routing after this operation. A cheaper implementation avoids materializing P: reshape the pair tensor to \([N,N]\), apply the species permutation by indexing, and mix it with the original tensor. The mathematical quantities are exact; \(u_t,s_t,\rho\) are learned or scheduled empirically. First test on a 4-layer, 8-expert Switch-style MLP on CIFAR-100 or a small language-model corpus, comparing standard router logits against this pair router at equal FLOPs. Measure auxiliary load-balancing loss, router entropy, gradient norm spikes, dropped-token rate, and validation loss. Success means fewer routing collapses or dropped tokens and smoother loss at equal perplexity or accuracy.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.