Paweł Labuda Portfolio
  • About me
  • Experience
  • Projects
  • Realizations
  • Blog
  1. You are here:  
  2. Blog
  3. Language & Agentic AI
  4. How Self-Attention Builds Contextual Representations in Transformers
Language & Agentic AI Mar 9, 2026 14 min read

How Self-Attention Builds Contextual Representations in Transformers

  • deep learning
  • large language models
  • rag and information retrieval

How Self-Attention Builds Contextual Representations in Transformers

Details
Category: Language & Agentic AI
  • large language models
  • rag and information retrieval
  • deep learning

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:

  1. Project every token representation into query, key, and value vectors.
  2. Compare each query with all permitted keys.
  3. 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)}"
)
Console output
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.

Multi-head attention with independent query, key, and value projections
Multi-head attention applies independent projections in parallel and combines the resulting head representations.

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

  1. Mikolov, T., Chen, K., Corrado, G., and Dean, J. (2013). Efficient Estimation of Word Representations in Vector Space.
  2. Pennington, J., Socher, R., and Manning, C. (2014). GloVe: Global Vectors for Word Representation.
  3. Vaswani, A., et al. (2017). Attention Is All You Need.
  4. Peters, M. E., et al. (2018). Deep Contextualized Word Representations.
  5. Devlin, J., Chang, M. W., Lee, K., and Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.
  6. Jain, S., and Wallace, B. C. (2019). Attention is not Explanation.
  7. Wiegreffe, S., and Pinter, Y. (2019). Attention is not not Explanation.
  8. Lewis, P., et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.
← Previous article Feature Engineering for Tabular Data: How to Use Target Encoding Without Leakage Next article A Model Can Stay Online and Still Fail: Monitoring Input Drift with PSI →
← Back to Blog

More articles

A Model Can Stay Online and Still Fail: Monitoring Input Drift with PSI

A machine learning service can return successful responses, remain within its latency target, and still become less reliable. No exception is required. The application code may be unchanged while the data reaching the model gradually moves away from the population represented in training and validation.

This creates a gap between software health and model health. Infrastructure monitoring can tell us whether an endpoint is available. It cannot tell us whether the assumptions behind the model still hold.

Read more …

From Notebook to Dependency: Where Data Science Ends and AI Engineering Begins

Data Science and AI Engineering overlap, but they are not interchangeable. The difference becomes visible when a model stops being an analytical result and becomes a dependency that other systems, teams, or customers rely on.

A data scientist may demonstrate that a useful signal exists, define how it should be measured, and estimate whether it generalizes beyond the training sample. An AI engineer must make that signal available under operational constraints: data contracts, interfaces, latency, deployment, observability, recovery, and long-term ownership.

Read more …

Causal Overlap Before Estimation: Diagnosing Support, Trimming, and Estimand Change

Causal inference can fail before an estimator is selected. The problem is not always the regression formula, matching algorithm, weighting method, or model family. It is often whether treated and untreated observations are sufficiently comparable for the intended causal question.

If one treatment group occupies regions of the covariate space that have almost no counterpart in the other group, the estimate starts depending on extrapolation. A model can still return a numerical effect and a narrow confidence interval, but the data may provide little direct support for the comparison.

Read more …

Paweł Labuda

AI engineering portfolio, personal projects, technical notes, and blog.

Be in touch mail pawel.labuda@itvix.pl

All rights reserved 
© 2026 IT Vix
Privacy Policy Cookie Policy
Built as a technical notebook for learning, building, and sharing.

This website uses cookies. Using the website means that you agree.

Privacy Policy Cookie Policy