feat(ride): displacement arm - ridge-projected teacher+base hidden regression (run-016)

TODO.sota-2026/01-r8-teacher-yallamorph.md

@@ -1,6 +1,8 @@
 # 01 — r8 teacher: run-009-yallamorph (YallaMorph/CamelMorph aux stream)
 
-Status: SPECIFIED (2026-10-01) — not launched
+Status: TRAINING IN FLIGHT (launched 2026-10-01 18:07, app ap-f1R8ChDikMKUGyBuJGO6;
+26,289 steps on A100-80GB; supervisor /tmp/r9-supervisor.sh v4 relaunches on
+kill-storm deaths — training is checkpoint-resumable and EVAL_DONE-idempotent)
 Literature basis: YallaMorph (arXiv 2609.10153, EMNLP 2026) — 663,804
 controlled morphological-generation instances over 4,795 lemmas,
 constructed from CamelMorph MSA via CAMeL Tools. The GitHub repo ships
@@ -44,16 +46,31 @@ Recipe = r7 verbatim (train_arabic_r7.py) + one new aux stream:
 
 ## Steps
 
-1. [ ] Clone CAMeL-Lab/YallaMorph; extract xlsx samples as the
+1. [x] Clone CAMeL-Lab/YallaMorph; extract xlsx samples as the
        validation set for our generated forms.
-2. [ ] Data build: camel-tools + camel_data MSA; generate paradigm
-       pairs; validate forms against YallaMorph samples (match rate
-       reported); dedupe; cap 300k aux lines; volume put to
-       /datasets/yallamorph-aux/lines.txt.
-3. [ ] TDD the line builder (pure function: feature dict → line pair).
-4. [ ] train_arabic_r9_yallamorph.py (r7 copy + stream M + gates).
-5. [ ] Launch `modal run --detach` (retry-loop supervisor per
-       [[modal-always-detach]]).
+2. [x] Data build: camel-tools 1.5.7 + **Camel Morph MSA v1.0**
+       (LREC-COLING 2024, CC BY 4.0 — the resource YallaMorph was
+       constructed from; camel_data's calima-msa-r13 CANNOT generate
+       mood/command forms — use the camel_morph repo DB). Key
+       interface lesson: generation must be UNDERSPECIFIED (pos +
+       proclitic variant only); fully-specified requests silently
+       reject cells with unmarked features (1st person gen='u').
+       13,000 lemmas (YallaMorph sample in-DB: v 513 / n 1,862 /
+       adj 594 + inventory top-up, seed 42) → 4,735,166 raw pairs →
+       300,000 lines (60/40 verb/nominal), 29MB, on volume
+       /datasets/yallamorph-aux/ (lines.txt + DONE + README).
+       Validation vs YallaMorph few-shot gold: **38/52 exact**;
+       residual divergence = proclitic-chain conventions
+       (hamzat-istifham أَلِـ, sin/lam stacking) + DB-version lemma
+       gaps — not incorrect forms.
+3. [x] TDD the line builder (build_yallamorph_aux.py, 12 tests
+       green; rababa PR #104).
+4. [x] train_arabic_r9_yallamorph.py (r7 copy + MORPH stream + gates;
+       rababa PR #104). Mix verified in launch logs: anchor=586,505,
+       news-mix=50,003, morph-used=197,169, **aux-share 25.00%**.
+5. [x] Launch `modal run --detach` (+ supervisor; one double-launch
+       incident from `modal app list` name truncation — grep prefix
+       "rababa-ara", dupes stopped, volume verified clean).
 6. [ ] ID gate: windowed zero-skip SadeedDiac-25 full 1,200-para DER
        ≤ 2.389 (r7 2.2864 + 0.1 tolerance).
 7. [ ] OOD gate: eval_wikinews_multiref improves over 17.3794/11.8273.

TODO.sota-2026/05-ride-sft-residual-extrapolation.md

@@ -1,6 +1,6 @@
 # 05 — RIDE-style SFT-residual extrapolation: probe-first arm
 
-Status: SPECIFIED (2026-10-01) — probe only; training arm gated on probe
+Status: PROBE PASSED (2026-10-01) — training arm spec'd below, launch pending owner call
 Literature basis: RIDE (arXiv 2609.36484) — extrapolate the
 teacher-over-base residual directly in representation space:
 student hidden states regressed toward
@@ -41,13 +41,49 @@ residual extrapolation of an SFT delta.
 
 ## Steps
 
-1. [ ] TDD pure computation: `residual_directions(h_base, h_teacher)`
-       and cosine sim on synthetic tensors (tests first, watch fail).
-2. [ ] Modal probe script (two models × 200 units × 2 domains;
-       A100 minutes, not hours).
-3. [ ] Run probe; write verdict + per-layer cosine table here.
-4. [ ] Gate decision: close, or spec the training arm separately.
+1. [x] TDD pure computation: `residual_directions(h_base, h_teacher)`
+       and cosine sim on synthetic tensors (7/7 green,
+       rababa/test_ride_probe_math.py).
+2. [x] Modal probe script (rababa/probe_ride_direction.py; A10G,
+       ~8 min, app ap-LePMpScA3RA8EN6TLDP2vN).
+3. [x] Run probe; verdict + per-layer cosine table below.
+4. [x] Gate decision: **TRANSFER** — training arm spec'd below.
 
-## Result
+## Result (measured 2026-10-01)
 
-(to be written only from measured numbers)
+Per-layer cos(d_classical, d_news), 200 units per domain, mean-pooled
+encoder hidden states, run-006-morph vs run-007-news:
+
+```
+L00 +0.9414  L01 +0.9513  L02 +0.9199  L03 +0.9002  L04 +0.8765
+L05 +0.8696  L06 +0.8593  L07 +0.7734  L08 +0.5931  L09 +0.4521
+L10 +0.2145  L11 -0.0015  L12 -0.1086  L13 -0.0842  L14 +0.0167
+L15 +0.1624  L16 +0.1971  L17 +0.2387  L18 +0.2775
+max +0.9513 >= 0.5  ->  TRANSFER
+```
+
+Reading: the r7-over-r6 SFT residual is strongly domain-general in
+early/mid encoder layers (L0–L8: 0.59–0.95) and idiosyncratic deep
+(L11+ ≈ noise). The news-mix fine-tune moved surface/orthographic
+processing in a direction that transfers to classical text — the
+RIDE displacement premise holds where representations are shared.
+Artefact: rababa-checkpoints:/ride_probe_r6_r7.json.
+
+## Training-arm spec (gated on this probe; launch = owner decision)
+
+The closure rule (TODO 03) permits this arm: it is a mechanism novel
+to the ledger (representation-space displacement; all prior student
+levers were loss/data/optimizer-side) and now has a measured transfer
+premise.
+
+- Infra: feature-regression aux loss in modal_distill — teacher/base
+  hidden states must be cached per layer subset. Restrict to L0–L8
+  (probe: only these transfer; deep-layer displacement would inject
+  domain idiosyncrasy).
+- Target: h_t' = h_teacher + λ(h_teacher − h_base), λ ∈ {0.5, 1.0}.
+- Loss: L = CE(labels) + β·MSE(h_student[L0-8], h_t'[L0-8]),
+  β tuned so MSE term ≈ 10% of total at start (pre-registration).
+- Single-variable off the 2.1 recipe (run-007 data, teacher labels,
+  Muon, seed 42); adopt gate ≥ 0.3pp DER improvement (E4-style bar).
+- Est. build: teacher/base hidden-state dump (one-off Modal job,
+  ~1h A100) + trainer loss path + spec; run cost ≈ one 2.1-recipe arm.

src/gpu/distill_specs.yaml

@@ -657,3 +657,45 @@ heb-diac-small-s46-layerdrop:
   val: hebrew-v4/val.jsonl
   out: rababa_hebrew_distill_small/run-003-s46-layerdrop
   note: Hebrew lite rung; cross-lingual depth-cut replication
+ara-diac-small-2-1-ride:
+  # TODO.sota-2026/05 arm: RIDE displacement. Probe passed 2026-10-01
+  # (r7-r6 residual direction domain-general in encoder L0-L8,
+  # max cos 0.9513). Single-variable off the 2.1 recipe (2-6ep
+  # verbatim): encoder-hidden regression toward ridge-projected
+  # h_t + lam*(h_t - h_b), frozen r7 teacher + r6 base, layers 0-8
+  # only (probe: deeper layers are domain-idiosyncratic).
+  # Gate: adopt iff >= 0.3pp better than 4.5701 (E4-style bar).
+  teacher: rababa_arabic_byt5/run-007-news/best
+  teacher_volume: rababa
+  out_volume: rababa
+  student_init: google/byt5-small
+  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-016-ride
+  labels_file: teacher_labels_r7.jsonl
+  labels_complete: 'true'
+  mode: sequence
+  optimizer: muon
+  ride_aux:
+    base: rababa_arabic_byt5/run-006-morph/best
+    layers:
+    - 0
+    - 1
+    - 2
+    - 3
+    - 4
+    - 5
+    - 6
+    - 7
+    - 8
+    lam: '1.0'
+    beta_frac: '0.1'
+    sub: '4'
+    fit_batches: '16'
+  note: 'TODO.sota-2026/05; gate vs 4.5701 (2.1) - adopt >= 0.3pp better'

src/gpu/modal_distill.py

@@ -789,6 +789,46 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict:
             f"sub={gkd_cfg.get('sample_sub', 2)} cap={gkd_cfg.get('sample_cap', 1024)}",
             flush=True,
         )
+    ride_cfg = spec.get("ride_aux")
+    ride = None
+    if ride_cfg:
+        # RIDE displacement arm (TODO.sota-2026/05; probe passed
+        # 2026-10-01: r7-r6 residual direction domain-general L0-L8).
+        # Regress student encoder hiddens toward ridge-projected
+        # h_t + lam*(h_t - h_b) from frozen teacher (r7) and base (r6).
+        _ensure_src_path()
+        from gpu.ride import displaced_targets, fit_ridge, masked_mse
+
+        base_path = str(
+            Path(VOLUME_MOUNTS[ride_cfg.get("base_volume", teacher_vol)]) / ride_cfg["base"]
+        )
+        base_tok = AutoTokenizer.from_pretrained(base_path)
+        base = (
+            AutoModelForSeq2SeqLM.from_pretrained(base_path)
+            .to("cuda", dtype=torch.float16)
+            .eval()
+        )
+        for p in base.parameters():
+            p.requires_grad_(False)
+        sample = "الْعَرَبِيَّةُ byte-parity 123"
+        ids_s = student_tok(sample).input_ids
+        for name, tok in (("teacher", teacher_tok), ("base", base_tok)):
+            if tok(sample).input_ids != ids_s:
+                raise RuntimeError(f"ride_aux: {name} tokenizer differs from student byte table")
+        ride = {
+            "base": base,
+            "layers": [int(x) for x in ride_cfg.get("layers", list(range(9)))],
+            "lam": float(ride_cfg.get("lam", 1.0)),
+            "beta_frac": float(ride_cfg.get("beta_frac", 0.1)),
+            "sub": int(ride_cfg.get("sub", 4)),
+            "beta": None,
+            "W": None,
+        }
+        print(
+            f"[{spec_id}] ride_aux: base={ride_cfg['base']} "
+            f"layers={ride['layers']} lam={ride['lam']} sub={ride['sub']}",
+            flush=True,
+        )
     student.train()
 
     class Pairs(Dataset):
@@ -1205,12 +1245,96 @@ def _usable(ck: Path) -> bool:
             scheduler.step()
         print(f"[{spec_id}] resume training from step-{step}", flush=True)
 
+    if ride is not None:
+        # displacement needs the teacher resident on cuda for both the
+        # ridge fit and every training step (labeling offloads it)
+        teacher.to("cuda")
+        teacher.eval()
+        # Ridge map per selected layer: teacher hidden space -> the
+        # student's INITIAL hidden space (linearity preserves the
+        # displacement through the projection). Fit on the first
+        # fit_batches batches, no grad, before any training step.
+        from gpu.ride import fit_ridge as _fit_ridge
+
+        ride_state_path = ckpts[-1] / "ride.pt" if ckpts else None
+        if ride_state_path is not None and ride_state_path.exists():
+            st = torch.load(ride_state_path, map_location="cpu", weights_only=True)
+            ride["W"] = {int(L): w.to("cuda") for L, w in st["W"].items()}
+            ride["beta"] = st["beta"]
+            print(f"[{spec_id}] ride resumed (beta={ride['beta']:.3e})", flush=True)
+        else:
+            it = iter(train_loader)
+            Ht = {L: [] for L in ride["layers"]}
+            Hs = {L: [] for L in ride["layers"]}
+            student.eval()
+            with torch.no_grad():
+                for _ in range(int(ride_cfg.get("fit_batches", 16))):
+                    try:
+                        b_ids, b_am, _lbl = next(it)
+                    except StopIteration:
+                        break
+                    b_ids, b_am = b_ids.to("cuda"), b_am.to("cuda")
+                    t_hs = teacher.encoder(
+                        input_ids=b_ids, attention_mask=b_am, output_hidden_states=True
+                    ).hidden_states
+                    s_hs = student.encoder(
+                        input_ids=b_ids, attention_mask=b_am, output_hidden_states=True
+                    ).hidden_states
+                    # mask-flatten to real positions: batches have
+                    # variable T, so keep (N_real, D) rows only
+                    keep = b_am.bool()
+                    for L in ride["layers"]:
+                        Ht[L].append(t_hs[L][keep].float().cpu())
+                        Hs[L].append(s_hs[L][keep].float().cpu())
+            student.train()
+            alpha = float(ride_cfg.get("ridge_alpha", 1.0))
+            ride["W"] = {
+                L: _fit_ridge(torch.cat(Ht[L]), torch.cat(Hs[L]), alpha=alpha).to("cuda")
+                for L in ride["layers"]
+            }
+            print(f"[{spec_id}] ride ridge fit on {len(Ht[ride['layers'][0]])} batches", flush=True)
+
     for _ in range(epochs):
         for ids, am, labels in train_loader:
             if step >= total_steps:
                 break
             ids, am, labels = ids.to("cuda"), am.to("cuda"), labels.to("cuda")
-            if mtp_head is not None:
+            if ride is not None:
+                s_out = student(
+                    input_ids=ids, attention_mask=am, labels=labels,
+                    output_hidden_states=True,
+                )
+                loss = s_out.loss
+                sub = ride["sub"]
+                with torch.no_grad():
+                    t_hs = teacher.encoder(
+                        input_ids=ids[:sub], attention_mask=am[:sub],
+                        output_hidden_states=True,
+                    ).hidden_states
+                    b_hs = ride["base"].encoder(
+                        input_ids=ids[:sub], attention_mask=am[:sub],
+                        output_hidden_states=True,
+                    ).hidden_states
+                mse_total = None
+                for L in ride["layers"]:
+                    tgt = displaced_targets(
+                        t_hs[L].float(), b_hs[L].float(), ride["lam"]
+                    ) @ ride["W"][L].T
+                    m = masked_mse(
+                        s_out.encoder_hidden_states[L][:sub].float(), tgt, am[:sub]
+                    )
+                    mse_total = m if mse_total is None else mse_total + m
+                if ride["beta"] is None:
+                    ride["beta"] = (
+                        ride["beta_frac"] * loss.detach() / (mse_total.detach() + 1e-12)
+                    ).item()
+                    print(
+                        f"[{spec_id}] ride beta calibrated: {ride['beta']:.3e} "
+                        f"(mse0={float(mse_total):.3f})",
+                        flush=True,
+                    )
+                loss = loss + ride["beta"] * mse_total
+            elif 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,
@@ -1270,6 +1394,12 @@ def _usable(ck: Path) -> bool:
                 (ck / "labels.sha").write_text(labels_digest)
                 torch.save(student.state_dict(), ck / "student.pt")
                 torch.save(optimizer.state_dict(), ck / "optim.pt")
+                if ride is not None and ride["beta"] is not None:
+                    torch.save(
+                        {"W": {L: w.cpu() for L, w in ride["W"].items()},
+                         "beta": ride["beta"]},
+                        ck / "ride.pt",
+                    )
                 CHECKPOINTS.commit()
                 SECRYST_CHECKPOINTS.commit()
                 PERSIAN_CHECKPOINTS.commit()

src/gpu/ride.py

@@ -0,0 +1,47 @@
+"""RIDE displacement-arm math (TODO.sota-2026/05).
+
+The probe (rababa/probe_ride_direction.py, 2026-10-01) measured the
+r7-over-r6 SFT residual direction as domain-general in encoder layers
+0-8 (cos 0.59-0.95 across classical/news text) and idiosyncratic
+deeper. These helpers implement the training arm: regress the
+student's encoder hidden states toward ridge-projected, extrapolated
+teacher targets h_t' = h_teacher + lam*(h_teacher - h_base).
+
+The ridge map is fit at arm start (teacher hiddens -> the student's
+initial hiddens on the same inputs); by linearity the displacement
+survives the projection: W(h_t + lam*(h_t - h_b)) = W h_t +
+lam*(W h_t - W h_b).
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def displaced_targets(h_teacher: torch.Tensor, h_base: torch.Tensor, lam: float) -> torch.Tensor:
+    return h_teacher + lam * (h_teacher - h_base)
+
+
+def masked_mse(pred: torch.Tensor, target: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    """Mean squared error over kept positions and all dims.
+
+    pred/target: (B, T, D); attention_mask: (B, T), 1 = kept.
+    Returns 0.0 when nothing is kept.
+    """
+    mask = attention_mask.to(pred.dtype).unsqueeze(-1)
+    diff2 = (pred - target).pow(2) * mask
+    n = mask.sum() * pred.shape[-1]
+    if n == 0:
+        return pred.new_tensor(0.0)
+    return diff2.sum() / n
+
+
+def fit_ridge(H_teacher: torch.Tensor, H_student: torch.Tensor, alpha: float) -> torch.Tensor:
+    """Closed-form ridge W minimizing ||H_teacher @ W.T - H_student||^2
+    + alpha*||W||^2. Returns W of shape (d_student, d_teacher)."""
+    Ht = H_teacher.double()
+    Hs = H_student.double()
+    d_t = Ht.shape[-1]
+    A = Ht.T @ Ht + alpha * torch.eye(d_t, dtype=Ht.dtype, device=Ht.device)
+    W = torch.linalg.solve(A, Ht.T @ Hs).T
+    return W.to(H_teacher.dtype)