Что такое Mixture of Experts (MoE) в машинном обучении?

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

Ключевые компоненты:

1. Эксперты (Experts): Небольшие, обычно идентичные по архитектуре нейронные сети (например, полносвязные слои или FFN в трансформерах). Каждый эксперт обучается специализироваться на определенном типе данных.
2. Шлюз (Gating Network): Небольшая сеть, которая на основе входных данных вычисляет вероятностное распределение (веса) по всем экспертам. Часто используется разреженная (sparse) активация, например, Top-K Gating, где для каждого примера выбираются только K экспертов с наибольшими весами.

Упрощенный пример на PyTorch:

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim)
        )
    def forward(self, x):
        return self.net(x)

class SparseMoE(nn.Module):
    def __init__(self, input_dim, num_experts=4, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList([Expert(input_dim, input_dim*4) for _ in range(num_experts)])
        self.gate = nn.Linear(input_dim, num_experts)  # Шлюз
        self.top_k = top_k

    def forward(self, x):
        # 1. Вычисление весов шлюза
        gate_logits = self.gate(x)  # [batch_size, num_experts]
        # 2. Разреженный Top-K выбор
        top_k_weights, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)  # Выбираем K лучших
        top_k_weights = F.softmax(top_k_weights, dim=-1)  # Нормализуем веса выбранных экспертов

        # 3. Нулевая маска для невыбранных экспертов (разреженность)
        zeros = torch.zeros_like(gate_logits)
        sparse_weights = zeros.scatter(-1, top_k_indices, top_k_weights)

        # 4. Агрегация выходов выбранных экспертов
        expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1)  # [batch, num_experts, dim]
        final_output = torch.einsum('bnd,bn->bd', expert_outputs, sparse_weights)
        return final_output

Преимущества и применение:

• Эффективность: При огромном общем числе параметров (1+ трлн) для обработки токена активируются лишь миллиарды (sparse computation).
• Масштабируемость: Позволяет значительно увеличивать емкость модели без пропорционального роста вычислительных затрат на инференс.
• Известные реализации: Google's Switch Transformer, GLaM, GShard. Предполагается, что подобные архитектуры используются в крупнейших моделях, таких как GPT-4.