README.md · Chiru20/ANLP_assignment2 at main

Upload ANLP Assignment 2 Model Card, Configs, Report, and Evaluation Deliverables

f90f02b 4 days ago

6.39 kB

---
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.*

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