The Transformer

Attention as a dynamic content-based lookup

In feedforward layers, weights are static connections learned during training. In self-attention, the weights connecting tokens are computed dynamically on the fly based on token compatibility. Each input token representation $x_{i} \in ℝ^{d}$ is projected into three subspaces:

Query (Q): what information token i is searching for.
Key (K): what features or context token j advertises.
Value (V): the actual content payload transmitted if attended to.

The matrix formulation evaluates all pairs simultaneously:

Attention (Q, K, V) = softmax (\frac{Q K^{T}}{\sqrt{d_{k}}} + M) V

Enable JavaScript to use this local playground. The explanation and sources below work without it.

Tensor shape progression across a single layer

For an input sequence of length T with model dimension d and H attention heads (head dimension $d_{k} = d / H$ ), shapes progress through the attention mechanism as follows:

Tensor	Expression	Shape	Description
`X`	Input embeddings	[`T`, `d`]	Batch of `T` token vectors of width `d`
`W`_Q, `W`_K, `W`_V	Projection weights	[`d`, `d`]	Learned linear projection matrices
`Q`_h, `K`_h, `V`_h	$X W_{Q, K, V}^{(h)}$	[`T`, `d`_k]	Per-head projected queries, keys, and values
`S`_h	$Q_{h} K_{h}^{T} / \sqrt{d_{k}} + M$	[`T`, `T`]	Pairwise scaled compatibility scores
`P`_h	$softmax (S_{h})$	[`T`, `T`]	Row-normalized attention probability weights (∑ⱼ `P`_ij = 1)
`O`_h	$P_{h} V_{h}$	[`T`, `d`_k]	Weighted combination of values for head `h`
`O`_concat	$[O_{1}, \dots, O_{H}]$	[`T`, `d`]	Concatenated outputs across all `H` heads
`Y`	$O_{concat} W_{O}$	[`T`, `d`]	Final linear output projection

Variance derivation: why divide by √d_k?

Assume components of query q and key k are independent, zero-mean, unit-variance random variables: $E [q_{i}] = E [k_{i}] = 0$ and $Var (q_{i}) = Var (k_{i}) = 1$ . The dot product is $q \cdot k = \sum_{i = 1}^{d_{k}} q_{i} k_{i}$ .

For each term, independence gives:

E [q_{i} k_{i}] = E [q_{i}] E [k_{i}] = 0

Var (q_{i} k_{i}) = E [q_{i}^{2} k_{i}^{2}] - {(E [q_{i} k_{i}])}^{2} = E [q_{i}^{2}] E [k_{i}^{2}] = 1 \cdot 1 = 1

Because the components are mutually independent, the variance of the sum is the sum of the variances:

Var (q \cdot k) = \sum_{i = 1}^{d_{k}} Var (q_{i} k_{i}) = d_{k}, StdDev (q \cdot k) = \sqrt{d_{k}}

As head dimension $d_{k}$ increases (e.g. $d_{k} = 128$ ), unscaled dot products have standard deviation $\sqrt{128} \approx 11.3$ . Dividing by $\sqrt{d_{k}}$ scales variance back to 1.0:

Var (\frac{q \cdot k}{\sqrt{d_{k}}}) = \frac{1}{d_{k}} Var (q \cdot k) = \frac{d_{k}}{d_{k}} = 1

The Softmax Jacobian and gradient saturation

Why does unscaled variance matter? The derivative of the row-wise softmax $P_{i j} = e^{S_{i j}} / \sum_{k} e^{S_{i k}}$ with respect to input score $S_{i m}$ is:

\frac{\partial P_{i j}}{\partial S_{i m}} = P_{i j} (δ_{j m} - P_{i m}) = {\begin{matrix} P_{i j} (1 - P_{i j}) & j = m \\ - P_{i j} P_{i m} & j \neq m \end{matrix}

When scores $S_{i j}$ have large variance, the largest score completely dominates the exponential sum, driving its probability $P_{i j} \to 1$ while all other probabilities $P_{i m} \to 0$ . At $P_{i j} \approx 1$ , $P_{i j} (1 - P_{i j}) \approx 0$ . At $P_{i m} \approx 0$ , $- P_{i j} P_{i m} \approx 0$ . The entire Jacobian matrix approaches zero, completely vanishing the gradient signal flowing back into Q and K. Scaling by $1 / \sqrt{d_{k}}$ prevents this saturation trap.

Causal masking

In autoregressive generation (GPT-style decoders), token i must not attend to future tokens $j > i$ . An additive attention mask M is applied to scores before softmax:

M_{i j} = {\begin{matrix} 0 & j \leq i \\ - \infty & j > i \end{matrix}

Causal Mask M (T = 4):                Softmax Attention Weights P:
     j=0   j=1   j=2   j=3                  j=0   j=1   j=2   j=3
i=0 ┌  0   -∞    -∞    -∞  ┐           i=0 ┌ 1.0    0     0     0  ┐
i=1 │  0    0    -∞    -∞  │    ──▶    i=1 │ p10   p11    0     0  │
i=2 │  0    0     0    -∞  │           i=2 │ p20   p21   p22    0  │
i=3 └  0    0     0     0  ┘           i=3 └ p30   p31   p32   p33 ┘

Because $e^{- \infty} = 0$ , future entries receive identically zero attention probability weight, preventing information leakage from future tokens during parallel teacher-forced training.

Layer FLOPs: when does attention dominate the FFN?

Consider a standard Transformer layer with sequence length T, model dimension d, and feedforward intermediate dimension $4 d$ . Counting multiply-accumulate operations (2 FLOPs per MAC):

Q, K, V projections: $3 \cdot 2 T d^{2} = 6 T d^{2}$
Attention score matrix ( $Q K^{T}$ ): $2 T^{2} d$
Value aggregation ( $P V$ ): $2 T^{2} d$
Output projection ( $W_{O}$ ): $2 T d^{2}$
Feedforward layer ( $W_{1} \in ℝ^{d \times 4 d}, W_{2} \in ℝ^{4 d \times d}$ ): $2 \cdot 2 T (d \cdot 4 d) = 16 T d^{2}$

Total layer FLOPs:

{FLOPs}_{layer} = \underset{O (T^{2}) attention}{\underset{⏟}{4 T^{2} d}} + \underset{O (T) linear projections}{\underset{⏟}{24 T d^{2}}}

Attention FLOPs exceed linear projection FLOPs when:

4 T^{2} d > 24 T d^{2} \Rightarrow T > 6 d

For standard 7B/8B models where d = 4096, attention compute only dominates once sequence length $T > 6 \times 4096 = 24 576$ tokens. At short contexts ( $T \leq 4096$ ), over 85% of total compute is spent in the feedforward projections.

The Pre-Norm Transformer block

Original Transformers placed LayerNorm after the residual addition (Post-LN: $x = LN (x + SubLayer (x))$ ), which suffered from gradient vanishing and required warmup learning rate schedules. Modern models adopt Pre-LN (or Pre-RMSNorm):

            x ──┬─────────────────────────────(+)──▶
                │                              ▲
                ▼                              │
             RMSNorm ──▶ Attention ──▶ W_O ────┘
                │
                ├──┬──────────────────────────(+)──▶ output
                │  │                           ▲
                │  ▼                           │
                │ RMSNorm ──▶ FFN (SwiGLU) ────┘
                ▼

x_{mid} = x + Attention (RMSNorm (x))

x_{out} = x_{mid} + FFN (RMSNorm (x_{mid}))

In Pre-LN, the identity highway $x + \dots$ is uninterrupted across the entire depth of the network. Gradients flow back from the loss directly to early token embeddings without attenuation.

Interactive Micro-Transformer Trainer

Train a 2-layer, 2-head Micro-Transformer on algorithmic sequence tasks (Sequence Reversal, Identity Copy, Sorting). Observe how attention weights self-organize along the anti-diagonal for sequence reversal, and inspect the impact of causal versus bidirectional masking.

Configuration & Controls

Algorithmic Task

Attention Mask M

Probe Input Sequence (4 Tokens from A–H)

Epoch

Token CE Loss

—

Token Accuracy

—

Multi-Head Attention Heatmaps (Head 1 & Head 2)

Training on Reverse causes attention to concentrate along the anti-diagonal (position 0 attends to position 3). Under a Causal Mask, position 0 cannot attend ahead, demonstrating why bidirectional context is required for non-autoregressive tasks.