Mixture of Experts (MoE): Глубокий разбор архитектуры

Введение

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

Как работает MoE

Базовая структура

                    Input Token
                         │
                    ┌────┴────┐
                    │  Router │
                    │  (Gating│
                    └────┬────┘
                         │
              ┌──────────┼──────────┐
              ↓          ↓          ↓
        ┌──────────┐ ┌──────────┐ ┌──────────┐
        │ Expert 1 │ │ Expert 2 │ │ Expert 3 │
        │ (FFN-a)  │ │ (FFN-b)  │ │ (FFN-c)  │
        └────┬─────┘ └────┬─────┘ └────┬─────┘
             │             │             │
             └─────────────┼─────────────┘
                           │
                    ┌──────┴──────┐
                    │  Combine    │
                    │  (weighted) │
                    └──────┬──────┘
                           │
                      Output Token

Математика gating

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

class MoELayer(nn.Module):
    """Базовый MoE слой"""
    
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.d_model = d_model
        
        # Router (gating network)
        self.gate = nn.Linear(d_model, num_experts)
        
        # Experts (identical architecture, different weights)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.GELU(),
                nn.Linear(d_model * 4, d_model)
            )
            for _ in range(num_experts)
        ])
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (batch, seq_len, d_model)
        """
        batch, seq_len, d_model = x.shape
        
        # Compute gating weights
        logits = self.gate(x)  # (batch, seq_len, num_experts)
        gates = F.softmax(logits, dim=-1)
        
        # Select top-k experts
        gate_values, expert_indices = torch.topk(gates, self.top_k, dim=-1)
        gate_values = F.softmax(gate_values, dim=-1)  # Renormalize top-k
        
        # Initialize output
        output = torch.zeros_like(x)
        
        # Route tokens to experts and compute
        for i, expert in enumerate(self.experts):
            # Find tokens routed to this expert
            mask = expert_indices == i  # (batch, seq_len, top_k)
            if mask.any():
                # Get the gate values for this expert
                expert_gates = torch.where(
                    mask, gate_values, 
                    torch.zeros_like(gate_values)
                )
                
                # Compute expert output
                expert_output = expert(x)  # (batch, seq_len, d_model)
                
                # Weighted sum
                output += expert_output * expert_gates.unsqueeze(-1)
        
        return output

Типы MoE

1. Standard MoE (Dense Gating)

class StandardMoE(nn.Module):
    """Стандартный MoE с плотным gating"""
    
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.GELU(),
                nn.Linear(d_model * 4, d_model)
            )
            for _ in range(num_experts)
        ])
    
    def forward(self, x):
        logits = self.gate(x)
        gates = F.softmax(logits, dim=-1)
        
        # Top-k selection
        top_gates, top_indices = torch.topk(gates, self.top_k, dim=-1)
        top_gates = F.softmax(top_gates, dim=-1)
        
        # One-hot encoding for efficient computation
        one_hot = F.one_hot(top_indices, num_classes=self.num_experts)
        one_hot = one_hot.permute(0, 2, 1).float()
        
        combined_gates = torch.bmm(one_hot, top_gates.unsqueeze(-1)).squeeze(-1)
        
        expert_outputs = torch.stack(
            [expert(x) for expert in self.experts], dim=1
        )
        
        output = torch.sum(
            expert_outputs * combined_gates.unsqueeze(-1),
            dim=1
        )
        
        return output

2. Sparse MoE с load balancing

class SparseMoE(nn.Module):
    """Sparse MoE с жёстким ограничением на нагрузку"""
    
    def __init__(self, d_model: int, num_experts: int, top_k: int = 1):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.gate = nn.Linear(d_model, num_experts)
        self.aux_loss_fn = AuxiliaryLoadBalancingLoss()
        
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.GELU(),
                nn.Linear(d_model * 4, d_model)
            )
            for _ in range(num_experts)
        ])
    
    def forward(self, x):
        batch, seq_len, d_model = x.shape
        logits = self.gate(x)
        gate_probs = F.softmax(logits, dim=-1)
        
        top_k_probs, top_k_indices = torch.topk(gate_probs, self.top_k, dim=-1)
        aux_loss = self.aux_loss_fn(gate_probs, top_k_indices)
        
        output = torch.zeros_like(x)
        for expert_idx, expert in enumerate(self.experts):
            expert_mask = (top_k_indices == expert_idx)
            if expert_mask.any():
                expert_out = expert(x)
                weights = top_k_probs[expert_mask].unsqueeze(-1)
                output[expert_mask] += expert_out[expert_mask] * weights[expert_mask]
        
        return output, aux_loss

3. GShard MoE

class GShardMoE(nn.Module):
    """GShard: MoE с capacity factor и balanced loss"""
    
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2, 
                 capacity_factor: float = 1.0):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts)
        
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.GELU(),
                nn.Linear(d_model * 4, d_model)
            )
            for _ in range(num_experts)
        ])
        
        self.aux_loss_fn = GShardAuxLoss()
    
    def forward(self, x):
        batch, seq_len, d_model = x.shape
        logits = self.gate(x)
        gate_probs = F.softmax(logits, dim=-1)
        
        # Top-k selection
        top_k_probs, top_k_indices = torch.topk(gate_probs, self.top_k, dim=-1)
        
        # Capacity: max tokens per expert
        capacity = int((batch * seq_len * self.top_k) / self.num_experts * self.capacity_factor)
        
        # Compute auxiliary loss for load balancing
        aux_loss = self.aux_loss_fn(gate_probs, top_k_indices)
        
        # Route tokens to experts with capacity constraint
        output = torch.zeros_like(x)
        importance = gate_probs.gather(-1, top_k_indices).max(dim=-1).values
        
        for expert_idx, expert in enumerate(self.experts):
            expert_mask = (top_k_indices == expert_idx)
            if expert_mask.any():
                # Apply capacity constraint
                expert_tokens = expert_mask.sum(dim=-1)
                if expert_tokens > capacity:
                    # Drop least important tokens
                    pass
                
                expert_out = expert(x)
                weights = top_k_probs[expert_mask]
                output[expert_mask] += expert_out[expert_mask] * weights.unsqueeze(-1)
        
        return output, aux_loss

Load Balancing

Auxiliary Loss для балансировки

class AuxiliaryLoadBalancingLoss(nn.Module):
    """Auxiliary loss для балансировки нагрузки между экспертами"""
    
    def __init__(self):
        super().__init__()
    
    def forward(self, gate_probs: torch.Tensor, top_k_indices: torch.Tensor) -> torch.Tensor:
        """
        gate_probs: (batch, seq_len, num_experts) - gating probabilities
        top_k_indices: (batch, seq_len, top_k) - selected expert indices
        """
        num_experts = gate_probs.shape[-1]
        batch, seq_len = gate_probs.shape[:2]
        
        # Fraction of tokens routed to each expert
        # Create one-hot encoding
        one_hot = F.one_hot(top_k_indices, num_classes=num_experts)  # (batch, seq_len, top_k, num_experts)
        token_fraction = one_hot.float().mean(dim=(0, 1))  # (num_experts, top_k)
        token_fraction = token_fraction.mean(dim=1)  # Average over top-k positions
        
        # Fraction of probability mass for each expert
        prob_fraction = gate_probs.mean(dim=(0, 1))  # (num_experts,)
        
        # Ideal: uniform distribution
        ideal = 1.0 / num_experts
        
        # Load balancing loss
        load_loss = torch.mean((prob_fraction - ideal) ** 2)
        
        # Token fraction loss (encourage uniform token distribution)
        token_loss = torch.mean((token_fraction - ideal) ** 2)
        
        return load_loss + token_loss

MoE в реальных моделях

Switch Transformer

class SwitchTransformerBlock(nn.Module):
    """Switch Transformer block (MoE с top-1 gating)"""
    
    def __init__(self, d_model: int, d_ff: int, num_experts: int, 
                 num_heads: int, dropout: float = 0.1):
        super().__init__()
        
        # Self-attention (dense)
        self.self_attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout)
        self.norm1 = nn.LayerNorm(d_model)
        
        # Switch FFN (sparse)
        self.switch_ffn = SwitchFFN(d_model, d_ff, num_experts)
        self.norm2 = nn.LayerNorm(d_model)
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        # Self-attention
        attn_output, _ = self.self_attn(x, x, x)
        x = x + self.dropout(attn_output)
        x = self.norm1(x)
        
        # Switch FFN
        switch_output, aux_loss = self.switch_ffn(x)
        x = x + self.dropout(switch_output)
        x = self.norm2(x)
        
        return x, aux_loss

class SwitchFFN(nn.Module):
    """Switch FFN: один эксперт на токен (top-1)"""
    
    def __init__(self, d_model: int, d_ff: int, num_experts: int):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model)
            )
            for _ in range(num_experts)
        ])
    
    def forward(self, x):
        logits = self.gate(x)  # (batch, seq_len, num_experts)
        probs = F.softmax(logits, dim=-1)
        
        # Top-1 selection
        indices = probs.argmax(dim=-1)  # (batch, seq_len)
        
        output = torch.zeros_like(x)
        batch, seq_len, d_model = x.shape
        
        for expert_idx, expert in enumerate(self.experts):
            mask = (indices == expert_idx)  # (batch, seq_len)
            if mask.any():
                output[mask] = expert(x[mask])
        
        return output

Mixtral 8x7B

class MixtralMoE(nn.Module):
    """Mixtral 8x7B: MoE с Dense Top-2 gating"""
    
    def __init__(self, d_model: int, num_experts: int = 8, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.SiLU(),
                nn.Linear(d_model * 4, d_model)
            )
            for _ in range(num_experts)
        ])
    
    def forward(self, x):
        logits = self.gate(x)  # (batch, seq_len, num_experts)
        
        # Gating with softmax
        scores = F.softmax(logits, dim=-1)
        
        # Top-2 gating
        top_scores, top_indices = torch.topk(scores, self.top_k, dim=-1)
        
        # Normalize top-2 scores
        top_scores = F.softmax(top_scores, dim=-1)
        
        # Compute output
        output = torch.zeros_like(x)
        
        for expert_idx, expert in enumerate(self.experts):
            # Find which tokens use this expert
            mask = (top_indices == expert_idx)  # (batch, seq_len, top_k)
            
            if mask.any():
                # Get the corresponding scores
                expert_scores = torch.where(mask, top_scores, torch.zeros_like(top_scores))
                # Sum over top-k dimension
                expert_weight = expert_scores.sum(dim=-1, keepdim=True)  # (batch, seq_len, 1)
                
                # Compute expert output
                expert_out = expert(x)
                
                # Add weighted contribution
                output += expert_out * expert_weight
        
        return output

Сравнение производительности

Модель	Параметры	Активные параметры	Experts	Top-K	Speedup
Dense Llama 3 8B	8B	8B	-	-	1x
Switch Transformer 1B	1B	1B	128	1	~1x
Mixtral 8x7B	467B	12.9B	8	2	~3x
Grok-1	335B	34B	64	4	~5x
Qwen2.5-MoE	32B	3.7B	64	2	~4x

Преимущества и недостатки

Преимущества

1. Масштабируемость — можно создавать модели с триллионами параметров
2. Эффективность инференса — активна лишь малая часть параметров
3. Параллелизм — эксперты легко распараллеливаются
4. Специализация — эксперты учатся специализироваться на разных задачах

Недостатки

1. Сложность обучения — проблемы с балансировкой нагрузки
2. Communication overhead — всеэксперты должны быть доступны
3. Непредсказуемость — разные токены используют разных экспертов
4. Деградация — при плохой балансировке качество падает

Best Practices

1. Используйте auxiliary loss для балансировки нагрузки
2. Capacity factor > 1.0 — дайте запас для предотвращения drop токенов
3. Top-2 gating — лучше Top-1 для стабильности
4. Начинайте с малого числа экспертов — 4-8, затем масштабируйте
5. Monitor expert utilization — следите за равномерностью загрузки

Заключение

MoE архитектура — это ключевой прорыв для создания больших эффективных моделей. Mixtral 8x7B и Qwen2.5-MoE показывают, что MoE позволяет получить преимущества больших моделей при затратах маленьких. Правильная реализация load balancing — ключ к успешному использованию MoE.