# Balanced KL projection for MoE routing

- ID: 2006
- Canonical URL: https://synthcore.org/idea/2006/balanced-kl-projection-for-moe-routing
- API JSON: https://synthcore.org/api/idea/2006.json
- API Markdown: https://synthcore.org/api/idea/2006.md
- Verification status: unverified
- Source: [arXiv:2608.08642](https://arxiv.org/abs/2608.08642)
- Category: training
- Solves: accuracy, scalability, stability
- ML areas: moe-routing, transformer, training
- Math tags: optimal-transport, convex-analysis, information-theory, linear-algebra
- Ratings: usefulness 6/10; difficulty 6/10; novelty 6/10

## Idea description

Use the shared-marginal KL projection to turn token-to-expert routing into a low-rank, exactly balanced assignment rather than relying only on an auxiliary load-balancing penalty. Tokens retain normalized routing distributions while the shared latent marginal enforces consistent aggregate usage across two independently learned routing factors.

## Mathematical statement

Let n tokens be assigned through nonnegative U,V in R^(n x r), with each row summing to one and with a common latent mass g=U^T 1_n=V^T 1_n. The construction W=U Diag(g)^(-1) V^T produces a nonnegative assignment kernel whose row and column sums are one. In a router, U can encode token-to-latent-group probabilities and V can encode expert-to-latent-group probabilities; the induced token-expert interaction need not be materialized. The exact KL projection uses generalized KL divergence, preserving positivity and producing multiplicative primal updates. Its dual Hessian is a sum of categorical covariance matrices C(p)=Diag(p)-p p^T, enabling matrix-free Newton steps in O(nr) time. The all-ones dual direction is redundant, so fixing one coordinate or imposing sum-zero gauge leaves r-1 variables.

## Key formulas

- $$W=U\operatorname{Diag}(g)^{-1}V^{\top},\qquad g=U^{\top}\mathbf{1}_n=V^{\top}\mathbf{1}_n.$$
- $$D_{\mathrm{KL}}(X\|X^0)=\sum_{i,k}\left[X_{ik}\log\frac{X_{ik}}{X^0_{ik}}-X_{ik}+X^0_{ik}\right].$$
- $$C(p)=\operatorname{Diag}(p)-pp^{\top},\qquad C(p)s=p\odot s-p\left(p^{\top}s\right).$$

## Implementation notes

Insert the method into a sparse MoE router immediately after router-logit computation and before top-k dispatch. For a batch of n tokens, produce factor logits A and B with shape n by r, then compute U0=softmax(A) and V0=softmax(B). Interpret U as token-to-latent routing and V as expert-to-latent routing, using a rectangular variant if the number of experts differs from r. Perform a KL projection toward the current factors subject to row-simplex and shared-marginal constraints. Use a damped Newton loop on the first r-1 dual coordinates: initialize dual variables at zero, compute normalized rows p_i and q_j by exponentiating and normalizing, evaluate marginal residuals, compute each Hessian-vector product with C(p)s=p elementwise-multiplied by (s-p dot s), solve the Newton system with 3 to 10 conjugate-gradient iterations, and backtrack until the KL objective decreases and all constraint residuals are below 1e-6. Use the projected factors to derive dispatch probabilities or top-k indices, and use straight-through gradients through the projection initially. The exact mathematics supplies positivity and marginal feasibility; choose capacity clipping, projection tolerance, and gradient treatment empirically. First benchmark an 8-expert Switch-Transformer-style model on a small language-model corpus, comparing standard auxiliary load loss, Sinkhorn balancing, and this projection at equal router FLOPs. Track expert-load coefficient of variation, dropped-token rate, validation perplexity, router gradient variance, and projection overhead. Success means lower dropped-token rate and expert imbalance without an auxiliary balancing loss, stable training, and less than 10 percent router-time overhead.

## Disclaimer

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