skimo2022/anlp-a1-transformers

ANLP Assignment 1 — Custom Transformers & Byte Latent Transformers

Encoder–decoder Transformers built from scratch in PyTorch (no nn.Transformer, no nn.MultiheadAttention) for the task of mapping encrypted binary sequences to plaintext, plus a controlled five-way architectural ablation.

Config	Change from base	Bit Acc (%)	Seq Acc (%)	Levenshtein	Params (M)	Peak GPU MB
C1		86.55	22.60	3.36	nan	358
C2		89.19	28.60	2.48	nan	358
C3		84.17	14.80	4.73	nan	346
C4		86.50	23.40	3.29	nan	337
C5		83.71	9.10	9.60	nan	533

Configurations

Config	Positional	Attention	Normalization	Tokenization
C1	Sinusoidal absolute	Multi-Head	LayerNorm	Subword (byte-level BPE)
C2	RoPE	Multi-Head	LayerNorm	Subword (byte-level BPE)
C3	Sinusoidal absolute	Grouped-Query	LayerNorm	Subword (byte-level BPE)
C4	Sinusoidal absolute	Multi-Head	RMSNorm	Subword (byte-level BPE)
C5	Sinusoidal absolute	Multi-Head	LayerNorm	BLT (token-free, entropy-patched bytes)

How C5 (BLT) works

• Bytes, not the text of bytes. The ciphertext is a binary sequence stored as ASCII 0/1 characters. Every 8 bits are packed into one byte value 0-255, so a 256-character cipher segment is 32 bytes, not 256.
• Entropy-based dynamic patching. Patches are variable length: a new patch opens where H(x_t | x_{t-2}, x_{t-1}) > theta. The entropy comes from a lightweight order-2 byte n-gram model fit on the training split only (entropy_model.json, also stored as buffers inside best.pt); no separate neural LM is trained. theta is calibrated to a target mean patch length, giving ~10 variable-length patches per 34-byte sequence.
• A standard learned 256-entry byte embedding represents the byte values.

Repository layout

C1/best.pt            checkpoint (state_dict + args + config)
C1/config.json        exact CLI arguments used
C1/metrics.json       test metrics, timings, peak memory
C1/tokenizer_*.json   the from-scratch byte-level BPE vocab + merges
...                   idem for C2..C5
C5/entropy_model.json the n-gram entropy model defining patch boundaries

Loading a checkpoint

import torch
from src.configs import get_config, build_model
import argparse

ck   = torch.load("C1/best.pt", map_location="cpu", weights_only=False)
cfg  = get_config(ck["config"]["name"])
args = argparse.Namespace(**ck["args"])
model = build_model(cfg, args, ck["src_vocab_size"], ck["tgt_vocab_size"])
model.load_state_dict(ck["model_state"])
model.eval()

Training/evaluation code: see the accompanying assignment submission (src/train.py, src/run_all.py).