The 2024 technical report “Mixtral of Experts” by Mistral AI introduced Mixtral 8x7B, a Sparse Mixture of Experts (SMO) model. While it has a total of 47B parameters, it only uses 13B active parameters during inference, allowing it to match or outperform much larger dense models like Llama 2 70B while being significantly faster.
1. The SMoE Architecture
In a standard Transformer, every token passes through the same Feed-Forward Network (FFN). In Mixtral, the dense FFN is replaced by a Mixture-of-Experts layer containing 8 “experts.” For every token at every layer, a Router selects the Top-2 experts to process it.
2. PyTorch Implementation: The MoE Sparse Layer
This is the core “Switch” that replaces the standard FeedForward block. It uses a gating network to distribute tokens to the correct experts.
Python
import torch
import torch.nn as nn
import torch.nn.functional as F
class MixtralSparseMoeBlock(nn.Module):
def __init__(self, hidden_dim, intermediate_dim, num_experts=8, top_k=2):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
# The Router (Gating Network)
self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
# The Experts: Each is a standard SwiGLU FFN
# In Mixtral, each expert has the same architecture as Mistral 7B's FFN
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(hidden_dim, intermediate_dim), # w1
nn.SiLU(), # Swish
nn.Linear(intermediate_dim, hidden_dim) # w2
) for _ in range(num_experts)
])
def forward(self, x):
# x: [batch_size, seq_len, hidden_dim]
batch_size, seq_len, hidden_dim = x.shape
x_flat = x.view(-1, hidden_dim) # [num_tokens, hidden_dim]
# 1. Get gating logits and top-k indices
logits = self.gate(x_flat) # [num_tokens, num_experts]
weights, selected_experts = torch.topk(logits, self.top_k, dim=-1)
# 2. Softmax over the selected experts only
weights = F.softmax(weights, dim=-1)
# 3. Process tokens through experts
# Note: In production, this is optimized via specialized CUDA kernels (e.g. Megablocks)
# to avoid looping and handle variable token counts per expert.
results = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts):
# Find tokens assigned to this expert
batch_idx, nth_expert = torch.where(selected_experts == i)
if batch_idx.numel() > 0:
# Get the specific tokens and their relative gating weights
tokens = x_flat[batch_idx]
expert_out = expert(tokens)
# Apply the weights and add to final result
# weights[batch_idx, nth_expert] is the softmax weight for this token/expert pair
results[batch_idx] += weights[batch_idx, nth_expert].unsqueeze(-1) * expert_out
return results.view(batch_size, seq_len, hidden_dim)
3. Key Technical Specifications
| Parameter | Value |
| Total Parameters | 46.7 Billion |
| Active Parameters | 12.9 Billion |
| Experts per Layer | 8 |
| Experts per Token | 2 |
| Context Length | 32k Tokens |
| Attention | Grouped-Query Attention (GQA) |
4. Implementation Nuance: Load Balancing
A common problem in MoE is that the router might “favor” one expert, causing it to overtrain while others stay idle. While the Mixtral report notes they didn’t need complex balancing, standard SMoE implementations often include an Auxiliary Loss to encourage equal expert utilization.
$$Loss_{aux} = \sum_{i=1}^N f_i \cdot P_i$$
- $f_i$: Fraction of tokens assigned to expert $i$.
- $P_i$: Average gating probability for expert $i$.
Why Mixtral is a 2024 Milestone
Mixtral proved that Expert Specialization is a viable way to scale model intelligence without scaling the power bill. By activating only a fraction of the weights per token, it achieves the reasoning capabilities of a 70B model with the latency of a 12B model. It also significantly outperforms Llama 2 70B on code and math, areas where specialized “experts” can thrive.