feat(e5): MTP-aux distillation rung — registered + implemented

EXPERIMENTS.md E5: control = run-006-r7-muon verbatim, single delta =
a 3-step multi-token-prediction head (~1.1M params, 0.4%) on the
student decoder, aux CE at beta 0.15, head discarded at inference so
the shipped artifact stays vanilla and size-identical. Source: Hy4's
native MTP mapped as a TRAINING auxiliary (TODO 07-hy4) — serving-side
speculation stays parked, decode measured non-binding.

gpu/mtp.py (MTPHead + build/mtp_named), distill_sequence wiring
(head before optimizer so Muon sees its matrices; mtp_head.pt saved
beside student.pt for resume; provenance-only copy under best/),
spec ara-diac-small-2-mtp -> run-007-r7-muon-mtp. 3 unit tests
(shift/mask semantics, uniform-logit baseline, naming).

Gate (pre-agreed): adopt at <=4.5218 full-set windowed DER; honest
report in [4.5218, 4.8218). Prediction: 4.5-4.75.

Author: Ronald Tse

docs/EXPERIMENTS.md

@@ -195,6 +195,31 @@ All rows passed the CER parity gate at release. Readings:
   teacher+0.5pp budget remains the disclosed north star.
 - **Prediction (registered):** 4.3–5.0, by E3's 4.829 on weaker labels.
 
+## E5 — MTP-aux distillation rung (run-007-r7-muon-mtp)
+
+- **Status:** REGISTERED 2026-09-01, launching.
+- **Source:** Tencent Hy4-preview carries a native multi-token
+  prediction layer; mapped to our stack as a TRAINING auxiliary (TODO
+  07-hy4) — per-position multi-step heads densify supervision on the
+  decode path; serving-side speculation stays parked (decode measured
+  non-binding at our sizes).
+- **Hypothesis:** the student's residual errors concentrate where the
+  single-step objective leaves the byte decision underconstrained;
+  forcing each decoder position to also predict t+1..t+3 regularizes
+  the hidden state toward the local sequence structure that
+  diacritization output exhibits (letter + haraqat pattern).
+- **Design:** control = run-006-r7-muon verbatim (same teacher labels,
+  corpus, limits, schedule, Muon groups); single delta = MTPHead
+  attached to the student decoder (3 steps, byte-vocab 259, ~1.1M
+  params ≈ 0.4%), auxiliary CE weighted β=0.15, head DISCARDED at
+  inference (zero serving cost/size delta in the shipped artifact).
+- **Pre-agreed gate:** adopt if full-set windowed DER ≤ 4.5218
+  (≥0.3pp over 4.8218, the E3-style bar); report honestly in
+  [4.5218, 4.8218); investigate if worse.
+- **Prediction (registered):** 4.5–4.75 — denser supervision helps
+  the tail, but the factorial attributes most of the remaining gap to
+  domain coverage, so the effect should be second-order.
+
 ## Parked
 
 - **Speculative decoding** (LongCat converts sparsity→speed): revisit

src/gpu/distill_specs.yaml

@@ -116,6 +116,29 @@ ara-diac-small-muon:
   labels_complete: 'true'
   mode: sequence
   note: vanilla ByT5-small + Muon (factorial cell 4)
+ara-diac-small-2-mtp:
+  # E5 (EXPERIMENTS.md): control = ara-diac-small-2 verbatim; single
+  # delta = MTP-aux head (3 steps, beta 0.15), discarded at inference
+  teacher: rababa_arabic_byt5/run-007-news/best
+  teacher_volume: rababa
+  student_init: google/byt5-small
+  optimizer: muon
+  muon_lr: '0.01'
+  train: r5-units/domain.txt
+  train_extra:
+  - r5-units/replay.txt
+  unit_limits:
+  - 24000
+  - 6000
+  max_len: 1450
+  label_beams: '1'
+  out: rababa_arabic_distill_small/run-007-r7-muon-mtp
+  labels_file: teacher_labels_r7.jsonl
+  mtp_aux:
+    steps: 3
+    beta: 0.15
+  mode: sequence
+
 ara-diac-small-2:
   teacher: rababa_arabic_byt5/run-007-news/best
   teacher_volume: rababa

src/gpu/modal_distill.py

@@ -429,6 +429,8 @@ def val_loss() -> float:
                 ck.mkdir(exist_ok=True)
                 torch.save(student.state_dict(), ck / "student.pt")
                 torch.save(optimizer.state_dict(), ck / "optim.pt")
+                if mtp_head is not None:
+                    torch.save(mtp_head.state_dict(), ck / "mtp_head.pt")
                 CHECKPOINTS.commit()
 
     vl = val_loss()
@@ -742,6 +744,14 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict:
         from gpu.pkm import inject_pkm
 
         inject_pkm(student, **spec["pkm"])
+    mtp_head = None
+    if spec.get("mtp_aux"):
+        _ensure_src_path()
+        from gpu.mtp import build_mtp
+
+        mtp_head = build_mtp(student, **spec["mtp_aux"])
+        n = sum(q.numel() for q in mtp_head.parameters()) / 1e6
+        print(f"[{spec_id}] mtp_aux head: {n:.2f}M params", flush=True)
     student.train()
 
     class Pairs(Dataset):
@@ -1043,7 +1053,12 @@ def __getitem__(self, i):
         _ensure_src_path()
         from gpu.muon import Muon, split_parameters
 
-        muon_params, adamw_params = split_parameters(student.named_parameters())
+        named = list(student.named_parameters())
+        if mtp_head is not None:
+            from gpu.mtp import mtp_named
+
+            named += list(mtp_named(mtp_head))
+        muon_params, adamw_params = split_parameters(named)
         optimizer = Muon(
             muon_params, lr=float(spec.get("muon_lr", 0.01)),
             momentum=0.95, weight_decay=0.01,
@@ -1083,6 +1098,16 @@ def _usable(ck: Path) -> bool:
         optimizer.load_state_dict(
             torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True)
         )
+        if mtp_head is not None and (ckpts[-1] / "mtp_head.pt").exists():
+            mtp_head.load_state_dict(
+                torch.load(ckpts[-1] / "mtp_head.pt", map_location="cpu", weights_only=True)
+            )
+        elif mtp_head is not None:
+            print(
+                f"[{spec_id}] WARNING: no mtp_head.pt at resume — "
+                "fresh head, aux dynamics reset",
+                flush=True,
+            )
         step = int(ckpts[-1].name.split("-")[1])
         for _ in range(step):
             scheduler.step()
@@ -1093,7 +1118,18 @@ def _usable(ck: Path) -> bool:
             if step >= total_steps:
                 break
             ids, am, labels = ids.to("cuda"), am.to("cuda"), labels.to("cuda")
-            loss = student(input_ids=ids, attention_mask=am, labels=labels).loss
+            if mtp_head is not None:
+                beta = float(spec["mtp_aux"].get("beta", 0.15))
+                s_out = student(
+                    input_ids=ids, attention_mask=am, labels=labels,
+                    output_hidden_states=True,
+                )
+                aux = mtp_head.aux_loss(
+                    s_out.decoder_hidden_states[-1], labels
+                )
+                loss = s_out.loss + beta * aux
+            else:
+                loss = student(input_ids=ids, attention_mask=am, labels=labels).loss
             loss.backward()
             torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
             optimizer.step()
@@ -1119,6 +1155,8 @@ def _usable(ck: Path) -> bool:
     best.mkdir(exist_ok=True)
     student.save_pretrained(str(best))
     student_tok.save_pretrained(str(best))
+    if mtp_head is not None:  # provenance only; never in the shipped artifact
+        torch.save(mtp_head.state_dict(), best / "mtp_head.pt")
     CHECKPOINTS.commit()
     SECRYST_CHECKPOINTS.commit()
     PERSIAN_CHECKPOINTS.commit()

src/gpu/mtp.py

@@ -0,0 +1,51 @@
+"""Multi-token-prediction auxiliary head (E5, TODO 07-hy4).
+
+Per-position multi-step prediction as a TRAINING auxiliary: each
+decoder position's hidden state additionally predicts the target
+tokens at t+1..t+k, densifying supervision on the decode path. The
+head is discarded at inference — the shipped student stays vanilla
+and size-identical; only the training run carries it (saved as
+mtp_head.pt beside student.pt for resume, never exported).
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class MTPHead(nn.Module):
+    def __init__(self, d_model: int, vocab: int, steps: int = 3):
+        super().__init__()
+        self.steps = steps
+        self.heads = nn.ModuleList(
+            nn.Linear(d_model, vocab) for _ in range(steps)
+        )
+
+    def aux_loss(self, hidden: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        """hidden [B, T, d] is the decoder's final hidden state whose
+        position t already predicts labels[t] through the main head;
+        step-k heads predict labels[t+k] from the same position."""
+        total = 0.0
+        n = 0
+        for k, head in enumerate(self.heads, start=1):
+            tgt = labels[:, k:]
+            m = tgt != -100
+            if m.any():
+                total = total + F.cross_entropy(
+                    head(hidden[:, :-k])[m].float(), tgt[m]
+                )
+                n += 1
+        return total / max(n, 1)
+
+
+def build_mtp(student, steps: int = 3) -> MTPHead:
+    return MTPHead(
+        student.config.d_model, student.config.vocab_size, steps=steps
+    ).to(student.device)
+
+
+def mtp_named(head: MTPHead):
+    for name, p in head.named_parameters():
+        yield f"mtp.{name}", p

tests/test_gpu_mtp.py

@@ -0,0 +1,51 @@
+import torch
+
+from gpu.mtp import MTPHead, mtp_named
+
+
+def _labels():
+    t = torch.full((2, 6), -100)
+    t[0, :5] = torch.tensor([10, 11, 12, 13, 14])
+    t[1, :3] = torch.tensor([4, 5, 6])
+    return t
+
+
+def test_mtp_head_shapes_and_shift():
+    torch.manual_seed(0)
+    head = MTPHead(d_model=8, vocab=16, steps=3)
+    hidden = torch.randn(2, 6, 8)
+    loss = head.aux_loss(hidden, _labels())
+    assert loss.ndim == 0 and float(loss) > 0
+
+
+def test_mtp_shift_targets_not_inputs():
+    torch.manual_seed(0)
+    head = MTPHead(d_model=8, vocab=16, steps=1)
+    head.heads[0].weight.data.zero_()
+    head.heads[0].bias.data.zero_()
+
+    labels = _labels()
+    hidden = torch.randn(2, 6, 8)
+    # uniform-logit head: CE is log(vocab) wherever any target is valid
+    loss = head.aux_loss(hidden, labels)
+    import math
+
+    assert math.isclose(float(loss), math.log(16), rel_tol=1e-4)
+
+    # a step-1 head must see labels[:, 1:] as targets: with all -100
+    # beyond position 0 there is nothing to predict
+    only_first = torch.full((1, 4), -100)
+    only_first[0, 0] = 7.0
+    assert float(head.aux_loss(torch.randn(1, 4, 8), only_first)) == 0.0
+
+
+def test_mtp_learnable_and_named():
+    head = MTPHead(d_model=8, vocab=16, steps=2)
+    named = dict(mtp_named(head))
+    assert set(named) == {
+        "mtp.heads.0.weight", "mtp.heads.0.bias",
+        "mtp.heads.1.weight", "mtp.heads.1.bias",
+    }
+    out = head.aux_loss(torch.randn(2, 6, 8), _labels())
+    out.backward()
+    assert head.heads[0].weight.grad is not None