Rook: Global Permutation Geometry for Neural Routing

Abstract

Modern optimizer stacks match geometry to parameter class, giving dense hidden matrices a spectral duality map and vectors coordinatewise adaptivity, yet doubly stochastic routing blocks, whose natural atoms are permutations, are still trained by applying AdamW to an unconstrained Sinkhorn score. We construct the missing member of the stack from those atoms. The central symmetrization of the Birkhoff polytope is the unit ball of a matching norm whose exact smooth duality map is a Gibbs expectation over permutations with a matrix permanent as its potential, and an exact identity splits that potential into Sinkhorn free energy plus permanent structure of the scaled matrix. Rook replaces the intractable part with a fractional-Bethe map whose coefficient \(\gamma_n=-(n-1)/n\) is uniquely fixed by exact permutation covariance, and which stays Lipschitz in the score where the hard assignment oracle jumps across ties. The construction pays off where routing has to adapt. After the routed streams are rewired from a shared checkpoint, Rook lowers post-shift cross-entropy from \(0.415\) to \(0.156\) against Muon+AdamW and recovers \(2.96\times\) sooner, on all ten paired seeds and at \(7.1\%\) lower step time. Under batch-32 gradient noise it reduces direction chatter by \(62.4\%\) relative to that oracle and halves final test loss. On stationary training AdamW on the scores stays ahead.

A matching norm for routing blocks

Some neural routers and residual-stream mixers are constrained to the Birkhoff polytope

\[\mathcal B_n=\{B\in\mathbb R^{n\times n}:B\ge0,\ B\mathbf 1=\mathbf 1,\ B^\top\mathbf 1=\mathbf 1\}=\operatorname{conv}(\mathcal P_n),\]

where \(\mathcal P_n\) is the set of permutation matrices. Their atoms are assignments, not coordinates and not orthogonal matrices, yet the common implementation forms \(B=\operatorname{Sinkhorn}(S/T)\) and applies AdamW to the score \(S\).

\(\mathcal B_n\) is not centrally symmetric, so it cannot serve as a signed unit ball. Its central symmetrization \(\mathcal K_n=\tfrac12(\mathcal B_n-\mathcal B_n)\) contains a neighborhood of zero in the tangent space and is therefore the unit ball of a norm on the tangent space, the matching norm. Its dual norm is the assignment spread

\[\|G\|_{\mathrm{match},*}=\tfrac12\Bigl[\max_{P\in\mathcal P_n}\langle G,P\rangle-\min_{P\in\mathcal P_n}\langle G,P\rangle\Bigr],\]

computable by two assignment solves, and its linear oracle is the signed difference of extreme assignments, which we call the hard LMO. This is the polytope analogue of choosing updates from the norm ball induced by a parameter class’s own atoms. Orthogonal matrices give the spectral ball and the polar factor (Muon), and permutations give \(\mathcal B_n\) and a difference of assignments.

The Sinkhorn–permanent decomposition

For \(\tau>0\) the smooth matching potential and its marginal are

\[\Psi_\tau(G)=\tau\log\operatorname{per}(e^{G/\tau}),\qquad M_\tau(G)=\nabla\Psi_\tau(G)=\mathbb E_{q_G}[P]\in\mathcal B_n,\]

with Hessian \(\tfrac1\tau\operatorname{Var}_{q_G}(\langle\cdot,P\rangle)\): the geometry of the block is the covariance of a Gibbs law over assignments. The omitted term is exact:

Theorem (Sinkhorn–permanent decomposition). For every \(G\) and \(\tau>0\), \[\Psi_\tau(G)=F_\tau(G)+\tau\log\operatorname{per}\bigl(B_\tau(G)\bigr),\qquad F_\tau(G)=\tau\log\operatorname{cap}\bigl(e^{G/\tau}\bigr).\] Consequently \(\tau\log\nu_n\le\Psi_\tau(G)-F_\tau(G)\le0\) with \(\nu_n=n!/n^n\), by the Van der Waerden bound. Capacity captures the Sinkhorn part of matching geometry and nothing else.

Differentiating gives the exact marginal correction, and each of its entries depends on a global \((n{-}1)\times(n{-}1)\) subpermanent, which is why coordinatewise corrections are structurally wrong.

One curvature-matched global map

Theorem (Unique curvature match). For tangent \(Z\), \(DR_{\tau,\gamma}(0)[Z]=\tfrac{n-1}{\tau n(n-1+\gamma)}Z\). The unique coefficient matching the exact response \(DM_\tau(0)[Z]=Z/[\tau(n-1)]\) is \[\gamma_n=-\frac{n-1}{n}.\]

Raw Bethe over-responds by \((n-1)^2/[n(n-2)]\); Rook removes that finite-width error, and since \(\gamma_n\in[-1,-\tfrac12]\) its potential stays inside the certified bracket containing the true permanent.

The signed map and its guarantees

Proposition (Lipschitz response versus tie discontinuity). Where the restricted fractional Hessian has curvature at least \(\mu>0\), the Rook marginal is \(1/\mu\)-Lipschitz in the dimensionless score, \(\|R(S_1)-R(S_2)\|_F\le\mu^{-1}\|S_1-S_2\|_F\), and the odd signed map inherits the bound. The hard LMO is piecewise constant and admits no finite local Lipschitz constant across a hyperplane where two assignments tie.

hard LMO \(\tfrac12(P_+-P_-)\)

signed Gibbs map \(\tfrac12\{M(\bar G)-M(-\bar G)\}\)

direction chatter \(1-\cos(d_t,d_{t-1})\)

score angle \(\theta\) 0.00 temperature \(\rho\) 1.00

Interactive figure. The two maps evaluated exactly at \(n=4\) by permutation enumeration, on a score \(G(\theta)\) rotating through two fixed tangent directions. Entries are signed (blue positive, orange negative). Where the score crosses a hyperplane on which two assignments tie, the hard LMO jumps between vertices and the chatter trace spikes; the signed Gibbs map turns continuously. In the paper’s experiments Rook matches the exact Gibbs reference to four decimals, and temperature interpolates between tangent-gradient geometry (\(\rho\to\infty\)) and the assignment oracle (\(\rho\downarrow0\)) inside one optimizer.

Experiments

Hyperparameters are selected on development seeds and frozen before held-out evaluation; exact permanents and Gibbs marginals appear only in validation oracles, never in training. Every comparison is paired. Arms are cloned from one checkpoint and receive identical minibatches, so only the optimizer assignment varies.

Fidelity to the exact map

For \(n\in\{4,6,8\}\), five temperatures spanning the deployed range, and 50 held-out tangent matrices per cell, we compare signed directions against exact permutation enumeration. Rook wins 748/750 paired cases, cutting Sinkhorn’s mean relative error from \(0.2169\) to \(0.0395\) and raw Bethe’s by \(72.8\%\).

Map	Mean cosine ↑	Minimum cosine ↑	Mean rel. error ↓
Sinkhorn (\(\gamma=0\))	0.987724	0.928120	0.216913
Classical Bethe (\(\gamma=-1\))	0.990092	0.859468	0.145470
Rook (\(\gamma_n\))	0.998802	0.983241	0.039506

Table 1. Fidelity to the exact signed Gibbs map over 750 held-out cases.

Adaptation after a rewiring

Four ordered image streams feed a learned \(4\times4\) Birkhoff router, a residual MLP, and a classifier; one Muon+AdamW checkpoint per seed is cloned into every arm and the stream order is rewired. Rook lowers mean post-shift CE by \(62.5\%\) and recovers \(2.96\times\) sooner, winning loss and recovery on 10/10 paired seeds at \(7.1\%\) lower step time, and matches the exact Gibbs reference to four decimals. The score-cap control rules out clipping.

Stack	Post-shift CE ↓	Recovery (ep.) ↓	Accuracy (%) ↑	ms/step
AdamW	.406±.009	6.9±.5	94.44±.36	2.882
Muon+AdamW	.415±.015	7.1±.5	95.91±.20	3.915
+ score cap	.270±.011	5.3±.3	96.42±.19	4.056
+ hard LMO	.174±.008	2.6±.2	97.42±.13	3.640
+ Rook	.156±.006	2.4±.3	97.16±.24	3.639
+ exact Gibbs	.156±.006	2.2±.2	97.09±.23	3.941

Table 2. Held-out rewiring from a shared checkpoint, mean ± s.e. over 10 seeds; CE averages eight adaptation epochs. Each arm isolates one explanation.

Hard versus smooth geometry under gradient noise

What separates the two matching-norm maps is continuity. We clone the same checkpoints, feed identical minibatches to both, and reduce the post-shift batch from 128 to 32 without retuning; chatter is the mean of \(1-\cos(d_t,d_{t-1})\) over consecutive directions. At batch 128 Rook reduces chatter by \(80.1\%\) with lower trajectory CE in 9/10 seeds; at batch 32 it reduces chatter by \(62.4\%\), halves final test loss, and improves accuracy by \(2.42\) points in all ten paired seeds, while the hard LMO reaches the recovery threshold \(0.3\) epoch sooner. The two maps therefore sit at different points of a speed and stability tradeoff, the smooth map trading a little threshold speed for continuity, steadier directions, and better final quality once gradients are noisy.

Batch	Map	Chatter ↓	Test CE ↓	Accuracy (%) ↑
128	Hard LMO	.178±.013	.110±.013	96.98±.59
128	Rook	.035±.002	.098±.005	97.47±.24
32	Hard LMO	.182±.004	.149±.034	95.42±1.06
32	Rook	.069±.001	.073±.008	97.84±.29

Table 3. Hard versus smooth map under gradient noise, mean ± standard error over ten seeds.

A 5M-parameter transformer

We repeated the protocol on a 5M-parameter byte-level model on enwik8, six layers of width 256 split into four 64-wide streams, each mixing through a learned \(4\times4\) Birkhoff block, over three seeds. Muon’s learning rate on the scores is flat across a \(30\times\) grid, so that baseline is not undertuned; Rook’s tuned step size here is \(0.25\), and at that value it reaches the best \(8{,}000\)-step endpoint of any arm. Under six recurring rewirings it pays a one-time cost in the first cycle, then leads in every later cycle by \(+0.16\) to \(+0.20\) mean CE.

Transformer-scale recovery after one stream permutation: validation cross-entropy and route sharpness over training steps — Figure 1. Transformer-scale recovery after one stream permutation, mean over three seeds. Left: validation CE; the dashed line is the pre-shift \(1.140\). Right: mean route sharpness, which separates the two optimizer classes throughout recovery.

Six recurring rewirings: validation cross-entropy and route sharpness over cycles — Figure 2. Six recurring rewirings of 750 steps, mean over three seeds; vertical lines mark the rewirings. Top: validation CE. Bottom: route sharpness, averaged over routes. The matching-norm arms re-adapt more cheaply at each later rewiring, while the arm trained on the score gradient pays a larger recovery every cycle. The two classes hold persistently different route sharpness.

Limitations

The solver targets widths up to a few dozen, the evidence isolates optimizer behavior under rewiring and noise, and the transformer study is a single 5M-parameter model, not a pretraining claim. Rook’s step size does not transfer across scale, and stationary runs favor AdamW. We do not compare against projected gradient or mirror descent on \(\mathcal B_n\), so the advantage cannot yet be separated from primal updating in general; rectangular transportation polytopes also remain open.

BibTeX

@misc{kandukoori2026rook,
  title  = {Rook: Global Permutation Geometry for Neural Routing},
  author = {Kandukoori, Sushaan and Patel, Aarya},
  year   = {2026},
  note   = {Preprint}
}