mars-wei / bertviz

Tool for visualizing attention in the Transformer model (BERT, GPT-2, and XLNet)

BertViz

Tool for visualizing attention in BERT, GPT-2, and XLNet. Extends Tensor2Tensor visualization tool by Llion Jones and pytorch-transformers from HuggingFace.

Blog posts:

• OpenAI GPT-2: Understanding Language Generation through Visualization
• Deconstructing BERT: Distilling 6 Patterns from 100 Million Parameters
• Deconstructing BERT Part 2: Visualizing the Inner Workings of Attention

Paper:

• A Multiscale Visualization of Attention in the Transformer Model

Attention-head view

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

BERT: [Notebook] [Colab]
GPT-2: [Notebook] [Colab]
XLNet: [Notebook] [Colab]

Model view

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

BERT: [Notebook] [Colab]
GPT-2 [Notebook] [Colab]
XLNet: [Notebook] [Colab]

Neuron view

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

BERT: [Notebook] [Colab]
GPT-2 [Notebook] [Colab]

Requirements

• PyTorch >=0.4.1
• Jupyter
• tqdm
• boto3
• IPython
• requests
• regex
• sentencepiece

(See requirements.txt)

Execution

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

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:

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