December 30, 2019

450 words 3 mins read

jessevig/bertviz

jessevig/bertviz

Tool for visualizing attention in the Transformer model (BERT, GPT-2, Albert, XLNet, RoBERTa, CTRL, etc.)

repo name jessevig/bertviz
repo link https://github.com/jessevig/bertviz
homepage
language Jupyter Notebook
size (curr.) 101894 kB
stars (curr.) 1430
created 2018-12-16
license Apache License 2.0

BertViz

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.

Blog post:

Paper:

Related blog posts:

Attention-head view

The attention-head view visualizes the attention patterns produced by one or more attention heads in a given transformer layer.

Attention-head view Attention-head view animated

The attention 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)

Model view

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

Model view

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)

Neuron view

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

Neuron view

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

Requirements

(See requirements.txt)

Execution

Running locally:
git clone https://github.com/jessevig/bertviz.git
cd bertviz
jupyter notebook

Click on any of the sample notebooks. Note that the sample notebooks do not cover all Huggingface models, but the code should be similar for those not included. The tool is designed for visualizing shorter sentences and may fail if the input text is very long.

Running from Colab:

Click on any of the Colab links above, and scroll to the bottom of the page. It should be pre-loaded with the visualization, so you don’t need to actually run anything.

If you write your own code for executing BertViz in Colab, note that some of the steps are different from those in the Jupyter notebooks (see Colab examples above).

Authors

Jesse Vig

Citation

When referencing BertViz, please cite this paper.

@article{vig2019transformervis,
  author    = {Jesse Vig},
  title     = {A Multiscale Visualization of Attention in the Transformer Model},
  journal   = {arXiv preprint arXiv:1906.05714},
  year      = {2019},
  url       = {https://arxiv.org/abs/1906.05714}
}

License

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

Acknowledgments

This project incorporates code from the following repos:

comments powered by Disqus