Arihant25/anlp-a1-transformer-ablations

ANLP Assignment 1 - Transformer Ablations & Byte Latent Transformer

Encoder-decoder Transformers built from scratch in PyTorch (no nn.Transformer, no nn.MultiheadAttention, no fused SDPA), trained to decrypt binary cipher sequences into English plaintext, with a controlled five-configuration ablation study. The BPE tokenizer is also implemented from scratch, with no tokenizer library.

Course: Advanced NLP, IIIT Hyderabad. Roll number 2023111026.

Task

5000 aligned line pairs. The cipher is a repeating-key XOR: int(cipher[8i:8i+8], 2) == ord(plain[i]) ^ KEY[i % 8] with KEY = b"ANLP2026", verified on all 5000 lines. Because the key is position-dependent, the task is alignment-heavy: the model must learn which source region each output position reads.

Configurations

Each of C2-C5 changes exactly one component from C1. All share d_model 256, 8 heads, 4+4 layers, FFN 1024, dropout 0.1, AdamW lr 3e-4 (1000 warmup, cosine), label smoothing 0.1, bf16, 200 epochs with early stopping, identical splits/seeds.

Config	Positional	Attention	Norm	Tokenization
C1-base	Sinusoidal	MHA	LayerNorm	BPE subword
C2-rope	RoPE	MHA	LayerNorm	BPE subword
C3-gqa	Sinusoidal	GQA (8Q/2KV)	LayerNorm	BPE subword
C4-rmsnorm	Sinusoidal	MHA	RMSNorm	BPE subword
C5-blt	Sinusoidal	MHA	LayerNorm	BLT (token-free)

C1-C4 use a from-scratch BPE on both sides: vocab 1024 over the raw 0/1 ciphertext, 8000 over the plaintext, each trained on the training split only.

C5 uses no vocabulary. Every 8 cipher bits are grouped into one byte value (0-255), and patch boundaries are placed by entropy, following the BLT paper: an order-2 byte n-gram model with backoff estimates H(next byte | context), and a new patch starts where the next byte is hard to predict. The threshold sits at the 75th percentile of the training entropy distribution (3.958 bits); 99.6% of boundaries come from entropy, and patches run 1-12 bytes averaging 4.0. Variable patches are pooled by cross-attention with a learned query.

Results (test split, greedy decoding, line-level)

Config	Bit acc.	Seq. acc.	Levenshtein	BLEU	ROUGE-1	ROUGE-2	ROUGE-L	Val loss
C1-base	0.712	0.004	75.91	61.74	0.786	0.652	0.785	1.238
C2-rope	0.780	0.046	19.30	84.31	0.920	0.859	0.920	0.497
C3-gqa	0.704	0.006	86.24	59.03	0.770	0.633	0.768	1.281
C4-rmsnorm	0.707	0.000	80.45	60.30	0.779	0.642	0.777	1.271
C5-blt	0.967	0.364	5.31	-	-	-	-	0.088

BLEU/ROUGE apply to tokenized models only. Mean test line is ~598 characters, so Levenshtein 5.31 is ~99.1% of characters correct.

Two findings worth noting. RoPE is the largest single-component win among the tokenized models, the opposite of what fixed-width tokenization showed earlier: with subword segmentation the source-to-target correspondence is a variable-length span relation, which relative offsets suit better than absolute positions. C5 dominates overall, reconstructing more than a third of test lines exactly.

Training efficiency

	Params	s/epoch	Samples/s	Tokens/s	Peak GPU mem	Best epoch
C1-base	9.68M	17.2	665.7	31.7k	3.00 GB	192
C5-blt	12.25M	65.9	174.3	36.8k	4.92 GB	80

C5 costs ~3.8x more per step but pushes more raw tokens per second, locating the cost in sequence length rather than model speed.

Files

C1-base/best.pt  C2-rope/best.pt  C3-gqa/best.pt  C4-rmsnorm/best.pt  C5-blt/best.pt
tokenizer.json           # plaintext BPE, vocab 8000 (from-scratch format)
src_tokenizer_1024.json  # ciphertext BPE, vocab 1024
entropy_model.json       # byte n-gram entropy model for BLT patching
splits.json              # exact 80/10/10 line indices (seed 42)
benchmarks.json          # throughput / memory / parameter counts

Each checkpoint is a dict with model, cfg, hp, epoch, val_loss. Loading requires the model classes from the assignment code.

Hardware: one NVIDIA GB10 (DGX Spark), PyTorch 2.13 + CUDA 13. Lines are split into phase-aligned 256-character windows for training and decoding; decoded windows are rejoined into full lines before metrics are computed.

Training logs: https://wandb.ai/arihanttr-iiit-hyderabad/anlp-a1