ANLP Assignment 2, Part 2: one dense transformer, five optimizers
The same 8-layer decoder-only transformer (d_model 512, 37.8M parameters; the dense configuration of Part 1) pretrained from scratch for next-token prediction on browndw/human-ai-parallel-corpus (revision b514ff64) with optimizers implemented from scratch (Wen et al. 2025, Fantastic Pretraining Optimizers, Appendix A). Every run uses the same initial weights, the same batches (one epoch, 37.7M tokens) and the same learning-rate schedule; hyperparameters come from a sweep on a quarter of the data (sweep_best/).
| folder | optimizer | tuned hyperparameters | val loss (1×) | test BLEU | test chrF | WandB |
|---|---|---|---|---|---|---|
adamw |
AdamW | lr=0.0006, betas=[0.9, 0.95], weight_decay=0.3 | 4.2442 | 3.16 | 16.43 | run |
lion |
Lion | lr=0.00015, betas=[0.9, 0.9], weight_decay=0.7 | 4.3576 | 3.08 | 15.71 | run |
mars |
MARS | lr=0.0006, betas=[0.9, 0.95], gamma=0.025, weight_decay=0.1 | 4.2679 | 3.14 | 16.06 | run |
muon |
Muon | lr=0.032, momentum=0.9, betas=[0.9, 0.95], weight_decay=0.1, adam_lr=0.0006 | 3.7484 | 3.34 | 16.72 | run |
sophia |
Sophia-H | lr=0.00025, betas=[0.95, 0.9], gamma=0.00625, k=10, weight_decay=0.1 | 4.2666 | 3.18 | 16.38 | run |
Validation loss: next-token cross-entropy over the whole validation split (414 documents). Test BLEU: greedy continuation (128 tokens) of the last 256 tokens of each test document's first human chunk, sacrebleu corpus BLEU against 7 references (the human continuation and 6 LLM continuations, first 96 words each).
Files
tokenizer.json: byte-level BPE (24,576 tokens,<pad>=0,<eos>=1), trained on the train split.splits.json: the document-level train/validation/test split.<run>/model.pt:{"model": state_dict, "model_cfg": dict, "optimizer", "step", "tokens_seen"}(fp32).<run>/curve.json: validation loss (overall and per source) and test BLEU / chrF at every 0.1× of the data.<run>/generations/stepNNNNN.jsonl: the greedy test continuations at each of those points.<run>/run_info.json: full config, parameter count, device, versions.<run>/final.json: last point.results.csv,curves.csv,sweep_results.csv,sweep_best/: result tables and the sweep winners.
Loading
Uses the model code from the assignment repository (src/part1/model.py, src/part2/):
import torch
from src.part1.model import Transformer, TransformerConfig
from src.part2.data import load_tokenizer
ck = torch.load("muon/model.pt", map_location="cpu", weights_only=True)
model = Transformer(TransformerConfig.from_dict(ck["model_cfg"]))
model.load_state_dict(ck["model"])
model.eval()
tok = load_tokenizer("tokenizer.json")
Greedy continuation with the KV cache: src/part2/evaluate.py (greedy_continue).