Redesenho de Roteadores Mixture-of-Experts com Iteração de Potência em Variedades

Resumo

O roteador é o componente fundamental dos modelos de Mistura de Especialistas (MoE). Servindo como proxies dos especialistas, as linhas da matriz do roteador calculam sua similaridade com as entradas do MoE para determinar qual subconjunto de especialistas é ativado. Idealmente, cada linha do roteador é projetada para codificar a matriz do especialista nesse vetor representativo, de modo que seu produto escalar com o token possa refletir melhor a afinidade token-especialista. No entanto, não existem princípios de design que imponham essa condensação. Neste artigo, propomos alinhar cada linha do roteador com a direção singular principal do especialista associado, pois essa direção fornece a descrição matemática mais expressiva de uma matriz. Com base nesse princípio, propomos um redesenho do roteador utilizando Iteração de Potência em Variedade (MPI, na sigla em inglês). Especificamente, introduz-se um paradigma de "Potência-depois-Contração", no qual um passo de iteração de potência é aplicado aos pesos do roteador, seguido por uma contração para impor uma restrição de norma, garantindo eficiência e estabilidade. Teoricamente, mostramos que o MPI faz com que as linhas do roteador convinjam em direção às direções singulares principais dos especialistas associados. Empiricamente, pré-treinamos modelos MoE em escalas de 1B a 11B parâmetros para confirmar que esse alinhamento facilita modelos MoE mais eficazes.

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.