anlp-a2

Model checkpoints and result files for Tasks 1 and 2 of Assignment 2 (Advanced NLP, Monsoon 2026). Every model here was trained from scratch with PyTorch on one RTX 5060 Ti (16 GB). Task 3 (decoding strategies) uses the pretrained EleutherAI/pythia-160m, so it has no checkpoints in this repository. The code is in the course repository on code.iiit.ac.in.

Layout

task1/                      Mixture-of-Experts translation models (Vietnamese/Japanese -> English)
  model.py                  model definition (needed to load the checkpoints)
  tokenizer.json            shared byte-level BPE tokenizer
  test_metrics.csv/.png     test perplexity and BLEU of the five variants
  val_loss.png              validation loss vs. training tokens
  v1/ ... v5/               one folder per FFN variant
task2/                      one dense model trained with four optimizers
  model.py, optimizers.py   model definition and the four optimizer implementations
  tokenizer.json            English byte-level BPE tokenizer
  final_metrics.json        final metrics of all runs
  training_dynamics.png     validation loss and test BLEU vs. fraction of the dataset seen
  adamw/ nadamw/ lion/ muon/

Every final.pt is {"config": <TransformerConfig fields>, "model": <state_dict>}. It holds the weights only, with no optimizer state.

import torch
from model import Transformer, TransformerConfig      # model.py from the task folder

ckpt = torch.load("v2/final.pt", map_location="cpu")  # e.g. task1/v2/final.pt
model = Transformer(TransformerConfig(**ckpt["config"]))
model.load_state_dict(ckpt["model"])
model.eval()
logits, aux_loss, expert_counts = model(input_ids, attention_mask, "cpu")   # input_ids, attention_mask: (B, N)

---

Task 1: Mixture of Experts

Decoder-only transformer (d_model 512, 8 heads, 6 layers, RMSNorm, sinusoidal positions, tied input/output embeddings, vocabulary 16,384). Attention, embeddings, data and training are identical across the five models; only the feed-forward block differs:

variant	FFN	params total / active per token
v1	dense 2-layer MLP (hidden 2048)	27.30M / 27.30M
v2	4 experts, top-1	27.32M / 17.86M
v3	4 experts, top-2	27.32M / 21.02M
v4	1 shared expert + 3 routed experts, top-1 routed	27.32M / 21.01M
v5	4 experts (twice the width), top-2, active parameters matched to v1	39.91M / 27.31M

Data and tokenizer. belumind/en-vi-ja-curated-500k-triplets: 446,252 training rows, each used twice (vi->en and ja->en), giving 892,504 samples; 24,792 validation and 24,792 test rows. Each sample is <BOS> source <SEP> English <EOS> (no language tags) and the loss covers only the English tokens. The tokenizer is a byte-level BPE (16,384 tokens, NFC normalisation) trained on the three language columns of the training split. Special ids: <PAD>=0, <BOS>=1, <EOS>=2, <SEP>=3.

Training. One epoch (13,946 steps, 64 sequences per step, about 39.1M tokens), AdamW (betas 0.9 / 0.95, weight decay 0.1), constant learning rate 3e-4, bf16 autocast, gradient clipping at 1.0, seed 0 (identical data order for every variant). The MoE variants add a Switch-style load-balancing loss with weight 0.01. Notes: v3 crashed with a CUDA out-of-memory error at step 4692 and was resumed from its step-4000 checkpoint. Batches above 12,288 padded tokens are split into chunks with gradient accumulation; this was used by v3 from step 4001 on and by v4 and v5, but not by v1 and v2 (the cross-entropy is unchanged by the split; only the load-balancing term is computed per chunk).

Results (test set, 24,792 sentences per direction; perplexity is exp of the mean cross-entropy on the English target tokens given the source; BLEU is sacreBLEU corpus BLEU of greedy decodes, "overall" pools both directions):

variant	final val loss	ppl vi->en	ppl ja->en	ppl overall	BLEU vi->en	BLEU ja->en	BLEU overall
v1	3.174	17.96	33.37	24.48	18.60	10.77	14.70
v2	2.866	13.59	24.52	18.26	22.53	13.51	18.00
v3	2.961	14.69	26.96	19.90	20.71	12.09	16.35
v4	2.928	14.27	25.70	19.15	21.52	12.32	16.84
v5	2.860	13.23	24.36	17.95	22.62	13.56	18.11

Files in each task1/vN/: final.pt, history.json (validation loss vs. tokens seen), test_metrics.json, samples.json (50 test translations per source language). For the MoE variants (v2 to v5) also expert_heatmap.png (expert usage per layer and source language, all tokens), expert_heatmap_source.png (source tokens only) and expert_counts.json / expert_counts_source.json (the raw routing counts behind the heatmaps).

---

Task 2: Optimizers

The dense architecture of Task 1 variant v1 (27.30M parameters, context 256, vocabulary 16,384), trained from scratch for next-token prediction on browndw/human-ai-parallel-corpus (66,320 texts). The texts are split 90 / 5 / 5 into train / validation / test by source document, so no version of a document appears in two splits (39,068,808 / 2,167,168 / 2,162,705 tokens). The English byte-level BPE tokenizer (16,384 tokens, <PAD>=0, <BOS>=1, <EOS>=2) is trained on the training split; documents are concatenated with <EOS> and cut into 256-token blocks.

Optimizers (optimizers.py, each subclasses torch.optim.Optimizer directly):

run	category	settings
adamw	baseline	betas (0.9, 0.95), eps 1e-8
nadamw	variance-reduced AdamW variants	Nesterov momentum, same settings as AdamW
lion	memory-efficient	betas (0.9, 0.99)
muon	matrix-based	momentum 0.95 (Nesterov), 5 Newton-Schulz steps, update scaled by sqrt(max(1, rows/cols)); used for the hidden 2-D weight matrices, while the embedding, norm gains and biases use an AdamW update

Training (identical for all four). One epoch (39.07M tokens), 16 sequences of 256 tokens per step, learning rate 3e-4 for every optimizer (Muon: 3e-4 for the matrices and 3e-4 for its AdamW path; no per-optimizer tuning), weight decay 0.1, linear warmup over the first 2% of steps and then linear decay to 10% of the peak, gradient clipping at 1.0, bf16 autocast, seed 0. Validation loss and test BLEU are recorded every 0.1x of the dataset.

Test BLEU is defined as follows: 500 fixed test documents; the model gets the first 128 tokens and greedily generates the next 64, which are scored with sacreBLEU against the true tokens 129 to 192.

Final results:

run	val loss	test loss	test BLEU
adamw	5.245	5.230	1.05
nadamw	5.084	5.068	1.17
lion	6.433	6.429	0.14
muon	4.675	4.653	1.95

Files in each task2/<run>/: final.pt, history.json (validation loss and test BLEU at each 0.1x of the epoch), final_metrics.json, bleu_examples.json (generated and reference continuations).