The Transformer Attention Mechanism

By Michael Samuel Naeem · February 9, 2026 · 5 min read

Before transformers, I spent weeks training an LSTM for document classification and watched it plateau at 82% accuracy while taking eight hours per epoch. Switching to a fine-tuned BERT — which is transformer all the way down — hit 91% in 40 minutes. The difference wasn't hyperparameter tuning. It was the attention mechanism letting every word see every other word directly, instead of compressing context through a fixed-size hidden state one token at a time. That architectural shift is why every modern LLM, vision transformer, and speech model uses attention as its core compute pattern.

The core idea: soft lookup over the sequence

Given a sequence of token embeddings, self-attention produces a new representation for each token by computing a weighted average of all tokens' values. The weights come from how relevant each token is to the current one.

For token i, the output is:

Attention(Q, K, V) = softmax(QK^T / √d_k) · V

Where Q (queries), K (keys), and V (values) are linear projections of the input embeddings. The dot product QK^T measures similarity between every query-key pair. Softmax normalizes those similarities into attention weights that sum to 1. Multiplying by V produces the weighted blend.

The √d_k scaling prevents dot products from growing large (which would push softmax into regions with tiny gradients).

A concrete walkthrough

Consider the sentence: "The cat sat on the mat because it was tired."

When processing "it", the model needs to know whether "it" refers to "cat" or "mat." Self-attention computes:

1. Project "it" to a query vector q_it
2. Project every token to key vectors k_the, k_cat, k_sat, ...
3. Compute dot products: q_it · k_cat, q_it · k_mat, ...
4. Softmax → weights: [0.02, 0.78, 0.01, 0.03, 0.05, 0.08, 0.03]
5. Output = weighted sum of value vectors, dominated by "cat"

The model learns these projections during training. No hand-coded coreference rules — just gradient descent on billions of examples.

Multi-head attention: parallel perspectives

A single attention head learns one type of relationship. Multi-head attention runs h independent attention operations in parallel, each with its own Q/K/V projections, then concatenates the results:

# Simplified multi-head attention
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        Q = self.W_q(x).view(batch, seq, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq, self.n_heads, self.d_k).transpose(1, 2)

        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
        weights = F.softmax(scores, dim=-1)
        out = weights @ V

        return self.W_o(out.transpose(1, 2).contiguous().view(batch, seq, -1))

With 12–32 heads, different heads specialize: one tracks syntactic dependencies, another handles positional proximity, another captures semantic similarity. The concatenated output gives the model a richer representation than any single head alone.

Positional encoding: restoring order

Self-attention has no inherent sense of position. The original transformer paper added sinusoidal positional encodings:

PE(pos, 2i)   = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

These are added to input embeddings before the first attention layer. Modern models use more sophisticated approaches:

• Learned positional embeddings — GPT-style, trained end-to-end
• RoPE (Rotary Position Embedding) — rotates Q and K vectors based on position; used in Llama, Mistral, and most open models
• ALiBi — adds a position-dependent bias to attention scores; extrapolates to longer sequences without retraining

RoPE is the dominant choice in 2026 because it handles context length extrapolation better than absolute embeddings.

The full transformer block

Each transformer layer combines attention with a feed-forward network:

x → MultiHeadAttention → Add & Norm → FeedForward → Add & Norm → output

The feed-forward sublayer is two linear transformations with a nonlinearity (ReLU or GELU) in between, applied independently to each token position. Residual connections (Add) and layer normalization (Norm) stabilize training across dozens of stacked layers.

Decoder-only models (GPT, Llama) use masked self-attention — each token can only attend to previous tokens, preserving autoregressive generation. Encoder-decoder models (original transformer, T5) add cross-attention where decoder queries attend to encoder keys and values.

Why attention won

The computational cost is O(n²) in sequence length — every token attends to every other token. This is why context windows have limits and why techniques like sparse attention, sliding windows, and KV caching matter for long sequences. But for the 2k–128k token ranges where most LLMs operate, the parallelism and representational power more than compensate.

What this means for practitioners

You don't implement attention from scratch — frameworks handle it. But understanding the mechanism explains:

• Why context length matters for cost — attention is quadratic in sequence length
• Why KV caching speeds up inference — past keys and values are reused, only new tokens compute new Q/K/V
• Why prompt structure affects quality — tokens attend globally, so relevant context should be positioned where the model can attend to it effectively
• Why attention visualization helps debugging — inspecting weights shows what the model focuses on

Resources

• Attention Is All You Need (Vaswani et al., 2017)
• The Illustrated Transformer (Jay Alammar)
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• PyTorch nn.MultiheadAttention
• KV Cache explained (Hugging Face)