Balanced KL projection for MoE routing

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.