Attention Flows: Analyzing and Comparing Attention Mechanisms in Language Models

TL;DR

This work tackles understanding how fine-tuning alters self-attention in Transformer-based language models, moving beyond static pre-training analyses to task-specific information flows. It introduces Attention Flows, a visual analytics tool that constructs an Attention Graph from per-head attention matrices $A_j^l$ across layers and uses a radial Token Rings layout to trace dependencies from input tokens to the classification token. The method defines token influence $I_l(w)$ using counts $c_l(w)$ with decay factor $\alpha$ and a threshold $\tau$ to connect tokens into bipartite graphs, enabling direct comparison between pre-trained and fine-tuned models via color and shared-vs-exclusive paths. Through use cases on textual entailment, question-answer verification, and paraphrase verification, plus user feedback, the approach demonstrates how fine-tuning reshapes attention to task-relevant cues and offers a practical tool for model interpretability and debugging.

Abstract

Advances in language modeling have led to the development of deep attention-based models that are performant across a wide variety of natural language processing (NLP) problems. These language models are typified by a pre-training process on large unlabeled text corpora and subsequently fine-tuned for specific tasks. Although considerable work has been devoted to understanding the attention mechanisms of pre-trained models, it is less understood how a model's attention mechanisms change when trained for a target NLP task. In this paper, we propose a visual analytics approach to understanding fine-tuning in attention-based language models. Our visualization, Attention Flows, is designed to support users in querying, tracing, and comparing attention within layers, across layers, and amongst attention heads in Transformer-based language models. To help users gain insight on how a classification decision is made, our design is centered on depicting classification-based attention at the deepest layer and how attention from prior layers flows throughout words in the input. Attention Flows supports the analysis of a single model, as well as the visual comparison between pre-trained and fine-tuned models via their similarities and differences. We use Attention Flows to study attention mechanisms in various sentence understanding tasks and highlight how attention evolves to address the nuances of solving these tasks.

Figures (12)

• Figure 1: We illustrate the computation involved in self-attention for the sequence "when was gegeen the emperor". Here the word "gegeen" attends to the rest of the sequence in independent attention heads, producing a separate embedding for each head. These are combined to produce a new embedding for "gegeen" at the next layer.
• Figure 2: The Attention Graph is built by adding edges between tokens that contain high attention. The [CLS] token (top) is the sole node for the last layer (12), and all extracted tokens ("not", "men", "?", "A" "Devo", "!") become new nodes for the prior layer (bottom).
• Figure 3: At the last layer, the tokens that [CLS] depends on (c.f. Fig. \ref{['fig:attgraph']}) are encoded in a ring (left), where the attention heads for [CLS] are shown in the center. This design is carried over to previous layers (right), where per-token attention information is depicted, along with the positions of attended tokens.
• Figure 4: Our design allows the user to compare model attention: here the user clicks on "against", a token that shares influence between models. Inspecting its dependent tokens in previous layers, we observe commonalities (in orange) and tokens unique to the fine-tuned model (purple), e.g. "dollar". Attention heads for tokens are, further, split between models where appropriate.
• Figure 5: Display of attention propagation while hovering over a token: (a) shows hovering over a single token and (b) shows hovering over a head glyph for propagation through only that head.