--- tags: - pytorch - mixture-of-experts - machine-translation - optimizers - language-modeling --- # 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": , "model": }`. It holds the weights only, with no optimizer state. ```python 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 ` source English ` (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: ``=0, ``=1, ``=2, ``=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, ``=0, ``=1, ``=2) is trained on the training split; documents are concatenated with `` 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//`:** `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).