Audit and strengthen baseline training recipes

[charlesmartin14]

Scope

Reviewed every reference experiment under baseline/ against its intended role, the repository's current implementation, and the primary optimizer/model recipes. This PR turns that review into an explicit versioned audit and corrects the baseline code where the existing recipe was materially weak or inconsistent.

The goal is a strong, reproducible reference recipe. It is not a claim that literature defaults or a finite sweep establish a globally optimal hyperparameter point; a true optimum remains an empirical validation-only question on the target hardware.

MNIST / MLP3

• Extends the reference budget to 30 epochs.
• Replaces fixed learning rates with optimizer-specific linear warm-up plus cosine decay to a non-zero floor.
• Uses SGD + Nesterov at 0.05 -> 5e-4, momentum 0.90, matrix-only weight decay 1e-4.
• Uses AdamW at 1e-3 -> 1e-5, betas (0.90, 0.999), matrix-only weight decay 1e-2.
• Corrects the historical sgd_momentum_muon arm: hidden matrices use Muon, while the classifier and biases now use auxiliary AdamW rather than auxiliary SGD.
• Muon uses 0.02 -> 0.002, momentum 0.95, Nesterov, five Newton-Schulz steps and weight decay 0.01; auxiliary AdamW uses 3e-4 -> 3e-5, betas (0.90, 0.95).
• Adds gradient clipping and persists the learning-rate values in per-epoch metrics/checkpoints.
• Keeps the historical result-directory key so existing result stores remain discoverable.

CIFAR-10 / small ViT

The earlier ViT control was under-regularized and did not have a separate validation set. The corrected reference now uses:

• deterministic 45,000/5,000 optimization/validation split from the official CIFAR-10 training set;
• protected official test set, marked monitoring-only;
• the same 4x4-patch, width-192, 6-block, 3-head model;
• 300 epochs;
• dropout 0, stochastic depth 0.10;
• RandAugment, color jitter, random erasing 0.25;
• mixup 0.80, CutMix 1.00, label smoothing 0.10;
• gradient clipping at 1.0;
• validation-loss-selected best checkpoints;
• full restart checkpoints with model, optimizer, data-generator and RNG state plus a protocol fingerprint.

Optimizer schedules are now explicit and non-zero-floor:

• SGD + Nesterov: 0.10 -> 0.001, 5-epoch warm-up, momentum 0.90, weight decay 5e-4.
• AdamW: 1.25e-4 -> 1e-5, 5-epoch warm-up, betas (0.90, 0.999), weight decay 0.05. The peak is the DeiT 5e-4 reference linearly scaled from effective batch 512 to 128.
• Muon: hidden matrices 0.02 -> 0.002; auxiliary AdamW 3e-4 -> 3e-5; 5-epoch warm-up, momentum 0.95, five Newton-Schulz steps.

One-head nanoGPT / FineWeb-Edu

Audited and retained the existing implementation:

• pinned FineWeb-Edu sample-10BT revision;
• exact document-disjoint 10M/1M/1M GPT-2-BPE train/validation/test splits;
• nanoGPT AdamW recipe (6e-4, betas (0.9, 0.95), weight decay 0.1, warm-up/cosine schedule);
• reference hidden-matrix Muon + auxiliary AdamW partition;
• restartable checkpoints and validation-selected best state.

The audit records one caveat: the GPT-2 embedding is large relative to a one-block model, so any parameter-scaling analysis must report or exclude embedding parameters explicitly. The matched optimizer control itself remains valid because every arm shares that embedding and Muon acts only on hidden matrices.

nanochat d12

Audited and retained the pinned native upstream recipe at commit 92d63d4e8bb4df75c3b71618f31ddde2378b2bcd:

• depth 12, width 768, context 2048;
• upstream data/tokenizer preparation and packing;
• native initialization and Muon/AdamW parameter groups;
• upstream depth/batch scaling rules;
• 40-step warm-up, long warmdown, Muon momentum schedule and cautious weight-decay schedule.

Updating the pinned upstream commit is explicitly treated as a new baseline version.

WeightWatcher contract

All strict reference baselines preserve direct output from:

watcher.analyze(ERG=True, randomize=True, ...)

The code requires/stores direct layer alpha, ERG_gap, and randomized-MP num_traps values. It does not create fallback alpha values, proxy trap counts, or synthesized ERG gaps. The ViT analysis runs on CPU matrix copies so unsupported SVD/RMT operations stay off Apple MPS.

Documentation and notebooks

• Adds baseline/BASELINE_RECIPE_AUDIT.md with the full data, initialization, optimizer, schedule, selection-policy and source audit.
• Rewrites the baseline and notebook READMEs around the corrected protocols.
• Updates the three MNIST training notebooks and comparison notebook.
• Updates the ViT notebook to show train/validation/test trajectories, LR schedules, layer alpha, ERG gap and trap counts with three-seed 95% Student-t intervals.

Validation performed

• Python syntax compilation for the changed package files.
• Notebook JSON validation and code-cell parsing.
• Finite one-step smoke tests for all MLP optimizer paths.
• Exact one-step agreement between the custom auxiliary AdamW path and torch.optim.AdamW for auxiliary parameters.
• Finite one-step smoke tests for SGD, AdamW and Muon on the small ViT.
• ViT checkpoint save/load test restoring model, optimizer and data-generator state.
• Adds regression tests for schedules, parameter partitioning, recipe invariants and ViT restart behavior.

The full three-seed MNIST and nine-run 300-epoch CIFAR-10 campaigns were not executed in this environment. No performance or accuracy result is claimed or fabricated.