[Paper Review] Visualizing Attention in Transformer-Based Language models

Jesse Vig|arXiv (Cornell University)|Apr 4, 2019

Topic Modeling3 references16 citations

TL;DR

This paper introduces an open-source tool for visualizing multi-head self-attention in Transformer-based language models across three granularities: attention heads, model architecture, and individual neurons. By applying the tool to BERT and GPT-2, the authors demonstrate its utility in detecting model bias, identifying recurring patterns, and linking neurons to behavioral outputs, enhancing interpretability of attention mechanisms.

ABSTRACT

We present an open-source tool for visualizing multi-head self-attention in Transformer-based language representation models. The tool extends earlier work by visualizing attention at three levels of granularity: the attention-head level, the model level, and the neuron level. We describe how each of these views can help to interpret the model, and we demonstrate the tool on the BERT model and the OpenAI GPT-2 model. We also present three use cases for analyzing GPT-2: detecting model bias, identifying recurring patterns, and linking neurons to model behavior.

Motivation & Objective

• To develop an open-source tool that enables visualization of self-attention mechanisms in Transformer models across multiple levels of granularity.
• To improve interpretability of attention mechanisms by analyzing attention at the head, model, and neuron levels.
• To demonstrate practical applications of the tool in analyzing model behavior, such as detecting bias and identifying recurring patterns in GPT-2.
• To link specific neurons to observable model behaviors, enhancing understanding of internal model dynamics.

Proposed method

• The tool visualizes attention at the attention-head level, showing how individual heads attend to tokens in the input sequence.
• It aggregates attention patterns across heads to provide a model-level view, revealing overall attention flow across layers.
• The tool enables neuron-level visualization by linking specific neurons to attention patterns and model outputs.
• It supports interactive exploration of attention distributions in BERT and GPT-2 using standardized attention matrices and layer-wise representations.
• The framework integrates with existing Transformer models and exposes attention weights through a modular, extensible interface.
• Use cases are implemented via post-hoc analysis of attention outputs, enabling detection of bias, pattern recurrence, and neuron-behavior correlations.

Experimental results

Research questions

• RQ1 How can multi-head self-attention be visualized effectively across different levels of granularity in Transformer models?
• RQ2 What insights into model behavior can be gained by analyzing attention at the neuron level?
• RQ3 Can attention visualization help detect bias in language models like GPT-2?
• RQ4 What recurring patterns emerge in attention distributions across sequences in GPT-2?
• RQ5 How can specific neurons be linked to observable model behaviors through attention analysis?

Key findings

• The tool successfully visualizes attention at the head, model, and neuron levels, enabling multi-scale interpretation of attention mechanisms.
• Attention visualization at the neuron level revealed specific neurons that consistently activate in response to particular linguistic patterns in GPT-2.
• The tool enabled detection of biased attention patterns in GPT-2, particularly in relation to gender and demographic associations.
• Recurring attention patterns were identified across multiple sequences, suggesting structured internal representations of common linguistic constructs.
• The integration of attention visualization with neuron-level analysis allowed for direct linking of specific neurons to model behaviors, enhancing model interpretability.