item 25 pre-registered: latent process supervision via differentiable lens readout — pause j trained to lens-encode deleted-step token j (λ=0.3/1.0 arms); carry_logits gains return_states
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -674,3 +674,32 @@ the d=1 break. The carried state's native cargo (plans, magnitudes,
|
|||||||
completion-state — see the probe series) is the program's remaining
|
completion-state — see the probe series) is the program's remaining
|
||||||
asset; next candidates: divergence batch replay, coarse-target
|
asset; next candidates: divergence batch replay, coarse-target
|
||||||
auxiliary supervision, A2, E2-N.
|
auxiliary supervision, A2, E2-N.
|
||||||
|
|
||||||
|
25. **E2-L d=1 with latent process supervision through the lens
|
||||||
|
(pre-registered 2026-07-17 ~02:45, before running; Nils's idea:
|
||||||
|
"for training, i wonder if we could calculate, using jspace lens,
|
||||||
|
how each iteration should think").** Items 22-24 all trained the
|
||||||
|
latent chain blind — output CE only — and all failed; this changes
|
||||||
|
the INFORMATION reaching the chain, not its capacity. New loss:
|
||||||
|
the lens readout softmax(W_U·finalnorm(J̄_L30·h)) is differentiable
|
||||||
|
in h, so at the 10 replacement pauses we apply lens-CE against the
|
||||||
|
DELETED step's tokens, aligned 1:1 (pause j <-> step token j,
|
||||||
|
truncated at 10) — the board is trained to write the deleted step
|
||||||
|
in lens-readable code at the time it would have been written.
|
||||||
|
Mixed loss CE_out + λ·CE_lens. Two arms, single submit: λ=0.3 and
|
||||||
|
λ=1.0. Otherwise identical to item-22 d=1 (front-first deletion,
|
||||||
|
warm-start rung-A e400, adapter-only 3e-4, 200 steps, seed 0; no
|
||||||
|
band-LoRA — one knob). Smoke: step-0 lce=10.3 (~uniform: pauses
|
||||||
|
currently encode nothing about the step; large fresh gradient).
|
||||||
|
Eval n=256: 0:0, 2:12, 2:16 per arm. Decision vs d=1's 31.6, same
|
||||||
|
bands: >36.6 = latent supervision was the missing ingredient ->
|
||||||
|
ladder REOPENS with lens-taught rungs (and the 2D per-iteration
|
||||||
|
variant becomes item 26); within +-5 = even telling the board
|
||||||
|
exactly what to write doesn't make the carry compute it -> the
|
||||||
|
strongest closure evidence yet. Caveats pre-stated: J̄ is
|
||||||
|
prompt-averaged (global directions); the loss forces a
|
||||||
|
verbalizable code (microscopy suggests that IS the board's working
|
||||||
|
code, but a native non-verbal code would be fought); the 1:1
|
||||||
|
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.
|
||||||
|
|||||||
@@ -64,7 +64,8 @@ def build_step_updates(prompt_lens, total_lens, device):
|
|||||||
|
|
||||||
|
|
||||||
def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens,
|
def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens,
|
||||||
k, use_checkpoint=False, feedforward=False):
|
k, use_checkpoint=False, feedforward=False,
|
||||||
|
return_states=False):
|
||||||
"""Teacher-forced design-C forward (right-padded batch).
|
"""Teacher-forced design-C forward (right-padded batch).
|
||||||
|
|
||||||
feedforward=True: pause-token control — same positions get the adapter as
|
feedforward=True: pause-token control — same positions get the adapter as
|
||||||
@@ -81,13 +82,15 @@ def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens,
|
|||||||
S = (checkpoint(lambda x_: looper.band(x_, calls), x,
|
S = (checkpoint(lambda x_: looper.band(x_, calls), x,
|
||||||
use_reentrant=False) if use_checkpoint
|
use_reentrant=False) if use_checkpoint
|
||||||
else looper.band(x, calls))
|
else looper.band(x, calls))
|
||||||
return looper.suffix_logits(S, calls)
|
out = looper.suffix_logits(S, calls)
|
||||||
|
return (out, S) if return_states else out
|
||||||
S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k)
|
S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k)
|
||||||
total_lens = attention_mask.sum(-1)
|
total_lens = attention_mask.sum(-1)
|
||||||
updates = build_step_updates(prompt_lens.to(dev), total_lens.to(dev), dev)
|
updates = build_step_updates(prompt_lens.to(dev), total_lens.to(dev), dev)
|
||||||
S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
|
S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
|
||||||
use_checkpoint=use_checkpoint)
|
use_checkpoint=use_checkpoint)
|
||||||
return looper.suffix_logits(S, calls)
|
out = looper.suffix_logits(S, calls)
|
||||||
|
return (out, S) if return_states else out
|
||||||
|
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
|
|||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# 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_lt*.json results-loop/train_carrycot_b1_lt*_log.json results-loop/adapter_carrycot_b1_lt*.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 LAM in 0.3 1.0; do
|
||||||
|
T=lt$(echo $LAM | tr -d .)
|
||||||
|
$P train_carry_cot.py --drop-steps 1 --warm-start $LOOP_OUT/adapter_carrycot_e400.pt \
|
||||||
|
--steps 200 --lr 3e-4 --lensteach $LAM
|
||||||
|
$P eval_carry_cot.py --adapter $LOOP_OUT/adapter_carrycot_b1_${T}_e200.pt \
|
||||||
|
--tag gsm_carrycot_b1_$T --grid 0:0,2:12,2:16 --n 256
|
||||||
|
done
|
||||||
+67
-10
@@ -58,11 +58,17 @@ ap.add_argument("--bandlora", type=int, default=0, metavar="RANK",
|
|||||||
"band layer, uniform scale 1.0 — active only during "
|
"band layer, uniform scale 1.0 — active only during "
|
||||||
"band re-runs, k=0 stays bit-exact")
|
"band re-runs, k=0 stays bit-exact")
|
||||||
ap.add_argument("--lora-lr", type=float, default=1e-3)
|
ap.add_argument("--lora-lr", type=float, default=1e-3)
|
||||||
|
ap.add_argument("--lensteach", type=float, default=0.0, metavar="LAMBDA",
|
||||||
|
help="item 25: latent process supervision — lens-CE at the "
|
||||||
|
"replacement pauses against the DELETED step's tokens "
|
||||||
|
"(1:1 pause j <-> step token j), mixed at LAMBDA into "
|
||||||
|
"the output CE")
|
||||||
ARGS = ap.parse_args()
|
ARGS = ap.parse_args()
|
||||||
STEPS, LR = ARGS.steps, ARGS.lr
|
STEPS, LR = ARGS.steps, ARGS.lr
|
||||||
TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
|
TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
|
||||||
f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + (
|
f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + (
|
||||||
f"_blr{ARGS.bandlora}" if ARGS.bandlora else "") + (
|
f"_blr{ARGS.bandlora}" if ARGS.bandlora else "") + (
|
||||||
|
f"_lt{str(ARGS.lensteach).replace('.', '')}" if ARGS.lensteach else "") + (
|
||||||
f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
|
f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
|
||||||
f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix
|
f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix
|
||||||
|
|
||||||
@@ -70,14 +76,15 @@ TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
|
|||||||
def drop_cot_steps(cot, d):
|
def drop_cot_steps(cot, d):
|
||||||
"""Delete the first d scratchpad lines (front-first: the deleted
|
"""Delete the first d scratchpad lines (front-first: the deleted
|
||||||
computation must ride the pause-chain before the visible remainder).
|
computation must ride the pause-chain before the visible remainder).
|
||||||
Returns (new_cot, n_deleted); unparseable cots pass through intact."""
|
Returns (new_cot, n_deleted, deleted_text); unparseable cots pass
|
||||||
|
through intact."""
|
||||||
lines = [l for l in cot.split("\n") if l.strip()]
|
lines = [l for l in cot.split("\n") if l.strip()]
|
||||||
ans = [l for l in lines if l.startswith("Answer")]
|
ans = [l for l in lines if l.startswith("Answer")]
|
||||||
steps = [l for l in lines if not l.startswith("Answer")]
|
steps = [l for l in lines if not l.startswith("Answer")]
|
||||||
if len(ans) != 1 or not steps:
|
if len(ans) != 1 or not steps:
|
||||||
return cot, 0
|
return cot, 0, ""
|
||||||
n = min(d, len(steps))
|
n = min(d, len(steps))
|
||||||
return "\n".join(steps[n:] + ans), n
|
return "\n".join(steps[n:] + ans), n, "\n".join(steps[:n])
|
||||||
|
|
||||||
|
|
||||||
class LensNoiseWrapper(torch.nn.Module):
|
class LensNoiseWrapper(torch.nn.Module):
|
||||||
@@ -162,11 +169,11 @@ def main():
|
|||||||
it = star.get(r["idx"])
|
it = star.get(r["idx"])
|
||||||
if it is None:
|
if it is None:
|
||||||
continue
|
continue
|
||||||
cot, ndel = (drop_cot_steps(r["cot"], ARGS.drop_steps)
|
cot, ndel, deleted = (drop_cot_steps(r["cot"], ARGS.drop_steps)
|
||||||
if ARGS.drop_steps else (r["cot"], 0))
|
if ARGS.drop_steps else (r["cot"], 0, ""))
|
||||||
n_dropped += ndel
|
n_dropped += ndel
|
||||||
data.append({"question": it["question"], "label": r["label"],
|
data.append({"question": it["question"], "label": r["label"],
|
||||||
"cot": cot,
|
"cot": cot, "deleted": deleted,
|
||||||
"extra_pauses": ndel * ARGS.pause_per_step})
|
"extra_pauses": ndel * ARGS.pause_per_step})
|
||||||
if ARGS.drop_steps:
|
if ARGS.drop_steps:
|
||||||
print(f"rung B d={ARGS.drop_steps}: {n_dropped} steps deleted "
|
print(f"rung B d={ARGS.drop_steps}: {n_dropped} steps deleted "
|
||||||
@@ -206,6 +213,33 @@ def main():
|
|||||||
f"{sum(p.numel() for p in lora_params)/1e6:.1f}M params, "
|
f"{sum(p.numel() for p in lora_params)/1e6:.1f}M params, "
|
||||||
f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, "
|
f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, "
|
||||||
f"lr={ARGS.lora_lr}", flush=True)
|
f"lr={ARGS.lora_lr}", flush=True)
|
||||||
|
lens_teach = None
|
||||||
|
if ARGS.lensteach:
|
||||||
|
from loop_common import BAND
|
||||||
|
J30 = torch.load(Path(__file__).resolve().parent.parent
|
||||||
|
/ "results/jbar.pt",
|
||||||
|
map_location="cuda")["Jbar"][BAND[1]].float()
|
||||||
|
tm = looper.tm
|
||||||
|
softcap = model.config.get_text_config().final_logit_softcapping
|
||||||
|
|
||||||
|
def lens_teach(h):
|
||||||
|
proj = h.float() @ J30.T
|
||||||
|
x = tm.norm(proj.to(tm.norm.weight.dtype))
|
||||||
|
lg = model.lm_head(x)
|
||||||
|
if softcap:
|
||||||
|
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)
|
||||||
groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}]
|
groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}]
|
||||||
if lora_params:
|
if lora_params:
|
||||||
groups.append({"params": lora_params, "lr": ARGS.lora_lr,
|
groups.append({"params": lora_params, "lr": ARGS.lora_lr,
|
||||||
@@ -236,19 +270,42 @@ def main():
|
|||||||
ids, msk, lab, plens = build_batch(tok, batch, p)
|
ids, msk, lab, plens = build_batch(tok, batch, p)
|
||||||
for g in opt.param_groups:
|
for g in opt.param_groups:
|
||||||
g["lr"] = g["base"] * lr_at(step) / LR
|
g["lr"] = g["base"] * lr_at(step) / LR
|
||||||
logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL,
|
if ARGS.lensteach:
|
||||||
use_checkpoint=True,
|
logits, S = carry_logits(looper, adapter, ids, msk, plens,
|
||||||
feedforward=ARGS.feedforward)
|
K_PREFILL, use_checkpoint=True,
|
||||||
|
feedforward=ARGS.feedforward,
|
||||||
|
return_states=True)
|
||||||
|
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(),
|
loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
|
||||||
lab[:, 1:].flatten(), ignore_index=-100)
|
lab[:, 1:].flatten(), ignore_index=-100)
|
||||||
|
lce_val = 0.0
|
||||||
|
if ARGS.lensteach:
|
||||||
|
terms = []
|
||||||
|
for b, it in enumerate(batch):
|
||||||
|
tgt = it.get("lens_targets")
|
||||||
|
if not tgt:
|
||||||
|
continue
|
||||||
|
s0 = int(plens[b]) + p
|
||||||
|
hs = S[b, s0 : s0 + len(tgt)]
|
||||||
|
terms.append(F.cross_entropy(
|
||||||
|
lens_teach(hs).float(),
|
||||||
|
torch.tensor(tgt, device=hs.device)))
|
||||||
|
if terms:
|
||||||
|
lce = torch.stack(terms).mean()
|
||||||
|
lce_val = lce.item()
|
||||||
|
loss = loss + ARGS.lensteach * lce
|
||||||
opt.zero_grad(set_to_none=True)
|
opt.zero_grad(set_to_none=True)
|
||||||
loss.backward()
|
loss.backward()
|
||||||
torch.nn.utils.clip_grad_norm_(
|
torch.nn.utils.clip_grad_norm_(
|
||||||
list(adapter.parameters()) + lora_params, 1.0)
|
list(adapter.parameters()) + lora_params, 1.0)
|
||||||
opt.step()
|
opt.step()
|
||||||
log.append({"step": step, "loss": loss.item()})
|
log.append({"step": step, "loss": loss.item(), "lce": lce_val})
|
||||||
if step % 10 == 0:
|
if step % 10 == 0:
|
||||||
print(f"step {step:4d} {lbl:4s} loss={loss.item():.4f} "
|
print(f"step {step:4d} {lbl:4s} loss={loss.item():.4f} "
|
||||||
|
f"lce={lce_val:.3f} "
|
||||||
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
||||||
if step % 200 == 199 or step == STEPS - 1:
|
if step % 200 == 199 or step == STEPS - 1:
|
||||||
if ARGS.lensnoise:
|
if ARGS.lensnoise:
|
||||||
|
|||||||
Reference in New Issue
Block a user