ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 3)
969c687 verified | """(3) Overfit/memorization test: ATOM recipe on a ~2 KB text, repeated passes. | |
| usage: overfit_atom.py NAME {fresh|ckpt} MINUTES [LR] | |
| Mirrors tools/run_atom_native.py per-step logic: AdamW(wd=.01), clip 1.0, stabilize_dynamics, | |
| efference every 10, atom flush every 64, last_atom_dropout .5, field loss .05, reset per pass. | |
| Eval: teacher-forced CE on the text + greedy regeneration from the first 32 bytes | |
| (generation path with train-like commit labels, i.e. the decode bug neutralised, and original).""" | |
| import json, sys, time | |
| from pathlib import Path | |
| import torch, torch.nn.functional as F | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| from common import CKPT, ROOT, load_model | |
| from src.atom_native import AtomNativeModel | |
| from src.io.atomizer import Atomizer | |
| OUT = Path(__file__).resolve().parent | |
| TEXT = OUT / "memorize_2k.txt" | |
| def make_text(): | |
| raw = (ROOT / "data/corpus_fr_medium.txt").read_bytes()[:2300] | |
| cut = raw[:2048].rfind(b"\n\n") + 2 | |
| TEXT.write_bytes(raw[:cut]) | |
| def fresh(): | |
| 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=0.3) | |
| return m | |
| def flush(m): # same as trainer --atom-flush-every 64 (AtomSet is unbounded; aggregation is O(n^2)/tick) | |
| m.core.atoms = type(m.core.atoms)(); m.core.energy_history = [] | |
| def tf_ce(m, data): | |
| m.eval(); m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data, reset=True) | |
| ls = [] | |
| for i in range(len(pk) - 1): | |
| ls.append(float(m.transition_loss(pk[i], pk[i + 1])[1]["byte_loss"])) | |
| if (i + 1) % 64 == 0: flush(m) | |
| m.train(); return sum(ls) / len(ls) | |
| def regen(m, data, train_labels=True, n=None): | |
| m.eval(); orig = Atomizer._infer_commit_boundary | |
| if train_labels: | |
| Atomizer._infer_commit_boundary = classmethod(lambda cls, p: "newline" if p[-1] in (10, 13) else "max_span") | |
| try: | |
| m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data[:32], reset=True) | |
| for q in pk[:-1]: m.forward_packet(q) | |
| cur = pk[-1]; out = bytearray(data[:32]); n = n or len(data) | |
| while len(out) < n: | |
| if len(out) % 64 == 0: flush(m) | |
| b = int(m.forward_packet(cur)["surface"]["byte_logits"][0].argmax()); out.append(b) | |
| cur = m.atomizer.packet_from_payload(bytes([b])) | |
| finally: | |
| Atomizer._infer_commit_boundary = orig | |
| m.train(); same = next((i for i, (u, v) in enumerate(zip(out, data)) if u != v), len(data)) | |
| return same, bytes(out[:260]).decode("utf-8", "replace") | |
| def main(): | |
| name, init, minutes = sys.argv[1], sys.argv[2], float(sys.argv[3]) | |
| lr = float(sys.argv[4]) if len(sys.argv) > 4 else 3e-4 | |
| torch.set_num_threads(1); torch.manual_seed(20260913) | |
| if not TEXT.exists(): make_text() | |
| data = TEXT.read_bytes() | |
| m = fresh() if init == "fresh" else load_model(CKPT) | |
| m.surface.payload_enabled = False; m.surface.last_atom_dropout = 0.5; m.field_loss_weight = 0.05; m._efference_every = 10 | |
| m.train() | |
| opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=lr, weight_decay=0.01) | |
| log = []; t0 = time.time(); step = 0; ep = 0 | |
| r0 = {"pass": 0, "tf_ce": tf_ce(m, data)}; log.append(r0); print(json.dumps(r0), flush=True) | |
| while time.time() - t0 < minutes * 60: | |
| m.reset_state(reset_atomizer=True); pk = m.atomizer.encode_bytes(data, reset=True); ep += 1 | |
| for i in range(len(pk) - 1): | |
| step += 1; opt.zero_grad(set_to_none=True) | |
| if step % 10 == 0: loss, info = m.efference_loss(pk[i], pk[i + 1]) | |
| else: loss, info = m.transition_loss(pk[i], pk[i + 1]) | |
| loss.backward(); torch.nn.utils.clip_grad_norm_([p for p in m.parameters() if p.requires_grad], 1.0); opt.step() | |
| if step % 64 == 0: | |
| m.core.atoms = type(m.core.atoms)(); m.core.energy_history = [] | |
| m.stabilize_dynamics_parameters() | |
| if ep % 5 == 0: | |
| r = {"pass": ep, "step": step, "min": round((time.time() - t0) / 60, 1), "tf_ce": tf_ce(m, data)} | |
| log.append(r); print(json.dumps(r), flush=True) | |
| s1, h1 = regen(m, data, True); s2, h2 = regen(m, data, False) | |
| res = {"name": name, "init": init, "lr": lr, "text_bytes": len(data), "passes": ep, "steps": step, | |
| "final_tf_ce": tf_ce(m, data), "regen_prefix_match_train_labels": s1, "regen_head_train_labels": h1, | |
| "regen_prefix_match_orig_labels": s2, "regen_head_orig_labels": h2, "log": log} | |
| (OUT / f"overfit_atom_{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, indent=1)) | |
| if __name__ == "__main__": | |
| main() | |