{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-1 fw=all -->\n",
    "\n",
    "# Mixture of Experts\n",
    "\n",
    "Every model in this chapter so far spends its parameters and its compute\n",
    "in lockstep: a token flowing through the GPT of that section\n",
    "multiplies against every weight the model owns, so a model with twice\n",
    "the parameters charges every token twice the floating-point operations.\n",
    "The scaling laws that close this chapter will say that parameters are\n",
    "precisely what we want more of, and the census of\n",
    "that section says where they would go, since the\n",
    "feed-forward network already holds two thirds of every block. A\n",
    "*mixture of experts* (MoE) breaks the lockstep at exactly that spot:\n",
    "keep $E$ copies of the FFN (the *experts*) and let a small learned\n",
    "*router* send each token to only $k$ of them. Parameters now scale\n",
    "with $E$ while per-token compute scales with $k$, and the two dials\n",
    "turn independently. The idea is old (a committee of specialist\n",
    "networks under a gating network dates to\n",
    "@Jacobs.Jordan.Nowlan.ea.1991, sparse routing at scale to\n",
    "@Shazeer.Mirhoseini.Maziarz.ea.2017), but it is having its\n",
    "decade: Mixtral, DeepSeek-V3, Qwen3, and gpt-oss all ship it\n",
    "[@Jiang.Sablayrolles.Roux.ea.2024; @Liu.Feng.Xue.ea.2024; @Yang.Li.Yang.ea.2025; @OpenAI.2025].\n",
    "This section does the accounting that makes the idea attractive, builds\n",
    "the layer, meets the failure mode that nearly killed it — routing\n",
    "collapse — and repairs it twice, once with an auxiliary loss and once\n",
    "without. We close by swapping the FFN of our GPT for a mixture and\n",
    "measuring what the extra parameters buy."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-mixture-of-experts"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import torch as d2l\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-2 fw=all -->\n",
    "\n",
    "## Conditional Computation\n",
    "\n",
    "The bet behind MoE is that a language model's FFN parameters do not all\n",
    "need to fire on every token. Predicting the next token after \"the\n",
    "integral of\" and after \"Act I, Scene\" exercise different knowledge; a\n",
    "dense FFN pays for both everywhere, a routed one lets each token consult\n",
    "the specialists it needs. What makes the bet so attractive is the\n",
    "asymmetry of the two costs. Stored parameters are cheap: they sit in\n",
    "memory (or on other devices) and idle experts charge no FLOPs per token.\n",
    "Compute is the expensive, per-token resource. An MoE layer holds\n",
    "$E$ experts' parameters but each token pays the FLOPs of $k$, so — at\n",
    "roughly 2 forward FLOPs per active parameter per token — the model's\n",
    "capacity and its serving cost decouple.\n",
    "\n",
    "The arithmetic deserves to be computed rather than asserted, and it is\n",
    "short enough to check against the published configurations of two\n",
    "deployed models: Mixtral 8x7B, with 8 full-width experts of which each\n",
    "token activates 2 [@Jiang.Sablayrolles.Roux.ea.2024], and\n",
    "DeepSeek-V3, with 256 narrow experts per layer of which each token\n",
    "activates 8, plus one always-on shared expert\n",
    "[@Liu.Feng.Xue.ea.2024]. Each expert is a SwiGLU FFN — three\n",
    "$d \\times \\textrm{width}$ matrices (that section) —\n",
    "and the router adds a $d \\times E$ matrix per layer. The first row is\n",
    "the configuration this section trains at the end."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-conditional-computation"
   },
   "source": [
    "def moe_accounting(name, d, width, num_experts, num_active, num_layers,\n",
    "                   shared=0, unit=1e9):\n",
    "    per_expert = 3 * d * width          # SwiGLU: three matrices\n",
    "    store = num_layers * (num_experts + shared) * per_expert\n",
    "    active = num_layers * ((num_active + shared) * per_expert\n",
    "                           + d * num_experts)      # + the router\n",
    "    tag = 'B' if unit == 1e9 else 'M'\n",
    "    print(f'{name:>12}: {store/unit:6.1f}{tag} in experts, '\n",
    "          f'{active/unit:5.1f}{tag} active per token, '\n",
    "          f'ratio {store/active:4.1f}x')\n",
    "    return store, active\n",
    "\n",
    "moe_accounting('ours', 256, 683, 8, 1, 4, unit=1e6)\n",
    "store, active = moe_accounting('Mixtral 8x7B', 4096, 14336, 8, 2, 32)\n",
    "attn = 32 * (2 * 4096 * 4096 + 2 * 4096 * 1024)  # GQA: 32 Q, 8 KV heads\n",
    "emb = 2 * 32000 * 4096                           # embedding + output head\n",
    "print(f'     with attention and embeddings: {(store+attn+emb)/1e9:.1f}B '\n",
    "      f'total, {(active+attn+emb)/1e9:.1f}B active')\n",
    "moe_accounting('DeepSeek-V3', 7168, 2048, 256, 8, 58, shared=1)"
   ],
   "execution_count": 2,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "        ours:   16.8M in experts,   2.1M active per token, ratio  8.0x\nMixtral 8x7B:   45.1B in experts,  11.3B active per token, ratio  4.0x\n     with attention and embeddings: 46.7B total, 12.9B active\n DeepSeek-V3:  656.5B in experts,  23.1B active per token, ratio 28.4x\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "(656463101952, 23095410688)"
     },
     "metadata": {},
     "execution_count": 2
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-3 fw=all -->\n",
    "\n",
    "For Mixtral, adding the non-expert parameters — grouped-query attention\n",
    "(that section) and untied embeddings — reproduces the model\n",
    "card exactly: 46.7B parameters stored, 12.9B active per token. The name\n",
    "\"8x7B\" suggests eight times the cost of a 7B model; the accounting says\n",
    "a token pays for less than two-sevenths of what it could read from.\n",
    "DeepSeek-V3 pushes the same lever much harder: its 58 expert layers\n",
    "store about 656B of the model's 671B parameters, yet a token activates\n",
    "only 23B of them, a 28-fold gap between capacity and per-token compute\n",
    "(the model card's 37B \"active\" adds attention, embeddings, and the three\n",
    "dense layers that our expert-only census leaves out). Our toy\n",
    "configuration keeps Mixtral's eight-expert layout but routes top-1\n",
    "rather than top-2, so its 8x store-to-active ratio is twice as sparse\n",
    "as Mixtral's 4x.\n",
    "\n",
    "Two caveats before we cash in. First, over any real batch *training*\n",
    "still touches nearly every expert: gradients flow to whichever experts\n",
    "the batch's tokens visited, and the whole store must live in accelerator\n",
    "memory (in practice sharded across\n",
    "devices, with tokens shipped to their experts and back — the systems\n",
    "half of MoE, which belongs to a systems chapter). Second, the accounting\n",
    "assumed the router spreads tokens evenly. Nothing so far makes that\n",
    "true, and most of this section is about what happens when it is false.\n",
    "\n",
    "## A Mixture-of-Experts Layer\n",
    "\n",
    "The layer replaces the FFN, and only the FFN. Attention still\n",
    "communicates between positions as before; normalization and the\n",
    "residual stream are untouched. Given a token's vector\n",
    "$\\mathbf{x} \\in \\mathbb{R}^d$, a linear router scores the $E$ experts,\n",
    "the top $k$ scores pick the experts, and their outputs are combined,\n",
    "weighted by the router's own probabilities:\n",
    "\n",
    "$$\n",
    "\\mathrm{MoE}(\\mathbf{x}) = \\sum_{i \\in \\mathcal{E}_k(\\mathbf{x})} p_i(\\mathbf{x})\\, \\mathrm{FFN}_i(\\mathbf{x}),\n",
    "\\qquad\n",
    "\\mathbf{p}(\\mathbf{x}) = \\mathrm{softmax}(\\mathbf{W}_r \\mathbf{x}),\n",
    "$$\n",
    "\n",
    "where $\\mathcal{E}_k(\\mathbf{x})$ is the set of $k$ largest entries of\n",
    "$\\mathbf{p}(\\mathbf{x})$. This is *token-choice* routing: each token\n",
    "independently chooses its experts, so two adjacent tokens in the same\n",
    "sentence may consult entirely different parameters. Keeping the raw\n",
    "probability $p_i$ as the mixture weight (rather than renormalizing over\n",
    "the selected $k$) follows the Switch transformer\n",
    "[@fedus2022switch], and it is why the router can learn at all: the\n",
    "selection itself, an argmax, has no gradient, but the *weight* on each\n",
    "chosen expert does. If expert $i$'s output helps, the loss's gradient\n",
    "raises $p_i$ and the router routes to it more eagerly — a point worth\n",
    "dwelling on, and the subject of the first exercise.\n",
    "the figure shows one token's pass through the layer.\n",
    "\n",
    "![A mixture-of-experts layer. The router scores all experts for each token; the top-$k$ (here one) expert runs and its output is weighted by the router probability. Unselected experts hold parameters but perform no computation for this token.](../img/mdl-transformers-moe.svg)\n",
    "\n",
    "One disclosure about the implementation below: it computes\n",
    "*every* expert on *every* token and multiplies by a gate that is zero\n",
    "for the unselected ones. The mathematics of the equation is\n",
    "exact, but the FLOPs savings of the accounting cell is deliberately not\n",
    "realized — at our scale the dense computation is simpler and plenty\n",
    "fast. Production systems realize the savings with scatter/gather\n",
    "kernels: tokens are physically regrouped per expert, each expert runs\n",
    "one dense batch of its own tokens, and fixed per-expert buffer sizes\n",
    "(the *capacity factor*, another exercise) keep the shapes static. The\n",
    "teaching implementation preserves the semantics; the accounting carries\n",
    "the economics."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-a-mixture-of-experts-layer-1"
   },
   "source": [
    "class MoELayer(nn.Module):\n",
    "    \"\"\"Mixture-of-experts FFN: a token-choice top-k router over E experts.\"\"\"\n",
    "    def __init__(self, num_hiddens, num_experts, num_active):\n",
    "        super().__init__()\n",
    "        self.num_experts, self.num_active = num_experts, num_active\n",
    "        self.router = nn.Linear(num_hiddens, num_experts, bias=False)\n",
    "        self.experts = nn.ModuleList([d2l.FeedForward(num_hiddens)\n",
    "                                      for _ in range(num_experts)])\n",
    "        self.register_buffer('expert_bias', torch.zeros(num_experts))\n",
    "        self.register_buffer('usage', torch.zeros(num_experts))\n",
    "\n",
    "    def forward(self, X):\n",
    "        probs = F.softmax(self.router(X), -1)             # (B, T, E)\n",
    "        scores = probs + self.expert_bias                 # selection only\n",
    "        idx = scores.topk(self.num_active, -1).indices    # (B, T, k)\n",
    "        mask = torch.zeros_like(probs).scatter(-1, idx, 1.0)\n",
    "        gates = probs * mask                              # weight = p_i\n",
    "        Y = torch.stack([e(X) for e in self.experts], -1)  # (B, T, d, E)\n",
    "        out = (Y * gates.unsqueeze(-2)).sum(-1)\n",
    "        frac = mask.sum((0, 1)) / mask.sum()              # realized load\n",
    "        self.usage += mask.sum((0, 1)).detach()\n",
    "        self.aux_loss = self.num_experts * (frac * probs.mean((0, 1))).sum()\n",
    "        return out"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-4 fw=all -->\n",
    "\n",
    "The experts are the `d2l.FeedForward` of\n",
    "that section, unchanged. Two buffers ride along for\n",
    "later: `expert_bias` enters the top-$k$ *selection* but never the\n",
    "mixture weights, and `usage` counts routed tokens — both sit idle until\n",
    "the balancing experiments below. The layer also reports `aux_loss`, a\n",
    "balance score whose meaning the next section derives; nothing uses it\n",
    "yet. Since `MoELayer` maps $(n, d)$ to $(n, d)$ like any FFN, it drops\n",
    "into `d2l.TransformerBlock` through the `ffn_factory` seam that\n",
    "that section built and that section already\n",
    "used for attention:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-a-mixture-of-experts-layer-2"
   },
   "source": [
    "torch.manual_seed(0)\n",
    "moe = MoELayer(256, num_experts=8, num_active=2)\n",
    "X = torch.randn(4, 100, 256)\n",
    "d2l.check_shape(moe(X), X.shape)\n",
    "blk = d2l.TransformerBlock(256, num_heads=8,\n",
    "                           ffn_factory=lambda: MoELayer(256, 8, 2))\n",
    "d2l.check_shape(blk(X), X.shape)\n",
    "count = lambda m: sum(p.numel() for p in m.parameters())\n",
    "expert = count(moe.experts[0])\n",
    "print(f'router {count(moe.router)}, per expert {expert}, total '\n",
    "      f'{count(moe)}, active '\n",
    "      f'{count(moe) - (moe.num_experts - moe.num_active) * expert}')\n",
    "print('usage at initialization: '\n",
    "      + ' '.join(f'{u:.2f}' for u in moe.usage / moe.usage.sum()))"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "router 2048, per expert 524544, total 4198400, active 1051136\nusage at initialization: 0.13 0.14 0.15 0.10 0.12 0.12 0.11 0.13\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-5 fw=all -->\n",
    "\n",
    "At initialization the picture is as benign as the accounting assumed: a\n",
    "random router spreads random tokens nearly evenly across the eight\n",
    "experts. The interesting question is whether training keeps it that way.\n",
    "\n",
    "## Routing Collapse and Load Balancing\n",
    "\n",
    "It does not.\n",
    "\n",
    "### The Rich Get Richer\n",
    "\n",
    "Follow the feedback loop. Early in training some expert — by pure\n",
    "initialization luck — is slightly better than its peers on the tokens it\n",
    "happens to receive. The loss gradient rewards it twice: its parameters\n",
    "improve on those tokens, *and* the router's gradient raises its\n",
    "selection probability, because weighting a helpful expert more heavily\n",
    "lowers the loss. More tokens mean more gradient signal, which means a\n",
    "better expert, which attracts more tokens. The experts that lose the\n",
    "early rounds see ever fewer tokens, learn ever more slowly, and their\n",
    "router scores sink. The stable end state is *routing collapse*: a couple\n",
    "of experts serve everything while the rest ride along as dead weight,\n",
    "the capacity we listed as the whole point of the architecture unused.\n",
    "@Shazeer.Mirhoseini.Maziarz.ea.2017 describe exactly this\n",
    "self-reinforcing imbalance, and every practical MoE since has shipped\n",
    "some countermeasure. We will see the collapse live in a moment; first,\n",
    "the two countermeasures we will test against it.\n",
    "\n",
    "### An Auxiliary Balancing Loss\n",
    "\n",
    "The classic repair adds a differentiable penalty for imbalance to the\n",
    "training loss. For a batch, let $f_i$ be the fraction of routed tokens\n",
    "that expert $i$ received (a count — no gradient flows through it), and\n",
    "let $\\bar{p}_i$ be the router's mean probability for expert $i$. The\n",
    "Switch transformer's loss [@fedus2022switch], descended from\n",
    "GShard's [@Lepikhin.Lee.Xu.ea.2021], is their correlation:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{\\mathrm{balance}} = E \\sum_{i=1}^{E} f_i\\, \\bar{p}_i,\n",
    "$$\n",
    "\n",
    "added to the language-modeling loss with a small weight $\\alpha$ (0.01\n",
    "in the Switch paper and below). When routing is uniform,\n",
    "$f_i = \\bar{p}_i = 1/E$ and the loss is exactly $1$; concentrating\n",
    "tokens on few experts makes $f$ and $\\bar{p}$ peak together and the\n",
    "correlation grow. Its gradient (through $\\bar{p}$ only) pushes\n",
    "probability away from overloaded experts *for every token, whatever the\n",
    "token needs* — and that is the cost. The balancing gradient is added to\n",
    "the language-modeling gradient at every step, and the two disagree\n",
    "whenever genuinely popular knowledge deserves a busy expert. Tune\n",
    "$\\alpha$ too high and routing quality degrades; too low and the\n",
    "imbalance persists. This interference is documented rather than\n",
    "hypothetical [@Wang.Gao.Zeng.ea.2024], and it motivated the second\n",
    "repair.\n",
    "\n",
    "### Balancing Without a Loss\n",
    "\n",
    "The `expert_bias` buffer in our layer implements the auxiliary-loss-free\n",
    "scheme of @Wang.Gao.Zeng.ea.2024, the one DeepSeek-V3 trains\n",
    "with [@Liu.Feng.Xue.ea.2024]. Selection — and only selection — sees\n",
    "a per-expert bias:\n",
    "\n",
    "$$\n",
    "\\mathcal{E}_k(\\mathbf{x}) = \\operatorname{argtop}_k \\big(p_i(\\mathbf{x}) + b_i\\big),\n",
    "\\qquad\n",
    "b_i \\leftarrow b_i + u\\, \\mathrm{sign}(\\bar{f} - f_i),\n",
    "$$\n",
    "\n",
    "where $\\bar{f} = 1/E$ is the ideal load and $u$ is a small update speed.\n",
    "This is not gradient descent; it is a thermostat. After each step, every\n",
    "overloaded expert's bias ticks down and every underloaded expert's ticks\n",
    "up, making the starved experts slightly easier to *select* — but because\n",
    "$b_i$ never enters the mixture weight $p_i$, the loss the model\n",
    "optimizes contains no balancing term at all, so there is no gradient\n",
    "interference to tune away. The controller steers where tokens go; the\n",
    "gradients remain purely about predicting text. One simplification to\n",
    "note: DeepSeek-V3's router scores experts with per-expert sigmoids and\n",
    "restricts each token's choices to a few expert groups (which also limits\n",
    "cross-device traffic); our layer keeps the softmax router and changes\n",
    "only the selection bias — the mechanism @Wang.Gao.Zeng.ea.2024\n",
    "isolate.\n",
    "\n",
    "### Three Runs, One Budget\n",
    "\n",
    "Both repairs slot into an ordinary training loop. The trainer below is\n",
    "`d2l.train_lm` (that section) plus the two mechanisms, each\n",
    "switched by one argument: `aux_weight` adds\n",
    "$\\alpha \\sum \\mathcal{L}_{\\mathrm{balance}}$ over the model's MoE\n",
    "layers, and `bias_rate` runs the thermostat of the equation\n",
    "after each optimizer step. To build the model we take `d2l.GPT`\n",
    "unchanged and replace each block's FFN, the same\n",
    "swap-after-construction move that section used to install\n",
    "grouped-query attention:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-three-runs-one-budget-1"
   },
   "source": [
    "def moe_gpt(vocab_size, num_hiddens, num_blks, num_experts, num_active,\n",
    "            num_heads=4):\n",
    "    \"\"\"A d2l.GPT whose blocks' FFNs are replaced by MoE layers.\"\"\"\n",
    "    model = d2l.GPT(vocab_size, num_hiddens=num_hiddens,\n",
    "                    num_heads=num_heads, num_blks=num_blks)\n",
    "    for blk in model.blks:\n",
    "        blk.ffn = MoELayer(num_hiddens, num_experts, num_active)\n",
    "    return model\n",
    "\n",
    "def train_moe_lm(model, data, optimizer, num_steps, aux_weight=0.0,\n",
    "                 bias_rate=0.0):\n",
    "    \"\"\"d2l.train_lm plus the two balancing mechanisms.\"\"\"\n",
    "    device = d2l.try_gpu()\n",
    "    model.to(device)\n",
    "    layers = [blk.ffn for blk in model.blks\n",
    "              if isinstance(blk.ffn, MoELayer)]\n",
    "    losses, step = [], 0\n",
    "    while step < num_steps:\n",
    "        for X, Y in data.train_dataloader():\n",
    "            X, Y = X.to(device), Y.to(device)\n",
    "            logits = model(X)\n",
    "            ce = F.cross_entropy(logits.reshape(-1, logits.shape[-1]),\n",
    "                                 Y.reshape(-1))\n",
    "            loss = ce\n",
    "            if aux_weight:\n",
    "                loss = loss + aux_weight * sum(l.aux_loss for l in layers)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if bias_rate:               # the thermostat: no gradients\n",
    "                with torch.no_grad():\n",
    "                    for l in layers:\n",
    "                        frac = l.usage / l.usage.sum()\n",
    "                        l.expert_bias += bias_rate * torch.sign(\n",
    "                            1 / l.num_experts - frac)\n",
    "                        l.usage.zero_()\n",
    "            losses.append(ce.item())\n",
    "            step += 1\n",
    "            if step >= num_steps:\n",
    "                return losses"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-6 fw=all -->\n",
    "\n",
    "Two measurement helpers: validation loss as in that section, and\n",
    "each layer's expert-usage distribution over a fixed set of validation\n",
    "batches."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-three-runs-one-budget-2"
   },
   "source": [
    "def val_loss(model, data):\n",
    "    device = d2l.try_gpu()\n",
    "    model.to(device).eval()\n",
    "    with torch.no_grad():\n",
    "        losses = [F.cross_entropy(\n",
    "            model(X.to(device)).flatten(0, 1), Y.to(device).flatten())\n",
    "            for X, Y in data.val_dataloader()]\n",
    "    model.train()\n",
    "    return sum(l.item() for l in losses) / len(losses)\n",
    "\n",
    "def usage_fractions(model, data, num_batches=10):\n",
    "    \"\"\"Per-layer expert usage over a fixed set of validation batches.\"\"\"\n",
    "    device = d2l.try_gpu()\n",
    "    layers = [blk.ffn for blk in model.blks\n",
    "              if isinstance(blk.ffn, MoELayer)]\n",
    "    for l in layers:\n",
    "        l.usage.zero_()\n",
    "    model.to(device).eval()\n",
    "    with torch.no_grad():\n",
    "        for i, (X, _) in enumerate(data.val_dataloader()):\n",
    "            if i >= num_batches:\n",
    "                break\n",
    "            model(X.to(device))\n",
    "    model.train()\n",
    "    return torch.stack([l.usage / l.usage.sum() for l in layers]).cpu()"
   ],
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-7 fw=all -->\n",
    "\n",
    "Now the experiment. One small model — two blocks of width 128, eight\n",
    "experts per block, one active per token (the Switch configuration,\n",
    "where the feedback loop bites hardest) — trained three times from the\n",
    "same initialization on the character-level Time Machine for the same\n",
    "800 steps. The only difference between the runs is the balancing: none,\n",
    "the auxiliary loss at $\\alpha = 0.01$, or the bias thermostat at\n",
    "$u = 0.01$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-three-runs-one-budget-3"
   },
   "source": [
    "data = d2l.TimeMachine(batch_size=64, num_steps=128, tokenization='char',\n",
    "                       num_train=100000, num_val=3000)\n",
    "results = {}\n",
    "for name, aux_weight, bias_rate in (('no balancing', 0.0, 0.0),\n",
    "                                    ('auxiliary loss', 0.01, 0.0),\n",
    "                                    ('bias', 0.0, 0.01)):\n",
    "    torch.manual_seed(0)\n",
    "    model = moe_gpt(len(data.vocab), num_hiddens=128, num_blks=2,\n",
    "                    num_experts=8, num_active=1)\n",
    "    losses = train_moe_lm(\n",
    "        model, data, torch.optim.AdamW(model.parameters(), lr=1e-3,\n",
    "                                       weight_decay=0.0),\n",
    "        800, aux_weight, bias_rate)\n",
    "    results[name] = usage_fractions(model, data)\n",
    "    print(f'{name:>14}: training loss {sum(losses[-100:])/100:.2f}, '\n",
    "          f'experts under 2% usage: {int((results[name] < 0.02).sum())} '\n",
    "          'of 16')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  no balancing: training loss 0.72, experts under 2% usage: 5 of 16\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "auxiliary loss: training loss 0.66, experts under 2% usage: 0 of 16\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "          bias: training loss 0.54, experts under 2% usage: 0 of 16\n"
     ]
    }
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-three-runs-one-budget-4"
   },
   "source": [
    "d2l.use_svg_display()\n",
    "fig, axes = d2l.plt.subplots(1, 3, figsize=(9, 2.8), sharey=True)\n",
    "x = torch.arange(8)\n",
    "for ax, (name, u) in zip(axes, results.items()):\n",
    "    ax.bar(x - 0.2, u[0], 0.4, label='layer 1')\n",
    "    ax.bar(x + 0.2, u[1], 0.4, label='layer 2')\n",
    "    ax.axhline(1 / 8, color='gray', ls='--', lw=1)\n",
    "    ax.set_xlabel('expert')\n",
    "    ax.set_title(name)\n",
    "axes[0].set_ylabel('fraction of tokens')\n",
    "axes[0].legend()"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "<matplotlib.legend.Legend at 0x7b98db7ec7a0>"
     },
     "metadata": {},
     "execution_count": 8
    },
    {
     "output_type": "display_data",
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<svg xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"556.140421pt\" height=\"215.106375pt\" viewBox=\"0 0 556.140421 215.106375\" xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n <metadata>\n  <rdf:RDF xmlns:dc=\"http://purl.org/dc/elements/1.1/\" xmlns:cc=\"http://creativecommons.org/ns#\" xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n   <cc:Work>\n    <dc:type rdf:resource=\"http://purl.org/dc/dcmitype/StillImage\"/>\n    <dc:date>2026-07-19T14:05:12.273148</dc:date>\n    <dc:format>image/svg+xml</dc:format>\n    <dc:creator>\n     <cc:Agent>\n      <dc:title>Matplotlib v3.10.8, https://matplotlib.org/</dc:title>\n     </cc:Agent>\n    </dc:creator>\n   </cc:Work>\n  </rdf:RDF>\n </metadata>\n <defs>\n  <style type=\"text/css\">*{stroke-linejoin: round; stroke-linecap: butt}</style>\n </defs>\n <g id=\"figure_1\">\n  <g id=\"patch_1\">\n   <path d=\"M 0 215.106375 \nL 556.140421 215.106375 \nL 556.140421 -0 \nL 0 -0 \nz\n\" style=\"fill: #ffffff\"/>\n  </g>\n  <g id=\"axes_1\">\n   <g id=\"patch_2\">\n    <path d=\"M 43.78125 177.550125 \nL 191.487132 177.550125 \nL 191.487132 22.318125 \nL 43.78125 22.318125 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"patch_3\">\n    <path d=\"M 50.495154 177.550125 \nL 57.381209 177.550125 \nL 57.381209 29.710125 \nL 50.495154 29.710125 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_4\">\n    <path d=\"M 67.710292 177.550125 \nL 74.596346 177.550125 \nL 74.596346 176.342537 \nL 67.710292 176.342537 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_5\">\n    <path d=\"M 84.925427 177.550125 \nL 91.811482 177.550125 \nL 91.811482 175.089033 \nL 84.925427 175.089033 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_6\">\n    <path d=\"M 102.140565 177.550125 \nL 109.026622 177.550125 \nL 109.026622 172.839154 \nL 102.140565 172.839154 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_7\">\n    <path d=\"M 119.355703 177.550125 \nL 126.241759 177.550125 \nL 126.241759 159.032245 \nL 119.355703 159.032245 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_8\">\n    <path d=\"M 136.570848 177.550125 \nL 143.456905 177.550125 \nL 143.456905 53.77465 \nL 136.570848 53.77465 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_9\">\n    <path d=\"M 153.785986 177.550125 \nL 160.672043 177.550125 \nL 160.672043 160.496962 \nL 153.785986 160.496962 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_10\">\n    <path d=\"M 171.001123 177.550125 \nL 177.88718 177.550125 \nL 177.88718 116.973283 \nL 171.001123 116.973283 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_11\">\n    <path d=\"M 57.381209 177.550125 \nL 64.267265 177.550125 \nL 64.267265 73.706733 \nL 57.381209 73.706733 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_12\">\n    <path d=\"M 74.596346 177.550125 \nL 81.482403 177.550125 \nL 81.482403 139.917459 \nL 74.596346 139.917459 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_13\">\n    <path d=\"M 91.811484 177.550125 \nL 98.697541 177.550125 \nL 98.697541 152.475456 \nL 91.811484 152.475456 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_14\">\n    <path d=\"M 109.026622 177.550125 \nL 115.912678 177.550125 \nL 115.912678 98.781408 \nL 109.026622 98.781408 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_15\">\n    <path d=\"M 126.241755 177.550125 \nL 133.127808 177.550125 \nL 133.127808 172.292755 \nL 126.241755 172.292755 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_16\">\n    <path d=\"M 143.456897 177.550125 \nL 150.342953 177.550125 \nL 150.342953 78.527902 \nL 143.456897 78.527902 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_17\">\n    <path d=\"M 160.672034 177.550125 \nL 167.558091 177.550125 \nL 167.558091 157.195609 \nL 160.672034 157.195609 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_18\">\n    <path d=\"M 177.887172 177.550125 \nL 184.773229 177.550125 \nL 184.773229 171.360662 \nL 177.887172 171.360662 \nz\n\" clip-path=\"url(#pf838c49614)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"matplotlib.axis_1\">\n    <g id=\"xtick_1\">\n     <g id=\"line2d_1\">\n      <defs>\n       <path id=\"md4ed840694\" d=\"M 0 0 \nL 0 3.5 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"57.381209\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_1\">\n      <!-- 0.0 -->\n      <g transform=\"translate(49.429646 192.148562) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-30\" d=\"M 2034 4250 \nQ 1547 4250 1301 3770 \nQ 1056 3291 1056 2328 \nQ 1056 1369 1301 889 \nQ 1547 409 2034 409 \nQ 2525 409 2770 889 \nQ 3016 1369 3016 2328 \nQ 3016 3291 2770 3770 \nQ 2525 4250 2034 4250 \nz\nM 2034 4750 \nQ 2819 4750 3233 4129 \nQ 3647 3509 3647 2328 \nQ 3647 1150 3233 529 \nQ 2819 -91 2034 -91 \nQ 1250 -91 836 529 \nQ 422 1150 422 2328 \nQ 422 3509 836 4129 \nQ 1250 4750 2034 4750 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-2e\" d=\"M 684 794 \nL 1344 794 \nL 1344 0 \nL 684 0 \nL 684 794 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_2\">\n     <g id=\"line2d_2\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"100.419053\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_2\">\n      <!-- 2.5 -->\n      <g transform=\"translate(92.46749 192.148562) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-32\" d=\"M 1228 531 \nL 3431 531 \nL 3431 0 \nL 469 0 \nL 469 531 \nQ 828 903 1448 1529 \nQ 2069 2156 2228 2338 \nQ 2531 2678 2651 2914 \nQ 2772 3150 2772 3378 \nQ 2772 3750 2511 3984 \nQ 2250 4219 1831 4219 \nQ 1534 4219 1204 4116 \nQ 875 4013 500 3803 \nL 500 4441 \nQ 881 4594 1212 4672 \nQ 1544 4750 1819 4750 \nQ 2544 4750 2975 4387 \nQ 3406 4025 3406 3419 \nQ 3406 3131 3298 2873 \nQ 3191 2616 2906 2266 \nQ 2828 2175 2409 1742 \nQ 1991 1309 1228 531 \nz\n\" transform=\"scale(0.015625)\"/>\n        <path id=\"DejaVuSans-35\" d=\"M 691 4666 \nL 3169 4666 \nL 3169 4134 \nL 1269 4134 \nL 1269 2991 \nQ 1406 3038 1543 3061 \nQ 1681 3084 1819 3084 \nQ 2600 3084 3056 2656 \nQ 3513 2228 3513 1497 \nQ 3513 744 3044 326 \nQ 2575 -91 1722 -91 \nQ 1428 -91 1123 -41 \nQ 819 9 494 109 \nL 494 744 \nQ 775 591 1075 516 \nQ 1375 441 1709 441 \nQ 2250 441 2565 725 \nQ 2881 1009 2881 1497 \nQ 2881 1984 2565 2268 \nQ 2250 2553 1709 2553 \nQ 1456 2553 1204 2497 \nQ 953 2441 691 2322 \nL 691 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-32\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_3\">\n     <g id=\"line2d_3\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"143.456897\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_3\">\n      <!-- 5.0 -->\n      <g transform=\"translate(135.505334 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-35\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_4\">\n     <g id=\"line2d_4\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"186.494741\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_4\">\n      <!-- 7.5 -->\n      <g transform=\"translate(178.543178 192.148562) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-37\" d=\"M 525 4666 \nL 3525 4666 \nL 3525 4397 \nL 1831 0 \nL 1172 0 \nL 2766 4134 \nL 525 4134 \nL 525 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-37\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_5\">\n     <!-- expert -->\n     <g transform=\"translate(101.419347 205.826687) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-65\" d=\"M 3597 1894 \nL 3597 1613 \nL 953 1613 \nQ 991 1019 1311 708 \nQ 1631 397 2203 397 \nQ 2534 397 2845 478 \nQ 3156 559 3463 722 \nL 3463 178 \nQ 3153 47 2828 -22 \nQ 2503 -91 2169 -91 \nQ 1331 -91 842 396 \nQ 353 884 353 1716 \nQ 353 2575 817 3079 \nQ 1281 3584 2069 3584 \nQ 2775 3584 3186 3129 \nQ 3597 2675 3597 1894 \nz\nM 3022 2063 \nQ 3016 2534 2758 2815 \nQ 2500 3097 2075 3097 \nQ 1594 3097 1305 2825 \nQ 1016 2553 972 2059 \nL 3022 2063 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-78\" d=\"M 3513 3500 \nL 2247 1797 \nL 3578 0 \nL 2900 0 \nL 1881 1375 \nL 863 0 \nL 184 0 \nL 1544 1831 \nL 300 3500 \nL 978 3500 \nL 1906 2253 \nL 2834 3500 \nL 3513 3500 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-70\" d=\"M 1159 525 \nL 1159 -1331 \nL 581 -1331 \nL 581 3500 \nL 1159 3500 \nL 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nz\nM 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-72\" d=\"M 2631 2963 \nQ 2534 3019 2420 3045 \nQ 2306 3072 2169 3072 \nQ 1681 3072 1420 2755 \nQ 1159 2438 1159 1844 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1341 3275 1631 3429 \nQ 1922 3584 2338 3584 \nQ 2397 3584 2469 3576 \nQ 2541 3569 2628 3553 \nL 2631 2963 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-74\" d=\"M 1172 4494 \nL 1172 3500 \nL 2356 3500 \nL 2356 3053 \nL 1172 3053 \nL 1172 1153 \nQ 1172 725 1289 603 \nQ 1406 481 1766 481 \nL 2356 481 \nL 2356 0 \nL 1766 0 \nQ 1100 0 847 248 \nQ 594 497 594 1153 \nL 594 3053 \nL 172 3053 \nL 172 3500 \nL 594 3500 \nL 594 4494 \nL 1172 4494 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-65\"/>\n      <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(59.773438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(118.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(182.429688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(243.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(285.066406 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_2\">\n    <g id=\"ytick_1\">\n     <g id=\"line2d_5\">\n      <defs>\n       <path id=\"m45dbd8aeb0\" d=\"M 0 0 \nL -3.5 0 \n\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </defs>\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"43.78125\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_6\">\n      <!-- 0.0 -->\n      <g transform=\"translate(20.878125 181.349344) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_2\">\n     <g id=\"line2d_6\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"43.78125\" y=\"139.935824\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_7\">\n      <!-- 0.1 -->\n      <g transform=\"translate(20.878125 143.735043) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-31\" d=\"M 794 531 \nL 1825 531 \nL 1825 4091 \nL 703 3866 \nL 703 4441 \nL 1819 4666 \nL 2450 4666 \nL 2450 531 \nL 3481 531 \nL 3481 0 \nL 794 0 \nL 794 531 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_3\">\n     <g id=\"line2d_7\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"43.78125\" y=\"102.321523\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_8\">\n      <!-- 0.2 -->\n      <g transform=\"translate(20.878125 106.120742) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_4\">\n     <g id=\"line2d_8\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"43.78125\" y=\"64.707222\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_9\">\n      <!-- 0.3 -->\n      <g transform=\"translate(20.878125 68.506441) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-33\" d=\"M 2597 2516 \nQ 3050 2419 3304 2112 \nQ 3559 1806 3559 1356 \nQ 3559 666 3084 287 \nQ 2609 -91 1734 -91 \nQ 1441 -91 1130 -33 \nQ 819 25 488 141 \nL 488 750 \nQ 750 597 1062 519 \nQ 1375 441 1716 441 \nQ 2309 441 2620 675 \nQ 2931 909 2931 1356 \nQ 2931 1769 2642 2001 \nQ 2353 2234 1838 2234 \nL 1294 2234 \nL 1294 2753 \nL 1863 2753 \nQ 2328 2753 2575 2939 \nQ 2822 3125 2822 3475 \nQ 2822 3834 2567 4026 \nQ 2313 4219 1838 4219 \nQ 1578 4219 1281 4162 \nQ 984 4106 628 3988 \nL 628 4550 \nQ 988 4650 1302 4700 \nQ 1616 4750 1894 4750 \nQ 2613 4750 3031 4423 \nQ 3450 4097 3450 3541 \nQ 3450 3153 3228 2886 \nQ 3006 2619 2597 2516 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-33\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_5\">\n     <g id=\"line2d_9\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"43.78125\" y=\"27.092921\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_10\">\n      <!-- 0.4 -->\n      <g transform=\"translate(20.878125 30.89214) scale(0.1 -0.1)\">\n       <defs>\n        <path id=\"DejaVuSans-34\" d=\"M 2419 4116 \nL 825 1625 \nL 2419 1625 \nL 2419 4116 \nz\nM 2253 4666 \nL 3047 4666 \nL 3047 1625 \nL 3713 1625 \nL 3713 1100 \nL 3047 1100 \nL 3047 0 \nL 2419 0 \nL 2419 1100 \nL 313 1100 \nL 313 1709 \nL 2253 4666 \nz\n\" transform=\"scale(0.015625)\"/>\n       </defs>\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-34\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_11\">\n     <!-- fraction of tokens -->\n     <g transform=\"translate(14.798438 143.721625) rotate(-90) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-66\" d=\"M 2375 4863 \nL 2375 4384 \nL 1825 4384 \nQ 1516 4384 1395 4259 \nQ 1275 4134 1275 3809 \nL 1275 3500 \nL 2222 3500 \nL 2222 3053 \nL 1275 3053 \nL 1275 0 \nL 697 0 \nL 697 3053 \nL 147 3053 \nL 147 3500 \nL 697 3500 \nL 697 3744 \nQ 697 4328 969 4595 \nQ 1241 4863 1831 4863 \nL 2375 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-61\" d=\"M 2194 1759 \nQ 1497 1759 1228 1600 \nQ 959 1441 959 1056 \nQ 959 750 1161 570 \nQ 1363 391 1709 391 \nQ 2188 391 2477 730 \nQ 2766 1069 2766 1631 \nL 2766 1759 \nL 2194 1759 \nz\nM 3341 1997 \nL 3341 0 \nL 2766 0 \nL 2766 531 \nQ 2569 213 2275 61 \nQ 1981 -91 1556 -91 \nQ 1019 -91 701 211 \nQ 384 513 384 1019 \nQ 384 1609 779 1909 \nQ 1175 2209 1959 2209 \nL 2766 2209 \nL 2766 2266 \nQ 2766 2663 2505 2880 \nQ 2244 3097 1772 3097 \nQ 1472 3097 1187 3025 \nQ 903 2953 641 2809 \nL 641 3341 \nQ 956 3463 1253 3523 \nQ 1550 3584 1831 3584 \nQ 2591 3584 2966 3190 \nQ 3341 2797 3341 1997 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-63\" d=\"M 3122 3366 \nL 3122 2828 \nQ 2878 2963 2633 3030 \nQ 2388 3097 2138 3097 \nQ 1578 3097 1268 2742 \nQ 959 2388 959 1747 \nQ 959 1106 1268 751 \nQ 1578 397 2138 397 \nQ 2388 397 2633 464 \nQ 2878 531 3122 666 \nL 3122 134 \nQ 2881 22 2623 -34 \nQ 2366 -91 2075 -91 \nQ 1284 -91 818 406 \nQ 353 903 353 1747 \nQ 353 2603 823 3093 \nQ 1294 3584 2113 3584 \nQ 2378 3584 2631 3529 \nQ 2884 3475 3122 3366 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-69\" d=\"M 603 3500 \nL 1178 3500 \nL 1178 0 \nL 603 0 \nL 603 3500 \nz\nM 603 4863 \nL 1178 4863 \nL 1178 4134 \nL 603 4134 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6f\" d=\"M 1959 3097 \nQ 1497 3097 1228 2736 \nQ 959 2375 959 1747 \nQ 959 1119 1226 758 \nQ 1494 397 1959 397 \nQ 2419 397 2687 759 \nQ 2956 1122 2956 1747 \nQ 2956 2369 2687 2733 \nQ 2419 3097 1959 3097 \nz\nM 1959 3584 \nQ 2709 3584 3137 3096 \nQ 3566 2609 3566 1747 \nQ 3566 888 3137 398 \nQ 2709 -91 1959 -91 \nQ 1206 -91 779 398 \nQ 353 888 353 1747 \nQ 353 2609 779 3096 \nQ 1206 3584 1959 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6e\" d=\"M 3513 2113 \nL 3513 0 \nL 2938 0 \nL 2938 2094 \nQ 2938 2591 2744 2837 \nQ 2550 3084 2163 3084 \nQ 1697 3084 1428 2787 \nQ 1159 2491 1159 1978 \nL 1159 0 \nL 581 0 \nL 581 3500 \nL 1159 3500 \nL 1159 2956 \nQ 1366 3272 1645 3428 \nQ 1925 3584 2291 3584 \nQ 2894 3584 3203 3211 \nQ 3513 2838 3513 2113 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-20\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-6b\" d=\"M 581 4863 \nL 1159 4863 \nL 1159 1991 \nL 2875 3500 \nL 3609 3500 \nL 1753 1863 \nL 3688 0 \nL 2938 0 \nL 1159 1709 \nL 1159 0 \nL 581 0 \nL 581 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n       <path id=\"DejaVuSans-73\" d=\"M 2834 3397 \nL 2834 2853 \nQ 2591 2978 2328 3040 \nQ 2066 3103 1784 3103 \nQ 1356 3103 1142 2972 \nQ 928 2841 928 2578 \nQ 928 2378 1081 2264 \nQ 1234 2150 1697 2047 \nL 1894 2003 \nQ 2506 1872 2764 1633 \nQ 3022 1394 3022 966 \nQ 3022 478 2636 193 \nQ 2250 -91 1575 -91 \nQ 1294 -91 989 -36 \nQ 684 19 347 128 \nL 347 722 \nQ 666 556 975 473 \nQ 1284 391 1588 391 \nQ 1994 391 2212 530 \nQ 2431 669 2431 922 \nQ 2431 1156 2273 1281 \nQ 2116 1406 1581 1522 \nL 1381 1569 \nQ 847 1681 609 1914 \nQ 372 2147 372 2553 \nQ 372 3047 722 3315 \nQ 1072 3584 1716 3584 \nQ 2034 3584 2315 3537 \nQ 2597 3491 2834 3397 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-66\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(35.205078 0)\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(76.318359 0)\"/>\n      <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(137.597656 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(192.578125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(231.787109 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(259.570312 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(320.751953 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(384.130859 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(415.917969 0)\"/>\n      <use xlink:href=\"#DejaVuSans-66\" transform=\"translate(477.099609 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(512.304688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(544.091797 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(583.300781 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6b\" transform=\"translate(644.482422 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(698.767578 0)\"/>\n      <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(760.291016 0)\"/>\n      <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(823.669922 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_10\">\n    <path d=\"M 43.78125 130.532249 \nL 191.487132 130.532249 \n\" clip-path=\"url(#pf838c49614)\" style=\"fill: none; stroke-dasharray: 3.7,1.6; stroke-dashoffset: 0; stroke: #808080\"/>\n   </g>\n   <g id=\"patch_19\">\n    <path d=\"M 43.78125 177.550125 \nL 43.78125 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_20\">\n    <path d=\"M 191.487132 177.550125 \nL 191.487132 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_21\">\n    <path d=\"M 43.78125 177.550125 \nL 191.487132 177.550125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_22\">\n    <path d=\"M 43.78125 22.318125 \nL 191.487132 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"text_12\">\n    <!-- no balancing -->\n    <g transform=\"translate(79.043879 16.318125) scale(0.12 -0.12)\">\n     <defs>\n      <path id=\"DejaVuSans-62\" d=\"M 3116 1747 \nQ 3116 2381 2855 2742 \nQ 2594 3103 2138 3103 \nQ 1681 3103 1420 2742 \nQ 1159 2381 1159 1747 \nQ 1159 1113 1420 752 \nQ 1681 391 2138 391 \nQ 2594 391 2855 752 \nQ 3116 1113 3116 1747 \nz\nM 1159 2969 \nQ 1341 3281 1617 3432 \nQ 1894 3584 2278 3584 \nQ 2916 3584 3314 3078 \nQ 3713 2572 3713 1747 \nQ 3713 922 3314 415 \nQ 2916 -91 2278 -91 \nQ 1894 -91 1617 61 \nQ 1341 213 1159 525 \nL 1159 0 \nL 581 0 \nL 581 4863 \nL 1159 4863 \nL 1159 2969 \nz\n\" transform=\"scale(0.015625)\"/>\n      <path id=\"DejaVuSans-6c\" d=\"M 603 4863 \nL 1178 4863 \nL 1178 0 \nL 603 0 \nL 603 4863 \nz\n\" transform=\"scale(0.015625)\"/>\n      <path id=\"DejaVuSans-67\" d=\"M 2906 1791 \nQ 2906 2416 2648 2759 \nQ 2391 3103 1925 3103 \nQ 1463 3103 1205 2759 \nQ 947 2416 947 1791 \nQ 947 1169 1205 825 \nQ 1463 481 1925 481 \nQ 2391 481 2648 825 \nQ 2906 1169 2906 1791 \nz\nM 3481 434 \nQ 3481 -459 3084 -895 \nQ 2688 -1331 1869 -1331 \nQ 1566 -1331 1297 -1286 \nQ 1028 -1241 775 -1147 \nL 775 -588 \nQ 1028 -725 1275 -790 \nQ 1522 -856 1778 -856 \nQ 2344 -856 2625 -561 \nQ 2906 -266 2906 331 \nL 2906 616 \nQ 2728 306 2450 153 \nQ 2172 0 1784 0 \nQ 1141 0 747 490 \nQ 353 981 353 1791 \nQ 353 2603 747 3093 \nQ 1141 3584 1784 3584 \nQ 2172 3584 2450 3431 \nQ 2728 3278 2906 2969 \nL 2906 3500 \nL 3481 3500 \nL 3481 434 \nz\n\" transform=\"scale(0.015625)\"/>\n     </defs>\n     <use xlink:href=\"#DejaVuSans-6e\"/>\n     <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(63.378906 0)\"/>\n     <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(124.560547 0)\"/>\n     <use xlink:href=\"#DejaVuSans-62\" transform=\"translate(156.347656 0)\"/>\n     <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(219.824219 0)\"/>\n     <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(281.103516 0)\"/>\n     <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(308.886719 0)\"/>\n     <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(370.166016 0)\"/>\n     <use xlink:href=\"#DejaVuSans-63\" transform=\"translate(433.544922 0)\"/>\n     <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(488.525391 0)\"/>\n     <use xlink:href=\"#DejaVuSans-6e\" transform=\"translate(516.308594 0)\"/>\n     <use xlink:href=\"#DejaVuSans-67\" transform=\"translate(579.6875 0)\"/>\n    </g>\n   </g>\n   <g id=\"legend_1\">\n    <g id=\"patch_23\">\n     <path d=\"M 117.857445 59.674375 \nL 184.487132 59.674375 \nQ 186.487132 59.674375 186.487132 57.674375 \nL 186.487132 29.318125 \nQ 186.487132 27.318125 184.487132 27.318125 \nL 117.857445 27.318125 \nQ 115.857445 27.318125 115.857445 29.318125 \nL 115.857445 57.674375 \nQ 115.857445 59.674375 117.857445 59.674375 \nz\n\" style=\"fill: #ffffff; opacity: 0.8; stroke: #cccccc; stroke-linejoin: miter\"/>\n    </g>\n    <g id=\"patch_24\">\n     <path d=\"M 119.857445 38.916562 \nL 139.857445 38.916562 \nL 139.857445 31.916562 \nL 119.857445 31.916562 \nz\n\" style=\"fill: #1f77b4\"/>\n    </g>\n    <g id=\"text_13\">\n     <!-- layer 1 -->\n     <g transform=\"translate(147.857445 38.916562) scale(0.1 -0.1)\">\n      <defs>\n       <path id=\"DejaVuSans-79\" d=\"M 2059 -325 \nQ 1816 -950 1584 -1140 \nQ 1353 -1331 966 -1331 \nL 506 -1331 \nL 506 -850 \nL 844 -850 \nQ 1081 -850 1212 -737 \nQ 1344 -625 1503 -206 \nL 1606 56 \nL 191 3500 \nL 800 3500 \nL 1894 763 \nL 2988 3500 \nL 3597 3500 \nL 2059 -325 \nz\n\" transform=\"scale(0.015625)\"/>\n      </defs>\n      <use xlink:href=\"#DejaVuSans-6c\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(27.783203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(89.0625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(148.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(209.765625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(250.878906 0)\"/>\n      <use xlink:href=\"#DejaVuSans-31\" transform=\"translate(282.666016 0)\"/>\n     </g>\n    </g>\n    <g id=\"patch_25\">\n     <path d=\"M 119.857445 53.594687 \nL 139.857445 53.594687 \nL 139.857445 46.594687 \nL 119.857445 46.594687 \nz\n\" style=\"fill: #ff7f0e\"/>\n    </g>\n    <g id=\"text_14\">\n     <!-- layer 2 -->\n     <g transform=\"translate(147.857445 53.594687) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-6c\"/>\n      <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(27.783203 0)\"/>\n      <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(89.0625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(148.242188 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(209.765625 0)\"/>\n      <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(250.878906 0)\"/>\n      <use xlink:href=\"#DejaVuSans-32\" transform=\"translate(282.666016 0)\"/>\n     </g>\n    </g>\n   </g>\n  </g>\n  <g id=\"axes_2\">\n   <g id=\"patch_26\">\n    <path d=\"M 221.028309 177.550125 \nL 368.734191 177.550125 \nL 368.734191 22.318125 \nL 221.028309 22.318125 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"patch_27\">\n    <path d=\"M 227.742213 177.550125 \nL 234.628268 177.550125 \nL 234.628268 125.206005 \nL 227.742213 125.206005 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_28\">\n    <path d=\"M 244.957351 177.550125 \nL 251.843405 177.550125 \nL 251.843405 125.39885 \nL 244.957351 125.39885 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_29\">\n    <path d=\"M 262.172486 177.550125 \nL 269.058541 177.550125 \nL 269.058541 123.856077 \nL 262.172486 123.856077 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_30\">\n    <path d=\"M 279.387624 177.550125 \nL 286.27368 177.550125 \nL 286.27368 135.541672 \nL 279.387624 135.541672 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_31\">\n    <path d=\"M 296.602761 177.550125 \nL 303.488818 177.550125 \nL 303.488818 132.552548 \nL 296.602761 132.552548 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_32\">\n    <path d=\"M 313.817907 177.550125 \nL 320.703964 177.550125 \nL 320.703964 130.568981 \nL 313.817907 130.568981 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_33\">\n    <path d=\"M 331.033045 177.550125 \nL 337.919101 177.550125 \nL 337.919101 141.143411 \nL 331.033045 141.143411 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_34\">\n    <path d=\"M 348.248182 177.550125 \nL 355.134239 177.550125 \nL 355.134239 129.990439 \nL 348.248182 129.990439 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_35\">\n    <path d=\"M 234.628267 177.550125 \nL 241.514323 177.550125 \nL 241.514323 134.770285 \nL 234.628267 134.770285 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_36\">\n    <path d=\"M 251.843405 177.550125 \nL 258.729462 177.550125 \nL 258.729462 134.233069 \nL 251.843405 134.233069 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_37\">\n    <path d=\"M 269.058543 177.550125 \nL 275.9446 177.550125 \nL 275.9446 111.940902 \nL 269.058543 111.940902 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_38\">\n    <path d=\"M 286.27368 177.550125 \nL 293.159737 177.550125 \nL 293.159737 130.316446 \nL 286.27368 130.316446 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_39\">\n    <path d=\"M 303.488814 177.550125 \nL 310.374866 177.550125 \nL 310.374866 135.578407 \nL 303.488814 135.578407 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_40\">\n    <path d=\"M 320.703956 177.550125 \nL 327.590012 177.550125 \nL 327.590012 141.005663 \nL 320.703956 141.005663 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_41\">\n    <path d=\"M 337.919093 177.550125 \nL 344.80515 177.550125 \nL 344.80515 128.888458 \nL 337.919093 128.888458 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_42\">\n    <path d=\"M 355.134231 177.550125 \nL 362.020287 177.550125 \nL 362.020287 127.524758 \nL 355.134231 127.524758 \nz\n\" clip-path=\"url(#pef14551727)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"matplotlib.axis_3\">\n    <g id=\"xtick_5\">\n     <g id=\"line2d_11\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"234.628268\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_15\">\n      <!-- 0.0 -->\n      <g transform=\"translate(226.676705 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_6\">\n     <g id=\"line2d_12\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"277.666112\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_16\">\n      <!-- 2.5 -->\n      <g transform=\"translate(269.714549 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-32\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_7\">\n     <g id=\"line2d_13\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"320.703956\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_17\">\n      <!-- 5.0 -->\n      <g transform=\"translate(312.752393 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-35\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_8\">\n     <g id=\"line2d_14\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"363.7418\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_18\">\n      <!-- 7.5 -->\n      <g transform=\"translate(355.790237 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-37\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_19\">\n     <!-- expert -->\n     <g transform=\"translate(278.666406 205.826687) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-65\"/>\n      <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(59.773438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(118.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(182.429688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(243.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(285.066406 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_4\">\n    <g id=\"ytick_6\">\n     <g id=\"line2d_15\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"221.028309\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_7\">\n     <g id=\"line2d_16\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"221.028309\" y=\"139.935824\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_8\">\n     <g id=\"line2d_17\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"221.028309\" y=\"102.321523\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_9\">\n     <g id=\"line2d_18\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"221.028309\" y=\"64.707222\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_10\">\n     <g id=\"line2d_19\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"221.028309\" y=\"27.092921\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_20\">\n    <path d=\"M 221.028309 130.532249 \nL 368.734191 130.532249 \n\" clip-path=\"url(#pef14551727)\" style=\"fill: none; stroke-dasharray: 3.7,1.6; stroke-dashoffset: 0; stroke: #808080\"/>\n   </g>\n   <g id=\"patch_43\">\n    <path d=\"M 221.028309 177.550125 \nL 221.028309 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_44\">\n    <path d=\"M 368.734191 177.550125 \nL 368.734191 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_45\">\n    <path d=\"M 221.028309 177.550125 \nL 368.734191 177.550125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_46\">\n    <path d=\"M 221.028309 22.318125 \nL 368.734191 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"text_20\">\n    <!-- auxiliary loss -->\n    <g transform=\"translate(255.659063 16.318125) scale(0.12 -0.12)\">\n     <defs>\n      <path id=\"DejaVuSans-75\" d=\"M 544 1381 \nL 544 3500 \nL 1119 3500 \nL 1119 1403 \nQ 1119 906 1312 657 \nQ 1506 409 1894 409 \nQ 2359 409 2629 706 \nQ 2900 1003 2900 1516 \nL 2900 3500 \nL 3475 3500 \nL 3475 0 \nL 2900 0 \nL 2900 538 \nQ 2691 219 2414 64 \nQ 2138 -91 1772 -91 \nQ 1169 -91 856 284 \nQ 544 659 544 1381 \nz\nM 1991 3584 \nL 1991 3584 \nz\n\" transform=\"scale(0.015625)\"/>\n     </defs>\n     <use xlink:href=\"#DejaVuSans-61\"/>\n     <use xlink:href=\"#DejaVuSans-75\" transform=\"translate(61.279297 0)\"/>\n     <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(124.658203 0)\"/>\n     <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(183.837891 0)\"/>\n     <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(211.621094 0)\"/>\n     <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(239.404297 0)\"/>\n     <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(267.1875 0)\"/>\n     <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(328.466797 0)\"/>\n     <use xlink:href=\"#DejaVuSans-79\" transform=\"translate(369.580078 0)\"/>\n     <use xlink:href=\"#DejaVuSans-20\" transform=\"translate(428.759766 0)\"/>\n     <use xlink:href=\"#DejaVuSans-6c\" transform=\"translate(460.546875 0)\"/>\n     <use xlink:href=\"#DejaVuSans-6f\" transform=\"translate(488.330078 0)\"/>\n     <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(549.511719 0)\"/>\n     <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(601.611328 0)\"/>\n    </g>\n   </g>\n  </g>\n  <g id=\"axes_3\">\n   <g id=\"patch_47\">\n    <path d=\"M 398.275368 177.550125 \nL 545.98125 177.550125 \nL 545.98125 22.318125 \nL 398.275368 22.318125 \nz\n\" style=\"fill: #ffffff\"/>\n   </g>\n   <g id=\"patch_48\">\n    <path d=\"M 404.989271 177.550125 \nL 411.875327 177.550125 \nL 411.875327 127.974736 \nL 404.989271 127.974736 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_49\">\n    <path d=\"M 422.204409 177.550125 \nL 429.090464 177.550125 \nL 429.090464 136.377342 \nL 422.204409 136.377342 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_50\">\n    <path d=\"M 439.419545 177.550125 \nL 446.3056 177.550125 \nL 446.3056 130.206247 \nL 439.419545 130.206247 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_51\">\n    <path d=\"M 456.634683 177.550125 \nL 463.520739 177.550125 \nL 463.520739 132.437758 \nL 456.634683 132.437758 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_52\">\n    <path d=\"M 473.84982 177.550125 \nL 480.735877 177.550125 \nL 480.735877 124.494309 \nL 473.84982 124.494309 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_53\">\n    <path d=\"M 491.064966 177.550125 \nL 497.951023 177.550125 \nL 497.951023 130.656222 \nL 491.064966 130.656222 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_54\">\n    <path d=\"M 508.280104 177.550125 \nL 515.16616 177.550125 \nL 515.16616 128.387979 \nL 508.280104 128.387979 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_55\">\n    <path d=\"M 525.495241 177.550125 \nL 532.381298 177.550125 \nL 532.381298 133.723404 \nL 525.495241 133.723404 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #1f77b4\"/>\n   </g>\n   <g id=\"patch_56\">\n    <path d=\"M 411.875326 177.550125 \nL 418.761382 177.550125 \nL 418.761382 126.225336 \nL 411.875326 126.225336 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_57\">\n    <path d=\"M 429.090464 177.550125 \nL 435.976521 177.550125 \nL 435.976521 133.213738 \nL 429.090464 133.213738 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_58\">\n    <path d=\"M 446.305602 177.550125 \nL 453.191658 177.550125 \nL 453.191658 126.482469 \nL 446.305602 126.482469 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_59\">\n    <path d=\"M 463.520739 177.550125 \nL 470.406796 177.550125 \nL 470.406796 133.952983 \nL 463.520739 133.952983 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_60\">\n    <path d=\"M 480.735873 177.550125 \nL 487.621925 177.550125 \nL 487.621925 131.882177 \nL 480.735873 131.882177 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_61\">\n    <path d=\"M 497.951014 177.550125 \nL 504.837071 177.550125 \nL 504.837071 137.245153 \nL 497.951014 137.245153 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_62\">\n    <path d=\"M 515.166152 177.550125 \nL 522.052209 177.550125 \nL 522.052209 130.771012 \nL 515.166152 130.771012 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"patch_63\">\n    <path d=\"M 532.38129 177.550125 \nL 539.267346 177.550125 \nL 539.267346 124.485128 \nL 532.38129 124.485128 \nz\n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: #ff7f0e\"/>\n   </g>\n   <g id=\"matplotlib.axis_5\">\n    <g id=\"xtick_9\">\n     <g id=\"line2d_21\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"411.875327\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_21\">\n      <!-- 0.0 -->\n      <g transform=\"translate(403.923764 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-30\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_10\">\n     <g id=\"line2d_22\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"454.91317\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_22\">\n      <!-- 2.5 -->\n      <g transform=\"translate(446.961608 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-32\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_11\">\n     <g id=\"line2d_23\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"497.951014\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_23\">\n      <!-- 5.0 -->\n      <g transform=\"translate(489.999452 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-35\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-30\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"xtick_12\">\n     <g id=\"line2d_24\">\n      <g>\n       <use xlink:href=\"#md4ed840694\" x=\"540.988858\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n     <g id=\"text_24\">\n      <!-- 7.5 -->\n      <g transform=\"translate(533.037296 192.148562) scale(0.1 -0.1)\">\n       <use xlink:href=\"#DejaVuSans-37\"/>\n       <use xlink:href=\"#DejaVuSans-2e\" transform=\"translate(63.623047 0)\"/>\n       <use xlink:href=\"#DejaVuSans-35\" transform=\"translate(95.410156 0)\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"text_25\">\n     <!-- expert -->\n     <g transform=\"translate(455.913465 205.826687) scale(0.1 -0.1)\">\n      <use xlink:href=\"#DejaVuSans-65\"/>\n      <use xlink:href=\"#DejaVuSans-78\" transform=\"translate(59.773438 0)\"/>\n      <use xlink:href=\"#DejaVuSans-70\" transform=\"translate(118.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-65\" transform=\"translate(182.429688 0)\"/>\n      <use xlink:href=\"#DejaVuSans-72\" transform=\"translate(243.953125 0)\"/>\n      <use xlink:href=\"#DejaVuSans-74\" transform=\"translate(285.066406 0)\"/>\n     </g>\n    </g>\n   </g>\n   <g id=\"matplotlib.axis_6\">\n    <g id=\"ytick_11\">\n     <g id=\"line2d_25\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"398.275368\" y=\"177.550125\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_12\">\n     <g id=\"line2d_26\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"398.275368\" y=\"139.935824\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_13\">\n     <g id=\"line2d_27\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"398.275368\" y=\"102.321523\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_14\">\n     <g id=\"line2d_28\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"398.275368\" y=\"64.707222\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n    <g id=\"ytick_15\">\n     <g id=\"line2d_29\">\n      <g>\n       <use xlink:href=\"#m45dbd8aeb0\" x=\"398.275368\" y=\"27.092921\" style=\"stroke: #000000; stroke-width: 0.8\"/>\n      </g>\n     </g>\n    </g>\n   </g>\n   <g id=\"line2d_30\">\n    <path d=\"M 398.275368 130.532249 \nL 545.98125 130.532249 \n\" clip-path=\"url(#p3cdad7409a)\" style=\"fill: none; stroke-dasharray: 3.7,1.6; stroke-dashoffset: 0; stroke: #808080\"/>\n   </g>\n   <g id=\"patch_64\">\n    <path d=\"M 398.275368 177.550125 \nL 398.275368 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_65\">\n    <path d=\"M 545.98125 177.550125 \nL 545.98125 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_66\">\n    <path d=\"M 398.275368 177.550125 \nL 545.98125 177.550125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"patch_67\">\n    <path d=\"M 398.275368 22.318125 \nL 545.98125 22.318125 \n\" style=\"fill: none; stroke: #000000; stroke-width: 0.8; stroke-linejoin: miter; stroke-linecap: square\"/>\n   </g>\n   <g id=\"text_26\">\n    <!-- bias -->\n    <g transform=\"translate(459.849871 16.318125) scale(0.12 -0.12)\">\n     <use xlink:href=\"#DejaVuSans-62\"/>\n     <use xlink:href=\"#DejaVuSans-69\" transform=\"translate(63.476562 0)\"/>\n     <use xlink:href=\"#DejaVuSans-61\" transform=\"translate(91.259766 0)\"/>\n     <use xlink:href=\"#DejaVuSans-73\" transform=\"translate(152.539062 0)\"/>\n    </g>\n   </g>\n  </g>\n </g>\n <defs>\n  <clipPath id=\"pf838c49614\">\n   <rect x=\"43.78125\" y=\"22.318125\" width=\"147.705882\" height=\"155.232\"/>\n  </clipPath>\n  <clipPath id=\"pef14551727\">\n   <rect x=\"221.028309\" y=\"22.318125\" width=\"147.705882\" height=\"155.232\"/>\n  </clipPath>\n  <clipPath id=\"p3cdad7409a\">\n   <rect x=\"398.275368\" y=\"22.318125\" width=\"147.705882\" height=\"155.232\"/>\n  </clipPath>\n </defs>\n</svg>\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-8 fw=all -->\n",
    "\n",
    "The left panel is the collapse, caught red-handed. Without balancing, a\n",
    "handful of experts absorb most of the traffic and several fall below 2%\n",
    "usage — dead weight exactly as the feedback-loop argument predicted.\n",
    "How hard the collapse bites varies with the seed (rerunning, we have\n",
    "seen anywhere from two dead experts to a layer served entirely by a\n",
    "single one), but it appears in every run of every seed in both\n",
    "frameworks. The other two panels show both repairs doing their job: the\n",
    "auxiliary loss and the thermostat flatten usage onto the dashed uniform\n",
    "line, with no expert dead and none dominant, and they do it equally\n",
    "well.\n",
    "\n",
    "The training-loss column tells the quieter half of the story, and it\n",
    "needs the more careful reading. A collapsed model trains with only a\n",
    "fraction of its FFN parameters participating, and the runs bear the\n",
    "prediction out directionally: in every seed we tried, in both\n",
    "frameworks, the unbalanced run finished above both balanced ones. But\n",
    "the margin is anywhere from a few hundredths to a few tenths of a nat\n",
    "depending on the seed — the same order as the run-to-run spread of any\n",
    "single variant — so take the direction from the loss column and the\n",
    "conviction from the histograms. Between the two repairs our experiment\n",
    "offers no verdict at all: their loss ranges overlap across seeds, and\n",
    "the summary our data supports is that both recover the lost capacity\n",
    "equally well here. The reason DeepSeek prefers the bias is the interference\n",
    "argument above, and that is a frontier-scale, small-margin effect\n",
    "[@Wang.Gao.Zeng.ea.2024], far below our noise floor; we say so\n",
    "rather than pretend our 800 steps could adjudicate it.\n",
    "\n",
    "## Fine-Grained and Shared Experts\n",
    "\n",
    "Mixtral's experts are full-width FFNs, eight complete copies with two\n",
    "active: big specialists, coarsely chosen. The DeepSeek line of models\n",
    "argues for slicing the same budget thinner\n",
    "[@Dai.Deng.Zhao.ea.2024]. Hold the stored parameters and the\n",
    "active parameters fixed, but make each expert several times narrower\n",
    "and multiply the count and $k$ accordingly: instead of choosing 2 of 8\n",
    "experts, choose 8 of 64 quarter-width ones. The active-parameter\n",
    "accounting is identical; what changes is *composability*. With $k$\n",
    "narrow experts a token assembles its FFN from\n",
    "$\\binom{64}{8} \\approx 4 \\times 10^{9}$ combinations rather than\n",
    "$\\binom{8}{2} = 28$, so specializations can mix and match — one slice\n",
    "for chemistry vocabulary and another for parsing nested clauses,\n",
    "combined on demand rather than fused into one wide monolith.\n",
    "\n",
    "The second DeepSeek refinement removes a redundancy that fine-grained\n",
    "routing exposes. Every token, whatever its topic, needs some common\n",
    "processing, and with purely routed experts each one must dedicate some\n",
    "of its width to those shared functions. So DeepSeek adds one *shared\n",
    "expert* that every token passes through unconditionally, alongside its\n",
    "$k$ routed choices: the common knowledge lives once, in the always-on\n",
    "path, and the routed experts spend their entire width on what actually\n",
    "distinguishes them. Neither refinement changes a line of our layer's\n",
    "logic — narrower experts change a constructor argument, and the shared\n",
    "expert is one more FFN added outside the router's jurisdiction (the\n",
    "final exercise builds it). We do not race granularities here: at our\n",
    "two-block scale the difference sits inside seed noise, and the design's\n",
    "justification is a frontier-scale ablation\n",
    "[@Dai.Deng.Zhao.ea.2024], not a claim our 180 KB corpus can test.\n",
    "What the deployed configurations converged on, however, is worth\n",
    "reading as data:\n",
    "\n",
    ":Routed-expert configurations of deployed MoE language models\n",
    "\n",
    "| Model | Experts per layer | Active $k$ | Expert width / $d$ | Shared expert |\n",
    "|---|---|---|---|---|\n",
    "| Mixtral 8x7B [@Jiang.Sablayrolles.Roux.ea.2024] | 8 | 2 | 3.5 | none |\n",
    "| DeepSeek-V3 [@Liu.Feng.Xue.ea.2024] | 256 | 8 | 0.29 | 1 |\n",
    "| Qwen3-235B [@Yang.Li.Yang.ea.2025] | 128 | 8 | 0.38 | none |\n",
    "| gpt-oss-120b [@OpenAI.2025] | 128 | 4 | 1.0 | none |\n",
    "\n",
    "The trend since Mixtral runs toward many narrow experts with moderate\n",
    "$k$; the shared expert has advocates (DeepSeek since V2\n",
    "[@DeepSeek-AI.2024]) and abstainers (Qwen3 dropped it after using\n",
    "one in its previous generation). As with the block architecture in\n",
    "that section, the field agrees on the axis and still\n",
    "argues about one notch on the dial.\n",
    "\n",
    "## A Mixture-of-Experts GPT\n",
    "\n",
    "Time to collect. We train two models on the same data for the same 600\n",
    "steps: the dense `d2l.GPT` of that section at width 256 with four\n",
    "blocks, and the same model with every FFN replaced by a mixture of\n",
    "eight experts, one active, balanced by the thermostat. Because the\n",
    "active expert is a full-width `FeedForward`, the two models do the\n",
    "same per-token compute and hold the same *active* parameters up to the\n",
    "router's rounding error; the MoE simply holds seven spare experts per\n",
    "block. Since that section showed this corpus drives models into\n",
    "memorization within a few hundred steps, we compare the *best*\n",
    "validation loss over the budget, as that section did."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "moe-a-mixture-of-experts-gpt"
   },
   "source": [
    "count = lambda m: sum(p.numel() for p in m.parameters())\n",
    "for name in ('dense', 'MoE'):\n",
    "    torch.manual_seed(0)\n",
    "    if name == 'dense':\n",
    "        model = d2l.GPT(len(data.vocab), num_hiddens=256, num_heads=8,\n",
    "                        num_blks=4)\n",
    "    else:\n",
    "        model = moe_gpt(len(data.vocab), num_hiddens=256, num_blks=4,\n",
    "                        num_experts=8, num_active=1, num_heads=8)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3,\n",
    "                                  weight_decay=0.0)\n",
    "    best = float('inf')\n",
    "    for chunk in range(6):\n",
    "        if name == 'dense':\n",
    "            losses = d2l.train_lm(model, data, optimizer, 100)\n",
    "        else:\n",
    "            losses = train_moe_lm(model, data, optimizer, 100,\n",
    "                                  bias_rate=0.01)\n",
    "        best = min(best, val_loss(model, data))\n",
    "    active = count(model)\n",
    "    if name == 'MoE':\n",
    "        active -= sum((blk.ffn.num_experts - blk.ffn.num_active)\n",
    "                      * count(blk.ffn.experts[0]) for blk in model.blks)\n",
    "    print(f'{name:>5}: {count(model)/1e6:5.2f}M parameters '\n",
    "          f'({active/1e6:.2f}M active), best validation {best:.2f}, '\n",
    "          f'final training loss {sum(losses[-50:])/50:.2f}')"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "dense:  3.16M parameters (3.16M active), best validation 1.51, final training loss 0.21\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "  MoE: 17.85M parameters (3.16M active), best validation 1.52, final training loss 0.19\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=moe-md-9 fw=all -->\n",
    "\n",
    "The mixture holds between five and six times the parameters of the\n",
    "dense baseline at the same active count, and the comparison comes out\n",
    "the way our data budget says it must. Best validation loss is\n",
    "indistinguishable: both land near 1.5, within the couple-of-hundredths\n",
    "spread that reseeding produces. The training loss ends slightly lower\n",
    "for the MoE, consistently across seeds and frameworks: the spare\n",
    "experts are real capacity, and on a corpus this small, real capacity\n",
    "spends itself on memorizing faster (that section). That is the\n",
    "claim our scale supports: matched on active compute, the MoE trains\n",
    "comparably while carrying several times the parameters, and nothing\n",
    "more. The regime where the extra capacity pays on *validation* is the\n",
    "one the accounting cell described: corpora that outweigh any dense\n",
    "model you could afford to serve, where MoEs reach a given loss at a\n",
    "fraction of the training FLOPs of dense models\n",
    "[@fedus2022switch; @Dai.Deng.Zhao.ea.2024] and dominate the\n",
    "deployed frontier. What our small-scale run demonstrates is the\n",
    "mechanism working end to end — routed capacity added, balance held by\n",
    "a controller, quality preserved at fixed per-token cost — which is\n",
    "precisely the part that does transfer across five orders of magnitude.\n",
    "\n",
    "## Summary\n",
    "\n",
    "A mixture of experts replaces the transformer block's FFN with $E$\n",
    "parallel FFNs and a linear router that sends each token to the top $k$\n",
    "of them, weighted by routing probability: parameters scale with $E$,\n",
    "per-token FLOPs with $k$. The accounting is the argument, computed\n",
    "from published configurations: Mixtral 8x7B stores 46.7B parameters\n",
    "and activates 12.9B per token, and DeepSeek-V3 stores 28 times what a\n",
    "token touches. Left alone, routing collapses: winners attract\n",
    "gradient, gradient makes winners, and a few experts end up serving\n",
    "everything while the rest die. The auxiliary balancing loss penalizes\n",
    "the correlation between load and routing probability at the price of a\n",
    "gradient that competes with the language-modeling objective; the\n",
    "auxiliary-loss-free bias steers only the top-$k$ selection through a\n",
    "gradient-free control loop. In our triptych both flatten usage\n",
    "completely and both beat the collapsed run, while differing from each\n",
    "other by less than seed noise. Modern designs slice capacity thinner\n",
    "(many narrow experts, sometimes one shared always-on expert), and our\n",
    "GPT accepts the whole apparatus through a swapped FFN: at matched\n",
    "active parameters, the mixture trains comparably to dense while\n",
    "holding several times the weights — the trade that has made MoE a\n",
    "major scaling strategy at the frontier.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. The top-$k$ selection is an argmax and has no gradient, yet the\n",
    "   router trains. Write the layer's output for a single token as\n",
    "   $\\mathbf{y} = \\sum_{i \\in \\mathcal{E}_k} p_i\\, \\mathrm{FFN}_i(\\mathbf{x})$\n",
    "   and derive\n",
    "   $\\partial \\mathcal{L} / \\partial z_j$ for the router logits\n",
    "   $\\mathbf{z} = \\mathbf{W}_r \\mathbf{x}$, treating $\\mathcal{E}_k$ as\n",
    "   constant. Which experts receive routing gradient? Now suppose the\n",
    "   gate weights are renormalized over the selected set,\n",
    "   $g_i = p_i / \\sum_{j \\in \\mathcal{E}_k} p_j$, as Mixtral does. Show\n",
    "   that for $k = 1$ the router's gradient vanishes identically — and\n",
    "   explain why Mixtral's $k = 2$ escapes this trap.\n",
    "2. Production MoE systems give each expert a fixed token buffer of\n",
    "   capacity $c \\cdot N k / E$ for a batch of $N$ tokens (the *capacity\n",
    "   factor* $c$); tokens routed to a full expert are *dropped* — they\n",
    "   skip the FFN and ride the residual stream unchanged\n",
    "   [@Lepikhin.Lee.Xu.ea.2021; @fedus2022switch]. Explain why static\n",
    "   buffers are needed at all (consider batching across devices), what\n",
    "   happens to dropped tokens' representations, and why a *batch-level*\n",
    "   balance objective matters for hardware efficiency even if the model\n",
    "   could tolerate statistical imbalance. What does $c = 1$ force the\n",
    "   router to be? What does large $c$ cost?\n",
    "3. Rerun the triptych with `num_active=2`. Does the unbalanced run\n",
    "   still collapse, and as hard? Explain the difference using the\n",
    "   feedback-loop argument: how does a second active expert change the\n",
    "   gradient received by the runners-up?\n",
    "4. Sweep the auxiliary-loss weight $\\alpha$ over\n",
    "   $\\{0, 10^{-3}, 10^{-2}, 10^{-1}, 1\\}$. Plot final training loss and\n",
    "   the number of dead experts against $\\alpha$. Where does balancing\n",
    "   stop improving and interference start hurting? Then sweep the bias\n",
    "   update speed $u$ over the same grid: which failure mode appears at\n",
    "   large $u$, and why is it gentler than large $\\alpha$?\n",
    "5. Our layer computes every expert densely and masks. Implement the\n",
    "   gathered alternative in PyTorch: for each expert, `torch.where` the\n",
    "   indices of its assigned tokens, run the expert on just those rows,\n",
    "   and scatter the results back. Verify it matches the dense layer's\n",
    "   output, then measure at what expert count the gathered version\n",
    "   becomes faster on your hardware.\n",
    "6. Add a DeepSeek-style shared expert: one additional `FeedForward`\n",
    "   applied to every token, its output added to the routed mixture\n",
    "   (no gate). Compare against the plain layer at matched *active*\n",
    "   parameters (shrink $k$ or the widths accordingly) on the triptych\n",
    "   task. Does the shared expert change how quickly the routed experts\n",
    "   specialize — measure usage entropy over training?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-pytorch",
   "display_name": "d2l (PyTorch)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
