The Attention Mechanism: How Models Learn to Focus

The attention mechanism is one of the most impactful ideas in modern deep learning. It allows a model to selectively weight different parts of its input when producing each output, rather than compressing everything into a fixed-size representation.

Originally introduced for sequence-to-sequence models (Bahdanau et al., 2015), attention became the core building block of the Transformer architecture (Vaswani et al., 2017), which underpins virtually all modern large language models.

The Problem Attention Solves

In early encoder-decoder architectures (e.g., for machine translation), the encoder compressed the entire input sentence into a single fixed-size vector. The decoder then had to generate the output using only that vector — a bottleneck that degraded performance on long sequences.

Attention solves this by allowing the decoder to look back at all encoder states at each decoding step and dynamically decide which ones to focus on.

Queries, Keys, and Values

Modern attention is described in terms of three components:

The attention output for a query is a weighted sum of the values, where the weights are determined by how well each key matches the query.

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V

• QK^T computes a similarity score (dot product) between the query and every key.
• Dividing by \sqrt{d_k} (key dimension) prevents the scores from growing too large, keeping gradients stable.
• Softmax converts scores to probabilities (weights that sum to 1).
• The output is a weighted sum of the values.

Self-Attention

Self-attention (intra-attention) is the variant used in Transformers. Here Q, K, and V all come from the same sequence. Each token attends to every other token in the sequence — including itself.

This allows the model to capture long-range dependencies in a single layer, regardless of token distance.

import torch
import torch.nn.functional as F
 
def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / d_k ** 0.5  # (batch, heads, seq, seq)
 
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
 
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights  # output + attention weights

Multi-Head Attention

Instead of running a single attention function, multi-head attention runs h attention operations in parallel on different linear projections of Q, K, V. The outputs are concatenated and projected again.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O
where head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)

Each head can learn to attend to different aspects of the input simultaneously — syntax, coreference, semantic roles, etc.

graph LR
    Input --> Q & K & V
    Q --> H1(Head 1) & H2(Head 2) & H3(Head h)
    K --> H1 & H2 & H3
    V --> H1 & H2 & H3
    H1 --> Concat
    H2 --> Concat
    H3 --> Concat
    Concat --> Linear --> Output

Attention Variants

Visualizing Attention

Attention weights can be visualized as a matrix over token pairs. High-weight cells reveal which tokens the model found most relevant.

            "The"  "cat"  "sat"  "on"  "the"  "mat"
"sat"  -->  0.01   0.15   1.00   0.05   0.01   0.03
"on"   -->  0.02   0.03   0.30   1.00   0.20   0.10

Tools like BertViz let you interactively explore what attention heads have learned.

Computational Complexity

Standard self-attention is O(n^2 d) in time and O(n^2) in memory, where n is sequence length and d is model dimension. This quadratic scaling is the main bottleneck for processing long documents.

Research directions addressing this include:

• Flash Attention — IO-aware exact attention (Dao et al., 2022)
• Linear Attention — approximates softmax attention in O(n)
• Sliding window / sparse patterns — Longformer, BigBird

Further Reading

• Attention Is All You Need (Vaswani et al., 2017) — the Transformer paper
• Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau et al., 2015) — original attention paper
• The Illustrated Transformer — Jay Alammar
• Flash Attention (Dao et al., 2022)
• BertViz — attention visualization