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.

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)

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

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