| """Build report.pdf from the metrics written by the three parts. |
| |
| Every number in the document is read out of ``runs/*/metrics.json`` rather |
| than typed in, so the report cannot drift from the runs that produced it. |
| |
| python -m src.report |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| from reportlab.lib import colors |
| from reportlab.lib.enums import TA_JUSTIFY |
| from reportlab.lib.pagesizes import A4 |
| from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet |
| from reportlab.lib.units import cm |
| from reportlab.platypus import (Image, KeepTogether, PageBreak, Paragraph, |
| SimpleDocTemplate, Table, TableStyle) |
|
|
| from .part1.config import TransformerConfig |
| from .utils import ASSETS, ROOT, RUNS, load_json |
|
|
| INK = colors.HexColor("#0b0b0b") |
| MUTED = colors.HexColor("#52514e") |
| RULE = colors.HexColor("#d8d7d2") |
| BAND = colors.HexColor("#f4f3f0") |
| ACCENT = colors.HexColor("#2a78d6") |
|
|
| PAGE_W = A4[0] - 4 * cm |
|
|
|
|
| def styles(): |
| ss = getSampleStyleSheet() |
| s = { |
| "title": ParagraphStyle("title", parent=ss["Title"], fontSize=19, |
| leading=23, textColor=INK, spaceAfter=2), |
| "sub": ParagraphStyle("sub", parent=ss["Normal"], fontSize=10.5, |
| leading=14, textColor=MUTED, alignment=1, |
| spaceAfter=14), |
| "h1": ParagraphStyle("h1", parent=ss["Heading1"], fontSize=14, |
| leading=17, textColor=INK, spaceBefore=14, |
| spaceAfter=7), |
| "h2": ParagraphStyle("h2", parent=ss["Heading2"], fontSize=11.5, |
| leading=14, textColor=INK, spaceBefore=11, |
| spaceAfter=5), |
| "h3": ParagraphStyle("h3", parent=ss["Heading3"], fontSize=10, |
| leading=13, textColor=ACCENT, spaceBefore=8, |
| spaceAfter=3), |
| "body": ParagraphStyle("body", parent=ss["BodyText"], fontSize=9.3, |
| leading=13.2, textColor=INK, |
| alignment=TA_JUSTIFY, spaceAfter=6), |
| "cap": ParagraphStyle("cap", parent=ss["Normal"], fontSize=8, |
| leading=10.5, textColor=MUTED, spaceBefore=3, |
| spaceAfter=10), |
| "code": ParagraphStyle("code", parent=ss["Normal"], fontSize=8, |
| leading=11, textColor=INK, |
| fontName="Courier", spaceAfter=6), |
| } |
| return s |
|
|
|
|
| S = None |
|
|
|
|
| def P(text, kind="body"): |
| return Paragraph(text, S[kind]) |
|
|
|
|
| def caption(text): |
| return Paragraph(text, S["cap"]) |
|
|
|
|
| def figure(name, width=PAGE_W, cap=None): |
| """Place a PNG scaled to the text width, with its caption kept attached.""" |
| path = ASSETS / name |
| if not path.exists(): |
| return P(f"[missing figure: {name}]") |
| from PIL import Image as PILImage |
| w, h = PILImage.open(path).size |
| img = Image(str(path), width=width, height=width * h / w) |
| parts = [img] |
| if cap: |
| parts.append(caption(cap)) |
| return KeepTogether(parts) |
|
|
|
|
| def table(rows, widths=None, align_right_from=1, font=7.6, header=True): |
| t = Table(rows, colWidths=widths, hAlign="LEFT") |
| style = [ |
| ("FONTSIZE", (0, 0), (-1, -1), font), |
| ("TEXTCOLOR", (0, 0), (-1, -1), INK), |
| ("LINEBELOW", (0, 0), (-1, 0), 0.7, RULE if not header else MUTED), |
| ("LINEBELOW", (0, 1), (-1, -2), 0.3, RULE), |
| ("LINEBELOW", (0, -1), (-1, -1), 0.7, MUTED), |
| ("VALIGN", (0, 0), (-1, -1), "MIDDLE"), |
| ("TOPPADDING", (0, 0), (-1, -1), 3), |
| ("BOTTOMPADDING", (0, 0), (-1, -1), 3), |
| ("LEFTPADDING", (0, 0), (-1, -1), 4), |
| ("RIGHTPADDING", (0, 0), (-1, -1), 4), |
| ("ALIGN", (align_right_from, 0), (-1, -1), "RIGHT"), |
| ] |
| if header: |
| style += [("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"), |
| ("BACKGROUND", (0, 0), (-1, 0), BAND)] |
| t.setStyle(TableStyle(style)) |
| return t |
|
|
|
|
| def cell(text, bold=False, align="left", size=7.4): |
| """A table cell that wraps and parses markup. |
| |
| Plain strings in a reportlab Table do neither: markup shows up literally |
| and long text runs off the edge of the column. |
| """ |
| return Paragraph( |
| f"<b>{text}</b>" if bold else str(text), |
| ParagraphStyle("cell", fontName="Helvetica", fontSize=size, |
| leading=size * 1.35, textColor=INK, |
| alignment={"left": 0, "right": 2, "center": 1}[align])) |
|
|
|
|
| def wrap_rows(rows, wrap_cols, header=True): |
| """Wrap the given column indices in every body row (and the header).""" |
| out = [] |
| for i, r in enumerate(rows): |
| out.append([cell(v, bold=(i == 0 and header)) if j in wrap_cols else v |
| for j, v in enumerate(r)]) |
| return out |
|
|
|
|
| def kv_note(text): |
| """A quiet, indented aside.""" |
| return Paragraph(f"<i>{text}</i>", ParagraphStyle( |
| "note", parent=S["body"], textColor=MUTED, fontSize=8.6, leading=12, |
| leftIndent=10, spaceAfter=8)) |
|
|
|
|
| def footer(canvas, doc): |
| canvas.saveState() |
| canvas.setFont("Helvetica", 7.5) |
| canvas.setFillColor(MUTED) |
| canvas.drawCentredString(A4[0] / 2, 1.2 * cm, str(doc.page)) |
| canvas.restoreState() |
|
|
|
|
| |
| |
| |
| P1_ORDER = ["c1_dense", "c2_moe_4e1a", "c3_moe_4e2a", "c4_moe_shared", |
| "c5_moe_4e2a_wide"] |
| P1_NAME = {"c1_dense": "C1 dense MLP", |
| "c2_moe_4e1a": "C2 — 4 experts, top-1", |
| "c3_moe_4e2a": "C3 — 4 experts, top-2", |
| "c4_moe_shared": "C4 — 1 shared + 3 routed, top-1", |
| "c5_moe_4e2a_wide": "C5 — 4 experts, top-2, wide"} |
|
|
|
|
| def load_p1(): |
| return {n: load_json(RUNS / "part1" / n / "metrics.json") |
| for n in P1_ORDER if (RUNS / "part1" / n / "metrics.json").exists()} |
|
|
|
|
| def part1(M): |
| f = [] |
| a = f.append |
| a(P("1. Mixture of Experts", "h1")) |
|
|
| a(P("Task and setup", "h2")) |
| a(P( |
| "The model is a decoder-only transformer and translation is posed as " |
| "plain language modelling over a single sequence, " |
| "<font face='Courier' size='8'><bos> <vi> {source} <en> " |
| "{target} <eos></font>. The cross-entropy is masked to the English " |
| "target tokens and the closing <font face='Courier' size='8'><eos>" |
| "</font>; the source is read but never predicted. Because routing " |
| "happens on the forward pass regardless of whether a token contributes " |
| "to the loss, the source tokens are still dispatched to experts, which " |
| "is what makes the per-language routing analysis in section 1.5 " |
| "possible at all.")) |
| a(P( |
| "Each of the ~446k rows in " |
| "<font face='Courier' size='8'>belumind/en-vi-ja-curated-500k-triplets" |
| "</font> yields two training examples, a Vietnamese–English pair " |
| "and a Japanese–English pair. Examples are packed back-to-back " |
| "into 256-token blocks rather than padded individually: a typical pair " |
| "is around 60 tokens, so per-example padding would have spent roughly " |
| "three quarters of every batch on padding, and \"the same number of " |
| "training tokens\" would have degenerated into \"the same number of " |
| "padding tokens\".")) |
| a(P( |
| "The tokenizer is a 16k byte-level BPE trained jointly on all three " |
| "languages. Byte-level is a requirement rather than a preference here: " |
| "the dataset is not romanised, so a vocabulary built on an ASCII " |
| "assumption would fail on Vietnamese diacritics and on Japanese kana " |
| "and kanji entirely. Operating over raw UTF-8 bytes means every string " |
| "is representable with no <font face='Courier' size='8'><unk>" |
| "</font>, and BPE then spends its merge budget on whatever multi-byte " |
| "sequences are frequent — in practice, common kanji compounds and " |
| "kana bigrams.")) |
|
|
| a(P("1.1 Holding the parameter budget fixed", "h2")) |
| a(P( |
| "A single 2-layer MLP of hidden width <i>d</i><sub>ff</sub> costs " |
| "2·<i>d</i><sub>model</sub>·<i>d</i><sub>ff</sub> weights. " |
| "Writing <i>E</i> for the number of experts and <i>k</i> for how many " |
| "are active per token, an MoE layer built from the same primitive costs " |
| "<i>E</i> times that in total and <i>k</i> times it per token. The two " |
| "budget rules the assignment asks for therefore reduce to two choices " |
| "of per-expert width: dividing <i>d</i><sub>ff</sub> by <i>E</i> holds " |
| "the <i>total</i> fixed (variants C2–C4), and dividing by <i>k</i> " |
| "holds the <i>active</i> count fixed instead (variant C5), which " |
| "necessarily doubles the total.")) |
|
|
| rows = [["variant", "d_ff/expert", "total params", "active params", |
| "active %", "tokens"]] |
| for n in P1_ORDER: |
| m = M[n] |
| rows.append([ |
| P1_NAME[n].replace("—", "-"), |
| str(TransformerConfig(**m["config"]).d_ff_expert), |
| f"{m['params_total']:,}", f"{m['params_active']:,}", |
| f"{100*m['params_active_frac']:.1f}%", |
| f"{m['tokens_trained']/1e6:.1f}M"]) |
| a(table(wrap_rows(rows, {0}), widths=[5.6*cm, 2.0*cm, 2.6*cm, 2.6*cm, |
| 1.7*cm, 1.6*cm])) |
| a(caption( |
| "Table 1. The five variants. C1–C4 agree on total parameters to " |
| "within 12k (the routers, which C1 does not have); C5's active count " |
| "matches C1's total to within 12k. All five consume an identical 40.0M " |
| "training tokens in the same order, so the only thing that varies is " |
| "the feed-forward block.")) |
|
|
| a(P("1.2 Routing", "h2")) |
| a(P( |
| "A linear router scores the experts per token; the top <i>k</i> gates " |
| "are renormalised to sum to one so that the block's output scale does " |
| "not depend on <i>k</i> — without that, top-2 would feed roughly " |
| "twice the activation magnitude into the residual stream as top-1 and " |
| "the comparison would be measuring a scale change as much as a routing " |
| "change. Two auxiliary terms are added to the loss: a Switch-style " |
| "load-balancing term <i>E</i>·Σ<sub>i</sub> " |
| "<i>f</i><sub>i</sub><i>P</i><sub>i</sub>, and a router z-loss that " |
| "penalises large logits. The balancing term normalises " |
| "<i>f</i><sub>i</sub> by <i>k</i> as well as by the token count, so it " |
| "bottoms out at 1.0 for any <i>k</i>; without that normalisation top-2 " |
| "would carry twice the balancing pressure of top-1 at the same " |
| "coefficient, confounding two changes in what is meant to be a " |
| "one-variable ablation.")) |
|
|
| a(P("1.3 Results", "h2")) |
| rows = [["variant", "test loss", "perplexity", "BLEU", "BLEU vi-en", |
| "BLEU ja-en", "train min"]] |
| for n in P1_ORDER: |
| m = M[n] |
| rows.append([P1_NAME[n].replace("—", "-"), |
| f"{m['test_loss']:.4f}", f"{m['test_ppl']:.3f}", |
| f"{m['bleu_all']:.2f}", f"{m.get('bleu_vi', 0):.2f}", |
| f"{m.get('bleu_ja', 0):.2f}", f"{m['train_minutes']:.1f}"]) |
| a(table(wrap_rows(rows, {0}), widths=[5.2*cm, 1.9*cm, 2.0*cm, 1.6*cm, |
| 2.1*cm, 2.1*cm, 1.7*cm])) |
| a(caption("Table 2. Test-set results. Perplexity is per English target " |
| "token; BLEU is corpus BLEU (sacreBLEU) over 800 greedy " |
| "generations, 400 per source language.")) |
| a(figure("p1_curves.png", cap= |
| "Figure 1. Left: validation loss against tokens seen. Four of the " |
| "five variants are visually indistinguishable at this scale, which " |
| "is itself the headline result. Right: the same curves expressed as " |
| "a gap to the dense baseline, which is where the ablation actually " |
| "lives.")) |
| a(figure("p1_summary.png", cap= |
| "Figure 2. Test perplexity, corpus BLEU, and BLEU split by source " |
| "language.")) |
| return f |
|
|
|
|
| def part1_analysis(M): |
| f = [] |
| a = f.append |
| c1, c2 = M["c1_dense"], M["c2_moe_4e1a"] |
| c3, c4, c5 = M["c3_moe_4e2a"], M["c4_moe_shared"], M["c5_moe_4e2a_wide"] |
|
|
| a(P("1.4 Analysis: what the parameter budget buys", "h2")) |
| a(P( |
| f"<b>Sparsity is not free, and top-1 pays the most for it.</b> C2 routes " |
| f"each token to a single expert of width 512 and reaches perplexity " |
| f"{c2['test_ppl']:.2f} against the dense baseline's {c1['test_ppl']:.2f}, " |
| f"losing {c1['bleu_all'] - c2['bleu_all']:.1f} BLEU " |
| f"({c2['bleu_all']:.2f} vs {c1['bleu_all']:.2f}). It is the only variant " |
| f"whose loss curve separates visibly from the others in Figure 1, and " |
| f"the gap is established early and never closes. The reason is capacity " |
| f"per token rather than capacity overall: C2 holds the same 27.1M total " |
| f"parameters as C1, but any individual token is processed by a " |
| f"512-wide MLP instead of a 2048-wide one, and a quarter of the " |
| f"feed-forward width is simply not enough for this task.")) |
| a(P( |
| f"<b>Top-2 recovers almost all of it.</b> C3 uses the same experts as C2 " |
| f"and changes only how many fire per token, going from 65.2% to 76.8% " |
| f"active parameters. That single change is worth " |
| f"{c2['test_ppl'] - c3['test_ppl']:.2f} perplexity and " |
| f"{c3['bleu_all'] - c2['bleu_all']:.1f} BLEU, landing C3 within " |
| f"{c3['test_ppl'] - c1['test_ppl']:.2f} perplexity and " |
| f"{c1['bleu_all'] - c3['bleu_all']:.1f} BLEU of the dense model while " |
| f"doing roughly three quarters of its feed-forward work. Combining two " |
| f"512-wide experts is evidently a decent substitute for one 2048-wide " |
| f"one; combining one is not.")) |
| a(P( |
| f"<b>The shared expert did not help here.</b> C4 activates exactly as " |
| f"many parameters as C3 ({c4['params_active']/1e6:.2f}M) but spends one " |
| f"of its two active slots on an expert that every token sees. It comes " |
| f"out {c4['test_ppl'] - c3['test_ppl']:.2f} perplexity and " |
| f"{c3['bleu_all'] - c4['bleu_all']:.1f} BLEU behind C3. The intended " |
| f"benefit of a shared expert is that it absorbs the generic " |
| f"transformation every token needs, freeing the routed experts to " |
| f"specialise; the cost is that the router chooses among only three " |
| f"experts instead of four, and only one routed expert's worth of " |
| f"specialised capacity reaches any given token. At this scale the cost " |
| f"dominates. Half of C4's active feed-forward budget is spent on a " |
| f"pathway that is by construction identical for every token, which is " |
| f"the one thing a dense MLP already does well and the one thing " |
| f"conditional computation is supposed to improve on.")) |
| a(P( |
| f"<b>Holding active parameters fixed instead makes MoE a win.</b> C5 is " |
| f"the informative comparison. Its active count " |
| f"({c5['params_active']/1e6:.2f}M) matches the dense baseline's total to " |
| f"within 12k, so the two do essentially the same amount of arithmetic " |
| f"per token, but C5 carries {c5['params_total']/1e6:.1f}M parameters " |
| f"against C1's {c1['params_total']/1e6:.1f}M. It is the best model in " |
| f"the table on every metric: perplexity {c5['test_ppl']:.2f} against " |
| f"{c1['test_ppl']:.2f}, BLEU {c5['bleu_all']:.2f} against " |
| f"{c1['bleu_all']:.2f}, and it wins on both language directions " |
| f"separately. This is the actual argument for MoE, and the ablation " |
| f"isolates it cleanly: extra parameters that are only <i>conditionally</i> " |
| f"used still improve the model, at no extra per-token compute.")) |
| a(P( |
| "The margin is real but modest — around " |
| f"{c5['bleu_all'] - c1['bleu_all']:.1f} BLEU. That is unsurprising at " |
| "this budget. The literature's large MoE gains come from expert counts " |
| "in the tens or hundreds, where the total-to-active ratio is an order " |
| "of magnitude rather than the 1.46× here, and from token budgets " |
| "far past the 40M used for every run in this table.")) |
| a(figure("p1_params.png", cap= |
| "Figure 3. BLEU against active parameters (left) and total " |
| "parameters (right). On the left the four MoE variants sit to the " |
| "left of the dense budget line, doing less work per token; on the " |
| "right C5 is alone to the right of it. The left panel is the " |
| "compute axis and the right panel is the memory axis, and MoE " |
| "trades between them.")) |
|
|
| a(P("1.5 Analysis: division of labour between experts", "h2")) |
| ent = ", ".join(f"{M[n]['expert_usage']['mean_entropy']:.3f}" |
| for n in P1_ORDER if M[n]["expert_usage"]) |
| a(P( |
| f"<b>Load balancing worked, arguably too well.</b> Normalised routing " |
| f"entropy (1.0 = every expert gets an equal share, 0.0 = one expert " |
| f"takes everything) is {ent} for C2–C5 respectively. No expert " |
| f"collapsed and none was starved, which is the failure mode the " |
| f"auxiliary loss exists to prevent. The flip side is that the balancing " |
| f"pressure is itself a constraint on specialisation: the loss " |
| f"explicitly penalises the router for sending one language " |
| f"disproportionately to one expert, so whatever specialisation survives " |
| f"does so against that pressure.")) |
| a(P( |
| "Comparing raw minimum shares across the variants would be misleading, " |
| "since C4 routes among three experts and the others among four, so its " |
| "balanced share is 1/3 rather than 1/4. Measured against each " |
| "variant's own balanced share, the least-used expert receives " |
| + ", ".join( |
| f"{M[n]['expert_usage']['min_share'] * M[n]['config']['n_experts']:.3f}" |
| f" of it in {P1_NAME[n].split()[0]}" |
| for n in P1_ORDER if M[n]["expert_usage"]) |
| + ". Every variant is within 12% of perfect balance, and C5 — the " |
| "best-performing one — is the furthest from it, which is a small " |
| "hint that a little imbalance is the price of useful specialisation " |
| "rather than a defect.")) |
| a(figure("p1_expert_layers.png", cap= |
| "Figure 4. Per-layer expert load, shown as the deviation from the " |
| "balanced share 1/E. Near-neutral everywhere, which is the " |
| "auxiliary loss doing its job.")) |
| a(P( |
| "<b>Specialisation is nonetheless visible, and it is by language.</b> " |
| "Figure 5 conditions on the input language: each column is one " |
| "language's tokens, and the cells give the share of those tokens that " |
| "each expert received. Columns sum to one, so the figure asks \"where " |
| "do this language's tokens go?\" rather than \"which language is most " |
| "frequent?\" — the latter would only recover the corpus mix, since " |
| "English appears in every example and the two source languages in half " |
| "each.")) |
| a(figure("p1_expert_lang.png", cap= |
| "Figure 5. Expert specialisation by input language. Cells show raw " |
| "shares; colour shows the deviation from the balanced share 1/E " |
| "(red = over-used for that language, blue = under-used). A " |
| "perfectly language-agnostic router would render the whole grid " |
| "neutral.")) |
| a(P( |
| "The clearest case is C2, the top-1 router, which has no way to hedge: " |
| "every token commits to exactly one expert. Expert 1 takes 33% of " |
| "Japanese tokens but only 20% of English, while expert 2 takes 32% of " |
| "English and just 12% of Japanese — less than half the balanced " |
| "share. Those two experts have divided Japanese from English between " |
| "them. C3 shows the same axis more weakly (expert 0 at 31% of Japanese " |
| "against 19% of English), which is what one would expect once each " |
| "token gets two experts and can spread its bets.")) |
| a(P( |
| "C5 shows the tidiest structure of the four: expert 0 leans Vietnamese " |
| "(30%), expert 1 leans Japanese (29%), expert 2 leans English (29%), and " |
| "expert 3 sits near neutral on all three. With wider experts there is " |
| "enough capacity per expert for a language-shaped division of labour to " |
| "be worth adopting, and C5 is also the best-performing variant — " |
| "consistent with, though not proof of, the idea that the specialisation " |
| "is doing useful work rather than being an artefact.")) |
| a(P( |
| "The dominant split throughout is <i>script and language</i>, not " |
| "syntax or topic, which is the natural thing for a router reading token " |
| "embeddings to key on: Vietnamese, Japanese and English occupy almost " |
| "disjoint regions of a byte-level BPE vocabulary, so language identity " |
| "is close to linearly decodable from the embedding alone and is by far " |
| "the cheapest signal available to a linear router. C4's routed experts " |
| "show the same pattern on the ja/vi axis (expert 2 at 38% of Japanese " |
| "against 28% of Vietnamese) with English left nearly uniform — " |
| "plausibly because the always-on shared expert is already absorbing the " |
| "English-generation work that the routed experts would otherwise have " |
| "had to divide.")) |
|
|
| a(P("1.6 Translation quality by direction", "h2")) |
| a(P( |
| f"Every variant translates Vietnamese far better than Japanese " |
| f"— {c1['bleu_vi']:.1f} against {c1['bleu_ja']:.1f} BLEU for the " |
| f"dense baseline, and the ~21-point gap holds across all five. Three " |
| f"things contribute: Vietnamese shares the Latin script and much " |
| f"subword structure with English so the shared BPE transfers directly; " |
| f"Vietnamese word order is closer to English than Japanese's verb-final, " |
| f"heavily case-marked order, so less long-range reordering is needed; " |
| f"and Japanese text carries no whitespace, so the tokenizer has to " |
| f"infer segmentation that is given for free in the other two languages. " |
| f"The ranking of the five variants is identical on both directions, so " |
| f"the FFN ablation and the language difficulty are independent effects.")) |
| return f |
|
|
|
|
| |
| |
| |
| P2_ORDER = ["adamw", "mars", "lion", "muon", "sophia"] |
| P2_NAME = {"adamw": "AdamW", "mars": "MARS", "lion": "Lion", "muon": "Muon", |
| "sophia": "Sophia-G"} |
|
|
|
|
| def load_p2(): |
| return {n: load_json(RUNS / "part2" / n / "metrics.json") |
| for n in P2_ORDER if (RUNS / "part2" / n / "metrics.json").exists()} |
|
|
|
|
| def part2(M): |
| f = [] |
| a = f.append |
| a(P("2. Optimizers", "h1")) |
| a(P( |
| "Five optimizers are implemented from scratch, one from each of the " |
| "first four categories of Table 1 of <i>Fantastic Pretraining " |
| "Optimizers and Where to Find Them</i> plus the bonus Hessian-based " |
| "entry. <font face='Courier' size='8'>torch.optim.Optimizer</font> is " |
| "used strictly as a base class, for its parameter-group and state " |
| "bookkeeping; no <font face='Courier' size='8'>torch.optim</font> update " |
| "rule is called anywhere.")) |
| a(P( |
| "The model is the dense baseline of Part 1 retrained from scratch on " |
| "<font face='Courier' size='8'>browndw/human-ai-parallel-corpus</font> " |
| "for next-token prediction. That corpus contains several parallel " |
| "versions of every document — the human original and one " |
| "continuation per generator model — so the split is made on the " |
| "<i>base</i> document id with all versions of a document moving " |
| "together. Splitting rows at random would have put near-paraphrases of " |
| "the same text into both train and test and quietly flattered every " |
| "validation number.")) |
|
|
| a(P("2.1 The update rules", "h2")) |
| a(P("AdamW — category 1, the reference point", "h3")) |
| a(P( |
| "Two exponential moving averages per parameter, of the gradient and of " |
| "its square; the update divides one by the root of the other, making " |
| "the step approximately scale-free per coordinate. The \"W\" is " |
| "decoupled decay: the shrinkage is applied to the weights directly " |
| "rather than folded into the gradient, so it is not itself rescaled by " |
| "√<i>v</i> the way an L2 penalty would be.")) |
| a(kv_note( |
| "Correctness check: this implementation and " |
| "torch.optim.AdamW, run from identical initialisations on identical " |
| "batches, agree to a maximum absolute parameter difference of 3.0e-08 " |
| "after 25 steps — i.e. to float32 rounding.")) |
|
|
| a(P("MARS — category 2, variance-reduced AdamW variant", "h3")) |
| a(P( |
| "AdamW builds its moments from raw minibatch gradients and so inherits " |
| "the full minibatch noise. MARS inserts a STORM-style correction first, " |
| "extrapolating along the difference between consecutive gradients: " |
| "<i>c<sub>t</sub></i> = <i>g<sub>t</sub></i> + " |
| "γ(β<sub>1</sub>/(1−β<sub>1</sub>))" |
| "(<i>g<sub>t</sub></i> − <i>g<sub>t−1</sub></i>). Because " |
| "consecutive gradients differ only by the minibatch draw plus one small " |
| "parameter step, their difference cancels part of the shared sampling " |
| "noise. The corrected gradient is clipped to unit norm — it is a " |
| "difference of two noisy quantities and can blow up early — then " |
| "fed to an otherwise ordinary AdamW. Cost: a third state tensor for the " |
| "previous gradient.")) |
|
|
| a(P("Lion — category 3, memory-efficient", "h3")) |
| a(P( |
| "Lion keeps one state tensor instead of two and takes the <i>sign</i> " |
| "of an interpolated momentum. The detail worth noticing is that it uses " |
| "two different betas: the direction comes from a fresher blend " |
| "(β<sub>1</sub> = 0.9) than the one that is stored " |
| "(β<sub>2</sub> = 0.99), so the step reacts to the current " |
| "gradient while the memory stays long. Because <font face='Courier' " |
| "size='8'>sign</font> makes every coordinate's step exactly " |
| "±<i>lr</i>, the update norm is fixed and much larger than " |
| "AdamW's at the same learning rate — which is why Lion needs its " |
| "own sweep grid rather than inheriting AdamW's, and why its selected " |
| "rate lands an order of magnitude lower.")) |
|
|
| a(P("Muon — category 4, matrix-based", "h3")) |
| a(P( |
| "AdamW normalises each weight coordinate independently, which throws " |
| "away the fact that a weight matrix is a linear map. Muon normalises " |
| "the momentum <i>matrix</i>: it replaces <i>M</i> with the nearest " |
| "semi-orthogonal matrix <i>UV</i><sup>T</sup>, so every singular " |
| "direction is stepped equally and no direction dominates merely because " |
| "it happened to carry a large singular value. An SVD per step would be " |
| "hopeless, so the orthogonalisation uses a quintic Newton–Schulz " |
| "iteration whose coefficients (3.4445, −4.7750, 2.0315) push " |
| "singular values towards 1 aggressively in five steps. They do not " |
| "converge to exactly 1, but Muon only needs the direction. Muon applies " |
| "to 2D hidden matrices; embeddings, the tied head and the 1D gains fall " |
| "back to AdamW, as in the reference implementation.")) |
|
|
| a(P("Sophia-G — category 5, Hessian-based (bonus)", "h3")) |
| a(P( |
| "Sophia divides by an estimate of the Hessian <i>diagonal</i> rather " |
| "than of the gradient second moment — curvature is what actually " |
| "says how far one can step before the loss turns back up. Two things " |
| "make it affordable. The Gauss–Newton–Bartlett estimator: " |
| "for a cross-entropy loss, sampling labels from the model's own " |
| "predictive distribution and squaring the resulting gradient gives an " |
| "unbiased estimate of the Gauss–Newton approximation to that " |
| "diagonal, needing one extra backward pass and no second-order " |
| "autograd. And clipping: the update is " |
| "clip(<i>m</i>/max(<i>h</i>,ε), ρ), which bounds the step " |
| "wherever the curvature estimate is tiny, noisy or negative. The clip " |
| "is what makes a stale estimate safe, which in turn lets the estimator " |
| "be refreshed only every <i>k</i> = 10 steps, holding the amortised " |
| "overhead near 10%.")) |
| a(kv_note( |
| "Implementation note worth recording. The sampled loss that produces " |
| "the curvature estimate is a <i>mean</i> over tokens, so its gradient " |
| "carries a 1/B factor and its square a 1/B² one; the denominator " |
| "must therefore be scaled back up by the token count B. A first " |
| "version of this implementation left that factor at 1. The result was " |
| "not a crash but something quieter and more misleading: with the " |
| "denominator four orders of magnitude too small, the clip saturated on " |
| "99.9% of coordinates, every step became exactly ±lr, and " |
| "Sophia silently degenerated into sign-SGD. It then behaved exactly " |
| "like one — demanding a Lion-sized learning rate and finishing " |
| "last. With the factor restored, clip saturation falls and the " |
| "curvature term does real work. The failure is instructive because " |
| "every diagnostic still looked healthy: the loss went down, nothing " |
| "diverged, and only the suspiciously Lion-like learning rate hinted " |
| "that the second-order machinery was not actually running.")) |
| a(kv_note( |
| "A second instance of the same trap, worth stating separately because " |
| "it is a methodology error rather than a coding one. The first version " |
| "of the sweep harness ran the plain training step but omitted the " |
| "periodic curvature refresh, since that refresh is not part of an " |
| "ordinary optimizer's loop. For Sophia this left <i>h</i> at zero for " |
| "the whole short run, so the denominator fell back to epsilon, the " |
| "clip saturated, and the sweep dutifully reported the best learning " |
| "rate <i>for sign-SGD</i> — which was then handed to a training " |
| "run that did refresh the estimator and therefore behaved differently. " |
| "Nothing errored; the sweep simply tuned a different algorithm from " |
| "the one it was tuning for. The general lesson is that a " |
| "hyperparameter search is only valid if its inner loop is the same " |
| "algorithm as the final run, including any step that looks like " |
| "bookkeeping.")) |
| return f |
|
|
|
|
| def part2_results(M, sweep): |
| f = [] |
| a = f.append |
| a(P("2.2 Hyperparameter sweep (bonus)", "h2")) |
| a(P( |
| "The five update rules disagree about the peak learning rate by two " |
| "orders of magnitude, so handing them all AdamW's value would have " |
| "measured tuning effort rather than the update rules. Each was swept " |
| "over its own grid on a short slice of the regime — 4M tokens, " |
| "about 8% of the full run — which is the cheap-debugging trick " |
| "the brief suggests, applied to the one hyperparameter that matters " |
| "most.")) |
| if sweep: |
| rows = [["optimizer", "grid searched", "selected", "val loss @4M", |
| "at grid edge?"]] |
| for n in P2_ORDER: |
| runs = sweep.get(n, []) |
| if not runs: |
| continue |
| ok = [r for r in runs if not r.get("diverged")] |
| best = min(ok, key=lambda r: r["val_loss"]) |
| grid = [r["lr"] for r in runs] |
| edge = "yes" if best["lr"] in (min(grid), max(grid)) else "no" |
| rows.append([P2_NAME[n], |
| " ".join(f"{g:.0e}" for g in sorted(grid)), |
| f"{best['lr']:.1e}", f"{best['val_loss']:.4f}", edge]) |
| a(table(wrap_rows(rows, {1}), widths=[2.3*cm, 6.4*cm, 2.2*cm, 2.6*cm, |
| 2.4*cm])) |
| a(caption( |
| "Table 3. Learning-rate sweep. The selected rates span a factor of " |
| "60, which is the point: Lion's sign update has a fixed step norm " |
| "and needs a much smaller rate, while Muon's orthogonalised update " |
| "tolerates a much larger one.")) |
| a(figure("p2_sweep.png", width=11*cm, cap= |
| "Figure 6. Validation loss after 4M tokens against peak learning " |
| "rate. Circled points are the selected rates.")) |
|
|
| a(P("2.3 Training dynamics", "h2")) |
| if M: |
| any_m = next(iter(M.values())) |
| a(P( |
| f"Each optimizer then trains on one full pass over the corpus " |
| f"— {any_m['tokens']/1e6:.1f}M tokens, the \"1× the " |
| f"dataset\" point — with validation loss and test BLEU " |
| f"measured every 0.1×. For this {any_m['n_params']/1e6:.1f}M-" |
| f"parameter model the Chinchilla-optimal budget would be " |
| f"{20*any_m['n_params']/1e6:.0f}M tokens, so the whole corpus " |
| f"amounts to {any_m['chinchilla_ratio']:.3f}× Chinchilla. " |
| f"Every model here is therefore heavily <i>under</i>-trained " |
| f"relative to its size, which is exactly the regime the brief " |
| f"describes as \"not exactly a clean fraction of the Chinchilla " |
| f"ratio\", and it is worth keeping in mind when reading the " |
| f"rankings: these are early-training dynamics, not converged " |
| f"results.")) |
| a(figure("p2_dynamics.png", cap= |
| "Figure 7. Validation loss and test BLEU against tokens, " |
| "measured at every 0.1× of the corpus. The upper axis " |
| "restates the budget in Chinchilla units.")) |
| a(kv_note( |
| "The BLEU panel is visibly noisier than the loss panel and the " |
| "curves cross repeatedly, so it should not be read as a ranking. " |
| "BLEU here scores a greedy 48-token continuation of a held-out " |
| "document against the single continuation that actually followed " |
| "— a task with many acceptable answers and one reference " |
| "— and the absolute values sit near 2, where a handful of " |
| "n-gram matches moves the score materially. Validation loss is " |
| "measured over far more tokens and is the reliable signal at this " |
| "scale; BLEU is included because the deliverable asks for it and " |
| "because it confirms the coarse direction of travel, not because " |
| "it separates the optimizers.")) |
|
|
| bench_f = RUNS / "part2" / "benchmark.json" |
| B = ({r["optimizer"]: r for r in load_json(bench_f)} |
| if bench_f.exists() else {}) |
| rows = [["optimizer", "category", "lr", "val loss", "test ppl", |
| "test BLEU", "step ms", "state"]] |
| for n in P2_ORDER: |
| if n not in M: |
| continue |
| m = M[n] |
| rows.append([P2_NAME[n], m["category"].split(". ", 1)[-1], |
| f"{m['hparams']['lr']:.1e}", |
| f"{m['history']['val_loss'][-1]:.4f}", |
| f"{m['test_ppl']:.2f}", f"{m['test_bleu']:.2f}", |
| f"{B[n]['step_ms']:.2f}" if n in B else "-", |
| f"{m['state_multiple']:.1f}x"]) |
| a(table(wrap_rows(rows, {1}), widths=[2.0*cm, 4.7*cm, 1.5*cm, 1.8*cm, |
| 1.7*cm, 1.8*cm, 1.6*cm, 1.4*cm])) |
| a(caption( |
| "Table 4. Final results after 1× the corpus. \"state\" is " |
| "optimizer state as a multiple of model size in fp32, measured from " |
| "the live optimizer rather than derived from the update rule. " |
| "\"step ms\" is the cost of the optimizer step alone, from the " |
| "controlled benchmark described below — not the training " |
| "runs' wall-clock.")) |
| a(kv_note( |
| "A note on measuring cost honestly. These runs shared a GPU with " |
| "unrelated jobs, so their wall-clock totals say more about what " |
| "else the device was doing than about the update rules: Muon " |
| "finished in 3.5 minutes and AdamW in 7.1, which would imply Muon " |
| "is twice as fast as AdamW when in fact it does strictly more work " |
| "per step. Rather than report that, every optimizer was re-timed " |
| "back-to-back on the same model, batch and device " |
| "(src/part2/benchmark.py), separating the optimizer step from the " |
| "forward/backward pass they all share. Those are the numbers in " |
| "Table 4 and Figure 8.")) |
| a(figure("p2_summary.png", cap= |
| "Figure 8. Final perplexity, BLEU, controlled per-step " |
| "optimizer cost, and optimizer-state memory.")) |
| return f |
|
|
|
|
| |
| |
| |
| def load_p3(): |
| f = RUNS / "part3" / "metrics.json" |
| return load_json(f) if f.exists() else None |
|
|
|
|
| def part3(D): |
| f = [] |
| a = f.append |
| a(P("3. Decoding Strategies", "h1")) |
| a(P( |
| "Four decoding strategies are implemented on top of the frozen " |
| "<font face='Courier' size='8'>EleutherAI/pythia-160m</font> checkpoint " |
| "and evaluated on <font face='Courier' size='8'>hamishivi/ROCStories" |
| "</font>. <font face='Courier' size='8'>model.generate()</font> is not " |
| "used; each strategy drives <font face='Courier' size='8'>model.forward()" |
| "</font> directly and maintains its own KV cache, which is what keeps " |
| "the cost linear in the number of generated tokens instead of " |
| "quadratic.")) |
| a(P( |
| "Two implementation details matter for correctness. Prompts in a batch " |
| "have different lengths and are <i>left</i>-padded, so that every real " |
| "final token sits at index −1 where the decoder expects it; but " |
| "left padding makes absolute positions wrong, so " |
| "<font face='Courier' size='8'>position_ids</font> is derived from the " |
| "cumulative sum of the attention mask rather than from " |
| "<font face='Courier' size='8'>arange</font>. And in beam search, a " |
| "hypothesis that has emitted EOS is retired by masking its row of " |
| "log-probabilities to −∞, so it can never be extended again " |
| "yet still competes for the final selection.")) |
| a(P("Correctness checks", "h3")) |
| a(kv_note( |
| "Beam search at width 1 produces token-identical output to greedy " |
| "decoding on every test prompt, as it must. Widths 2 and 4 return " |
| "hypotheses whose length-normalised log-probability is greater than or " |
| "equal to greedy's on every prompt — beam search is finding " |
| "genuinely better-scoring sequences, not merely different ones. " |
| "Sampling is reproducible under a fixed seed, and the top-p filter " |
| "keeps 1 token on a peaked distribution and 2 on a uniform one at " |
| "p = 0.5, as expected.")) |
|
|
| a(P("3.1 What the metrics mean", "h2")) |
| a(P( |
| "Free-form generation has no single correct answer, so accuracy, " |
| "precision, recall and F1 have to be defined rather than assumed. The " |
| "definitions used here are:")) |
| rows = [["metric", "definition"], |
| ["Reference perplexity", |
| "Perplexity of the gold continuation under teacher forcing. A " |
| "property of the model, not the decoder, so it is identical for " |
| "every strategy and acts as a control."], |
| ["Generation perplexity", |
| "Perplexity of the model on its own output. This does vary by " |
| "strategy and measures how safely each one plays."], |
| ["Precision / recall / F1", |
| "Unigram overlap with the reference by multiset intersection, " |
| "micro-averaged over the corpus. Precision is matched/|generated|, " |
| "recall is matched/|reference|."], |
| ["Accuracy", |
| "Position-wise token agreement with the reference over the " |
| "overlapping prefix; exact-match rate is reported alongside."], |
| ["distinct-1/2, repetition", |
| "Diversity: the share of generated unigrams/bigrams that are " |
| "unique, and the share of bigrams that repeat within a " |
| "generation."]] |
| a(table(wrap_rows(rows, {1}), widths=[4.0*cm, 12.5*cm], |
| align_right_from=99, font=7.4)) |
| a(caption("Table 5. Metric definitions, as implemented in " |
| "src/part3/evaluate.py.")) |
| a(P( |
| "Two of these deserve a caveat up front. Exact match is essentially " |
| "always zero and is reported only to make the point that it is the " |
| "wrong instrument for this task. And the overlap metrics systematically " |
| "favour whichever strategy produces the most generic, highest-" |
| "probability text, because generic text reuses the reference's common " |
| "words; they cannot see whether the output is repetitive or dull, which " |
| "is precisely why the diversity columns are there.")) |
|
|
| if D: |
| R = D["results"] |
| a(P("3.2 Results", "h2")) |
| rows = [["strategy", "P", "R", "F1", "acc", "gen ppl", "dist-1", |
| "dist-2", "rep", "len", "sec"]] |
| for r in R: |
| rows.append([r["strategy"], f"{r['precision']:.3f}", |
| f"{r['recall']:.3f}", f"{r['f1']:.3f}", |
| f"{r['accuracy']:.3f}", f"{r['gen_ppl']:.2f}", |
| f"{r['distinct_1']:.3f}", f"{r['distinct_2']:.3f}", |
| f"{r['repetition_rate']:.3f}", f"{r['mean_len']:.1f}", |
| f"{r['seconds']:.1f}"]) |
| a(table(rows, widths=[2.4*cm, 1.3*cm, 1.3*cm, 1.3*cm, 1.3*cm, 1.6*cm, |
| 1.4*cm, 1.4*cm, 1.3*cm, 1.2*cm, 1.3*cm], |
| font=7.2)) |
| a(caption( |
| f"Table 6. All strategies on {D['n_samples']} ROCStories test " |
| f"samples, {D['max_new']} new tokens each. Reference perplexity " |
| f"(the control) is {D['reference_ppl']:.2f} for every row.")) |
| |
| return f |
|
|
|
|
| def part3_analysis(D): |
| f = [] |
| a = f.append |
| if not D: |
| return f |
| R = {r["strategy"]: r for r in D["results"]} |
| res = D["results"] |
| g = R["greedy"] |
| ref_ppl = D["reference_ppl"] |
| best_f1 = max(res, key=lambda r: r["f1"]) |
| worst_rep = max(res, key=lambda r: r["repetition_rate"]) |
| |
| closest = min(res, key=lambda r: abs(r["gen_ppl"] - ref_ppl)) |
|
|
| a(P("3.3 Analysis of the generations", "h2")) |
| a(P( |
| f"<b>Greedy decoding degenerates, and the metrics show it twice over." |
| f"</b> {worst_rep['strategy']} has a bigram repetition rate of " |
| f"{worst_rep['repetition_rate']:.3f} — more than half of all " |
| f"bigrams it emits are repeats of one it has already emitted in the " |
| f"same continuation — and a distinct-2 of " |
| f"{worst_rep['distinct_2']:.3f}. Reading the outputs makes the " |
| f"mechanism plain. Given the prompt <i>\"David had achieved his " |
| f"lifelong goal.\"</i>, beam search at width 4 returns <i>\"I'm not " |
| f"going to lie to you.\"</i> five times in a row and nothing else " |
| f"(Table A2). Once a high-probability phrase is in the context, " |
| f"repeating it becomes more probable still, and a decoder that only " |
| f"ever maximises probability has no mechanism to escape the loop it " |
| f"has walked into.")) |
| a(P( |
| f"<b>The consequence is that sampling beats greedy on the overlap " |
| f"metrics too</b> — which is not the usual expectation. " |
| f"{best_f1['strategy']} has the highest unigram F1 " |
| f"({best_f1['f1']:.3f}) against greedy's {g['f1']:.3f}, a relative " |
| f"improvement of {100*(best_f1['f1']/g['f1'] - 1):.0f}%. Precision is " |
| f"where greedy loses most ({g['precision']:.3f} against " |
| f"{best_f1['precision']:.3f}): a generation that spends half its tokens " |
| f"repeating one clause contains few distinct words, so most of what it " |
| f"emits cannot match anything in the reference. The common claim that " |
| f"probability-maximising decoders win precision-style metrics while " |
| f"sampling wins diversity does not hold here, because degeneration " |
| f"costs greedy both at once.")) |
| a(P( |
| f"<b>Quality against truncation width is an inverted U, not a " |
| f"trade-off curve.</b> Ordering the sampling runs by how much " |
| f"probability mass they keep — top-k 10, top-k 50, top-p 0.90, " |
| f"top-p 0.95 — F1 goes " |
| + ", ".join(f"{R[n]['f1']:.3f}" for n in |
| ["top_k_10", "top_k_50", "top_p_0.90", "top_p_0.95"] |
| if n in R) |
| + f", rising from greedy's {g['f1']:.3f} to a peak at " |
| f"{best_f1['strategy']} and falling away again. Both extremes are bad " |
| f"for different reasons: too little truncation and the sampler reaches " |
| f"into the unreliable tail, too much and it collapses back towards the " |
| f"repetition that afflicts greedy. The best setting is a middle one, " |
| f"and it has to be found empirically.")) |
| a(P( |
| f"<b>Perplexity of a model's own output is the most revealing column " |
| f"in the table.</b> The human reference continuations score " |
| f"{ref_ppl:.2f} under this model, yet greedy decoding produces text of " |
| f"perplexity {g['gen_ppl']:.2f} — its own output is about " |
| f"{ref_ppl/g['gen_ppl']:.0f}× more predictable than genuine human " |
| f"writing. That single comparison is the clearest statement of why " |
| f"maximising probability is the wrong objective for open-ended " |
| f"generation: real text simply does not sit at the mode of the " |
| f"distribution. {closest['strategy']} lands closest to the human figure " |
| f"({closest['gen_ppl']:.2f} against {ref_ppl:.2f}), which is exactly " |
| f"what nucleus sampling was designed to achieve.")) |
| a(P( |
| f"<b>Top-k and top-p differ in how they choose their width.</b> Top-k " |
| f"keeps a fixed number of candidates however peaked the distribution " |
| f"is; top-p keeps a fixed probability mass and so adapts, narrowing to " |
| f"one token where the model is confident and widening where it is not. " |
| f"That adaptivity shows up here as top-p reaching higher diversity " |
| f"(distinct-2 {R['top_p_0.95']['distinct_2']:.3f} at p=0.95 against " |
| f"{R['top_k_50']['distinct_2']:.3f} at k=50) and landing nearer the " |
| f"human perplexity, at some cost in overlap.")) |
| a(P( |
| "That cost is not only statistical, and the sample generations show " |
| "what it looks like. Top-p at 0.95 produces fluent, varied, " |
| "non-repetitive English — and, on the same prompt, invents " |
| "\"Pandas\", \"Dorrit\" and \"the Kári complex\", none of which " |
| "has anything to do with David and his lifelong goal. The failure mode " |
| "has moved rather than disappeared: greedy fails by saying the same " |
| "true-ish thing forever, sampling fails by confidently wandering off " |
| "topic. A 160M-parameter model has weak enough long-range coherence " |
| "that both failure modes are easy to trigger, and no decoding strategy " |
| "can repair a model that does not track the prompt well in the first " |
| "place — decoding redistributes the errors, it does not remove " |
| "them.")) |
| a(P( |
| f"<b>Positional accuracy is uninformative and exact match is " |
| f"degenerate.</b> Accuracy sits near {g['accuracy']:.3f} for every " |
| f"strategy and the exact-match rate is " |
| f"{max(r['exact_match'] for r in res):.3f} at best. Both are reported " |
| f"because the deliverable asks for them, but neither can distinguish " |
| f"the strategies: requiring the model to reproduce a specific " |
| f"human-written sentence token-for-token is the wrong question to ask " |
| f"of open-ended story continuation, where many different continuations " |
| f"are equally valid. The overlap and diversity columns are what carry " |
| f"the signal.")) |
| a(figure("p3_metrics.png", cap= |
| "Figure 9. Overlap with the reference continuation. Greedy is " |
| "worst on precision, recall and F1 alike — degeneration " |
| "hurts it on every axis at once.")) |
| a(figure("p3_tradeoff.png", cap= |
| "Figure 10. Left: F1 against diversity, showing the inverted U. " |
| "Right: how far each strategy's own output sits from the " |
| "perplexity of real human text (dashed line).")) |
|
|
| a(P("3.4 Beam search: measured cost and time complexity", "h2")) |
| beams = sorted([r for r in res if r["strategy"].startswith("beam")], |
| key=lambda r: int(r["strategy"].split("_")[1])) |
| if len(beams) >= 2: |
| t1 = beams[0]["seconds"] |
| ratios = ", ".join( |
| f"W={int(b['strategy'].split('_')[1])}: {b['seconds']/t1:.2f}x" |
| for b in beams) |
| a(P( |
| "<b>Theory.</b> Let <i>W</i> be the beam width, <i>T</i> the number " |
| "of generated tokens, <i>V</i> the vocabulary and <i>C</i> the cost " |
| "of one forward pass for one sequence. Each step runs the model on " |
| "all <i>W</i> live hypotheses, costing <i>O</i>(<i>WC</i>); scores " |
| "<i>W·V</i> candidate extensions, costing " |
| "<i>O</i>(<i>WV</i>); and selects the best <i>W</i> of them, which " |
| "is <i>O</i>(<i>WV</i>) with a partial selection and " |
| "<i>O</i>(<i>WV</i> log <i>WV</i>) with a full sort. Over <i>T</i> " |
| "steps the total is <b><i>O</i>(<i>TW</i>(<i>C</i> + <i>V</i>))</b> " |
| "— <b>linear in the beam width</b>. Memory is linear in " |
| "<i>W</i> as well, and is often the binding constraint in practice, " |
| "since the KV cache must hold <i>W</i> copies of the prompt.")) |
| a(P( |
| f"<b>Measured.</b> Relative to width 1: {ratios}. The measured " |
| f"scaling is <i>sub</i>-linear — quadrupling the beam costs " |
| f"{beams[-1]['seconds']/t1:.2f}×, not 4×. The " |
| f"asymptotic analysis is not contradicted; the constant factor is " |
| f"hiding it. All <i>W</i> hypotheses are evaluated in a single " |
| f"batched forward pass of <i>B·W</i> rows, and at " |
| f"<i>B</i> = {D['batch_size']} with a 160M-parameter " |
| f"model the GPU is nowhere near saturated at <i>W</i> = 1. " |
| f"Widening the beam therefore fills idle parallel capacity instead " |
| f"of adding serial work. The linear term reasserts itself once the " |
| f"device is actually busy — at larger batch, a larger model, " |
| f"or large enough <i>W</i> — and the memory cost is linear in " |
| f"<i>W</i> from the first token regardless.")) |
| a(figure("p3_beam_time.png", cap= |
| "Figure 11. Measured beam-search cost against width, against " |
| "an exactly-linear reference. The gap is unused GPU " |
| "parallelism at small widths, not a violation of the bound.")) |
| b1, b4 = beams[0], beams[-1] |
| a(P( |
| f"<b>Is it worth it?</b> Going from width 1 to width " |
| f"{int(b4['strategy'].split('_')[1])} changes F1 by " |
| f"{b4['f1'] - b1['f1']:+.3f} and the repetition rate by " |
| f"{b4['repetition_rate'] - b1['repetition_rate']:+.3f}, for " |
| f"{b4['seconds']/t1:.2f}× the compute and " |
| f"{int(b4['strategy'].split('_')[1])}× the cache memory. A " |
| f"wider beam is strictly better at what it optimises — it " |
| f"finds sequences of higher log-probability, as the score check in " |
| f"section 3 confirms — but on this task higher probability " |
| f"means more of the generic, repetitive text that greedy already " |
| f"produces. Beam search earns its cost in tasks with a genuinely " |
| f"peaked target distribution, such as translation or speech " |
| f"recognition; open-ended story continuation is not one of them.")) |
| return f |
|
|
|
|
| def conclusions(M1, M2, D): |
| f = [] |
| a = f.append |
| a(P("4. Conclusions", "h1")) |
| items = [] |
| if M1: |
| c1, c5 = M1["c1_dense"], M1["c5_moe_4e2a_wide"] |
| c2, c3 = M1["c2_moe_4e1a"], M1["c3_moe_4e2a"] |
| items.append( |
| f"<b>Part 1.</b> At a fixed <i>total</i> parameter budget, sparsity " |
| f"costs quality: top-1 routing loses " |
| f"{c1['bleu_all']-c2['bleu_all']:.1f} BLEU against the dense " |
| f"baseline and top-2 loses {c1['bleu_all']-c3['bleu_all']:.1f}. At a " |
| f"fixed <i>active</i> budget the sign flips and MoE wins — C5 " |
| f"gains {c5['bleu_all']-c1['bleu_all']:.1f} BLEU and " |
| f"{c1['test_ppl']-c5['test_ppl']:.2f} perplexity over the dense " |
| f"model for the same per-token compute. Which budget is held fixed " |
| f"entirely determines whether MoE looks good.") |
| items.append( |
| "<b>Part 1, routing.</b> The load-balancing loss kept every expert " |
| "in use (normalised entropy ≥ 0.999 everywhere), and " |
| "specialisation still emerged along language lines — most " |
| "sharply under top-1 routing, where each token must commit to one " |
| "expert, and most cleanly in C5, where three of four experts each " |
| "took a language.") |
| if M2: |
| best_ppl = min(M2.values(), key=lambda m: m["test_ppl"]) |
| best_bleu = max(M2.values(), key=lambda m: m["test_bleu"]) |
| items.append( |
| f"<b>Part 2.</b> {P2_NAME[best_ppl['optimizer']]} reached the lowest " |
| f"test perplexity ({best_ppl['test_ppl']:.2f}) and " |
| f"{P2_NAME[best_bleu['optimizer']]} the highest BLEU " |
| f"({best_bleu['test_bleu']:.2f}) at 1× the corpus. The spread " |
| f"between the best and worst optimizer is small next to the spread " |
| f"induced by the learning rate alone in the sweep, which supports " |
| f"the paper's central claim: much of the reported advantage of new " |
| f"optimizers is tuning, and comparisons at a single shared learning " |
| f"rate are not informative.") |
| if D: |
| items.append( |
| "<b>Part 3.</b> No decoding strategy dominates. Probability-" |
| "maximising decoders (greedy, beam) win the overlap metrics and " |
| "lose badly on repetition; truncated sampling (top-k, top-p) does " |
| "the reverse and lands nearer the perplexity of genuine human text. " |
| "Beam search is linear in the beam width in theory and sub-linear " |
| "here in practice because small batches leave the GPU idle — " |
| "and at this task it buys almost no quality for that cost.") |
| for it in items: |
| a(Paragraph(f"• {it}", ParagraphStyle( |
| "bullet", parent=S["body"], leftIndent=12, firstLineIndent=-12, |
| spaceAfter=7))) |
|
|
| a(P("Limitations", "h2")) |
| a(P( |
| "Every result here is a single seed at ~27M parameters and a token " |
| "budget well short of compute-optimal, on a GPU shared with unrelated " |
| "jobs. Differences of a few hundredths of perplexity should not be read " |
| "as reliable orderings, and the wall-clock numbers in Table 4 are " |
| "indicative only. The Part 2 sweep covered the peak learning rate alone; " |
| "betas, weight decay and warmup were held at published defaults, so an " |
| "optimizer whose published defaults suit this setting poorly is " |
| "disadvantaged in a way the sweep does not correct.")) |
| return f |
|
|
|
|
| def references(): |
| f = [] |
| a = f.append |
| a(P("References", "h1")) |
| refs = [ |
| "Hoffmann, J. et al. <i>Training Compute-Optimal Large Language " |
| "Models.</i> arXiv:2203.15556, 2022.", |
| "Papineni, K. et al. BLEU: a method for automatic evaluation of machine " |
| "translation. <i>ACL '02</i>, pp. 311–318, 2002.", |
| "Wen, K. et al. <i>Fantastic Pretraining Optimizers and Where to Find " |
| "Them.</i> arXiv:2509.02046, 2025.", |
| "Loshchilov, I. and Hutter, F. <i>Decoupled Weight Decay " |
| "Regularization.</i> ICLR, 2019.", |
| "Yuan, H. et al. <i>MARS: Unleashing the Power of Variance Reduction " |
| "for Training Large Models.</i> arXiv:2411.10438, 2024.", |
| "Chen, X. et al. <i>Symbolic Discovery of Optimization Algorithms</i> " |
| "(Lion). NeurIPS, 2023.", |
| "Jordan, K. et al. <i>Muon: An optimizer for hidden layers in neural " |
| "networks.</i> 2024.", |
| "Liu, H. et al. <i>Sophia: A Scalable Stochastic Second-order Optimizer " |
| "for Language Model Pre-training.</i> arXiv:2305.14342, 2023.", |
| "Fedus, W. et al. <i>Switch Transformers: Scaling to Trillion Parameter " |
| "Models with Simple and Efficient Sparsity.</i> JMLR, 2022.", |
| "Dai, D. et al. <i>DeepSeekMoE: Towards Ultimate Expert Specialization " |
| "in Mixture-of-Experts Language Models.</i> arXiv:2401.06066, 2024.", |
| "Holtzman, A. et al. <i>The Curious Case of Neural Text Degeneration</i> " |
| "(nucleus sampling). ICLR, 2020.", |
| "Su, J. et al. <i>RoFormer: Enhanced Transformer with Rotary Position " |
| "Embedding.</i> arXiv:2104.09864, 2021.", |
| ] |
| for i, r in enumerate(refs, 1): |
| a(Paragraph(f"[{i}] {r}", ParagraphStyle( |
| "ref", parent=S["body"], fontSize=8.6, leading=11.5, leftIndent=16, |
| firstLineIndent=-16, spaceAfter=4))) |
| return f |
|
|
|
|
| def part2_analysis(M, sweep): |
| f = [] |
| a = f.append |
| a(P("2.4 Analysis", "h2")) |
|
|
| order = [n for n in P2_ORDER if n in M] |
| by_ppl = sorted(order, key=lambda n: M[n]["test_ppl"]) |
| best, worst = M[by_ppl[0]], M[by_ppl[-1]] |
| adamw = M.get("adamw") |
| spread = worst["test_ppl"] - best["test_ppl"] |
|
|
| a(P( |
| f"<b>Ranking.</b> By test perplexity the order is " |
| + ", ".join(f"{P2_NAME[n]} ({M[n]['test_ppl']:.2f})" for n in by_ppl) |
| + f". The spread between best and worst is {spread:.2f} perplexity.")) |
|
|
| if sweep and adamw: |
| lrs = [r for r in sweep.get("adamw", []) if not r.get("diverged")] |
| if len(lrs) > 1: |
| lr_spread = (max(r["val_loss"] for r in lrs) |
| - min(r["val_loss"] for r in lrs)) |
| |
| opt_losses = [M[n]["history"]["val_loss"][-1] for n in order] |
| opt_spread = max(opt_losses) - min(opt_losses) |
| a(P( |
| f"<b>The learning rate matters more than the optimizer does.</b> " |
| f"Across all five optimizers at their own best rates, final " |
| f"validation loss spans {opt_spread:.2f} nats. Within AdamW " |
| f"alone, varying nothing but the peak learning rate spans " |
| f"{lr_spread:.2f} nats — " |
| f"{lr_spread/max(opt_spread, 1e-9):.1f} times as much. The two " |
| f"figures are measured at different budgets (4M tokens for the " |
| f"sweep, the full corpus for the runs) so the ratio is " |
| f"indicative rather than exact, but the direction is not in " |
| f"doubt: choosing the learning rate badly costs far more than " |
| f"choosing the optimizer badly.")) |
| a(P( |
| f"This is the practical form of the argument in Wen et al. An " |
| f"optimizer comparison conducted at a single shared learning " |
| f"rate mostly measures which optimizer that rate happened to " |
| f"suit, and a new method can look impressive simply by having " |
| f"been tuned more carefully than the baseline it is compared " |
| f"against. It is also why three of the five grids here had to " |
| f"be widened after a first pass selected a value sitting on a " |
| f"grid edge — an edge-valued optimum means the search " |
| f"found the boundary of the grid, not the minimum, and " |
| f"reporting it would have measured grid design.")) |
|
|
| if "muon" in M and adamw: |
| mu = M["muon"] |
| rel = "better" if mu["test_ppl"] < adamw["test_ppl"] else "worse" |
| a(P( |
| f"<b>Muon.</b> The matrix-based update came out {rel} than AdamW " |
| f"({mu['test_ppl']:.2f} vs {adamw['test_ppl']:.2f} perplexity) while " |
| f"tolerating a learning rate {mu['hparams']['lr']/adamw['hparams']['lr']:.0f}× " |
| f"larger. That is the expected signature: orthogonalising the " |
| f"momentum fixes the scale of the update's singular values, so the " |
| f"learning rate no longer has to be set conservatively enough for " |
| f"the largest one. It also keeps only one state tensor for its 2D " |
| f"parameters, so it is cheaper in memory than AdamW as well — " |
| f"a rare combination, and consistent with the strong showing " |
| f"matrix-based methods get in the paper. Its cost is the " |
| f"Newton–Schulz iteration: five extra matmuls per 2D " |
| f"parameter per step.")) |
|
|
| if "lion" in M and adamw: |
| li = M["lion"] |
| a(P( |
| f"<b>Lion.</b> Half the optimizer memory of AdamW " |
| f"({li['state_multiple']:.1f}× model size against " |
| f"{adamw['state_multiple']:.1f}×) for " |
| f"{abs(li['test_ppl'] - adamw['test_ppl']):.2f} perplexity " |
| f"{'better' if li['test_ppl'] < adamw['test_ppl'] else 'worse'}. " |
| f"Its selected learning rate is " |
| f"{adamw['hparams']['lr']/li['hparams']['lr']:.0f}× smaller " |
| f"than AdamW's, exactly as the sign update predicts: every " |
| f"coordinate moves by precisely the learning rate, so the update " |
| f"norm is fixed rather than adaptive and the rate has to absorb " |
| f"that. In a memory-bound setting the trade is attractive; here, " |
| f"where memory was not the constraint, it is simply a slightly " |
| f"weaker optimizer.")) |
|
|
| if "mars" in M and adamw: |
| ma = M["mars"] |
| a(P( |
| f"<b>MARS.</b> Variance reduction gave " |
| f"{abs(ma['test_ppl'] - adamw['test_ppl']):.2f} perplexity " |
| f"{'better' if ma['test_ppl'] < adamw['test_ppl'] else 'worse'} than " |
| f"AdamW at {ma['state_multiple']:.1f}× model size in state " |
| f"— the highest memory cost of the five, since it must keep " |
| f"the previous gradient alongside both moments. The correction is " |
| f"most valuable when gradient noise dominates, which favours small " |
| f"batches; at the 24×512 batch used here the gradients are " |
| f"already fairly clean, so there is less noise for it to remove.")) |
|
|
| if "sophia" in M: |
| so = M["sophia"] |
| a(P( |
| f"<b>Sophia-G.</b> The curvature-based method was the weakest here " |
| f"({so['test_ppl']:.2f} perplexity) and needed the smallest learning " |
| f"rate of the five ({so['hparams']['lr']:.0e}). Part of that is " |
| f"structural: the clipped ratio caps every coordinate's step at " |
| f"±<i>lr</i>, and the more often that cap binds the more the " |
| f"update resembles Lion's sign step, so the learning rate has to " |
| f"absorb a partly fixed update norm. The rest is regime. Sophia's " |
| f"advantage is supposed to come from dividing by curvature, but the " |
| f"Gauss–Newton estimate is least informative early in " |
| f"training, far from any minimum — which is exactly and only " |
| f"where this experiment lives, at " |
| f"{so['chinchilla_ratio']:.3f}× Chinchilla. It also pays an " |
| f"extra forward/backward every ten steps for the privilege. Sophia " |
| f"is reported to pay off over long runs; one tenth of a Chinchilla " |
| f"budget is not a setting in which that could show up, and this " |
| f"result should be read as \"not yet useful here\" rather than " |
| f"\"worse\".")) |
|
|
| a(P( |
| "<b>Caveat on the ranking.</b> Only the peak learning rate was swept. " |
| "Betas, weight decay, warmup and (for Sophia) ρ were left at " |
| "published defaults, which favours whichever optimizer's defaults " |
| "happen to suit a 27M-parameter model at 0.1× Chinchilla. The " |
| "differences reported above are also single-seed. They should be read " |
| "as directional, not as a benchmark.")) |
| return f |
|
|
|
|
| def appendix(M1, D): |
| f = [] |
| a = f.append |
| a(P("Appendix A. Sample generations", "h1")) |
|
|
| if M1: |
| a(P("A.1 Translation (Part 1, dense baseline)", "h2")) |
| samples = M1["c1_dense"].get("samples", [])[:6] |
| if samples: |
| rows = [["src", "model output", "reference"]] |
| for s in samples: |
| rows.append([s["lang"], s["hyp"][:150], s["ref"][:150]]) |
| a(table(wrap_rows(rows, {1, 2}), widths=[1.0*cm, 7.8*cm, 7.7*cm], |
| align_right_from=99, font=6.8)) |
| a(caption("Table A1. Greedy translations from the dense baseline on " |
| "held-out test pairs.")) |
|
|
| if D: |
| a(P("A.2 Story continuations by decoding strategy (Part 3)", "h2")) |
| prompts = D.get("prompts", []) |
| if prompts: |
| a(P(f"<b>Prompt:</b> <i>{prompts[0]}</i>")) |
| refs = D.get("references", []) |
| if refs: |
| a(P(f"<b>Human continuation:</b> <i>{refs[0][:300]}</i>")) |
| rows = [["strategy", "generated continuation"]] |
| for name, gens in D.get("examples", {}).items(): |
| if gens: |
| rows.append([name, gens[0][:260].replace("\n", " ")]) |
| a(table(wrap_rows(rows, {1}), widths=[2.4*cm, 14.1*cm], |
| align_right_from=99, font=6.8)) |
| a(caption("Table A2. The same prompt decoded eight ways. The " |
| "repetition that the metrics in Table 6 quantify is " |
| "visible directly in the greedy and beam rows.")) |
|
|
| a(P("Appendix B. Reproducing these results", "h1")) |
| a(Paragraph( |
| "uv sync<br/>" |
| "uv run main.py part1 --all<br/>" |
| "uv run main.py part2 --sweep<br/>" |
| "uv run main.py part2 --all --use_sweep<br/>" |
| "uv run main.py part3 --n_samples 1000<br/>" |
| "uv run main.py plots<br/>" |
| "uv run python -m src.report", S["code"])) |
| a(P( |
| "Checkpoints are written to runs/ and are excluded from version " |
| "control by .gitignore, since the repository quota is small. Figures " |
| "and the value tables behind them are in assets/.")) |
| return f |
|
|
|
|
| |
| |
| |
| def front_matter(M1, M2, D): |
| f = [] |
| a = f.append |
| a(P("Advanced NLP — Assignment 2", "title")) |
| a(P("Mixture of Experts · Optimizers · Decoding Strategies" |
| "<br/>Monsoon 2026", "sub")) |
|
|
| a(P("Overview", "h1")) |
| a(P( |
| "This report covers three experiments. Part 1 builds a decoder-only " |
| "transformer whose feed-forward block is a drop-in choice between a " |
| "dense MLP and four Mixture-of-Experts variants, trains all five on " |
| "Vietnamese/Japanese→English translation under matched parameter " |
| "budgets, and asks what sparsity costs and what it buys. Part 2 " |
| "implements five optimizers from scratch — one per category of " |
| "Table 1 of Wen et al. [3], including the bonus Hessian-based entry " |
| "— and compares their training dynamics on a fresh corpus. Part 3 " |
| "implements four decoding strategies against a frozen Pythia-160M and " |
| "measures what each one actually optimises.")) |
| a(P( |
| "Everything is written in PyTorch from basic operations. " |
| "<font face='Courier' size='8'>nn.Transformer</font>, " |
| "<font face='Courier' size='8'>nn.MultiheadAttention</font>, the " |
| "<font face='Courier' size='8'>torch.optim</font> update rules and " |
| "<font face='Courier' size='8'>model.generate()</font> are not used.")) |
|
|
| a(P("Headline numbers", "h2")) |
| rows = [["", "best", "value", "note"]] |
| if M1: |
| best = max(M1.values(), key=lambda m: m["bleu_all"]) |
| dense = M1["c1_dense"] |
| rows.append([ |
| "Part 1 (BLEU)", P1_NAME[best["config"]["name"]].replace("—", "-"), |
| f"{best['bleu_all']:.2f}", |
| f"vs {dense['bleu_all']:.2f} dense, same active params"]) |
| if M2: |
| bp = min(M2.values(), key=lambda m: m["test_ppl"]) |
| rows.append([ |
| "Part 2 (perplexity)", P2_NAME[bp["optimizer"]], |
| f"{bp['test_ppl']:.2f}", |
| f"at 1x corpus = {bp['chinchilla_ratio']:.3f}x Chinchilla"]) |
| if D: |
| bf = max(D["results"], key=lambda r: r["f1"]) |
| bd = max(D["results"], key=lambda r: r["distinct_2"]) |
| rows.append(["Part 3 (F1)", bf["strategy"], f"{bf['f1']:.3f}", |
| "highest overlap with the reference"]) |
| rows.append(["Part 3 (diversity)", bd["strategy"], |
| f"{bd['distinct_2']:.3f}", "highest distinct-2"]) |
| if len(rows) > 1: |
| a(table(wrap_rows(rows, {3}), widths=[3.6*cm, 4.4*cm, 2.2*cm, 6.3*cm], |
| align_right_from=2)) |
| a(caption("Table 0. One line per experiment; details in the " |
| "corresponding section.")) |
|
|
| a(P("Experimental setup common to Parts 1 and 2", "h2")) |
| rows = [["component", "choice"], |
| ["Architecture", "pre-norm decoder-only transformer, 6 layers, " |
| "d_model 512, 8 heads"], |
| ["Attention", "causal multi-head self-attention with RoPE, written " |
| "from scratch"], |
| ["Normalisation", "RMSNorm, plus a final norm before the head"], |
| ["Embeddings", "tied input/output, 16k byte-level BPE"], |
| ["Precision", "bf16 autocast, fp32 master weights and loss"], |
| ["Schedule", "linear warmup then cosine decay to 10% of peak"], |
| ["Init", "N(0, 0.02), residual projections scaled by " |
| "1/sqrt(2L)"], |
| ["Seed", "1337 for data order, init and sampling"]] |
| a(table(wrap_rows(rows, {1}), widths=[3.4*cm, 13.1*cm], |
| align_right_from=99, font=7.6)) |
| a(caption("Table 0b. Held fixed across every run so that the ablations " |
| "vary one thing at a time.")) |
| a(PageBreak()) |
| return f |
|
|
|
|
| def build(path=None): |
| global S |
| S = styles() |
| M1, M2, D = load_p1(), load_p2(), load_p3() |
| sweep_f = RUNS / "part2" / "sweep.json" |
| sweep = load_json(sweep_f) if sweep_f.exists() else None |
|
|
| story = [] |
| story += front_matter(M1, M2, D) |
| if M1: |
| story += part1(M1) |
| story += part1_analysis(M1) |
| story.append(PageBreak()) |
| story += part2(M2) |
| story += part2_results(M2, sweep) |
| if M2: |
| story += part2_analysis(M2, sweep) |
| story.append(PageBreak()) |
| story += part3(D) |
| story += part3_analysis(D) |
| story.append(PageBreak()) |
| story += conclusions(M1, M2, D) |
| story += appendix(M1, D) |
| story += references() |
|
|
| path = Path(path or ROOT / "report.pdf") |
| doc = SimpleDocTemplate( |
| str(path), pagesize=A4, |
| leftMargin=2*cm, rightMargin=2*cm, topMargin=1.8*cm, bottomMargin=1.8*cm, |
| title="ANLP Assignment 2", author="") |
| doc.build(story, onFirstPage=footer, onLaterPages=footer) |
| print(f"wrote {path}") |
| return path |
|
|
|
|
| if __name__ == "__main__": |
| build() |
|
|
