dhyeythummar9/ANLP-A2-Part2-Optimizers

ANLP Assignment 2 — Part 2: Optimizer Comparison

This repository contains the final trained checkpoints for Part 2: Optimizers of Advanced NLP Assignment 2.

A decoder-only Transformer with a dense feed-forward network was pretrained for next-token prediction using the assignment dataset.

The experiment compares optimizer implementations written from scratch for the assignment.

Optimizers

The following optimizers were evaluated:

File	Optimizer
AdamW.pt	AdamW
NAdamW.pt	NAdamW
Lion.pt	Lion
Muon.pt	Muon
Sophia.pt	Sophia

The required optimizer implementations were written directly in PyTorch by inheriting from torch.optim.Optimizer, rather than using predefined optimizer implementations for the assignment comparison.

Sophia was included as the bonus Hessian-based optimizer.

Training setup

All final runs used the same model architecture, dataset split, tokenizer, token budget, initialization protocol, and evaluation schedule.

Final token budget:

39,080,642 tokens

Evaluation was performed at approximately 0.1× dataset-token intervals.

The final runs were evaluated on:

• validation loss
• validation perplexity
• test BLEU
• training loss
• optimizer-specific diagnostics

BLEU evaluation used 200 examples.

Final 1.0× results

Optimizer	Validation loss	Validation perplexity	Test BLEU	Train loss
AdamW	4.57476	97.0051	2.35506	4.42750
NAdamW	4.55052	94.6817	2.29319	4.40318
Lion	4.60710	100.1930	2.37691	4.45246
Muon	4.29563	73.3783	3.19020	4.15267
Sophia	4.46071	86.5486	2.49855	4.30476

At the final 1.0× token budget, Muon obtained the lowest validation loss and perplexity and the highest BLEU among the evaluated optimizer configurations.

Supporting files

This repository may also include:

• tokenizer.json
• experiment_config.json

These files describe the tokenizer and final experimental configuration used for the checkpoints.

Checkpoint format

The checkpoints were saved using PyTorch and are intended to be loaded by the Part 2 implementation in the course repository.

Example:

import torch

checkpoint = torch.load("AdamW.pt", map_location="cpu")