Rediseño de Enrutadores de Mezcla de Expertos con Iteración de Potencia en Variedades

Resumen

El enrutador es el componente fundamental de los modelos de Mezcla de Expertos. Al actuar como representantes de los expertos, las filas de la matriz del enrutador calculan su similitud con las entradas de la Mezcla de Expertos para determinar qué subconjunto de expertos se activa. Idealmente, cada fila del enrutador está diseñada para codificar la matriz del experto en este vector representativo, de modo que su producto escalar con el token pueda reflejar mejor la afinidad token-experto. Sin embargo, no existen principios de diseño que impongan esta condensación. En este artículo, proponemos alinear cada fila del enrutador con la dirección singular principal del experto asociado, ya que esta dirección proporciona la descripción matemática más expresiva de una matriz. Basándonos en este principio, proponemos un rediseño del enrutador mediante la Iteración de Potencia en Variedades (MPI, por sus siglas en inglés). Específicamente, introduce un paradigma de "Potencia y luego Retracción", donde se realiza un paso de iteración de potencia en los pesos del enrutador, seguido de una retracción para imponer una restricción de norma que garantice tanto la eficiencia como la estabilidad. Teóricamente, demostramos que MPI hace que las filas del enrutador converjan hacia las direcciones singulares principales de los expertos asociados. Empíricamente, preentrenamos modelos de Mezcla de Expertos en escalas desde 1B hasta 11B parámetros para confirmar que esta alineación facilita modelos de Mezcla de Expertos más efectivos.

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.