Transformer Architecture: The Math Behind Attention

The Problem Attention Solves

Language is full of dependencies. In the sentence “The animal did not cross the road because it was tired,” the word “it” refers to “animal.” A model cannot understand the sentence properly unless it connects distant words.

Older models such as RNNs processed tokens in sequence. This made long-range dependencies difficult because information had to travel step by step. Transformers approach the problem differently. They let each token look directly at all other tokens in the sequence.

That “looking” is not emotional attention. It is mathematical weighting. A token assigns scores to other tokens, converts those scores into probabilities, and uses them to mix information from the sequence.

Attention in One Formula

The central formula of attention in the Transformer is:

This formula looks compact, but it contains the main idea of the Transformer. First, compare queries with keys. Second, scale the scores. Third, apply softmax to turn scores into weights. Fourth, use those weights to combine values.

In plain words, attention says: “For every token, calculate how much it should care about every other token, then build a new representation using those importance weights.”

Understanding Q, K, and V

The letters Q, K, and V stand for Query, Key, and Value. They come from the same input tokens but are projected through different learned matrices. This means each token is transformed into three different roles.

A query represents what a token is looking for. A key represents what a token offers for matching. A value represents the actual information that will be passed forward if the token is attended to.

This separation is powerful. A token can be useful for matching in one way and useful for information transfer in another way. The model learns these transformations during training.

Step 1: Dot Products Measure Similarity

The term QKᵀ calculates similarity between every query and every key. If a query and key point in a similar direction in vector space, their dot product is large. If they are unrelated, the dot product is smaller.

Suppose we have five tokens. Each token creates one query and one key. Multiplying Q by Kᵀ creates a 5 × 5 attention score matrix. Each row represents one token asking, “How relevant is every token to me?”

The dot product is efficient because it can be computed as matrix multiplication. This is one reason Transformers train well on modern hardware. The model can compare many tokens in parallel rather than processing them strictly one after another.

Step 2: Why Divide by √dₖ?

The scaling term √dₖ is easy to ignore, but it matters. Here, dₖ is the dimension of the key vectors. As vector dimensions grow, dot products can become large in magnitude.

Large dot products create very sharp softmax outputs. When softmax becomes too sharp, one token may receive almost all the attention while others receive almost zero. That can make gradients small and training less stable.

Dividing by √dₖ keeps the scores in a healthier range. It does not change the basic meaning of attention. It simply prevents the dot-product scores from becoming too extreme as the vector dimension increases.

Step 3: Softmax Turns Scores into Weights

After scaling, the model applies softmax. Softmax converts raw scores into probabilities. Every row of the attention matrix sums to 1. This means each token distributes its attention across the sequence.

If token A strongly relates to token B, the softmax weight for B becomes high. If token C is irrelevant, its weight becomes low. The model then uses these weights to combine value vectors.

This is why attention is often described as a weighted average. The output for each token is not copied from one place. It is a learned mixture of value vectors from relevant tokens.

Step 4: Values Carry the Information

The final multiplication by V creates the attention output. The softmax weights decide how much each value contributes. If a word attends strongly to another word, that other word’s value vector contributes more to the output representation.

This updated representation now contains context. The token is no longer represented in isolation. It carries information from other tokens that the model judged relevant.

That is the core of self-attention. Every token updates itself by looking at the sequence around it.

Self-Attention vs Cross-Attention

In self-attention, Q, K, and V come from the same sequence. This is used heavily inside Transformer encoders and decoders. Each token attends to other tokens in the same input.

In cross-attention, queries come from one sequence while keys and values come from another sequence. This is common in encoder-decoder models. For example, a decoder generating a translation may use queries from the target language and keys-values from the encoded source sentence.

Multi-Head Attention: Why One Attention Is Not Enough

A single attention operation can learn one kind of relationship pattern. But language has many relationship types. One head may focus on subject-verb agreement. Another may track pronouns. Another may learn position-based dependencies. Another may focus on punctuation or phrase boundaries.

Multi-head attention solves this by running several attention operations in parallel. Each head has its own learned projections for Q, K, and V. After all heads produce outputs, the results are concatenated and projected again.

The important idea is representation diversity. The model does not force all attention into one similarity space. It allows different heads to look for different patterns at the same time.

Attention Masks: Controlling What a Token Can See

Attention is powerful because every token can attend to every other token. But sometimes that is not allowed. In language generation, a model should not look at future tokens when predicting the next token.

This is where masks are used. A causal mask blocks access to future positions. The model can attend only to current and previous tokens. This is essential for autoregressive language models.

Padding masks are another common type. When sequences are padded to the same length, the model should ignore padding tokens. Otherwise, it may treat meaningless padding as real context.

• Causal mask: prevents a token from seeing future tokens during generation.
• Padding mask: prevents the model from attending to artificial padding positions.
• Custom mask: restricts attention based on task-specific rules or structure.

Where Positional Encoding Fits

Attention compares tokens, but by itself it does not know order. If we shuffle the tokens, pure attention has no natural sense of first, second, or third position. This is a problem because language depends on order.

Positional encoding adds order information to token embeddings. In the original Transformer, sinusoidal positional encodings were used. Modern models may use learned position embeddings or newer relative and rotary position methods.

The key point is simple: attention tells the model which tokens matter, while positional information tells the model where those tokens are located.

The Transformer Block Around Attention

Attention is the heart of the Transformer, but it is not the entire block. A typical Transformer block contains multi-head attention, residual connections, layer normalization, and a feed-forward network.

The attention layer mixes information across tokens. The feed-forward network transforms each token representation independently. Residual connections help gradients flow. Layer normalization keeps activations stable.

Stacking many such blocks gives the model depth. Lower layers may learn local or syntactic patterns, while higher layers may capture more abstract relationships. This layered attention is one reason Transformers work well across language, vision, audio, code, and multimodal tasks.

A Minimal Attention Implementation

The following code shows the core scaled dot-product attention operation in PyTorch-style logic. It is not a full Transformer, but it captures the mathematical engine.

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q: [batch, heads, query_len, d_k]
    # K: [batch, heads, key_len, d_k]
    # V: [batch, heads, value_len, d_v]

    d_k = Q.size(-1)

    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    attention_weights = F.softmax(scores, dim=-1)

    output = torch.matmul(attention_weights, V)

    return output, attention_weights

This small function contains the essence of attention. Matrix multiplication creates scores. Scaling stabilizes them. Masking blocks invalid positions. Softmax creates weights. The final multiplication mixes values.

The Shape of Attention

Understanding tensor shapes makes attention less mysterious. If the sequence has n tokens and each key has dimension dₖ, then Q has shape n × dₖ and K has shape n × dₖ. The product QKᵀ has shape n × n.

That n × n matrix is important. It means every token has a score against every token. For long sequences, this becomes expensive because the attention matrix grows quadratically with sequence length.

This is one reason long-context Transformers require careful engineering. Many efficient attention variants try to reduce memory or computation while preserving useful long-range connections.

Common Misunderstandings About Attention

The first misunderstanding is that attention is the same as explanation. Attention weights can show which tokens influenced a representation, but they are not always a complete explanation of model behavior.

The second misunderstanding is that attention alone understands language. Attention is a mechanism. It becomes powerful because it is trained with large data, learned projections, feed-forward layers, normalization, and optimization.

The third misunderstanding is that multi-head attention simply repeats the same calculation. In practice, each head learns different projection matrices, so each head can represent a different relationship space.

Why the Math Matters

Knowing the math behind attention helps you understand why Transformers behave the way they do. It explains why context length affects cost, why masks matter for generation, why embeddings matter for similarity, and why multi-head attention improves representation power.

It also helps when reading papers, debugging models, or studying advanced topics such as efficient attention, retrieval-augmented generation, encoder-decoder models, and multimodal Transformers. Once Q, K, V, softmax, and masking are clear, the rest of the architecture becomes much easier to follow.

If you want to connect this concept with broader reasoning techniques, Codeayan’s guide on Chain-of-Thought Prompting is a useful next read. For agent workflows that build on LLM reasoning, explore ReAct prompting.

Key Takeaways

Conclusion

Further reading: Review the Attention Is All You Need paper, The Annotated Transformer, and the PyTorch scaled dot-product attention documentation.