item 29 pre-registered (Nils's design): trajectory teacher-forcing — 10 waypoint transitions supervised independently (T[i-1]→T[i]), TF and TF+FR arms; job 1 d=1 step-span, job 2 answer-only full-CoT span

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

results-loop/PROTOCOL_UNIFIED.md

@@ -867,3 +867,36 @@ at inference via the lens/embedding-basis machinery, no training, and
 measure) is item 29's natural content and would assign blame causally.
 Nils's morning decision: clamp test vs pivot to the hybrid/A2 line
 (where the 57.4 and the drop-bucket reach already live).
+
+29. **Trajectory teacher-forcing of the burst (pre-registered
+    2026-07-17 ~11:50, before running; Nils's design: "measure the CoT
+    run's L30 at 10 evenly spaced intervals and teacher force them
+    into the loop — we know both the previous that goes into the merge
+    layer and the next L30 state").** The key upgrade over item 28:
+    the teacher trajectory decomposes the burst's learning into TEN
+    INDEPENDENT SUPERVISED TRANSITIONS — iteration i receives teacher
+    waypoint T[i-1] as its merge input and its band output is pulled
+    onto T[i] (cosine, λ=5.0) — dense gradient at every iteration, no
+    long BPTT chain, no compounding; and T[0] == the student's own
+    settled anchor state exactly (shared prompt+settle), so the chain
+    starts aligned by construction. Exposure gap handled by a second
+    arm adding the free-running loss (burst's own s_i onto T[i],
+    λ=5.0). The free-running burst ALWAYS runs last and seeds the
+    answer scan (training matches inference); TF transitions are
+    scaffolding erased from context afterward. Teacher = frozen
+    warm-start adapter on the full cot, 10 evenly spaced L30 states.
+    Two jobs, per Nils (both chosen): JOB 1 (zzz_u) d=1, waypoints
+    across the deleted step, arms tf-only and tf+fr; reference 31.6,
+    same bands; eval 0:0 / 2:0 ii10 / 2:0 ii0 per arm. JOB 2 (zzz_v)
+    THE HEADLINE: d=all (answer-only output, rung C), waypoints across
+    the FULL CoT — the burst as a 10-waypoint compressed rehearsal of
+    the entire reasoning; references: cold answer-only 9.4, curriculum
+    plateau 19.1, base 10.9; same arms and cells. Smokes: ltf/lfr
+    baseline 0.116 both modes; d=99 deletes 1118 steps across 427.
+    Priors, stated: the transition-decomposition is the first
+    mechanism that trains the loop's DYNAMICS rather than its
+    endpoints — if the loop can compute at all, this is its best
+    chance; if job 2 beats 19.1 meaningfully, compressed-trajectory
+    rehearsal becomes the internalization method. All prior caveats
+    (position coloring, teacher=warm-start quality) carry over.
+    Jobs: scripts/jobs/zzz_u_traj_d1.sh, scripts/jobs/zzz_v_traj_full.sh.

scripts/jobs/zzz_u_traj_d1.sh

@@ -0,0 +1,15 @@
+# gpuq-in:  results-loop/star_data.json results-loop/gsm_cot_data.json results-loop/adapter_carrycot_e400.pt
+# gpuq-out: results-loop/eval_gsm_carrycot_b1_tj*.json results-loop/train_carrycot_b1_*tj*_log.json results-loop/adapter_carrycot_b1_*tj*_e200.pt
+git pull origin main -q 2>/dev/null
+P=/home/nils/jspace/.venv/bin/python
+export JLENS_MODEL=google/gemma-4-E2B-it LOOP_OUT=/home/nils/jspace/results-loop
+cd /home/nils/jspace/scripts
+for FR in 0.0 5.0; do
+  T=tjs50_$(echo $FR | tr -d .)
+  $P train_carry_cot.py --drop-steps 1 --traj-span step --traj-tf 5.0 --traj-fr $FR \
+     --pause-per-step 0 --base-pauses 0 --inner-iters 10 \
+     --warm-start $LOOP_OUT/adapter_carrycot_e400.pt --steps 200 --lr 3e-4
+  A=$LOOP_OUT/adapter_carrycot_b1_ii10_${T}_p0_e200.pt
+  $P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b1_$T --grid 0:0,2:0 --n 256 --inner-iters 10
+  $P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b1_${T}_ablate --grid 2:0 --n 256 --inner-iters 0
+done

scripts/jobs/zzz_v_traj_full.sh

@@ -0,0 +1,15 @@
+# gpuq-in:  results-loop/star_data.json results-loop/gsm_cot_data.json results-loop/adapter_carrycot_e400.pt
+# gpuq-out: results-loop/eval_gsm_carrycot_b99_tj*.json results-loop/train_carrycot_b99_*tj*_log.json results-loop/adapter_carrycot_b99_*tj*_e200.pt
+git pull origin main -q 2>/dev/null
+P=/home/nils/jspace/.venv/bin/python
+export JLENS_MODEL=google/gemma-4-E2B-it LOOP_OUT=/home/nils/jspace/results-loop
+cd /home/nils/jspace/scripts
+for FR in 0.0 5.0; do
+  T=tjf50_$(echo $FR | tr -d .)
+  $P train_carry_cot.py --drop-steps 99 --traj-span full --traj-tf 5.0 --traj-fr $FR \
+     --pause-per-step 0 --base-pauses 0 --inner-iters 10 \
+     --warm-start $LOOP_OUT/adapter_carrycot_e400.pt --steps 200 --lr 3e-4
+  A=$LOOP_OUT/adapter_carrycot_b99_ii10_${T}_p0_e200.pt
+  $P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b99_$T --grid 0:0,2:0 --n 256 --inner-iters 10
+  $P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b99_${T}_ablate --grid 2:0 --n 256 --inner-iters 0
+done

scripts/train_carry_cot.py

@@ -77,6 +77,18 @@ 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("--traj-tf", type=float, default=0.0, metavar="LAMBDA",
+                help="item 29: teacher-forced TRANSITION loss — iteration i "
+                     "gets teacher waypoint T[i-1] as merge input, its band "
+                     "output is pulled onto T[i] (cosine); 10 independent "
+                     "supervised one-step regressions")
+ap.add_argument("--traj-fr", type=float, default=0.0, metavar="LAMBDA",
+                help="item 29: free-running trajectory loss — the burst's "
+                     "own state s_i pulled onto waypoint T[i] (exposure-gap "
+                     "term)")
+ap.add_argument("--traj-span", choices=("step", "full"), default="step",
+                help="waypoints sampled evenly across the deleted step "
+                     "(job 1, d=1) or the full CoT (job 2, answer-only)")
 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 "
@@ -94,6 +106,9 @@ TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
     f"_ii{ARGS.inner_iters}" if ARGS.inner_iters else "") + (
     f"_ts{str(ARGS.teachstate).replace('.', '')}"
     if ARGS.teachstate else "") + (
+    f"_tj{ARGS.traj_span[0]}{str(ARGS.traj_tf).replace('.', '')}"
+    f"_{str(ARGS.traj_fr).replace('.', '')}"
+    if (ARGS.traj_tf or ARGS.traj_fr) 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
@@ -173,6 +188,49 @@ def gen_staging_targets(tok, cot):
     return out
 
 
+def traj_burst_forward(looper, adapter, ids, msk, plens, m, T, tf_on, k):
+    """Item 29 forward: optional teacher-forced transition predictions,
+    then the free-running burst (whose final state seeds the answer scan,
+    matching inference), then the visible-token carry. Returns
+    (logits, tf_preds, fr_states)."""
+    from torch.utils.checkpoint import checkpoint
+    from carry_common import prompt_prefill, build_step_updates, carry_steps
+    dev = ids.device
+    calls, _ = looper.capture(ids, msk, logits_to_keep=1)
+    e = looper._hin[looper.l0].detach()
+    B = ids.shape[0]
+    ar = torch.arange(ids.shape[1], device=dev)
+    pmask = ar[None, :] < plens[:, None].to(dev)
+    S, X = prompt_prefill(looper, adapter, e, calls, pmask, k)
+    rows = torch.arange(B, device=dev)
+    anchor = (plens - 1).to(dev)
+    s0 = S[rows, anchor]
+
+    def upd(S, X, seed):
+        x_new = adapter(e[rows, anchor], seed)
+        X = X.clone()
+        X[rows, anchor] = x_new.to(X.dtype)
+        S = checkpoint(lambda X_: looper.band(X_, calls), X,
+                       use_reentrant=False)
+        return S, X
+
+    tf_preds, fr_states = [], []
+    if tf_on:
+        for i in range(m):
+            seed = s0 if i == 0 else T[:, i - 1].to(s0.dtype)
+            S, X = upd(S, X, seed)
+            tf_preds.append(S[rows, anchor])
+    for i in range(m):
+        seed = s0 if i == 0 else fr_states[-1]
+        S, X = upd(S, X, seed)
+        fr_states.append(S[rows, anchor])
+    total = msk.sum(-1)
+    updates = build_step_updates(plens.to(dev), total.to(dev), dev)
+    S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
+                       use_checkpoint=True)
+    return looper.suffix_logits(S, calls), tf_preds, fr_states
+
+
 def lr_at(step):
     if step < WARMUP:
         return LR * (step + 1) / WARMUP
@@ -321,6 +379,35 @@ 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.traj_tf or ARGS.traj_fr:
+        assert ARGS.warm_start and ARGS.inner_iters \
+            and ARGS.base_pauses == 0, \
+            "traj mode needs --warm-start, --inner-iters, --base-pauses 0"
+        data = [it for it in data if it["deleted"]]
+        m = ARGS.inner_iters
+        t0_ = time.time()
+        with torch.no_grad():
+            for i in range(0, len(data), BATCH):
+                chunk = data[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, 0)
+                _, S_ = carry_logits(looper, adapter, ids_, msk_, plens_,
+                                     K_PREFILL, return_states=True)
+                for b, c in enumerate(chunk):
+                    if ARGS.traj_span == "step":
+                        span = len(tok(c["deleted"],
+                                       add_special_tokens=False)["input_ids"])
+                    else:
+                        span = len(tok(c["deleted"] + "\n" + c["cot"],
+                                       add_special_tokens=False)["input_ids"])
+                    idxs = [round(j * (span - 1) / (m - 1)) for j in range(m)]
+                    pos = [int(plens_[b]) + jj for jj in idxs]
+                    c["traj_states"] = S_[b, pos].float().clone()
+        print(f"trajectory waypoints: {len(data)} items x {m} states "
+              f"({ARGS.traj_span} span, {time.time()-t0_:.0f}s; frozen "
+              f"warm-start teacher, full cot)", flush=True)
     if ARGS.teachstate:
         assert ARGS.warm_start and ARGS.inner_iters, \
             "teachstate needs --warm-start (frozen teacher) + --inner-iters"
@@ -375,26 +462,34 @@ def main():
         ids, msk, lab, plens = build_batch(tok, batch, p)
         for g in opt.param_groups:
             g["lr"] = g["base"] * lr_at(step) / LR
-        ckw, itstates = {}, None
+        ckw, itstates, tfp, frs = {}, None, None, None
-        if ARGS.inner_iters:
+        if ARGS.traj_tf or ARGS.traj_fr:
-            extras = torch.tensor([b_.get("extra_pauses", 0) for b_ in batch],
+            T = torch.stack([torch.as_tensor(b_["traj_states"])
-                                  device=plens.device)
+                             for b_ in batch]).cuda()
-            assert (extras == extras[0]).all(), \
+            logits, tfp, frs = traj_burst_forward(
-                "inner-iters needs a batch-uniform pause block"
+                looper, adapter, ids, msk, plens, ARGS.inner_iters, T,
-            ckw = dict(inner_iters=ARGS.inner_iters,
+                bool(ARGS.traj_tf), K_PREFILL)
-                       inner_at=plens + p + extras - 1)
-            if ARGS.lensteach or ARGS.teachstate:
-                itstates = []
-                ckw["iter_states"] = itstates
-        if ARGS.lensteach or ARGS.lensteach_gen or ARGS.inner_iters:
-            logits, S = carry_logits(looper, adapter, ids, msk, plens,
-                                     K_PREFILL, use_checkpoint=True,
-                                     feedforward=ARGS.feedforward,
-                                     return_states=True, **ckw)
         else:
-            logits = carry_logits(looper, adapter, ids, msk, plens,
+            if ARGS.inner_iters:
-                                  K_PREFILL, use_checkpoint=True,
+                extras = torch.tensor([b_.get("extra_pauses", 0)
-                                  feedforward=ARGS.feedforward)
+                                       for b_ in batch],
+                                      device=plens.device)
+                assert (extras == extras[0]).all(), \
+                    "inner-iters needs a batch-uniform pause block"
+                ckw = dict(inner_iters=ARGS.inner_iters,
+                           inner_at=plens + p + extras - 1)
+                if ARGS.lensteach or ARGS.teachstate:
+                    itstates = []
+                    ckw["iter_states"] = itstates
+            if ARGS.lensteach or ARGS.lensteach_gen or ARGS.inner_iters:
+                logits, S = carry_logits(looper, adapter, ids, msk, plens,
+                                         K_PREFILL, use_checkpoint=True,
+                                         feedforward=ARGS.feedforward,
+                                         return_states=True, **ckw)
+            else:
+                logits = carry_logits(looper, adapter, ids, msk, plens,
+                                      K_PREFILL, use_checkpoint=True,
+                                      feedforward=ARGS.feedforward)
         loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
                                lab[:, 1:].flatten(), ignore_index=-100)
         lce_val, lgen_val = 0.0, 0.0
@@ -458,17 +553,32 @@ def main():
                 lt = torch.stack(tterms).mean()
                 lts_val = lt.item()
                 loss = loss + ARGS.teachstate * lt
+        ltf_val, lfr_val = 0.0, 0.0
+        if ARGS.traj_tf and tfp:
+            ltf = torch.stack([
+                (1 - F.cosine_similarity(p_.float(), T[:, i], dim=-1)).mean()
+                for i, p_ in enumerate(tfp)]).mean()
+            ltf_val = ltf.item()
+            loss = loss + ARGS.traj_tf * ltf
+        if ARGS.traj_fr and frs:
+            lfr = torch.stack([
+                (1 - F.cosine_similarity(s_.float(), T[:, i], dim=-1)).mean()
+                for i, s_ in enumerate(frs)]).mean()
+            lfr_val = lfr.item()
+            loss = loss + ARGS.traj_fr * lfr
         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, "lts": lts_val})
+                    "lgen": lgen_val, "lts": lts_val, "ltf": ltf_val,
+                    "lfr": lfr_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"lts={lts_val:.3f}  ltf={ltf_val:.3f}  "
+                  f"lfr={lfr_val:.3f}  "
                   f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
         if step % 200 == 199 or step == STEPS - 1:
             if ARGS.lensnoise: