Self-Attention Concept

In the previous chapter, we explored attention in encoder-decoder models where the decoder attends to encoder states. This is cross-attention: one sequence attending to a different sequence. But what if a sequence could attend to itself? This simple shift in perspective gives rise to self-attention, the mechanism that powers transformer models and modern language AI.

Self-attention allows each position in a sequence to directly interact with every other position. Instead of processing tokens one at a time through recurrent connections, self-attention computes relationships between all pairs of tokens simultaneously. This architectural change enables parallelization during training and captures long-range dependencies without the vanishing gradient problems that plague RNNs.

To understand why self-attention matters so deeply, consider what it means for a word to have meaning. The word "run" can mean a physical act of running, a computer program execution, a hole in a stocking, or the flow of a color dye. None of these interpretations are accessible from the word alone. Meaning emerges from context, from the words that surround it, from the syntactic role it plays, from the broader discourse it inhabits. Any model that processes words independently, assigning each a fixed representation regardless of its surroundings, is fundamentally limited. Self-attention is the mechanism that allows a model to build truly contextual representations by letting every word examine every other word and ask: "How relevant are you to understanding me?"

The shift from cross-attention to self-attention may sound like a minor technical adjustment, but it changes how we model language. Cross-attention is fundamentally about translation: a decoder reaching backward to an encoder to retrieve relevant source information. Self-attention is about comprehension within a single sequence: every part of the text simultaneously building its meaning by consulting every other part. This is closer to how we understand language. When you read the sentence "The trophy didn't fit in the suitcase because it was too big," you don't process each word and forget the previous ones. You hold the whole sentence in mind, and when you hit the word "it," you immediately search backward to find the right referent. Self-attention formalizes this process computationally.

The story of self-attention is also a story about the limitations of the architectures that came before it. Recurrent neural networks process sequences one token at a time, passing a hidden state forward. This sequential design means that two distant tokens can only interact through a long chain of intermediate states, and each step risks losing or corrupting information. Convolutional models can aggregate local neighborhoods of tokens, but connecting tokens that are far apart requires stacking many layers. Self-attention sidesteps both problems by computing direct connections between any two positions in a single operation. That direct access lets a model learn patterns that require reasoning about long-range structure.

In encoder-decoder attention, we have two distinct sequences: the encoder outputs serve as keys and values, while the decoder state provides the query. The decoder asks, "Which parts of the input are relevant to what I'm generating right now?"

Self-attention simplifies this setup. Instead of attending across two different sequences, a single sequence attends to itself. Every token in the sequence serves simultaneously as a query, a key, and a value. The question becomes: "For each token in this sequence, which other tokens are most relevant to understanding its meaning in context?"

Consider the sentence "The animal didn't cross the street because it was too tired." To understand what "it" refers to, you need to connect it back to "animal" rather than "street." Self-attention enables this by allowing "it" to attend strongly to "animal," building a representation that captures this coreference relationship.

The key insight is that the meaning of a word depends heavily on its context. The word "bank" means something different in "river bank" versus "bank account." Self-attention lets each word gather information from surrounding words to disambiguate and enrich its representation.

Think of self-attention as a mechanism where every word in a sentence gets to vote on how much it contributes to every other word's meaning. The word "tired" votes heavily for "animal" when "it" is being interpreted, because "tired" is more naturally predicated of animals than streets. The word "cross" votes for "street" because streets are things you cross. All of this happens simultaneously, in a single computational pass over the sequence. The result is a set of enriched representations where each word's embedding combines its inherent meaning with its meaning in this specific sentence.

The transition from cross-attention to self-attention also has architectural implications. In an encoder-decoder model with cross-attention, the encoder and decoder are separate modules. The encoder reads the full source sequence, and the decoder generates the target sequence one token at a time, querying the encoder at each step. Self-attention enables a different design: a single stack of layers that can read a sequence, build increasingly rich representations, and either produce a classification or generate output. This is the architecture underlying BERT, GPT, and the entire modern transformer ecosystem.

The defining characteristic of self-attention is all-pairs interaction. For a sequence of $n$ tokens, self-attention computes interactions between all $n^2$ pairs of positions. Why $n^2$? Each of the $n$ tokens must compute its relationship with every other token, including itself. Token 1 interacts with tokens 1 through $n$ (that's $n$ interactions), token 2 does the same (another $n$ interactions), and so on for all $n$ tokens, giving us $n \times n = n^2$ total interactions.

This is fundamentally different from recurrent models, which process tokens sequentially and must pass information through intermediate states.

In an RNN, if you want information from position 1 to reach position 100, it must flow through 99 intermediate hidden states. Each step risks losing or distorting information. Self-attention eliminates this problem by computing direct connections between any two positions.

The all-pairs property gives self-attention its expressive power. Consider trying to model long-distance syntactic agreement. In English, a subject and verb must agree in number: "The keys to the cabinet are on the table" versus "The key to the cabinets is on the table." The head noun "keys" or "key" determines the verb form, but in a complex sentence with intervening phrases, the distance between them can be substantial. An RNN must carry the number information through every intervening token without losing it, relying on the hidden state to serve as a memory buffer. Self-attention can compute a direct edge from the subject to the verb, regardless of distance.

Similarly, consider coreference: understanding that "it" in "The cat knocked over the vase because it was clumsy" refers to the cat, not the vase. The model needs to compare "it" to all potential antecedents simultaneously, evaluate which one fits semantically, and link the pronoun's representation to the correct noun's representation. An RNN would need to carry candidate antecedent information through the hidden state until the pronoun is encountered. Self-attention computes this comparison directly.

Notice that all-pairs interaction implies that self-attention is permutation equivariant without additional components. If you shuffled the order of tokens in the input, you would get the same attention weights between the same pairs of tokens, just in a different arrangement. This is why positional encodings are necessary: self-attention itself contains no positional information, and we must add it externally. We will return to this point in later sections and chapters, but the result follows directly from the all-pairs design.

The contrast is striking. In the RNN heatmap (left), path lengths grow with distance: connecting position 0 to position 9 requires 9 steps, shown by the dark red corners. The self-attention heatmap (right) is nearly uniform: every position connects to every other position in exactly 1 step (with 0 for self-connections on the diagonal). This constant path length helps capture long-range dependencies without information degradation.

The path length difference has a deeper implication for gradient flow during training. In an RNN trained with backpropagation through time, gradients must also flow backward through these same long paths. If a gradient signal needs to travel from position 100 back to position 1 to update early weights, it must pass through 99 matrix multiplications, each potentially shrinking it toward zero (vanishing gradient) or inflating it toward infinity (exploding gradient). LSTM and GRU cells were invented specifically to alleviate this problem, and they help substantially, but they do not eliminate it. Self-attention, because it computes direct connections, gives gradients equally short paths in both the forward and backward passes. The model can update its parameters based on distant dependencies just as easily as local ones.

Self-attention excels at building contextual representations. A word embedding from Word2Vec or GloVe assigns the same vector to "bank" regardless of context. Self-attention creates contextualized embeddings where each token's representation incorporates information from surrounding tokens.

The process works as follows. Each token starts with an initial embedding. Through self-attention, each token gathers information from all other tokens, weighted by relevance. The output is a new set of representations where each token "knows about" the full sequence context.

This contextual awareness enables several capabilities that static embeddings lack:

• Word sense disambiguation: The representation of "bank" differs based on surrounding words
• Coreference resolution: Pronouns can gather information from their antecedents
• Syntactic awareness: Verbs can connect to their subjects and objects regardless of distance
• Semantic composition: Phrases and sentences build meaning from their constituent words

To appreciate how powerful this is, it helps to compare it against what came before. Pre-transformer NLP relied on static word embeddings like Word2Vec and GloVe, which encode a word's meaning as a fixed vector that never changes regardless of context. These embeddings capture distributional statistics: "bank" ends up somewhere between "river" and "money" because it co-occurs with both kinds of words in training data. But this is a single point in embedding space, a compromise between two entirely different meanings, not a true representation of either. When the model processes "river bank," it uses the same "bank" vector as when it processes "bank account." Any downstream task must somehow disentangle these meanings using only context provided by neighboring tokens.

Self-attention solves this problem directly. When a transformer processes "river bank," the representation of "bank" is computed as a weighted sum of all tokens in the sentence, with "river" receiving very high weight. The resulting vector is shifted toward the "river"-related region of embedding space. When the same model processes "bank account," the high-weight neighbors are "account," "deposit," "money," and the representation shifts in the financial direction. The same input token produces different output representations depending on context. This is what makes transformer embeddings so much more useful for downstream tasks.

The mechanism also enables compositional understanding in a way that static embeddings cannot. Consider the phrase "not happy." In static embeddings, you might try to negate the "happy" vector, but there is no obvious way to do this without training a separate component. In self-attention, "not" directly attends to "happy" with high weight, and "happy" attends back to "not." The resulting representation of "happy" in this context incorporates information from "not," enabling the model to learn that negated words have different downstream behavior. This is the basis for handling composition, including cases such as negation and intensification.

Let's visualize how self-attention might weight different tokens when processing a simple sentence.

The visualization shows that when building a representation for "sat," the model attends most strongly to "cat" (the subject performing the action) and "mat" (the location). Function words like "the" and "on" receive less attention. This weighted combination creates a rich representation that captures the verb's relationship to other sentence elements.

This pattern, a verb attending to its syntactic arguments, is precisely the kind of structural awareness that transformers develop during training. The attention weights are not hand-coded. They emerge from optimization: the model learns that representing a verb well requires knowing who performed the action (subject) and what it was performed on or where it occurred (object or adjunct). In a real trained model, you would see systematic patterns: determiners attending to their head nouns, adjectives attending to the nouns they modify, verbs forming attention arcs to their subjects and objects. The attention mechanism provides a soft, differentiable way to implement the kinds of syntactic relationships that linguists have described for decades.

Now that we understand why self-attention is useful, let's examine how it works. The core idea is elegantly simple: to build a contextual representation for any token, we take a weighted average of all tokens in the sequence. The weights reflect how relevant each token is to the one we're updating.

Think of it like asking every word in a sentence: "How much should I pay attention to you when trying to understand myself?" Words that are more relevant get higher weights, and their information contributes more to the final representation.

The elegance of this design is that relevance is computed from the same embeddings that we are trying to enrich. There are no separate relevance signals or external knowledge sources. The model learns, from data, which kinds of tokens should attend to which other kinds of tokens. A verb learns to attend to nouns. A pronoun learns to attend to candidate antecedents. An adjective learns to attend to the noun it modifies. All of this emerges from training on text, supervised only by the task objective.

In practice, this means that self-attention is a form of dynamic, input-dependent aggregation. Unlike a convolution, which always applies the same kernel to local neighborhoods regardless of content, self-attention computes different aggregation weights for every input. Two sentences with different words will produce entirely different attention patterns. The same word in two different sentences will aggregate information from different neighbors with different weights. This input-dependence is what allows self-attention to be sensitive to context in a way that convolutions and bag-of-words models cannot be.

Self-attention follows a consistent three-stage pattern:

1. Compute similarity scores: For each pair of positions, calculate a number indicating how relevant one is to the other. High scores mean strong relevance.

2. Normalize to weights: Raw scores can be any real number. We need to convert them into proper weights that sum to 1, so we can interpret them as "how much attention to pay." The softmax function handles this conversion.

3. Aggregate values: Finally, compute a weighted sum of all positions' representations. Positions with higher weights contribute more to the output.

This pipeline runs for every position simultaneously, producing a new set of contextual embeddings where each token has gathered information from the entire sequence.

These three stages map onto a natural information-theoretic story. In stage one, we are asking a question: "How related are these two tokens?" In stage two, we are converting raw relatedness scores into a probability distribution over which tokens to consult. In stage three, we are reading the answer: a blend of all tokens' information, weighted by how likely each one is to be the right source.

The beauty of this formulation is that it is fully differentiable. Every step can be computed using standard matrix operations, and gradients flow cleanly through all three stages. This is what makes it possible to train self-attention end-to-end with gradient descent. The model does not need to know in advance which tokens should attend to which. It learns this from data by adjusting the embeddings so that tokens that should attend to each other have high dot products.

Let's formalize this intuition. Suppose we have a sequence of $n$ tokens, each represented by a $d$-dimensional embedding vector. We can write the input as a matrix:

where each $\mathbf{x}_i \in \mathbb{R}^d$ is the embedding for token $i$.

Stage 1: Similarity Scores

How do we measure relevance between two tokens? The simplest approach uses the dot product. If two embedding vectors point in similar directions, their dot product is large and positive. If they're orthogonal (unrelated), the dot product is zero. This gives us a natural measure of similarity.

The dot product is not the only possible similarity measure. You could use cosine similarity, a learned bilinear form, or even a neural network. But the dot product has the virtue of simplicity, efficiency, and a clean geometric interpretation. It is also naturally suited to the way embeddings are trained: if you want two tokens to have high similarity under dot product, you train their embeddings to point in the same direction. This is exactly how the model learns to encode semantic relationships.

For positions $i$ and $j$, the similarity score is:

where:

• $s_{ij}$: the raw similarity score between positions $i$ and $j$
• $\mathbf{x}_i \cdot \mathbf{x}_j$: the dot product of the two embedding vectors
• $d$: the embedding dimension
• $x_{i,k}$: the $k$-th component of the embedding at position $i$

The dot product captures semantic similarity: tokens with similar meanings tend to have similar embeddings, producing high dot products. This is only true if the embeddings have been trained appropriately, but in practice, transformer embeddings rapidly learn to encode semantic similarity in ways that the dot product can exploit.

Stage 2: Softmax Normalization

Raw dot products can be any real number, positive or negative, with no upper bound. To use them as weights for averaging, we need to transform them into a probability distribution. The softmax function accomplishes this:

where:

• $\alpha_{ij}$: the attention weight from position $i$ to position $j$
• $\exp(s_{ij})$: the exponential function applied to the similarity score. This keeps positivity
• $\sum_{k=1}^{n} \exp(s_{ik})$: the normalizing constant that makes all weights from position $i$ sum to 1

Why softmax? It has two useful properties:

1. Positivity: The exponential function maps any real number to a positive value, so all weights are positive.
2. Normalization: Dividing by the sum ensures $\sum_{j=1}^{n} \alpha_{ij} = 1$, giving us valid weights for a weighted average.

Additionally, softmax amplifies differences: if one score is much higher than others, its weight dominates. This allows the model to focus sharply on the most relevant tokens when appropriate.

There is an important practical consideration here: the scale of the dot products. If the embeddings have dimension $d$ and their components are approximately unit variance, then the dot product of two random embeddings will have variance approximately $d$. For large $d$, this means dot products can be very large in magnitude, pushing the softmax into a regime where its gradient is nearly zero. This is the motivation for the scaled dot product that you will encounter in the full query-key-value formulation in the next chapter: dividing by $\sqrt{d}$ keeps the dot products in a reasonable range where softmax remains informative. For now, the basic formulation is sufficient to understand the core mechanism.

The softmax function also has an interesting normalization property: it computes a local competition. Each token $i$ does not have a global attention budget that it must divide among all other tokens in the sequence (which would be a simplex constraint across all $i$ and $j$). Instead, for each query position $i$, the weights from $i$ to all key positions $j$ must sum to 1. This means position $i$ always attends to something, even if no token is truly relevant. In practice, when no token is highly relevant, the attention spreads uniformly across all tokens, and the output is close to the average of all embeddings.

Stage 3: Weighted Aggregation

With attention weights in hand, we compute the output representation for each position as a weighted sum of all input embeddings:

where:

• $\mathbf{y}_i$: the output representation for position $i$, now enriched with contextual information
• $\alpha_{ij}$: the attention weight determining how much position $j$ contributes to position $i$'s output
• $\mathbf{x}_j$: the input embedding at position $j$

This formula is the heart of self-attention. Each output $\mathbf{y}_i$ is a blend of all inputs, with the blending proportions determined by the learned attention weights. Tokens that are highly relevant to position $i$ contribute strongly; irrelevant tokens contribute little.

Notice that the output lives in the same vector space as the input: $\mathbf{y}_i \in \mathbb{R}^d$, just like $\mathbf{x}_i$. This is important for stacking multiple layers of self-attention. The output of one layer can serve directly as the input to the next, with no change in dimension. Each successive layer can build on the contextual representations computed by the previous one, creating increasingly abstract and task-relevant representations. This is exactly how transformers work: multiple self-attention layers applied in sequence, each one refining the representations produced by the previous layer.

The complete self-attention computation flows naturally from these three stages. For every position $i$ in the sequence:

1. Compute dot products with all positions: $s_{ij} = \mathbf{x}_i \cdot \mathbf{x}_j$ for $j = 1, \ldots, n$
2. Apply softmax to get attention weights: $\alpha_{ij} = \text{softmax}(s_{i1}, s_{i2}, \ldots, s_{in})_j$
3. Compute weighted sum: $\mathbf{y}_i = \sum_{j=1}^{n} \alpha_{ij} \mathbf{x}_j$

Because each position's computation is independent of the others, we can process all $n$ positions in parallel. This parallelism is what makes self-attention so efficient on modern GPU hardware.

In matrix form, the entire computation collapses into three operations: a matrix multiplication to compute scores, a row-wise softmax, and another matrix multiplication to compute outputs. The notation also exposes the implementation: the entire self-attention operation can be executed as two large matrix multiplications, which GPUs and TPUs are specifically designed to perform as fast as possible. The theoretical $O(n^2 d)$ complexity, which sounds expensive, is quite manageable in practice because modern accelerators can execute these matrix operations with extreme parallelism.

Let's translate this mathematical framework into code. We'll implement a simple self-attention function that takes a sequence of embeddings and returns the contextual outputs along with the attention weights.

The attention weight matrix is $n \times n$, where entry $(i, j)$ tells us how much position $i$ attends to position $j$. Each row sums to exactly 1.0, confirming that softmax produces valid probability distributions. The output has the same shape as the input: each of our 4 tokens now has a new 8-dimensional representation that incorporates information from all other tokens.

A heatmap provides an intuitive view of self-attention patterns. Rows represent query positions (which token is gathering information), and columns represent key positions (which tokens are being attended to). Darker cells indicate stronger attention.

In this random example, we see some diagonal emphasis where tokens attend to themselves, plus distributed attention to other positions. In trained models, these patterns become meaningful: related words attend strongly to each other, syntactic structures emerge, and semantic relationships become visible.

Researchers have spent considerable effort interpreting what trained attention patterns mean. In models like BERT, some attention heads learn to attend to the next token, some to the previous token, some to the sentence separator tokens. Other heads develop more linguistically interpretable patterns: there are heads that attend from verbs to subjects, from pronouns to antecedents, from nouns to modifying adjectives. This does not mean that each head has a single clean function. In practice, the patterns are messy and context-dependent. But the existence of interpretable patterns suggests that self-attention does learn to exploit syntactic and semantic structure.

The diagonal is informative. Random embeddings tend to produce some diagonal emphasis because the dot product of a vector with itself is always larger than with a random other vector. In trained models, the diagonal is often strong, meaning tokens attend significantly to themselves. This makes sense: a token's own embedding contains the most concentrated information about that token. The off-diagonal weights then represent how much additional context each token gathers from its neighbors.

Self-attention and recurrence represent fundamentally different approaches to sequence modeling. Understanding their trade-offs clarifies why transformers have largely replaced RNNs for most NLP tasks.

Parallelization: RNNs process tokens sequentially because each hidden state depends on the previous one. Self-attention computes all pairwise interactions simultaneously, enabling massive parallelization on GPUs. This difference dramatically accelerates training on modern hardware.

Long-range dependencies: In an RNN, information must flow through many timesteps to connect distant positions. Gradients can vanish or explode along this path. Self-attention connects any two positions directly, making it easier to learn long-range dependencies.

Computational complexity: Self-attention computes $n^2$ pairwise interactions for a sequence of length $n$, where $n$ is the number of tokens. RNNs have linear complexity, meaning their computational cost grows proportionally to $n$. For very long sequences, the quadratic cost of self-attention becomes prohibitive: doubling the sequence length quadruples the computation. This motivates efficient attention variants.

Positional information: RNNs inherently encode position through their sequential processing. Self-attention treats all positions symmetrically and requires explicit positional encodings to distinguish word order. We'll cover positional encodings in a later chapter.

It is worth pausing on the parallelization advantage, because it is easy to underestimate how important it is. The difference between $O(n)$ and $O(n^2)$ operations sounds like RNNs should be faster, but this analysis ignores the sequential bottleneck. An RNN processing a sequence of length $n$ must execute $n$ sequential steps, because step $t$ depends on the output of step $t-1$. No amount of hardware parallelism can overcome this dependency. A modern GPU with thousands of cores cannot use most of them to process a single RNN sequence, because at each step, only one computation is active.

Self-attention, by contrast, computes all $n^2$ pairwise interactions simultaneously. With sufficient hardware, these can all be computed in parallel, making the wall-clock time essentially constant regardless of sequence length (up to memory constraints). In practice, the $n^2$ operations are organized as two matrix multiplications, each of which maps perfectly onto GPU tensor cores. This is why transformers can be trained on sequences of thousands of tokens in reasonable time, while training deep RNNs on sequences longer than a few hundred tokens was extremely slow.

The practical consequence of this difference was the scaling revolution. Transformers could be trained on much larger datasets in less time than RNNs, because more of the compute could be utilized in parallel. This, combined with the architectural improvements from self-attention, is why models like BERT and GPT-2 represented such dramatic jumps in capability compared to their RNN-based predecessors.

Memory requirements: Self-attention must store the entire $n \times n$ attention weight matrix during computation, requiring $O(n^2)$ memory. For a sequence of 10,000 tokens, this is 100 million weights, which requires several gigabytes of memory even in 16-bit precision. RNNs require only $O(nd)$ memory, where $d$ is the hidden state dimension, because only the current hidden state needs to be stored. This memory bottleneck is one of the key practical constraints on self-attention and has motivated research into memory-efficient attention variants.

Inductive biases: RNNs have a strong inductive bias toward local, recency-weighted processing. They naturally place more emphasis on recent tokens (the hidden state is more directly influenced by recent inputs) and require explicit mechanisms like LSTMs to maintain long-term information. Self-attention has a weaker, more uniform inductive bias: all positions are treated symmetrically, and the model must learn from data which positions are relevant. This makes self-attention more flexible but also means it requires more data to learn good patterns from scratch.

The quadratic growth of self-attention is dramatic: at sequence length 1000, self-attention requires 1 million operations compared to just 1000 for an RNN. However, self-attention's operations are independent and can run in parallel on GPUs, while RNN operations must be sequential. This trade-off explains why transformers dominate despite higher theoretical complexity: parallelism wins on modern hardware.

To put concrete numbers on this: a typical transformer layer with sequence length 512 and embedding dimension 768 spends roughly 40% of its compute on the attention matrix multiplication and 60% on the feed-forward layers. The quadratic bottleneck is real but manageable for sequences of a few hundred to a few thousand tokens. For the long-context models that have emerged recently, with context windows of 32,000 or 128,000 tokens, efficient attention algorithms (such as FlashAttention, which reorganizes the computation to minimize memory bandwidth rather than total operations) are essential.

The formulas become clearer when we trace through a concrete example step by step. Let's process a three-word sequence, "I love coffee," using tiny 2-dimensional embeddings. With only 2 dimensions, we can easily verify each calculation by hand and visualize what's happening geometrically.

We'll assign each word an embedding that points in a different direction:

The geometric arrangement makes the dot product intuition clear: vectors pointing in similar directions have high dot products, while perpendicular vectors have zero dot product.

First, we compute the dot product between every pair of words. Remember, the dot product $\mathbf{x}_i \cdot \mathbf{x}_j$ measures how much two vectors point in the same direction.

Let's interpret this matrix:

• Diagonal entries (1.00, 0.50, 1.00): These are self-similarities. "I" and "coffee" have unit-length embeddings, so their self-dot-products are 1.0. "love" has length $\sqrt{0.5^2 + 0.5^2} = \sqrt{0.5} \approx 0.71$, giving a self-dot-product of 0.5.

• "I" vs "coffee" (0.00): These vectors are perpendicular (orthogonal), so their dot product is zero. In embedding space, this means they share no directional similarity.

• "love" vs others (0.50): "love" points between "I" and "coffee," giving it moderate similarity to both. The dot product of 0.5 reflects this intermediate relationship.

The asymmetry between "love" and itself (0.50) versus "I" and "coffee" and themselves (1.00) comes from vector length. The dot product is not normalized, so longer vectors have higher self-similarity. This is one reason why the scaled dot product version of self-attention (which we will study in the next chapter) sometimes normalizes embedding lengths. In this toy example, the unnormalized dot product is sufficient to see the main intuition.

Raw dot products aren't suitable weights for averaging because they don't sum to 1. The softmax function maps each row into a probability distribution:

Now each row sums to 1.0, and we can interpret the values as "attention percentages":

• "I" attends most to itself (about 58%), moderately to "love" (31%), and least to "coffee" (11%). This makes sense: "I" had the highest dot product with itself, moderate with "love," and zero with "coffee."

• "love" distributes attention more evenly across all three words (about 33% each). Its embedding lies between the others, giving it similar dot products with everyone.

• "coffee" mirrors "I": it attends mostly to itself, moderately to "love," and barely to "I."

The softmax has transformed our raw similarities into a meaningful attention distribution. Higher similarity scores become higher attention weights, but even zero-similarity pairs get some weight (the exponential of 0 is 1, not 0). This is an important property: softmax never produces exactly zero weights. Every token always attends at least a little to every other token, which prevents the model from completely ignoring any part of the input. The practical implication is that the information flow in self-attention is always fully dense, not sparse.

Notice also the symmetry in this particular example: the attention weight from "I" to "coffee" (0.11) equals the weight from "coffee" to "I" (0.11). This symmetry arises because our scores matrix is symmetric: $s_{ij} = \mathbf{x}_i \cdot \mathbf{x}_j = \mathbf{x}_j \cdot \mathbf{x}_i = s_{ji}$. In the full query-key-value formulation of self-attention (covered in the next chapter), separate learned projections break this symmetry, allowing "I" attending to "coffee" to have a different weight than "coffee" attending to "I." This asymmetry is linguistically meaningful: a pronoun attending to its antecedent is not the same relationship as the antecedent attending to the pronoun.

Finally, we use these attention weights to compute a weighted average of all embeddings for each position:

The visualization captures the essence of self-attention geometrically. Each word's representation moves toward the center of mass of all words, weighted by attention. "I" and "coffee" shift significantly toward each other (via "love"), while "love" barely moves because it was already centrally positioned. After this transformation, all three words carry information about the full context.

This worked example, despite its simplicity, illustrates several important properties of self-attention that generalize to real models.

First, the outputs are always convex combinations of the inputs. Because attention weights are non-negative and sum to 1, the output $\mathbf{y}_i$ is a weighted average: it lies in the convex hull of all input embeddings. This means self-attention cannot produce representations that are "outside" the range of the inputs; it can only blend and interpolate. This is both a strength (it is a stable, well-behaved operation) and a limitation (it cannot extrapolate beyond what is present in the input). The feed-forward layers that follow self-attention in a transformer block are responsible for introducing the non-linearity and capacity to produce outputs that go beyond simple blending.

Second, the transformation is context-dependent. The attention weights, and therefore the output representations, change with every different input sequence. Run the same three-word sequence through a random permutation (say, "coffee I love") and the attention weights will be different, and the output representations will be different. This context-dependence is what makes self-attention fundamentally more powerful than any fixed transformation.

Third, the mechanism is symmetric in the simplified version here, but asymmetric in the full QKV formulation. When queries and keys are computed by separate learned projections (rather than using the raw embeddings for both), the relationship "how much does position $i$ attend to position $j$" can differ from "how much does position $j$ attend to position $i$." This asymmetry is linguistically important and is one reason why the full QKV formulation is more expressive.

While the worked example used random or hand-crafted embeddings, trained transformer models develop rich, interpretable attention patterns. Understanding what these patterns look like in practice helps build intuition for what self-attention is doing in models like BERT and GPT.

In BERT, researchers have visualized and analyzed the attention patterns of individual heads across its twelve layers. Several patterns emerge consistently. Some heads learn to attend predominantly to the next token or the previous token, effectively encoding local context. These heads are similar in function to a bigram model: they give each token information about its immediate neighbor. Other heads attend strongly to the separator tokens (the [CLS] and [SEP] tokens in BERT), which serve as aggregate representations of the entire sequence. Still other heads develop more linguistically structured patterns, with verbs attending to their subjects, determiners attending to their head nouns, and coreferent mentions attending to each other.

The emergence of these patterns is not guaranteed and is not the result of any explicit supervision signal telling the model "this head should learn dependency parsing." It arises entirely from the combination of the attention mechanism, the pretraining objective, and the optimization pressure of fitting a large training corpus. The model discovers that representing language well requires tracking syntactic and semantic relationships, and it distributes this work across its many attention heads.

In GPT-style autoregressive models, the attention patterns have a different character because causal masking prevents each position from attending to future positions. The pattern must be lower-triangular: position $i$ can only attend to positions $1, \ldots, i$. Within this constraint, heads still develop specialized functions: some heads track long-range subject-verb agreement, some attend to the most recent noun phrase (tracking the current topic), and some attend to specific function words that signal discourse structure.

In practice, the learned attention patterns are noisier and more complex than the idealized examples in textbooks. Many heads attend to a diffuse mixture of positions. This contributes to the representation without any single head being primarily responsible for a specific linguistic function. The "division of labor" across heads is more statistical than categorical. But the overall picture is clear: self-attention learns to implement something resembling the kinds of dependency structures that linguists have described, because these structures are what the training data demands.

Self-attention does not operate in isolation. In a transformer, each self-attention layer is followed by a position-wise feed-forward network, layer normalization, and residual connections. Understanding how these components interact helps clarify what self-attention's specific contribution is.

The residual connection is particularly important for understanding self-attention's role. Rather than replacing the input representations entirely, self-attention adds to them. The output of a self-attention layer is added back to its input:

where $\text{SelfAttn}(\mathbf{x}_i, \mathbf{X})$ denotes the self-attention output for position $i$ given the full sequence $\mathbf{X}$. This means each token retains its original representation and supplements it with contextual information gathered from the sequence. The self-attention output is a correction or augmentation, not a replacement.

This residual formulation has important consequences for learning and for interpretation. It means the gradient always has a direct path from the output to any earlier representation, without passing through the attention operation. This prevents vanishing gradients and allows very deep networks to be trained. It also means that in the early layers of a transformer, the representations are still relatively close to the original embeddings, with self-attention providing small contextual corrections. In later layers, the residual additions accumulate, producing representations that are substantially richer than the original embeddings.

The feed-forward network that follows each self-attention layer plays a complementary role. While self-attention blends information across positions, the feed-forward network operates independently on each position, applying the same learned transformation to each token's representation. The feed-forward network introduces non-linearity that self-attention lacks, and it provides the capacity to perform complex transformations on individual token representations after they have been enriched with context.

Layer normalization, applied before or after each sub-layer (depending on the specific architecture variant), keeps the activations in a stable range throughout the network. This stabilizes training in deep transformers: without normalization, the accumulation of residual additions would cause the representations to grow in magnitude, eventually making the softmax in the attention mechanism saturate and gradients vanish.

Together, these components create a system where self-attention handles the cross-position information flow and the feed-forward network handles the per-position transformation. Self-attention is responsible for asking and answering the question "what information from the sequence is relevant to this position?" while the feed-forward network is responsible for asking "given this enriched representation, what feature should I extract or computation should I perform?"

Self-attention transformed NLP by enabling parallel processing and direct long-range connections. The transformer architecture, built on self-attention, powers models such as BERT and GPT. These models achieve state-of-the-art results across virtually all NLP benchmarks.

However, self-attention has significant limitations. The quadratic complexity in sequence length makes it expensive for long documents. Processing a 10,000-token document requires computing 100 million pairwise interactions. This has motivated research into efficient attention variants, including Longformer and BigBird alongside linear-attention mechanisms that reduce complexity while preserving much of self-attention's power.

The quadratic complexity also creates a memory bottleneck. During training, the entire attention weight matrix must be stored for the backward pass. For a batch of 32 sequences, each 2048 tokens long, with 12 attention heads, this is approximately 32 × 12 × 2048 × 2048 × 4 bytes ≈ 6 gigabytes for attention weights alone, before accounting for activations and gradients. This is one reason why very long context models require specialized hardware and training techniques.

Self-attention also lacks inherent positional awareness. Unlike RNNs, which process tokens in order, self-attention treats the sequence as a set. The sentence "dog bites man" would produce the same attention weights as "man bites dog" without explicit positional information. Transformers address this with positional encodings, but the need for this additional component reveals a fundamental limitation of the attention mechanism itself.

The position-agnostic nature of self-attention is also a limitation for generalization. A model trained on sequences up to length 512 may struggle with sequences of length 1024, not because 1024-length sequences are inherently harder, but because the positional encodings for positions 513-1024 were never seen during training. Research into relative positional encodings and rotary position embeddings (RoPE) has addressed some of these generalization issues, and modern models like GPT-4 can handle very long contexts, but the fundamental tension between self-attention's position-agnosticism and the importance of position in language remains an active area of research.

Self-attention is also not well-suited to online or streaming inference. Because the computation requires the full attention matrix, you cannot process a sequence incrementally as tokens arrive without either recomputing from scratch or caching the key and value matrices. The key-value cache is the standard solution for autoregressive generation (used in every GPT-style model), but it requires storing $O(n)$ intermediate representations, which can be memory-intensive for very long sequences.

Despite these challenges, self-attention's benefits outweigh its costs for most NLP applications. The ability to train on massive datasets with full parallelization, combined with the capacity to model arbitrary dependencies, has made transformer-based models the dominant paradigm in modern language AI. The limitations have motivated a rich line of research into efficient attention, and many of the most important advances in the field since 2017 can be traced directly to the need to scale or extend the basic self-attention mechanism.

The impact of self-attention extends beyond NLP. Vision transformers (ViTs) apply self-attention to sequences of image patches, achieving state-of-the-art performance on image classification tasks. Transformers have been applied to protein structure prediction (AlphaFold2), audio synthesis (AudioLM), video understanding, and reinforcement learning. The mechanism's generality, its ability to model arbitrary pairwise relationships in any kind of structured sequence, is what makes it so broadly applicable. Self-attention applies beyond language; it is a general-purpose mechanism for building representations of structured data.

Self-attention is the mechanism that allows a sequence to attend to itself, enabling each position to gather information from all other positions. This simple idea changes how we build language models.

Key takeaways from this chapter:

• Self-attention vs cross-attention: In cross-attention, one sequence attends to another. In self-attention, a sequence attends to itself, with each token serving in a QKV role simultaneously.
• All-pairs interaction: Self-attention computes direct connections between all $n^2$ pairs of positions (where $n$ is the sequence length), eliminating the need for information to flow through intermediate states.
• Contextual representations: By aggregating information from surrounding tokens, self-attention creates embeddings that capture word meaning in context, enabling disambiguation, coreference resolution, and syntactic awareness.
• Parallelization: Unlike recurrent models, self-attention computes all interactions simultaneously, enabling efficient training on modern hardware and allowing the field to scale to much larger models and datasets.
• Quadratic complexity: The all-pairs computation scales as $O(n^2)$, meaning computational cost grows with the square of sequence length $n$. This makes self-attention expensive for very long sequences and has motivated research into efficient attention variants.
• Position agnostic: Self-attention treats sequences as sets, requiring explicit positional encodings to capture word order. Without positional information, permuting the input tokens produces the same attention weights.
• Residual integration: In transformer blocks, self-attention output is added to the input via residual connections, meaning each layer enriches rather than replaces the existing representations.

The next chapter covers how self-attention computes these interactions through QKV projections. These learned linear transformations give self-attention its full expressiveness, allowing the model to compute asymmetric, task-specific attention patterns and to specialize different projections for different aspects of context. Understanding the QKV formulation is essential for implementing and reasoning about transformer models.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about self-attention.