✗ Failed on benchmark 2026

KL Mirror-Prox for coupled routing

Implementation & benchmark of arXiv:2608.10293 — Kullback-Leibler Mirror-Prox for Measure-Valued Variational Inequalities and Mean-Field Equilibria

Usefulness8/10

Difficulty5/10

Novelty5/10

Optimization Stability Accuracy Sample efficiency Moe routing Attention Optimizer Optimization Convex analysis Information theory

Source paper: Kullback-Leibler Mirror-Prox for Measure-Valued Variational Inequalities and Mean-Field Equilibria arXiv:2608.10293 ⓘ · analyzed Sep 1, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Replace a standard softmax-gradient update for a probability vector with a two-stage KL Mirror-Prox update. The predictor evaluates the population-dependent cost at the current distribution, and the corrector evaluates it at the predicted distribution, reducing oscillation when routing or attention costs are coupled across tokens or samples.

Formulas

$$\langle J(\mu,\cdot),\eta-\mu\rangle\ge 0,\qquad \forall\eta\in\mathcal{P}(\Omega).$$

$$\langle J(\mu^{\varepsilon},\cdot),\eta-\mu^{\varepsilon}\rangle\ge-\varepsilon,\qquad\forall\eta\in\mathcal{P}(\Omega).$$

$$\operatorname{KLProx}_{p}(g;\eta)=\arg\min_{q\in\Delta_n}\left\{\langle g,q\rangle+\frac{1}{\eta}\mathrm{KL}(q\|p)\right\},\qquad \mathrm{KL}(q\|p)=\sum_{i=1}^{n}q_i\log\frac{q_i}{p_i}.$$

$$z_{t,i}=\frac{p_{t,i}e^{-\eta c_i(p_t)}}{\sum_jp_{t,j}e^{-\eta c_j(p_t)}},\qquad p_{t+1,i}=\frac{p_{t,i}e^{-\eta c_i(z_t)}}{\sum_jp_{t,j}e^{-\eta c_j(z_t)}}.$$

Mathematical statement

The paper defines a variational inequality on probability measures: for a cost operator J and equilibrium measure $\mu^*$, $\langle J(\mu^*,\cdot),\eta-\mu^*\rangle\ge 0$ for every probability measure $\eta$. On a finite discretization with $n$ states, represent a measure by $p\in\Delta_n=\{p_i\ge0,\sum_i p_i=1\}$, and represent the cost by $c(p)\in\mathbb{R}^n$. The paper's $\varepsilon$-VI condition is $\langle J(\mu^\varepsilon,\cdot),\eta-\mu^\varepsilon\rangle\ge-\varepsilon$, while its $\varepsilon$-MVI condition is $\langle J(\eta,\cdot),\eta-\mu^\varepsilon\rangle\ge-\varepsilon$. The KL prox map is $\operatorname{KLProx}_{p}(g;\eta)=\arg\min_{q\in\Delta_n}\{\langle g,q\rangle+\eta^{-1}\mathrm{KL}(q\|p)\}$, where $\eta>0$ is the step size and $\mathrm{KL}(q\|p)=\sum_i q_i\log(q_i/p_i)$. Solving the constrained first-order conditions gives the explicit exponential reweighting $q_i=p_i\exp(-\eta g_i)/\sum_jp_j\exp(-\eta g_j)$. The adapted recursion is $z_t=\operatorname{KLProx}_{p_t}(c(p_t);\eta)$, followed by $p_{t+1}=\operatorname{KLProx}_{p_t}(c(z_t);\eta)$.

Implementation notes

(1) Integration point: use this in the router or attention probability update when each token's cost depends on aggregate routing or attention usage. For an MoE experiment, let $p_t\in\Delta_E$ be the global expert-load distribution over $E$ experts, and let $c(p)$ be the measured per-expert cost, combining negative validation reward, load-balancing penalty, and capacity overflow. For attention, use $p$ as a head or memory-slot distribution and define $c(p)$ from the next minibatch loss plus a congestion penalty. Do not replace the normal forward softmax initially; maintain a positive probability vector as a slowly updated routing prior or gate bias. (2) Pseudocode: initialize $p_i=1/n$; at each minibatch compute $c=\texttt{cost}(p)$; set $z_i=\exp(\log p_i-\eta c_i)$, normalize $z$; recompute the cost using the predicted distribution, $\tilde c=\texttt{cost}(z)$; set $p_i\leftarrow\exp(\log p_i-\eta\tilde c_i)$, normalize; use $\log p$ as an additive gate bias or attention-logit prior. Stop gradients through the two optimizer updates unless intentionally learning the cost model. (3) Compute the exponential reweighting and normalization exactly from the KL formula; estimate $c(p)$ empirically from batch losses, expert loads, overflow, or attention congestion. Use log-sum-exp for numerical stability and clip costs to a robust percentile range. Track the empirical Minty residual $\max_{k\in\mathcal{B}}\langle c(p_k),p_k-p_t\rangle$ over sampled competitor distributions $p_k$, and the VI residual against one-hot competitors. (4) First experiment: train a 4-expert, top-1 MoE MLP on CIFAR-100 or a small language-model subset, comparing Adam-updated softmax gates, entropic mirror descent, and KL Mirror-Prox at equal forward FLOPs. Sweep $\eta\in\{0.01,0.03,0.1,0.3\}$. Success means lower expert-load variance and fewer routing oscillations at equal validation loss, or faster validation-loss reduction; a strong signal is a 20% reduction in load variance or overflow without reducing accuracy.

Verification

Failed on benchmark

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗

Stage 2 · Practical benchmark result: No measurable effect

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 9/10

Built a numerically stable KL prox implementation and a coupled three-state cyclic routing proxy comparing entropic mirror descent with KL Mirror-Prox. The three mechanism predictions were confirmed: zero coupling produced exactly zero update, local amplification/contraction factors matched theory with maximum errors 5.60e-7 (MD) and 1.02e-6 (MP), and the predicted MP transition at x=eta*Lambda/3=1 was observed between x=0.667 and 1.333, with factor 1.00000003 at x=1. In the mini-run at x=2/3, MD diverged in radius while MP contracted from 0.04315 to 8.57e-6 and reduced the VI residual from 0.05658 to 1.61e-5.

Agent confidence: 9/10

Baseline: Entropic mirror descent: local factor matched sqrt(1+x^2); at eta=1, Lambda=2, radius 0.04315 -> 0.81081 and VI residual 0.05658 -> 1.13835 after 60 steps.
Idea: KL Mirror-Prox: local factor matched sqrt(1-x^2+x^4); at eta=1, Lambda=2, radius 0.04315 -> 8.57e-6 and VI residual 0.05658 -> 1.61e-5 after 60 steps. Maximum local prediction error was 1.02e-6.

Limitations: Only a deterministic three-state cyclic operator was tested; this is a mechanism-level toy verification rather than a trained MoE, attention, CIFAR, or language-model benchmark. No wall-clock or equal-FLOP neural routing comparison was performed, and behavior outside the local uniform-equilibrium analysis was not systematically swept.

How to run: python3 kl_mirror_prox_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.10293) Sep 2, 2026

Ahead of the baseline (−0.9%) but not statistically significant (p=0.2021, wins 5 / 8 paired seeds).

Benchmark: Correlated_token_moe_regression
Model: custom_moe_shared
Paired seeds: 8
Baseline mean: 0.0409
Idea mean: 0.0406
Effect (Δ): -0.0004 (−0.9%; negative = idea better)
Wins: 5 / 8 paired seeds
p-value: 0.2021 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±1.4%
Mechanism: Not confirmed ✗
Practical effect: No effect
Baseline tuning: swept over 9 configs

Limitations:

The benchmark used a small registered correlated-token MoE regression track with 400 training and 200 test examples, three experts, 18 epochs, and a custom optimizer loop because the intervention changes the training update. It did not test larger MoE systems, top-1 dispatch capacity overflow, CIFAR-scale routing, longer training, or wall-clock overhead. The baseline eta field is inert for the one-stage update but was included in the shared search-space parity grid.

How to run: python3 run_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Sep 1, 2026 latest Failed ✗

Agent confidence: 9/10

KL Mirror-Prox was implemented as a two-stage KL prior update in a matched three-expert MoE system on the registered correlated_token_moe_regression track. It produced a large reduction in observed expert-load variance, but the primary test MSE improvement was not statistically significant, so the bench verdict is no measurable effect and the idea did not win under the required criterion.

Baseline: 8-seed test MSE mean 0.04093059757724404; tuned baseline sweep best cfg {lr:0.006, eta:0.03}.
Idea: 8-seed test MSE mean approximately 0.04056028323248028; paired delta -0.0003703143447637558; 5/8 wins; permutation p=0.20205. Routing-load variance was substantially lower for the idea.

Benchmark result