MoE layer

1. 1

   Et lag i et neuralt netværk der anvender en blanding af eksperter (sub-netværk) og en port-mekanisme til at dirigere input til en sparsomt aktiveret delmængde af eksperterne, hvilket muliggør skalerbarhed.

   • Mixtral 8x7B-modellen anvender MoE-lag med otte eksperter, hvor kun to aktiveres per token. — forskning, 2024
   • Switch Transformer introducerede en forenklet MoE-arkitektur hvor gate-mekanismen vælger én enkelt ekspert per input. — Switch Transformer, 2021

MoE-lag bruges primært i store sprogmodeller for at skalere parametertallet uden tilsvarende stigning i inferensomkostninger. De er centrale i arkitekturer som Mixtral 8x7B og Switch Transformer.

y = sum_{i=1}^n G(x)_i * E_i(x), where G(x) = softmax(W_g · x) and typically only top-k experts are selected (others set to zero).

class MoELayer(nn.Module):
    def __init__(self, d_model, num_experts, top_k):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts)
        self.experts = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(num_experts)])
        self.top_k = top_k

    def forward(self, x):
        gate_scores = F.softmax(self.gate(x), dim=-1)
        values, indices = torch.topk(gate_scores, self.top_k, dim=-1)
        values = values / (values.sum(dim=-1, keepdim=True) + 1e-8)
        out = torch.zeros_like(x)
        for i, expert in enumerate(self.experts):
            mask = (indices == i).any(dim=-1, keepdim=True).float()
            out += values[..., i].unsqueeze(-1) * expert(x) * mask
        return out

Forenklet implementation af et MoE-lag i PyTorch. Gate-mekanismen vælger top-k eksperter, og outputtet vægtes med de normaliserede gate-scores.

MoE står for 'Mixture of Experts', en idé der stammer fra statistisk modellering i 1990'erne.

• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• Mixtral of Experts