Source-linked AI summary
Attention Flows: Analyzing and Comparing Attention Mechanisms in Language Models
Joseph F DeRose, Jiayao Wang, Matthew Berger
TL;DR
Understanding how attention changes as language models are fine-tuned for downstream tasks remains an open question. Attention Flows traces and compares attention across words, layers, and heads to examine classification decisions. Its use cases show task-specific attention patterns, while user feedback supports the interface's usability and potential insights.
Problem
How attention mechanisms change when pre-trained models are fine-tuned for sentence-understanding tasks remains unclear, and understanding those changes can inform pre-training design.
Source: §1 · p. 2Method
Attention Flows traces classification-relevant attention from the final layer backward across words, layers, and heads, and supports comparison of pre-trained and fine-tuned models.
Source: §1 · p. 2, §3.2 · p. 4Results
Takeaways & Limitations
The examples indicate that self-attention evolves during fine-tuning to address task-specific details in classification.
Source: §6 · p. 9The interface compares one sample across models, and its simple classification-based browsing does not summarize or browse collections of samples.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The introduction frames the gap between pre-trained and fine-tuned BERT as a challenge for understanding sentence classification. Attention Flows visualizes attention across words, layers, and heads to trace classification decisions and compare task-specific patterns.
1 INTRODUCTION
- Fine-tuned BERT's generalization to sentence-understanding tasks remains unclear, although understanding model changes can inform pre-training design.Source: §1 · p. 2
- Unlike analyses centered on known linguistic relations or isolated heads or layers, this work follows how multiple heads and layers aggregate into a classification output.Source: §1 · p. 2
- Attention Flows traces classification-relevant attention backward from the final classification token through words, layers, and heads in a radial view.Source: §1 · p. 2
- Use cases and feedback show task-specific patterns: inference targets the Five Ws, while non-paraphrase attention centers on differentiating phrases.Source: §1 · p. 2
2. Our visualization supports the comparison of self-attention mech-
Attention Flows supports comparison of self-attention mechanisms in pre-trained and fine-tuned models. Use cases and user feedback examine fine-tuning changes across sentence-understanding tasks.
2. Our visualization supports the comparison of self-attention mech-
- Attention Flows compares self-attention mechanisms across pre-trained and fine-tuned models, emphasizing their differences and similarities.Source: §2 · p. 2
- Use cases and user feedback examine fine-tuning changes in textual entailment, question answering, and sentence paraphrasing.Source: §2 · p. 2
- The authors report that the interface offers insight into changes in self-attention for inference tasks.Source: §2 · p. 2
2 RELATED WORK
Related work covers language modeling, model interpretability, visual analytics for NLP, and comparative visualization. It motivates a holistic approach to comparing attention across layers and models for downstream tasks.
2. RELATED WORK
- Prior work in language modeling and Transformers describes contextualized representations, self-attention, and pre-training for downstream supervised NLP tasks.BERT uses masked word prediction and next sentence prediction before fine-tuning on supervised tasks.Source: §2.1 · p. 2
2.1 Language Modeling
- Language models commonly learn contextualized representations from text corpora, while Transformers use self-attention over input embeddings and multiple independent heads.Source: §2.1 · p. 2
2.2 Model Interpretability
- Interpretability studies identify linguistic relations in pre-trained models and changes after fine-tuning, but often lack instance-level explanations and cross-layer views.Prior work reports attention heads capturing dependency relations and embeddings capturing grammatical features; other studies examine fine-tuning loss landscapes or attention patterns.Source: §2.2 · pp. 2–3
2.3 Visual Analytics for Interpreting Models
- Visual analytics for NLP includes work on recurrent-model states, machine translation, classification, and Transformer attention, often focused on individual heads or layers.This work instead seeks a holistic view for visually analyzing and comparing models on downstream sentence-understanding tasks.Source: §2.3 · p. 3
2.4 Comparative Visualization
- Comparative visualization has used graph merging, color, heatmaps, grouped bars, and node-link diagrams, while this work compares how models reason over a given input.Prior deep-network comparison has often focused on model performance rather than reasoning over the same input.Source: §2.4 · p. 3
3 OVERVIEW AND TASKS
BERT builds contextualized token embeddings through stacked self-attention layers, while downstream classification uses the final-layer [CLS] embedding. The section frames Attention Flows’ objectives as tracing classification attention and examining how fine-tuning changes it.
3.1 A Summary of BERT
- BERT forms each layer’s contextualized token embeddings through self-attention, per-head vectors, and their combination.Attention weights encode how much one token’s output depends on another token’s input; per-head vectors are concatenated before further transformations.Source: §3.1 · p. 3
- Pre-training uses masked-token prediction and next-sentence prediction; fine-tuning then adapts model weights for a supervised sentence-understanding task.The downstream task typically classifies using the final-layer [CLS] embedding.Source: §3.1 · pp. 3–4
3.2 Objectives and Tasks
- The analysis focuses on how self-attention across heads and layers organizes the final classification embedding, while noting gradients as another way to quantify influence.The authors argue gradients are not particularly meaningful for pre-trained models before task-specific [CLS] parameters are updated.Source: §3.2 · p. 4
- The objectives are to understand attention’s role in classification and assess what fine-tuning teaches the model for a specific task.The authors treat understanding classification as necessary for identifying changes in attention due to fine-tuning.Source: §3.2 · p. 4
- The tasks are to trace and query attention across layers and heads, discover head functionality, and compare pre-trained with fine-tuned attention flows.Comparison includes locating shared or distinct heads and tracing global flow differences across layers.Source: §3.2 · p. 4
4 VISUALIZATION DESIGN
The section introduces the information extracted for Attention Flows and then explains the visualization design.
4 VISUALIZATION DESIGN
- The design section first presents extracted information, then describes the visualization.Source: §4 · p. 4
4.1 Attention Graph
The Attention Graph links token dependencies across adjacent layers to trace attention flows. Edges can be based on high attention in any head or restricted to selected heads.
4.1 Attention Graph
- For each attention matrix, the method creates a bipartite graph with an edge when a token’s attention exceeds an interactively adjustable threshold.The edge represents a token’s strong influence on another token’s embedding at the next layer.Source: §4.1 · p. 4
- The graph is built backward from [CLS] at the final layer, adding token dependencies across successive layers.Source: §4.1 · p. 5
- Dependencies can include a link present in any head by default or be limited to links passing through user-selected heads.Source: §4.1 · p. 5
- The Attention Graph connects tokens only across adjacent layers; paths indicate possible influence, with multiple paths treated as stronger evidence than a single path.The authors characterize a long single path as indicating minimal influence because a token’s identity is likely to be lost between layers.Source: §4.1 · p. 5
4.2 Attention Overview
Attention Overview assigns each token an influence score by tracing head-based dependencies backward from the [CLS] token. It averages scores across layers with decay and visualizes influence for inspection and model comparison.
4.2 Attention Overview
- Head counts from tokens attending to [CLS] are propagated backward through layers to compute token influence.At each preceding layer, the count reflects how many heads connect a token to tokens in the next layer.Source: §4.2 · p. 5
- Influence averages scores across layers, exponentially downweighting earlier layers with α = 0.5.The authors report that different α values do not affect relative model comparisons.Source: §4.2 · p. 5
- Five-circle ratings above tokens encode influence, while color distinguishes shared and model-specific influence in comparisons.The displayed influence is clamped to 5; Fig. 5 also shows shared dependent tokens and tokens unique to the fine-tuned model.Source: §4.2 · p. 5, Figure 5 · p. 5
4.3 Attention Flows Design
Attention Flows presents token dependencies, head information, and model comparisons in a layered view. Its radial layout traces influence toward [CLS], while glyphs and sparklines expose attention-head patterns.
4.3 Attention Flows Design
- The Attention Flow view depicts token dependencies, head information, and model comparisons across multiple layers.The central radial layout places [CLS] at the innermost ring and adds tokens from preceding layers that influence it, tracing token influence through the model.Source: §4.3 · p. 5, §4.3.1 · p. 5
4.3.1 Token Rings
- Each radial ring represents a layer, and preceding rings add tokens that influence tokens in the layer closer to [CLS].The auxiliary radial view preserves original sentence positions; the design notes that token counts tend to grow by at most 5–10 per preceding layer in practice.Source: §4.3.1 · p. 5
4.3.2 Attention Heads
- Twelve head glyphs and sparklines encode head counts and where dependencies fall in the input sequence.Fig. 6 contrasts hovering over a token with restricting propagation to one head; Fig. 7 highlights heads connecting selected tokens across adjacent layers.Source: §4.3.2 · p. 6, Figure 6 · p. 6, Figure 7 · p. 6
4.3.3 Model Comparison
- The combined display uses model-specific colors and split glyphs to distinguish model-only from shared attention heads.It takes the union of tokens attended to by either model and carries the color scheme into the Sentence Context view.Source: §4.3.3 · p. 6
4.4 Interacting with Attention
Attention Flows supports interactive exploration of attention subgraphs and head functionality. Users can trace dependencies within or across layers, constrain paths by head, and compare traversals between models.
4.4.1 Exploring Attention Subgraphs
- Selecting tokens reveals dependent tokens across layers, while selecting a head restricts traversal to that head at its layer.Brushing token sequences shows the intersection of their attention graphs, which can expose shared dependencies.Source: §4.4.1 · p. 6
- The combined view traverses each model separately and colors the union of resulting tokens to show similarities and differences.Split head glyphs provide separate hover targets for each model, supporting model-specific or comparative traversal.Source: §4.4.1 · p. 6
4.4.2 Querying Head Functionality
- Selecting tokens in different layers highlights intervening tokens and heads on paths between them.When multiple tokens are selected, the view shows intermediate tokens attended to by all selected tokens; Fig. 8 illustrates tokens in “wireless power transmission” sharing heads in the traversal.Source: §4.4.2 · pp. 6–7, Figure 8 · p. 6
- In the combined model view, paths are computed separately for each model and need not connect all selected tokens.Fig. 9 illustrates attention flowing between “receive” and “go” across the sentences in an entailment example.Source: §4.4.2 · p. 7, Figure 9 · p. 7
5 EXPERIMENTS
The experiments apply Attention Flows to three GLUE sentence-pair tasks using fine-tuned BERTBASE, then use task examples and participant feedback to examine attention patterns and interface usability.
5 EXPERIMENTS
5.1 Datasets
- The study examines RTE, MRPC, and QNLI, three of GLUE’s nine sentence-pair binary-classification tasks.RTE tests entailment, MRPC tests whether two sentences have the same meaning, and QNLI tests whether the second sentence contains the answer to the first sentence’s question.Source: §5.1 · p. 7
5.2 Use Case: Textual Entailment
- The fine-tuned model focuses on task-relevant words and relational language in entailment, question answering, and paraphrase examples, while pre-trained attention often differs.For entailment, fine-tuning emphasizes relevant nouns, prepositions, verbs, and possessive phrases; for question answering, question aspects bridge to answers; for negative paraphrases, it attends to distinguishing spans.Source: §5.2 · p. 9, §5.3 · p. 9, §5.4 · p. 9
5.3 Use Case: Question Answer Verification
- For positive QNLI examples, question words often bridge to the answer span, a pattern the authors say is learned through fine-tuning even though answer spans are not required.Examples also show answer tokens depending on question tokens across layers, while the pre-trained model may fail to attend to the answer across multiple layers.Source: §5.3 · p. 9
5.4 Use Case: Paraphrasing Verification
- For negative MRPC examples, fine-tuned BERT attends to sentence-distinguishing words, and deeper layers retain token identities that enable this attention.By contrast, the pre-trained model attends to matched spans or nearby tokens rather than the distinguishing spans in the examples.Source: §5.4 · p. 9
5.5 User Feedback
- Three NLP researchers found the visualization useful for model comparison, while noting that its central visualization could be complex to understand.Researchers also expressed interest in broader model comparisons beyond fine-tuning.Source: §5.5 · p. 9
- All crowdsourced participants found fine-tuned attention better at handling task-relevant details, though some found interactions difficult or unclear.Some participants said the interface became easy to use after an initial learning period; others did not find the purpose of certain interactions evident.Source: §5.5 · p. 9
6 CONCLUSIONS AND FUTURE WORK
The conclusion presents Attention Flows as a tool for exploring how self-attention changes during fine-tuning and informs classification, while identifying limits in model scope and sample browsing.
6 CONCLUSIONS AND FUTURE WORK
- Attention Flows supports exploration of self-attention changes during fine-tuning and how attention informs classification, with use cases and feedback indicating task-specific patterns and interface usability.The authors propose linking attention with word embeddings and extending comparisons to models with different pre-training objectives or optimization schemes.Source: §6 · p. 9
- The current design considers self-attention matrices but not other Transformer components such as contextualized embeddings, and compares only a single sample across models.The authors identify collection-level sample browsing and comparison as future work.Source: §6 · p. 9
Related papers
This paper cites
- Attention Is All You Need2017
- BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding2018
- ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION2014
- Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation2014
- Distributed Representations of Words and Phrases and their Compositionality2013
- GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding2018