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