Dinesh/ Blog
learning

R22 · Week 22 Recall & Drill

Week 22 revision: exact toy BPE, attention variance and tensor shapes, permutation equivariance, positional mechanisms and fixed-width versus KV-cache budgets.

40 min readChecking device speech…
Lesson preparation & details

Level: intermediate

By the end, you should be able to

  • Execute deterministic symbol-level BPE and distinguish vocabulary coverage, byte fallback and tokenizer versions.
  • Explain encoder bottlenecks, attention routing and the boundary between parallel prefill and sequential generation.
  • Derive square-root attention scaling under explicit independence and variance assumptions.
  • Trace Q/K/V and multi-head shapes, masks and key-axis normalization.
  • Compare fixed-width projection budgets with MHA/GQA/MQA cache and execution tradeoffs.
  • Prove attention permutation equivariance and RoPE relative-dot behavior without assuming length extrapolation.

Bring with you

Sources checked: . A source check does not change the publication or update date.

Editorial review: · What review means

In this article · 17 sections

Assessment scope

Tokenization defines positions; attention routes information between those positions; heads partition routing subspaces; positional mechanisms make order available. This revision includes actual token-sequence merges, attention outputs, permutation checks and a KV-cache budget. It does not infer model quality from a shape table or present a toy tokenizer as a vendor billing tokenizer.

Sessions covered

90-min structure

BlockMinutesWork
Warm-up5Trace text → tokens → Q/K/V → output.
Reconstruction30Derive assumptions and tensor shapes.
Drill30Execute merges, counts, attention and budget checks.
Quiz15Separate routing, attribution and memory.
Gap analysis10Retest a mask or tokenizer boundary.

The time split is pacing guidance; it does not predict the runtime of a trained model.

Blank-page reconstruction · 30 min

S106 · Tokenization

  1. Compare character, word and subword vocabularies.
  2. Walk BPE training and inference, and distinguish neighboring algorithms.
  3. Explain byte fallback and vocabulary edits.

Solutions. Character/byte units bound vocabulary but can produce long sequences; word units shorten familiar words but need a large vocabulary and a policy for unseen words. Subwords allocate reusable pieces, balancing vocabulary and length. They do not intrinsically guarantee every possible input is representable: that depends on normalization, base alphabet, unknown-token handling and byte fallback.

In the word-segmentation BPE version of Sennrich et al., start with characters plus an end-of-word marker, count adjacent symbol pairs weighted by word frequency, merge a most-frequent pair, and repeat to a vocabulary budget. Store a deterministic tie rule and merge order. Inference applies the learned merges to the same base representation; it does not retrain on each request. Merges must operate on adjacent symbols, not arbitrary substrings inside an already merged token.

WordPiece is a related subword method with its own vocabulary construction and segmentation rules; it is not universally “BPE with a different name.” SentencePiece is a tokenizer framework that can train directly on raw sentences and supports different model families such as BPE and unigram. A unigram model scores segmentations over a learned vocabulary instead of using the same deterministic pair-merge procedure. Keep normalization, whitespace treatment, special tokens, vocabulary and merge/model version together.

A byte-level alphabet or configured byte fallback can represent otherwise unknown UTF-8 text as bytes, at the cost of more tokens; the text normalizer may still alter the original string. Adding a token changes segmentation and requires a model embedding/output entry. Its initialization is implementation-specific (not necessarily random); the pretrained model has not learned the intended new token semantics merely because the vocabulary entry exists. Version caches by tokenizer/model, resize the relevant matrices correctly and evaluate adaptation before deployment.

S107 · Attention Intuition

  1. What bottleneck did early encoder-decoder attention address?
  2. Explain soft dictionary lookup and its limits as an explanation.
  3. Which sequence work parallelizes, and which remains sequential?

Solutions. A fixed-vector encoder-decoder compressed an entire source into one state. Attention allowed a decoder to select a weighted combination of source states relevant to its current query, rather than relying only on that bottleneck. This did not mean RNNs universally “failed”; it changed access paths and performance tradeoffs for particular tasks.

A query represents the current information request; keys define compatibility scores; values are the content combined with normalized weights. For fixed weights, the output is a linear combination of values. A large weight on a zero value can contribute zero, and residual paths can carry influential information around attention. Attention maps are useful diagnostics but not automatically causal attributions. Controlled input/value interventions and the whole computation are needed to investigate influence.

For a known input/prefix in training or prefill, attention computes position-pair scores in batched matrix operations instead of waiting on a recurrent hidden state. Autoregressive generation still must generate the next token before knowing the following token's input. Standard dense attention uses quadratic pairwise arithmetic and, if materialized, a quadratic score matrix; memory-efficient exact implementations can avoid storing that full matrix without removing all pairwise work.

S108 · Q/K/V Math

  1. Write scaled dot-product attention with shapes and a mask.
  2. Derive the square-root factor with assumptions.
  3. What does softmax over the wrong axis break, and are Q/K meanings unique?

Solutions. In one head, Q is (B,Tq,dk), K (B,Tk,dk), V (B,Tk,dv). Scores Q @ K.swapaxes(-1,-2) have shape (B,Tq,Tk). Add a mask of zero for allowed pairs and negative infinity for forbidden ones before softmax. Normalize each query row across key positions and multiply by V to obtain (B,Tq,dv):

Attention⁡(Q,K,V)=softmax⁡keys ⁣(QK⊤dk+M)V.\operatorname{Attention}(Q,K,V)=\operatorname{softmax}_{\mathrm{keys}}\!\left(\frac{QK^\top}{\sqrt{d_k}}+M\right)V.

All-masked rows require an explicit policy; blindly softmaxing negative infinity minus negative infinity can produce NaNs. The lab rejects them. Normalizing across queries would make keys distribute mass over queries, not give each query a unit-mass lookup over values.

If all q_i and k_i are mutually independent, zero mean and unit variance, E[q_i*k_i]=0, Var(q_i*k_i)=1, and independent coordinate products have zero covariance. Thus Var(sum(q_i*k_i))=dk; dividing by sqrt(dk) gives variance 1. Learned activations need not satisfy those assumptions exactly, so this is an initialization-scale rationale, not an invariant for every trained head. Dividing by dk instead makes variance 1/dk. Large score gaps can saturate softmax and reduce its sensitivity; near-one-hot heads may also be legitimate learned behavior, not usually provable bugs from a heatmap alone.

For plain self-attention without position-dependent transforms, scores contain X Wq Wk.T X.T, a factorized bilinear form. Reparameterizing Wq/Wk can preserve that product, so “query means search, key means identity” is intuition, not unique semantics. Separate projections retain rank/compute structure, generalize to cross-attention, and with RoPE the position-dependent rotations enter between them; the simple single-matrix statement no longer tells the whole story. V supplies the mixed content, but later output/residual operations also determine influence.

S109 · Multi-Head Attention

  1. Trace a fixed-width multi-head block.
  2. Explain the output projection and parameter budget.
  3. Distinguish MHA, GQA and MQA.

Solutions. For (B,T,D)=(2,4,12) and h=3 with dk=dv=4, dense Q/K/V projections each yield (2,4,12), reshaped/transposed to (2,3,4,4). Scores are (2,3,4,4) where the last two axes are query/key lengths; head outputs have the same shape here, concatenate to (2,4,12), then project through Wo (12,12).

With D fixed, h dividing D and square dense projections without biases, Wq/Wk/Wv/Wo contain 4*D² parameters. Leading dense attention arithmetic across all heads is proportional to T²*D, not h*T²*D with head width ignored. But score storage if materialized is h*T², and kernel overhead/efficiency can change. Holding parameter count fixed does not prove equal expressive capacity: different heads allow distinct routing distributions while smaller head width changes their constraints. There is no universal “head width below 16 is useless” boundary.

Wo learns combinations across concatenated head channels. Omitting it does not make all computation undefined; it changes the architecture and removes that learned mixing at this point. Standard MHA has one K/V head per query head. MQA shares one K/V head among all query heads; GQA uses an intermediate number of K/V groups. Query count and head width can remain fixed while K/V projection sizes/cache shrink. If there are h query heads, g KV heads and shared head width d=D/h, the bias-free projections contain D² + 2*D*g*d + D² = 2*D² + 2*D*g*d weights: Q and output remain square, but K/V each project to g*d. At D=12, h=3 and g=1 this is 384 weights, versus 576 for g=3 MHA. The KV cache shrinks by three here, but projection parameters shrink by only one third; these are different budgets. The earlier 4*D² claim applies to the specified full-MHA layout, not unchanged to GQA/MQA. Quality, bandwidth and parallel placement require evaluation; changing grouping in a pretrained model is not a free configuration toggle.

S110 · Positional Encoding

  1. Prove permutation equivariance, distinguishing invariance.
  2. Explain sinusoidal, learned and rotary positions.
  3. Why rotate Q/K, and why does definition beyond training length not prove quality?

Solutions. Without position-dependent information/masks, let P permute input positions. Q, K and V become PQ, PK, PV; scores become P S P.T. Row-wise softmax commutes with this simultaneous permutation, and softmax(P S P.T) PV = P softmax(S)V. The output is correspondingly permuted: equivariance. Pooling over positions can then be invariant. Causal/padding masks carry structure; a fixed causal mask does not commute with an arbitrary permutation, so “every transformer without positional embeddings is just bag-of-words” is too broad.

Sinusoidal position vectors use sin/cos frequencies and are added to token embeddings; learned absolute tables assign trainable vectors to a finite supported index range. A table can be larger than the lengths actually trained, leaving defined but poorly trained entries; beyond allocated size it needs extension/interpolation or a different mechanism. A formula being defined for every index does not guarantee generalization to unseen positions or distances.

In standard RoPE, rotate pairs of Q/K coordinates by position-dependent angles. For a 2D rotation R, (R_m q).T (R_n k) = q.T R_(n-m) k, so the dot product depends on relative displacement. Values need not rotate for this score mechanism; variants may differ. Unseen distance/frequency patterns and changed attention distributions can degrade long-context behavior. Test tasks over target lengths and positions rather than inferring success from a formula or declaring every extrapolation must fail.

Hands-on drill · 30 min

Run the four fences in order. The word BPE and byte BPE are deliberately small, fully specified tokenizers. All attention/cache results are local numerical or arithmetic checks, not LLM accuracy or serving benchmarks.

Step 1 — the merge algorithm and the token ratio (8 min)

from collections import Counter
corpus="low low low low low lower lower newest newest newest widest widest"
vocab=Counter({tuple(word)+('</w>',):count for word,count in Counter(corpus.split()).items()})
def pair_counts(vocab):
    counts=Counter()
    for symbols,freq in vocab.items():
        for pair in zip(symbols,symbols[1:]):counts[pair]+=freq
    return counts
def merge_symbols(symbols,pair):
    out=[];i=0
    while i<len(symbols):
        if i+1<len(symbols) and (symbols[i],symbols[i+1])==pair:
            out.append(symbols[i]+symbols[i+1]);i+=2
        else:out.append(symbols[i]);i+=1
    return tuple(out)
merges=[]
for step in range(6):
    counts=pair_counts(vocab)
    best=min(counts,key=lambda pair:(-counts[pair],pair))  # lexical tie break
    merges.append(best);new=Counter()
    for symbols,freq in vocab.items():new[merge_symbols(symbols,best)]+=freq
    vocab=new
    print("merge",step+1,best,"weighted pair count",counts[best])
assert sum(vocab.values())==12
assert merge_symbols(('ab','c','a','b'),('a','b'))==('ab','c','ab')
for symbols,count in vocab.items():print(symbols,count)

The first maximum-frequency pair is ('l','o') with count 7 under this lexical tie rule (other equally frequent pairs can yield a different valid trace). Sequence merging avoids accidentally merging substrings across symbol boundaries. All word frequencies survive; end-of-word markers prevent ambiguity when reconstructing words. Six merges cannot establish the best vocabulary size or performance on another language.

Step 2 — tokens per word is a property of your data (7 min)

# A byte-level BPE tokenizer: all 256 bytes are representable; no normalization or special tokens.
training=["the river and the road", "the quick brown fox", "the river flows"]
byte_vocab=Counter({tuple(bytes([b]) for b in text.encode('utf-8')):1 for text in training})
byte_merges=[]
for _ in range(12):
    counts=pair_counts(byte_vocab)
    if not counts:break
    pair=min(counts,key=lambda pair:(-counts[pair],pair));byte_merges.append(pair)
    new=Counter()
    for symbols,freq in byte_vocab.items():new[merge_symbols(symbols,pair)]+=freq
    byte_vocab=new
def tokenize(text):
    symbols=tuple(bytes([b]) for b in text.encode('utf-8'))
    for pair in byte_merges:symbols=merge_symbols(symbols,pair)
    return symbols
samples={'prose':'the river and the road','json':'{"userId":"a3f9-2c1b","amountCents":12995}',
         'code':'def get_user(id): return db.get(id)','identifier':'550e8400-e29b-41d4-a716-446655440000',
         'non-Latin':'నమస్కారం ప్రపంచం'}
for label,text in samples.items():
    tokens=tokenize(text);assert b''.join(tokens).decode('utf-8')==text
    words=len(text.split());ratio=len(tokens)/words
    print(label,'whitespace words',words,'TOY byte-BPE tokens',len(tokens),'ratio',round(ratio,3),
          'illustrative words per 8192 tokens',int(8192/ratio))
assert tokenize('')==()

This replaces the regex proxy with an actual, though tiny, tokenizer. The count is correct for this toy merge table, not for an OpenAI, SentencePiece or other deployed model. The whitespace-word denominator is itself language-dependent. For real context/cost budgeting, use the exact tokenizer/version on representative data and include chat templates, system text, tools, retrieved context and reserved generation tokens. The 8192 division is arithmetic under a constant measured ratio, not a guarantee that a concatenation will retain that ratio.

Step 3 — why the scaling constant is a square root (8 min)

import numpy as np
from itertools import product
rng=np.random.default_rng(0)
def softmax(z):
    e=np.exp(z-z.max(axis=-1,keepdims=True));return e/e.sum(axis=-1,keepdims=True)
for d in [8,64,512]:
    q=rng.normal(size=(2000,d));k=rng.normal(size=(2000,d));dots=(q*k).sum(-1)
    scores=dots.reshape(100,20)
    print('width',d,'empirical variance/d',float(dots.var()/d),
          'mean max weight raw/scaled',float(softmax(scores).max(-1).mean()),float(softmax(scores/np.sqrt(d)).max(-1).mean()))
# Independent exact finite-distribution oracle: all independent +/-1 components for d=2.
products=[sum(qi*ki for qi,ki in zip(q,k)) for q in product([-1,1],repeat=2) for k in product([-1,1],repeat=2)]
assert np.mean(products)==0 and np.var(products)==2
assert np.isclose(np.var(np.array(products)/np.sqrt(2)),1)
def attention(Q,K,V,allowed=None):
    S=Q@K.T/np.sqrt(Q.shape[1])
    if allowed is not None:
        if not allowed.any(axis=1).all():raise ValueError('all-masked query')
        S=np.where(allowed,S,-np.inf)
    P=softmax(S)
    return P@V,P
Q=np.array([[1.,0.],[0.,1.],[1.,1.]])
K=Q.copy();V=np.array([[1.,2.],[3.,4.],[5.,6.]])
out,P=attention(Q,K,V);perm=np.array([2,0,1])
permuted,_=attention(Q[perm],K[perm],V[perm])
assert np.allclose(permuted,out[perm]) and np.allclose(P.sum(1),1)
causal=np.tril(np.ones((3,3),dtype=bool))
co,cp=attention(Q,K,V,causal)
assert np.array_equal(cp[np.triu_indices(3,1)],np.zeros(3))
assert np.array_equal(co[0],V[0])
def rotation(angle):
    return np.array([[np.cos(angle),-np.sin(angle)],[np.sin(angle),np.cos(angle)]])
q=np.array([2.,1.]);k=np.array([-1.,3.]);m=.2;n=.7
assert np.isclose((rotation(m)@q)@(rotation(n)@k),q@(rotation(n-m)@k))
print('attention row sums',P.sum(1),'causal first output',co[0],'RoPE relative-dot identity checked')

The empirical variances fluctuate; the finite sign enumeration proves the stated variance for its distribution exactly. The causal example includes the current position and blocks future positions; an all-masked row is rejected. The permutation test is equivariance, not identical output at unchanged positions. These checks establish the toy operator, not the interpretability of trained attention maps.

Step 4 — head budgets and cache size (7 min)

D=512
for heads in [1,2,8,32,128]:
    assert D%heads==0
    width=D//heads;params=4*D*D
    assert heads*width==D
    print('heads',heads,'width',width,'dense projection parameters no biases',params)
def cache_bytes(layers,kv_heads,width,seq,batch=1,bytes_per=2):
    if min(layers,kv_heads,width,seq,batch,bytes_per)<0:raise ValueError('negative shape')
    return 2*layers*kv_heads*width*seq*batch*bytes_per
for seq in [2048,32768]:
    full=cache_bytes(80,64,128,seq);grouped=cache_bytes(80,8,128,seq)
    assert full==8*grouped
    print('length',seq,'MHA bytes',full,'GQA bytes',grouped,'decimal GB',full/1e9,grouped/1e9)
small=np.zeros((2,3,4,5,6),dtype=np.float32) # K/V, layers, heads, sequence, width
assert cache_bytes(3,4,6,5,bytes_per=4)==small.nbytes
assert cache_bytes(3,4,6,0,bytes_per=4)==0

The 80-layer, 64-KV-head example stores 5,368,709,120 bytes at length 2048 and 85,899,345,920 at 32768, before metadata, allocator overhead, padding, quantization scales or sharding/replication. Eight KV groups reduce these particular counts by eight. These are decimal GB when divided by 1e9; divide by 2³⁰ for GiB. They do not include weights, temporary activations or request scheduling, and do not show that a particular service fits a GPU.

Answer constraints and corrections

Misconception: “An attention heatmap explains the prediction.” It shows routing over value vectors in one part of a larger computation. A large weight can multiply negligible content; residual and other heads/layers can carry information. Use it to inspect masks/scaling, not as standalone causal evidence.

Quiz answers: (1) A token is a unit in a particular tokenizer, not a universal word fraction; measure the exact model's data distribution and framing. (2) Known positions can be processed in parallel, whereas future generated tokens are unknown; exact fused attention changes storage/IO, not all dense pairwise arithmetic. (3) Independence, zero means and unit variances yield dot-product variance dk, hence sqrt(dk) normalization. (4) More heads at fixed D keep this projection parameter count fixed but alter routing and implementation costs; capability is an empirical architectural question. (5) During generation account separately for weights, per-layer K/V cache, activations and metadata. Grouped K/V heads, supported cache quantization, window/eviction policies and bounded concurrent context can reduce memory, with quality/implementation costs. Memory-efficient attention avoids the full score matrix but does not eliminate retained K/V. Benchmark the intended length/batch distribution; this lesson performs no GPU serving benchmark.

Gap analysis + next week preview · 10 min

Retest an unseen Unicode string, an all-masked attention row, a permutation under a fixed causal mask or a cache estimate in GiB. State which assumption changes the answer. The next sessions assemble the transformer block, compare model families, study decoding and scaling, and distinguish exact from approximate efficient attention.

Corrective source notes

Primary references checked 2026-10-07: Sennrich et al. BPE, SentencePiece, encoder-decoder attention, Attention Is All You Need v7, RoFormer/RoPE, GQA v3, Transformers 4.57.1 KV cache, and FlashAttention. The finite operator identities and byte counts are derived and executed here; no trained tokenizer quality, attention attribution or long-context quality result is claimed.

Additional tokenizer detail: WordPiece implementation explanation. The rotation equations are checked against RoFormer v5 full text.

Part of the 6-month evergreen learning plan.

Pause / Recall / Apply

Can you explain it without the page?

Close the example. Reconstruct the core idea, then change one assumption. Mark complete when you’re ready; you can always undo it.

Stored in this browser only. No account, no sync. Clearing browser data removes your record.