Transformers and Attention in Practice

The attention mechanism is the reason transformers beat everything before them - but most engineers who use transformer models have never looked at what attention actually computes. That is fine until something goes wrong: the model ignores a key fact in the context, fixates on irrelevant tokens, or produces outputs that are hard to debug. Understanding attention at the implementation level gives you a genuine diagnostic tool.

This article walks through the scaled dot-product attention math, shows runnable Python for multi-head attention, and explains what you actually see when you visualize attention weights.

Scaled Dot-Product Attention

The core formula is:

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

Q (queries), K (keys), and V (values) are linear projections of the input. The dot products between queries and keys measure similarity; scaling by sqrt(d_k) prevents the softmax from saturating when the dimension is large. The output is a weighted sum of values.

Here is a minimal numpy implementation:

pythonCopy
import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / np.sqrt(d_k)  # (batch, heads, seq, seq)
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)
    weights = softmax(scores, axis=-1)
    return weights @ V, weights

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

Multi-Head Attention in PyTorch

Real transformer models split the model dimension across multiple heads. Each head learns a different type of relationship.

pythonCopy
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2)

    def forward(self, x: torch.Tensor, mask=None):
        Q = self.split_heads(self.W_q(x))
        K = self.split_heads(self.W_k(x))
        V = self.split_heads(self.W_v(x))
        scale = self.d_k ** 0.5
        scores = torch.matmul(Q, K.transpose(-2, -1)) / scale
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        weights = F.softmax(scores, dim=-1)
        out = torch.matmul(weights, V)
        out = out.transpose(1, 2).contiguous().view(x.shape[0], -1, self.num_heads * self.d_k)
        return self.W_o(out), weights

Visualizing Attention Weights

weights from the forward pass has shape (batch, heads, seq_len, seq_len). Each row is a probability distribution over source tokens for one output position.

pythonCopy
import matplotlib.pyplot as plt
import seaborn as sns

def plot_attention(weights, tokens, head=0, layer=0):
    w = weights[layer][head].detach().cpu().numpy()
    plt.figure(figsize=(8, 6))
    sns.heatmap(w, xticklabels=tokens, yticklabels=tokens,
                cmap='Blues', vmin=0, vmax=1)
    plt.xlabel("Keys (source)")
    plt.ylabel("Queries (target)")
    plt.title(f"Attention weights - layer {layer}, head {head}")
    plt.tight_layout()
    plt.show()

What you see matters. In BERT-style models early heads often attend to punctuation; later heads in the final layers tend to carry semantic signal. In generative models, look at how much attention the final token places on key facts you gave in the prompt. If a critical fact is barely attended to, the model probably will not use it.

Common Mistakes

Forgetting the causal mask for decoder models. Without masking, each token can attend to future tokens during training, which is a data leak. The mask should be a lower-triangular boolean matrix.

Assuming all heads are meaningful. Many heads are nearly uniform - they are not doing anything useful. Pruning heads is a real technique (see Michel et al., 2019). Do not over-interpret every head.

Confusing cross-attention and self-attention. In encoder-decoder architectures, the decoder's cross-attention uses encoder outputs as keys and values. The attention map in cross-attention shows which encoder positions the decoder is pulling from - very useful for translation debugging.

Not normalizing when visualizing. Raw logits before softmax can be misleading. Always visualize the post-softmax weights.

Practical Debugging with Attention

When a model misses a fact from the prompt: check whether the answer token's row in the final-layer attention map has non-trivial weight on the tokens containing that fact. If it does not, the fact is likely buried too deep in context or the model has routed it through a head that handles syntax, not semantics.

When a model hallucinates a number: inspect attention on the decoder's final layer cross-attention. If the model's attention on the source context is diffuse, it is generating from prior rather than from your text.

Where to Go Next

See also: [transformer-inference-context-engineering], [inference-optimization-performance], [rag-foundations-retrieval-quality]