ANLP Assignment 2 · Part 1 — MoE FFN variants (v1–v5)
Decoder-only transformer trained for vi→en and ja→en translation across five parameter-matched feed-forward variants: one dense MLP (v1) and four MoE variants (v2–v5) — top-1, top-2, shared+top-1, and top-2-wide expert layouts. All five share the same embeddings/attention, data stream, budget, and schedule; differences are attributable to the FFN variant alone.
Models in this repo
| variant | folder | structure | params | val ppl | test ppl | BLEU | vi→en | ja→en | empty preds |
|---|---|---|---|---|---|---|---|---|---|
| v1 | v1-dense/model.pt |
dense MLP FFN, full width | 26.7M | 11.88 | 11.93 | 21.34 | 25.07 | 17.40 | 80 / 600 (13.3%) |
| v2 | v2-top1-quarter/model.pt |
4 experts × d_ff 384, top-1 per token; active FFN = ¼ dense | 26.7M | 15.05 | 15.14 | 21.00 | 25.10 | 16.76 | 2 / 600 (0.3%) |
| v3 | v3-top2-half/model.pt |
4 experts × d_ff 384, top-2 per token; active FFN = ½ dense | 26.7M | 12.38 | 12.45 | 23.45 | 27.75 | 19.01 | 3 / 600 (0.5%) |
| v4 | v4-shared-top1-half/model.pt |
3 routed experts × d_ff 384 + 1 shared expert (always active), top-1 routing | 26.7M | 12.55 | 12.59 | 23.47 | 27.06 | 19.75 | 5 / 600 (0.8%) |
| v5 | v5-top2-wide/model.pt |
4 experts × d_ff 768, top-2 per token; active FFN ≈ dense, 2× stored FFN | 36.1M | 11.91 | 11.96 | 22.06 | 27.40 | 16.63 | 1 / 600 (0.2%) |
Shared architecture: d_model 384, 8 layers, 6 heads, d_ff 1536 (dense reference), tied embeddings, learned positional embeddings, dropout 0, context 512. MoE variants: experts of d_ff None, routed per token by a learned gate (n_active = 1 or 2); v4 adds 1 always-active shared expert.
Dataset
- belumind/en-vi-ja-curated-500k-triplets (license: CC-BY-4.0) — EN-VI-JA translation triplets curated from OPUS parallel corpora (train 446,252 / val 24,792 / test 24,792 rows).
- byte-level BPE, vocab 32,000, special tokens:
Training details
- Optimizer: AdamW, lr 8e-4, betas (0.9, 0.98), weight decay 0.01, grad clip 1.0
- Precision: fp16 AMP (fp32 master params)
- Schedule: warmup 1.5M tokens + cosine decay to 10% (final lr 8.94e-5)
- Seed 42; identical shuffle stream, budget, and steps (23,382; ~30M tokens) for all five variants
- Val every 3M tokens (10 checkpoints); best checkpoint = final (monotone curves)
Evaluation
Protocol:
- 600 greedy translations (300 test rows × 2 languages) via a hand-rolled decoder (no model.generate)
- BLEU: sacrebleu corpus BLEU, empty predictions included
- test/ppl over the test split; val/ppl over the val split
BLEU = sacrebleu corpus BLEU over 600 greedy translations (300 test rows × 2 langs), empty predictions included; val/ppl at the best (== final) checkpoint. Results in the table above and, per model, in eval_result.json.
Files
Root: tokenizer.json (+ configuration) shared by all variants. Per model: <folder>/model.pt (= copy of part1-vN_best.pt, self-contained: state dict + config + tokens/step/val-ppl), config.json, eval_result.json, checkpoint_meta.json.
Usage
Custom decoder-only architecture (course repository src/part1/model.py, TransformerConfig + ffn_variant): ckpt = torch.load('<variant-folder>/model.pt', weights_only=False) and rebuild the model with ckpt['config']. tokenizer.json is a tokenizers BPE (Tokenizer.from_file or PreTrainedTokenizerFast(tokenizer_file=...)).
Limitations
- Small-scale:
39M (v1–v4) / 48M (v5) params, ~30M tokens (0.1× Chinchilla regime) - v1 (dense) shows an empty-prediction degeneracy (13.3% of greedy translations are empty)
- Per-variant ppl-vs-BLEU trade-offs documented in the report
Provenance
Part 1 of Advanced NLP Assignment 2. Training runs part1-v1, part1-v2, part1-v3, part1-v4, part1-v5 logged in the WandB project suryamanojphy31-iiit-hyderabad/anlp-assignment2.