Как работает MoE (Mixture of Experts) внутри LLM (Mixtral, GPT-4)?

Краткий тезис

Mixture of Experts (MoE) — это архитектурный приём, при котором вместо одной большой Feed-Forward Network (FFN) в каждом слое трансформера используется несколько «экспертов» (отдельных FFN), а специальный router (gating network) для каждого токена выбирает top-k экспертов и комбинирует их выходы. Это позволяет иметь огромное количество параметров (например, 8×7B = 56B у Mixtral), но при инференсе активировать лишь часть (2 эксперта), что даёт вычислительную стоимость (FLOPs) как у модели 12–14B. Такой подход называется sparse computation — разреженные вычисления.

1. Термин: MoE (Mixture of Experts)

MoE — это метод ансамблирования нейронных сетей, где несколько подсетей («экспертов») специализируются на разных подзадачах, а gating network (роутер) решает, кого из экспертов вызвать для каждого входного примера. В контексте LLM MoE применяется внутри слоёв трансформера, заменяя плотный FFN-слой на набор разреженно активируемых FFN.

Ключевая идея: увеличить число параметров без пропорционального роста вычислительных затрат. Модель становится «шире» (больше параметров), но не «глубже» (FLOPs почти не растут).

2. Архитектура: эксперты (Feed-Forward Networks)

В стандартном трансформере каждый слой состоит из self-attention и FFN (обычно два линейных преобразования с нелинейностью, например, ReLU или SwiGLU). В MoE-слое вместо одного FFN ставится N экспертов — каждый эксперт — это полноценный FFN со своей матрицей весов.

Пример: Mixtral 8x7B использует 8 экспертов, каждый размером ~7B параметров (итого 56B), но активируется только 2.

Структура эксперта

expert_i(x) = W2_i * activation(W1_i * x + b1_i) + b2_i

где W1_i, W2_i — веса i-го эксперта.

3. Router (gating network)

Router — это небольшой обучаемый линейный слой, который для каждого токена вычисляет логиты (score) для каждого эксперта. Затем применяется softmax (или top-k softmax) для получения вероятностей.

Формула router

scores = softmax(router_weights * x)

где router_weights — матрица размера (d_model, N), N — число экспертов.

Top-k routing выбираются только k экспертов с наибольшими scores (обычно k=2). Остальным экспертам присваивается нулевой вес. Это и даёт разреженность.

Пример кода (PyTorch-like):

import torch
import torch.nn.functional as F

class MoELayer(torch.nn.Module):
    def __init__(self, d_model, num_experts, k=2):
        super().__init__()
        self.num_experts = num_experts
        self.k = k
        self.router = torch.nn.Linear(d_model, num_experts, bias=False)
        self.experts = torch.nn.ModuleList([
            torch.nn.Sequential(
                torch.nn.Linear(d_model, 4*d_model),
                torch.nn.GELU(),
                torch.nn.Linear(4*d_model, d_model)
            ) for _ in range(num_experts)
        ])

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        batch, seq, d = x.shape
        x_flat = x.view(-1, d)  # (batch*seq, d)
        router_logits = self.router(x_flat)  # (batch*seq, num_experts)
        router_probs = F.softmax(router_logits, dim=-1)
        topk_vals, topk_idx = torch.topk(router_probs, self.k, dim=-1)
        # topk_idx: (batch*seq, k) — индексы выбранных экспертов
        # topk_vals: (batch*seq, k) — их веса

        # Собираем выходы экспертов
        final_output = torch.zeros_like(x_flat)
        for i in range(self.num_experts):
            mask = (topk_idx == i).any(dim=-1)  # токены, выбравшие эксперта i
            if mask.any():
                expert_out = self.experts[i](x_flat[mask])
                # Вес для эксперта i: topk_vals[mask, позиция i в topk]
                # Упрощённо: берём вес из topk_vals
                weights = topk_vals[mask][topk_idx[mask] == i]  # нужно аккуратно
                final_output[mask] += expert_out * weights.unsqueeze(-1)
        return final_output.view(batch, seq, d)

4. Top-k routing и sparse computation

Sparse computation означает, что для каждого токена вычисляется только малая часть экспертов (k из N). Остальные эксперты не участвуют в прямом проходе, их веса не загружаются в память для этого токена. Это даёт:

• Больше параметров (capacity) — модель может запомнить больше фактов.
• Меньше FLOPs — фактически выполняется k раз FFN, а не N.
• Эффективное использование GPU — загрузка только активных экспертов.

Сравнение dense vs MoE

5. Load balancing и auxiliary loss

Проблема: router может начать отправлять все токены к одним и тем же экспертам, перегружая их и недоиспользуя остальных. Для равномерной загрузки вводят auxiliary loss (вспомогательную функцию потерь).

Популярный подход — Importance + Load loss (из Switch Transformer):

• Importance loss штрафует за дисбаланс суммы вероятностей по экспертам.
• Load loss штрафует за дисбаланс доли токенов, назначенных каждому эксперту.

Формула (упрощённо):

aux_loss = alpha * (CV(importance)^2 + CV(load)^2)

где CV — коэффициент вариации, alpha — гиперпараметр (обычно 0.01).

Также применяют z-loss (стабилизация router) и expert choice routing (где эксперты выбирают токены, а не наоборот).

6. Преимущества MoE

• Масштабирование параметров без роста FLOPs — можно увеличить число экспертов до сотен.
• Специализация экспертов — разные эксперты обучаются обрабатывать разные типы токенов (синтаксис, семантика, редкие слова).
• Эффективность инференса — для каждого токена активируется лишь k экспертов, что снижает latency.
• Возможность распределения — эксперты могут быть размещены на разных GPU (model parallelism).

7. Примеры: Mixtral 8x7B и GPT-4

Mixtral 8x7B (Mistral AI):

• 8 экспертов, каждый ~7B параметров (итого 56B).
• Активируется 2 эксперта на токен.
• FLOPs соответствуют модели ~12B.
• Превосходит Llama 2 70B по качеству при меньшей стоимости инференса.

GPT-4 (слухи и косвенные данные):

• По неподтверждённой информации, GPT-4 использует MoE с 16 экспертами, каждый ~111B, активируется 2.
• Это объясняет, почему GPT-4 имеет огромное количество параметров (1.76T), но работает с задержкой, сравнимой с меньшими моделями.
• Microsoft и OpenAI патентуют технологии MoE для распределённого инференса.

Другие модели Switch Transformer (Google, 1.6T параметров, 2048 экспертов), GLaM (Google, 1.2T), DeepSeek-MoE.

8. Ограничения и проблемы MoE

• Неравномерная загрузка — требуется тщательная настройка auxiliary loss.
• Коммуникационные издержки — при распределении экспертов по GPU нужна пересылка данных (all-to-all).
• Сложность обучения — нестабильность, чувствительность к learning rate.
• Память — хотя FLOPs малы, все параметры экспертов должны храниться в памяти (или на диске с offloading).
• Инференс на маленьких батчах — overhead router может быть заметен.

9. Сравнение с dense model (таблица)

10. Обучение и инференс MoE

Обучение

• Стандартный backpropagation, но с учётом auxiliary loss.
• Используют mixed precision (FP16/BF16) для экономии памяти.
• Для больших MoE применяют expert parallelism — каждый эксперт на отдельном GPU, router на всех.

Инференс

• Для каждого токена router вычисляет top-k экспертов.
• Загружаются только веса этих экспертов (если они не в памяти, то с диска).
• Выходы экспертов взвешиваются и суммируются.
• Batch inference токены в батче могут требовать разных экспертов, что усложняет эффективную реализацию (нужна динамическая маршрутизация).

Пет-проект для закрепления

Задача Реализовать простой MoE-слой на PyTorch и обучить его на задаче классификации текстов (например, IMDB).

Инструменты PyTorch, Hugging Face Datasets, Weights & Biases для логирования.

Шаги:

1. Загрузить датасет IMDB, токенизировать (BERT tokenizer).
2. Определить MoE-слой с 4 экспертами, k=2, d_model=128.
3. Построить простую модель: Embedding → MoE → Linear → классификация.
4. Обучить с auxiliary loss (коэффициент 0.01).
5. Сравнить с dense-версией (один FFN той же размерности) по accuracy и FLOPs.
6. Визуализировать распределение токенов по экспертам (гистограмма).

Ожидаемый результат MoE-модель покажет схожее качество при меньшем числе FLOPs, а гистограмма покажет равномерное распределение (благодаря auxiliary loss).

Связь с другими вопросами (из списка 736)

Навигация

• Предыдущий: 281
• Следующий: 283
• Индекс: 00. Индекс разборов