diff --git a/results-loop/PROTOCOL_UNIFIED.md b/results-loop/PROTOCOL_UNIFIED.md index 9d51f85..bd597db 100644 --- a/results-loop/PROTOCOL_UNIFIED.md +++ b/results-loop/PROTOCOL_UNIFIED.md @@ -703,3 +703,32 @@ auxiliary supervision, A2, E2-N. temporal alignment is one choice among several (bag-of-tokens, result-digits-only are untested alternatives if this null's). Job: scripts/jobs/zzz_q_rungb_lt.sh. + IN-FLIGHT NOTE (~03:20): λ=0.3 arm's lens-CE collapsed 10.3 -> 2.3 + by step 40 — world 1 (cannot encode) is eliminated before the eval; + the open question is world 2 (writes but doesn't compute) vs + world 3 (reopens the ladder). + +26. **Result-staging supervision during generation (pre-registered + 2026-07-17 ~03:25, before running; Nils: "can we try to train the + loop using jspace during generation?").** The microscopy's + signature of the healthy board: each line's RESULT is on the + carried state before it is emitted ("1/twelve" at the '=', + P=0.997). This trains that signature directly. For every visible + scratchpad line, at the pre-'=' positions (the result is NOT yet + in the causal context there — nothing can be copied; low loss + requires computation) lens-CE the L30 carried state against the + line's result tokens; Answer-line analog via the ':' separator. + Spans computed in token space (gen_staging_targets; 1102 spans on + the d=1-transformed 427). Two arms, one submit, both otherwise + item-22 d=1 (warm-start rung-A e400, 200 steps, adapter-only): + (a) lg=0.3 alone — staging supervision without pause supervision; + (b) lt=0.3 + lg=0.3 — the full lens curriculum (deleted step + taught at the pauses, results staged during emission). Eval n=256: + 0:0, 2:12, 2:16 per arm; decision vs 31.6, same bands as items + 23-25. Distinct from item 25's question: 25 teaches WHAT the + silent chain should hold; 26 teaches the compute-before-emit + discipline on the visible chain (the behavior that separated + carry from FF at the 430 divergence). Caveat: on-policy ("during + generation" in the free-running sense) is NOT this — that's a + rollout-based follow-up (latent DAgger) if either lens arm moves. + Job: scripts/jobs/zzz_r_rungb_lg.sh. diff --git a/scripts/jobs/zzz_r_rungb_lg.sh b/scripts/jobs/zzz_r_rungb_lg.sh new file mode 100644 index 0000000..f84f4b3 --- /dev/null +++ b/scripts/jobs/zzz_r_rungb_lg.sh @@ -0,0 +1,16 @@ +# 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_lg*.json results-loop/eval_gsm_carrycot_b1_lt03_lg*.json results-loop/train_carrycot_b1_l*_log.json results-loop/adapter_carrycot_b1_l*_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 +# arm a: staging-only +$P train_carry_cot.py --drop-steps 1 --warm-start $LOOP_OUT/adapter_carrycot_e400.pt \ + --steps 200 --lr 3e-4 --lensteach-gen 0.3 +$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_b1_lg03_e200.pt \ + --tag gsm_carrycot_b1_lg03 --grid 0:0,2:12,2:16 --n 256 +# arm b: full lens curriculum (pauses + staging) +$P train_carry_cot.py --drop-steps 1 --warm-start $LOOP_OUT/adapter_carrycot_e400.pt \ + --steps 200 --lr 3e-4 --lensteach 0.3 --lensteach-gen 0.3 +$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_b1_lt03_lg03_e200.pt \ + --tag gsm_carrycot_b1_lt03_lg03 --grid 0:0,2:12,2:16 --n 256 diff --git a/scripts/train_carry_cot.py b/scripts/train_carry_cot.py index ba4a68d..2d575b7 100644 --- a/scripts/train_carry_cot.py +++ b/scripts/train_carry_cot.py @@ -63,12 +63,19 @@ ap.add_argument("--lensteach", type=float, default=0.0, metavar="LAMBDA", "replacement pauses against the DELETED step's tokens " "(1:1 pause j <-> step token j), mixed at LAMBDA into " "the output CE") +ap.add_argument("--lensteach-gen", type=float, default=0.0, metavar="LAMBDA", + help="item 26: result-staging supervision during generation " + "— at each visible scratchpad line's pre-'=' positions " + "(result not yet in causal context), lens-CE the " + "carried state against that line's result tokens") ARGS = ap.parse_args() STEPS, LR = ARGS.steps, ARGS.lr TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + ( f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + ( f"_blr{ARGS.bandlora}" if ARGS.bandlora else "") + ( f"_lt{str(ARGS.lensteach).replace('.', '')}" if ARGS.lensteach else "") + ( + f"_lg{str(ARGS.lensteach_gen).replace('.', '')}" + if ARGS.lensteach_gen else "") + ( f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + ( f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix @@ -112,6 +119,41 @@ class LensNoiseWrapper(torch.nn.Module): return self.base(e, s) +def gen_staging_targets(tok, cot): + """Item 26: per-line result-staging spans, computed in TOKEN space. + For each cot line with '=', the pre-'=' positions (result not yet in + causal context) target the line's result tokens; the Answer line's + 'Answer:' positions target the answer tokens. + Returns [(rel_positions, result_token_ids), ...] relative to tok(cot).""" + ids = tok(cot, add_special_tokens=False)["input_ids"] + decoded = [tok.decode([t]) for t in ids] + out, line_start = [], 0 + for i, d in enumerate(decoded + ["\n"]): + if "\n" not in d and i < len(ids): + continue + line = list(range(line_start, min(i, len(ids)))) + line_start = i + 1 + if not line: + continue + text = "".join(decoded[j] for j in line) + seps = [j for j in line if decoded[j].strip() == "="] + if seps: + sep = seps[-1] + elif text.strip().startswith("Answer"): + colons = [j for j in line if decoded[j].strip() == ":"] + if not colons: + continue + sep = colons[-1] + else: + continue + result = [ids[j] for j in line if j > sep + and decoded[j].strip()] + pre = [j for j in line if j <= sep] + if result and pre: + out.append((pre, result)) + return out + + def lr_at(step): if step < WARMUP: return LR * (step + 1) / WARMUP @@ -214,7 +256,7 @@ def main(): f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, " f"lr={ARGS.lora_lr}", flush=True) lens_teach = None - if ARGS.lensteach: + if ARGS.lensteach or ARGS.lensteach_gen: from loop_common import BAND J30 = torch.load(Path(__file__).resolve().parent.parent / "results/jbar.pt", @@ -230,16 +272,28 @@ def main(): lg = softcap * torch.tanh(lg / softcap) return lg - n_tgt = 0 - for it in data: - if it["deleted"]: - it["lens_targets"] = tok( - it["deleted"], add_special_tokens=False - )["input_ids"][: it["extra_pauses"]] - n_tgt += bool(it["lens_targets"]) - print(f"lens-teach λ={ARGS.lensteach}: targets on {n_tgt} items " - f"(pause j <-> deleted-step token j, lens at L{BAND[1]})", - flush=True) + if ARGS.lensteach: + n_tgt = 0 + for it in data: + if it["deleted"]: + it["lens_targets"] = tok( + it["deleted"], add_special_tokens=False + )["input_ids"][: it["extra_pauses"]] + n_tgt += bool(it["lens_targets"]) + print(f"lens-teach λ={ARGS.lensteach}: targets on {n_tgt} " + f"items (pause j <-> deleted-step token j, lens at " + f"L{BAND[1]})", flush=True) + if ARGS.lensteach_gen: + n_gt, n_spans = 0, 0 + for it in data: + gt = gen_staging_targets(tok, it["cot"]) + if gt: + it["gen_targets"] = gt + n_gt += 1 + n_spans += len(gt) + 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) groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}] if lora_params: groups.append({"params": lora_params, "lr": ARGS.lora_lr, @@ -270,7 +324,7 @@ 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 - if ARGS.lensteach: + if ARGS.lensteach or ARGS.lensteach_gen: logits, S = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL, use_checkpoint=True, feedforward=ARGS.feedforward, @@ -281,7 +335,7 @@ def main(): feedforward=ARGS.feedforward) loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), lab[:, 1:].flatten(), ignore_index=-100) - lce_val = 0.0 + lce_val, lgen_val = 0.0, 0.0 if ARGS.lensteach: terms = [] for b, it in enumerate(batch): @@ -297,15 +351,33 @@ def main(): lce = torch.stack(terms).mean() lce_val = lce.item() loss = loss + ARGS.lensteach * lce + if ARGS.lensteach_gen: + gterms = [] + for b, it in enumerate(batch): + base = int(plens[b]) + p + it.get("extra_pauses", 0) + for pre, res in it.get("gen_targets", []): + posl = torch.tensor([base + r for r in pre], + device=S.device) + lg = lens_teach(S[b, posl]).float() + gterms.append(torch.stack([ + F.cross_entropy( + lg, torch.full((len(pre),), r, dtype=torch.long, + device=S.device)) + for r in res]).mean()) + if gterms: + gl = torch.stack(gterms).mean() + lgen_val = gl.item() + loss = loss + ARGS.lensteach_gen * gl 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}) + log.append({"step": step, "loss": loss.item(), "lce": lce_val, + "lgen": lgen_val}) if step % 10 == 0: print(f"step {step:4d} {lbl:4s} loss={loss.item():.4f} " - f"lce={lce_val:.3f} " + f"lce={lce_val:.3f} lgen={lgen_val:.3f} " f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) if step % 200 == 199 or step == STEPS - 1: if ARGS.lensnoise: