The Architecture of Context in Transformer Models

The Architecture of Context in Transformer Models

The behavior of any large language model is fundamentally determined by a single mechanism: how it decides what to look at when producing each output token. This mechanism—self-attention—is not an implementation detail but the core operation around which all transformer architecture is designed. Understanding how self-attention works, what constraints it imposes, and how different architectural choices navigate those constraints is essential for anyone engineering applications that depend on context.

The fundamental question this chapter answers is deceptively simple: given a sequence of tokens in the context window, how does the model decide which previous tokens to consider when generating the next one? The answer requires examining the mathematical structure of attention, the architectural patterns that implement it at scale, and the practical implications for what information can and cannot influence model outputs. This understanding becomes the foundation for every strategy discussed in subsequent chapters—whether you're optimizing prompt design, managing retrieval-augmented generation pipelines, or engineering systems that process documents far exceeding any model's native context window.

The Mathematics of Self-Attention

From Words to Vectors: The Foundation

Before attention can operate, tokens must be represented as numerical vectors. This transformation happens through embedding layers that map discrete token identifiers to dense vector representations in a high-dimensional space, typically with hundreds or thousands of dimensions. Each dimension captures some aspect of semantic meaning, grammatical function, or positional relationship learned during pretraining.

The embedding process is straightforward in principle: a vocabulary of size $V$ is mapped through an embedding matrix $W_E$ of dimensions $d_{model} \times V$, producing a sequence of embedding vectors. These vectors are not arbitrary—they encode the statistical regularities the model learned about language during pretraining on vast text corpora. Tokens with similar meanings or functions tend to cluster in this embedding space, though the exact geometry reflects the particular biases of the training data and objectives.

Once embedded, each token vector undergoes a linear transformation to produce three distinct representations: the query ($Q$), the key ($K$), and the value ($V$). These transformations are learned parameters of the model, implemented as weight matrices $W_Q$, $W_K$, and $W_V$. The terminology derives from information retrieval: each token's query asks "what information am I looking for?" while other tokens' keys answer "what information do I contain?" The value represents the actual information to be retrieved if the query-key match is strong.

The Scaled Dot-Product Attention Equation

The attention computation itself follows a precise mathematical form that balances computational efficiency with expressive power. The scaled dot-product attention mechanism computes attention weights through the following operations:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Breaking this down step by step reveals why each component exists. First, the dot product $QK^T$ computes the similarity between every pair of query and key vectors. For a sequence of length $n$, this produces an $n \times n$ matrix where entry $(i, j)$ represents how much token $i$ "queries" token $j$. Higher values indicate stronger relevance relationships.

The division by $\sqrt{d_k}$—where $d_k$ is the dimensionality of the key vectors—serves a critical purpose. Without this scaling, the dot products grow large in magnitude as $d_k$ increases. The softmax function that follows is sensitive to scale: when its inputs have large magnitudes, it produces nearly one-hot outputs (extreme attention weights) rather than smooth distributions. This leads to gradients that are either vanishingly small or explosively large during training. The scaling factor $\sqrt{d_k}$ keeps the dot products in a range where the softmax produces well-behaved probability distributions.

Finally, the softmax normalizes each row of the attention matrix to sum to one, and the resulting attention weights multiply the value matrix $V$. This produces output vectors that are weighted sums of the value vectors, where the weights reflect learned relevance relationships.

Multi-Head Attention: Parallel Perspectives on Context

Single-head attention, while mathematically sound, captures only one type of relationship between tokens. Language is inherently multidimensional: a token's relevance might depend on syntactic relationships, semantic content, discourse structure, or coreference chains simultaneously. Multi-head attention addresses this by running several attention operations in parallel, each with its own learned query, key, and value projections.

Given $h$ attention heads, each head $i$ computes:

$$\text{head}_i = \text{Attention}(QW_Q^i, KW_K^i, VW_V^i)$$

The outputs of all heads are concatenated and linearly transformed:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W_O$$

This architecture allows different heads to specialize: some might attend to syntactic neighbors, others to semantically similar content, and others to coreferential relationships. The concatenated output preserves information from all perspectives before the final projection.

The computational implications are significant. If each head operates on $d_{model}/h$ dimensions (rather than $d_{model}$), the total computation remains comparable to single-head attention on the full dimension—each head does less work, but they run in parallel. However, the memory requirements for storing attention matrices scale quadratically with sequence length regardless of the number of heads: the $n \times n$ attention matrix must be computed and stored for every head in every layer.

```mermaid flowchart TD subgraph Input["Input Sequence"] T1[Token 1] T2[Token 2] T3[Token 3] Tn[Token n] end

subgraph EmbeddingLayer["Embedding + Position Encoding"] E1[Embedding Vector] E2[Embedding Vector] E3[Embedding Vector] En[Embedding Vector] end

subgraph LinearProjections["Linear Projections"] Q[Query Matrix Q] K[Key Matrix K] V[Value Matrix V] end

subgraph AttentionComputation["Attention Computation"] DotProduct[QK^T Matrix] Scale[÷ √d_k] Softmax[Row-wise Softmax] Weights[Attention Weights] end

subgraph Output["Output"] WeightedSum[V × Weights] Linear[Final Linear Transform] end

T1 & T2 & T3 & Tn --> EmbeddingLayer EmbeddingLayer --> LinearProjections LinearProjections --> AttentionComputation AttentionComputation --> Output

style Input fill:#e1f5fe style Output fill:#e8f5e8 style AttentionComputation fill:#fff3e0 ```

Residual Connections and Layer Normalization

The attention output does not flow directly to subsequent layers. Modern transformer architectures wrap the attention computation within a residual block: the attention output is added to its input (the skip connection) and then layer-normalized. This pattern:

$$\text{output} = \text{LayerNorm}(x + \text{Attention}(x))$$

serves multiple purposes. The skip connection preserves gradient flow during training, allowing deep networks to optimize effectively. Layer normalization stabilizes the scale of activations across the batch and sequence dimensions. And critically, the residual connection means that the attention mechanism can learn to modify information flow without completely transforming the representation—identities are easy to learn, and the network can selectively attend to modify only what matters.

A feed-forward network (typically two linear layers