KL Mirror-Prox for coupled routing

What is implemented

kl_mirror_prox_experiment.py implements the KL prox update

prox(p,c,eta) = normalize(exp(log(p) - eta*c))

and compares one-stage entropic mirror descent (MD), p_next = prox(p, A p, eta), against the proposed two-stage KL Mirror-Prox (MP), z = prox(p, A p, eta); p_next = prox(p, A z, eta).

The coupled-routing proxy uses a three-state cyclic skew operator A, whose tangent-plane eigenvalues are +/- i Lambda; the uniform distribution is the equilibrium.

Quantitative mechanism checks

At the uniform distribution, the KL prox linearization is I - (eta/3) A. Defining x = eta*Lambda/3, the predicted local radius factors are:

• MD: sqrt(1 + x^2) (always amplifies a cyclic coupling).
• MP: sqrt(1 - x^2 + x^4) (contracts for 0 < x < 1, is neutral at x=1, and amplifies for x>1).
• At Lambda=0, both methods are exactly the identity.

The deterministic sweep in results.json measured maximum absolute errors of 5.60e-7 for MD and 1.02e-6 for MP. The zero-coupling change was exactly 0.0; the measured transition was bracketed by x=0.667 (contracting) and x=1.333 (expanding), with the x=1.0 measurement equal to 1.00000003.

Mini experiment

At eta=1, Lambda=2 (x=2/3), after 60 updates from the same perturbed uniform distribution:

• MD radius: 0.04315 -> 0.81081; VI residual: 0.05658 -> 1.13835.
• MP radius: 0.04315 -> 0.00000857; VI residual: 0.05658 -> 0.0000161.

This is a mechanism-level toy result, not a trained MoE benchmark.

Reproduce

/home/maxwelhelp/main/bin/python3 kl_mirror_prox_experiment.py

Outputs are written to results.json.