Attention Visualization › Attention Mechanisms: A Brief Recap

Building on the transformer architecture covered in the earlier parts of this book, recall that attention computes a weighted combination of values, where the weights reflect how relevant each position in the input is to the current query. This chapter assumes you are familiar with that mechanism. The key point for visualization purposes is that attention is not a single scalar per layer but a matrix per head per layer. This dimensionality is both a strength and a challenge: a model like BERT produces so many attention matrices that there is too much to examine by eye, which is why systematic, entropy-guided analysis is necessary rather than optional.

For a sequence of $n$ tokens, each attention head produces an $n \times n$ weight matrix. Entry $(i, j)$ in this matrix represents the degree to which position $i$ attends to position $j$ when the model processes that head. With a model that has $L$ layers and $H$ heads per layer, a single forward pass generates $L \times H$ such matrices. For BERT-base, which has 12 layers and 12 heads per layer, that is 144 attention matrices for every sentence. Each of those 144 matrices encodes a different view of the same sequence, learned independently through gradient descent. This redundancy is a feature: it allows the model to simultaneously reason about multiple types of relationships within the same input.

The attention weight matrix for a single head is computed as:

$$A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$$

where:

• $Q$: the query matrix, shape $(n \times d_k)$, obtained by projecting the input sequence through a learned weight matrix $W^Q$
• $K$: the key matrix, shape $(n \times d_k)$, obtained through a learned projection $W^K$
• $d_k$: the dimension of the key vectors, used to scale the dot products and prevent gradient saturation from large values
• $\text{softmax}$: applied row-wise, so each row of $A$ sums to 1 and can be interpreted as a probability distribution over positions

The resulting matrix $A$ has shape $(n \times n)$, and row $i$ is the attention distribution for token $i$: it tells you how much weight token $i$ places on each token in the sequence when computing its updated representation.

Why does this formulation matter for visualization? The softmax normalization means that every row of $A$ is a proper probability distribution. You can interpret each row as a question: "Given that I am at position $i$ and I want to update my representation, which other positions should I consult, and by how much?" The answer varies by head and by layer, and that variation is exactly what attention visualization tries to expose.

Attention weights occupy a privileged position among neural network internals: they are explicitly computed as an intermediate variable during the forward pass, before being used to weight the value vectors. Unlike the raw numerical weights of a dense layer, which transform inputs in high-dimensional space and resist direct interpretation, attention weights sit between zero and one and sum to one across each row. That probabilistic structure makes them look interpretable even when they may not fully reflect causal reasoning.

This accessibility explains why attention visualization emerged so early in transformer research. By 2018, papers like "Visualizing Attention in Transformer-Based Language Representation Models" and the original transformer paper's supplementary material were already displaying attention heatmaps. The technique requires minimal tooling, produces easy-to-read outputs, and maps onto intuitive concepts like "which words are important?" Whether the visual appeal tracks the underlying computation is the deeper question this chapter addresses.