Blame for experiments/train_forgetting.py

#!/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()