diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index 71f3bb9..2b99bc9 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/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. diff --git a/scripts/jobs/zzz_u_traj_d1.sh b/scripts/jobs/zzz_u_traj_d1.sh new file mode 100644 index 0000000..509fec4 --- /dev/null +++ b/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 diff --git a/scripts/jobs/zzz_v_traj_full.sh b/scripts/jobs/zzz_v_traj_full.sh new file mode 100644 index 0000000..51dc894 --- /dev/null +++ b/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 diff --git a/scripts/train_carry_cot.py b/scripts/train_carry_cot.py index 8c32923..e01688f 100644 --- a/scripts/train_carry_cot.py +++ b/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 - if ARGS.inner_iters: - extras = torch.tensor([b_.get("extra_pauses", 0) 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) + ckw, itstates, tfp, frs = {}, None, None, None + if ARGS.traj_tf or ARGS.traj_fr: + T = torch.stack([torch.as_tensor(b_["traj_states"]) + for b_ in batch]).cuda() + logits, tfp, frs = traj_burst_forward( + looper, adapter, ids, msk, plens, ARGS.inner_iters, T, + bool(ARGS.traj_tf), K_PREFILL) else: - logits = carry_logits(looper, adapter, ids, msk, plens, - K_PREFILL, use_checkpoint=True, - feedforward=ARGS.feedforward) + if ARGS.inner_iters: + extras = torch.tensor([b_.get("extra_pauses", 0) + 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: