Neugestaltung von Mixture-of-Experts-Routern mit Mannigfaltigkeits-Potenziteration

Zusammenfassung

Der Router ist die grundlegende Komponente von Mixture-of-Experts-Modellen. Als Experten-Proxys berechnen die Zeilen der Router-Matrix ihre Ähnlichkeit mit den MoE-Eingaben, um zu bestimmen, welche Teilmenge der Experten aktiviert wird. Idealerweise ist jede Router-Zelle so konzipiert, dass sie die Experten-Matrix in diesen repräsentativen Vektor kodiert, sodass ihr Skalarprodukt mit dem Token die Token-Experten-Affinität besser widerspiegeln kann. Es gibt jedoch keine Gestaltungsprinzipien, die diese Kondensierung erzwingen. In diesem Papier schlagen wir vor, jede Router-Zelle mit der hauptsächlichen singulären Richtung des zugehörigen Experten auszurichten, da diese Richtung die ausdrucksstärkste mathematische Beschreibung einer Matrix liefert. Basierend auf diesem Prinzip schlagen wir eine Neugestaltung des Routers mittels Manifold Power Iteration (MPI) vor. Konkret führt dies ein "Power-then-Retract"-Paradigma ein, bei dem ein Potenzierungsschritt auf die Router-Gewichte angewendet wird, gefolgt von einer Retraktion, um eine Normbedingung aufzuerlegen, die sowohl Effizienz als auch Stabilität gewährleistet. Theoretisch zeigen wir, dass MPI die Router-Zellen zur Konvergenz in Richtung der hauptsächlichen singulären Richtungen der zugehörigen Experten treibt. Empirisch trainieren wir MoE-Modelle in Größenordnungen von 1B bis 11B Parametern vor, um zu bestätigen, dass diese Ausrichtung effektivere MoE-Modelle ermöglicht.

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.