insights/architecture.md

DeepSeek MoE Architecture

This document provides a detailed explanation of the DeepSeek Mixture of Experts (MoE) architecture and what makes it unique compared to other MoE implementations.

Core Concepts of MoE

At a high level, Mixture of Experts (MoE) is a neural network architecture that divides computation across specialized "expert" networks. Rather than passing all inputs through the entire network, MoE selectively activates only a subset of experts for each input. This approach enables scaling models to have significantly more parameters while maintaining reasonable computational costs, as only a fraction of the network is active for any given input.

DeepSeek MoE Architecture: Key Innovations

DeepSeek's MoE implementation has several key innovations that distinguish it from previous MoE approaches such as GShard and Switch Transformers:

1. Hybrid Expert Structure: Shared + Routed Experts

One of the most distinctive features of DeepSeek MoE is its hybrid architecture that combines:

• Shared Experts: Process all tokens, providing a baseline transformation
• Routed Experts: Process only specific tokens they specialize in

The feed-forward network (FFN) output for token t is calculated as:

$${\widehat{h}}_t=u_t+\sum _i^{N_s}FF{N}_t^s(u_t)+\sum _i^{N_r}{g}_i(t)FF{N}_i^r(u_t)\backslash hat\{h\}\_t\; =\; u\_t\; +\; \backslash sum^\{N\_s\}\_i\; FFN^s\_t\; (u\_t)\; +\; \backslash sum^\{N\_r\}\_i\; g\_i(t)\; FFN^r\_i\; (u\_t)$$

Where:

• $u_t$ is the original input token representation
• $FFN^s_i$ is the i-th shared expert
• $FFN^r_i$ is the i-th routed expert
• $g_i(t)$ is the gate value determining how much each routed expert contributes

This hybrid approach has several advantages:

• Shared experts maintain global information flow
• Routed experts can specialize in specific patterns
• Residual connection (u_t term) preserves the original token information
• Reduces knowledge redundancy among experts

2. Token-Expert Affinity Calculation

The router determines which experts should process each token using a similarity-based mechanism:

$$s_{i,t}={Softmax}_i(u_t^T{e}_i)s\_\{i,t\}\; =\; \{Softmax\}\_i(u\_t^T\; e\_i)$$

Where:

• $s_{i,t}$ is the token-to-expert affinity
• $e_i$ is the centroid of the i-th routed expert
• $u_t$ is the token representation

This is conceptually similar to the attention mechanism's query-key dot product (QK^T). It measures the similarity between a token vector and expert centroids:

• Similar vectors (token and expert specialty) → large dot product
• Different vectors → small dot product
• Softmax converts these similarities into a probability distribution

The router then selects the top-K experts for each token:

g_i,t = {
  s_i,t, if s_i,t ∈ Topk({s_j,t|1 ≤ j ≤ N_r}, K_r),
  0, otherwise
}

This approach combines soft routing (through the affinity scores) and hard routing (through the TopK selection), allowing for more nuanced expert specialization.

3. Multi-Level Load Balancing

DeepSeek MoE implements a cascading auxiliary loss structure to ensure balance at three different levels:

Expert-Level Balance Loss

$${\mathcal{L}}_{ExpBal}={\alpha }_1\sum (f_i{P}_i)\backslash mathcal\{L\}\_\{ExpBal\}\; =\; \backslash alpha\_1\; \backslash sum(f\_i\; P\_i)$$

Where:

• $f_i$ is the fraction of tokens routed to expert i
• $P_i$ is the average routing probability for expert i

This prevents "expert collapse" where only a few experts get consistently used.

Device-Level Balance Loss

$${\mathcal{L}}_{DevBal}={\alpha }_2\sum (f_i^{\mathrm{\prime }}{P}_i^{\mathrm{\prime }})\backslash mathcal\{L\}\_\{DevBal\}\; =\; \backslash alpha\_2\; \backslash sum(f\text{'}\_i\; P\text{'}\_i)$$ Where:

• $f'_i$ is the average fraction of tokens routed to experts on device i
• $P'_i$ is the sum of routing probabilities for experts on device i

This ensures computation is evenly distributed across hardware devices.

Communication Balance Loss

$${\mathcal{L}}_{CommBal}={\alpha }_3\sum (f_i^{\mathrm{\prime }\mathrm{\prime }}{P}_i^{\mathrm{\prime }\mathrm{\prime }})\backslash mathcal\{L\}\_\{CommBal\}\; =\; \backslash alpha\_3\; \backslash sum(f\text{'}\text{'}\_i\; P\text{'}\text{'}\_i)$$

Where:

• $f''_i$ measures the fraction of tokens sent to device i
• $P''_i$ is the sum of routing probabilities for experts on device i

This manages network traffic patterns between devices, which is critical for distributed training.

The multi-level approach is particularly effective because imbalance at any level causes inefficiency:

• Expert imbalance → wasted model capacity
• Device imbalance → some hardware sits idle
• Communication imbalance → network congestion

4. Device-Limited Routing

For distributed training, DeepSeek MoE implements a device-limited routing mechanism that bounds communication costs:

1. For each token, select M devices that have experts with the highest affinity scores
2. Perform top-K selection only among experts on these M devices

This approach ensures that each token's computation is limited to a manageable number of devices, reducing cross-device communication overhead. Empirically, setting M ≈ 3 achieves performance comparable to unrestricted routing.

5. Token-Dropping Strategy

To further optimize computation, DeepSeek MoE implements a device-level token-dropping strategy:

1. Compute the average computational budget for each device (capacity factor = 1.0)
2. Drop tokens with the lowest affinity scores on each device until reaching the budget
3. Ensure tokens from approximately 10% of training sequences are never dropped

This approach provides flexibility to adjust computation vs. quality tradeoffs during inference while maintaining consistency between training and inference.

Sequence Understanding with Token-Level Routing

Despite routing happening at the token level, DeepSeek MoE maintains sequence understanding through several mechanisms:

1. Self-attention layers before and after the MoE layer process the whole sequence together
2. The residual connection preserves the original token information
3. Shared experts process all tokens, providing a base transformation
4. Layer normalization helps integrate the different expert contributions

This design allows each token to get specialized processing from relevant experts while the attention layers ensure these individually-processed tokens still work together to understand the sequence as a whole.

Comparison with Other MoE Implementations

vs. Switch Transformers

• Routing Granularity: Switch routes each token to exactly one expert; DeepSeek routes to multiple experts (top-K)
• Expert Structure: Switch uses standard FFNs; DeepSeek uses both shared and routed experts
• Load Balancing: DeepSeek uses a more sophisticated multi-level balancing approach

vs. GShard

• Expert Specialization: DeepSeek uses finer granularity for better specialization
• Knowledge Sharing: DeepSeek's shared experts reduce redundancy
• Load Balancing: DeepSeek's cascade of balance losses provides more robust load distribution
• Token Handling: DeepSeek uses a simplified but effective token-dropping strategy

Integration in the Transformer Architecture

DeepSeek MoE layers replace the standard feed-forward networks in transformer blocks, while keeping the attention mechanism intact:

Transformer Block
  ├── RMS Norm
  ├── Attention
  ├── Residual Connection
  ├── RMS Norm
  ├── DeepSeekMoE Layer
  │   ├─┬─ Shared Experts (process all tokens)
  │   │ │
  │   │ ├─ Router → Top-K Selection
  │   │ │
  │   │ └─ Routed Experts (process tokens via routing)
  │   │
  │   └── Combine outputs (residual + shared + routed)
  └── Residual Connection

Conclusion

The DeepSeek MoE architecture represents a sophisticated approach to building large-scale language models that balance parameter count and computational efficiency. By using a hybrid expert structure, intelligent routing, and multi-level load balancing, DeepSeek MoE achieves better performance than previous MoE implementations with the same computational budget.

The design reflects careful consideration of both theoretical aspects (how experts specialize and share knowledge) and practical engineering challenges (distributed training efficiency, communication patterns). This makes DeepSeek MoE not just an academic advancement but a practical approach for deploying large language models efficiently.