Mixture of Experts

Идея: разорвать связь между памятью и вычислениями

В обычной сети число параметров и объём вычислений на токен — это одно и то же число: каждый токен проходит через все веса. MoE эту связь разрывает.

$$y = \sum_{e \in \text{top-}k} \htmlData{k=gate}{g_e(x)}\, \htmlData{k=expert}{E_e(x)}, \qquad g(x) = \text{softmax}(W_r x)$$

Сумма идёт только по $k$ выбранным экспертам из $E$, поэтому на токен приходится $k/E$ доля параметров слоя. При $E = 8$, $k = 2$ это $25\%$: восемь наборов веса, работы — как от двух.

$E$	$k$	активная доля
$8$	$2$	$25\%$
$8$	$1$	$12.5\%$
$16$	$4$	$25\%$
$64$	$2$	$3.1\%$

Mixtral 8×7B — прямая иллюстрация: $46.7$ млрд параметров всего, $12.9$ млрд активны на токен. Память нужна под все, арифметика — под четверть.

Почему это трудно

Проблема не в формуле, а в том, что гейт обучается вместе со всем остальным, и у него есть вырожденное решение: отправлять всё нескольким экспертам. Оно выгодно на ранних шагах — эксперт, получивший больше токенов, обучается быстрее, становится лучше, и роутер отправляет ему ещё больше.

Это положительная обратная связь, и её отказ тихий: лосс продолжает падать, ошибок нет, просто большая часть параметров не используется, и модель по сути имеет размер активной части.

Покрутите разброс логитов роутера. Уверенный роутер — несбалансированный роутер.

эксперт токенов

вместимость: 20

экспертов на токен k

2

разброс логитов

1

capacity factor

1.25

активная доля

25%

перегрузка

×1.63

balance loss

1.239

отброшено

16

Часть токенов отброшена: эксперт исчерпал вместимость, и для этих токенов слой сработал как тождественное отображение — они прошли по residual-связи без обработки. Ошибки при этом не возникает.

Два числа под графиком стоит понимать точно.

Перегрузка — отношение максимальной загрузки к средней. При идеальном балансе единица; при $\times 3$ на восьми экспертах треть слоя простаивает, пока один считает за трёх.

Balance loss — вспомогательное слагаемое, добавляемое к основной цели:

$\mathcal{L}_{\text{aux}} = E \sum_{e=1}^{E} f_e \cdot P_e$

где $f_e$ — доля токенов, попавших эксперту $e$, а $P_e$ — его средний гейт. Устроено так, что при равномерном распределении величина равна единице (каждое слагаемое $\frac1E \cdot \frac1E$, сумма $\frac{1}{E}$, умноженная на $E$), и растёт при перекосе.

Почему в формуле произведение двух величин, а не одна: $f_e$ считается через argmax и по нему нельзя дифференцировать, а $P_e$ — гладкая функция параметров роутера. Произведение даёт градиент, который толкает роутер понижать гейт у перегруженных экспертов. Приём общий и стоит запомнить: недифференцируемую величину используют как вес при дифференцируемой.

Capacity factor и отброшенные токены

Каждому эксперту заранее выделяется буфер:

$\text{вместимость} = \frac{\text{токенов} \cdot k}{E} \cdot \text{capacity factor}$

Причина этого ограничения — не математика, а форма тензоров: чтобы считать батчем, буферы должны быть равными и известными заранее. Токены сверх вместимости отбрасываются — проходят слой по residual-связи без обработки.

Заметьте, что происходит: это не ошибка и не исключение. Для отброшенного токена MoE-слой работает как тождественное отображение, и обнаружить это можно только по счётчику. Ещё один тихий режим отказа, и второй в этом уроке.

Capacity factor $1.25$ означает запас в четверть. Больше — меньше потерь и больше памяти под буферы, меньше — наоборот.

Что ещё ломается

проблема	проявление	обычное решение
коллапс роутера	все токены к нескольким экспертам	balance loss
отброшенные токены	тихий пропуск слоя	capacity factor выше единицы
нестабильность обучения	расходимость на больших моделях	z-loss на логиты роутера
дисбаланс между устройствами	одно ждёт остальные	балансировка + capacity
переобучение экспертов	каждый видит мало данных	больше данных, меньше $E$

Строка про устройства стоит отдельного внимания, потому что она превращает статистику в инженерную проблему. Эксперты обычно распределены по устройствам, и раз шаг синхронный, все ждут самое загруженное. Значит перегрузка $\times 3$ — это не «треть простаивает в среднем», а шаг втрое дольше. Баланс тут нужен не для качества, а для скорости.

Что MoE даёт и чего не даёт

	обычная сеть	MoE
параметров	$N$	$E \cdot N_{\text{экс}}$
вычислений на токен	$\propto N$	$\propto k \cdot N_{\text{экс}}$
памяти под веса	$\propto N$	под все эксперты
обмен между устройствами	нет	all-to-all на каждом слое
качество при равных вычислениях	базовое	обычно выше
качество при равной памяти	базовое	обычно ниже

Две последние строки — самое полезное для практики. MoE выгодна, когда узкое место — вычисления, и невыгодна, когда память: держать восемь экспертов, чтобы использовать два, приходится целиком. Отсюда и типичная область применения: обучение и сервинг больших моделей на многих устройствах, где памяти суммарно достаточно, а времени всегда мало.

Итог

• MoE разрывает связь «параметры = вычисления»: активная доля равна $k/E$.
• Главная проблема — статистическая: роутер сходится к вырожденному решению, и отказ тихий.
• Balance loss равен единице при равномерном распределении и работает через произведение недифференцируемой доли на дифференцируемый гейт.
• Отброшенные токены — не ошибка, а пропуск слоя; capacity factor задаёт запас.
• Выгода в вычислениях, плата — памятью и обменом между устройствами.