BertViz is a tool for visualizing attention in the Transformer model, supporting all models from the transformers library (BERT, GPT-2, XLNet, RoBERTa, XLM, CTRL, etc.). It extends the Tensor2Tensor visualization tool by Llion Jones and the transformers library from HuggingFace.

🕹️ Colab tutorial

✍️ Blog post

📖 Paper

The head view visualizes the attention patterns produced by one or more attention heads in a given transformer layer. It is based on the excellent Tensor2Tensor visualization tool by Llion Jones.

Try out this interactive Colab Notebook with the head view pre-loaded.

The head view supports all models from the Transformers library, including:
BERT: [Notebook] [Colab]
GPT-2: [Notebook] [Colab]
XLNet: [Notebook]
RoBERTa: [Notebook]
XLM: [Notebook]
ALBERT: [Notebook]
DistilBERT: [Notebook]
(and others)

The model view provides a birds-eye view of attention across all of the model’s layers and heads.

Try out this interactive Colab Notebook with the model view pre-loaded.

The model view supports all models from the Transformers library, including:
BERT: [Notebook] [Colab]
GPT2: [Notebook] [Colab]
XLNet: [Notebook]
RoBERTa: [Notebook]
XLM: [Notebook]
ALBERT: [Notebook]
DistilBERT: [Notebook]
(and others)

The neuron view visualizes the individual neurons in the query and key vectors and shows how they are used to compute attention.

Try out this interactive Colab Notebook with the neuron view pre-loaded (requires Chrome).

The neuron view supports the following three models:
BERT: [Notebook] [Colab]
GPT-2 [Notebook] [Colab]
RoBERTa [Notebook]

git clone https://github.com/jessevig/bertviz.git
cd bertviz
<Install from requirements.txt as needed>
jupyter notebook

Click on any of the sample notebooks. You can view a notebook's cached output visualizations by selecting File > Trust Notebook (and confirming in dialog) or you can run the notebook yourself. Note that the sample notebooks do not cover all Huggingface models, but the code should be similar for those not included.

(If running in Jupyter Lab or Colab, see additional instructions below).

First load a Huggingface model, either a pre-trained model as shown below, or your own fine-tuned model. Be sure to set output_attention=True.

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased", output_attentions=True)
inputs = tokenizer.encode("The cat sat on the mat", return_tensors='pt')
outputs = model(inputs)
attention = outputs[-1]  # Output includes attention weights when output_attentions=True
tokens = tokenizer.convert_ids_to_tokens(inputs[0]) 

Then display the returned attention weights using the BertViz head_view or model_view function:

from bertviz import head_view
head_view(attention, tokens)

Just be sure that a copy of the bertviz directory is contained in the same folder or has been added to the PYTHONPATH. For more advanced use cases, e.g., specifying a two-sentence input to the model, please refer to the sample notebooks. The neuron view has a more constrained API (see Limitations); refer to the sample notebooks for examples.

Include the following code prior to calling BertViz:

from IPython.core.display import HTML
display(HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>'))

Include the following code in the beginning of the cell where BertViz is called:

from IPython.core.display import HTML
display(HTML('<script src="/static/components/requirejs/require.js"></script>'))

See the Colab links above for full working examples.

For the head view, you may pre-select a specific layer and collection of heads, e.g.:

head_view(attention, tokens, layer=2, heads=[3,5])

You may also pre-select a specific layer and single head for the neuron view.

The model view and neuron view support dark (default) and light modes. You may turn off dark mode in these views using the display_mode parameter:

model_view(attention, tokens, display_mode="light")

The head_view and model_view functions may technically be used to visualize self-attention for any Transformer model, as long as the attention weights are available and follow the format specified in model_view and head_view (which is the format returned from Huggingface models). In some case, Tensorflow checkpoints may be loaded as Huggingface models as described in the Huggingface docs.

• The visualizations works best with shorter inputs (e.g. a single sentence) and may run slowly if the input text is very long, especially for the model view.
• When running on Colab, some of the visualizations will fail (runtime disconnection) when the input text is long.
• The neuron view only supports BERT, GPT-2, and RoBERTa models. This view needs access to the query and key vectors, which required modifying the model code (see transformers_neuron_view directory), which has only been done for these three models. Also, only one neuron view may be included per notebook.

Visualizing attention weights illuminates a particular mechanism within the model architecture but does not necessarily provide a direct explanation for model predictions. See [1], [2], [3].

Jesse Vig

When referencing BertViz, please cite this paper.

@inproceedings{vig-2019-multiscale,
    title = "A Multiscale Visualization of Attention in the Transformer Model",
    author = "Vig, Jesse",
    booktitle = "Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: System Demonstrations",
    month = jul,
    year = "2019",
    address = "Florence, Italy",
    publisher = "Association for Computational Linguistics",
    url = "https://www.aclweb.org/anthology/P19-3007",
    doi = "10.18653/v1/P19-3007",
    pages = "37--42",
}

This project is licensed under the Apache 2.0 License - see the LICENSE file for details

We are grateful to the authors of the following projects, which are incorporated into this repo:

• https://github.com/tensorflow/tensor2tensor
• https://github.com/huggingface/pytorch-pretrained-BERT