Exact Rank-Space KL Projection for Shared-Marginal Low-Rank Factors: Application to Doubly Stochastic Clustering
arXiv:2608.08642
2026
Architecture
2 ideas extracted · analyzed Sep 1, 2026
What the math gives to ML
The paper gives a constructive way to represent and optimize an exactly doubly stochastic nonnegative matrix without materializing its dense entries. Its key asset is the shared latent marginal: two independent row-simplex factors U and V can induce W=U Diag(g)^(-1) V^T while preserving both row and column sums exactly, and the KL projection can be solved in only r-1 effective dual dimensions. This is directly transferable to low-rank attention, token mixing, or MoE assignment matrices where conservation of mass and balanced usage are desirable. The most promising experiment is a low-rank doubly stochastic attention layer whose factor updates are followed by the paper's exact KL projection, with associative multiplication retaining O(nr) memory and compute rather than O(n^2).
Ideas from this paper
△ Mechanism confirmed, baseline not beaten
2026
Replace an n-by-n attention or token-mixing matrix with two nonnegative rank-r factors having row-simplex constraints and a shared latent column marginal. The induced matrix is exactly doubly stochastic at every accepted update, while applying it to values uses two thin matrix multiplications and never constructs the dense attention matrix.
Useful7/10
Difficulty6/10
Novelty7/10
Unverified
2026
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.
Useful6/10
Difficulty6/10
Novelty6/10