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 kernels), trained to decrypt binary cipher sequences into English plaintext, with a controlled five-configuration ablation study.

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

Task

The dataset pairs 5000 binary cipher lines with plaintext lines. 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, zero mismatches). Models learn cipher → plaintext. Because the key is position-dependent, the task is alignment-heavy: the model must learn that output position i reads source position i, and that the substitution depends on i mod 8.

Configurations

Each of C2–C5 changes exactly one component from the C1 base. All configurations share depth (4+4 layers), width (d_model 256, 8 heads, FFN 1024), dropout 0.1, AdamW lr 3e-4 with 1000-step warmup + cosine, label smoothing 0.1, bf16, 80 epochs with early stopping, and identical data splits and seeds.

Config	Positional	Attention	Norm	Tokenization
C1-base	Sinusoidal absolute	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 bytes)

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

Config	Bit acc.	Seq. acc.	Levenshtein	BLEU	ROUGE-1	ROUGE-2	ROUGE-L
C1-base	0.720	0.004	54.50	67.12	0.818	0.702	0.817
C2-rope	0.704	0.004	83.75	61.98	0.779	0.652	0.776
C3-gqa	0.709	0.006	83.23	57.49	0.765	0.622	0.763
C4-rmsnorm	0.719	0.004	54.37	66.79	0.821	0.702	0.820
C5-blt	0.990	0.296	4.02	—	—	—	—

BLEU/ROUGE are reported for tokenized models only, per the assignment. Mean reference line length is ~598 characters, so Levenshtein ≈ 54 corresponds to roughly 91% of characters correct; sequence accuracy is strict exact match over the whole line.

Training efficiency (C5 vs C1)

	Params	s/epoch	Samples/s	Tokens/s	Peak GPU mem	Epochs to best
C1-base	9.49M	24.2	475	22.7k	3.15 GB	79
C5-blt	12.51M	95.1	121	25.5k	6.06 GB	15

Per step the token-free model is the more expensive one: ~3.9× fewer sentences per second at ~1.9× the peak memory (byte-level local attention dominates step time; patching amortizes the global transformer). Measured as wall-clock to the best checkpoint, however, C5 is cheaper — ~24 min versus ~32 min for C1 — and it converges to a far better model.

C5 outperforms the subword configurations by a wide margin on this task. The cipher is positional, and the token-free model gets identity alignment (output byte i ↔ source patch i, one patch per character), whereas subword models must learn a variable-length span mapping. This should not be generalized to natural-language tasks, where subword vocabularies break no comparable alignment.

Files

C1-base/best.pt     C2-rope/best.pt     C3-gqa/best.pt
C4-rmsnorm/best.pt  C5-blt/best.pt
tokenizer.json      # byte-level BPE, vocab 8000, trained on the train split
splits.json         # the exact 80/10/10 line indices (seed 42)
benchmarks.json     # throughput / memory / parameter counts
results_table.md    # final metric table

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

Training details

Hardware: one NVIDIA GB10 (DGX Spark, 121 GB unified memory), PyTorch 2.13 + CUDA 13. Lines are split into phase-aligned 256-character windows (a multiple of the 8-character key period) for training and decoding; decoded windows are concatenated back into full lines before metrics are computed, so all reported numbers are line-level.

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