Перепроектирование маршрутизаторов смеси экспертов с использованием степенной итерации на многообразии

Аннотация

Маршрутизатор является ключевым компонентом моделей смеси экспертов (MoE). Выступая в качестве прокси экспертов, строки матрицы маршрутизатора вычисляют свое сходство с входными данными MoE, чтобы определить, какое подмножество экспертов активируется. В идеале каждая строка маршрутизатора предназначена для кодирования матрицы эксперта в репрезентативный вектор, так чтобы ее скалярное произведение с токеном лучше отражало сродство токена и эксперта. Однако не существует принципов проектирования, обеспечивающих такое сжатие. В данной работе мы предлагаем выравнивать каждую строку маршрутизатора по главному сингулярному направлению соответствующего эксперта, поскольку это направление дает наиболее выразительное математическое описание матрицы. Основываясь на этом принципе, мы предлагаем перепроектирование маршрутизатора с помощью итерации степени на многообразии (Manifold Power Iteration, MPI). В частности, вводится парадигма «степенная итерация с последующей ретракцией» (Power-then-Retract), в которой сначала выполняется шаг степенной итерации для весов маршрутизатора, а затем ретракция для наложения ограничения нормы, что обеспечивает как эффективность, так и стабильность. Теоретически мы показываем, что MPI заставляет строки маршрутизатора сходиться к главным сингулярным направлениям соответствующих экспертов. Эмпирически мы предобучаем модель MoE в масштабах от 1B до 11B параметров, чтобы подтвердить, что такое выравнивание способствует созданию более эффективных моделей MoE.

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.