High-girth Bethe assignment layer
Implementation & benchmark of arXiv:2609.02017 — Optimal girth-dependent bounds for the Bethe approximation of the permanent
Source paper: Optimal girth-dependent bounds for the Bethe approximation of the permanent arXiv:2609.02017 ⓘ · analyzed Sep 3, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace independent softmax routing or attention normalization with a differentiable approximate distribution over one-to-one assignments, using the Bethe permanent as the partition-function surrogate. Constrain the allowed token-to-expert or query-to-key support graph to have high girth, which gives an explicit bound on the approximation error and reduces short-cycle-induced correlations.
Formulas
Mathematical statement
For a nonnegative matrix $A\in\mathbb{R}_{\ge 0}^{n\times n}$, the permanent is $\operatorname{per}(A)=\sum_{\sigma\in\operatorname{Sym}(n)}\prod_{i=1}^n A_{i,\sigma(i)}$, equivalently the weighted sum over perfect matchings of the bipartite support graph $G_A$. Let $\mathcal{B}_n=\{\beta\in\mathbb{R}_{\ge0}^{n\times n}:\sum_j\beta_{ij}=1,\ \sum_i\beta_{ij}=1\}$ be the Birkhoff polytope, and let $E=\{(i,j):A_{ij}>0\}$. The Bethe approximation can be written as $\operatorname{Bethe}(A)=\exp(-\min_{\beta\in\mathcal{B}_n,\ \beta_{ij}=0\ \text{if }A_{ij}=0}F_B(\beta;A))$, with Bethe free energy $F_B(\beta;A)=\sum_{(i,j)\in E}\left[\beta_{ij}\log\frac{\beta_{ij}}{A_{ij}}-(1-\beta_{ij})\log(1-\beta_{ij})\right]$. The paper proves the universal comparison $\operatorname{Bethe}(A)\leq\operatorname{per}(A)\leq2^{n/2}\operatorname{Bethe}(A)$ and, if the bipartite support graph has even girth at least $g\ge4$, the sharper bound $\operatorname{per}(A)\leq2^{2n/g}\operatorname{Bethe}(A)$. Thus the log partition-function error is at most $(2n/g)\log2$. Here $n$ is the number of left and right vertices, $g$ is the length of the shortest cycle in the support graph, and $A_{ij}$ is the positive edge weight produced by the neural logits.
Implementation notes
Integrate this as a sparse assignment-normalization layer between token/query representations and experts/keys. For a batch of $n$ tokens and $n$ experts, compute logits $s_{ij}$ only on edges of a fixed sparse bipartite graph $E$, set $A_{ij}=\exp(s_{ij}/\tau)$ on those edges and $A_{ij}=0$ elsewhere, and require each row and column to sum to one through the Bethe optimization variable $\beta$. Use projected mirror descent or damped belief-propagation-style updates to minimize $F_B(\beta;A)$: initialize positive feasible $\beta$ by sparse Sinkhorn scaling; repeatedly compute the gradient $\partial F_B/\partial\beta_{ij}=\log(\beta_{ij}/A_{ij})+\log(1-\beta_{ij})+2$; take a mirror or exponentiated-gradient step; then approximately project onto the sparse doubly-stochastic constraints with 5--20 Sinkhorn iterations. Return $\beta$ as routing weights, optionally with top-1 or top-2 dispatch after balancing. The theorem is used to choose the graph: construct a random or algebraically generated sparse bipartite graph, measure its shortest even cycle with BFS, and retain patterns with girth $g\ge6$ or $8$. The mathematical quantity computed exactly is the certified bound $(2n/g)\log2$; the Bethe minimizer and its gradients are numerical approximations. For a first experiment, compare a 12-layer small Transformer MoE with standard independently normalized top-2 routing, sparse Sinkhorn routing, and high-girth Bethe routing on CIFAR-100 or WikiText-2. Match active edges and FLOPs. Measure training loss at equal steps, expert-load coefficient of variation, routing entropy, and the gap between Bethe and exact permanent on small $n\leq10$ instances. Success means lower load imbalance and more stable loss, or improved loss at equal dispatch cost; the strongest falsifiable signal is that the observed log-permanent error decreases as the support girth increases in accordance with the $(2n/g)\log2$ ceiling.
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.