Transformers from scratch — and one component at a time.

§0The assignment at a glance

The dataset is a parallel corpus of plaintext lines from the Brown corpus and their ciphertexts, each of which is a binary string that is 8× longer than the plaintext. The cipher is a repeating-key XOR (KEY = b"ANLP2026"): every plaintext character c at position i is written as the 8-bit big-endian encoding of ord(c) XOR KEY[i % 8]. The task is to learn a sequence-to-sequence model that maps the binary cipher back to the plaintext.

The work has two parts:

The viva is worth half the marks. For that reason, these notes focus on why each design was chosen, what would change under an alternative, and how the implementation connects to the course material.

StartFrom Python code to a learning system

Ignore the word “transformer” for a moment. Your program receives a list of encrypted strings and should write a list of plaintext strings. The model is a large, configurable Python function that turns integer arrays into scores. Training adjusts the numbers stored inside that function until the correct outputs receive high scores.

The six words the rest of the page uses

One example, before and after batching

Suppose the plaintext were "Hi". The supplied cipher turns each character into eight bits, so the paired source has 16 characters made only of 0 and 1. The model never receives either Python string directly:

# Conceptual example; the exact IDs depend on the trained BPE vocabulary.
plain  = "Hi"
cipher = "...16 bits..."

# TokenizedCipherDataset.__init__
src = src_tokenizer.encode(cipher)                    # e.g. [81, 407, 19]
tgt = [BOS] + tgt_tokenizer.encode(plain) + [EOS]    # e.g. [2, 94, 3]

# collate pads several examples into rectangular LongTensors.
src.shape == (B, Ls)
tgt.shape == (B, Lt)

B is batch size, Ls is the longest source length in that batch, and Lt is the longest target length. Shorter rows receive PAD IDs. C1–C4 train separate BPE tokenizers for the binary ciphertext and the plaintext; C5 instead groups every eight cipher bits into one integer byte and uses raw plaintext bytes.

The whole C1–C4 model as a data pipeline

src IDs (B, Ls)
  → source embedding lookup                 (B, Ls, 256)
  → encoder: source positions exchange info (B, Ls, 256)  # "memory"

previous target IDs (B, Lt)
  → target embedding lookup                 (B, Lt, 256)
  → decoder self-attention: read earlier output positions
  → decoder cross-attention: read encoder memory
  → linear layer                            (B, Lt, target_vocab_size)
                                               ↑ logits

The encoder reads the complete encrypted input and builds a useful vector for each source position. The decoder writes plaintext left to right. At each output position it can inspect earlier plaintext positions and the encoder's representation of the cipher. “Attention” is simply the operation that calculates how much information to take from each of those positions.

What training does

During training, the code already knows the correct plaintext. It gives the decoder every correct target token except the last and asks it to predict every correct target token except the first. This one-position shift is called teacher forcing:

# For tgt = [BOS, "H", "i", EOS]:
tgt[:, :-1]  # model input:  [BOS, "H", "i"]
tgt[:, 1:]   # correct next: ["H", "i", EOS]

logits = model(src, tgt[:, :-1])
loss = cross_entropy(logits, tgt[:, 1:])
optimizer.zero_grad()
loss.backward()       # fill each parameter's .grad
optimizer.step()      # slightly change parameters to lower this loss

At test time the correct plaintext is unavailable. Decoding starts with BOS, chooses the highest-scoring next token, appends it, and calls the decoder again until it emits EOS. That repeated loop is greedy autoregressive decoding.

§1The problem

Start with the data. Its structure determines what the model needs to learn and which design choices make sense.

Three consequences follow directly from that definition and shape almost everything downstream:

• The map is positional. The correct bit-flip pattern for a character depends only on i mod 8. So the model does not need to translate ideas across a sentence; it needs to align byte i of the target to bit range [8i, 8i+8) of the source and apply the right XOR mask. This is a copy-with-transformation task, not a translation task.
• Absolute position is meaningful modulo 8. Whatever positional encoding you pick, it needs to let the model learn phase-locked behaviour with period 8. Sinusoidal encodings already have periodic components, but the ones aligned to period 8 sit at specific dimensions of the encoding; RoPE gives the model directly usable relative-phase information.
• Alignment is long-range but regular. A line up to 2670 plaintext characters becomes a cipher string of up to 21360 bits, and quadratic attention over that is expensive. The two responses in the code are chunking (CHUNK_CHARS = 256) and, for C5, windowed local attention.

§2Task 1 · The scratch modules

Task 1 is the modules. They live in src/models/: attention.py holds scaled dot-product attention, MHA, and GQA; positional.py holds sinusoidal and RoPE; norm.py holds LayerNorm and RMSNorm; __init__.py assembles the FFN, encoder/decoder layers, and the Seq2SeqTransformer; and blt.py holds the byte-level pieces for C5. This section walks through each in the order the lectures cover them.

§2.1 · Scaled dot-product attention

This function lets every position build a weighted average of other positions. q, k, and v are not three different input sentences. They are three learned views of the same vectors: query means “what am I looking for?”, key means “what do I contain?”, and value means “what information should I send if I am selected?” The surrounding MHA class creates those views with three nn.Linear layers.

class ScaledDotProductAttention(nn.Module):
    def forward(self, q, k, v, mask=None):
        d_k = q.size(-1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min)
        attn = torch.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        return torch.matmul(attn, v)

Why divide by sqrt(d_k)? A dot product adds d_k products. As d_k grows, the scores grow in magnitude, softmax becomes almost all zeroes and one one, and its gradients become unhelpfully small. The division keeps score sizes roughly stable.

Why this mask implementation? The mask fill value is torch.finfo(scores.dtype).min (the dtype's most negative finite value), not -inf. If every position in a row is blocked (a fully-padded key row, which does happen for empty BLT patches downstream), softmax over a row of -inf produces NaN; with the dtype minimum it degrades to a uniform distribution, which is harmless because those rows are also going to be masked out of the loss.

The mask convention throughout the codebase is bool, True = BLOCKED, broadcast against (B, H, Lq, Lk). Key-padding masks arrive as (B, 1, 1, Lk); causal masks arrive as (1, 1, Lq, Lk). Broadcasting does the rest, so you can OR the two together and pass the result in.

§2.2 · Multi-Head Attention

MultiHeadAttention is the adapter around the function above. It accepts ordinary model vectors of shape (B, L, 256), constructs Q/K/V, runs eight smaller attentions, and returns the same shape it received. Keeping the input and output width at 256 is what allows attention to be dropped into a residual expression such as x + self_attn(x, x, x).

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1, rope=None):
        ...
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.o_proj = nn.Linear(d_model, d_model)
        self.attention = ScaledDotProductAttention(dropout)
        self.rope = rope

    def forward(self, query, key, value, mask=None):
        q = self._split(self.q_proj(query), self.n_heads)   # (B, H, Lq, d_k)
        k = self._split(self.k_proj(key),   self.n_heads)
        v = self._split(self.v_proj(value), self.n_heads)
        if self.rope is not None:
            q, k = self.rope(q, k)
        out = self.attention(q, k, v, mask)
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, Lq, -1))

With this assignment's values, each projection maps 256 floats to 256 floats. _split then reinterprets those 256 as 8 heads × 32 features; it does not create eight copies. Moving the head axis forward gives (B, 8, L, 32), and PyTorch's batched matrix multiplication runs all heads in parallel. After attention, transpose(...).contiguous().view(...) reverses that rearrangement. contiguous() is needed because transpose changes stride metadata and view requires contiguous storage. Finally, o_proj lets information from the eight heads mix instead of leaving them as unrelated 32-float blocks.

The rope hook is what lets the same class serve C1 (rope=None, sinusoidal PE added at the embeddings) and C2 (rope=RotaryPositionalEmbedding(...), no absolute PE at all, rotation applied to Q and K inside every self-attention layer). For cross-attention the caller passes rope=None, because the encoder and decoder positions live on different sequences that share no coordinate system (see the DecoderLayer assembly below).

§2.3 · Grouped-Query Attention (C3)

GQA (Ainslie et al. 2023) is a compromise between MHA and MQA (Multi-Query Attention). MQA has all heads share a single K/V, which makes the KV cache h× smaller but empirically loses quality. GQA groups the query heads into g groups and gives each group its own K/V pair, so K/V memory shrinks by h/g× while the query heads still see distinct-enough K/V to keep quality high.

Translate that to this code: MHA stores eight Q heads, eight K heads, and eight V heads. C3 still stores eight Q heads, but only two K and two V heads. Query heads 0–3 share K/V head 0; query heads 4–7 share K/V head 1. The constructor requires n_heads % n_kv_heads == 0 so those groups are equal.

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, dropout=0.1, rope=None):
        ...
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim)   # smaller!
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim)
        self.o_proj = nn.Linear(d_model, d_model)

    def forward(self, query, key, value, mask=None):
        q = self.q_proj(query).view(B, Lq, self.n_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(key)  .view(B, Lk, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(value).view(B, Lk, self.n_kv_heads, self.head_dim).transpose(1, 2)
        if self.rope is not None:
            q, k = self.rope(q, k)
        repeat = self.n_heads // self.n_kv_heads
        k = k.repeat_interleave(repeat, dim=1)   # broadcast K/V up to n_heads
        v = v.repeat_interleave(repeat, dim=1)
        out = self.attention(q, k, v, mask)
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, Lq, -1))

§2.4 · Sinusoidal positional encoding (C1, C3, C4, C5)

The classic Vaswani et al. encoding. It is added to token embeddings once, before the first layer, and never appears again. This is necessary because an embedding lookup only knows a token ID. The same token at positions 2 and 20 initially gets the exact same vector, and attention by itself has no notion of list order. Adding a different known vector at every index makes those occurrences distinguishable.

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=4096, dropout=0.1):
        super().__init__()
        position = torch.arange(max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float32)
            * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe, persistent=False)

Different columns oscillate at different speeds, giving every position a distinctive 256-number pattern while also exposing short and long distances. The table is fixed rather than learned because the ablation specifically asks for sinusoidal encoding and fixed values can work at positions not seen often during training. register_buffer makes pe travel with the module (moves to GPU on .cuda(), saves to state dict) without being registered as a parameter, since sinusoidal PE has no learnable weights. persistent=False keeps it out of the checkpoint because we can recompute it from d_model and max_len alone.

§2.5 · Rotary Position Embeddings (C2)

RoPE (Su et al. 2021, RoFormer) replaces the additive absolute encoding entirely. Instead of adding a position vector at the input, it rotates the Q and K vectors inside every attention layer by an angle that depends on their position. The head dimension is split in half; each half is rotated by m · θ_i with θ_i = base^−2i/d.

In code terms, C2 sets self.pos_enc = None, constructs one RotaryPositionalEmbedding(head_dim=32), and passes it into each self-attention module. Therefore the token vectors are not changed at embedding time. Instead, immediately after Q and K are split into heads, their pairs of features are rotated using the row of cos/sin for that sequence position. Values are not rotated because position should affect which value is selected, not alter the information carried by that value.

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, head_dim, max_len=4096, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
        t = torch.arange(max_len, dtype=torch.float32)
        freqs = torch.outer(t, inv_freq)                 # (max_len, head_dim/2)
        self.register_buffer("cos", freqs.cos(), persistent=False)
        self.register_buffer("sin", freqs.sin(), persistent=False)

    @staticmethod
    def _rotate(x, cos, sin):
        half = x.size(-1) // 2
        x1, x2 = x[..., :half], x[..., half:]
        return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)

    def forward(self, q, k):
        L = q.size(-2)
        cos = self.cos[:L].to(dtype=q.dtype)
        sin = self.sin[:L].to(dtype=q.dtype)
        return self._rotate(q, cos, sin), self._rotate(k, cos, sin)

The important part is what happens to Q^TK after both are rotated. Because rotation preserves inner products of matched components, the dot product between a rotated query at position m and a rotated key at position n depends only on m − n. That means relative offsets appear directly in the attention scores, with no additional bias term, no additional learnable parameters, and no need to add anything to the embeddings.

Two engineering notes. First, RoPE only applies to self-attention. DecoderLayer passes rope=None to its cross attention because encoder and decoder positions live on different sequences that share no coordinate system for rotation. Second, head_dim must be even for the half-split; the assertion catches this at construction.

§2.6 · LayerNorm and RMSNorm (C1-C3, C5 · vs · C4)

Both are per-token normalisations over the last dimension. LayerNorm subtracts the mean and divides by the standard deviation; RMSNorm (Zhang & Sennrich 2019) drops the mean-centering and the bias entirely.

A model vector can become very large or very small after many residual additions. Normalization rescales each token's 256-float vector before another sublayer processes it. The code uses dim=-1, so statistics are computed independently for every (batch, position) pair. It does not mix examples or sequence positions, and the shape remains (B, L, 256).

class LayerNorm(nn.Module):
    def forward(self, x):
        dtype = x.dtype
        x = x.float()                                       # stats in fp32
        mean = x.mean(dim=-1, keepdim=True)
        var  = x.var(dim=-1, keepdim=True, unbiased=False)
        y = (x - mean) * torch.rsqrt(var + self.eps)
        return (y * self.gamma.float() + self.beta.float()).to(dtype)

class RMSNorm(nn.Module):
    def forward(self, x):
        dtype = x.dtype
        x = x.float()
        y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (y * self.gamma.float()).to(dtype)

§2.7 · Position-wise feed-forward network

The other sub-block of every transformer layer, and the one that does most of the actual non-linear computation. It is applied position-independently: the same MLP runs at every sequence position with no cross-position mixing (that's attention's job).

class FeedForward(nn.Module):
    def __init__(self, d_model, ffn_dim, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, ffn_dim)      # d_model -> 4·d_model
        self.linear2 = nn.Linear(ffn_dim, d_model)      # 4·d_model -> d_model
        self.activation = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(self.activation(self.linear1(x))))

PyTorch applies nn.Linear to the last dimension, so this code needs no loop over B or L. For every token, linear1 expands 256 features to 1024, GELU makes the mapping non-linear, and linear2 compresses it back to 256 so the result can be added to x. Without GELU, the two linear layers would collapse algebraically into one linear layer and add little expressive power. Dropout randomly removes some hidden activations only during training to discourage dependence on one path.

With d_model=256, ffn_dim=1024 the FFN's parameter share dominates the model. Position-wise MLPs are typically 60%+ of a transformer's parameters, as the L5 param-count table shows. GELU is the modern default; the original Transformer paper used ReLU.

§2.8 · Encoder/decoder assembly (Pre-LN)

This is where the small modules become the model. Read an encoder layer as two updates to the same variable: first let positions share information, then let each position transform what it learned.

# EncoderLayer.forward; x stays (B, Ls, 256) throughout.
h = self.norm1(x)
x = x + self.dropout(self.self_attn(h, h, h, mask))
x = x + self.dropout(self.ffn(self.norm2(x)))
return x

The Vaswani et al. paper puts LayerNorm after the residual add: x + Sublayer(x) → LayerNorm (Post-LN). Every modern codebase (this one included) uses Pre-LN instead: the norm sits inside the residual branch, on the input to the sublayer.

class EncoderLayer(nn.Module):
    def forward(self, x, mask=None):
        h = self.norm1(x)
        x = x + self.dropout(self.self_attn(h, h, h, mask))
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

class DecoderLayer(nn.Module):
    def forward(self, x, memory, self_mask=None, cross_mask=None):
        h = self.norm1(x)
        x = x + self.dropout(self.self_attn(h, h, h, self_mask))
        x = x + self.dropout(self.cross_attn(self.norm2(x), memory, memory, cross_mask))
        x = x + self.dropout(self.ffn(self.norm3(x)))
        return x

The decoder has one extra update. Its self-attention reads the target prefix, cross-attention reads memory from the encoder, and the FFN transforms each resulting position. In cross_attn(norm2(x), memory, memory, src_mask), decoder vectors are Q while encoder vectors are K and V. That argument choice is the entire mechanical difference between self- and cross-attention.

Pre-LN is easier to train from scratch. With Post-LN the residual branch adds unbounded activations that immediately get normalised, so the residual is quickly wiped out and gradients flow poorly without a warm-up schedule. Pre-LN keeps a clean residual highway through the network (the norm only sees the residual-branch input), so gradient scale stays sane at deep depths.

Two more assembly details. Each stack ends with a final norm (self.final_norm) because Pre-LN's residual highway leaves the top-layer output unnormalised, and one final norm fixes that. Cross-attention is always MHA without RoPE regardless of the configuration; the ablated components (attention type in C3, RoPE in C2) concern self-attention only, and cross-attention query/key positions belong to different sequences.

§2.9 · Seq2SeqTransformer.forward, end to end

def _embed(self, ids, table):
    x = table(ids) * math.sqrt(self.d_model)       # IDs → 256 floats each
    return self.pos_enc(x) if self.pos_enc else self.embed_dropout(x)

def encode(self, src):
    src_mask = padding_mask(src, self.src_pad_id)  # do not read PAD columns
    memory = self.encoder(self._embed(src, self.src_embed), src_mask)
    return memory, src_mask

def decode(self, tgt_in, memory, src_mask):
    self_mask = causal_mask(tgt_in.size(1), tgt_in.device) \
              | padding_mask(tgt_in, self.tgt_pad_id)
    h = self.decoder(self._embed(tgt_in, self.tgt_embed),
                     memory, self_mask, src_mask)
    return self.generator(h)                       # 256 → target vocabulary

There are separate source and target embedding tables because a ciphertext BPE ID and a plaintext BPE ID belong to different vocabularies. Multiplication by sqrt(256) = 16 puts the deliberately small initialized embeddings at roughly unit scale, so the added sinusoidal values are neither negligible nor overwhelming. The encoder is called once per source. During greedy decoding, its returned memory is reused while the target prefix grows.

The two decoder masks solve different bugs. The causal upper triangle blocks target position i from reading correct future target tokens during teacher forcing. The padding mask blocks fake PAD positions added only to make a rectangular batch. Boolean OR combines both rules through broadcasting.

§3Task 2 · The five configurations

Every configuration shares the same hyperparameters (d_model=256, n_heads=8, enc_layers=dec_layers=4, ffn_dim=1024, dropout=0.1, lr=3e-4, batch token-budget=8192), so any difference in results is attributable to the one component that changed. That is the whole point of a controlled ablation.

C1 through C4 all share the same Seq2SeqTransformer assembly (src/models/__init__.py); they differ only in the flags passed to pos_encoding, attn_type, n_kv_heads, and norm. C5 is a separate BLTModel class because token-free processing changes the input pipeline shape enough to make a shared class awkward. But its global transformer stack is literally the same Encoder and Decoder classes as C1, only operating on patch representations instead of token embeddings.

§4C5 · The Byte Latent Transformer

C5 is the token-free configuration, following Pagnoni et al. (2024). Instead of running BPE, the source ciphertext is grouped 8 bits at a time into raw byte values 0–255 (the underlying data, not the ASCII codes of the '0'/'1' characters), and the target is the raw plaintext bytes. There is no vocabulary beyond the 256 possible byte values plus PAD/BOS/EOS specials.

§4.1 · Why token-free is worth trying here

BPE has an obvious pathology on this task: the ciphertext alphabet is only two symbols ('0', '1'), so BPE merges happily cross the 8-bit character boundaries of the cipher. A source token might cover 6 bits of character i and 4 bits of character i+1, which forces the model to un-mix those two characters again through cross-attention before it can produce either plaintext byte. Byte-level input side-steps that: every 8-bit group is one input unit, aligned perfectly to the cipher's period.

§4.2 · The LocalEncoder: bytes → patches

C5 cannot simply replace one embedding table in Seq2SeqTransformer. Raw bytes give longer sequences than BPE, and ordinary global attention allocates an L × L score matrix for every head. C5 therefore first does cheap local processing, compresses groups of bytes into patches, and runs the shared global encoder on that shorter patch sequence.

The LocalEncoder is a small transformer that runs over raw bytes with windowed local self-attention: the sequence is split into non-overlapping windows of 128 positions, attention runs independently inside each window, and long-range context is carried by the global stack on top. This keeps the byte-side cost O(L · window) instead of O(L²), which matters for long inputs. A maximum full line is roughly 21,000 bits, or 2,670 grouped cipher bytes; training chunks are at most 256 grouped bytes because make_chunks runs first.

class LocalTransformerLayer(nn.Module):
    def _local_attention(self, x, pad_mask):
        B, L, D = x.shape
        W = self.window
        xp = _pad_to_multiple(x, W)                # pad to a multiple of W
        mp = ...pad mask likewise...
        nb = xp.size(1) // W
        xb = xp.view(B * nb, W, D)                 # each window is its own "batch"
        mask = mp.view(B * nb, 1, 1, W)            # key padding per window
        if self.causal:
            mask = mask | causal_mask(W, x.device)
        out = self.attn(xb, xb, xb, mask)
        return out.view(B, nb * W, D)[:, :L]

The reshape trick does the work: batching windows across the batch dimension lets the ordinary MHA class run over each window in parallel, without ever computing a full-length attention matrix. The causal variant is used on the decoder side (§4.4).

Why non-overlapping windows rather than a boolean local mask over the full sequence? A full mask would still make MHA construct the entire L × L score tensor and only hide most entries afterward. Reshaping to (B × number_of_windows, W, D) means those unwanted scores are never allocated. The tradeoff is that adjacent bytes on opposite sides of a window cannot communicate locally; the later global patch transformer provides the long-range route.

§4.3 · Entropy-based patch boundaries

After the byte-level layers, the LocalEncoder pools variable groups of adjacent byte vectors into single patch vectors. Following the paper, the boundaries are not a fixed stride; they come from a byte n-gram entropy model. A new patch starts wherever the next byte is hard to predict: information-dense regions get many short patches, predictable stretches collapse into a few long ones.

“Entropy” here is not a loss and is not trained with gradients. The code counts which byte follows each zero-, one-, or two-byte context in the training split. entropy_curve uses the longest context with enough observations and backs off when it has too little evidence. A high entropy means many different next bytes followed that context, so the next position is a sensible place to spend a new patch.

class ByteEntropyModel:
    """Order-k byte n-gram model used to place BLT patch boundaries."""

    def patch_boundaries(self, seq, threshold, max_patch=16):
        curve = self.entropy_curve(seq)
        starts, run = [0], 1
        for i in range(1, len(seq)):
            if curve[i] > threshold or run >= max_patch:
                starts.append(i)
                run = 1
            else:
                run += 1
        return starts

Two design choices. First, the threshold is set at the 75th percentile of the training entropy distribution (entropy_percentile=75 in SHARED), not as an absolute number of bits. A fixed cutoff of "7 bits" would never fire on ciphertext that tops out at ~6 bits of entropy, and every patch would silently hit the length cap. Second, max_patch=12 caps how long a predictable run may grow; without it, a highly-predictable stretch could swallow a whole sequence, and pooling cost becomes unbounded.

§4.4 · PatchPooling: cross-attention with a learned query

Once boundaries are known, each variable-length group of byte vectors needs to collapse to one patch vector. The BLT paper uses cross-attention with a learned query, and this codebase does the same.

Each example has a different number of differently sized patches, so they cannot be represented as one ordinary rectangular tensor of patch contents. collate_bytes instead constructs member[b, k, i]: it is true exactly when source byte i belongs to patch k. The pooling code turns that membership into an attention mask. This is a tensor-friendly encoding of a ragged list[list[list[int]]].

class PatchPooling(nn.Module):
    def __init__(self, d_model, n_heads=4, dropout=0.1):
        super().__init__()
        self.query = nn.Parameter(torch.zeros(d_model))
        nn.init.normal_(self.query, std=0.02)
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.norm = make_norm("layernorm", d_model)

    def forward(self, x, member):
        # member: (B, K, L) bool, True = byte in patch
        B, K, L = member.shape
        D = x.size(-1)
        q  = self.query.to(x.dtype).expand(B * K, 1, D)
        kv = x.unsqueeze(1).expand(B, K, L, D).reshape(B * K, L, D)
        mask = (~member).reshape(B * K, 1, 1, L)
        pooled = self.attn(q, kv, kv, mask).view(B, K, D)
        return self.norm(pooled).masked_fill(empty.unsqueeze(-1), 0.0), empty

Two points. First, patches are batched over B · K so a single MHA call handles the whole layer. Second, and the docstring flags this from a bug that took a while to find, order-preserving pooling matters. An earlier version used a mean of the byte vectors, which is permutation-invariant, and the model could not recover a character from what amounted to a bag of its bits. Cross-attention with a learned query gives the pooling access to token position information carried in the byte vectors and fixes this immediately.

§4.5 · Patch-causal cross-attention and the source-memory path

The global decoder operates at the patch level: pooled target patches are shifted right by one slot (behind a learned start_patch) and run through the transformer with causal self-attention plus cross-attention to the encoder's source patches. The LocalDecoder then refines that patch-level signal back down to per-byte logits.

The shift by one patch is the patch-level equivalent of tgt[:, :-1] versus tgt[:, 1:]: a global target-patch state must summarize only bytes that are already known. The local decoder then combines three sources for every byte: earlier target bytes, allowed earlier global target-patch states, and encoder source-patch memory.

The tricky part is preventing information leakage between target bytes. Global state g_j summarises target positions ≤ j · patch_tgt − 1, so byte position i may attend states j ≤ (i+1) // patch_tgt without ever seeing its own byte or any future one.

def patch_causal_cross_mask(self, Lt, K, device):
    """(1, 1, Lt, K) bool: byte position i may attend global slot j iff
    j <= (i + 1) // patch (see module docstring); True = blocked."""
    i = torch.arange(Lt, device=device)[:, None]
    j = torch.arange(K, device=device)[None, :]
    return (j > (i + 1) // self.patch)[None, None]

A second cross-attention was added inside LocalDecoderLayer that attends the encoder source-patch memory directly. The docstring records the reason: on a positional cipher, without a direct byte→source path, the only route for source information is the pooled-target-patch chain, which empirically never learns the alignment. An 80-epoch C5 run without this cross-mem attention converged to a source-free language model. Source memory carries no target information, so this extra path does not break causality.

§5The data pipeline

build_data turns two parallel text files into dictionaries of padded tensors. The order matters: verify pairs, split whole lines, chunk selected lines, train preprocessing only on training data, wrap chunks in a Dataset, and batch similar lengths together. Each step prevents a specific failure.

plain, cipher = load_pairs()
verify_dataset(plain, cipher)       # lengths, binary alphabet, XOR round-trip
splits = get_splits(len(plain))     # line indices: 80% train, 10% val, 10% test
chunked = {name: make_chunks(plain, cipher, ids)
           for name, ids in splits.items()}

# C1-C4 only: learn vocabularies from training text, then reuse for all splits.
tgt_tokenizer = train_bpe_tokenizer(training_plaintext, 8000, ...)
src_tokenizer = train_bpe_tokenizer(training_ciphertext, 1024, ...)

Splitting occurs before tokenizer or entropy-model training. Otherwise preprocessing would inspect validation and test text, a form of data leakage that makes evaluation less honest. The saved splits.json and cached tokenizers also ensure all five configurations see the same examples rather than benefiting from a lucky new random split.

§5.1 · Chunking with a phase-aligned window

Lines run up to 2670 plaintext characters, or 21360 cipher bits. Full lines make the cross-attention copy mechanism hard to learn, because the model must simultaneously discover the phase-8 XOR pattern and align long sequences. Chunking with CHUNK_CHARS = 256 shortens the alignment range and multiplies the number of training pairs.

CHUNK_CHARS is a multiple of 8, so every chunk starts at XOR-key phase 0 and the per-position cipher mapping is identical across chunks. If you chunked at 250, chunk 1 would start at key phase 0 but chunk 2 at key phase (2000 mod 8) = 0 by accident; try 253 and every chunk starts on a different phase and the model has to relearn eight rotations of the same pattern.

The indexing in make_chunks is worth reading literally: plaintext slice [start:start+256] pairs with ciphertext slice [8*start:8*(start+256)]. Metadata stores the original line index and chunk position. Evaluation decodes each chunk, sorts by that position, and joins them, so reported metrics are still for complete original lines.

§5.2 · BPE from scratch, with a linked-list corpus

The tokenizers for C1-C4 are learned per-side BPEs, trained with a pure-Python implementation. Sennrich et al. (2016) describes the algorithm in five lines: start from the character alphabet, count adjacent symbol pairs, merge the most frequent, repeat until the vocab is full. The engineering challenge is making that fast enough on a corpus of ~1.9M source tokens.

The trick is a doubly-linked-list representation over one flat symbol array plus lazy delete-on-pop counting:

• Pair counts are updated only around each merge, not recomputed from scratch.
• A heap tracks "most frequent pair"; stale entries are discarded on pop instead of being deleted eagerly (a full delete would need the pair's heap index).
• Encoding memoises per pre_tokenize unit, since units repeat constantly in natural text.

A second decision: the ciphertext BPE is deliberately small (src_bpe_vocab_size=1024). The cipher alphabet has only two symbols, so 1024 merges already yield ~10-bit tokens, and a larger table would both inflate the source embedding well beyond the other configurations and leave most of its rows undertrained on ~1.9M source tokens.

§5.3 · The byte entropy model

For C5's patch boundaries. An order-2 Markov model over training bytes: H(next | context) is estimated from conditional distributions of the following byte given the previous two, backing off to shorter contexts (finally to the unigram) when a context was seen too few times.

The paper trains a small byte-level LM for this. Pure Markov works here for two reasons: bytes have a tiny state space, and the metric the boundaries care about, which is relative unpredictability at a position, does not need a deep model to be well-calibrated.

§5.4 · Length-bucketed batching

Sorting sequences by length and grouping similar-length examples together minimises padding waste. The sampler:

class BucketBatchSampler(Sampler[list[int]]):
    def __init__(self, lengths, token_budget, max_batch_size=64, ...):
        order = sorted(range(len(lengths)), key=lambda i: lengths[i], reverse=True)
        # greedy: fill batches under a source-token budget, capped at max_batch_size

The token budget (8192) caps the biggest batch, and a hard max_batch_size=64 keeps short-line batches from ballooning into thousands of examples. Batches shuffle every epoch (deterministically from seed + epoch) so the model does not see the same length order twice.

collate finds the longest source and target in the chosen batch, allocates filled PAD tensors, and copies each shorter 1D tensor into a row. Bucketing similar lengths means fewer of those allocated cells are useless padding. A token budget is more predictable than a fixed batch count: 64 short sequences may fit where only a handful of long sequences fit.

§5.5 · Weight initialisation is not decorative

init_transformer_weights uses Xavier for linears and N(0, d^−1/2) for embeddings. The docstring flags why this matters:

PyTorch's default N(0, 1) embedding init is catastrophic here. Combined with the √d_model embedding scale it produces std~16 activations that (a) drown the unit-amplitude sinusoidal positional encoding, destroying the positional information this task depends on, and (b) yield hugely over-confident initial logits through the tied output head.

Get this wrong and every configuration starts training in a hole that only C2 and C5 have any chance of climbing out of, because RoPE does not depend on additive positional signal and BLT does not depend on the sinusoidal-embedding balance at all.

§5.6 · The training loop, line by line

logits = model(src, tgt[:, :-1], *extra)
loss = loss_fn(logits, tgt[:, 1:], pad_id, label_smoothing=0.1)

optimizer.zero_grad(set_to_none=True)
loss.backward()
clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.step()

Label smoothing uses 0.1 only for training: it discourages infinitely confident logits by assigning a little probability mass to other tokens. Validation passes smoothing 0 so its cross-entropy remains a direct measure of prediction quality. At each epoch the code saves last.pt, saves best.pt only when validation improves, and stops after ten non-improving epochs. Test metrics are computed once from that best checkpoint, not whichever epoch happened to run last.

§6Results

All five configurations trained under identical hyperparameters, early stopping on validation loss with patience 10. Below are the final test metrics and per-configuration training benchmarks. Everything is greedy-decoded, as the assignment requires.

§7Viva question bank

The viva is 50% of the mark. What follows is a long list of questions grouped by topic: a mix of "explain your code," "explain the theory," "what would break if…," and "compare against the alternative." Answering all of these out loud, without hesitating on any one, is the bar to aim for.

Every question below has a fold-out answer. Read the question, answer it out loud, then click to reveal.

§7.1 · Attention fundamentals

1. Write down the scaled dot-product attention formula. Why do we divide by √d_k and not d_k?

   Attention(Q, K, V) = softmax(Q Kᵀ / √d_k) V. The dot product of two random d_k-dim vectors with unit-variance entries has variance d_k and standard deviation √d_k. Dividing by √d_k restores unit-variance logits so softmax does not saturate as d_k grows. Dividing by d_k would over-shrink the logits and blunt attention.

2. What are Q, K, V geometrically? Why do we need three projections instead of just Q = K = V?

   Think of attention as content-addressable retrieval. Q is "what am I looking for?", K is "what do I advertise?", and V is "what payload do I hand over if selected?" With three separate projections the model can advertise one thing (K) but deliver another (V), and it can ask about a role from a different subspace than the one it stores under. If Q = K = V you collapse retrieval and payload into the same vector and the layer is provably weaker: cosine self-similarity dominates and attention degenerates to soft copying.

3. Why does your implementation fill masked positions with torch.finfo(scores.dtype).min rather than -inf?

   If every position in a row is blocked (a fully padded key row, or an empty BLT patch), softmax over a row of -inf is 0 / 0 = NaN, and the NaN then poisons downstream matmuls and the loss. The dtype's most negative finite value gives softmax a well-defined uniform output on that row, which is harmless because those rows are already excluded from the loss.

4. What is the shape and semantics of the mask tensor in your code? What convention did you pick, and why?

   Bool, True = blocked, broadcast against (B, H, L_q, L_k). Key-padding masks arrive as (B, 1, 1, L_k); causal masks as (1, 1, L_q, L_k). A boolean lets you OR the two together and pass the result in. "True = blocked" is chosen over "True = keep" because masked_fill takes a "positions to fill" argument, so the semantics line up directly with PyTorch.

5. Trace the shape of Q, K, V through MultiHeadAttention.forward.

   Input query has shape (B, L_q, d_model). q_proj keeps that shape. _split reshapes to (B, L_q, n_heads, head_dim) and transposes to (B, n_heads, L_q, head_dim) — heads become an independent batch dim. Attention returns (B, n_heads, L_q, head_dim). The tail transposes back to (B, L_q, n_heads, head_dim), .view(B, Lq, d_model) concatenates heads, and o_proj keeps that shape.

6. Why does MHA need the output projection o_proj? What would happen if you dropped it?

   After concatenation the head outputs occupy disjoint slices of d_model. Without o_proj, downstream layers can never mix information from different heads within the same residual stream — heads run as if in parallel silos. o_proj is the mixer that lets one head's answer feed another head at the next layer.

7. What is the computational complexity of scaled dot-product attention in L? Which term dominates for long sequences?

   O(L² · d_k) for the score matrix plus O(L² · d_v) for the value read-out. Both are quadratic in L; the L² · d factor dominates any linear-in-L term. For long sequences this is the reason to reach for windowed / sparse / linear attention.

8. Which does more computation for the same d_model: attention or the FFN?

   At short-to-medium L, the FFN dominates — its FLOPs are 2 · L · d_model · ffn_dim ≈ 8 · L · d_model² (with ffn_dim = 4·d_model). Attention is 4 · L² · d_model (Q/K/V/O projections) plus 2 · L² · d_model (scores + values). FFN > attention while L < 2 · d_model; past that L, attention overtakes because of the L² term.

§7.2 · Multi-Head vs Grouped-Query Attention

1. Why does MHA use multiple heads instead of one big head of dimension d_model?

   Because softmax(Q Kᵀ / √d) V produces one attention distribution per token pair — with a single head, every position must pick one thing to attend to. Splitting into heads gives the model several distributions in parallel, so one head can track syntactic dependency, another can attend to the previous token, another to a specific role. FLOPs are almost identical (the head_dim shrinks by the same factor as heads grow), so multi-head is a strictly-more-expressive rearrangement of the same compute.

2. What is the parameter and FLOP difference between C1 (MHA, 8 heads) and C3 (GQA, 8 query heads, 2 KV heads)?

   C1's K/V projections are each d_model × d_model. C3's are d_model × (n_kv_heads · head_dim) = d_model × (d_model / 4) — a quarter of the parameters. Total model params: C1 = 9.68M, C3 = 8.89M (≈ 8% fewer). FLOPs during the score/value matmuls are unchanged, because K/V are repeat-interleaved back up to 8 heads before attention — GQA's win is on the K/V projections and (mainly at inference) the KV cache footprint.

3. Explain repeat_interleave in your GQA forward. Why not expand?

   K/V start with shape (B, 2, L, 32), while Q has eight heads. The code needs the mapping [kv0, kv0, kv0, kv0, kv1, kv1, kv1, kv1]. A plain expand can only enlarge a dimension whose size is one, so it cannot directly turn this two-head axis into that grouped eight-head mapping. repeat_interleave(4, dim=1) states the mapping clearly and produces the exact shape expected by the unchanged attention function. It does materialise the repeated activation; the persistent saving comes from projecting and caching only two distinct K/V heads.

4. Why does GQA save memory at inference but not at training in your benchmarks?

   At training you still materialise the full (B, H, L_q, L_k) score tensor for every layer — that dominates activation memory and is unchanged by GQA. At inference (autoregressive decoding), you cache K and V for every prior token, and the cache is what GQA shrinks by h / g×. This task's benchmark is a training pass, so C3's peak GPU (2.99 GB) is essentially the same as C1's (3.00 GB).

5. What is MQA? Where does it sit relative to MHA and GQA?

   Multi-Query Attention (Shazeer 2019) is GQA with n_kv_heads = 1: all query heads share a single K/V pair. Its KV cache is h× smaller than MHA — the extreme end of the trade-off — but the sharing empirically hurts quality on many tasks. GQA (Ainslie et al. 2023) is the middle ground: h/g groups of query heads each get their own K/V, keeping most of the memory win with almost all of MHA's quality.

6. What is Multi-head Latent Attention (MLA), and how is it different from GQA?

   MLA (DeepSeek) caches a low-rank latent projection of K and V instead of the K/V themselves. The heads are still fully independent — no sharing — but the on-disk cache is a much smaller latent tensor, decompressed on-the-fly per step. It attacks the same cache-size problem as GQA but keeps every head's degrees of freedom, at the cost of a small decompression matmul per step.

7. Why did C3's quality drop slightly compared to C1 on this task, despite the same query head count?

   Cutting the number of independent K/V heads from 8 to 2 forces four query heads to share each K/V pair — they can still ask different questions (their Q projections differ) but the "index" they are searching over is coarser. On a positional-cipher task where the alignment is precise, that loss of K/V diversity translates directly into slightly noisier attention. It is not a big drop (bit-acc 0.712 → 0.704) but it is real.

§7.3 · Positional encoding

1. Write down the sinusoidal PE formula. Why sin/cos at logarithmically-spaced frequencies?

   PE(pos, 2i) = sin(pos / 10000^2i/d), PE(pos, 2i+1) = cos(...). Logarithmically-spaced frequencies span wavelengths from 2π up to 10000·2π, so together the dimensions form a continuous binary-like encoding: low-index dimensions swing fast (fine position), high-index dimensions swing slowly (coarse position). No single frequency has to resolve the full range.

2. Why is sinusoidal PE additive to token embeddings, not concatenated?

   Addition keeps the input dimensionality at d_model — every downstream layer stays the same size — and lets the model choose per-dimension how much of the residual budget to spend on position versus content. Concatenation would grow the input and pin position to fixed dimensions the rest of the network can never overwrite. Empirically addition works, so there is no reason to pay concatenation's cost.

3. Why does your code scale embeddings by √d_model before adding sinusoidal PE?

   The embeddings are initialised at N(0, d^−1/2), so a raw embedding vector has ~1/√d amplitude per dimension. The sinusoidal PE is unit-amplitude per dimension. Without the √d rescale, PE would dominate the sum and swamp the content signal. Multiplying by √d brings the embedding amplitude back to unit scale, so the two components are comparable.

4. What does register_buffer(persistent=False) mean, and why is it the right choice for the sinusoidal PE table?

   register_buffer declares a non-trainable tensor that still travels with the module: it moves to GPU with .to(), saves and loads with the state dict. persistent=False means it is not written to the state dict — perfect for the sinusoidal table because it is fully determined by d_model and max_len and can be recomputed at load time, so there is no point paying for the checkpoint bytes.

5. Explain RoPE. Why does it rotate Q and K only, not V?

   RoPE encodes position by rotating the Q and K vectors at each position by an angle proportional to that position. Position enters the attention score — Q^TK — as a relative offset, without any additive PE at the input. V is the payload attention hands over once the score is computed; the retrieval already correctly weighted by position, so rotating V would just re-encode position inside the payload with no benefit.

6. Prove informally that after RoPE, the dot product Q^TK depends only on the relative offset m − n.

   The rotation at position m is a block-diagonal orthogonal matrix R(m·θ). RoPE-Q at m is R(m·θ)·q_m; RoPE-K at n is R(n·θ)·k_n. Their dot product is q_m^T R(m·θ)^T R(n·θ) k_n = q_m^T R((n − m)·θ) k_n, using R(a)^TR(b) = R(b − a). The score depends on positions only through n − m — relative offset.

7. Why is RoPE applied only to self-attention, not cross-attention, in your code?

   Cross-attention queries live on the decoder sequence, keys on the encoder sequence. Their positions belong to two different coordinate systems that share no meaningful "offset" — decoder token 3 and encoder token 3 have no positional relationship. Rotating both would inject a spurious relative-position signal. Cross-attention passes rope=None.

8. Your C2 (RoPE) beats C1 by a huge margin on this task. Why is that unsurprising?

   The cipher is a period-8 XOR: the correct transformation for character i depends on i mod 8. That is a purely-relative-position pattern. RoPE places relative offsets directly in the QK dot product, so the model can learn "attend to keys whose relative offset satisfies i mod 8 = c" essentially in one step. Sinusoidal PE requires the model to reconstruct relative offsets from the difference of two additive encodings through matmuls — same information, harder path.

9. What is NoPE, and when does it work?

   NoPE = no positional encoding at all. Kazemnejad et al. (2023) showed that decoder-only transformers with causal masking can learn position implicitly, because the mask breaks permutation symmetry — layer l's output at position i already carries a summary of "everything before i". It works on modest-length causal LMs and often generalises better to longer sequences than sinusoidal, because there is no learned schedule to break. It fails on bidirectional encoders (no mask, so no positional asymmetry) and on precise-alignment tasks like this cipher.

10. What is the half-split RoPE convention? What is the alternative pair-wise convention, and why prefer half-split?

    Half-split: split the head dim into two halves (x1, x2), rotate as (x1·cos − x2·sin, x2·cos + x1·sin) at angle m·θ_i. Pair-wise: interleave (x0, x1, x2, x3, …) into pairs and rotate each pair independently. They are algebraically equivalent up to an index permutation. Half-split is preferred in implementations because it can be expressed as two contiguous slices and one torch.cat — no interleave, no strided indexing — so it is faster and easier to fuse.

§7.4 · Normalisation

1. Write down LayerNorm and RMSNorm side by side. What did RMSNorm drop?

   LayerNorm: y = (x − mean) / √(var + ε) · γ + β. RMSNorm: y = x / √(mean(x²) + ε) · γ. RMSNorm dropped the mean-centering and the bias β. It normalises by the root-mean-square of the activations rather than the standard deviation.

2. Why do your norm modules compute statistics in fp32 even under bf16 autocast?

   bf16 has 8 mantissa bits — the variance of a small activation vector can underflow the representable range, or lose enough precision that var + ε rounds to just ε. Casting to fp32 for the mean/var/rsqrt keeps the statistics accurate; the final .to(dtype) puts us back in bf16 for the rest of the pipeline. This is the standard trick used in every serious LLM codebase.

3. What is Pre-LN vs Post-LN? Why did you pick Pre-LN?

   Post-LN (the original Vaswani et al. paper): x = LayerNorm(x + Sublayer(x)) — norm sits after the residual add. Pre-LN: x = x + Dropout(Sublayer(LayerNorm(x))) — norm sits inside the residual branch, on the sublayer's input. Pre-LN is trained without warm-up much more easily because the residual highway is not renormalised at every step, so gradients flow cleanly through depth. Post-LN needs a careful warm-up schedule and often diverges without it.

4. Why does each of your stacks end with a final norm layer?

   Pre-LN normalises the input of each sublayer, so the residual highway's output at the top of the stack is never normalised. If you fed that unnormalised activation to a linear head, you would see the head's output scale drift with depth. One extra norm at the top fixes that.

5. Why is RMSNorm cheaper than LayerNorm?

   FLOPs: RMSNorm skips the mean subtraction (one reduction per row) and the bias add (one broadcast add per output). Memory: no β parameter, so per-layer parameter count drops by d_model. Bandwidth: mean-centering requires two passes over the row (mean, then centred squared sum); RMSNorm needs one pass over the squared values. Small per-layer savings, but LLMs have hundreds of norm calls per forward pass.

6. Modern LLMs (Llama, Mistral, T5) use RMSNorm. Why has it become the default?

   Because the ablations kept saying "no accuracy loss, small speed win". Zhang & Sennrich (2019) argued the mean-centering does not actually do anything useful in transformer residual streams — the residual dominates the scale, not the mean. Once the community trusted that empirically, the free speed-up became a default. Your C1 → C4 result (bit-acc 0.712 → 0.707, memory 3.00 GB → 2.85 GB) is a clean example.

7. Why did the C1 → C4 swap change quality so little on this task?

   Because on a well-behaved residual stream, mean-centering does almost nothing — the residual amplitude dominates, and the mean of a d-dim activation from a random projection is close to zero anyway. RMSNorm and LayerNorm produce essentially the same normalised outputs in that regime. The "aha" is that the C1 → C4 wash confirms RMSNorm is a strict cost improvement here.

8. What is the role of the residual connection? What breaks if you drop it at 24 layers?

   The residual creates an identity path from input to output that gradients can travel down unimpeded. Without it, gradients through 24 nonlinear layers vanish or explode, and the loss landscape looks like the "without skip connections" panel of the Li et al. (2018) visualisation from L2 — a jagged surface with no basin. With residuals the same landscape has a wide, smooth bowl around the optimum, and training just works.

§7.5 · Tokenisation and BPE

1. Describe the BPE training algorithm in five lines.

   (1) Initialise the vocabulary with every distinct character in the corpus. (2) Count every adjacent symbol pair. (3) Merge the most frequent pair into a new symbol; add that symbol to the vocab. (4) Update pair counts in the neighbourhood of each merge site. (5) Repeat until the vocabulary reaches the target size.

2. Why does your BPE use a doubly-linked-list corpus and lazy heap? What is the naïve alternative and why is it too slow here?

   Naïve BPE recounts every pair from scratch after each merge — O(corpus × merges). At 1.9M source tokens and 1024 merges that is ~2 · 10⁹ operations in Python. The linked-list-plus-incremental-count approach updates only the pair counts around each merge site (constant work per merge), and the lazy heap keeps the "most frequent pair" query cheap without needing an eager delete on a heap that does not support it. Together these turn training from hours to seconds.

3. Why does your source (ciphertext) BPE have vocab 1024 while the target (plaintext) BPE has vocab 8000?

   The ciphertext alphabet has only two symbols (0, 1), so 1024 merges are enough to build ~10-bit compound tokens covering the whole training distribution. A larger table would (a) inflate the source embedding well beyond what the other configurations use — undermining the "same hyperparameters" contract of the ablation — and (b) leave most rows undertrained on ~1.9M source tokens. The plaintext side has a full English character set, so 8000 subword pieces is where the compression asymptote sits.

4. What is pre_tokenize doing? Why do you split on whitespace, and what would happen without it?

   pre_tokenize splits the input into "merge units" — either a leading run of whitespace + the following word, or a run of pure whitespace. Merges are then confined to within each unit. Without it, BPE happily merges across spaces and learns multi-word phrases as single tokens (the classic "Church of Scientology " token). Those tokens occur a handful of times in training, memorise particular strings, and are useless at test time.

5. Ciphertext contains no whitespace. Does your pre_tokenize still do anything on it?

   No — it returns the whole ciphertext as one unit, so merges can span the entire string. That is fine here because the ciphertext is a low-entropy binary sequence: cross-boundary merges are exactly what you want to compress it. The same pre-tokeniser correctly handles the plaintext side (which does need word-boundary discipline) and passes through the ciphertext unchanged.

6. What is the difference between BPE, WordPiece, and Unigram tokenisation?

   BPE is greedy bottom-up: merge the pair with highest raw frequency. WordPiece (BERT) is greedy bottom-up too but merges the pair with highest likelihood ratio — count(ab) / (count(a) · count(b)) — which favours merges that "explain" the corpus, not just frequent ones. Unigram LM (SentencePiece) is top-down: start with a huge candidate vocab, treat each vocab item as a probability, and iteratively drop the pieces that hurt the corpus likelihood least. Unigram is probabilistic, the other two are deterministic.

7. What is byte-level BPE? Why isn't your ciphertext BPE already byte-level despite operating on binary strings?

   Byte-level BPE runs BPE over UTF-8 bytes — 256 base symbols, guaranteed to encode any Unicode input without an <unk>. Your ciphertext BPE runs over the ASCII characters '0' and '1', so its base alphabet has 2 symbols, not 256. It happens to also be byte-level in the trivial sense that each character is one byte — but semantically it is a two-symbol BPE, not a byte-level BPE. C5's BLT is the one that goes to real bytes: it groups every 8 characters of the cipher into one byte value 0–255 and works over those.

§7.6 · The Byte Latent Transformer (C5)

1. What does "token-free" mean in your C5 model? What is the vocabulary of the byte side?

   "Token-free" means there is no learned tokeniser between raw bytes and the model — no BPE merges, no <unk>. The vocabulary is exactly 256 possible byte values (0–255) plus three specials: PAD, BOS, EOS. Total: 259 embeddings.

2. Why is 8 bits of the ciphertext grouped into a byte value, not fed as ASCII codes of '0'/'1'?

   The ciphertext is a binary sequence written out as text. The underlying data is one byte per 8 characters. Feeding ASCII would give the model 2-symbol input at 8× the length and force it to relearn the 8-bit grouping structure that is baked into the cipher. Grouping into byte values 0–255 aligns the input with the cipher's period and cuts the sequence length by 8×.

3. Walk through the forward pass of BLTModel from cipher bytes to next-byte logits.

   (1) Source bytes → LocalEncoder: byte embedding + windowed local self-attention layers + PatchPooling using entropy-based membership → source patches (B, K_s, D). (2) Source patches → sinusoidal PE at patch level → global_encoder stack → memory. (3) Target bytes → byte embedding + PE → tgt_pooling gives target patches, shifted right by one slot behind a learned start_patch → global_decoder (causal self-attention + cross-attention to source memory) → global states g. (4) Byte embeddings + LocalDecoder: causal windowed self-attention, cross-attention to g under patch-causal mask, cross-attention to source memory → per-byte logits (B, L_t, 259).

4. What is windowed local attention and why is it necessary here? What does the trade-off cost you?

   The sequence is split into non-overlapping windows of W positions; attention runs independently inside each window. Cost is O(L · W) instead of O(L²). A maximum full line has about 21,000 cipher bits, or 2,670 grouped bytes; training chunks contain at most 256 grouped bytes. The cost is that long-range interactions inside the local layer are impossible — a byte in window k cannot attend a byte in window k+1 directly. Long-range context is carried by the global stack on top, which sees pooled patch representations.

5. Explain patch-causal cross-attention. Why does byte position i attend global slot j iff j ≤ (i+1) / patch_tgt?

   Global state g_j is produced from target patches shifted right by one, so g_j summarises target input positions ≤ j · patch_tgt − 1. Byte i lives in patch p = i // patch_tgt. For the cross-attention to stay causal, byte i must see g_j only when g_j's content is strictly older than i — that means j · patch_tgt − 1 < i, or equivalently j ≤ (i + 1) // patch_tgt. Any slot beyond that would leak i's own token or a future one.

6. Why did the naive "pool bytes by mean" fail? What replaced it?

   A mean is permutation-invariant: (b₁, b₂) and (b₂, b₁) pool to the same vector. On this task where the source-target alignment is bit-precise, that erased the information the model needed — the pooled patch became a "bag of bits" and the model could not recover a character. It was replaced by cross-attention with a learned query vector, which sees the byte vectors in order and uses their positional encoding to distinguish them.

7. Why does your LocalDecoderLayer have two cross-attentions?

   One attends the target-side global states g (needed for the language-model-like flow of BLT). The other attends the encoder source-patch memory directly. On a positional cipher, without a direct byte→source path, the only route for source information is g's pooled-target-patch chain — which empirically never learns the alignment. An 80-epoch run without the second cross-attention converged to a source-free language model. Adding the source-memory path fixed it, and it does not break causality because source memory carries no target information.

8. Explain the entropy-based patch boundaries. Why is the threshold a percentile of the training entropy distribution, not an absolute bit count?

   A small byte n-gram model gives H(next byte | context) at every position. A new patch starts wherever the next byte is hard to predict — information-dense regions get many short patches, predictable regions collapse into few long ones. The threshold is set at the 75th percentile of training-time entropies because an absolute figure in bits is not portable: on this cipher H tops out around 6 bits, so a "reasonable" cutoff of 7 bits would never fire and every patch would silently hit the length cap. Percentiles are calibrated to the distribution and generalise across data.

9. What does max_patch guard against?

   Runaway patches. A stretch of highly-predictable bytes could otherwise grow into a single patch of thousands of positions — the cross-attention pooling cost is O(patch_length) and would explode. max_patch = 12 forces a boundary at least every 12 bytes and bounds the worst-case cost.

10. Why is the BLT's global transformer literally the same class as C1's? What would break in the ablation if it silently used RMSNorm or RoPE?

    The ablation contract is "each of C2–C5 changes exactly one component from C1". C5 changes the tokenisation. If the C5 global transformer also swapped LayerNorm for RMSNorm or Sinusoidal for RoPE, C5 would be a two-variable change and the comparison to C1 would tell you nothing about tokenisation alone — you would not know whether the win came from bytes or from RoPE. Reusing the C1 Encoder/Decoder classes verbatim (only the input pipeline differs) preserves the single-variable property.

§7.7 · Assembly and training

1. Draw the forward pass of a Pre-LN EncoderLayer as an equation. Now do DecoderLayer.

   Encoder: x ← x + Dropout(SelfAttn(Norm(x))); x ← x + Dropout(FFN(Norm(x))). Decoder: x ← x + Dropout(SelfAttn(Norm(x), causal)); x ← x + Dropout(CrossAttn(Norm(x), memory)); x ← x + Dropout(FFN(Norm(x))). Three sublayers on the decoder, two on the encoder, each Pre-LN inside its own residual branch.

2. Why does your decoder cross-attention use MHA even when the self-attention is GQA (C3)?

   The ablation only substitutes the self-attention. Cross-attention is a different operation (queries and keys live on different sequences with no shared coordinate system) and it is kept fixed across all five configurations so that the C3 vs C1 comparison isolates the self-attention change alone.

3. Why does your model tie the output projection weights with the target embedding?

   Weight tying (Press & Wolf 2017) uses the target embedding matrix as the transpose of the output projection. Two wins: (a) it halves the parameter count of a large vocab head — meaningful when the target vocab is 8000+ pieces — and (b) it improves generalisation because the same "meaning" of a token must serve both as input embedding and as output logit direction, which is a strong regulariser.

4. Why did the default PyTorch embedding initialisation break training on this task?

   PyTorch initialises nn.Embedding at N(0, 1). Combined with the √d_model = √256 = 16 scaling used in Vaswani-style transformers, the scaled embeddings end up with std ~16 per dimension — vastly larger than the unit-amplitude sinusoidal PE being added on top. The PE gets drowned out, the positional signal is destroyed, and the tied output head produces massively over-confident initial logits that need many warm-up steps just to unblow up. The fix is N(0, d^−1/2) embedding init, which makes the scaled embeddings unit-variance and comparable to the PE.

5. What is label smoothing and why is it set to 0.1 here?

   Label smoothing replaces the one-hot target distribution with a mixture: (1 − ε) on the correct class, ε / (V − 1) on every other class. It prevents the model from becoming overconfident on training examples — hard 1.0 targets push logits toward ±∞, which hurts calibration and often generalisation. 0.1 is the Vaswani et al. default and a widely-tested sweet spot for machine translation-shaped tasks.

6. Explain the warm-up schedule. Why do transformers famously need warm-up?

   At initialisation the LayerNorm gains and Adam's second-moment estimate are both far from their steady-state values, so a full learning rate in the first few hundred steps causes huge parameter updates and blows up the loss. Warm-up linearly ramps the LR from 0 to peak over the first N steps (here 1000) so early updates are gentle enough that the optimiser statistics can converge before the LR hits its target.

7. What does BucketBatchSampler do that DataLoader(shuffle=True) would not?

   Regular shuffling mixes short and long sequences within a batch, so every batch is padded to the longest example and short sequences waste compute on padding. Bucketed sampling sorts by length and greedily fills batches under a source-token budget, so all sequences in a batch are similar length and padding waste is near zero. Batches are shuffled every epoch (deterministic from seed + epoch) so the model still sees varied length orderings across training.

8. Why is CHUNK_CHARS a multiple of 8? What if you used 250 instead?

   The cipher key has period 8. If CHUNK_CHARS = 256, every chunk starts at key phase 0 and the per-position XOR mapping is identical across chunks — the model learns one pattern once. With 250, chunk k starts at key phase (250k mod 8) = (2k mod 8), so the model would have to learn 4 different phase-aligned mappings and figure out from context which one applies to each chunk. Same information, much harder learning problem.

9. What is the causal mask in a decoder, exactly? Which positions does position i attend?

   An upper-triangular boolean mask with True above the diagonal — blocking all positions j > i. Position i attends only positions 0, 1, …, i (itself included). This is what enables teacher-forcing training: the target sequence is fed in one shot but each position's prediction depends only on prior positions, so training in parallel matches autoregressive generation at inference.

10. What is the difference between a key-padding mask and a causal mask? Can they be applied together?

    A causal mask blocks future positions and depends only on (L_q, L_k) — the same for every batch element. A key-padding mask blocks pad positions and depends on the batch content — shape (B, L_k). Yes, they are combined by OR-ing (blocked-in-either is still blocked) and broadcasting to (B, 1, L_q, L_k) before being passed into the score fill.

§7.8 · Metrics and evaluation

1. Define each of the four metrics you report. Which ones are order-sensitive?

   Bit accuracy — fraction of matching bits between the re-encoded prediction and the reference cipher (order-sensitive). Sequence accuracy — fraction of predictions that exactly match the reference (order-sensitive, all-or-nothing). Levenshtein — character-level edit distance (order-sensitive, but forgiving of insertions and deletions). BLEU/ROUGE — n-gram overlap; BLEU counts precision of n-grams up to length 4 with a brevity penalty, ROUGE counts recall. Both are partially order-sensitive (order matters inside an n-gram window, not globally).

2. Why do you re-encode the predicted plaintext to bits before computing bit accuracy?

   The model outputs plaintext characters, but the metric wants the fraction of correct cipher bits. Re-encoding the predicted plaintext through the known cipher scheme (text_to_bits) recovers the bit sequence the model would need to have decoded from. If you compared plaintext characters directly you would lose the sub-character resolution that bit accuracy is designed to expose (a one-character error is 8 wrong bits, not 1).

3. Why is sequence accuracy so low (0–5%) on C1–C4 while BLEU is 60+?

   Sequence accuracy on a 256-character line requires every one of those characters to be correct after greedy decoding. Even a 99%-per-char model would produce a perfect line only 0.99²⁵⁶ = 7.7% of the time. BLEU sees the same predictions but rewards partial n-gram overlap, so it stays high while sequence accuracy collapses. The gap tells you the model is nearly right most places, wrong occasionally — not that it "learned nothing".

4. Why does the assignment restrict BLEU/ROUGE to tokenised models?

   BLEU and ROUGE are defined on tokens, and their scores are not comparable across different tokenisations. Bytes and BPE pieces are different units, so a BLEU score from the byte-level BLT is not on the same scale as a BLEU from the C1 baseline. Bit accuracy, sequence accuracy, and Levenshtein are all defined on the raw characters and remain comparable, so those three are reported for every configuration.

5. Explain the two-row Levenshtein implementation. What is its space complexity?

   Full Levenshtein DP builds an (m+1) × (n+1) matrix; each cell depends only on cell (i−1, j), (i, j−1), (i−1, j−1). So only the current row and the previous row are ever needed. The two-row implementation keeps a length-(n+1) "prev" array and a length-(n+1) "curr" array and rolls them each i. Space is O(min(m, n)) after the swap-to-put-shorter-second at the top; time is unchanged O(m · n).

6. Why is greedy decoding used across all configurations? What would beam search change?

   The assignment specifies greedy for consistency across models — beam search widths interact with each model's calibration differently, so a beam comparison would mix decoding strategy with architecture in the same measurement. Beam search would help most where greedy is stuck on a locally-plausible-but-globally-wrong token (typically raising sequence accuracy the most). It would help C1–C4 more than C5, since C5 is already very confident.

7. How does greedy_decode_blt differ from greedy_decode_subword? Where is the encoder memory reused?

   Structurally the same loop: encode once, then decode step-by-step, feeding the argmax back until EOS or max_len. The BLT version threads the entropy-based patch_member tensor into the encode call — the source-side patch boundaries are computed once from the input and passed in. Both cache the encoder output (memory) outside the loop and reuse it across every decode step, so encode cost is amortised over the full generation. Neither implements a KV cache on the decoder side; the target is recomputed each step.

§7.9 · The ablation, holistically

1. State the guarantee an ablation study is supposed to give. What does "changing exactly one component at a time" buy you methodologically?

   The guarantee is attributability: any measured difference between C1 and Ck is caused by the one component that changed, not by a confound. Changing multiple components at once destroys that guarantee — you cannot tell whether the win came from the tokeniser or the norm. Fixed hyperparameters, fixed data, fixed seed, fixed evaluation are the other legs of the same claim.

2. C2, C3, C4 all changed one component. Only C2 improved. What assumption did the C3/C4 outcomes teach you to distrust?

   The assumption that "component upgrades transfer across tasks". GQA and RMSNorm are widely-reported wins on general LLMs (KV-cache pressure for GQA, cheaper stats for RMSNorm), but their gains are inference- or scale-related — they do not translate to quality wins at this scale, on this task. The ablation says: pick a substitution because it targets a bottleneck you actually have, not because it worked in someone else's paper.

3. C5 replaces the tokeniser and roughly quadruples training time. Would you still pick C5 here? For a general translation task?

   Here, yes — bit-acc jumped from 0.71 to 0.97 and sequence accuracy from 0.4% to 36%, which is a categorical improvement even at 4× cost. For a general translation task, probably not: BPE fits natural text well, the tokeniser's inductive bias is helpful rather than harmful, and paying 4× training cost for a small (or zero) quality gain is bad economics. The lesson is that token-free wins when your tokeniser's bias is fighting the data.

4. If you had one more configuration slot, what would you change from C1, and why?

   Three good candidates: (a) SwiGLU FFN — swap the GELU MLP for the Shazeer/Llama-style gated linear unit; consistent small quality wins on LM tasks at negligible cost. (b) Deeper vs wider at fixed params — e.g., 6 layers of d_model=192 vs 4 × 256; the crossover point between depth and width is worth measuring. (c) Learned absolute PE — a nn.Embedding of position ids, which would give a proper reference for how much of C1's shortfall is "fixed" absolute PE vs the general absolute-PE class. I would pick (c) here because it isolates a knob directly relevant to the task's positional structure.

5. You reported peak GPU memory for C1 and C5. What memory does that measurement include and exclude?

   torch.cuda.max_memory_allocated reports peak allocator-tracked memory: parameters, gradients, optimiser state, activations, and workspace tensors PyTorch created. It excludes CUDA context overhead (roughly 500 MB), memory held by other processes on the same device, and any fragmentation held by cudaMalloc but not currently allocated. So the reported figure is a lower bound on what nvidia-smi shows.

6. Which of the five configurations has the fewest parameters? Why?

   C3 (GQA), at 8.89M. GQA shrinks the K and V projections from d_model × d_model to d_model × (n_kv_heads · head_dim), which at 2 KV heads is a quarter of the parameters — saved four times per attention layer, across all encoder and decoder layers. C4 (RMSNorm) is the runner-up because it drops the LayerNorm β bias, but that saves only d_model params per norm layer, which is a much smaller absolute number.