Attention Weight Analysis

Description

Attention weight analysis reads the weights an attention-based neural network assigns to the parts of its input, to see which positions it draws on at each step. The weights come from an ordinary forward pass and are plotted as heatmaps, one row per output step or token. In recurrent encoder-decoder models, such as early neural machine translation or image captioning, each output step has a row over the source words or image regions, read as a soft alignment in which weight is shared across several source positions. In transformers, each attention head (one of several parallel attention units in a layer) has its own matrix at every layer; the analyst views these per head and layer, or averages them over many inputs to profile a head, for example by how much attention it gives each kind of token, or how often its strongest attention lands on a labelled relation such as a verb and its object. The heatmaps show where each step looks; on their own they do not measure how much each input contributed to the prediction, because the weights leave out what is done with the attended values and information mixes between layers.

Example Use Cases

Explainability

A localisation team validating an in-house recurrent encoder-decoder translation model for equipment manuals plots the attention matrix for sample sentences and checks that each technical term draws most of its weight from its source counterpart. Sentences whose weight falls mainly on punctuation or the end-of-sentence marker go to a translator; the team first tests this screen on sentences with known errors, and does not treat a clean alignment as proof of a correct translation.

A team assessing an image-captioning model for an alt-text service overlays each generated word's attention weights on the image to check that nouns such as 'wheelchair' or 'crossing' fall on the matching objects. Captions whose nouns attend mainly to background regions are sent for manual review before release, as a screen only: a match does not show that the object produced the word.

A clinical NLP team reviewing a transformer that extracts medication details from discharge letters uses head and model views on a sample of letters to find heads whose attention links each drug name to its dose and frequency, then checks whether those heads behave the same way on letters from a second hospital before describing in the model documentation where those heads look, without claiming that the extraction depends on them.

A model-risk team reviewing a contract-clause classifier uses attention maps as a first look at which phrases the model draws on, then tests the highlighted phrases with an attribution or ablation method before reporting them, because its validation policy does not accept attention weights alone as evidence of what drove a decision.

Limitations

A bright cell in a heatmap does not show that the input behind it caused the output. In recurrent text classifiers, attention often disagreed with other measures of which words mattered, and very different patterns gave the same prediction (Jain and Wallace 2019). Wiegreffe and Pinter (2019) reply that attention can be a believable account without being a faithful one. Neither tested transformers. Confirm any reading by masking the highlighted inputs and checking that the output changes.
In transformers each layer mixes information between tokens, so one layer's weights do not show which inputs a later representation depends on. In a small six-layer transformer, the classification token's weights agreed with deletion-based importance only in the first layer; rollout, which combines weights across layers, agreed better later (Abnar and Zuidema 2020). For a fine-tuned DistilBERT model the gain was small: rollout improves on raw weights but does not fix them.
Each input gives one matrix per head per layer (144 in BERT-base), and the most visible pattern is often uninformative. In BERT, over half the attention in layers 6 to 10 falls on the separator token, which Clark et al. (2019) suggest acts as a 'do nothing' option; Kobayashi et al. (2020) found these weights carry little information once vector size is counted. Profiling a head needs averaging over many inputs, since a few hand-picked examples can mislead.
The weights mean different things in different architectures. In recurrent translation models source attention reads as a word alignment (Bahdanau et al. 2015). In transformer translation models, averaging source attention over all layers gave poor alignments; better ones needed a chosen layer or a measure that counts vector size (Garg et al. 2019; Kobayashi et al. 2020). Check any alignment reading against reference alignments for the model being assessed.
Attention can be trained to look harmless while the model still uses what it hides. Pruthi et al. (2020) trained occupation classifiers to put almost no attention on gender words with little loss of accuracy; the models still used gender, yet reviewers shown the manipulated highlights judged no predictions gender-influenced, against 66% with the original attention. A heatmap showing no weight on a sensitive input is weak evidence; test by changing the input.

Resources

Research Papers

Neural Machine Translation by Jointly Learning to Align and Translate

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio•May 19, 2016

Neural machine translation is a recently proposed approach to machine translation. Unlike the traditional statistical machine translation, the neural machine translation aims at building a single neural network that can be jointly tuned to maximize the translation performance. The models proposed recently for neural machine translation often belong to a family of encoder-decoders and consists of an encoder that encodes a source sentence into a fixed-length vector from which a decoder generates a translation. In this paper, we conjecture that the use of a fixed-length vector is a bottleneck in improving the performance of this basic encoder-decoder architecture, and propose to extend this by allowing a model to automatically (soft-)search for parts of a source sentence that are relevant to predicting a target word, without having to form these parts as a hard segment explicitly. With this new approach, we achieve a translation performance comparable to the existing state-of-the-art phrase-based system on the task of English-to-French translation. Furthermore, qualitative analysis reveals that the (soft-)alignments found by the model agree well with our intuition.

Software Packages

bertviz

Dec 16, 2018

BertViz: Visualize Attention in NLP Models (BERT, GPT2, BART, etc.)

Tutorials

BertViz Interactive Tutorial (Colab)

NLP From Scratch: Translation with a Sequence to Sequence Network and Attention — PyTorch Tutorials

Documentations

Model outputs — Hugging Face Transformers documentation

We’re on a journey to advance and democratize artificial intelligence through open source and open science.

Related Techniques

Name	Description	Assurance Goals
Gradient-weighted Class Activation Mapping	Grad-CAM (Gradient-weighted Class Activation Mapping) draws a heatmap over an image showing which regions most raise a convolutional neural network's score for a chosen class, which need not be the one the model predicted. The last convolutional layer produces a stack of feature maps, each a coarse grid recording where one learned pattern appears. Grad-CAM takes the gradient of the class score with respect to each map (how much the score changes as that map's values change), averages it to give the map a weight, and adds the weighted maps together. It keeps only the positive values, so the result marks regions that count in favour of the class, then enlarges the coarse grid to the image's size and lays it over the image, so the output shows broad regions rather than pixels. In a bias-diagnosis use, the maps are laid over the errors of a classifier already known to err more for one group, to see which image region the model relies on; the group-wise errors and the training-data counts come from separate checks.	Explainability Fairness
Neuron Activation Analysis	Neuron activation analysis examines when individual neurons in a neural network respond strongly. The analyst runs a set of inputs through the model, finds the inputs that activate a chosen neuron most and least, and looks for the concept that separates them, such as buildings in an image classifier or French text in a language model. Some methods test the match automatically, for example by scoring how closely a neuron's activation lines up with labelled examples of a concept. The result is a list of neurons with candidate meanings and the evidence behind each. These meanings are correlations between inputs and activations, so they show what a neuron responds to, not what it does to the model's output.	Explainability
Sparse Autoencoder Feature Analysis	Sparse autoencoder feature analysis looks for human-interpretable units inside a neural network, usually a language model, by learning from the activations at one layer a large dictionary of directions, called features, such that each activation vector can be rebuilt from only a few of them. Unlike neuron activation analysis it does not take the network's own neurons as the units, since single neurons often respond to several unrelated things, and unlike concept activation vectors it needs no labelled examples of a concept chosen in advance, because the features are found without supervision. The analyst records activations at the chosen layer over a large body of ordinary text and trains an autoencoder, a network with one hidden layer usually several times wider than the layer studied, to reconstruct them while a sparsity penalty keeps most hidden units at zero; the decoder weights of each hidden unit then give one feature direction. Each feature is labelled from the inputs on which it fires most strongly, by a person or by a language model that writes a description, which is then scored on how well it predicts the feature's activations on other text, and the label is tested by removing or amplifying the feature inside the running model and checking the change in output. The output is a catalogue of labelled features with example activations, interpretability scores and measured effects, plus a reconstruction score saying how much of the layer's behaviour the dictionary accounts for; the features can then serve as units for causal mediation analysis, which traces pathways, while this method finds and checks the units themselves.	Explainability