Add four-head large-data memorization pilot

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.

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."
+}

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):

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