L2R: Low-Rank and Lipschitz-Controlled Routing for Mixture-of-Experts

TL;DR

The paper tackles fundamental limitations in MoE routing, notably representation mismatch, angular concentration, and scale-sensitive scoring that hinder expert specialization. It introduces Low-rank & Lipschitz-controlled Routing (L2R), which maps inputs to a shared low-rank routing space, represents experts with anchors in that space, and uses Saturated Inner-Product Scoring (SIPS) to bound magnitude effects while preserving directional information. A multi-anchor head mechanism further enhances expressiveness with minimal parameter overhead, and the entire design is optimized under standard MoE objectives. Empirical results across large-language MoE pretraining (OLMoE) and ViT-based ImageNet MoE demonstrate improved routing stability, more coherent expert usage, faster convergence, and superior downstream performance, indicating broad applicability in both language and vision domains.

Abstract

Mixture-of-Experts (MoE) models scale neural networks by conditionally activating a small subset of experts, where the router plays a central role in determining expert specialization and overall model performance. However, many modern MoE systems still adopt linear routers in raw high-dimensional representation spaces, where representation mismatch, angular concentration, and scale-sensitive scoring can jointly undermine routing discriminability and stable expert specialization. In this work, we propose Low-rank \& Lipschitz-controlled Routing (L2R), a unified routing framework that reshapes both the routing space and scoring geometry. L2R performs expert assignment in a shared low-rank latent routing space and introduces Saturated Inner-Product Scoring (SIPS) to explicitly control the Lipschitz behavior of routing functions, yielding smoother and more stable routing geometry. In addition, L2R incorporates a parameter-efficient multi-anchor routing mechanism to enhance expert expressiveness. Extensive experiments on a large-scale language MoE model and a vision MoE setting on ImageNet demonstrate that L2R consistently improves routing stability, expert specialization, and overall model performance.

Figures (17)

• Figure 1: Comparison between the linear router (left) and the proposed L2R framework (right). (a) Routing space: Tokens are projected from high-dimensional raw representations into a low-rank routing space as $\bm{q}$, where experts are represented by learnable anchors $\bm{k}$. (b) Scoring mode: Compared to dot-product scoring, SIPS reshapes the score landscape into a bounded and smoother geometry, improving routing stability and expert specialization.
• Figure 2: Variance of pairwise cosine similarity in routing spaces. We compare the layer-averaged variance of pairwise token cosine similarities in the latent routing space for Linear, X-MoE chi2022on, and L2R-SIPS applied to OLMoE muennighoff2025olmoe. For L2R-SIPS, we use rank $r{=}2$, hence the isotropic reference corresponds to the 2D value ($0.5$). While Linear and X-MoE exhibit near-zero variance, L2R-SIPS attains a much higher variance, indicating substantially improved angular diversity.
• Figure 3: Score landscapes under fixed expert anchor. Heatmaps visualize routing logits as a function of query location $\bm{q}=(Q_x,Q_y)$ with a fixed anchor $\bm{k}=[2,2]$. Standard dot-product yields a linear half-space separation with unbounded magnitude effects, while SIPS reshapes the landscape into a bounded, angle-sensitive geometry that is more amenable to stable routing.
• Figure 4: OLMoE training dynamics. Curves show MMLU and HellaSwag accuracies, C4 raffel2020exploring validation cross-entropy (CE), training CE, and load-balance loss. L2R exhibits clear convergence over 10B tokens and consistently improves MMLU/HellaSwag.
• Figure 5: Training dynamics (train cross-entropy loss) for ablations. Both lower rank and more heads yield faster convergence over 10B tokens.