Self-Attention Concept › The Computational Pattern

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:

$$\mathbf{X} = [\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_n]$$

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:

$$s_{ij} = \mathbf{x}_i \cdot \mathbf{x}_j = \sum_{k=1}^{d} x_{i,k} \cdot x_{j,k}$$

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:

$$\alpha_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{n} \exp(s_{ik})}$$

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:

$$\mathbf{y}_i = \sum_{j=1}^{n} \alpha_{ij} \mathbf{x}_j$$

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.