Today's Journey

The Scaling Dilemma

Larger models perform better... but at what cost?

Training a 175B parameter model (GPT-3 scale):

• Cost: $4-12 million in compute
• Energy: Equivalent to 120 homes for a year
• Time: Weeks to months on thousands of GPUs
• Inference: Slow and expensive ($0.002-0.02 per 1K tokens)
• Environmental: Massive carbon footprint

Key Insight: Not all parameters need to be active for every input!

What is Mixture of Experts?

Key Components:

1. Experts: Multiple specialized feed-forward networks
2. Router/Gate: Learned function that assigns inputs to experts
3. Sparse Activation: Only top-k experts process each input

Intuition:

• Different experts become good at different things
• Math expert, code expert, language expert, etc.
• Router learns which expert(s) to use for each input

Reference: Shazeer et al. (2017) - "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer"

MoE Architecture

Concrete Example: Processing "def factorial(n):"

1# Router computes scores for each expert
2router_logits = router(token_embedding)  # Shape: (8,) for 8 experts
3# [0.1, 0.9, 0.3, 0.2, 0.1, 0.1, 0.1, 0.1]  # Expert 2 (Code) scores highest!
4
5# Select top-2 experts
6top_k_indices = [1, 2]   # Expert 2 (Code) and Expert 3 (Language)
7top_k_weights = [0.75, 0.25]  # Normalized weights
8
9# Only these 2 experts process the token
10output = 0.75 * expert_2(x) + 0.25 * expert_3(x)
11# Other 6 experts do NO computation for this token!

Typical setup: N=8-64 experts, but only top-2 are active per token

The Router Mechanism

Step-by-step routing for a single token:

1import torch.nn.functional as F
2
3def route_token(x, router_weights, num_experts=8, k=2):
4    """Route a token to top-k experts"""
5    # Step 1: Compute routing scores (linear projection)
6    # router_weights: (hidden_dim, num_experts)
7    logits = x @ router_weights  # Shape: (num_experts,)
8    # Example: [-0.5, 2.1, 1.3, 0.2, -0.1, 0.0, -0.3, 0.1]
9
10    # Step 2: Convert to probabilities
11    probs = F.softmax(logits, dim=-1)
12    # Example: [0.04, 0.52, 0.24, 0.08, 0.03, 0.03, 0.03, 0.03]
13
14    # Step 3: Select top-k experts
15    top_k_probs, top_k_indices = torch.topk(probs, k)
16    # indices: [1, 2], probs: [0.52, 0.24]
17
18    # Step 4: Normalize selected probabilities
19    top_k_probs = top_k_probs / top_k_probs.sum()
20    # [0.68, 0.32]  # Now sums to 1

The Router Mechanism

21
22    return top_k_indices, top_k_probs
23    # Expert 1 handles 68%, Expert 2 handles 32%

MoE in PyTorch (Simplified)

21            for _ in range(num_experts)
22        ])
23
24    def forward(self, x):
25        # x: (batch_size, seq_len, d_model)
26        batch_size, seq_len, d_model = x.shape
27
28        # Compute routing scores
29        router_logits = self.gate(x)  # (batch, seq, num_experts)
30        router_probs = F.softmax(router_logits, dim=-1)
31
32        # Select top-k experts
33        top_k_probs, top_k_indices = torch.topk(router_probs, self.k, dim=-1)
34        # top_k_probs: (batch, seq, k)
35        # top_k_indices: (batch, seq, k)

MoE Forward Pass (cont.)

1# Initialize output
2        output = torch.zeros_like(x)
3
4        # Route to experts
5        for i in range(self.k):
6            # Get expert indices for this position
7            expert_idx = top_k_indices[:, :, i]  # (batch, seq)
8            expert_weight = top_k_probs[:, :, i]  # (batch, seq)
9
10            # Process through each expert
11            for expert_id in range(self.num_experts):
12                # Mask for tokens routed to this expert
13                mask = (expert_idx == expert_id)
14
15                if mask.any():
16                    # Get tokens for this expert
17                    expert_input = x[mask]
18
19                    # Process through expert
20                    expert_output = self.experts[expert_id](expert_input)

MoE Forward Pass (cont.)

21
22                    # Add weighted output
23                    output[mask] += expert_weight[mask].unsqueeze(-1) * expert_output
24
25        return output

Note: This is simplified. Production implementations handle batching and load balancing more efficiently.

Load Balancing Problem

Problem: Without balancing, some experts get overused!

1Expert Usage Distribution (Unbalanced):
2
3E1: ████████████████████████████████ 40%  ← Overloaded!
4E2: ██████████████████               22%
5E3: ████████████                     15%
6E4: ████████                         10%
7E5: ████                              5%
8E6: ███                               4%
9E7: ██                                3%  ← Undertrained
10E8: █                                 1%  ← "Dead" expert

Desired (Balanced):

1E1: ████████████  12.5%
2E2: ████████████  12.5%
3E3: ████████████  12.5%
4... (all equal)

Why imbalance is bad:

• "Dead" experts waste parameters (never trained properly)
• Popular experts become bottlenecks during parallel inference
• Model loses expressiveness it paid for in memory

Load Balancing Solutions

Solution 1: Auxiliary Loss (Penalize Imbalance)

1def compute_load_balance_loss(router_probs, expert_assignments, alpha=0.01):
2    """Add penalty to main loss for imbalanced routing"""
3    num_experts = router_probs.shape[-1]
4
5    # f_i: fraction of tokens actually sent to expert i
6    tokens_per_expert = expert_assignments.sum(dim=0)  # Count per expert
7    f = tokens_per_expert / tokens_per_expert.sum()    # [0.4, 0.22, ...]
8
9    # P_i: average routing probability for expert i
10    P = router_probs.mean(dim=0)                       # [0.35, 0.2, ...]
11
12    # Loss: encourages f and P to both be uniform (1/N each)
13    aux_loss = alpha * num_experts * (f * P).sum()
14    return aux_loss  # Added to main training loss

Solution 2: Expert Capacity Limits

1capacity = (batch_size * seq_len) // num_experts * capacity_factor  # e.g., 1.25x
2# If Expert 1 already has `capacity` tokens, overflow goes to Expert 2

Training Instability

Challenges in training MoE:

1. Routing Collapse
   • All tokens routed to one or few experts

• Solution: Load balancing loss, initialization

2. Expert Imbalance
   • Some experts undertrained

• Solution: Balanced batching, capacity limits

3. High Variance Gradients
   • Discrete routing decisions

• Solution: Softmax gating, larger batch sizes

4. Dead Experts
   • Expert never gets activated

• Solution: Expert dropout, random routing

Best Practice: Start with dense model, then convert to MoE for fine-tuning

Memory and Communication

MoE Memory Requirements:

Example: Mixtral 8x7B

• 8 experts × 7B parameters each = 56B total
• But only 2 experts active → 14B active parameters
• Memory: Need to store all 56B parameters
• Compute: Only process 14B parameters per token

Solutions:

1. Expert Parallelism: Distribute experts across GPUs
2. Offloading: Keep inactive experts on CPU/disk
3. Quantization: Reduce precision (int8, int4)
4. Shared Experts: Some experts always active for all tokens

Mixtral 8x7B

Architecture:

• Total parameters: 47B (8 experts × 7B, minus shared layers)
• Active parameters: 13B (only 2 experts active per token)
• Layers: 32 transformer blocks
• Context: 32K tokens
• Vocabulary: 32K tokens (SentencePiece)

Training:

• Open weights (Apache 2.0 license)
• Multilingual (English, French, German, Spanish, Italian)
• Code-capable

Reference: Jiang et al. (2024) - "Mixtral of Experts"

What Do Experts Learn?

Experts naturally specialize without explicit supervision!

Analyzed routing patterns in Mixtral:

1Token Type              → Most Active Experts
2─────────────────────────────────────────────
3"def", "class", "import" → Expert 2 (Code)
4"la", "le", "français"   → Expert 5 (French)
5"∑", "∫", "theorem"      → Expert 7 (Math)
6"the", "is", "and"       → Expert 1 (Common words)
7"neural", "gradient"     → Expert 3 (Technical)

Concrete Example: Sentence routing

1"The neural network learns via backpropagation"
2 │    │      │        │     │    │
3 E1   E3     E3       E1    E1   E3
4
5Different tokens in the same sentence use different experts!

State Space Models: Mamba

Alternative to Transformers with linear-time complexity

1Sequence length:   1K   4K   16K   64K
2Compute (relative): 1   16   256   4096
3                         ↑
4              Gets expensive fast!

1Sequence length:   1K   4K   16K   64K
2Compute (relative): 1    4    16    64
3                              ↑
4              Linear scaling!

1# State evolves with each token
2state = initial_state  # Fixed size!
3for token in sequence:
4    # Update state (no attention)
5    state = A @ state + B @ token
6    output = C @ state
7    # Always O(1) per token!

Reference: Gu & Dao (2023) - "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"

Environmental Impact

The carbon cost of large language models:

Factors affecting carbon footprint:

• Model size (more parameters = more compute)
• Training time
• Hardware efficiency (newer GPUs more efficient)
• Energy source (coal vs solar)
• Location of data center

Source: Strubell et al. (2019) - "Energy and Policy Considerations for Deep Learning in NLP"

Democratizing Access

Should powerful AI be accessible to everyone, or only large organizations?

Current reality:

• Training GPT-3 scale model: $4-12M
• Requires 1000+ GPUs
• Only big tech companies can afford
• Creates AI "haves" and "have-nots"

Efficiency enables democratization:

• Smaller models on consumer hardware
• Open-source weights (Llama, Mixtral, Mistral)
• Efficient fine-tuning (LoRA, QLoRA)
• Cloud access with pay-per-use pricing
• Edge deployment (on-device models)

Future of Efficient LLMs

Where is the field heading?

1. Hybrid Architectures
   • Combine MoE + attention + SSMs

• Use different mechanisms for different parts
• Dynamic architecture selection

2. Conditional Computation
   • Adjust depth/width based on input difficulty

• Early exit for easy examples
• More compute for hard problems

3. Neural Architecture Search
   • Automatically find efficient architectures

• Multi-objective optimization (quality + speed + size)

4. Hardware Co-Design
   • Design models for specific hardware

• Custom chips for LLMs (TPUs, Groq, Cerebras)
• Edge-specific models

Key Takeaways

1. MoE enables efficient scaling
   • More parameters, fewer active per token

• Better quality/compute ratio

2. Challenges exist but solvable
   • Load balancing, training instability

• Solutions: auxiliary losses, capacity limits

3. Mixtral proves MoE works at scale
   • 70B-quality with 13B-cost

• Open weights accelerate adoption

4. Many paths to efficiency
   • Quantization, pruning, distillation

• Flash attention, speculative decoding
• SSMs, hybrid architectures

5. Efficiency enables democratization
   • Lower costs, environmental impact

• Accessible to more researchers

Readings

Required:

1. Shazeer et al. (2017): Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
   [arXiv]
2. Jiang et al. (2024): Mixtral of Experts
   [arXiv]

Recommended:

• Fedus et al. (2022): Switch Transformers [arXiv]
• Gu & Dao (2023): Mamba [arXiv]
• Dao et al. (2022): FlashAttention [arXiv]
• Strubell et al. (2019): Energy and Policy Considerations [arXiv]
• HuggingFace MoE Blog [Blog]

Questions?

Next: Ethics, Bias, and Safety!