Visualizing Self-Attention Scores

The Scaled Dot-Product Attention formula, $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$, yields attention scores. Observing the distribution of these scores across a sequence offers valuable understanding of a model's internal workings. Visualizing self-attention scores clarifies which input sequence parts a model deems most informative when processing a specific word.

Recall that self-attention allows each position in the input sequence to attend to all positions (including itself) in the same sequence. The calculated attention scores, resulting from the $\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$ part of the formula, represent these relationships. A higher score between word $i$ and word $j$ indicates that the model places greater importance on word $j$ when computing the representation for word $i$.

Attention Score Matrices

A common way to visualize these scores is through a heatmap or an attention matrix. In such a visualization:

• Each row typically corresponds to a token in the sequence generating the query ($Q$). This is the token whose representation is currently being calculated.
• Each column corresponds to a token in the sequence providing the ($K$) and value ($V$). These are the tokens being attended to.
• The intensity or color of the cell at row $i$, column $j$ represents the attention score from token $i$ to token $j$. Darker or brighter colors usually indicate higher attention scores.

Let's consider a simple example sequence: "The quick brown fox". When the model calculates the updated representation for the word "quick", it computes attention scores indicating how much it should focus on "The", "quick", "brown", and "fox". Similarly, when processing "fox", it calculates scores for attending to each of the other words.

Imagine we're calculating the attention scores for the word "quick" in our example sentence. The model might produce scores like this:

• "quick" attending to "The": 0.10
• "quick" attending to "quick": 0.35
• "quick" attending to "brown": 0.45
• "quick" attending to "fox": 0.10

This suggests that when processing "quick", the model pays most attention to "brown" and itself ("quick"), and less attention to "The" and "fox".

Example Visualization

We can represent the attention scores for the entire sequence "The quick brown fox" as a matrix. Below is a visualization of the attention scores for each word attending to all other words in the sequence.

Self-attention scores for the sequence "The quick brown fox". Each row shows the attention distribution for that word (Query) over all words in the sequence (Key). Darker blue indicates higher attention scores.

In this example:

• Row 1 ("The"): Attends most strongly to itself.
• Row 2 ("quick"): Attends strongly to "brown" and itself.
• Row 3 ("brown"): Attends strongly to itself, moderately to "quick" and "fox".
• Row 4 ("fox"): Attends most strongly to itself, moderately to "brown".

Interpreting Patterns

By examining these visualizations, we can often identify interesting patterns:

• Identity Attention: Words often attend strongly to themselves.
• Local Context: Words may attend strongly to adjacent words, capturing local dependencies.
• Syntactic Relationships: Verbs might attend to their subjects or objects, even if they are distant in the sequence. Nouns might attend to modifying adjectives.
• Global Context: Sometimes, specific words (or special tokens like [CLS] in BERT-like models) might attend broadly across the sequence, acting as aggregators of information.

It's important to remember that this visualization usually depicts one attention head. In Multi-Head Attention, each head computes its own set of attention scores. Visualizing different heads often reveals that each head learns to focus on different types of relationships or representation subspaces simultaneously. For instance, one head might focus on local syntax, while another captures longer-range semantic dependencies.

Visualizing attention scores is not just aesthetically interesting; it's a valuable diagnostic tool. It provides a window into the model's reasoning process, helping us understand how it relates different parts of the input to form meaningful representations, a task that was much harder to interpret in earlier sequence models like RNNs.