# KL Mirror-Prox for coupled routing

- ID: 2078
- Canonical URL: https://synthcore.org/idea/2078/kl-mirror-prox-for-coupled-routing
- API JSON: https://synthcore.org/api/idea/2078.json
- API Markdown: https://synthcore.org/api/idea/2078.md
- Verification status: failed_benchmark
- Source: [arXiv:2608.10293](https://arxiv.org/abs/2608.10293)
- Category: optimization
- Solves: stability, accuracy, sample-efficiency
- ML areas: moe-routing, attention, optimizer
- Math tags: optimization, convex-analysis, information-theory
- Ratings: usefulness 8/10; difficulty 5/10; novelty 5/10

## 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.

## 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)\).

## Key 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)}}.$$

## 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

- Status: failed_benchmark
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: no_effect
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 9/10
- 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.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.20205
- Paired wins: 5/8
- Benchmark verdict: no measurable effect

## Artifacts

- [README.md](https://synthcore.org/code/775/README.md)
- [bench_report.json](https://synthcore.org/code/775/bench_report.json)
- [kl_mirror_prox_experiment.py](https://synthcore.org/code/775/kl_mirror_prox_experiment.py)
- [report.md](https://synthcore.org/code/775/report.md)
- [report_bench_2026-09-01T100708.md](https://synthcore.org/code/775/report_bench_2026-09-01T100708.md)
- [results.json](https://synthcore.org/code/775/results.json)
- [run_bench.py](https://synthcore.org/code/775/run_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/775)

## Disclaimer

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