만폴드 파워 이터레이션을 활용한 전문가 혼합 모델 라우터 재설계

Redesign Mixture-of-Experts Routers with Manifold Power Iteration

라우터는 전문가 혼합(Mixture-of-Experts, MoE) 모델의 핵심 구성 요소입니다. 라우터 행은 전문가 프록시 역할을 하며, 입력 데이터와 각 전문가 간의 유사도를 계산하여 어떤 전문가 집합이 활성화될지를 결정합니다. 이상적으로, 각 라우터 행은 해당 전문가 행렬을 대표 벡터로 인코딩하여 토큰과 전문가 간의 친화력을 더 잘 반영해야 합니다. 그러나 이러한 인코딩을 강제하는 설계 원칙은 존재하지 않습니다. 본 논문에서는 각 라우터 행을 관련 전문가의 주 특이 방향(principal singular direction)에 맞추는 방법을 제안합니다. 이는 행렬을 가장 잘 설명하는 표현 방식이기 때문입니다. 이 원칙에 기반하여, 만폴드 파워 이터레이션(MPI)을 활용한 라우터 재설계 방안을 제시합니다. 구체적으로, MPI는 '파워-리트랙' 패러다임을 도입하며, 먼저 라우터 가중치에 대해 파워 이터레이션을 수행하고, 이후 정규화 제약 조건을 적용하여 효율성과 안정성을 확보합니다. 이론적으로, MPI가 라우터 행을 관련 전문가의 주 특이 방향으로 수렴하도록 유도한다는 것을 증명했습니다. 실험적으로, 10억 개에서 110억 개의 파라미터를 가진 MoE 모델을 사전 학습하면서, 이러한 정렬이 더욱 효과적인 MoE 모델을 구축하는 데 기여한다는 것을 확인했습니다.

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.