Skip to content

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

License

Notifications You must be signed in to change notification settings

siddhantrele/bertviz

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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 posts:

Paper:

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

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

NOTE: If you wish to run BertViz using Colab, please see the example Colab scripts above, as they differ slightly from the Jupyter notebook versions.

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:

About

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

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Jupyter Notebook 98.9%
  • Other 1.1%