parth12-ui/anlp-a2-optim-muon

ANLP Assignment 2, Part 2: muon

Dense decoder-only transformer (33.4M parameters, 8 layers, d_model 512, context 256) pretrained from scratch for next-token prediction on browndw/human-ai-parallel-corpus, for 1x the dataset (39,075,840 tokens), using a from-scratch muon optimizer (category: matrix-based).

Optimizer settings: lr=0.02, momentum=0.95, nesterov=True, ns_steps=5, weight_decay=0.01, adamw_lr=0.0006, adamw_betas=[0.9, 0.95], adamw_eps=1e-08.

metric at 1x dataset	value
validation loss	3.8926
validation perplexity	49.04
test BLEU (greedy 64-token continuation)	1.18

train_log.jsonl holds validation loss and test BLEU every 0.1x dataset tokens.

Loading

import json
from safetensors.torch import load_model
from model_src.config import TransformerConfig
from model_src.model import Transformer
model = Transformer(TransformerConfig.from_dict(json.load(open("config.json"))))
load_model(model, "model.safetensors")