ANLP Assignment 2 - Part 2: from-scratch optimizers
The dense Part 1 architecture pre-trained for next-token prediction on
browndw/human-ai-parallel-corpus
for one full pass (1x the dataset), once per optimizer. Every optimizer is
implemented from scratch (only torch.optim.Optimizer is used, as a base
class), one per category of Table 1 of Fantastic Pretraining Optimizers and
Where to Find Them; Sophia-G is the bonus Hessian-based optimizer.
| Optimizer | Category | Folder | LR | Final val loss | Final test BLEU | Optimizer state |
|---|---|---|---|---|---|---|
| AdamW | 1. AdamW | adamw/ |
0.0012 | 4.0925 | 1.25 | 2.00x |
| Lion (memory-efficient) | 3. Memory-efficient | lion/ |
0.0002 | 4.2090 | 1.19 | 1.00x |
| MARS (variance-reduced) | 2. Variance-reduced AdamW variant | mars/ |
0.0012 | 4.0731 | 1.40 | 3.00x |
| Muon (matrix-based) | 4. Matrix-based | muon/ |
0.04 | 4.1127 | 1.40 | 1.37x |
| Sophia-G (Hessian-based) | 5. Hessian-based (bonus) | sophia/ |
0.0003 | 4.2017 | 1.24 | 2.00x |
Each folder holds final.pt (weights after 1x the dataset), summary.json
(hyperparameters and the validation-loss / test-BLEU curve at every 0.1x of
the dataset) and train_log.jsonl. lr_sweep/ has the learning-rate sweep
results that chose each LR. All runs share tokenizer.json (16k BPE) and part2.yaml.
Training curves: WandB (group part2-optimizers; the learning-rate sweep trials are in group part2-lr-sweep).
from src.part1.model import TransformerLM
model, extra = TransformerLM.load("muon/final.pt", map_location="cpu")