What Is a Transformer? Self-Attention Explained
The transformer architecture is a function that takes a matrix of token embeddings and returns a matrix of the same shape, enriched by information from every other position in the sequence. The mechanism that does that enrichment is called self-attention, and understanding it precisely removes most of the mystery around how LLMs work.
The Problem Attention Solves
Consider the sentence "The animal did not cross the street because it was too tired." The word "it" refers to "animal," not "street." A system that processes each word in isolation cannot resolve this. It needs to relate "it" to the correct antecedent. That is what attention does: it allows each token to gather information from other tokens weighted by relevance.
In a recurrent network, this would happen indirectly: information about "animal" would have to travel through every subsequent hidden state until it reached "it." If the sequence is long, that signal can degrade. The transformer's attention mechanism allows "it" to directly query "animal" in a single step, regardless of how many tokens separate them.
Queries, Keys, and Values
Self-attention starts with three learned linear projections of the input. For an input matrix X of shape (sequence length, d_model), three weight matrices W_Q, W_K, W_V each of shape (d_model, d_k) produce:
K = X W_K (keys)
V = X W_V (values)
The intuition is borrowed from database retrieval. A query represents what a token is looking for. A key represents what a token is offering. The dot product Q K^T measures compatibility between every query-key pair, producing an (n, n) matrix of raw scores where n is sequence length.
Those scores are scaled by 1/sqrt(d_k) to prevent the dot products from becoming very large when d_k is large, which would push the softmax into regions where gradients are nearly zero. Softmax is then applied row-wise to produce attention weights: non-negative values that sum to one for each query position.
The final output is a weighted sum of the value vectors. Each token's output is a blend of all value vectors in the sequence, with the blend determined by how much attention it paid to each position.
The Causal Mask
During training, the model sees the full sequence and predicts all positions simultaneously, but it should not be able to look ahead. A token at position i should only attend to positions 0 through i. This is enforced by adding a mask to the attention scores before softmax: positions in the future are set to negative infinity, which softmax maps to zero attention weight.
During inference, the model generates one token at a time, so there are no future positions to mask. But the forward pass still runs through the same architecture. The KV cache exploits this: since the keys and values for all prior positions are already computed and unchanged when a new token is added, they can be cached and reused, avoiding redundant computation.
Multi-Head Attention
A single attention head can only express one "mode" of relating tokens to each other. Multi-head attention runs h parallel attention heads, each with its own Q, K, V projections to a reduced dimension d_k = d_model / h. The outputs of all heads are concatenated and projected back to d_model with a learned weight matrix W_O.
Different heads learn to attend to different types of relationships. Some heads attend to adjacent tokens (local syntax), others attend to semantically related tokens far in the sequence, and others appear to track coreference or subject-verb agreement. The multi-head design gives the model the capacity to express all of these simultaneously.
Attention weights are not an explanation of what the model is doing. They are a description of where the model looked. Those are different things.
The Feed-Forward Sub-Layer
Each transformer layer contains two sub-layers: attention and a position-wise feed-forward network (FFN). The FFN applies two linear transformations with a nonlinearity between them:
The inner dimension is typically 4 times d_model, so the FFN expands the representation, applies a nonlinearity, and contracts it again. Modern models often replace ReLU with GELU or SwiGLU, which have smoother gradients and empirically perform better at large scale.
Research suggests the FFN layers act as key-value memories: the rows of W_1 are keys and the rows of W_2 are values, and an input activates the rows most similar to it, retrieving the associated stored pattern. This may be where much of the factual knowledge of a language model is stored, separate from the contextual aggregation done by attention.
Residual Connections and Layer Normalization
Both sub-layers in a transformer layer are wrapped with a residual connection: the output of the sub-layer is added to the sub-layer's input, then layer normalization is applied. This is the pre-norm variant used in most modern models:
x = x + FFN(LayerNorm(x))
Residual connections allow gradients to flow directly through the network without passing through the attention or FFN operations, which makes training deep networks much more stable. Layer normalization stabilizes the activation distribution within each layer, reducing sensitivity to hyperparameter choices and enabling higher learning rates.
Why Depth Matters
A single transformer layer can only combine information from context in one step. Stacking N layers allows progressively more abstract representations to emerge. Early layers tend to capture local syntactic patterns; later layers tend to represent semantic content and long-range relationships. The depth of the network is what allows the model to build up the hierarchical representations that support complex language understanding.
Modern LLMs use anywhere from 12 layers (small models) to 96 or more layers (frontier models). The optimal depth depends on the parameter budget and the compute available for training, but in general, depth and width should scale together. A very wide but shallow network and a very narrow but deep network will both underperform a balanced configuration at the same parameter count.
Positional Encodings: Giving the Model a Sense of Order
Self-attention operates on a set of vectors. Without additional information, it is permutation-invariant: shuffling the input tokens in any order produces the same output, because attention weights depend on content similarity, not position. Positional encodings break this symmetry by adding position-dependent information to the token embeddings before the first layer.
The original transformer used fixed sinusoidal functions of position, with different frequencies for different embedding dimensions. This approach allows the model to reason about relative distances through the trigonometric identities that relate sin and cos at different positions. Learned positional embeddings, used in early GPT variants, assign a trainable vector to each absolute position. Both approaches struggle to generalize to sequences longer than those seen during training.
Rotary Position Embedding (RoPE), introduced by Su et al. (2021) and adopted in Llama, Mistral, and most modern architectures, takes a different approach: it encodes position as a rotation applied to the query and key vectors in the attention computation. The rotation angle is proportional to position and frequency, so the dot product between a query at position i and a key at position j depends only on their relative distance i minus j, not on their absolute positions. This relative encoding property allows RoPE to generalize more naturally to longer sequences, and position interpolation techniques can extend it further.
The KV Cache: Making Inference Efficient
During autoregressive generation, the model generates one token at a time. At each step, the full context (prompt plus previously generated tokens) is passed through the transformer stack. Without optimization, this means recomputing the key and value vectors for every prior token at every generation step, which is quadratically expensive in the number of tokens generated.
The KV cache eliminates this redundancy. Because the keys and values for token positions 0 through t are determined by those tokens alone (in a causal model that does not look ahead), they are identical regardless of what token is generated at position t+1. The KV cache stores these tensors after they are first computed, so each generation step only requires computing the keys and values for the new token. The attention then queries the cached keys and values for all prior positions.
The memory cost of the KV cache scales with context length, number of layers, number of attention heads, and head dimension. For a model with 32 layers, 32 heads, and a head dimension of 128, running at 16-bit precision, each token in the cache requires 32 multiplied by 2 multiplied by 32 multiplied by 128 multiplied by 2 bytes, which is approximately 0.5 MB per token. For a 128,000-token context, the KV cache alone requires roughly 64 GB of memory. This is the primary constraint limiting context length at deployment scale.
Grouped-Query Attention and Mixture of Experts
Two architectural modifications have become standard in frontier models because they address the cost constraints of inference at scale.
Grouped-query attention (GQA) reduces the KV cache memory requirement by sharing key and value projections across groups of query heads. In standard multi-head attention, each of the h heads has its own W_K and W_V projections. GQA uses g groups where g is smaller than h: all heads within a group share the same key and value projections. With g=1 (multi-query attention, MQA), all heads share a single K and V, reducing KV cache memory by a factor of h. GQA with g between 1 and h finds a balance between memory savings and model quality. Llama 3 uses GQA across its model family.
Mixture of experts (MoE) addresses parameter count and computational cost separately. A MoE model replaces each dense FFN layer with a set of E expert FFN networks and a router that selects the top-k experts for each token. Because only k of E experts are activated per token, the active compute per forward pass is much lower than a dense model with E times as many parameters. The tradeoff is that all E experts must fit in memory, even though only k are used per step. Sparse MoE architectures have enabled scaling to very large parameter counts without proportional increases in training or inference compute.
From Architecture to Capability: What the Transformer Actually Learns
The transformer architecture is a framework. The capabilities that appear in trained models, the ability to translate languages, write code, solve math problems, and reason about novel situations, are not properties of the architecture itself. They emerge from applying that architecture to massive amounts of text under the next-token prediction objective.
This distinction matters for interpreting research findings. When a paper claims that a model can perform a task, it is making a claim about the trained model, not about the architecture. Changing the training data, the training procedure, or the fine-tuning strategy can substantially alter what tasks the model can perform, even with an identical architecture. The architecture sets the ceiling; training determines what fraction of that ceiling is reached for any specific capability.
Mechanistic interpretability research, which attempts to reverse-engineer what specific attention heads and MLP neurons have learned, has made progress in identifying structures like induction heads (which copy patterns from earlier in the context) and factual recall circuits (which retrieve specific associations stored during training). These findings suggest that the transformer's capabilities are not opaque magic: they are the result of learnable, identifiable computational structures that emerge reliably from training on text at sufficient scale.
Choosing the Right Transformer Configuration
Several architectural hyperparameters must be set before training begins, and they are difficult to change afterward. The number of layers, the hidden dimension, the number of attention heads, the head dimension, and the FFN expansion factor all determine the final model size and the distribution of parameters across layers. For a given parameter budget, the right configuration depends on the inference cost constraints and the target deployment scenario.
Models intended for edge deployment or low-latency inference favor wider and shallower configurations: fewer layers means fewer sequential operations that must complete before the next token can be generated. Models intended for high-accuracy, latency-tolerant applications can afford more depth, which enables richer hierarchical representations at the cost of higher minimum latency per token.
The attention head configuration also matters for the KV cache requirements at inference time. Models with many small heads (high head count, low head dimension) have smaller KV cache entries per layer than models with few large heads, which can make a significant difference when serving thousands of concurrent users. Grouped-query attention reduces this further by sharing keys and values across groups of query heads, with measurable quality impact only at high reduction ratios.
For practitioners who are fine-tuning or adapting existing models rather than training from scratch, these choices have already been made. But understanding them informs decisions about which base model to select for a given use case, how to configure serving infrastructure, and what performance envelope to expect from any given deployment configuration.
The transformer architecture has proven durable because it is modular and parallelizable. Each architectural decision, from head count to FFN expansion factor to positional encoding, can be studied in isolation. Each improvement, whether GQA, RoPE, or SwiGLU activations, addressed a specific, well-understood limitation of the original design. This incremental improvability is one reason the transformer has remained the dominant architecture for large-scale language modeling for nearly a decade: it is not merely effective but also legible and modifiable in principled ways. Understanding these principles, not just the surface-level API, is what enables engineers to make good decisions when the default configuration is not appropriate for a specific task.
The self-attention mechanism is the central innovation that makes transformers powerful, but it is the combination of attention with residual connections, layer normalization, and deep stacking that makes them trainable at scale. Each of these components solves a specific engineering problem that arises when training very deep networks: vanishing gradients, unstable activations, and the difficulty of learning long-range dependencies through sequential computation. The transformer's value is not in any single component but in how these components interact to allow stable, scalable, parallelizable training on the next-token prediction objective. That combination, applied to text at the scale that modern compute enables, is what produces the language capabilities that practitioners rely on today.
References
- Vaswani et al. "Attention Is All You Need." NeurIPS 2017. arXiv:1706.03762
- Ba et al. "Layer Normalization." arXiv 2016. arXiv:1607.06450
- Geva et al. "Transformer Feed-Forward Layers Are Key-Value Memories." EMNLP 2021. arXiv:2012.14913
- Su et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding." arXiv 2021. arXiv:2104.09864
- Shazeer. "GLU Variants Improve Transformer." arXiv 2020. arXiv:2002.05202
- Press et al. "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation." ICLR 2022. arXiv:2108.12409
Go Deeper with the Full Course
Module 1 of LLMs from Scratch includes an interactive attention score visualizer where you can tune query/key angles and see how attention weights change. Free.
Open Module 1