README.md

metadata

language:
  - en
  - vi
  - ja
tags:
  - nlp
  - mixture-of-experts
  - moe
  - transformer
  - optimizers
  - decoding-strategies
  - adamw
  - adan
  - lion
  - muon
  - sophia
license: mit
datasets:
  - belumind/en-vi-ja-curated-500k-triplets
  - browndw/human-ai-parallel-corpus
  - hamishivi/ROCStories
metrics:
  - bleu
  - perplexity
pipeline_tag: text-generation

Advanced Natural Language Processing (ANLP) - Assignment 2

Author: Chiranjeevi Kumaravadivel | Roll Number: 2023122010
Monsoon 2026 | International Institute of Information Technology, Hyderabad (IIIT-H)

---

🔗 Official Submission Deliverables

• Hugging Face Model Repository: https://huggingface.co/Chiru20/ANLP_assignment2
• Weights & Biases (WandB) Dashboard: https://forge.coreweave.com/wandb/chiranjeevi-kumaravadivel-iiit-hyderabad/ANLP-Assignment-2
• GitHub Source Code: https://github.com/Chiranjeevi-2020/ANLP_assignment2
• Official Technical Report (PDF): report.pdf (8 pages, 11pt, matching IIIT-H assignment formatting)

---

📖 Model & Task Overview

This repository hosts the official experimental deliverables for ANLP Assignment 2 (Monsoon 2026) at IIIT Hyderabad:

1. Mixture of Experts (MoE) Architecture (Part 1)

• Task: Multilingual translation from non-romanized Vietnamese (<vi>) and Japanese (<ja>) to English (<en>) using belumind/en-vi-ja-curated-500k-triplets.
• 5 Feed-Forward Network (FFN) Ablations:
  1. Dense Baseline: Single shared feed-forward layer ($D_{\text{ff}} = 1024$).
  2. Top-1 MoE: 4 experts ($d_e = 256$), 1 active per token.
  3. Top-2 MoE: 4 experts ($d_e = 256$), 2 active per token with weighted softmax combination.
  4. Shared + Routed MoE: 1 dedicated shared expert + 1 routed expert from 3 candidates.
  5. Matched-Active MoE: 4 wide experts ($d_e = 1024$), active parameters match Dense baseline exactly (525K active per layer).
• Load Balancing: Switch Transformers auxiliary loss ($\alpha = 0.01$): $$\mathcal{L}{\text{aux}} = \alpha \cdot N \sum{i=1}^N f_i \cdot P_i$$
• Language Specialization Heatmap: Preserved in assets/part1/expert_heatmap_matched_active.png.

FFN Ablation Variant	Total Parameters	Active Parameters	Final Validation Loss	Test BLEU
Dense Baseline	8.90M	8.90M	2.88	13.82
Top-1 MoE	8.91M (+0.12%)	7.86M (-11.7%)	2.96	13.15
Top-2 MoE	8.91M (+0.12%)	8.21M (-7.8%)	2.96	13.10
Shared + Routed	8.91M (+0.12%)	8.21M (-7.8%)	2.96	13.08
Matched-Active MoE	12.06M	8.91M (+0.08%)	2.94	13.30

---

2. Custom Optimizers from Scratch (Part 2)

• Task: Dense Transformer pretraining dynamics on browndw/human-ai-parallel-corpus replicating Figure 2 of Wen et al. (2025).
• Implemented Strictly from Scratch via torch.optim.Optimizer:
  • Category 1 (AdamW): Decoupled weight decay $\theta_{t} = \theta_{t-1}(1 - \eta \lambda) - \frac{\eta}{\sqrt{v_t} + \epsilon} m_t$.
  • Category 2 (Adan): Adaptive Nesterov momentum with gradient difference tracking $v_k = (1 - \beta_2) v_{k-1} + \beta_2 (g_k - g_{k-1})$.
  • Category 3 (Lion): Sign-based momentum updating $\theta_t = \theta_{t-1} - \eta \cdot \text{sign}(c_t)$, cutting second-moment memory in half.
  • Category 4 (Muon): Matrix polar decomposition via Quintic Newton-Schulz polynomial iterations ($a=3.4445, b=-4.7750, c=2.0315$).
  • Bonus (Sophia-G): Lightweight Hessian curvature estimation with coordinate-wise clipping $\text{clip}(m_t / \max(h_t, \epsilon), \rho)$.

Optimizer	Final Val Loss	Test BLEU	Memory Per Param	Relative Convergence Rate
CustomMuon	3.89	10.85	8 bytes (2D only)	Fastest (1.35x vs AdamW)
CustomAdamW	4.12	10.12	8 bytes	1.00x (Baseline)
CustomAdan	4.05	10.38	12 bytes	1.15x
CustomLion	4.18	9.85	4 bytes (50% savings)	0.95x
CustomSophia	4.08	10.25	8 bytes	1.10x

---

3. Autoregressive Decoding Strategies (Part 3)

• Task: 1,000 story completions on hamishivi/ROCStories evaluated on EleutherAI/Pythia-160M.
• Constraint: Implemented exclusively via model.forward() (zero calls to model.generate()).
• Methods:
  • Greedy Decoding: Deterministic $\text{argmax}{w} P(w \mid x{<t})$.
  • Top-k Sampling: Sampling from truncated top-$k$ vocabulary ($k=40, T=0.8$).
  • Top-p (Nucleus) Sampling: Sampling from dynamic nucleus mass ($p=0.9, T=0.8$).
  • Beam Search ($B \in {1, 2, 4}$): Length-normalized candidate beam tracking with strict $\mathcal{O}(B)$ latency scaling.

Strategy	Perplexity ($\downarrow$)	Token Accuracy (%)	SacreBLEU ($\uparrow$)	Latency / Sample
Greedy Search	18.42	42.15%	14.85	45.2 ms
Top-k ($k=40$)	24.15	36.80%	11.20	46.8 ms
Top-p ($p=0.9$)	22.80	38.10%	12.05	47.1 ms
Beam Search ($B=1$)	18.42	42.15%	14.85	45.8 ms
Beam Search ($B=2$)	16.95	44.30%	15.40	89.4 ms ($1.95\times$)
Beam Search ($B=4$)	15.80	45.90%	16.10	176.2 ms ($3.85\times$)

---

🚀 How to Run in Google Colab Pro

Open colab/ANLP_Assignment2_Colab_Runner.ipynb directly in Google Colab. The notebook clones the codebase, authenticates WandB, runs the mathematical unit tests, executes all three tasks, and saves artifacts automatically.

---

Submitted for evaluation in Monsoon 2026, IIIT Hyderabad.