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:
Nils
Claude Fable 5
+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 re-runs, k=0 stays bit-exact")
|
||||
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()
|
||||
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"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
|
||||
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):
|
||||
"""Delete the first d scratchpad lines (front-first: the deleted
|
||||
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()]
|
||||
ans = [l for l in lines if l.startswith("Answer")]
|
||||
steps = [l for l in lines if not l.startswith("Answer")]
|
||||
if len(ans) != 1 or not steps:
|
||||
return cot, 0
|
||||
return cot, 0, ""
|
||||
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):
|
||||
@@ -162,11 +169,11 @@ def main():
|
||||
it = star.get(r["idx"])
|
||||
if it is None:
|
||||
continue
|
||||
cot, ndel = (drop_cot_steps(r["cot"], ARGS.drop_steps)
|
||||
if ARGS.drop_steps else (r["cot"], 0))
|
||||
cot, ndel, deleted = (drop_cot_steps(r["cot"], ARGS.drop_steps)
|
||||
if ARGS.drop_steps else (r["cot"], 0, ""))
|
||||
n_dropped += ndel
|
||||
data.append({"question": it["question"], "label": r["label"],
|
||||
"cot": cot,
|
||||
"cot": cot, "deleted": deleted,
|
||||
"extra_pauses": ndel * ARGS.pause_per_step})
|
||||
if ARGS.drop_steps:
|
||||
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"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, "
|
||||
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}]
|
||||
if lora_params:
|
||||
groups.append({"params": lora_params, "lr": ARGS.lora_lr,
|
||||
@@ -236,19 +270,42 @@ 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
|
||||
logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL,
|
||||
use_checkpoint=True,
|
||||
feedforward=ARGS.feedforward)
|
||||
if ARGS.lensteach:
|
||||
logits, S = carry_logits(looper, adapter, ids, msk, plens,
|
||||
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(),
|
||||
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)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(
|
||||
list(adapter.parameters()) + lora_params, 1.0)
|
||||
opt.step()
|
||||
log.append({"step": step, "loss": loss.item()})
|
||||
log.append({"step": step, "loss": loss.item(), "lce": lce_val})
|
||||
if step % 10 == 0:
|
||||
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)
|
||||
if step % 200 == 199 or step == STEPS - 1:
|
||||
if ARGS.lensnoise:
|
||||
|
||||
Reference in New Issue
Block a user