Attention Is All You Need: a careful walkthrough of the Transformer

In June 2017, eight researchers at Google posted a paper whose title was also its thesis: Attention Is All You Need by Vaswani et al. It proposed the Transformer, a sequence-to-sequence model for machine translation that used no recurrence and no convolution. Every interaction between positions in a sentence went through attention. Almost every large language model in use today is a descendant of that architecture, usually with modifications the paper did not contain.

This article walks through the paper carefully: the problem it addressed, the attention formula and why it is scaled, multi-head attention, the encoder–decoder layout, positional encoding, the complexity argument, and the training recipe. Every numeric example comes from code shown here, and every reported result is attributed to the paper. A NumPy implementation of one encoder layer closes the technical part. If you want to train a small decoder-only model end to end in PyTorch afterwards, the language-model-from-scratch guide does that; this article stays closer to the original paper.

The problem: sequence-to-sequence with recurrence

Machine translation maps a source sequence $x_1,\ldots,x_n$ to a target sequence $y_1,\ldots,y_m$ of a different length. Before 2017, the strongest neural systems were encoder–decoder recurrent networks, usually built from LSTM or GRU cells. The encoder reads the source one token at a time and updates a hidden state:

The decoder then generates the target one token at a time, conditioned on what the encoder produced. Two problems follow from that recurrence.

The first is sequential computation. $h_t$ cannot be computed until $h_{t-1}$ exists, so a sentence of length $n$ requires $n$ dependent steps per layer. GPUs are good at large parallel matrix products and comparatively poor at long chains of small dependent ones. Batching across sentences helps, but memory limits how far that goes for long sequences.

The second is the path between distant positions. Information about $x_1$ reaches the representation of $x_{50}$ only after passing through 49 state updates. Gradients make the same journey backwards. LSTMs mitigate vanishing gradients, but a long path is still a harder learning problem than a short one.

Early encoder–decoders also compressed the whole source into a single fixed-size vector. Bahdanau, Cho and Bengio (2014) removed that bottleneck with attention: at each decoding step, the decoder computes a weighted average of all encoder states, with weights produced by a small learned scoring network. Luong, Pham and Manning (2015) studied simpler scoring functions, including a plain dot product. In both, attention was an addition to a recurrent network. The Transformer's move was to ask whether the recurrent part was necessary at all.

flowchart LR
  subgraph RNN["Recurrent encoder: n dependent steps"]
    R1["h1"] --> R2["h2"] --> R3["h3"] --> R4["h4"]
  end
  subgraph ATT["Self-attention layer: one parallel step"]
    A1["x1"] --> B1["z1"]
    A1 --> B2["z2"]
    A1 --> B3["z3"]
    A2["x2"] --> B1
    A2 --> B2
    A2 --> B3
    A3["x3"] --> B1
    A3 --> B2
    A3 --> B3
  end

In a self-attention layer, every output position reads every input position directly, and all outputs can be computed at once. The price, discussed later, is that the number of pairwise interactions grows with the square of the sequence length.

Scaled dot-product attention

The paper describes attention as mapping a query and a set of key–value pairs to an output. Stack $n_q$ queries as rows of $Q\in\mathbb R^{n_q\times d_k}$, $n_k$ keys as rows of $K\in\mathbb R^{n_k\times d_k}$, and the matching values as rows of $V\in\mathbb R^{n_k\times d_v}$. Then

Read it in four steps, tracking shapes:

1. Scores: $S = QK^\top$ has shape $n_q\times n_k$. Entry $S_{ij}=q_i\cdot k_j$ measures how well query $i$ matches key $j$.
2. Scale: divide every entry by $\sqrt{d_k}$. The shape is unchanged.
3. Normalize: apply softmax to each row, giving weights $A$ of shape $n_q\times n_k$ whose rows are non-negative and sum to one.
4. Mix: $AV$ has shape $n_q\times d_v$. Output row $i$ is a weighted average of the value rows, using row $i$ of $A$ as weights.

A query asks a question, keys advertise what each position can answer, and values carry the content returned. In self-attention all three are linear projections of the same sequence, so a position can retrieve information from any other position in a single operation. The whole computation is two matrix multiplications and a softmax, which the paper notes is faster and more memory-efficient in practice than additive (Bahdanau-style) attention, because it uses highly optimized matrix-multiplication code.

flowchart TD
  Q["Q: n_q x d_k"] --> MM1["MatMul: Q times K transpose"]
  K["K: n_k x d_k"] --> MM1
  MM1 --> SC["Scale by 1 over sqrt d_k"]
  SC --> MK["Optional mask: set disallowed scores to minus infinity"]
  MK --> SM["Row-wise softmax: n_q x n_k weights"]
  SM --> MM2["MatMul: weights times V"]
  V["V: n_k x d_v"] --> MM2
  MM2 --> O["Output: n_q x d_v"]

Why divide by the square root of d_k

Suppose the components of a query $q$ and a key $k$ are independent random variables with mean $0$ and variance $1$. Their dot product is

Each term has mean $\mathbb E[q_i]\,\mathbb E[k_i]=0$. Its variance is

The terms are independent, so their variances add:

This is the argument in the paper's footnote. Without scaling, a head with $d_k=64$ produces scores with a standard deviation of about $8$. Softmax is sensitive to differences between its inputs, not their absolute level, so scores spread over a range of tens push nearly all the weight onto the largest entry. In that saturated region the softmax gradient is tiny: the Jacobian $\operatorname{diag}(p)-pp^\top$ approaches zero when $p$ is nearly one-hot. Dividing by $\sqrt{d_k}$ keeps the score scale roughly independent of head width, so learning starts in a region where gradients flow.

A quick simulation confirms both the variance and the saturation:

import numpy as np

rng = np.random.default_rng(0)
for d_k in (4, 64, 512):
    q = rng.standard_normal((100_000, d_k))
    k = rng.standard_normal((100_000, d_k))
    dots = (q * k).sum(axis=1)
    print(f"d_k={d_k:4d}  var(q.k)={dots.var():8.2f}  "
          f"var(q.k/sqrt(d_k))={(dots / np.sqrt(d_k)).var():.3f}")

# What an unscaled score spread does to softmax
z = rng.standard_normal(10)
for s in (1.0, 8.0):
    p = np.exp(s * z - (s * z).max()); p /= p.sum()
    print(f"std scale {s:3.0f}: max weight {p.max():.4f}")

On this run the output was:

d_k=   4  var(q.k)=    4.03  var(q.k/sqrt(d_k))=1.007
d_k=  64  var(q.k)=   64.09  var(q.k/sqrt(d_k))=1.001
d_k= 512  var(q.k)=  514.18  var(q.k/sqrt(d_k))=1.004
std scale   1: max weight 0.3935
std scale   8: max weight 0.9997

Ten scores with unit spread give a largest weight of about $0.39$. Multiply the same scores by $8$, the standard deviation an unscaled $d_k=64$ head would produce, and the largest weight becomes $0.9997$. Real learned queries and keys are not independent unit-variance variables, so treat this as a statement about initialization and scale, not a theorem about trained models.

Masking with negative infinity

Sometimes a query must not see certain keys: padding tokens, or, in the decoder, future positions. The paper implements this inside the softmax by setting the disallowed scores to $-\infty$ before normalization. Since $e^{-\infty}=0$, those positions receive exactly zero weight, and the remaining weights still sum to one. Formally, add a mask matrix $M$ with $M_{ij}=0$ where attention is allowed and $M_{ij}=-\infty$ where it is not:

A causal mask for length $n$ allows $j\le i$ only, so $M$ is $-\infty$ strictly above the diagonal. Every row must keep at least one allowed entry; a row of all $-\infty$ would divide zero by zero.

A worked numeric example

Take three tokens with $d_k=d_v=2$. To keep the arithmetic readable, the queries and keys are identical here; in a real layer they come from different learned projections.

The raw scores are the dot products of every query row with every key row:

Dividing by $\sqrt 2\approx 1.4142$ gives entries $0$, $0.7071$ and $1.4142$. For the first row, the exponentials are $e^{0.7071}\approx 2.028$, $e^{0}=1$ and $2.028$ again, with sum $5.056$. The weights are therefore about $0.4011$, $0.1978$ and $0.4011$. The first output is

The third query matches the third key most strongly, with a scaled score of $1.4142$. Its weights are $0.2483$, $0.2483$ and $0.5035$, and its output $(3.5105,\ 4.5105)$ leans toward the third value row. The code reproduces every number:

import numpy as np

# Three tokens, d_k = d_v = 2. Rows are positions.
Q = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [1.0, 1.0]])
K = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [1.0, 1.0]])
V = np.array([[1.0, 2.0],
              [3.0, 4.0],
              [5.0, 6.0]])

d_k = Q.shape[-1]
scores = Q @ K.T                      # (3, 3)
scaled = scores / np.sqrt(d_k)        # (3, 3)
weights = np.exp(scaled - scaled.max(axis=-1, keepdims=True))
weights /= weights.sum(axis=-1, keepdims=True)
out = weights @ V                     # (3, 2)

np.set_printoptions(precision=4, suppress=True)
print("scores\n", scores)
print("scaled\n", scaled)
print("weights\n", weights)
print("row sums", weights.sum(axis=-1))
print("output\n", out)

scores
 [[1. 0. 1.]
 [0. 1. 1.]
 [1. 1. 2.]]
scaled
 [[0.7071 0.     0.7071]
 [0.     0.7071 0.7071]
 [0.7071 0.7071 1.4142]]
weights
 [[0.4011 0.1978 0.4011]
 [0.1978 0.4011 0.4011]
 [0.2483 0.2483 0.5035]]
row sums [1. 1. 1.]
output
 [[3.     4.    ]
 [3.4067 4.4067]
 [3.5105 4.5105]]

Notice that every output is a convex combination of value rows, so it lies inside their convex hull. Attention alone cannot produce a vector outside that region. The learned projections, the output projection, and the feed-forward layers supply the rest of the expressiveness.

Multi-head attention

A single attention distribution averages. If a token needs to attend to its syntactic head and, separately, to a coreferent noun, one set of weights must compromise between them. The paper's remedy is to run several attention functions in parallel on different learned projections:

The projection matrices have shapes $W_i^Q,W_i^K\in\mathbb R^{d_{\mathrm{model}}\times d_k}$, $W_i^V\in\mathbb R^{d_{\mathrm{model}}\times d_v}$ and $W^O\in\mathbb R^{hd_v\times d_{\mathrm{model}}}$. The base model uses $h=8$ heads with $d_{\mathrm{model}}=512$ and $d_k=d_v=d_{\mathrm{model}}/h=64$.

Because each head works in a 64-dimensional subspace, the total cost is similar to one head with full width. The projections for all heads together are still $d_{\mathrm{model}}\times d_{\mathrm{model}}$ matrices, and the score computation does $h$ products of size $n\times 64$ by $64\times n$, the same multiply count as one $n\times 512$ by $512\times n$ product. What changes is that there are $h$ separate softmaxes, so the model can keep $h$ distinct attention patterns per layer. In implementation, the eight $W_i^Q$ are stored as one $512\times 512$ matrix and the result is reshaped into heads.

The paper's ablations (Table 3, rows A) found that a single head was worse by 0.9 BLEU than the best setting on the development set, and that too many heads (32, each with $d_k=16$) also hurt. Heads sometimes learn interpretable patterns, and the paper's appendix shows examples, but later analysis found that many heads can be pruned with little loss. "Each head learns one linguistic role" is a hypothesis to check, not a design guarantee.

flowchart TD
  X["Input: n x 512"] --> P["Linear projections Q, K, V<br/>split into 8 heads of width 64"]
  P --> H1["Head 1: scaled dot-product attention"]
  P --> H2["Head 2: scaled dot-product attention"]
  P --> H3["..."]
  P --> H8["Head 8: scaled dot-product attention"]
  H1 --> C["Concat: n x 512"]
  H2 --> C
  H3 --> C
  H8 --> C
  C --> WO["Linear W_O: 512 x 512"]
  WO --> Y["Output: n x 512"]

The encoder–decoder architecture

The Transformer keeps the encoder–decoder shape of earlier translation models. The encoder maps the source tokens to a sequence of continuous representations $z_1,\ldots,z_n$. The decoder generates the target one token at a time, auto-regressively, consuming its own previous outputs.

The encoder is a stack of $N=6$ identical layers. Each layer has two sublayers: multi-head self-attention, then a position-wise feed-forward network. The decoder is also a stack of $N=6$ layers, each with three sublayers: masked self-attention, encoder–decoder attention over the encoder output, and a feed-forward network. All sublayers and embeddings produce vectors of width $d_{\mathrm{model}}=512$, which is what makes the residual additions possible.

flowchart TD
  subgraph ENC["Encoder: N = 6 layers"]
    SE["Source embedding x sqrt d_model<br/>plus positional encoding"] --> ESA["Multi-head self-attention<br/>Add and Norm"]
    ESA --> EFF["Feed-forward<br/>Add and Norm"]
  end
  subgraph DEC["Decoder: N = 6 layers"]
    TE["Shifted target embedding x sqrt d_model<br/>plus positional encoding"] --> DSA["Masked multi-head self-attention<br/>Add and Norm"]
    DSA --> DCA["Encoder-decoder attention<br/>Add and Norm"]
    DCA --> DFF["Feed-forward<br/>Add and Norm"]
  end
  EFF -->|"keys and values"| DCA
  DFF --> LIN["Linear, tied to embeddings"]
  LIN --> SMX["Softmax over vocabulary"]
  SMX --> OUT["Next-token probabilities"]

Three uses of attention

The same multi-head function appears in three roles, distinguished only by where $Q$, $K$ and $V$ come from and which mask is applied.

Encoder self-attention. $Q$, $K$ and $V$ all come from the previous encoder layer. Every source position can attend to every other source position, apart from padding. This builds context-dependent representations of the source sentence.

Masked decoder self-attention. $Q$, $K$ and $V$ come from the previous decoder layer, but position $i$ may attend only to positions $j\le i$. During training, the whole target sentence is fed in at once, shifted right by one position, and the causal mask prevents position $i$ from reading the token it is supposed to predict. This is what lets training be parallel across target positions while generation remains left to right.

Encoder–decoder attention. Queries come from the decoder; keys and values come from the final encoder output. Each target position can read any source position. This is the direct descendant of Bahdanau attention, now multi-headed and repeated in every decoder layer. Here $n_q$ is the target length and $n_k$ is the source length, so the score matrix is rectangular.

Feed-forward layers, residuals and normalization

Attention mixes information across positions. The feed-forward network then transforms each position separately and identically:

with $W_1\in\mathbb R^{512\times 2048}$ and $W_2\in\mathbb R^{2048\times 512}$, so $d_{ff}=2048$. The paper also describes it as two convolutions with kernel size 1. The parameters differ from layer to layer but are shared across positions within a layer. In the base model the two feed-forward matrices hold about 2.1 million weights per layer, roughly twice the attention projections' 1.05 million.

Every sublayer is wrapped in a residual connection followed by layer normalization:

The residual connection, from He et al. (2015), lets each sublayer learn a correction to its input instead of a full replacement, and gives gradients a direct path through a deep stack. Layer normalization (Ba, Kiros and Hinton, 2016) normalizes each position's vector to zero mean and unit variance across features, then applies a learned gain and bias. Dropout with rate $P_{drop}=0.1$ in the base model is applied to each sublayer output before it is added to the residual.

flowchart LR
  X["x"] --> SUB["Sublayer: attention or FFN"]
  SUB --> DR["Dropout"]
  X --> ADD["Add: x plus Sublayer output"]
  DR --> ADD
  ADD --> LN["LayerNorm"]
  LN --> Y["Output to next sublayer"]

This placement, normalization after the residual addition, is now called post-LN. Most later large models use pre-LN instead, $x+\operatorname{Sublayer}(\operatorname{LayerNorm}(x))$, which leaves an unnormalized residual stream running through the network. Xiong et al. (2020) analyzed the difference and argued that post-LN produces large gradients near the output at initialization, which is one reason it needs a learning-rate warmup, while pre-LN can often train with little or none. Neither choice is universally better; post-LN is what the paper used and what the code below implements.

Embeddings and the output layer

Input and output tokens are mapped to $d_{\mathrm{model}}$-dimensional vectors by learned embeddings. The decoder output is mapped to next-token probabilities by a linear layer and a softmax. The paper shares one weight matrix among the source embedding, the target embedding and the pre-softmax linear layer, following Press and Wolf. Sharing source and target embeddings requires a vocabulary common to both languages, which is what the paper's shared subword vocabulary for English–German provides. Inside the embedding layers, the shared weights are multiplied by $\sqrt{d_{\mathrm{model}}}$. One reasonable reading is that this keeps the embedding scale comparable to the positional encodings, whose entries lie in $[-1,1]$; the paper states the scaling without giving a reason.

Sinusoidal positional encoding

Self-attention is permutation-equivariant: shuffle the input positions and the outputs are shuffled in the same way, with no other change. Without extra information the model sees a bag of tokens. The paper adds a fixed positional encoding to the embeddings at the bottom of both stacks:

Each pair of dimensions $(2i,2i+1)$ is a sine and cosine at angular frequency $\omega_i=10000^{-2i/d_{\mathrm{model}}}$. The wavelengths form a geometric progression from $2\pi$ to $10000\cdot 2\pi$. Low dimensions change quickly with position; high dimensions change slowly, like the hands of a clock running at many speeds.

The paper chose this form because it hypothesized it would let the model attend by relative position: for any fixed offset $k$, $PE_{pos+k}$ is a linear function of $PE_{pos}$. For one frequency pair the proof is the angle-addition formula. Write $a=\omega_i\,pos$ and $b=\omega_i k$. Then

The matrix depends on the offset $k$ and the frequency, but not on $pos$. Stacking one such $2\times 2$ rotation per frequency gives a block-diagonal matrix $R_k$ with $PE_{pos+k}=R_k\,PE_{pos}$ for every position. This shows that relative offsets are representable by a fixed linear map; it does not prove the model learns to use them. The check below confirms the identity numerically:

import numpy as np

d_model, k = 512, 7
i = 5                                   # which sin/cos pair
w = 1.0 / 10000 ** (2 * i / d_model)
R = np.array([[ np.cos(w * k), np.sin(w * k)],
              [-np.sin(w * k), np.cos(w * k)]])

for pos in (0, 3, 40):
    pe_pos  = np.array([np.sin(w * pos),       np.cos(w * pos)])
    pe_next = np.array([np.sin(w * (pos + k)), np.cos(w * (pos + k))])
    print(pos, np.allclose(R @ pe_pos, pe_next))

It prints True for each position. The paper also tried learned positional embeddings and found nearly identical results (Table 3, row E), and it preferred sinusoids partly because they might extrapolate to sequences longer than those seen in training. Later models have used learned absolute positions, relative position biases, and rotary embeddings, which apply exactly this kind of rotation to queries and keys instead of adding it to the input.

Why self-attention: the complexity comparison

Section 4 of the paper compares layer types that map a sequence of $n$ vectors of width $d$ to another such sequence. Table 1 summarizes three criteria; $k$ is a convolution kernel size and $r$ is the neighborhood size of restricted attention.

Complexity per layer counts arithmetic. Self-attention forms an $n\times n$ score matrix, each entry a $d$-dimensional dot product, giving $n^2 d$. A recurrent layer does a $d\times d$ matrix–vector product at each of $n$ steps, giving $n d^2$. Self-attention is cheaper when $n<d$. The paper notes that this is most often the case for the sentence representations used in machine translation at the time, such as word-piece and byte-pair encodings. This row counts only the attention mixing itself; the $Q$, $K$, $V$ and output projections and the feed-forward network add $O(n d^2)$ work in a Transformer layer too.

Sequential operations is the minimum number of steps that must happen one after another. A recurrent layer needs $n$; self-attention needs a constant number, so all positions are processed in parallel. This is the property that made training fast on GPUs.

Maximum path length is the longest chain of computations a signal must traverse between any two input and output positions. Shorter paths make long-range dependencies easier to learn. Self-attention connects every pair in one step. A stack of convolutions needs $O(n/k)$ layers with contiguous kernels, or $O(\log_k(n))$ with dilated ones. Restricted self-attention, attending only within a window of size $r$, trades a longer path for linear cost; the paper mentions it as future work for very long sequences.

The table also reveals the cost. At $n=512$ and $d=512$ the two terms are equal. At $n=8192$ the attention term is sixteen times the recurrent one, and the score matrix per head holds 67 million entries.

Training details and results from the paper

The paper evaluated on the WMT 2014 English–German dataset, about 4.5 million sentence pairs encoded with byte-pair encoding into a shared vocabulary of about 37,000 tokens, and on WMT 2014 English–French, 36 million sentences with a 32,000-token word-piece vocabulary. Batches contained approximately 25,000 source tokens and 25,000 target tokens, grouped by similar length. Training ran on one machine with eight NVIDIA P100 GPUs.

The optimizer was Adam with $\beta_1=0.9$, $\beta_2=0.98$ and $\epsilon=10^{-9}$. The learning rate followed a warmup-then-decay schedule:

with $\mathit{warmup}=4000$ steps. For the first 4000 steps the second term is smaller, so the rate grows linearly. After that, it decays proportionally to the inverse square root of the step number. The two terms are equal at the warmup step, which is the peak.

d_model, warmup = 512, 4000

def lrate(step):
    return d_model ** -0.5 * min(step ** -0.5, step * warmup ** -1.5)

for step in (1, 100, 1000, 4000, 10000, 50000, 100000):
    print(f"step {step:6d}  lrate {lrate(step):.3e}")

step      1  lrate 1.747e-07
step    100  lrate 1.747e-05
step   1000  lrate 1.747e-04
step   4000  lrate 6.988e-04
step  10000  lrate 4.419e-04
step  50000  lrate 1.976e-04
step 100000  lrate 1.398e-04

The peak is about $7\times 10^{-4}$ for $d_{\mathrm{model}}=512$. The factor $d_{\mathrm{model}}^{-0.5}$ lowers the peak for wider models.

Two regularizers were used. Residual dropout, $P_{drop}=0.1$ for the base model, was applied to each sublayer output and to the sum of embeddings and positional encodings. Label smoothing with $\epsilon_{ls}=0.1$ replaced each one-hot target with a mixture that puts a little probability mass on other tokens. The paper reports that this hurt perplexity, because the model learns to be less certain, but improved accuracy and BLEU.

At inference the paper used beam search with beam size 4 and length penalty $\alpha=0.6$, and averaged the weights of the last 5 checkpoints for base models and the last 20 for big models.

The two main configurations, from Table 3 of the paper:

The paper reports that the base model trained in about 12 hours (0.4 seconds per step) and the big model in 3.5 days (1.0 seconds per step) on the eight P100s. Its Table 2 reports newstest2014 BLEU scores of 27.3 (base) and 28.4 (big) on English–German, and 38.1 (base) and 41.8 (big) on English–French. The big English–French model used $P_{drop}=0.1$ rather than 0.3. The paper presented 28.4 as more than 2 BLEU above the best previously reported results, including ensembles, at a fraction of their estimated training cost. The first arXiv version listed 41.0 for English–French; later revisions report 41.8. BLEU depends on tokenization and evaluation scripts, so cross-paper comparisons should be read with care. The paper also showed the model transferring to English constituency parsing with little task-specific tuning.

A NumPy implementation of one encoder layer

The code below implements scaled dot-product attention with optional masking, multi-head attention, a causal mask, sinusoidal positional encoding, a post-LN encoder layer, and a tied output projection. It uses the paper's widths with a tiny batch and random weights. It is a forward pass only; there is no training, dropout or learned layer-norm gain.

import numpy as np

rng = np.random.default_rng(0)

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)   # stability; -inf stays -inf
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Q: (..., n_q, d_k), K: (..., n_k, d_k), V: (..., n_k, d_v).
    mask: boolean, broadcastable to (..., n_q, n_k); True = may attend."""
    d_k = Q.shape[-1]
    assert K.shape[-1] == d_k and K.shape[-2] == V.shape[-2]
    scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(d_k)   # (..., n_q, n_k)
    if mask is not None:
        scores = np.where(mask, scores, -np.inf)
    weights = softmax(scores, axis=-1)
    out = weights @ V                                    # (..., n_q, d_v)
    assert out.shape == Q.shape[:-1] + (V.shape[-1],)
    return out, weights

def causal_mask(n):
    return np.tril(np.ones((n, n), dtype=bool))          # (n, n)

def sinusoidal_pe(n_pos, d_model):
    pos = np.arange(n_pos)[:, None]                      # (n_pos, 1)
    i = np.arange(d_model // 2)[None, :]                 # (1, d_model/2)
    angle = pos / (10000 ** (2 * i / d_model))
    pe = np.zeros((n_pos, d_model))
    pe[:, 0::2] = np.sin(angle)
    pe[:, 1::2] = np.cos(angle)
    return pe

def init(shape):
    return rng.standard_normal(shape) / np.sqrt(shape[0])

class MultiHeadAttention:
    def __init__(self, d_model, h):
        assert d_model % h == 0
        self.h, self.d_k = h, d_model // h
        self.W_q, self.W_k, self.W_v, self.W_o = (init((d_model, d_model)) for _ in range(4))

    def split(self, x):                  # (B, n, d_model) -> (B, h, n, d_k)
        B, n, _ = x.shape
        return x.reshape(B, n, self.h, self.d_k).transpose(0, 2, 1, 3)

    def __call__(self, x_q, x_kv, mask=None):
        B, n_q, d_model = x_q.shape
        Q = self.split(x_q @ self.W_q)
        K = self.split(x_kv @ self.W_k)
        V = self.split(x_kv @ self.W_v)
        heads, _ = scaled_dot_product_attention(Q, K, V, mask)   # (B, h, n_q, d_k)
        concat = heads.transpose(0, 2, 1, 3).reshape(B, n_q, d_model)
        return concat @ self.W_o

def layer_norm(x, eps=1e-6):
    mu = x.mean(axis=-1, keepdims=True)
    var = x.var(axis=-1, keepdims=True)
    return (x - mu) / np.sqrt(var + eps)     # learnable gain/bias omitted (=1, 0)

class FeedForward:
    def __init__(self, d_model, d_ff):
        self.W1, self.b1 = init((d_model, d_ff)), np.zeros(d_ff)
        self.W2, self.b2 = init((d_ff, d_model)), np.zeros(d_model)
    def __call__(self, x):
        return np.maximum(0, x @ self.W1 + self.b1) @ self.W2 + self.b2

class EncoderLayer:
    """Post-LN, as in the paper: LayerNorm(x + Sublayer(x)). No dropout here."""
    def __init__(self, d_model, h, d_ff):
        self.attn, self.ffn = MultiHeadAttention(d_model, h), FeedForward(d_model, d_ff)
    def __call__(self, x, mask=None):
        x = layer_norm(x + self.attn(x, x, mask))
        return layer_norm(x + self.ffn(x))

# Paper-sized widths, tiny batch and sequence.
B, n, d_model, h, d_ff, vocab = 2, 10, 512, 8, 2048, 1000
E = rng.standard_normal((vocab, d_model)) / np.sqrt(d_model)   # shared embedding
tokens = rng.integers(0, vocab, size=(B, n))
x = E[tokens] * np.sqrt(d_model) + sinusoidal_pe(n, d_model)    # (B, n, d_model)
assert x.shape == (B, n, d_model)

layer = EncoderLayer(d_model, h, d_ff)
y = layer(x)
assert y.shape == (B, n, d_model)
print("encoder output", y.shape, "mean %.3f std %.3f" % (y.mean(), y.std()))

logits = y @ E.T                   # tied pre-softmax projection
assert logits.shape == (B, n, vocab)
print("logits", logits.shape)

# Causal check: change only the last token; earlier outputs must not move.
mha = MultiHeadAttention(d_model, h)
m = causal_mask(n)
out1 = mha(x, x, m)
x2 = x.copy(); x2[:, -1] += rng.standard_normal(d_model)
out2 = mha(x2, x2, m)
print("max change, positions 0..n-2:", np.abs(out1[:, :-1] - out2[:, :-1]).max())
print("max change, last position:  ", np.abs(out1[:, -1] - out2[:, -1]).max())
# Same experiment with no mask: earlier positions now do change.
leak = np.abs(mha(x, x)[:, :-1] - mha(x2, x2)[:, :-1]).max()
print("no mask, max change, positions 0..n-2:", leak)

On this run it printed:

encoder output (2, 10, 512) mean -0.000 std 1.000
logits (2, 10, 1000)
max change, positions 0..n-2: 0.0
max change, last position:   1.4952215602361845
no mask, max change, positions 0..n-2: 1.5112725161901035

The shape assertions pass. The encoder output has mean zero and standard deviation one per position because the final operation is a layer norm without a learned gain. The causal check is the important line. Perturbing the last token changes the last output, but the change at every earlier position is exactly 0.0, not merely small: masked weights are exactly zero, so the perturbed value and key rows never enter those sums. Remove the mask and the same perturbation moves earlier outputs by about 1.5. This is the property that lets a decoder be trained on all target positions in parallel without seeing the answer.

Three implementation details are worth noticing. The softmax subtracts the row maximum before exponentiating; this prevents overflow and leaves $-\infty$ entries at $-\infty$. The heads are produced by reshaping one big projection rather than by eight separate matrices, which is equivalent and faster. And the embedding table is initialized with standard deviation $1/\sqrt{d_{\mathrm{model}}}$ and multiplied by $\sqrt{d_{\mathrm{model}}}$, so token vectors have entries of roughly unit scale, comparable to the positional encoding. A decoder layer would add a masked self-attention call with causal_mask(n) and a cross-attention call attn(x_dec, encoder_out) between the two sublayers shown here.

Descendants and limitations

The encoder–decoder split turned out to be optional. Three families followed.

Encoder-only models such as BERT (Devlin et al., 2018) keep the bidirectional encoder stack and train it to recover masked tokens. Without a causal mask, every position sees the whole input, which suits classification, tagging and retrieval rather than free generation.

Decoder-only models such as GPT drop the encoder and cross-attention, keeping a stack of masked self-attention and feed-forward layers trained to predict the next token. Most current large language models are in this family, usually with pre-LN or RMSNorm, rotary or other relative position schemes, and different activation functions. The language-model guide linked in the introduction builds one of these.

Encoder–decoder models such as T5 keep the original layout and cast many tasks as text-to-text problems.

The main limitation is visible in Table 1: attention cost and memory grow as $n^2$. For a 512-token sentence this is modest; for long documents it dominates. Later work has attacked the problem from several directions. Sparse, local and low-rank attention variants approximate the full matrix. Exact methods such as FlashAttention reorganize the computation to avoid materializing the $n\times n$ matrix in slow memory, reducing memory traffic without changing the result. During generation, a cache of past keys and values avoids recomputing them, though the cache itself grows with context length. These are engineering and modeling refinements around the core operation; scaled dot-product attention itself remains largely as the paper defined it.

There are also limits the paper did not claim to solve. Positional encodings that work on training lengths do not automatically extrapolate to much longer inputs. Attention weights are not reliable explanations of model behavior on their own. And the paper is a translation paper: its results show a better sequence-to-sequence model at its scale, not the later scaling phenomena that made the architecture famous.

For a line-by-line PyTorch reproduction of the original model, The Annotated Transformer from Harvard NLP is a careful companion to the paper.

Practice exercises

1. Redo the worked example by hand with $V$ replaced by the identity matrix. Explain why the output now equals the attention weight matrix. Then change the third query to $(2,2)$ and predict, before running code, whether the third row's weights become more or less concentrated.

2. Measure saturation. Using the code in the scaling section, initialize random $Q$ and $K$ with $d_k=256$ and compute the average entropy of the attention rows with and without the $1/\sqrt{d_k}$ factor. Relate the result to the softmax Jacobian.

3. Verify the cost claim. Count the multiply–add operations of the score computation for one head with $d_k=512$ and for eight heads with $d_k=64$, at $n=100$. Then count the projection costs. Which part of a full layer dominates when $n=100$ and when $n=4000$?

4. Build a decoder layer. Extend the NumPy code with masked self-attention, cross-attention over the encoder output, and a feed-forward sublayer. Use a source length of 7 and a target length of 5, and assert the shape of the cross-attention score matrix. Repeat the causal check for the decoder.

5. Extend the positional proof. Build the full block-diagonal $R_k$ for $d_{\mathrm{model}}=16$ and verify $PE_{pos+k}=R_k\,PE_{pos}$ for several positions. Then show that the dot product $PE_{pos}\cdot PE_{pos+k}$ depends only on $k$.

6. Swap post-LN for pre-LN. Change the encoder layer to $x+\operatorname{Sublayer}(\operatorname{LayerNorm}(x))$, stack 24 layers with random weights, and compare the standard deviation of the residual stream across depth for both versions. Write down what you would expect to happen to gradients and test the expectation in a framework with automatic differentiation.

7. Read Table 3 of the paper and pick one ablation row. State what was varied, what was held fixed, and whether the reported difference is large compared with the variation you would expect between random seeds.