item 28 pre-registered (Nils's variant): teacher-state distillation — frozen full-cot teacher's band-exit state at step end, cosine into burst s^10; λ amended 1.0→5.0 pre-run (smoke: baseline cos-dist 0.113)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-17 00:46:19 +02:00
parent 62ce30e557
commit bbf8a02c6b
3 changed files with 96 additions and 2 deletions
@@ -77,6 +77,12 @@ ap.add_argument("--base-pauses", type=int, default=-1, metavar="P",
                help="override P_BY_LABEL with a fixed pause count; 0 = NO "
                     "pause tokens at all (inner iterations anchor on the "
                     "last prompt position)")
+ap.add_argument("--teachstate", type=float, default=0.0, metavar="LAMBDA",
+                help="item 28 (Nils's variant): teacher-state distillation "
+                     "— frozen warm-start adapter runs the FULL cot (step "
+                     "visible); its band-exit state at the deleted step's "
+                     "last token becomes the target; burst final iterate "
+                     "trained to it by cosine, weighted LAMBDA")
 ARGS = ap.parse_args()
 STEPS, LR = ARGS.steps, ARGS.lr
 TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
@@ -86,6 +92,8 @@ TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
    f"_lg{str(ARGS.lensteach_gen).replace('.', '')}"
    if ARGS.lensteach_gen else "") + (
    f"_ii{ARGS.inner_iters}" if ARGS.inner_iters else "") + (
+    f"_ts{str(ARGS.teachstate).replace('.', '')}"
+    if ARGS.teachstate else "") + (
    f"_p{ARGS.base_pauses}" if ARGS.base_pauses >= 0 else "") + (
    f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
    f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix
@@ -313,6 +321,29 @@ def main():
            print(f"lens-teach-gen λ={ARGS.lensteach_gen}: staging targets "
                  f"on {n_gt} items ({n_spans} line-spans; pre-'=' "
                  f"positions target the line result)", flush=True)
+    if ARGS.teachstate:
+        assert ARGS.warm_start and ARGS.inner_iters, \
+            "teachstate needs --warm-start (frozen teacher) + --inner-iters"
+        t0_ = time.time()
+        todo = [it for it in data if it["deleted"]]
+        pt = ARGS.base_pauses if ARGS.base_pauses >= 0 else 0
+        with torch.no_grad():
+            for i in range(0, len(todo), BATCH):
+                chunk = todo[i:i + BATCH]
+                full_items = [{"question": c["question"],
+                               "cot": c["deleted"] + "\n" + c["cot"],
+                               "extra_pauses": 0} for c in chunk]
+                ids_, msk_, _, plens_ = build_batch(tok, full_items, pt)
+                _, S_ = carry_logits(looper, adapter, ids_, msk_, plens_,
+                                     K_PREFILL, return_states=True)
+                for b, c in enumerate(chunk):
+                    nstep = len(tok(c["deleted"],
+                                    add_special_tokens=False)["input_ids"])
+                    pos = int(plens_[b]) + pt + nstep - 1
+                    c["teacher_state"] = S_[b, pos].float().clone()
+        print(f"teacher states: {len(todo)} captured ({time.time()-t0_:.0f}s;"
+              f" frozen warm-start adapter, full cot, band-exit at the "
+              f"deleted step's last token)", flush=True)
    groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}]
    if lora_params:
        groups.append({"params": lora_params, "lr": ARGS.lora_lr,
@@ -352,7 +383,7 @@ def main():
                "inner-iters needs a batch-uniform pause block"
            ckw = dict(inner_iters=ARGS.inner_iters,
                       inner_at=plens + p + extras - 1)
-            if ARGS.lensteach:
+            if ARGS.lensteach or ARGS.teachstate:
                itstates = []
                ckw["iter_states"] = itstates
        if ARGS.lensteach or ARGS.lensteach_gen or ARGS.inner_iters:
@@ -414,16 +445,30 @@ def main():
                gl = torch.stack(gterms).mean()
                lgen_val = gl.item()
                loss = loss + ARGS.lensteach_gen * gl
+        lts_val = 0.0
+        if ARGS.teachstate and itstates:
+            tterms = []
+            for b, it in enumerate(batch):
+                T = it.get("teacher_state")
+                if T is None:
+                    continue
+                tterms.append(1 - F.cosine_similarity(
+                    itstates[-1][b].float(), T, dim=0))
+            if tterms:
+                lt = torch.stack(tterms).mean()
+                lts_val = lt.item()
+                loss = loss + ARGS.teachstate * lt
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            list(adapter.parameters()) + lora_params, 1.0)
        opt.step()
        log.append({"step": step, "loss": loss.item(), "lce": lce_val,
-                    "lgen": lgen_val})
+                    "lgen": lgen_val, "lts": lts_val})
        if step % 10 == 0:
            print(f"step {step:4d}  {lbl:4s} loss={loss.item():.4f}  "
                  f"lce={lce_val:.3f}  lgen={lgen_val:.3f}  "
+                  f"lts={lts_val:.3f}  "
                  f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
        if step % 200 == 199 or step == STEPS - 1:
            if ARGS.lensnoise: