ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 3)
969c687 verified | """Long-memory bench (docs/MIX_DATA_LONG_MEMORY.md §3): fresh ATOM, same per-tick training logic as | |
| docs/artifacts/decfix_bptt/memorize_bptt.py (AdamW wd .01, clip 1, stabilize after each step, atom flush | |
| every 64, field loss .05, last-atom dropout .5, efference off, --field-bptt window rule), 1 thread. | |
| task recall : an endless stream of episodes "Nom: <NAME>.\n<filler of G bytes>\nQui ? <NAME>.\n\n" | |
| (16 names with distinct first letters, G uniform in [8, 160], filler = random slices of the | |
| mixed training corpus). The first name byte after "Qui ? " is only predictable from the | |
| name seen G+~15 bytes earlier: chance = ln 16 = 2.77 nats / 6.25 % acc. Eval on 256 fixed | |
| held-out episodes (other filler, other seed) bucketed by distance. | |
| task mem2k : the 2 KB memorisation test (root_cause/memorize_2k.txt, greedy regen past a 32-byte prefix). | |
| usage: longmem_bench.py NAME --task recall|mem2k --minutes M [--k K] [--slow-block F] [--slow-cap C] | |
| [--slow-gain ema|none] [--lr 2e-3] [--leak .05] | |
| """ | |
| import argparse, json, math, random, sys, time | |
| from pathlib import Path | |
| import torch | |
| HERE = Path(__file__).resolve().parent | |
| ROOT = HERE.parents[2] | |
| sys.path.insert(0, str(ROOT)); sys.path.insert(0, str(ROOT / "tools")); sys.path.insert(0, str(HERE.parent / "decfix_bptt")) | |
| from src.atom_native import AtomNativeModel # noqa: E402 | |
| from src.io.atomizer import Atomizer # noqa: E402 | |
| import memorize_bptt as MB # noqa: E402 (flush, tf_ce, regen) | |
| NAMES = ["Alice", "Bruno", "Cyril", "Denis", "Emile", "Fanny", "Gaston", "Hugo", | |
| "Irene", "Jules", "Karim", "Louise", "Marc", "Nadia", "Oscar", "Paul"] | |
| OUTDIR = ROOT / "checkpoints/mix_long_memory_bench" | |
| def filler_source(): | |
| raw = b"".join(p.read_bytes() for p in sorted((ROOT / "data").glob("corpus_fr_mix_[0-9]*.txt"))) | |
| return raw.decode("utf-8", "ignore").replace("\n", " ") | |
| def episode(rng, src, gmin=8, gmax=160, gap=None): | |
| name = rng.choice(NAMES); g = gap if gap is not None else rng.randint(gmin, gmax) | |
| o = rng.randrange(0, len(src) - 400); fill = src[o:o + g] | |
| head = f"Nom: {name}.\n"; q = "\nQui ? " | |
| text = head + fill + q + name + ".\n\n" | |
| ans = len((head + fill + q).encode()) # byte index of the first answer byte | |
| dist = ans - len("Nom: ".encode()) # bytes between the name's first byte and the answer | |
| return text.encode(), ans, dist, name | |
| def build(a): | |
| torch.manual_seed(a.seed) | |
| m = AtomNativeModel(d_model=32, n_modes=32, n_atoms_max=64, max_payload_bytes=1, | |
| atomizer=Atomizer(max_span_bytes=1), field_max_rms=3.0, enable_merge=False, | |
| field_loss_weight=0.05, field_obligatory_hard=True, last_atom_readout=True, | |
| field_leak=a.leak) | |
| m.field_leak = float(a.leak) | |
| m.set_field_leak_learnable(True, init=float(a.leak)) | |
| if a.slow_block > 0: | |
| m.set_field_leak_slow_block(a.slow_block, cap=a.slow_cap, gain=a.slow_gain, ref=a.leak) | |
| m.set_field_bptt(int(a.k)) | |
| m.surface.payload_enabled = False; m.surface.last_atom_dropout = a.ladrop; m.field_loss_weight = 0.05 | |
| m._efference_every = 0 | |
| return m | |
| def eval_recall(m, eps): | |
| m.eval(); m.reset_state(reset_atomizer=True) | |
| buckets = {}; tick = 0 | |
| for text, ans, dist, name in eps: | |
| pk = m.atomizer.encode_bytes(text, reset=False) | |
| for i in range(len(pk) - 1): | |
| out = m.forward_packet(pk[i]); tick += 1 | |
| if tick % 64 == 0: MB.flush(m) | |
| if i + 1 == ans: | |
| lp = torch.log_softmax(out["surface"]["byte_logits"][0], -1) | |
| tgt = text[ans] | |
| b = "d<32" if dist < 32 else "d32-63" if dist < 64 else "d64-127" if dist < 128 else "d128+" | |
| r = buckets.setdefault(b, {"n": 0, "ce": 0.0, "acc": 0}) | |
| r["n"] += 1; r["ce"] += float(-lp[tgt]); r["acc"] += int(int(lp.argmax()) == tgt) | |
| m.forward_packet(pk[-1]); tick += 1 | |
| m.train() | |
| return {b: {"n": r["n"], "ce": round(r["ce"] / r["n"], 4), "acc": round(r["acc"] / r["n"], 4)} for b, r in sorted(buckets.items())} | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("name"); ap.add_argument("--task", default="recall"); ap.add_argument("--minutes", type=float, default=45) | |
| ap.add_argument("--k", type=int, default=128); ap.add_argument("--leak", type=float, default=0.05) | |
| ap.add_argument("--slow-block", type=float, default=0.0); ap.add_argument("--slow-cap", type=float, default=0.02) | |
| ap.add_argument("--slow-gain", default="ema"); ap.add_argument("--lr", type=float, default=2e-3) | |
| ap.add_argument("--ladrop", type=float, default=0.5); ap.add_argument("--seed", type=int, default=20260913) | |
| ap.add_argument("--eval-min", type=float, default=10.0) | |
| a = ap.parse_args() | |
| torch.set_num_threads(1) | |
| m = build(a); m.train() | |
| params = [p for p in m.parameters() if p.requires_grad] | |
| opt = torch.optim.AdamW(params, lr=a.lr, weight_decay=0.01) | |
| t0 = time.time(); step = 0; updates = 0; capped = 0; log = []; nan = False | |
| if a.task == "recall": | |
| src = filler_source(); half = len(src) // 2 | |
| train_src, eval_src = src[:half], src[half:] | |
| erng = random.Random(777) | |
| eval_eps = [episode(erng, eval_src) for _ in range(256)] | |
| rng = random.Random(a.seed) | |
| stream = iter(lambda: episode(rng, train_src)[0], None) | |
| r0 = {"min": 0, "recall": eval_recall(m, eval_eps)}; log.append(r0); print(json.dumps(r0), flush=True) | |
| else: | |
| data = MB.TEXT.read_bytes() | |
| r0 = {"min": 0, "tf_ce": MB.tf_ce(m, data)}; log.append(r0); print(json.dumps(r0), flush=True) | |
| next_eval = a.eval_min * 60; acc = None; n = 0 | |
| m.reset_state(reset_atomizer=True) | |
| def tick(cur, tgt): | |
| nonlocal step, updates, capped, acc, n, nan | |
| step += 1 | |
| loss, info = m.transition_loss(cur, tgt) | |
| capped += info["field_scale"] < 0.999999 | |
| acc = loss if acc is None else acc + loss; n += 1 | |
| if n >= a.k: | |
| window_step() | |
| if step % 64 == 0: MB.flush(m) | |
| def window_step(): | |
| nonlocal updates, acc, n, nan | |
| if acc is None: return | |
| opt.zero_grad(set_to_none=True); (acc / n).backward() | |
| gn = torch.nn.utils.clip_grad_norm_(params, 1.0) | |
| if not math.isfinite(float(gn)): nan = True; raise FloatingPointError("nan grad") | |
| opt.step(); updates += 1; m.detach_field_graph(); acc = None; n = 0 | |
| m.stabilize_dynamics_parameters() | |
| def leak_stats(): | |
| lam = m.learnable_leak_rates().detach().reshape(-1) | |
| out = {"leak_min_mean_max": [round(float(lam.min()), 4), round(float(lam.mean()), 4), round(float(lam.max()), 4)]} | |
| if m.field_leak_slow_block is not None: | |
| out["slow_leak"] = [round(float(lam[i]), 4) for i in m.field_leak_slow_block["modes"]] | |
| return out | |
| try: | |
| while time.time() - t0 < a.minutes * 60: | |
| if a.task == "recall": | |
| chunk = b"".join(next(stream) for _ in range(8)) | |
| pk = m.atomizer.encode_bytes(chunk, reset=False) | |
| prev = getattr(m, "_bench_prev", None) | |
| if prev is not None: tick(prev, pk[0]) | |
| for i in range(len(pk) - 1): tick(pk[i], pk[i + 1]) | |
| m._bench_prev = pk[-1] | |
| else: | |
| m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data, reset=True) | |
| for i in range(len(pk) - 1): tick(pk[i], pk[i + 1]) | |
| window_step() | |
| if time.time() - t0 >= next_eval: | |
| window_step(); next_eval += a.eval_min * 60 | |
| r = {"min": round((time.time() - t0) / 60, 1), "step": step, "updates": updates, "cap_frac": round(capped / step, 4), **leak_stats()} | |
| if a.task == "recall": | |
| saved = m.core.state.alpha.clone(); r["recall"] = eval_recall(m, eval_eps) | |
| m.reset_state(reset_atomizer=True); m._bench_prev = None | |
| else: | |
| r["tf_ce"] = MB.tf_ce(m, data); r["regen"] = MB.regen(m, data)[0] - 32 | |
| log.append(r); print(json.dumps(r), flush=True) | |
| except FloatingPointError as exc: | |
| print("STOP:", exc, flush=True) | |
| window_step() | |
| res = {"name": a.name, "args": vars(a), "steps": step, "updates": updates, "minutes": round((time.time() - t0) / 60, 1), | |
| "field_cap_fraction": capped / max(step, 1), "nan": nan, **leak_stats(), | |
| "dynamics": {k: v for k, v in m.stabilize_dynamics_parameters().items() if k in ("energy_decay", "coupling_scale", "phase_sync")}} | |
| if a.task == "recall": | |
| res["final_recall"] = eval_recall(m, eval_eps) | |
| else: | |
| same, head = MB.regen(m, data) | |
| res.update(final_tf_ce=MB.tf_ce(m, data), regen_bytes_past_prefix=same - 32, regen_head=head[:300]) | |
| res["log"] = log | |
| OUTDIR.mkdir(parents=True, exist_ok=True) | |
| ck = OUTDIR / f"{a.name}.pt" | |
| if not ck.exists(): | |
| m.save(ck); res["checkpoint"] = str(ck.relative_to(ROOT)) | |
| (HERE / f"bench_{a.name}.json").write_text(json.dumps(res, indent=1, ensure_ascii=False)) | |
| print(json.dumps({k: v for k, v in res.items() if k != "log"}, ensure_ascii=False), flush=True) | |
| if __name__ == "__main__": | |
| main() | |