Hugging Face's logo

Spaces:
Ex0bit
/
mLattice
Running

App Files Community
mLattice / experiments /train_forgetting.py
Ex0bit's picture
Ex0bit
Docker Space: actual demo server, 3 checkpoints, readable dark text
2f06744 verified
Raw History Blame Contribute Delete
11.9 kB
#!/usr/bin/env python3
"""Forgetting training: after overwrite, queries must return the UPDATED value.
Protocol (CONTINUE.md step 2): sequences where 50% of keys appear twice (first
the original value, later an updated value). The loss is on the query positions
AFTER the updates; the delta rule must learn that the LAST value wins.
Warm-started from checkpoints/capacity8k.pt (the capacity curriculum already
mixes 25% overwrite batches, so this run specialises an existing skill rather
than bootstrapping one).
Curriculum: 64 -> 256 -> 1024 -> 4096 pairs, overwrite fraction cycling through
{0.3, 0.5, 0.7}, both dense (f=0) and 131K-protocol (f=14) batches. Final phase
trains the exact acceptance protocol (N=5,461, f=14, 50% overwrite) and quits
when the held-out probe (500 updated + 500 untouched queries) is >= 0.998 on
both query classes.
"""
from __future__ import annotations
import json
import os
import sys
import time
import torch
import torch.nn.functional as F
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, "..")
sys.path.insert(0, os.path.join(ROOT, "src"))
from mulattice.protocol import (SynthCfg, build_model, eval_batch, force_fast_scan, logits_at,
spread_loss, synth_batch)
DEVICE = os.environ.get("DEVICE", "cuda")
CFG = SynthCfg()
CKPT_WARM = os.environ.get("WARM", os.path.join(ROOT, "checkpoints", "capacity8k.pt"))
CKPT_OUT = os.path.join(ROOT, "checkpoints", "forgetting.pt")
HEARTBEAT = os.path.join(ROOT, "logs", "forgetting.heartbeat")
HISTORY = os.path.join(ROOT, "results", "forgetting_training.json")
LEVELS = [
dict(pairs=64, steps=1500),
dict(pairs=256, steps=2000),
dict(pairs=1024, steps=2500),
dict(pairs=4096, steps=4000),
]
FINAL = dict(pairs=5461, steps=6000, filler=14) # the acceptance protocol
OW_CYCLE = [0.3, 0.5, 0.7, 0.5]
TOKEN_BUDGET = 20000
if os.environ.get("SMOKE"): # fast local validation of the exact code path
LEVELS = [dict(pairs=8, steps=5), dict(pairs=16, steps=5)]
FINAL = dict(pairs=24, steps=6, filler=3)
def save(m, path, extra=None):
# record the LIVE mixer scale (the FORK ramp mutates it; writing the stale
# CFG value made scale-128 weights masquerade as scale 96 on reload)
live = float(m.blocks[0].mixer.cfg.addr_scale)
payload = {"model": m.state_dict(), "proto_cfg": {
"NK": CFG.NK, "NV": CFG.NV, "d_model": CFG.d_model, "n_heads": CFG.n_heads,
"d_k": CFG.d_k, "d_v": CFG.d_v, "addr_scale": live,
"retention_init": CFG.retention_init}, "cfg": "forgetting overwrite curriculum"}
if extra:
payload.update(extra)
torch.save(payload, path)
def held_out(m, pairs, filler, n_u=200, n_k=200, seed=777):
"""Acceptance-shaped probe: half queries on updated, half on untouched keys.
Query counts are clamped to the pool sizes (n_u <= n_upd, n_k <= n_pairs-n_upd)."""
g = torch.Generator().manual_seed(seed + pairs + filler)
n_upd = pairs // 2
perm = torch.randperm(pairs, generator=g)
upd_idx = perm[:n_upd]
upd_set = set(upd_idx.tolist())
updated_pool = [i for i in perm.tolist() if i in upd_set]
untouched_pool = [i for i in perm.tolist() if i not in upd_set]
qpos = torch.tensor(updated_pool[:min(n_u, len(updated_pool))]
+ untouched_pool[:min(n_k, len(untouched_pool))]).unsqueeze(0)
b = synth_batch(1, pairs, qpos.shape[1], CFG, filler=filler, overwrite_frac=0.5,
generator=g, qpos=qpos, updated_idx=upd_idx.unsqueeze(0))
return eval_batch(m, b, DEVICE)
def main():
dev = DEVICE
torch.manual_seed(0)
m = build_model(CFG).to(dev)
warm = CKPT_WARM if os.path.exists(CKPT_WARM) else os.path.join(
ROOT, "checkpoints", "capacity8k_latest.pt")
if os.path.exists(warm):
ck = torch.load(warm, map_location="cpu", weights_only=False)
m.load_state_dict(ck["model"])
# restore the TRAINED grid scale (addr_scale is config, not weights).
# Without this the forgetting lineage ran at scale 64 while its
# capacity ancestor was trained at 96. At N=5,461 + 50% overwrites the
# sequence stores 8,192 pair-blocks; at scale 64 this exceeds the
# density ceiling (see LIMITATIONS.md).
saved_scale = ck.get("proto_cfg", {}).get("addr_scale")
if saved_scale:
m.blocks[0].mixer.cfg.addr_scale = float(saved_scale)
CFG.addr_scale = float(saved_scale)
print(f"warm-start from {warm} (scale {m.blocks[0].mixer.cfg.addr_scale:.0f})",
flush=True)
else:
print("WARNING: no warm-start checkpoint at all, training from scratch", flush=True)
force_fast_scan(m)
opt = torch.optim.AdamW(m.parameters(), lr=1e-3)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(
opt, T_max=sum(l["steps"] for l in LEVELS) + FINAL["steps"], eta_min=1e-4)
g = torch.Generator().manual_seed(2)
step = 0
t0 = time.time()
history = []
def train_steps(pairs, n_steps, filler_cycle, probe_f=None):
nonlocal step
good = 0
for s in range(n_steps):
ow = OW_CYCLE[step % len(OW_CYCLE)]
f = filler_cycle[step % len(filler_cycle)]
Q = 8 if step % 4 else 64
T_est = int(pairs * (1 + ow) * (f + 2)) + 2 * Q
B = max(1, min(16, TOKEN_BUDGET // T_est))
b = synth_batch(B, pairs, Q, CFG, filler=f, overwrite_frac=ow, generator=g)
tok = b["tokens"].to(dev)
logits = logits_at(m, tok, b["ans_pos"].to(dev))
ce = F.cross_entropy(logits.reshape(-1, CFG.V), b["gold"].to(dev).reshape(-1))
kp = b["key_pos"]
if b["upd_key_pos"] is not None:
kp = torch.cat([kp, b["upd_key_pos"]], 1)
sl = spread_loss(m, tok, kp.to(dev), CFG, gate_subsample=1024)
loss = ce + sl
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
opt.step()
sched.step()
step += 1
if s % 100 == 0 or s == n_steps - 1:
sc = m.blocks[0].mixer.scan
print(f" p{pairs:5d} f{f:2d} ow{ow:.1f} step {step:6d} ce {float(ce.detach()):.3f} "
f"spread {float(sl.detach()):.3f} rounds {sc.last_rounds} "
f"{time.time()-t0:6.0f}s", flush=True)
with open(HEARTBEAT, "w") as fh:
fh.write(f"{step}\n")
if (s + 1) % 500 == 0 or s == n_steps - 1:
r = held_out(m, pairs, filler_cycle[-1])
history.append({"step": step, "pairs": pairs, "f": filler_cycle[-1],
"acc_updated": round(r["acc_updated"], 4),
"acc_untouched": round(r["acc_untouched"], 4)})
with open(HISTORY, "w") as fh:
json.dump(history, fh, indent=2)
print(f" EVAL p{pairs} upd {r['acc_updated']:.4f} keep {r['acc_untouched']:.4f}",
flush=True)
print(f"FORGETTING CURRICULUM warm={warm} device={dev}", flush=True)
START = int(os.environ.get("START_LEVEL", "0"))
ladder = [lv for lv in LEVELS if lv["pairs"] >= START]
for lv in ladder:
train_steps(lv["pairs"], lv["steps"], [0, 0, 14])
save(m, os.path.join(ROOT, "checkpoints", f"forgetting_lvl{lv['pairs']}.pt"), {"step": step})
# FINAL phase: a LADDER, not a jump. Measured: the overwrite skill scales
# down with pair count (1.0000 @64/256, ~0.98 @1024, ~0.95 @4096) exactly
# like capacity did before its ladder; and f=14 steps cost ~10s (T=131K) so
# a single-N final is both too shallow and too slow. Train at f in {6,8,14}
# (the mechanism is filler-agnostic; the acceptance distribution gets its
# own align steps), N ladder 2048->3548->5461, probe on the exact protocol.
# FORK envs (fork-and-flush, cf. arXiv:2610.07447): a parallel trajectory
# with DIFFERENT basin assumptions -- a bigger grid for double-writes
# (FORK_SCALE ramps during the final ladder) and long-gap erasure training
# (FORK_FILLERS raises the filler mix, stretching store->update gaps).
FS = float(os.environ.get("FORK_SCALE", "0")) or None
FF = os.environ.get("FORK_FILLERS", "6,8,14,8")
fcyc_fork = [int(x) for x in FF.split(",")]
print(f"FINAL phase: ladder 2048->3548->5461, filler mix {fcyc_fork}"
+ (f", scale ramp ->{FS:.0f}" if FS else ""), flush=True)
FO = os.environ.get("FINAL_ONLY", "")
if FO:
FINAL_LADDER = [(5461, int(FO))]
else:
FINAL_LADDER = [(2048, 1200), (3548, 1500), (5461, 3000)]
scale0_f = float(m.blocks[0].mixer.cfg.addr_scale)
LR0 = float(os.environ.get("FINAL_LR", "0")) # entry LR for the final
# ladder; without it the phase runs at the optimizer's construction LR
# (1e-3), which EXPLODES a converged model at B=1 (measured: probes
# 0.988 -> 0.754 in 200 steps -- Law 3 again). Warm up over 100 steps.
WARM = 100
import math as _m
for FL_N, FL_STEPS in FINAL_LADDER:
for s in range(FL_STEPS):
if LR0:
frac = min(1.0, (s + 1) / WARM)
cur = LR0 * (0.05 + 0.95 * frac)
tail = LR0 * (1.0 - 0.75 * min(1.0, s / max(1, FL_STEPS - 1)))
lr = min(cur, tail) if s >= WARM else cur
for pg in opt.param_groups:
pg["lr"] = lr
if FS:
frac = min(1.0, step / 2000.0)
m.blocks[0].mixer.cfg.addr_scale = scale0_f + (FS - scale0_f) * frac
ow = OW_CYCLE[s % len(OW_CYCLE)]
f = fcyc_fork[s % len(fcyc_fork)]
b = synth_batch(1, FL_N, 64, CFG, filler=f,
overwrite_frac=ow, generator=g)
tok = b["tokens"].to(dev)
logits = logits_at(m, tok, b["ans_pos"].to(dev))
ce = F.cross_entropy(logits.reshape(-1, CFG.V), b["gold"].to(dev).reshape(-1))
kp = torch.cat([b["key_pos"], b["upd_key_pos"]], 1)
sl = spread_loss(m, tok, kp.to(dev), CFG, gate_subsample=1024)
loss = ce + sl
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
opt.step()
sched.step()
step += 1
if s % 100 == 0 or s == FL_STEPS - 1:
print(f" final N{FL_N} f{f} {s:4d} ce {float(ce.detach()):.3f} T{b['L']} "
f"{time.time()-t0:6.0f}s", flush=True)
if (s + 1) % 200 == 0 or s == FL_STEPS - 1:
r = held_out(m, 5461, 14, n_u=500, n_k=500, seed=777001)
history.append({"step": step, "final": FINAL["pairs"] if False else FL_N,
"acc_updated": round(r["acc_updated"], 4),
"acc_untouched": round(r["acc_untouched"], 4)})
with open(HISTORY, "w") as fh:
json.dump(history, fh, indent=2)
print(f" FINAL PROBE upd {r['acc_updated']:.4f} keep {r['acc_untouched']:.4f}",
flush=True)
save(m, CKPT_OUT, {"step": step})
if r["acc_updated"] >= 0.998 and r["acc_untouched"] >= 0.998:
print(f"TARGET REACHED step {step}", flush=True)
break
else:
continue
break # TARGET REACHED inside the inner loop -> stop the ladder too
save(m, CKPT_OUT, {"step": step})
r = held_out(m, 64, 0, n_u=256, n_k=256, seed=777003)
print(f"SELF-CHECK dense N=64: upd {r['acc_updated']:.4f} keep {r['acc_untouched']:.4f}",
flush=True)
print(f"saved {CKPT_OUT} total {time.time()-t0:.0f}s", flush=True)
if __name__ == "__main__":
main()