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

TL;DR

L2R introduces a unified routing framework for Mixture-of-Experts models that improves expert discrimination and model performance by reshaping routing space and controlling scoring geometry.

Contribution

It proposes Low-rank & Lipschitz-controlled Routing (L2R), a novel method combining low-rank routing space and Lipschitz control for enhanced MoE routing stability.

Findings

L2R improves routing geometry and expert discrimination.

L2R enhances overall model performance on language and vision tasks.

Experiments demonstrate consistent benefits across different MoE settings.

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 an OLMoE-based language MoE model and a vision MoE setting on ImageNet demonstrate that L2R consistently improves routing geometry, expert discrimination, and overall model performance. Code will be released.