Surya4/anlp2-part2

ANLP Assignment 2 · Part 2 — four optimizers on one epoch of HAP-E

Decoder-only transformer (dense MLP FFN, ~20.5M params) trained for English next-token prediction / text continuation on the human-authored side of the HAP-E corpus, for exactly one epoch (1× the dataset = 29.5M tokens) — once with each of four optimizers from the assignment's reference paper (Appendix A): AdamW (baseline), NadamW (variance-reduced), Muon (matrix-based), Lion (memory-efficient). Identical stream, budget, steps, and schedule across runs; differences are attributable to the update rule + its tuned hyperparameters.

Models in this repo

run	folder	optimizer	lr / wd / betas	state	final val ppl	test ppl	final BLEU	best BLEU
part2-adamw	adamw/model.pt	AdamW (baseline)	lr 5e-4, wd 0.1, betas (0.9, 0.95), state 2.00×	97.04	96.37	0.814	0.905
part2-lion	lion/model.pt	Lion (memory-efficient, sign updates)	lr 1.5e-4, wd 0.6, betas (0.9, 0.95), state 1.00×	84.99	84.20	1.040	1.131
part2-muon	muon/model.pt	Muon (matrix-based)	NS lr 0.03 (momentum 0.95) + AdamW branch lr 3e-4, wd 0.1, betas (0.9, 0.95)	63.13	62.42	0.573	0.803
part2-nadamw	nadamw/model.pt	NadamW (variance-reduced AdamW)	lr 5e-4, wd 0.1, betas (0.9, 0.95), state 2.00×	91.93	91.27	1.069	1.069

Per-optimizer notes:

• AdamW (baseline) — The Table-1 baseline: decoupled weight decay + bias-corrected first/second moments.
• Lion (memory-efficient, sign updates) — Sign-of-interpolation updates with weight decay folded in — single momentum buffer (1.00× state), per the Lion paper.
• Muon (matrix-based) — Matrix-based: Newton–Schulz orthogonalization for 2-D matrices with an AdamW branch for vector parameters (assignment's reference paper); state ≈1.31×.
• NadamW (variance-reduced AdamW) — Variance-reduced variant of AdamW from the assignment's reference paper.

Shared architecture: d_model 384, 8 layers, 6 heads, d_ff 1536, tied embeddings, learned positional embeddings, dropout 0, context 512, vocab 16,000.

Full curves (tokens vs val ppl / test BLEU, 10 checkpoints + final)

tokens (×1e6)	AdamW	NadamW	Muon	Lion
2.96	337.2	450.7	351.1	329.8
5.91	238.6	248.5	215.8	231.5
8.85	197.1	191.7	156.9	189.4
11.81	164.9	150.9	126.9	157.1
14.77	141.0	125.5	106.1	132.6
17.71	123.9	109.6	91.2	116.8
20.67	112.4	99.2	80.5	106.1
23.61	104.9	92.3	72.6	99.1
26.57	100.1	87.8	66.8	94.8
29.52	97.0	85.0	63.1	91.9

tokens (×1e6)	AdamW	NadamW	Muon	Lion
2.96	0.541	0.454	0.651	0.514
5.91	0.499	0.539	0.421	0.487
8.85	0.700	0.717	0.696	0.871
11.81	0.676	0.620	0.790	0.657
14.77	0.905	1.131	0.627	1.046
17.71	0.522	0.884	0.700	0.673
20.67	0.748	1.007	0.803	0.747
23.61	0.612	0.981	0.710	0.845
26.57	0.820	0.939	0.741	0.951
29.52	0.814	1.040	0.573	1.069

Dataset

• browndw/human-ai-parallel-corpus (license: MIT) — HAP-E: human-authored English text (seed ~500-word chunks and their true continuations). Used for language modeling: doc-grouped 90/5/5 splits by document root; English side tokenized with a 16k byte-level BPE.
• byte-level BPE, vocab 16,000 (English-only), trained on the train split

Training details

• One epoch = 1× the dataset (29,511,235 tokens; 29,523,280 seen incl. spill batch)
• Steps: 1,866 (batch 32 × 512 tokens)
• Schedule: per-group warmup 5% + cosine decay to 10% lr; seed 42
• Precision: bf16 autocast (fp32 master params), grad clip 1.0
• Identical stream/budget/steps/schedule across all four optimizer runs

  run	tokens seen	steps	checkpoint val ppl
  part2-adamw	29,523,280	1866	97.04
  part2-lion	29,523,280	1866	84.99
  part2-muon	29,523,280	1866	63.13
  part2-nadamw	29,523,280	1866	91.93

Evaluation

Protocol:

• val/ppl: val split, every 2,951,123 tokens (10 checkpoints + final)
• test BLEU: 415 human chunk1→chunk2 pairs from the TEST documents, 128-token greedy continuations, single human reference
• Full per-checkpoint curves (tokens vs val ppl vs BLEU) in each model folder's eval_result.json

BLEU uses a single human reference (test split, 415 chunk1→chunk2 pairs, 128-token greedy continuations); val/ppl comes from the val split. Full 11-point data per run in each folder's eval_result.json.

Files

Root: tokenizer.json (+ configuration) shared by all runs. Per run: <optimizer>/model.pt (= copy of part2-<name>_best.pt, self-contained), config.json, eval_result.json, checkpoint_meta.json.

Usage

Custom decoder-only architecture (course repository src/part1/model.py, dense ffn_variant=1, n_vocab=16000): ckpt = torch.load('<optimizer>/model.pt', weights_only=False) and rebuild the model with ckpt['config']. tokenizer.json is a tokenizers BPE.

Limitations

• Tiny model (~20.5M params) at ~0.1× Chinchilla tokens — single pass, no repetition
• Continuation BLEU is a noisy, single-reference signal at this scale (~0.5–1.1); val/ppl is the primary metric
• ppl↔BLEU dissociate at this scale (best ppl does not imply best BLEU)

Provenance

Part 2 of Advanced NLP Assignment 2: four optimizers from the assignment's reference paper (Appendix A), each sub-classing torch.optim.Optimizer directly. Training runs part2-adamw, part2-lion, part2-muon, part2-nadamw in the WandB project suryamanojphy31-iiit-hyderabad/anlp-assignment2.