Mixture-of-Experts and Scaling

What This Is

Mixture-of-Experts (MoE) is an architectural trick that decouples total parameters from activated parameters. A MoE layer has many parallel "experts" (usually feed-forward sub-networks), a lightweight router decides which expert(s) each token is sent to, and only the chosen experts run for that token. A 1-trillion-parameter MoE model may activate only 30B parameters for any given forward pass.

Scaling laws are the complementary idea: a quantitative account of how model quality improves as parameters, data, and compute increase. Chinchilla (Hoffmann et al.) and follow-ups reframed what "scaling" should mean — not "make the model as big as possible" but "spend compute so that data and parameters scale together."

This topic treats the two together because modern frontier models use them together. MoE is the shape that gets the most capacity out of a fixed compute budget; scaling laws are the decision tool that tells you how to spend that budget.

When You Use It

• you are training a large model and every unit of compute matters
• you have more parameters you could train than your compute lets you activate on every token
• your model needs to route specialized capacity (domain-specific expert, language-specific expert)
• you are studying modern LLM architectures and need to know what the papers are doing

Do Not Use It When

• you are training a small or medium model — MoE overhead (routing, load balancing, all-to-all communication) does not pay off below a certain scale
• you cannot afford the engineering complexity — MoE changes training, serving, and evaluation all at once
• your compute is the bottleneck per token — MoE helps with parameter count at fixed per-token compute, but adds its own overhead

MoE In One Picture

  token embedding
        │
        ▼
   ┌────── router (softmax over N experts) ──────┐
   │       pick top-k experts per token          │
   └──────┬──────────┬──────────┬─────────┬──────┘
          ▼          ▼          ▼         ▼
     expert 1    expert 2    expert 3  ...  expert N
          │          │          │         │
          └──────────┴──────────┴─────────┘
                       │
                       ▼
              weighted sum of expert outputs
                       │
                       ▼
                residual add-in

The router is usually a single linear layer that outputs N logits, followed by a top-k selection (k=1 or k=2 typically). Only the chosen experts do any work on that token. Load balance is maintained via auxiliary losses so no single expert becomes the "popular" one.

A Minimal MoE Layer In PyTorch

Readable rather than performant:

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

class Expert(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model)
        )
    def forward(self, x): return self.net(x)

class TopKMoE(nn.Module):
    def __init__(self, d_model, d_ff, n_experts=8, k=2):
        super().__init__()
        self.k = k
        self.n_experts = n_experts
        self.router = nn.Linear(d_model, n_experts, bias=False)
        self.experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(n_experts)])

    def forward(self, x):                       # x: (B, T, D)
        B, T, D = x.shape
        x_flat = x.reshape(B * T, D)
        logits = self.router(x_flat)            # (B*T, N)
        topk_vals, topk_idx = logits.topk(self.k, dim=-1)  # (B*T, k)
        gates = F.softmax(topk_vals, dim=-1)    # normalize within top-k

        out = torch.zeros_like(x_flat)
        for e_idx, expert in enumerate(self.experts):
            # tokens routed to this expert (in any of the k slots)
            mask = (topk_idx == e_idx)                     # (B*T, k)
            token_mask = mask.any(dim=-1)                  # (B*T,)
            if not token_mask.any():
                continue
            tokens = x_flat[token_mask]
            expert_out = expert(tokens)
            # gating weight for this token/expert pair
            gate = (gates * mask.float()).sum(dim=-1)[token_mask].unsqueeze(-1)
            out[token_mask] += gate * expert_out
        return out.view(B, T, D)

This snippet is a teaching tool, not a production implementation. Real MoEs need:

• efficient all-to-all token shuffling across GPUs
• load-balancing auxiliary losses (see below)
• capacity limits (a maximum number of tokens per expert per step)
• dropping or padding overflow tokens

Use frameworks (fairscale, DeepSpeed-MoE, tutel, megablocks) for real runs.

Load Balance — The Thing That Breaks MoEs

Without explicit pressure, MoE routers collapse to "always send to expert 0." The auxiliary load-balancing loss, added to the main loss:

L_aux = n_experts · Σ_e  (fraction_of_tokens_to_e · avg_probability_of_e)

Intuition: punish the product of how often an expert is actually chosen and how much probability it gets. When both are high on a single expert, the aux loss penalizes it; when load is balanced across experts, the product is minimized.

Inspections that catch load-balance failures:

• routing histogram — fraction of tokens per expert per step. Should be roughly uniform.
• expert utilization over time — plot per-expert token count; catastrophic concentrations show up within the first few thousand steps
• capacity overflow rate — when experts hit the per-step token cap, surplus tokens are dropped. High overflow means the router is misaligned with the capacity.

Switch Transformer And Mixtral

Two reference points students should know:

• Switch Transformer (Fedus et al.) — k=1 routing (one expert per token), simpler and faster than top-2, trades a bit of quality for more tractable engineering
• Mixtral-8×7B (Mistral) — open-weight decoder LLM with k=2 routing over 8 experts per layer; activates roughly 13B of ~47B total parameters per token

The lesson in Mixtral's success: a well-engineered MoE with modest k and modest N can be released as open weights and run on modest hardware. Open MoEs are no longer a research curiosity.

Scaling Laws — The Decision Tool

Kaplan et al. (2020) noted that loss decreases as a power law in parameters, data, and compute, with roughly predictable exponents. Hoffmann et al. (2022, "Chinchilla") showed that earlier large models were under-trained — they had too many parameters for the tokens they saw.

The Chinchilla-optimal rule of thumb:

• for a given compute budget, parameters and training tokens should scale at roughly 1:1 ratio (e.g., ~20 tokens per parameter is a rough Chinchilla point — this exact constant moves as the papers refine)
• a model with too many parameters and too little data is strictly worse than a smaller model trained on more tokens with the same compute

The practical consequences:

• a 70B model trained on 1T tokens underperforms a 13B model trained on 3T tokens if compute is the same
• in deployment, smaller well-trained models are often preferable to larger under-trained ones — they are faster and more accurate per unit compute

Scaling Laws + MoE

How do the two ideas combine?

• for a fixed activated-parameter compute budget, MoEs let you pack more total parameters into the network, which generally improves quality
• scaling laws still apply on the activated parameter count — you cannot trade data for "total parameters" the way you can with dense models
• the right question for frontier training becomes: "under a compute budget, how many activated parameters should I have, how many total parameters, and how many tokens?"

No universal formula exists. The working answer at frontier labs in 2025 is: "MoE with small k (1–2), N chosen to fit hardware, total params ~3–5× activated params, trained for as many tokens as you can afford." Read current papers — the constants move.

What To Inspect

• router entropy — should be well above zero; collapse signals a dead router
• per-expert utilization distribution — spikes early are diagnostic of load-balance failure
• capacity overflow rate — if >10% of tokens are dropped, the capacity is too low or the aux loss is too weak
• training loss vs. compute curve — deviations from a predicted scaling curve flag engineering bugs more often than research breakthroughs
• validation loss per expert (for models that can be probed) — severely unbalanced quality across experts is a signal to redesign

Failure Pattern

The canonical MoE failure is router collapse: one expert wins all the traffic, the rest become dead weight. Diagnosable in the first few thousand steps via the routing histogram. Fix: increase aux-loss weight, enable noise in the router, or switch to expert-choice routing (where experts pick tokens rather than tokens picking experts).

The canonical scaling failure is ignoring the Chinchilla rule: training a very large model on too few tokens. The loss looks fine — it just underperforms a smaller well-trained peer on the same compute. This is hard to see without a reference run at a different scale.

Common Mistakes

• setting k = 1 or k = 2 without load-balancing aux loss — router collapse
• treating MoE as "free compute" — the router, all-to-all communication, and auxiliary losses have real costs
• reporting MoE quality against dense models at equal total parameters rather than equal activated parameters or equal FLOPs
• scaling parameters aggressively without scaling tokens — produces under-trained Chinchilla-suboptimal models
• routing at every layer when a handful of MoE layers interleaved with dense layers is often sufficient
• training MoE without capacity caps — uncontrolled token skew
• assuming MoE always helps — at small scale the overhead exceeds the benefit

Decision: Dense vs. MoE

option	when it wins	trade
dense	small/medium models, simple serving, evaluation clarity	fixed parameter/compute ratio
MoE (k=1, Switch-style)	very large models, fewer routing hops per layer	small quality gap to k=2
MoE (k=2)	large models, better quality	more compute per step; harder to balance
MoE with expert-choice	load-balance-critical deployments	newer, more engineering

A practical rule: below roughly 1B activated parameters, dense dominates for almost every real task. Above that, the MoE vs. dense decision should follow the compute budget, the serving constraints, and whether you have an engineering team that can ship the all-to-all communication.

Practice

1. Implement the minimal MoE layer above. Feed it random tokens. Confirm that the routing histogram over a few thousand tokens is approximately uniform when the router is freshly initialized.
2. Train a tiny MoE transformer on a synthetic task. Inspect the routing histogram every 1000 steps. Induce collapse by weakening the aux loss.
3. Add the auxiliary load-balancing loss. Watch utilization re-balance.
4. Measure activated vs. total parameters. Confirm you match what you expect from the top-k setting.
5. For scaling laws: take two dense models (a 10M-param and a 50M-param baseline) and train each on a small synthetic task for fixed compute. Fit a loss-vs-compute curve. Extrapolate; compare to a prediction at 25M params.
6. Read the Chinchilla paper's data-vs-params table once. Explain in writing why a 70B/1T-token run is worse than a 13B/3T-token run at the same compute.

Runnable Example

Real MoE training requires multi-GPU all-to-all communication and is far beyond the browser runner. The CPU-scale snippet above is enough to validate shapes and routing behavior. For reference implementations, megablocks and DeepSpeed-MoE are the modern open-source starting points.

Longer Connection

MoE and scaling sit at the intersection of several academy topics:

• Attention and Transformers — MoE layers replace the feed-forward sublayer of a transformer block; the attention half is unchanged
• PEFT and LoRA — PEFT is the small-scale complement to MoE; LoRA adapts a frozen dense base, MoE builds capacity directly into the architecture
• Quantization, Distillation, and Deployment — MoE models are expensive to serve; distillation to a dense student is a common deployment path
• Text Generation and Language Models — Mixtral is an MoE that students will meet in practice

For the decision frame — "should I be using MoE at all?" — almost no students should, because almost no students are training models at the scale where MoE is necessary. The reason to learn it is to read modern papers and to make informed decisions when the scale eventually matches. Read Baseline-First Task Solving one more time; the same rule applies: dense first, MoE only when dense is demonstrably the bottleneck.