nils/jspace
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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)

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-17 00:46:19 +02:00
co-authored by Claude Fable 5
parent 62ce30e557
commit bbf8a02c6b
3 changed files with 96 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
results-loop/PROTOCOL_UNIFIED.md
+35
View File
@@ -776,3 +776,38 @@ after it (job renamed zzz_t).
(forces state-borne computation rather than KV re-reading) and
anchor-at-prompt (iterations see the full settled question).
Job: scripts/jobs/zzz_s_rungb_ii.sh.
IN-FLIGHT NOTE (~02:20): lce 9.4 -> 2.9 by step 60 — the tape-free
burst also encodes the step trajectory; encoding is never the
obstacle. Accuracy pending.
28. **Teacher-state distillation into the burst (pre-registered
2026-07-17 ~02:30, before running; Nils's variant: "meaningful CoT
chunks yield internal state that we then teacher-force into the
loop-only model").** Items 25/27 supervise the VERBAL SHADOW of
the deleted step (token targets through the lens); this forces the
FULL state. Teacher = the frozen warm-start (rung-A e400) adapter
running the complete cot (deleted step visible), zero pauses, same
carry architecture; capture its band-exit state at the deleted
step's LAST token — "the state of having finished thinking the
step" — one no-grad pass over the 413 parseable items at startup
(targets fixed, no moving teacher). Student: identical zero-pause
M=10 burst as item 27, but the loss is cosine distance between the
burst's FINAL iterate s^10 and the teacher state, plus output CE.
Weight AMENDED pre-run λ=1.0 -> 5.0 after the smoke measured the
starting cosine distance at only 0.113 (nearby band-exit states
share most structure); at 1.0 the term would be ~10x weaker than
the CE and a null would be underpowered — 5x0.113 puts the two
terms at comparable initial scale. No lens loss (one knob vs item 27: full-state
targets instead of verbal-shadow targets; trajectory supervision
dropped — only the endpoint is forced). Known approximation,
stated in advance: teacher state is captured ~10 positions later
in the sequence than the student anchor (RoPE position coloring);
cosine + the carry machinery's routine state transplantation
across nearby positions make this tolerable, but a null could
partly hide here. Eval n=256: 0:0, 2:0 ii10 (matched), 2:0 ii0
(ablation). Decision vs 31.6, same bands; the informative
three-way is 25 (verbal, tape) vs 27 (verbal, no tape) vs 28
(full state, no tape) — if 28 moves where 25/27 don't, the
computation-carrying content lives OUTSIDE the verbalizable
subspace; if none move, the read-side clamp test (item 29
candidate) decides. Job: scripts/jobs/zzz_sa_rungb_ts.sh.
scripts/jobs/zzz_sa_rungb_ts.sh
+14
View File
@@ -0,0 +1,14 @@
# 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_ts*.json results-loop/train_carrycot_b1_ii10_ts50_p0_log.json results-loop/adapter_carrycot_b1_ii10_ts50_p0_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
$P train_carry_cot.py --drop-steps 1 --pause-per-step 0 --base-pauses 0 \
--inner-iters 10 --teachstate 5.0 \
--warm-start $LOOP_OUT/adapter_carrycot_e400.pt --steps 200 --lr 3e-4
A=$LOOP_OUT/adapter_carrycot_b1_ii10_ts50_p0_e200.pt
$P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b1_ts10 \
--grid 0:0,2:0 --n 256 --inner-iters 10
$P eval_carry_cot.py --adapter $A --tag gsm_carrycot_b1_ts10_ablate \
--grid 2:0 --n 256 --inner-iters 0
scripts/train_carry_cot.py
+47 -2
View File
@@ -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: