Inside the Transformer

"The keys to the cabinet beside the stairs are missing." You know that "are" belongs with "keys," even though "cabinet" and "stairs" sit closer to it. A language model needs a way to preserve that relationship while processing all the other words. Otherwise, the nearest noun can become a very convincing distraction.

The Transformer gives each token a direct way to draw information from other tokens. That is the useful idea behind attention. We can follow this sentence through the machinery to see how the model builds context, why word order still needs special treatment, and why being able to reach the right word doesn't guarantee using it correctly.

Getting information across a sentence

A recurrent network processes a sequence one step after another. Information about "keys" must survive in a hidden state as the network passes through "cabinet," "beside," and "stairs." LSTMs were designed to retain information over longer spans, but that information still travels through successive states. The dependency between timesteps also limits parallel computation across positions, even though work within a timestep and across examples can run in parallel.

Self-attention lets a position consult earlier positions directly. During training, the known tokens can be processed together within a layer, with a mask controlling which connections are allowed. Generating new text still requires successive prediction steps. This combination of shorter information paths and more parallel training was central to the original Transformer paper.

For the example, treat each word as one token and ignore punctuation. Real tokenizers can split words into smaller pieces. Each token starts as a vector, a list of numbers learned during training. Through successive layers, that vector becomes a representation of the token in its context.

What a token asks, matches, and receives

Consider the representation at "are." We'd like it to carry information about "keys." Attention creates three versions of each input vector:

A query, which determines what it matches against.
A key, which other queries can match.
A value, which supplies information when a match receives weight.

The names suggest a lookup, but there is no exact-match dictionary or hand-written rule saying "find the subject." The model learns projections that make useful relationships influence its predictions. Given the input matrix $X$ , those projections are:

Q = XW_Q \qquad K = XW_K \qquad V = XW_V

For $n$ tokens, $Q, K \in \mathbb{R}^{n \times d_k}$ and $V \in \mathbb{R}^{n \times d_v}$ . The weight matrices are learned during training. At inference, the same learned matrices produce different queries, keys, and values for different inputs.

The complete attention calculation is:

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

Read it as a sequence of operations. $QK^\top$ compares every query with every key using a dot product, producing an $n \times n$ score matrix. The entry $(i, j)$ is a score for how strongly position $i$ matches position $j$ . The scaling factor $1/\sqrt{d_k}$ controls the scale of these scores as the vector dimension $d_k$ grows; excessively large scores can push softmax into regions with small gradients. Softmax then turns each row into nonnegative weights that sum to one. Finally, multiplying by $V$ mixes the value vectors in those proportions.

Imagine, purely as a toy example, that the row for "are" assigns 0.60 to "keys," 0.10 to "cabinet," 0.10 to "stairs," and 0.20 spread across the remaining allowed positions. Its output contains 60% of the value vector for "keys," 10% from each of the other named tokens, and the remaining contributions. These are illustrative weights, not a recording from a trained model.

Notice what the calculation produces: another vector. It doesn't output a grammatical explanation, and 0.60 isn't a 60% probability that "keys" is the subject. Later operations must use the mixed information to make a prediction. Attention determines which information gets combined; the usefulness of that combination is learned.

Why several heads help

The sentence contains more than subject agreement. "Cabinet" relates to "keys," while "beside the stairs" locates the cabinet. One weighted mixture can entangle information useful for different relationships.

Multi-head attention gives the layer several separately learned sets of queries, keys, and values. Each head calculates its own mixture; the outputs are concatenated and projected back into the model's vector dimension. In the implementation below, $h$ heads each use $d_{model}/h$ dimensions, so their combined width stays fixed.

Heads can develop recognizable patterns. Clark and colleagues' analysis of BERT found heads associated with positional offsets, syntax, and coreference. That doesn't assign every head a tidy job title. A head can participate in several patterns, and different heads can overlap. Our imagined "keys" connection explains a possible computation, not a guaranteed subject-detection head in every Transformer.

Order and the words a model may see

"The cabinet contains the keys" and "The keys contain the cabinet" use the same words. Mixing their vectors without information about order would miss something fairly important about furniture.

Unmasked self-attention without positional information is permutation-equivariant: reorder the input vectors and the output vectors reorder with them. The mechanism alone does not encode which token came first. The original Transformer adds positional encodings to the embeddings:

PE_{(pos,\,2i)} = \sin\!\left(\frac{pos}{10000^{2i/d_{model}}}\right) \qquad PE_{(pos,\,2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d_{model}}}\right)

Different pairs of dimensions oscillate at different frequencies, giving positions distinct patterns. For a fixed offset $k$ , the encoding at $p+k$ can be expressed as a linear transformation of the encoding at $p$ . This gives the network a way to learn relationships between positions, including nearby ones.

Other designs use learned position embeddings, as BERT does, or rotate queries and keys, as rotary position embeddings, or RoPE, do. RoPE makes relative position part of the query-key comparison. Its distance-related properties don't mean every attention weight decreases monotonically with distance, or that a model will work reliably beyond its training lengths. The LLaMA architecture is one example that uses RoPE.

A causal mask imposes a different constraint: which positions are visible. To predict "are" from "The keys to the cabinet beside the stairs," a decoder uses the representation at the last available position, "stairs." It cannot inspect the answer token "are" before predicting it. Once "are" is part of the input, its representation can attend back to "keys" while helping predict "missing."

During training, the entire example is already known, so a triangular mask allows position $i$ to attend only to positions 0 through $i$ . That prevents it from reading future answers while still permitting parallel computation over the known sequence. Position encodings describe order; the mask restricts access.

What happens after attention

Attention mixes information across positions. A feed-forward network then transforms each position's vector independently, using the same learned network at every position. With "keys" information now present in a later vector, that transformation can contribute features useful for predicting what follows.

The original Transformer uses two linear layers with a ReLU between them, expanding the width by 4x before projecting back. Other designs use different nonlinearities or gated feed-forward layers. Residual connections add a sublayer's output to its input, giving information and gradients a direct route through the stack. Layer normalization controls the scale of each position's features.

Omitting dropout, the original block's post-normalization arrangement is:

\text{h} = \mathrm{LayerNorm}\big(x + \mathrm{MultiHeadAttention}(x)\big)

\text{output} = \mathrm{LayerNorm}\big(\text{h} + \mathrm{FFN}(\text{h})\big)

Other Transformers place normalization before the sublayers. The important pattern is repeated contextual mixing and transformation, with residual connections carrying information forward. A final projection and softmax convert the resulting representation into probabilities over the output vocabulary.

Feed-forward layers also matter for learned knowledge. Geva and colleagues describe them as key-value memories that associate input patterns with output tendencies. ROME demonstrated edits to particular factual associations by changing feed-forward weights. These findings identify useful mechanisms and intervention points. They don't establish that all model knowledge occupies a separate, neatly editable fact store. Prediction depends on computations across the network.

One mechanism, different arrangements

The original encoder-decoder Transformer has two stacks. Its encoder reads the full input with bidirectional self-attention. Its decoder uses causal self-attention over the output so far, plus cross-attention: decoder queries match keys and values derived from the encoder's output. In translation, that lets each emerging output token consult the source sentence.

A decoder-only model such as LLaMA uses causal attention over a single sequence. A BERT-style encoder reads both left and right context, which suits tasks such as classification and extraction. BERT predicts masked tokens during pretraining; it isn't configured as a standard left-to-right generator. Encoder-decoder models such as T5 retain separate input and output sequences for text-to-text tasks. The arrangement determines what information a prediction can use.

Why the design scales, and where it costs

Parallel training and direct connections between positions made Transformers attractive for larger models. Empirical work by Kaplan and colleagues found predictable relationships between language-model loss, parameters, data, and compute. Hoffmann and colleagues showed why balancing model size with training tokens matters under a fixed compute budget. Those are measured trends in training and evaluation, not a promise that increasing one resource improves every answer.

The direct connections also have a price. Dense self-attention performs a quadratic number of pairwise comparisons in sequence length, $O(n^2)$ at fixed head dimension. For 100K tokens, a full attention score matrix has 10 billion entries per head. Doubling the sequence length quadruples that number.

Sparse, sliding-window, and linear attention methods change or approximate the attention computation to reduce its cost. FlashAttention computes exact attention while reducing memory traffic and avoiding storage of the full matrix. It improves how the calculation runs; it doesn't remove the quadratic arithmetic of dense attention.

The calculation in code

These two compact implementations expose the operations discussed above. The first returns both the mixed values and the attention weights. The second projects the input into multiple heads, applies attention, then recombines them. They are teaching implementations; the attention weights are useful to inspect but don't by themselves explain a model's full decision.

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

The mask must broadcast to the score tensor, and every query must retain at least one allowed key. A lower-triangular mask including the diagonal meets that condition for ordinary causal self-attention. Masking every entry in a row would make this softmax undefined.

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_Q = torch.nn.Linear(d_model, d_model)
        self.W_K = torch.nn.Linear(d_model, d_model)
        self.W_V = torch.nn.Linear(d_model, d_model)
        self.W_O = torch.nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()

        # Project and reshape into (batch, heads, seq_len, d_k)
        Q = self.W_Q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Attention per head
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask)

        # Concatenate heads and project
        attn_output = attn_output.transpose(1, 2).contiguous().view(
            batch_size, seq_len, -1
        )
        return self.W_O(attn_output)

For this reshape to work, d_model must be divisible by n_heads. A shared causal mask of shape (seq_len, seq_len) broadcasts across batches and heads; padding masks need dimensions compatible with (batch, heads, seq_len, seq_len). This example implements self-attention because Q, K, and V all come from x. Cross-attention takes queries from the decoder and keys and values from the encoder.

Having access is the beginning

Return to the missing keys. We can now identify a computational path: earlier token vectors supply keys and values, a later query determines their weights, and the resulting mixture carries context into subsequent transformations. Position information preserves order, while causal masking keeps prediction from looking ahead. Repeating those operations allows a model to develop richer representations of the sentence.

That path explains how a distant word can influence a prediction. Correct interpretation still depends on what the model learned and how it uses the available context. In the Lost in the Middle experiments, moving relevant information within long inputs changed question-answering and retrieval performance, including for models designed to accept long contexts.

When evaluating a model for a real task, I'd therefore test the relationship I need it to preserve: move the relevant fact, add plausible distractions, and check whether the answer stays correct. A context window tells us how much text fits. The useful question is whether the model can keep track of the keys when the cabinet and the stairs get in the way.