{"id":6713,"date":"2023-07-16T14:37:03","date_gmt":"2023-07-16T22:37:03","guid":{"rendered":"https:\/\/live-cometml.pantheonsite.io\/?p=6713"},"modified":"2025-06-18T10:28:01","modified_gmt":"2025-06-18T10:28:01","slug":"explainable-ai-for-transformers","status":"publish","type":"post","link":"https:\/\/www.comet.com\/site\/blog\/explainable-ai-for-transformers\/","title":{"rendered":"Explainable AI: Visualizing Attention in Transformers"},"content":{"rendered":"\n
\"pink
Photo by Jeffery Ho<\/a> on Unsplash<\/a>, edited by author.<\/figcaption><\/figure>\n\n\n\n
\n
Follow along with the Colab<\/a><\/div>\n\n\n\n
Create a free Comet account<\/a><\/div>\n<\/div>\n\n\n\n

In this article we explore one of the most popular tools for visualizing the core distinguishing feature of transformer architectures: the attention mechanism. Keep reading to learn more about BertViz and how you can incorporate this attention visualization tool into your NLP and MLOps workflow with Comet.  <\/span><\/p>\n\n\n\n

Feel free to follow along with the <\/span>full-code tutorial here<\/span><\/a>, or, if you can\u2019t wait, check out <\/span>the final project here<\/span><\/a>.<\/span><\/p>\n\n\n\n

Introduction<\/span><\/h2>\n\n\n\n

Transformers have been described as the single most important technological development to NLP in recent years, but their processes remain largely opaque. This is a problem because, as we continue to make major machine learning advancements, we can\u2019t always explain how or why\u2013 which can lead to issues like undetected model bias, model collapse, and other ethical and reproducibility issues. Especially as models are more frequently deployed to sensitive areas like healthcare, law, finance, and security, model explainability is critical. <\/span><\/p>\n\n\n\n

\"Horizontal
Gender and race projections across professions, as calculated by Word2Vec. These learned biases could have a variety of negative consequences depending on the application of such a model. Image from Bias in NLP Embeddings<\/a> by Simon Warchal.<\/figcaption><\/figure>\n\n\n\n

What is BertViz?<\/span><\/h2>\n\n\n\n

BertViz is an open source tool that visualizes the attention mechanism of transformer models at multiple scales, including model-level, attention head-level, and neuron-level. But BertViz isn\u2019t new. In fact, early versions of BertViz have been around since as early as 2017. <\/span><\/p>\n\n\n\n

So, why are we still talking about BertViz? <\/span><\/p>\n\n\n\n

BertViz is an explainability tool in a field (NLP) that is otherwise notoriously opaque. And, despite its name, BertViz doesn\u2019t only work on BERT. The BertViz API supports many transformer language models, including GPT2, T5, and <\/span>most HuggingFace models<\/span><\/a>.<\/span><\/p>\n\n\n\n

\"BertViz
Despite its name, BertViz supports a wide variety of models. On the left, we visualize a question-answering task using an encoder-only model, and on the right, a text generation task using a decoder-only model. GIF by author.<\/figcaption><\/figure>\n\n\n\n

As transformer architectures have increasingly dominated the machine learning landscape in recent years, they\u2019ve also revived an old but important debate regarding interpretability and transparency in AI. So, while BertViz may not be new, its application as an explainability tool in the AI space is more relevant now than ever. <\/span><\/p>\n\n\n\n

But first, transformers<\/span><\/h2>\n\n\n\n

To explain BertViz, it helps to have a basic understanding of transformers and self-attention. If you\u2019re already familiar with these concepts, feel free to skip ahead to the section where we start coding.<\/span><\/p>\n\n\n\n

We won\u2019t go into the nitty gritty details of transformers here, as that\u2019s a little beyond the scope of this article, but we will cover some of the basics. I also encourage you to check out the additional resources at the end of the article. <\/span><\/p>\n\n\n\n

In the beginning (the prehistoric era of NLP)<\/span><\/h2>\n\n\n\n

So, how, exactly, does a computer \u201clearn\u201d natural language? In short, they can\u2019t\u2013 at least not directly. Computers can only understand and process numerical data, so the first step of NLP is to break down sentences into \u201ctokens,\u201d which are assigned numerical values. The question driving NLP then becomes \u201chow can we accurately reduce language and communication processes to computations?\u201d <\/span><\/p>\n\n\n\n

Some of the first NLP models included feed-forward neural networks like the Multi-Layer Perceptron (MLP) and even CNNs, which are more popularly used today for computer vision. These models worked for some simple classification tasks (like sentiment analysis) but had a major drawback: their feed-forward nature meant that at each point in time, the network only saw one word as its input. Imagine trying to predict the word that follows \u201cthe\u201d in a sentence. How many possibilities are there?<\/span><\/p>\n\n\n\n

\"A
Without much context, next word prediction can become extremely difficult. Graphic by author.<\/figcaption><\/figure>\n\n\n\n

To solve this problem, Recurrent Neural Networks (RNNs) and Long Short-Term Memory Networks (LSTMs like Seq2Seq) allowed for feedback, or cycles. This meant that each computation was informed by the previous computation, allowing for more context. <\/span><\/p>\n\n\n\n

This context was still limited, however. If the input sequence was very long, the model would tend to forget the beginning of the sequence by the time it got to the end of the sequence. Also, their sequential nature didn\u2019t allow for parallelization, making them extremely inefficient. RNNs also suffered notoriously from exploding gradients.<\/span><\/p>\n\n\n\n

Introducing transformers<\/span><\/h2>\n\n\n\n

Transformers are sequence models that abandon the sequential structure of RNNs and LSTMs and adopt a fully attention-based approach. Transformers were initially developed for text processing, and are central to relatively all state-of-the-art NLP neural networks today, but they can also be used with image, video, audio, or virtually any other sequential data. <\/span><\/p>\n\n\n\n

The key differentiating feature of transformers from previous NLP models was the attention mechanism, as popularized in the <\/span>Attention Is All You Need<\/span><\/a> paper. This allowed for parallelization, which meant faster training and optimized performance. Attention also allowed for much larger contexts than recurrence, meaning transformers could craft more coherent, relevant, and complex outputs.<\/span><\/p>\n\n\n\n

\"The
The original transformer architecture, as visualized in the 2017 paper that made them famous, Attention Is All You Need<\/a>.<\/figcaption><\/figure>\n\n\n\n

Transformers are made up of encoders and decoders, and the tasks we can perform with them depend on whether we use either or both of these components. Some common transformer tasks for NLP include text classification, named entity recognition, question-answering, text summarization, fill-in-the-blanks, next word prediction, translation, and text generation.<\/span><\/p>\n\n\n\n

\"Chart
Transformers are made up of encoders and decoders, and the tasks we can perform with them depend on whether we use either or both of these components. Graphic by author.<\/figcaption><\/figure>\n\n\n\n

How do transformers fit into the larger ecosystem of NLP models?<\/span><\/h2>\n\n\n\n

You\u2019ve probably heard of Large Language Models (LLMs) like ChatGPT or LLaMA. The transformer architecture is a fundamental building block of LLMs, which use self-supervised learning on vast amounts of unlabelled data. These models are also sometimes referred to as \u201cfoundation models\u201d because they tend to generalize well to a wide range of tasks, and in some cases are also available for more specific fine-tuning. BERT is an example of this category of model.<\/span><\/p>\n\n\n\n

\"A
Not all LLMs or foundation models use transformers, but they usually do. Likewise, not all foundation models are LLMs, but they usually are. Finally, not all transformers are LLMs or FMs. The important takeaway is that all transformer models use attention. Graphic by author.<\/figcaption><\/figure>\n\n\n\n

That\u2019s a lot of information but the important takeaway here is that the key differentiating feature of the transformer model (and by extension all transformer-based foundational LLMs) is the concept of <\/span>self-attention<\/b>, which we\u2019ll go over next. <\/span><\/p>\n\n\n\n

Attention<\/span><\/h2>\n\n\n\n

Generally speaking, attention describes the ability of a model to pay attention to the important parts of a sentence (or image, or any other sequential input). It does this by assigning weights to input features based on their importance and their position in the sequence. <\/span><\/p>\n\n\n\n

Remember that attention was the concept that improved the performance of previous NLP models (like RNNs and LSTMs) by lending itself to parallelization. But attention isn\u2019t just about optimization. It also plays a pivotal role in broadening the context a language model is able to consider while processing and generating language. This enables a model to produce contextually appropriate and coherent texts in much longer sequences.<\/span><\/p>\n\n\n\n

\"A
In this example, GPT-2 finished the input sequence with the word \u201cscared.\u201d How did the model know what \u201cit\u201d was? By examining the attention heads, we learn the model associated \u201cit\u201d with \u201cthe animal\u201d (instead of, for example, \u201cthe street\u201d). Image by author.<\/figcaption><\/figure>\n\n\n\n

If we break transformers down into a  \u201ccommunication\u201d phase and a \u201ccomputation\u201d phase, attention would represent the \u201ccommunication\u201d phase. In another analogy, attention is a lot like a search-retrieval problem, where given a <\/span>query<\/b>, <\/span>q<\/span><\/i>, we want to find the set of <\/span>keys<\/b>, <\/span>k<\/span><\/i>, most similar to <\/span>q<\/span><\/i> and return the corresponding <\/span>values<\/b>, <\/span>v<\/span><\/i>.<\/span><\/p>\n\n\n\n

    \n
  • Query:<\/b> What are the things I am looking for?<\/span><\/li>\n\n\n\n
  • Key:<\/b> What are the things that I have?<\/span><\/li>\n\n\n\n
  • Value:<\/b> What are the things that I will communicate? <\/span><\/li>\n<\/ul>\n\n\n\n
    \"A
    Visualization of attention calculation. Image from Erik Storrs<\/a>.<\/figcaption><\/figure>\n\n\n\n

    Types of attention<\/h3>\n\n\n\n

    Self-attention<\/b> refers to the fact that every node produces a key, query, and a value from that individual node. <\/span>Multi-headed attention<\/b> is just self-attention that is applied multiple times in parallel with different initialized weights. <\/span>Cross-attention<\/b> means that the queries are still produced from a given decoder node, but the keys and the values are produced as a function of the nodes in the encoder.<\/span><\/p>\n\n\n\n

    This is an oversimplified summary of transformer architectures, and we\u2019ve glossed over quite a few details (like <\/span>positional encodings<\/span><\/a> and <\/span>attention masks<\/span><\/a>). For more information, check out the additional resources below.<\/span><\/p>\n\n\n\n

    Visualizing attention before BertViz<\/span><\/h2>\n\n\n\n

    Transformers are not <\/span>inherently interpretable<\/span><\/a>, but there have been many attempts to contribute <\/span>post-hoc explainability<\/span><\/a> tools to attention-based models. <\/span><\/p>\n\n\n\n

    Previous attempts to visualize attention were often overly complicated and didn\u2019t translate well to non-technical audiences. They could also vary greatly from project to project and use-case to use-case.<\/span><\/p>\n\n\n\n

    \"A
    Previous attempts to visualize attention weren\u2019t standardized and were often overly confusing. Graphic compiled by author from Interactive visualization and manipulation of attention-based neural machine translation<\/a> (2017) and A Visual Debugging Tool for Sequence-to-Sequence Models<\/a> (2018).<\/figcaption><\/figure>\n\n\n\n

    Some successful attempts to explain attention behavior included attention-matrix heat maps and bi-partite graph representations, both of which are still used today. But these methods also have some major limitations. <\/span><\/p>\n\n\n\n

    \"A
    The attention-matrix heatmap (left) shows us that the model is not translating word-for-word, but considering a larger context for word order. But it\u2019s missing a lot of the finer details of the attention mechanism.<\/figcaption><\/figure>\n\n\n\n

    BertViz ultimately gained popularity for its ability to illustrate low-level, granular details of self-attention, while still remaining remarkably simple and intuitive to use. <\/span><\/p>\n\n\n\n

    \"GIF
    BertViz ultimately gained popularity for its ability to illustrate low-level, granular details of self-attention, while still remaining remarkably simple and intuitive to use. GIF by author.<\/figcaption><\/figure>\n\n\n\n

    That\u2019s a nice, clean visualization. But, what are we actually looking at?<\/span><\/p>\n\n\n\n

    How BertViz Breaks It All Down<\/span><\/h2>\n\n\n\n

    BertViz visualizes the attention mechanism at multiple <\/span>local scales<\/span><\/a>: the neuron-level, attention head-level, and model-level. Below we break down what that means, starting from the lowest, most granular level, and making our way up.<\/span><\/p>\n\n\n\n

    \"A
    BertViz visualizes attention at multiple scales, including the model level, attention head level, and neuron layer. Graphic by author.<\/figcaption><\/figure>\n\n\n\n

    Visualizing BertViz With Comet<\/span><\/h3>\n\n\n\n

    We\u2019ll log our BertViz plots to Comet, an experiment tracking tool, so we can compare our results later on. To get started with Comet, <\/span>create a free account here<\/span><\/a>, grab your API key, and run the following code:<\/span><\/p>\n\n\n\n