Transformer Math and Implementation Deep Dive

This note is the bridge from "I understand the diagram" to "I can implement, debug, and reason about the block."

Primary sources:

• Attention Is All You Need
• RoFormer / RoPE
• RMSNorm
• GLU Variants Improve Transformer

Decoder-Only Forward Pass

Input token ids:

idx: [B, T]

Embedding lookup:

x = token_embedding(idx)
x: [B, T, d_model]

Then for each layer:

x = x + attention(norm1(x))
x = x + mlp(norm2(x))

Finally:

logits = x @ W_vocab.T
logits: [B, T, vocab_size]

Attention Shapes

Given normalized hidden states:

h: [B, T, d_model]

Project:

q = h @ Wq
k = h @ Wk
v = h @ Wv

Reshape:

q: [B, n_q_heads, T, d_head]
k: [B, n_kv_heads, T, d_head]
v: [B, n_kv_heads, T, d_head]

If using standard MHA:

n_q_heads == n_kv_heads

If using GQA:

n_q_heads > n_kv_heads

The implementation repeats or maps KV heads across query head groups.

Attention Equation

For each head:

scores = q @ k.transpose(-2, -1) / sqrt(d_head)
scores = causal_mask(scores)
weights = softmax(scores)
out = weights @ v

Then:

out: [B, n_heads, T, d_head]
out -> [B, T, d_model]
out = out @ Wo

Why Scale By sqrt(d_head)

The dot product of random vectors grows in variance with dimensionality. Dividing by sqrt(d_head) keeps logits in a softmax-friendly range.

Without this:

• Softmax can saturate.
• Gradients become poor.
• Training becomes less stable.

Causal Masking

The causal mask sets future-token scores to negative infinity before softmax:

scores[i, j] = -inf if j > i

This preserves the autoregressive factorization:

P(x_1, ..., x_T) = product_t P(x_t | x_<t)

RoPE Intuition

RoPE rotates query and key vectors by an angle determined by token position. The dot product between rotated query/key vectors naturally encodes relative position.

Research-level intuition:

• Position is injected into attention matching, not added as a separate token embedding.
• Relative distances affect attention scores.
• Extrapolation depends on frequency choices and training distribution.

RMSNorm

RMSNorm normalizes by root mean square without subtracting the mean.

Simplified:

rms(x) = sqrt(mean(x^2) + eps)
y = scale * x / rms(x)

Why it matters:

• Simpler than LayerNorm.
• Often efficient and stable.
• Common in modern LLMs.

SwiGLU / Gated MLP

A common modern MLP:

gate = SiLU(x @ W_gate)
up = x @ W_up
hidden = gate * up
out = hidden @ W_down

Why gated MLPs matter:

• More expressive than a plain activation MLP.
• Strong empirical performance.
• MLP sublayers hold much of the model's parameter capacity.
• MoE often replaces this MLP with multiple routed expert MLPs.

Minimal Pseudocode

def block(x):
    x = x + attention(rmsnorm(x))
    x = x + swiglu_mlp(rmsnorm(x))
    return x

def attention(h):
    q = project_q(h)
    k = project_k(h)
    v = project_v(h)
    q, k = apply_rope(q, k)
    scores = q @ k.transpose(-2, -1) / sqrt(d_head)
    scores = apply_causal_mask(scores)
    weights = softmax(scores)
    y = weights @ v
    return project_out(y)

Implementation Bugs To Expect

• Wrong tensor layout after reshape/transpose.
• Mask broadcast shape wrong.
• Softmax over the wrong dimension.
• RoPE applied to values instead of query/key.
• KV cache concatenated along the wrong axis.
• GQA head repetition wrong.
• Loss computed on unshifted labels.
• Dropout accidentally enabled during generation.
• Sampling from logits before temperature/top-p transforms.

What To Implement

1. Token embedding and tied output embedding.
2. RMSNorm.
3. Causal attention.
4. RoPE.
5. SwiGLU.
6. KV cache.
7. GQA.
8. Sampling: temperature, top-k, top-p.
9. Tiny training loop.
10. Evaluation loss and perplexity.

How This Connects

• Transformer Block Anatomy gives the conceptual layout.
• Attention Mechanics and KV Cache explains serving-time memory.
• Training Optimization and Stability Deep Dive explains why the block trains.
• Mixture of Experts Architectures modifies the MLP side.
• Long Context and Efficient Sequence Models modifies the sequence side.