How Self-Attention Builds Contextual Representations in Transformers
How Self-Attention Builds Contextual Representations in Transformers
- Details
- Category: Language & Agentic AI
Modern NLP changed when models began representing a token according to the context in which it appears. The word "bank" should not receive the same final representation in a sentence about a loan and in a sentence about a river. Its meaning depends on the surrounding words and on the relationships between them.
Self-attention provides one mechanism for constructing such context-dependent representations. For every position in a sequence, it calculates which other positions should contribute information and how strongly their representations should be combined.
This article explains the core operation without treating it as a black box. The objective is to build a precise mental model of self-attention, understand the roles of queries, keys, and values, and identify what an attention matrix can and cannot tell us about a model.
From static embeddings to contextual representations
Methods such as Word2Vec and GloVe represented words as points in a continuous vector space. This was a major improvement over sparse representations such as one-hot vectors because distances and directions in the embedding space could capture useful statistical relationships.
These methods typically assigned one main learned vector to each vocabulary entry. As a result, the initial representation of "bank" did not change between different sentences. A downstream model could still use surrounding words to resolve the meaning, but the embedding itself was not context-dependent.
Contextual models changed this arrangement. They still begin with token embeddings, usually obtained from a learned lookup table, but they repeatedly transform those embeddings using information from the sequence. The hidden state associated with a token therefore depends on the sentence in which it occurs.
Models such as ELMo demonstrated the value of deep contextualized word representations, while BERT showed how bidirectional Transformer layers could produce representations conditioned on both left and right context.
The important distinction is:
- A static embedding represents a vocabulary item.
- A contextual representation represents a token occurrence inside a particular sequence.
An opening example: what does "it" refer to?
Consider the sentence:
"The vessel changed course because it detected an obstacle."
The token "it" is not sufficiently informative in isolation. To represent it usefully, the model needs information from other parts of the sentence. In this example, "vessel" is a plausible antecedent because a vessel can detect an obstacle and change course.
A trained Transformer may represent this relationship by allowing the position associated with "it" to collect information from the position associated with "vessel". Some attention heads may emphasize this dependency, while other heads may capture different relationships, such as local syntax or phrase boundaries.
This does not mean that one attention weight contains the complete grammatical interpretation. Meaning is distributed across heads, layers, residual connections, and feed-forward transformations. Self-attention supplies a path through which information can move; the complete model determines how that information is used.
A useful model: project, compare, and combine
Self-attention can be understood as a three-stage operation:
- Project every token representation into query, key, and value vectors.
- Compare each query with all permitted keys.
- Use the resulting weights to combine the corresponding values.
The familiar descriptions of queries, keys, and values are analogies rather than literal definitions:
- Query: the representation used by a position to search for relevant information.
- Key: the representation used to determine whether another position matches that search.
- Value: the information transferred when a position receives attention.
The projections are learned during training. A token does not contain one permanent query or key. Each layer and attention head constructs its own projections for the current hidden states.
The scaled dot-product attention equation
Assume that a sequence contains \(n\) tokens. Let:
$$X \in \mathbb{R}^{n \times d_{model}}$$
represent the input hidden states. Each row corresponds to one token position. In a Transformer, these states also need some form of positional information because self-attention alone does not encode token order.
The query, key, and value matrices are created using learned projections:
$$Q = XW_Q$$
$$K = XW_K$$
$$V = XW_V$$
where \(W_Q\), \(W_K\), and \(W_V\) are parameter matrices learned during training.
Scaled dot-product attention is then defined as:
$$Attention(Q, K, V) = softmax\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$
Here:
- \(QK^T\) contains the pairwise compatibility scores between queries and keys,
- \(d_k\) is the dimensionality of each query and key,
- \(M\) is an optional attention mask,
- the softmax function is applied independently to every row,
- multiplication by \(V\) produces a weighted combination of value vectors.
For a single attention head, the element in row \(i\) and column \(j\) of the attention matrix indicates how strongly position \(i\) mixes the value from position \(j\). It should not automatically be interpreted as the importance of token \(j\) to the model's final prediction.
Why divide by the square root of the key dimension?
The scaling factor \(\sqrt{d_k}\) prevents dot products from becoming excessively large as the vector dimension increases.
Suppose the components of a query and a key are independent random variables with mean zero and variance one. Their dot product is a sum of \(d_k\) products:
$$q \cdot k = \sum_{r=1}^{d_k} q_r k_r$$
Under these simplifying assumptions, the variance of the dot product grows approximately with \(d_k\). Large score magnitudes can push the softmax function toward highly concentrated distributions in which one position receives almost all of the weight. In those regions, gradients may become very small.
Dividing by \(\sqrt{d_k}\) keeps the scale of the scores more stable before the softmax operation. It does not guarantee uniform attention, nor is uniform attention the objective. It improves the numerical conditions under which the attention distribution is learned.
Masks determine which relationships are permitted
The basic equation compares every query with every key, but not every architecture allows unrestricted access to the sequence.
Bidirectional attention
Encoder models such as BERT generally allow a token to attend to positions on both sides. The representation of a word can therefore depend on preceding and following context.
Causal attention
Autoregressive language models use a causal mask. When predicting the next token, a position may attend only to itself and earlier positions. Future tokens are blocked by adding a very large negative value, conceptually \(-\infty\), to the corresponding scores before applying softmax.
Padding masks
Batches often contain sequences of different lengths. Shorter sequences are padded to a common size, and a padding mask prevents those artificial positions from contributing to attention.
Masking is not an implementation detail. It defines which information paths exist in the model and therefore changes the learning problem.
Self-attention does not know token order by itself
The attention calculation operates on sets of query-key comparisons. Without additional position information, permuting the input tokens produces a corresponding permutation of the outputs. The mechanism can compare content, but it cannot distinguish "the dog chased the cat" from "the cat chased the dog" using token identity alone.
Transformer architectures solve this by injecting information about position. Different models use absolute position embeddings, sinusoidal encodings, relative positions, or rotary position representations. The exact method varies, but the purpose is the same: make token order available to the network.
This leads to a more precise statement:
Self-attention models relationships between token representations, while positional mechanisms tell the model where those tokens occur.
A NumPy implementation
The following example implements one unmasked attention head. It includes separate projections for queries, keys, and values.
The token representations and projection matrices are generated randomly with a fixed seed. They are synthetic and have not been trained on language. The example demonstrates the matrix operations and output dimensions, not a learned relationship between "it" and "vessel".
import numpy as np
def softmax(values, axis=-1):
shifted = values - np.max(values, axis=axis, keepdims=True)
exponentials = np.exp(shifted)
return exponentials / exponentials.sum(axis=axis, keepdims=True)
def scaled_dot_product_attention(queries, keys, values, mask=None):
key_dimension = queries.shape[-1]
scores = queries @ keys.T / np.sqrt(key_dimension)
if mask is not None:
scores = np.where(mask, scores, -np.inf)
attention_weights = softmax(scores, axis=-1)
context = attention_weights @ values
return context, attention_weights
random_generator = np.random.default_rng(42)
tokens = [
"the",
"vessel",
"changed",
"course",
"because",
"it",
"detected",
"obstacle",
]
model_dimension = 16
key_dimension = 8
value_dimension = 8
token_states = random_generator.normal(
size=(len(tokens), model_dimension)
)
query_projection = random_generator.normal(
scale=1 / np.sqrt(model_dimension),
size=(model_dimension, key_dimension),
)
key_projection = random_generator.normal(
scale=1 / np.sqrt(model_dimension),
size=(model_dimension, key_dimension),
)
value_projection = random_generator.normal(
scale=1 / np.sqrt(model_dimension),
size=(model_dimension, value_dimension),
)
queries = token_states @ query_projection
keys = token_states @ key_projection
values = token_states @ value_projection
context, attention_weights = scaled_dot_product_attention(
queries,
keys,
values,
)
print(f"attention_weights_shape={attention_weights.shape}")
print(f"context_shape={context.shape}")
print(
"row_sums="
f"{np.round(attention_weights.sum(axis=1), 6)}"
)
attention_weights_shape=(8, 8)
context_shape=(8, 8)
row_sums=[1. 1. 1. 1. 1. 1. 1. 1.]
The attention matrix has shape \(8 \times 8\) because the sequence contains eight token positions. Each row describes how one query distributes its weight across eight keys.
The rows sum to one because softmax converts each row of scores into a probability-like distribution. These values are non-negative and normalized, although they should not be interpreted as calibrated probabilities of linguistic relationships.
The context matrix has one output row for each input position. Every output row is a weighted mixture of the value vectors.
What the small example proves - and what it does not
The implementation demonstrates that:
- queries and keys produce an \(n \times n\) score matrix,
- softmax normalizes each row independently,
- attention weights determine how values are mixed,
- the output contains one contextualized vector per token position.
It does not demonstrate that the model understands the sentence. The parameters are random, so any apparent pattern in the attention matrix is accidental. Meaningful linguistic behavior appears only after the parameters have been optimized using a training objective and suitable data.
This distinction matters when presenting small attention demonstrations. A matrix produced by untrained random vectors illustrates the computation, not the semantic capabilities of a trained Transformer.
Why Transformers use multiple attention heads
A single attention head produces one set of query-key comparisons. Multi-head attention repeats this operation using several independent sets of learned projections:
$$head_h = Attention(XW_Q^{(h)}, XW_K^{(h)}, XW_V^{(h)})$$
The head outputs are concatenated and projected back into the model dimension:
$$MultiHead(X) = Concat(head_1, \ldots, head_H)W_O$$
Different heads can learn different interaction patterns. One head may emphasize nearby tokens, another may connect syntactically related positions, and another may distribute information more broadly. These interpretations are tendencies observed in some trained models, not fixed roles assigned in advance.
A Transformer layer also contains components that are absent from the small NumPy example:
- residual connections,
- normalization,
- a position-wise feed-forward network,
- dropout or other regularization mechanisms,
- multiple stacked layers.
The contextual representation produced by a model is therefore the result of the entire computational path, not attention alone.
Attention weights are not a complete explanation
Attention visualizations are attractive because they convert an internal matrix into an intuitive heatmap. A strong connection between two tokens may reveal a plausible interaction learned by a particular head.
However, an attention matrix does not show the complete decision process. Several factors limit its interpretation:
- the values being mixed may contain information from earlier layers,
- different heads can carry different or redundant signals,
- residual connections allow information to bypass the attention output,
- feed-forward layers transform the representation after attention,
- different attention distributions can sometimes produce similar outputs.
Research on whether attention provides an explanation has produced a more nuanced conclusion than either "attention always explains the model" or "attention is never useful". Attention weights can support diagnostics and hypothesis generation, but their explanatory value depends on the model, task, definition of explanation, and validation method.
For high-stakes interpretation, an attention heatmap should be combined with other analyses, such as input perturbation, gradient-based methods, counterfactual tests, error analysis, and controlled interventions.
How self-attention relates to real NLP systems
Classification, extraction, and routing
Many production NLP tasks involve classifying support requests, routing documents, detecting intent, extracting entities, or identifying relevant passages. These tasks often depend on relationships between words rather than on isolated keywords.
For example, the phrase "the payment was not approved" differs from "the payment was approved". A contextual model can combine the negation with the relevant event instead of treating both sentences as nearly identical collections of terms.
Summarization and generation
In generation tasks, self-attention allows each output position to combine information from the available context. Causal masking ensures that the model cannot access tokens that have not yet been generated.
Longer context windows make more tokens available, but they do not guarantee that the model will use all of them reliably. Irrelevant evidence, repeated passages, conflicting instructions, and weak document structure can still degrade the result.
Retrieval-augmented generation
Retrieval and self-attention should not be treated as the same mechanism.
In a typical RAG system, an external retriever searches a document collection and selects passages that appear relevant to the query. The selected passages are then inserted into the language model's context. Self-attention operates inside the model, combining information from the tokens that the retriever supplied.
Both stages may use vector similarity, but they solve different problems:
- Retrieval decides which external documents or chunks enter the context.
- Self-attention decides how token representations interact inside that context.
A strong generator cannot fully compensate for missing or irrelevant evidence. RAG quality therefore depends on document preparation, chunking, metadata, embedding quality, retrieval strategy, reranking, prompt construction, and evaluation - not on self-attention alone.
Domain-specific language
Technical domains contain specialized terminology, abbreviations, identifiers, and recurring document structures. A general-purpose model may represent common language well while handling domain expressions inconsistently.
Self-attention gives the model the capacity to combine domain terms with their context, but useful performance still depends on tokenization, training data, task supervision, retrieval quality, and evaluation on representative examples.
Common implementation and interpretation errors
Confusing an available context with an effectively used context
A token may technically be inside the context window without having a reliable effect on the output. Context length specifies capacity, not guaranteed recall or reasoning quality.
Ignoring positional information
Self-attention without a positional mechanism cannot represent sequence order correctly. Removing or misconfiguring position information changes the model into a system that primarily processes an unordered set of token representations.
Using the wrong mask
An incorrect causal mask can expose future tokens during training. An incorrect padding mask can allow artificial padding positions to influence the representation. Both errors may produce apparently valid tensor shapes while invalidating the model behavior.
Treating a heatmap as a causal explanation
A visually strong attention weight does not prove that the associated token caused the prediction. The relationship should be tested by changing, removing, or intervening on the relevant input or internal component.
Ignoring tokenization
Models process tokens rather than human-defined words. Rare terminology, product identifiers, units, source code, and abbreviations may be split into several pieces. This changes sequence length and affects how information is distributed across positions.
Evaluating the model without evaluating the system
Good benchmark performance does not guarantee a useful production pipeline. Classification thresholds, document quality, retrieval errors, missing metadata, latency, and domain drift can dominate the final result.
Key conclusions
- Contextual models produce a different hidden representation for a token depending on the sequence in which it appears.
- Self-attention constructs those representations by comparing learned queries and keys and combining learned values.
- The scaling factor keeps query-key scores in a more stable range before softmax.
- Masks control which token interactions are permitted, while positional mechanisms provide information about sequence order.
- Multi-head attention is only one component of a Transformer layer; residual paths and feed-forward transformations also shape the result.
- Attention matrices can support analysis, but they should not be treated as complete causal explanations.
- In RAG, retrieval selects external evidence and self-attention processes the evidence placed inside the model context.
References
- Mikolov, T., Chen, K., Corrado, G., and Dean, J. (2013). Efficient Estimation of Word Representations in Vector Space.
- Pennington, J., Socher, R., and Manning, C. (2014). GloVe: Global Vectors for Word Representation.
- Vaswani, A., et al. (2017). Attention Is All You Need.
- Peters, M. E., et al. (2018). Deep Contextualized Word Representations.
- Devlin, J., Chang, M. W., Lee, K., and Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.
- Jain, S., and Wallace, B. C. (2019). Attention is not Explanation.
- Wiegreffe, S., and Pinter, Y. (2019). Attention is not not Explanation.
- Lewis, P., et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.