| --- |
| 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": <TransformerConfig fields>, "model": <state_dict>}`. 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 `<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). |
|
|
