Attention Is All You Need

Vaswani et al. · NeurIPS 2017 · 16 min read

A recurrent encoder passes information along a chain of states. Self-attention lets each position form a weighted combination of other positions directly. The Transformer organizes that operation into a sequence model with separate rules for reading an input and generating an output.

This companion builds a tiny, hand-set example: three positions, four coordinates, and two attention heads. You can reproduce every weight, change a future input, and see exactly when a causal mask blocks its influence. These teaching weights are not learned parameters from the paper.

What changes when the recurrent chain disappears?

Consider a translation whose source contains three tokens. In a recurrent encoder, the third hidden state depends on the second, which depends on the first. In a self-attention layer, all three positions can construct their queries, keys, and values from the same layer input. Their pairwise comparisons can be computed together.

That removes a dependency within this layer; it does not remove every sequential dependency. Layer two still needs layer one. During ordinary autoregressive generation, the next token is still unknown until the preceding token has been generated. The authors’ Google Research explanation describes the motivation and parallel training advantage.

The original base configuration uses six encoder and six decoder layers, model width 512, eight heads, and feed-forward width 2048. Its residual sublayers use normalization after addition. Modern variants need not preserve those choices. Paper §3.

Queries choose weights; values supply the content

Let $X$ contain one input vector per row. A head has three learned projections:

A query is the vector used to ask what this position should read. Keys are the vectors it compares against. Values contain the coordinates that will actually be combined. All three come from $X$ in self-attention, but their projection matrices can differ. A high query–key score does not mean the key vector itself becomes the output.

For query row $i$ and key row $j$, form a score, normalize across allowed keys, and combine values:

Here $d_k$ is the number of query/key coordinates and $\mathcal A_i$ is the set of allowed keys. Forbidden weights are zero. The denominator is per query row: each row sums to one, while a column need not. In matrix notation:

Here $M$ is zero at permitted entries and negative infinity at forbidden ones.

Self-attention has n = m because both sides are the same sequence. Cross-attention may have different source and target lengths. Batch and head dimensions sit outside these matrices; softmax must still operate along the key dimension.

An original three-position, two-head calculation

Take these inputs to one attention sublayer. Treat them as already supplied representations; this example does not train embeddings or add a separate position vector.

Head 1 uses coordinates 1–2 for its queries and keys, and 3–4 for its values. Head 2 does the reverse. Each head therefore has $d_k=d_v=2$. These projections select coordinates, a particularly transparent special case of a learned linear map.

At position 3, head 1’s query is [1, 1]. Its three dot products are [1, 1, 2]; after division by √2, the scores are [0.707107, 0.707107, 1.414214]. Subtracting the maximum before exponentiating leaves [−0.707107, −0.707107, 0], so the weights are approximately [0.248255, 0.248255, 0.503490].

The values are [1, 0], [0, 1], and [0, 0]. Despite receiving the largest weight, the third value contributes a zero vector. Thus head 1 returns [0.248255, 0.248255]. Head 2’s query is [0, 0], giving three equal scores, uniform weights, and the average of [1, 0], [0, 1], and [1, 1]: [2/3, 2/3].

Multi-head attention concatenates the two outputs and then applies an output projection. Our identity output projection leaves [0.248255, 0.248255, 0.666667, 0.666667]. Averaging the heads would instead give only two coordinates and erase their separate outputs. A learned $W_O$ can subsequently mix the concatenated coordinates; it does not make concatenation and averaging the same operation.

Test the causal rule by changing the future

Select position 2 and keep the causal mask. Its allowed keys are positions 1 and 2, so both heads give weights approximately [0.330238, 0.669762, 0]. Now change the final input row. The selected output must remain identical: that row is a forbidden key/value, and position 2’s query has not changed.

Switch to “All positions.” The same future change can now alter position 2’s output. That is appropriate for bidirectional encoding, but would leak future target information in an autoregressive decoder. Position 3 can change in either mode because the edited row is its own input, not its future.

A causal mask alone is insufficient if the target inputs are misaligned. For the desired targets A B END, teacher-forced decoder inputs should be START A B. Position 2 reads START and A to predict B. Feeding B at that same input position would reveal the answer even though the mask blocks positions after 2.

Scaling changes the distribution, not just the units

Disable scaling at position 3. Head 1 now uses [1, 1, 2], yielding [0.211942, 0.211942, 0.576117]. More weight goes to the largest score. Because its value is zero, each output coordinate decreases from 0.248255 to 0.211942. This connects a change in logits to an actual change in retrieved content.

A simple variance calculation explains the square-root factor. If each query/key coordinate is independent, centered, and has variance 1, then each product has variance 1 and the sum of $d_k$ products has variance $d_k$. Dividing by $\sqrt{d_k}$ returns variance 1 under those assumptions. Learned coordinates need not satisfy the assumptions exactly; scaling is not a guarantee that attention remains diffuse.

Subtracting the row maximum is different: softmax is unchanged by adding the same constant to every allowed score. It prevents overflow without changing the modeled probabilities. Dividing scores changes their differences and therefore changes those probabilities.

Attention is one sublayer, not the whole model

Our position-3 attention output is added to its incoming vector [1, 1, 0, 0]. The residual sum is [1.248255, 1.248255, 0.666667, 0.666667]. With normalization across its four coordinates, mean 0.957461, scale parameters γ = 1 and β = 0, and ε = 10⁻⁶, the normalized vector is approximately [0.999994, 0.999994, −0.999994, −0.999994].

This operation normalizes a position’s features, not all tokens in the sequence together. A constant vector becomes zero before the learned offset; the epsilon prevents division by zero. The LayerNorm definition specifies the mean, population-variance denominator, and learned affine terms.

The feed-forward sublayer then applies a shared nonlinear function to each row, for example $\operatorname{ReLU}(zW_1+b_1)W_2+b_2$. “Shared” means the same weights are used at every position, not that the positions have identical outputs. Attention mixes information between positions; this sublayer transforms each contextual vector’s coordinates. Another residual connection preserves a direct route around that transformation.

The toy stops after the first residual normalization. It does not include a feed-forward fit, cross-attention, vocabulary projection, or training. Those omissions make the arithmetic inspectable; they must not be mistaken for a complete Transformer implementation.

Why position information is needed

Without position information or an order-dependent mask, permuting the input rows permutes the self-attention output rows in the same way. Pairwise dot products and weighted sums have no independent notion of “first” or “third.” This is useful symmetry for a set, but insufficient to represent arbitrary word order.

The paper adds sinusoidal position vectors. For each frequency pair, advancing position rotates a sine/cosine pair by a fixed angle; that is a property of the representation, not a promise that learned translation behavior extrapolates to any sequence length. The cyclical-encoding example develops the same paired-coordinate geometry. A causal mask supplies an ordering restriction too, but it is not the same mechanism as adding positional features.

Parallel work still has a cost

With n positions and total model width d, projections take on the order of nd² arithmetic operations. Computing dense query–key comparisons and mixing values takes on the order of n²d across heads when their widths sum to d. A straightforward implementation stores h matrices of size n × n for h heads.

For n = 1,024 and eight heads, that is 8,388,608 attention entries: 16 MiB if each stored entry uses two bytes, before gradients or other activations. Doubling n quadruples that entry count. Implementations can avoid materializing the whole matrix, but “parallel” alone does not remove the pairwise work.

During teacher forcing, every shifted target input is already available, so causal rows can be computed in parallel. At generation time, the next input token depends on the latest prediction. Keeping previous keys and values avoids recomputing their projections; it does not make unknown future tokens available.

Read the reported results in their original setting

These are machine-translation results under the paper’s data, tokenization, training, checkpoint, and decoding setup. They are not accuracy percentages or results from this toy. The larger model’s gains were 1.1 and 2.9 BLEU points over the base model in that table. Original proceedings, Table 2.

Version matters: the later arXiv v7 Table 2 gives 41.8 for big English→French, while its nearby prose still says 41.0. The table above follows the original conference PDF so the source version is explicit.

The useful reading question is what the comparison demonstrates: the proposed encoder–decoder could train and translate competitively without a recurrent sequence operator. It does not establish that every task benefits equally, that all heads learn distinct linguistic roles, or that attention weights alone identify a causal explanation.

Reproduce and challenge the calculation

Download the complete standard-library Python example and all 24 computed states. The script records every projected vector, allowed score, weight, head output, residual sum, and normalized vector. The following is the standalone masked normalization at its core:

import math

def masked_softmax(scores):
    # None marks a forbidden key, not a numerical score of zero.
    allowed = [score for score in scores if score is not None]
    if not allowed:
        raise ValueError("No permitted key for this query")
    maximum = max(allowed)
    numerator = [0.0 if s is None else math.exp(s - maximum)
                 for s in scores]
    denominator = sum(numerator)
    return [n / denominator for n in numerator]

weights = masked_softmax([0.0, 1 / math.sqrt(2), None])
assert abs(sum(weights) - 1) < 1e-12
assert weights[2] == 0
print(weights)  # [0.3302384507, 0.6697615493, 0.0]

Useful checks go beyond a row sum of one. Change a forbidden future row and compare an earlier output. Verify that reversing the head-concatenation order changes the coordinates unless the output projection changes accordingly. Check that a constant offset leaves softmax unchanged while scaling does not. Finally, compare head outputs, not only their heatmaps: the values determine what information is carried forward.

Where to go next

Additive attention explains learned alignment while retaining a recurrent decoder. This chapter replaces the sequence operator and separates three attention roles. BERT and GPT-2 then make different choices about permitted context and training tasks; they should not be read as interchangeable applications of one mask.

For a broader implementation-oriented explanation, continue to the Transformer and pretraining textbook chapter. When reading any implementation, write down its input shift, mask convention, normalization order, dimensions, and output units before matching it to a paper diagram.

References

• Vaswani et al., Attention Is All You Need — Original architecture and reported translation experiments.
• Google Research: Transformer — The authors’ explanation of sequence dependencies and parallel processing.
• PyTorch Transformer and LayerNorm — Shape, masking, and normalization conventions to inspect when implementing the calculation.