From ee6b1f75c65a632064f1a975ea92b967d0feee01 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 18 Sep 2026 10:15:30 -0700 Subject: [PATCH 1/4] Allow alternate memorization protocols --- .../nanogpt_alpha_memorization/run_study.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/baseline/experiments/nanogpt_alpha_memorization/run_study.py b/baseline/experiments/nanogpt_alpha_memorization/run_study.py index fbde480a..010f15a4 100644 --- a/baseline/experiments/nanogpt_alpha_memorization/run_study.py +++ b/baseline/experiments/nanogpt_alpha_memorization/run_study.py @@ -71,7 +71,11 @@ def main(argv=None): parser.add_argument('command',choices=['run','plan','report','spectra','export','worker'],nargs='?',default='run') parser.add_argument('--root'); parser.add_argument('--device',choices=['mps','cpu','cuda'],default='mps'); parser.add_argument('--resume',action='store_true') parser.add_argument('--arm',choices=['adamw','muon']); parser.add_argument('--seed',type=int); parser.add_argument('--steps',type=int); parser.add_argument('--no-plots',action='store_true'); parser.add_argument('--force',action='store_true') - args=parser.parse_args(argv); cfg=json.loads((HERE/'protocol.json').read_text()) + parser.add_argument('--protocol',default='protocol.json',help='Protocol JSON in this experiment directory.') + args=parser.parse_args(argv) + protocol_path=(HERE/args.protocol).resolve() + if protocol_path.parent!=HERE or not protocol_path.is_file(): parser.error('Protocol must name a JSON file in this experiment directory.') + cfg=json.loads(protocol_path.read_text()) if args.steps is not None: cfg['steps']=args.steps if cfg['steps']<2: parser.error('At least two updates required.') try: @@ -102,7 +106,10 @@ def main(argv=None): with (root/'.queue.lock').open('a') as lock: fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB); LATEST.write_text(str(root)+'\n'); (root/'logs').mkdir(exist_ok=True) os.environ.setdefault('MPLCONFIGDIR',str(root/'cache'/'matplotlib')) - print(f'Results: {root}\n10 planned runs: AdamW x5 first, then Muon x5. Online WeightWatcher OFF.',flush=True) + plan=jobs(cfg) + counts={arm:sum(1 for a,_ in plan if a==arm) for arm in cfg['arms']} + description=', '.join(f'{arm} x{counts[arm]}' for arm in cfg['arms']) + print(f'Results: {root}\n{len(plan)} planned runs: {description}. Online WeightWatcher OFF.',flush=True) if not args.resume: preflight(cfg,root,args.device) statuses=[] for arm,seed in jobs(cfg): From 729c9369b3b5cc2b207e9c2dfdb80734e6a34d89 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 18 Sep 2026 10:15:40 -0700 Subject: [PATCH 2/4] Add four-head large-data memorization pilot --- .../protocol_four_head_large.json | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 baseline/experiments/nanogpt_alpha_memorization/protocol_four_head_large.json diff --git a/baseline/experiments/nanogpt_alpha_memorization/protocol_four_head_large.json b/baseline/experiments/nanogpt_alpha_memorization/protocol_four_head_large.json new file mode 100644 index 00000000..c1930bae --- /dev/null +++ b/baseline/experiments/nanogpt_alpha_memorization/protocol_four_head_large.json @@ -0,0 +1,35 @@ +{ + "version": 3, + "name": "memorization_four_head_large_data_pilot_v3", + "parent_commit": "63adbf337427966dd11fcccd4b705a7953928522", + "seeds": [1337], + "arms": ["adamw", "muon"], + "source_recipe": "baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml", + "source_blobs": { + "baseline/experiments/nanogpt_one_head_2026_08_21_baseline/configs/baseline.yaml": "7fd3c592afc7fdfea952b1f1aa0f9b44cdc53a2b", + "baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py": "029b42675ce230c6f1a142afaa8730b9f4d4ced3", + "baseline/nanogpt_one_head/src/rg_nanogpt_one_head/optimizers.py": "b1a97ad684afd882ffcff4ae0fe5d70140e7d22c" + }, + "model_overrides": { + "vocab_size": 512, + "n_head": 4 + }, + "data_seed": 20260918, + "steps": 30000, + "behavior_every": 1000, + "expensive_every": 5000, + "checkpoint_every": 1000, + "doses": [0, 1, 4, 16, 64], + "long_per_dose": 32, + "short_per_dose": 8, + "prefix_tokens": 64, + "suffix_tokens": 32, + "exposure_alphabet": [272,273,274,275,276,277,278,279,280,281,282,283,284,285,286,287], + "exposure_length": 3, + "prefix_grid": [8,16,32,64], + "modulus": 251, + "noise_fraction": 0.25, + "rule_audit_limit": 2048, + "online_weightwatcher": false, + "notes": "Four-head large-data pilot. One matched AdamW/Muon seed, 31,500 training rule examples from modulus 251, 15,750 validation and 15,751 test combinations, 30,000 updates (~30.5 effective epochs at batch 32). Same pinned optimizer profiles as the original study. Online WeightWatcher remains off; checkpoints are retained every 1,000 updates for post-hoc spectral analysis." +} From 28b20d8d7b61eaf3e4c9c65c62c2f251109a0674 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 18 Sep 2026 10:15:55 -0700 Subject: [PATCH 3/4] Document four-head large-data pilot --- .../nanogpt_alpha_memorization/README.md | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/baseline/experiments/nanogpt_alpha_memorization/README.md b/baseline/experiments/nanogpt_alpha_memorization/README.md index e4d71c62..0fa6d566 100644 --- a/baseline/experiments/nanogpt_alpha_memorization/README.md +++ b/baseline/experiments/nanogpt_alpha_memorization/README.md @@ -71,3 +71,24 @@ The protocol fingerprint must match. Completed seeds are not retrained. ## Scientific interpretation The study does **not** assume AdamW will produce alpha below 2 or that Muon will remain above 2. The behavioral study establishes what each optimizer memorizes and how generalization changes. The post-hoc WeightWatcher pass then tests whether those transitions coincide with a valid alpha-below-two regime. + + +## Four-head large-data pilot + +The original synthetic rule study has only 480 training combinations. For a more meaningful generalization test, `protocol_four_head_large.json` changes the model to four attention heads and expands the modular rule universe to modulus 251: 31,500 training combinations, 15,750 validation combinations, and 15,751 test combinations. The effective batch remains 32, so 30,000 updates are about 30.5 passes over the training set. + +The pilot runs one matched seed with ordinary AdamW followed by ordinary Muon. It keeps the same pinned optimizer profiles, 25% fixed random-label corruption, and canary memorization probes. Canary counts are increased to 32 long and 8 short examples per dose. Behavioral audits run every 1,000 updates; checkpoints are saved every 1,000; expensive exposure/compression audits run every 5,000. Online WeightWatcher remains disabled. + +Run the pilot with: + +```bash +caffeinate -dimsu python run_study.py run --protocol protocol_four_head_large.json +``` + +Preview it with: + +```bash +python run_study.py plan --protocol protocol_four_head_large.json +``` + +This is a pilot, not a five-seed statistical comparison. If both arms train sensibly and held-out accuracy improves, replicate the frozen protocol across the remaining seeds rather than changing hyperparameters after seeing the result. From f5f1c04bf8e47728b4157b8f09d5670ff05edb65 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 18 Sep 2026 10:16:03 -0700 Subject: [PATCH 4/4] Test four-head large-data protocol --- .../tests/test_four_head_large_protocol.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 baseline/experiments/nanogpt_alpha_memorization/tests/test_four_head_large_protocol.py diff --git a/baseline/experiments/nanogpt_alpha_memorization/tests/test_four_head_large_protocol.py b/baseline/experiments/nanogpt_alpha_memorization/tests/test_four_head_large_protocol.py new file mode 100644 index 00000000..df2c28e4 --- /dev/null +++ b/baseline/experiments/nanogpt_alpha_memorization/tests/test_four_head_large_protocol.py @@ -0,0 +1,23 @@ +import json +from pathlib import Path + +from am_data import Dataset +from am_runtime import baseline, make_model + +HERE = Path(__file__).resolve().parents[1] + + +def test_four_head_large_protocol_shape_and_data(): + cfg = json.loads((HERE / "protocol_four_head_large.json").read_text()) + source = baseline(cfg) + data = Dataset(cfg, 1337, source["training"]["batch_size"] * source["training"]["grad_accum_steps"]) + + assert cfg["model_overrides"]["n_head"] == 4 + assert len(data.train) == 31500 + assert sum(r.cohort == "validation_clean" for r in data.audit) == 15750 + assert sum(r.cohort == "test_clean" for r in data.audit) == 15751 + assert data.withdrawal == 15000 + + model = make_model(source, cfg, 1337, "cpu") + assert model.blocks[0].attn.n_head == 4 + assert model.cfg.n_embd % 4 == 0