README.md

---
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)

[![Weights & Biases](https://img.shields.io/badge/Weights%20%26%20Biases-Dashboard%20Live-orange)](https://forge.coreweave.com/wandb/chiranjeevi-kumaravadivel-iiit-hyderabad/ANLP-Assignment-2)
[![GitHub](https://img.shields.io/badge/GitHub-Repository-blue)](https://github.com/Chiranjeevi-2020/ANLP_assignment2)
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Chiranjeevi-2020/ANLP_assignment2/blob/main/colab/ANLP_Assignment2_Colab_Runner.ipynb)

---

## 🔗 Official Submission Deliverables
- **Hugging Face Model Repository:** [https://huggingface.co/Chiru20/ANLP_assignment2](https://huggingface.co/Chiru20/ANLP_assignment2)
- **Weights & Biases (WandB) Dashboard:** [https://forge.coreweave.com/wandb/chiranjeevi-kumaravadivel-iiit-hyderabad/ANLP-Assignment-2](https://forge.coreweave.com/wandb/chiranjeevi-kumaravadivel-iiit-hyderabad/ANLP-Assignment-2)
- **GitHub Source Code:** [https://github.com/Chiranjeevi-2020/ANLP_assignment2](https://github.com/Chiranjeevi-2020/ANLP_assignment2)
- **Official Technical Report (PDF):** [`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`](./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.*