Классическая смесь экспертов (Mixture of Experts, MoE)

Классическая смесь экспертов (MoE, Mixture of Experts) — это архитектурный подход в машинном обучении, в котором несколько моделей (экспертов) обучаются на разные подзадачи, а специальная сеть маршрутизации (gating network) определяет, какие эксперты должны быть активированы для обработки конкретного входа.

Основные компоненты MoE:

1. Эксперты (Experts):

   • Это обычно нейронные сети или простые модели.

   • Каждый эксперт обучается на подмножестве данных или на специализированной задаче.

   • Все эксперты получают один и тот же вход, но только некоторые из них используются на каждом шаге.

2. Сеть маршрутизации (Gating Network):

   • Выдает веса (обычно через softmax) для каждого эксперта на основе входа.

   • Эти веса определяют вклад каждого эксперта в финальный ответ.

   • Может быть:

     • Soft routing: используется несколько экспертов одновременно, с разными весами.

     • Hard routing: выбирается фиксированное число лучших экспертов (например, топ-1 или топ-2).

3. Агрегация выходов:

   • Ответы активированных экспертов комбинируются в итоговый выход модели, обычно как взвешенная сумма.

Классическая формализация:

Пусть:

• $x$ — вход,

• $E_i(x)$ — выход i-го эксперта,

• $G(x)_i$ — вес, выданный gater'ом для эксперта $i$.

Тогда итоговый выход модели:

$$\mathrm{<span\; style="font-size:16px;"><em>y</em></span>}<span\; style="font-size:16px;"><em>=</em></span>\underset{\mathrm{<span\; style="font-size:16px;"><em>i</em></span>}<span\; style="font-size:16px;"><em>=</em></span>\mathrm{<span\; style="font-size:16px;"><em>1</em></span>}}{\overset{\mathrm{<span\; style="font-size:16px;"><em>N</em></span>}}{<span\; style="font-size:16px;"><em>\sum </em></span>}}\mathrm{<span\; style="font-size:16px;"><em>G</em></span>}<span\; style="font-size:16px;"><em>(</em></span>\mathrm{<span\; style="font-size:16px;"><em>x</em></span>}{<span\; style="font-size:16px;"><em>)</em></span>}_{\mathrm{<span\; style="font-size:16px;"><em>i</em></span>}}<span\; style="font-size:16px;"><em>\cdot </em></span>{\mathrm{<span\; style="font-size:16px;"><em>E</em></span>}}_{\mathrm{<span\; style="font-size:16px;"><em>i</em></span>}}<span\; style="font-size:16px;"><em>(</em></span>\mathrm{<span\; style="font-size:16px;"><em>x</em></span>}<span\; style="font-size:16px;"><em>)</em></span>$$

Если используется hard routing с top-k, то только k экспертов участвуют в суммировании, а веса нормализуются среди них.

Цель и мотивация:

• Масштабируемость: позволяет использовать очень большие модели, активируя лишь малую часть параметров при каждом вызове (sparse computation).

• Модулярность: каждый эксперт может специализироваться на своей подзадаче, улучшая общую эффективность.

• Универсальность: MoE может использоваться в NLP, CV и других областях.

Пример: Использование MoE в NLP

В трансформерах, как Switch Transformer или GShard, MoE применяется к feed-forward слоям. Например:

• Вместо одного FFN блока в трансформере, MoE блок содержит десятки экспертов.

• На каждый токен выбирается, скажем, top-2 эксперта.

• Это позволяет масштабировать трансформер до триллионов параметров без пропорционального увеличения вычислений.

Проблемы и вызовы:

1. Балансировка нагрузки:

   • Некоторые эксперты могут переиспользоваться, другие — не использоваться.

   • Используются специальные потери (auxiliary loss), чтобы равномерно распределять входы между экспертами.

2. Обратное распространение (Backpropagation):

   • Hard routing делает обучение недифференцируемым (иногда используют REINFORCE или другие подходы).

   • Soft routing более стабильный, но менее эффективный в вычислениях.

3. Коммуникационные издержки (в распределённых системах):

   • В больших MoE-моделях (например, в Google GShard) возникают сложности с маршрутизацией между GPU.

Историческая справка:

• Введён в 1991 году Jacobs et al. в работе "Adaptive Mixtures of Local Experts".

• Вновь получил популярность после работ Google:

  • GShard (2020),

  • Switch Transformer (2021),

  • Pathways и PaLM (2022) — масштаб до 540B параметров.

Заключение

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