Blame for src/report.py

"""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()


# ---------------------------------------------------------------------------
# Part 1
# ---------------------------------------------------------------------------
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 &mdash; 4 experts, top-1",
           "c3_moe_4e2a": "C3 &mdash; 4 experts, top-2",
           "c4_moe_shared": "C4 &mdash; 1 shared + 3 routed, top-1",
           "c5_moe_4e2a_wide": "C5 &mdash; 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'>&lt;bos&gt; &lt;vi&gt; {source} &lt;en&gt; "
        "{target} &lt;eos&gt;</font>. The cross-entropy is masked to the English "
        "target tokens and the closing <font face='Courier' size='8'>&lt;eos&gt;"
        "</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&ndash;English pair "
        "and a Japanese&ndash;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'>&lt;unk&gt;"
        "</font>, and BPE then spends its merge budget on whatever multi-byte "
        "sequences are frequent &mdash; 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&middot;<i>d</i><sub>model</sub>&middot;<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&ndash;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("&mdash;", "-"),
            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&ndash;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> &mdash; 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>&middot;&Sigma;<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("&mdash;", "-"),
                     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 &mdash; 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&times; 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&ndash;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 &mdash; the "
        "best-performing one &mdash; 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?\" &mdash; 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 &mdash; 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 &mdash; "
        "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 &mdash; "
        "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"&mdash; {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


# ---------------------------------------------------------------------------
# Part 2
# ---------------------------------------------------------------------------
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 &mdash; the human original and one "
        "continuation per generator model &mdash; 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 &mdash; 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 "
        "&radic;<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 &mdash; i.e. to float32 rounding."))

    a(P("MARS &mdash; 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> + "
        "&gamma;(&beta;<sub>1</sub>/(1&minus;&beta;<sub>1</sub>))"
        "(<i>g<sub>t</sub></i> &minus; <i>g<sub>t&minus;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 &mdash; it is a "
        "difference of two noisy quantities and can blow up early &mdash; then "
        "fed to an otherwise ordinary AdamW. Cost: a third state tensor for the "
        "previous gradient."))

    a(P("Lion &mdash; 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 "
        "(&beta;<sub>1</sub> = 0.9) than the one that is stored "
        "(&beta;<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 "
        "&plusmn;<i>lr</i>, the update norm is fixed and much larger than "
        "AdamW's at the same learning rate &mdash; 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 &mdash; 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&ndash;Schulz "
        "iteration whose coefficients (3.4445, &minus;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 &mdash; 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 &mdash; curvature is what actually "
        "says how far one can step before the loss turns back up. Two things "
        "make it affordable. The Gauss&ndash;Newton&ndash;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&ndash;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>,&epsilon;), &rho;), 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&sup2; 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 &plusmn;lr, and "
        "Sophia silently degenerated into sign-SGD. It then behaved exactly "
        "like one &mdash; 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> &mdash; 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 &mdash; 4M tokens, "
        "about 8% of the full run &mdash; 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"&mdash; {any_m['tokens']/1e6:.1f}M tokens, the \"1&times; the "
            f"dataset\" point &mdash; with validation loss and test BLEU "
            f"measured every 0.1&times;. 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}&times; 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&times; 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 "
            "&mdash; a task with many acceptable answers and one reference "
            "&mdash; 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&times; 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 &mdash; 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


# ---------------------------------------------------------------------------
# Part 3
# ---------------------------------------------------------------------------
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 &minus;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 &minus;&infin;, 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 &mdash; 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."))
        # Figures 9 and 10 belong with the discussion in 3.3, not here.
    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"])
    # Strategy whose own-output perplexity sits closest to real human text.
    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} &mdash; more than half of all "
        f"bigrams it emits are repeats of one it has already emitted in the "
        f"same continuation &mdash; 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> &mdash; 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 &mdash; top-k 10, top-k 50, top-p 0.90, "
        f"top-p 0.95 &mdash; 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} &mdash; its own output is about "
        f"{ref_ppl/g['gen_ppl']:.0f}&times; 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 &mdash; and, on the same prompt, invents "
        "\"Pandas\", \"Dorrit\" and \"the K&aacute;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 &mdash; 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 &mdash; 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&middot;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> "
            "&mdash; <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 &mdash; quadrupling the beam costs "
            f"{beams[-1]['seconds']/t1:.2f}&times;, not 4&times;. 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&middot;W</i> rows, and at "
            f"<i>B</i>&nbsp;=&nbsp;{D['batch_size']} with a 160M-parameter "
            f"model the GPU is nowhere near saturated at <i>W</i>&nbsp;=&nbsp;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 &mdash; at larger batch, a larger model, "
            f"or large enough <i>W</i> &mdash; 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}&times; the compute and "
            f"{int(b4['strategy'].split('_')[1])}&times; the cache memory. A "
            f"wider beam is strictly better at what it optimises &mdash; it "
            f"finds sequences of higher log-probability, as the score check in "
            f"section 3 confirms &mdash; 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 &mdash; 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 &ge; 0.999 everywhere), and "
            "specialisation still emerged along language lines &mdash; 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&times; 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 &mdash; "
            "and at this task it buys almost no quality for that cost.")
    for it in items:
        a(Paragraph(f"&bull;&nbsp; {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&ndash;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}]&nbsp; {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))
            # Compare like with like: both spreads in nats of validation loss.
            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 &mdash; "
                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 &mdash; 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}&times; "
            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 &mdash; "
            f"a rare combination, and consistent with the strong showing "
            f"matrix-based methods get in the paper. Its cost is the "
            f"Newton&ndash;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}&times; model size against "
            f"{adamw['state_multiple']:.1f}&times;) 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}&times; 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}&times; model size in state "
            f"&mdash; 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&times;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"&plusmn;<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&ndash;Newton estimate is least informative early in "
            f"training, far from any minimum &mdash; which is exactly and only "
            f"where this experiment lives, at "
            f"{so['chinchilla_ratio']:.3f}&times; 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) &rho; were left at "
        "published defaults, which favours whichever optimizer's defaults "
        "happen to suit a 27M-parameter model at 0.1&times; 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


# ---------------------------------------------------------------------------
# Front matter and assembly
# ---------------------------------------------------------------------------
def front_matter(M1, M2, D):
    f = []
    a = f.append
    a(P("Advanced NLP &mdash; Assignment 2", "title"))
    a(P("Mixture of Experts &middot; Optimizers &middot; 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&rarr;English translation under matched parameter "
        "budgets, and asks what sparsity costs and what it buys. Part 2 "
        "implements five optimizers from scratch &mdash; one per category of "
        "Table 1 of Wen et al. [3], including the bonus Hessian-based entry "
        "&mdash; 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("&mdash;", "-"),
            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()