Repenser les routeurs à mélange d'experts avec l'itération de puissance sur variété

Résumé

Le routeur est le composant fondamental des modèles Mixture d'Experts. Servant de proxies d'experts, les lignes de la matrice du routeur calculent leur similarité avec les entrées MoE afin de déterminer quel sous-ensemble d'experts est activé. Idéalement, chaque ligne du routeur est conçue pour condenser la matrice d'expert en ce vecteur représentatif, de telle sorte que son produit scalaire avec le jeton puisse mieux refléter l'affinité jeton-expert. Cependant, il n'existe aucun principe de conception imposant cette condensation. Dans cet article, nous proposons d'aligner chaque ligne du routeur sur la direction singulière principale de l'expert associé, cette direction fournissant la description mathématique la plus expressive d'une matrice. Sur la base de ce principe, nous proposons une refonte du routeur avec l'Itération de Puissance sur Variété (MPI). Plus précisément, elle introduit un paradigme « Puissance puis Rétractation », dans lequel une étape d'itération de puissance est effectuée sur les poids du routeur, suivie d'une rétractation imposant une contrainte de norme pour garantir à la fois efficacité et stabilité. Théoriquement, nous montrons que MPI conduit les lignes du routeur à converger vers les directions singulières principales des experts associés. Empiriquement, nous pré-entraînons un modèle MoE à différentes échelles, de 1B à 11B paramètres, pour confirmer que cet alignement favorise des modèles MoE plus efficaces.

English

Router is the cornerstone component to the Mixture-of-Experts models. Serving as expert proxies, the rows of the router matrix compute their similarity to the MoE inputs to determine which subset of experts is activated. Ideally, each router row is designed to encode the expert matrix into this representative vector, such that its dot-product with token can better reflect token-expert affinity. However, there exists no design principles to enforce this condensation. In this paper, we propose to align each router row with the principal singular direction of the associated expert, as this direction provides the most expressive mathematical description of a matrix. Based on this principle, we propose a router redesign with Manifold Power Iteration (MPI). Specifically, it introduces a "Power-then-Retract" paradigm, where a power iteration step is performed on the router weights, followed by a retraction to impose a norm constraint to ensure both efficiency and stability. Theoretically, we show that MPI drives router rows to converge toward the principal singular directions of associated experts. Empirically, we pretrain MoE model across scales from 1B to 11B parameters to confirm that this alignment facilitates more effective MoE models.