AttentionViz A Global View of Transformer Attention

摘要¶

transformer 模型正在彻底改变机器学习领域，但其内部工作机制仍然颇具神秘色彩。在本文中，我们提出了一种新的可视化技术，旨在帮助研究人员理解 transformer 中的 self-attention 机制，这一机制使得这些模型能够学习序列元素之间丰富的、上下文相关的关系。我们方法的主要思想是可视化 transformer 模型在计算注意力时使用的 Query 和 Key 向量的联合 embedding。与之前的注意力可视化技术不同，我们的方法能够分析多个输入序列之间的全局模式。我们基于这些联合 Query - Key embedding 创建了一个交互式可视化工具，AttentionViz（演示：http://attentionviz.com），并使用它来研究语言和 ViT 中的注意力机制。我们通过几种应用场景和专家反馈展示了我们方法在改善模型理解和提供关于 Query - Key 交互新见解方面的实用性。关 Key 词—transformer，注意力，自然语言处理，计算机视觉，视觉分析。

Fig. 1: AttentionViz，我们的交互式可视化工具，允许用户通过为 Query 和 Key 创建联合 embedding 空间来大规模探索 transformer 的 self-attention。(a) 在语言 transformer 中，这些可视化呈现出引人注目的视觉痕迹，可以与注意力模式相关联。散点图中的每个点代表一个单词的 Query 或 Key 版本，点的颜色表示。用户可以探索单个 Attention-head（左）或放大以获得注意力的“全局”视图（右）。(b) 我们的可视化还揭示了 ViT 中的有趣见解，例如通过色调和亮度对图像块进行分组的 Attention-head。边框颜色表示一个图像块的 Query embedding（绿色）或 Key embedding（粉色）。(c) 提供了用于参考的样本输入句子（来自 [23]）和 (d) 图像（合成数据集）。

1 引言¶

transformer 神经网络架构 [52] 对从自然语言处理（NLP）[13, 42] 到计算机视觉 [14] 等领域产生了重大影响。事实上，transformer 现已被部署在由数亿人使用的大型实际系统中（例如，Stable Diffusion, ChatGPT, Microsoft Copilot）。然而，这些成功背后的机制仍有些神秘，特别是随着模型复杂性和规模的不断增加，新的能力持续涌现 [11, 60]。深入理解 transformer 模型可以帮助我们构建更可靠的系统，排查问题，并提出改进途径。在这项工作中，我们描述了一种旨在更好理解 transformer 运作方式的新可视化技术。（我们在第二节简要介绍 transformer。）我们分析的目标是 transformer 特有的 self-attention 机制，该机制使这些模型能够学习并使用输入元素之间的丰富关系集。尽管注意力模式已经被密集研究，但先前的技术通常一次只可视化与单一输入序列（例如，一句话或一张图片）相关的信息。典型的方法为给定输入序列创建二部图 [51, 53] 或热图 [20, 30] 表示的注意力权重。

我们的方法提供了一个更高层次的视角，我们可以一次观察许多输入序列的 self-attention 模式。这种方法的灵感之一来自于诸如激活图谱 [5] 之类的工具的成功，这些工具允许研究者“缩小”以获得神经网络的概览，然后深入了解细节。在我们的案例中，我们寻求构建一种“注意力图谱”，它可以为研究人员提供关于 transformer 不同 Attention-head 运作的丰富而详细的视图。主要的新技术是可视化 transformer 使用的 Query 和 Key 向量的联合 embedding，为个别 Attention-head 创建视觉签名。

为了说明我们的技术，我们实现了 AttentionViz，一个交互式可视化工具，允许用户探索语言和 ViT 中的注意力。AttentionViz 通过多个细节层次（图 1）进行探索，提供一次性查看所有 Attention-head 的全局视图以及放大查看单个 Attention-head 或输入序列的细节的能力。

我们通过几个应用场景和领域专家访谈展示了我们技术的实用性。为具体说明，我们重点关注可视化可以揭示的一些广泛使用的 transformer：BERT [13]，GPT-2 [41]，和 ViT [14]。我们发现与 BERT 的注意力模式相关的几个可识别的“视觉痕迹”，在 ViT 的视觉注意力机制中检测到新颖的色调/频率行为，并揭示了 GPT-2 中潜在的异常行为。

用户反馈还支持我们的方法在大规模可视化其他 embedding 时的广泛适用性。总结而言，这项工作的贡献包括：

一种基于联合 Query - Key embedding 探索 transformer 模型注意力趋势的可视化技术。
AttentionViz，一个交互式工具，应用我们的技术在多个规模上研究视觉和语言 transformer 的 self-attention。
应用场景和专家反馈显示，AttentionViz 如何揭示 transformer 注意力模式的洞见。

2 transformer 模型背景¶

transformer，首次介绍于文献 [52]，是一种设计用于处理序列输入的神经网络架构。完整描述 transformer 超出了本文的范围，但有几个概念对于理解我们的工作至关重要。首先，transformer 接收一组向量（通常称为 embedding）作为输入。embedding 可以代表多种输入类型。在基于文本的 transformer 中，它们对应于单词或单词片段；在 ViT 中，它们 encoding 像素块。

网络通过一系列注意力层迭代地转换这些向量，每个层都在 embedding 对之间移动信息。" 注意力 " 这一名字表明并非所有 embedding 都同等相关；某些对会更强烈地交互——即，彼此之间更加“关注”。注意力层决定哪些对应当互动，以及信息应如何在它们之间流动。

例如，在处理句子“The brown capybara is sleeping now”中的 transformer，人们可能期望在“capybara”和“is”之间的 embedding 具有高度的注意力（和信息流动），但“brown”和“now”之间则不会。self-attention 机制，是本文的重点，允许 transformer 学习并使用序列元素之间的丰富关系集，从而在各种 NLP 和计算机视觉任务中显著提升性能 [13, 14, 41]。

embedding 对之间的注意力可能有不同的理由。例如，在我们的示例句子中，“brown”和“capybara”通过形容词 - 名词关系连接，而“capybara”和“is”形成主语 - 谓语关系。为了允许多种关系类型，transformer 注意力层包含多个 Attention-head，每个头都可以表示不同的注意力和信息流动模式。

每个 Attention-head 使用 Query 权重矩阵 (\(W_Q\)) 和 Key 权重矩阵 (\(W_K\)) 计算其自身的注意力模式。具体来说，对于两个 embedding 向量 (\(x\)) 和 (\(y\))，注意力与 Query 向量 (\(W_Q\) x) 和 Key 向量 (\(W_K y\)) 的缩放内积相关。设 (\(W_K y\)) 的维度为 (\(d\))，我们有：

\[ f(x, y) = \frac{1}{\sqrt{d}} \langle W_Q x, W_K y \rangle \]

给定 embedding 向量 (\(\{x_1, x_2, \ldots, x_n\}\))，我们使用 softmax 函数计算 (\(x_i\)) 与其他向量之间的 self-attention：

\[ \text{attn}(x_i, x_j) = \text{softmax}_j\left(f(x_i, x), \ldots, f(x_i, x_n)\right) = \frac{e^{f(x_i, x_j)}}{\sum_k e^{f(x_i, x_k)}} \]

关 Key 的是，这个公式显示了 Query 和 Key 向量之间的点积越大，最终的注意力值就越高，这是我们在联合 embedding 可视化中依赖的一个事实。transformer 架构还有更多内容，但我们在此没有涵盖。特别是，我们只描述了 embedding 对之间的注意力权重，而没有描述在它们之间流动的具体信息。（如后文讨论，这是一个适合进一步研究的领域。）

最后值得一提的技术点是，它将帮助理解本文稍后的图像。transformer 接收的初始 embedding 通常包含它们的排序（对于一维序列）或空间配置（对于网格，如在 ViT 中）的向量表示。对于序列，这些位置向量使用三角函数定义，并位于高维空间中的类螺旋曲线上（见文献 [52]）。

2.1 本文研究的模型¶

我们研究了三种 transformer 模型：BERT（语言），GPT-2（语言）和 ViT（视觉）。每种模型都是机器学习社区研究的重要对象，这三种模型涵盖了 transformer 的不同架构和应用范围。BERT，即双向 encoding 器表示 transformer [13]，是一个多层 transformer encoding 器。作为一个双向模型，BERT 可以同时关注句子中任一方向的标记（即输入元素）。GPT-2，即生成式预训练 transformer 2 [42]，是一个多层 transformer 解码器。GPT-2 是一个单向模型，意味着它只关注前面的标记。ViT，即 ViT [14]，通过将图像分割成“块”，并像处理句子中的标记一样处理这些块，采用基于 self-attention 的 transformer 架构。类似于 BERT，ViT 是一个多层的双向 transformer encoding 器。在这项工作中，我们关注 ViT 在 16x16（ViT-16）和 32x32（ViT-32）块大小上的表现。

3 相关工作¶

许多研究者试图探索 transformer 的内部机制。文献 [8, 34] 旨在通过探索学习到的语言表示来理解基于 transformer 的语言模型的性能提升，而文献 [49] 观察到 BERT 重现了自然语言分析中的经典步骤，从词性标注到关系分类。另一种流行的方法是机械解释性，即反向工程 transformer 模型（例如，[15, 16, 37]）。注意力，transformer 的支柱，也被密集研究。例如，NLP 系统中的注意力似乎与句法结构相关 [9, 57]，ViT 中的注意力与格式塔式分组相似 [33]。研究人员还将 ViT 的视觉注意力机制与卷积滤波器进行了比较，发现注意力在图像遮挡、损坏和高频噪声方面更为稳健 [35, 40]。在我们讨论相关工作时，我们将重点放在研究 transformer 注意力的视觉方法上。

3.1 在单一输入序列中可视化注意力¶

注意力模式自然适合于可视化，在语言和 ViT 中均是如此 [4,12,21,31,39]。这些可视化主要集中在使用二部图（例如，[30, 48, 51, 53]）或热图（例如，[1, 10, 20, 22, 25, 30, 43]）可视化单一输入序列中 Query 和 Key 标记之间的注意力权重。

一些可视化提案允许比较多个模型或层。例如，Attention Flows [12] 支持用户比较 BERT 的不同层内以及不同 Attention-head 在给定单个句子中的注意力。Dodrio [59] 使用网格视图，应用于单一输入，使得直接比较不同的 Attention-head 成为可能。另一个系统，VisQA [21]，通过显示语言 self-attention、视觉 self-attention 和语言 - 视觉交叉注意力的热图，可视化视觉问答任务中不同头的注意力。即使在这些模型比较系统中，分析师也必须一次查看不同的输入，以识别和验证给定 Attention-head 的模式。

3.2 超越单一输入：可视化 embedding 和激活最大化¶

寻找在多个输入中保持一致的模式是很自然的行为。一种被证明对这一目标有效的技术是可视化来自多个输入序列的 embedding 向量集合 [3, 18, 19, 46, 47, 58]。例如，文献 [43] 可视化了在许多不同上下文中使用相同单词的 BERT embedding，并发现了与词义相对应的群集。在探索语法处理方面，文献 [8] 可视化了多语言 BERT 模型的 embedding，再次发现了有助于解释的有意义的群集。LMFingerprints [45] 使用基于树的径向布局来比较不同语言模型的 embedding 向量。

第二种技术，在文献 [17,61] 中用于 ViT，旨在找到最大化特定单元激活的图像。应用于 embedding 向量时，这种技术产生了清晰可解释的结果。然而，作者指出，当应用于 Query 和 Key 向量时，这种技术似乎并未产生有用的结果。

图 2：为单个 Attention-head 创建一个联合 Query - Key embedding 空间。在自然语言处理的情况下，给定一个输入句子，我们首先将每个标记转换为其对应的 Query 向量和 Key 向量。然后，我们使用 t-SNE/UMAP/PCA 将这些 1×d 向量投影到 2D/3D 散点图坐标中。对于使用的 BERT、GPT-2 和 ViT 模型，d = 64。

3.3 文献中的空白¶

我们指出现有文献中存在三个空白，这些空白激励了我们的工作。首先，虽然已显示可视化 embedding 向量是分析多个输入中的模式的有效技术，但我们不知道有系统尝试可视化 transformer 模型中的 Query 和 Key embedding 的。文献 [6] 也认为 self-attention 的中间产物，如 Query 和 Key，被探索不足。这些观察激发了我们的联合 Query - Key embedding 技术。

其次，尽管已提出了比较多个 embedding 的可视化技术（例如，[2, 3, 26]），这些方法通常限于少数 embedding，无法满足我们比较不同 transformer 头和层次的 embedding 的需求。因此，我们设计了一个全局矩阵视图来大规模可视化 Query - Key embedding。

最后，二部图表示已在分析基于 NLP 的 transformer 时证明是有帮助的，但我们尚未看到它们应用于视觉任务。我们通过创建二部图风格的可视化来探索这一方向，以研究 ViT 中的图像注意力模式。

4 目标与任务¶

本工作的总体目标是设计一种新颖的可视化技术，该技术允许探索 transformer 模型中的全局注意力趋势。为了收集关于这一想法的初步反馈并更多了解用户需求，我们与 5 位对模型可解释性感兴趣的机器学习（ML）研究人员（4 名博士生和 1 名教授）进行了交谈。在这些个别访谈中，我们请专家描述他们在使用 transformer 时的当前实践和挑战，以及注意力可视化如何帮助他们实现研究目标。我们将这些专家称为 E1-5。

4.1 目标¶

最终，与专家的对话产生了三个主要目标：

G1 - 理解 self-attention 如何告知模型行为。
总的来说，所有 5 位专家都希望更好地理解不同 Attention-head 的行为以及 transformer 模型通过其特征 self-attention 机制学到了什么。因此，他们表达了希望能够快速且轻松地探索注意力模式的愿望。E2 解释说：“注意力仍然是相当封闭的黑盒，还有很多谜团”，因此深入理解 transformer 的注意力模式可能提供洞见，例如“为什么大型语言模型在推理任务和数学上失败”。
G2 - 比较和对比 Attention-head。E5 提到，可视化 Attention-head 的差异有助于假设生成，这是他们研究过程的第一步：“可视化可以帮助形成要测试的假设，并直观感受 transformer 在做什么。”此外，3 位专家（E1、E2、E5）注意到，Attention-head 的比较对于模型修剪和编辑目的很有用。也就是说，如果两个 Attention-head 表现相似，或许可以移除一个而不会显著影响模型性能。用 E1 的话来说，比较头部可能让我们“找到模型中实际有用的部分”。
G3 - 识别注意力异常。四位研究人员（E2-5）希望通过探索注意力模式来识别 transformer 的不规则性和潜在行为问题。然后可以使用这些信息进行模型调试。例如，E4 说：“可视化注意力可以帮助你注意到模型关注的错误事物，即使结果是正确的。”E3 同意，重申在模型训练的背景下调试的重要性：“训练经常失败并死掉，但很难理解它为什么失败或产生意外行为。”

4.2 任务¶

鉴于这些目标，我们开发了以下设计任务： T1 - 大规模可视化 Attention-head。为了帮助用户快速探索模型行为 [G1] 并轻松比较和对比注意力模式 [G2]，我们的工具同时可视化 transformer 层中的自 Attention-head。

T2 - 探索 Query - Key 交互。E1 和 E4 表达了希望更好地理解 Query - Key 配对信息以提高对 transformer self-attention 的理解的愿望。因此，我们的工具通过可视化 Query - Key 交互进一步支持注意力模式比较 [G2] 和异常检测 [G3]。

T3 - 在多个层次上探测注意力。我们的工具通过在句子/图像、头部和模型层次提供可视化，允许进行局部和全局的注意力比较 [G2]。在单一界面中切换多个视图的灵活性也有助于知识发现 [G1]，并帮助用户识别模型的不规则性 [G3]。

T4 - 自定义模型和数据输入。AttentionViz 容易扩展到新的 transformer 和数据集，提供快速的视觉比较 [G2] 和不同模型及模态（语言和视觉）之间的注意力模式综合 [G1]。

5 Query/Key embedding 及 AttentionViz 的设计¶

为了解决这些目标和任务，我们构建了一个名为 AttentionViz 的工具。我们工具使用的主要技术是可视化每个 Attention-head 的 Query 和 Key 向量的联合 embedding。在这一节中，我们首先描述这种技术背后的动机和数学基础，然后讨论完整应用的设计。

5.1 可视化 Query/Key embedding¶

AttentionViz 背后的技术相对简单，虽然如下所述，它需要两个数学技巧才能有效。请回想一下，每个 transformer Attention-head 通过应用矩阵 \(W_Q\) 和 \(W_K\) 分别将输入 embedding 转换为 Query 向量和 Key 向量（见第 2 节）。这些矩阵将原始向量 embedding 投影到一个较低的维度空间，本质上选择了高维向量 embedding 中的特定类型信息。因此，通过检查 Query 和 Key 向量，我们可能希望学习 \(W_Q\) 和 \(W_K\) 选择了什么信息。

一个核心的观察是，Query 和 Key 向量的相对位置可以提供关于注意力如何分配的线索，因为注意力系数依赖于 Query 和 Key 之间的点积。为了理解原因，请考虑一个假设情况，其中 Query 和 Key 向量总是具有相同的范数。然后，更接近的距离将直接关联到更高的注意力系数。实际上，Query 和 Key 向量的范数是变化的，因此点积和距离之间的关系不是精确的。然而，如下一节所描述，我们可以安排这种关系出奇地接近。

图 2 用语言 transformer 中单个 Attention-head 的合成示例来说明我们的技术。为了创建联合 embedding，我们首先获取给定句子中每个标记的 Query 和 Key 向量表示（见第 2 节）。然后，我们使用三种降维方法之一将这些高维向量投影到一个共享的低维子空间：t-SNE [50]、UMAP [32] 或 PCA [24]。这些降维算法的输出是一个 \(2D/3D\) 散点图，其中每个点代表一个单一的 Query 或 Key 标记。同样的过程可以用于为 ViT Attention-head 创建联合 embedding，其中每个标记是一个图像块。默认情况下，我们用绿色可视化 Query，用粉色可视化 Key。然而，用户可以从多种颜色 encoding 中选择（见第 5.2 节），其他调色板也可以轻松地替换到系统中。

我们的联合 embedding 允许用户探索 Query 和 Key 之间的细粒度交互，这些图的形状经常可以作为潜在 self-attention 模式的视觉指示器（见第 7 节）。每种降维技术为给定数据集产生不同的模式，提供独特的见解和用例。

图 3：左图：联合 embedding 空间中的原始 Query 和 Key。右图：移动 Key 以使 Query 和 Key 的质心对齐后，重叠增加。

图 4：(a) 理想的距离 - 注意力关系，其中 Query - Key 对的点积越高，它们在联合 embedding 空间中的距离越近。 (b) 一个示例 Attention-head，在 BERT 中 Query - Key 距离与点积之间存在强烈的负相关（-0.983）。

5.1.1 向量标准化¶

在设计 AttentionViz 时，我们注意到两个“自由参数”，这些参数可以在不丢失任何信息的情况下变化。调整这些参数可以使 embedding 距离与注意力权重之间的关系更紧密，并大大提高可视化的可读性。这些标准化在降维之前应用（见图 2）。

Key 位移：在我们的可视化中，Query 和 Key 向量有时在空间上分离得很开（见图 3，左）。这种分离使得直接比较 Query 和 Key embedding 变得困难。然而，一个简单的数学技巧允许我们将这些 embedding 拉近，而不影响任何给定输入序列的 self-attention 计算。特别是，注意 softmax 函数是平移不变的：即，对于任何常数 \(a\)，我们有 \(\operatorname{softmax}_j(x_1+a, x_2+a, \ldots) = \operatorname{softmax}_j(x_1, x_2, \ldots)\)。现在，考虑一个 Query 向量 \(x\) 和 Key 向量 \(y_1, \ldots, y_n\)。对于任何向量 \(v\)，我们有：

\[ \begin{aligned} \operatorname{attention}_j(x) & =\operatorname{softmax}_j(\langle x, y_1 \rangle, \langle x, y_2 \rangle, \ldots) \\ & =\operatorname{softmax}_j(\langle x, y_1 \rangle + \langle x, v \rangle, \langle x, y_2 \rangle + \langle x, v \rangle, \ldots) \\ & =\operatorname{softmax}_j(\langle x, y_1+v \rangle, \langle x, y_2+v \rangle, \ldots) \end{aligned} \]

其中第二步是由平移不变性推导出来的。这意味着在不改变任何给定输入的注意力模式的情况下，我们可以平移所有 Key 向量，使得每个 Attention-head 的 Query 和 Key 分布具有相同的质心。这使得比较 Query 和 Key 变得更容易（见图 3，右）。

缩放 Query 和 Key：在某些 transformer 中，例如 GPT-2，我们观察到平均 Query 范数与平均 Key 范数非常不同的情况。这种差异使得在联合 embedding 中解释 Key - Query 关系变得困难。数学上，这表明点积与距离之间的关系较差；在视觉上，这意味着 Query 可能是一个紧密的簇，周围是一团松散的 Key 云。

幸运的是，缩放是系统的另一个“自由参数”。self-attention 水平仅依赖于 Query 和 Key 向量的点积，因此，如果我们将所有 Query 向量按因子 \(c \neq 0\) 缩放，将所有 Key 向量按因子 \(c^{-1}\) 缩放，注意力值则保持不变。这使得在我们的联合可视化中高注意力的 Query - Key 对能够更靠近，如图 4a 所示。（一个微妙的点是：单独的缩放不改变余弦距离；然而，与平移标准化结合时，它具有重要的效果。）

为了确定 \(c\) 的最佳值，我们可以定义一个加权相关度量，该度量对距离较小的 Query - Key 对给予更重的权重，因为在联合可视化中，我们最关心邻近的 Query 和 Key。因此，我们可以选择一个缩放因子 \(c\)，使得 Query - Key 点积和距离之间的加权相关性最大化。这种缩放方法允许联合 embedding 空间中的距离最准确地代表 Query 和 Key 之间的实际注意力值。

5.1.2 以距离为注意力的代理¶

如上所述，理想情况下，如果一个 Query - Key 对有一个大的、正的点积（对应于一个高的最终 self-attention 值），它们应该在 embedding 空间中被放置得更接近，反之亦然（见图 4a）。因此，我们期望在我们的联合 Query - Key embedding 中，距离与注意力呈反比相关。我们通过计算每个 Attention-head 在 BERT、GPT-2 和 ViT 中余弦距离与点积之间的斯皮尔曼等级相关性来研究这一潜在联系。我们还尝试在创建 Query 和 Key 的 t-SNE 和 UMAP 投影时使用欧几里得距离作为距离度量，但这通常导致距离 - 点积相关性较弱。

在多个数据集和模型中，距离和注意力之间的关系相当稳定。例如，使用 Wiki-Auto 数据 [23]，BERT 的 Query - Key 距离与点积的平均相关性为 -0.938，GPT 的为 -0.792。BERT 的一个示例结果如图 4b 所示。在使用的 COCO 图像集 [28] 上，ViT-32 的平均相关性为 -0.873，ViT-16 的为 -0.884。

图 5：在 BERT 中将形式与功能联系起来。 (a) 在矩阵视图中，第三层有几个螺旋形的图表。 (b) 通过使用单视图放大查看其中一个头部（L3 H9），我们可以使用从浅到深的颜色方案来看到位置注意力模式，该颜色方案 encoding 了输入序列中的位置。 (c) 通过探索句子级可视化，可以确认这些模式。

5.2 颜色 encoding¶

为了可视化 Query 和 Key 的不同属性，AttentionViz 提供了多种颜色 encoding。默认选项按标记类型着色，即 Query 或 Key。对于 ViT，用户可以按图像块的行或列着色，以可视化位置模式（见图 10）。由于图像自带颜色信息，我们也允许用户查看没有额外样式元素的原始块（见图 8）。

对于语言 transformer，我们支持两种位置颜色方案：标准化和离散。为了计算标准化位置，我们将每个标记在句子中的位置除以句子长度，产生一个连续的颜色刻度。较浅的色调表示靠近句子开头的标记（见图 5b）。我们的离散位置 encoding 取每个标记的位置，并应用模运算符得到除以 5 的余数。因此，第 1 个和第 6 个标记接收相同的颜色，第 2 个和第 7 个标记接收相同的颜色，等等。我们使用相同的五种颜色来 encoding 不同位置的 Query 和 Key，对前者使用较深的色调。虽然这种方案对于长度大于 5 的句子引入了歧义，但我们发现我们的离散着色有助于看到基于小偏移位置的关系（例如，如图 11 左所示，Query 关注一步之遥的 Key）。歧义也可以通过悬停在标记上轻松解决，在我们的探索中，我们没有看到大于 2 或 3 的偏移模式。用户也可以按 Query/Key 的范数着色（见图 12a）。

图 6：通过全局搜索探索注意力模式。 (a) 搜索结果较少聚类的头部往往表现出更多的语义行为，而结果分散的头部则更多地关注标记位置。 (b) 放大查看 L2 H6，一个主要结果聚类的头部，我们确实看到了一大群语义相关的 Query 和 Key 标记。

5.3 视图¶

AttentionViz 提供了三种主要的交互式视图用于探索注意力：矩阵视图、单一视图和句子/图像视图。

5.3.1 矩阵视图¶

在 AttentionViz 中的初始视图是矩阵视图，该视图使用小倍数一次性可视化 transformer 中的所有 Attention-head（图 5a），直接解决了 [T1] 和 [T3]。每行对应于一个模型层，从界面顶部的早期层次到底部的后期层次。通过这种“全局”视角，用户可以比单图（例如，[53]）或实例级可视化（例如，[3, 47]）更容易地扫描不同 transformer 层和头部之间的模式。本工作中使用的所有模型都具有相同的架构：12 层 x 每层 12 个头 = 共 144 个 Attention-head，但我们的系统可以扩展到其他维度。

在矩阵视图中，用户可以查看使用 t-SNE、UMAP 或 PCA 创建的联合 Query - Key embedding。他们还可以在模型类型（即 BERT、GPT-2、ViT-16/32）或数据集之间切换 [T4]，探索不同的配色方案，并以 2D 或 3D 查看结果图。矩阵视图支持全局搜索功能（图 6a），这有助于突出显示不同头部和位置中标记的模式，并提供另一种在大规模下分析注意力的方法（见第 7 节）。

5.3.2 单一视图¶

用户可以点击矩阵视图中的任何图表来放大至单一视图（图 5b），这允许更详细地探索单个 Attention-head [T3]。与矩阵视图类似，用户可以在单一视图中切换颜色、维度、投影模式、数据集和模型 [T4]；所有图形变化在视图间同步，以便比较。用户可以点击一个点以突出显示相应输入序列中的所有标记，突出显示我们联合 embedding 空间中相关的 Query 和 Key。用户还可以选择将注意力线投影到散点图上，这些线连接 Query 和 Key 标记（图 5c）。

我们仅显示每个标记的前 2 个注意力权重以增强可读性。我们的注意力线功能支持 [T2]，并提供了一种在头部级别可视化 transformer 中注意力模式的新方式。在单一视图中，用户还可以搜索标记，并使用我们的标签功能来揭示数据中的语义模式，类似于 [47]。例如，在图 6b 中，搜索显示，具有相似含义的 Query/Key 标记在这个 BERT 头的联合 embedding 中放置在一起，表明它们之间有强烈的注意力（第 5.1.2 节）。

图 7：ViT 图像视图。 (a) 原始图像。 (b) 透明度注意力热图。我们用绿色边框突出显示选中的标记（Query）。 (c) 叠加的注意力箭头。 (d) 全局注意力流动。方形图标表示图像块与 \<CLS> 标记有最强的注意力连接，而这个标记并不在原始图像中。

图 8：在 ViT-32 矩阵视图中，我们发现了两个有趣的视觉 Attention-head：一个头根据亮度对黑白图像标记进行排序，而另一个头则根据色调对彩色补丁进行对齐。图像视图中显示的注意力模式确认了具有相同亮度的补丁之间的注意力流动。

5.3.3 句子/图像视图¶

句子/图像视图允许对单个句子或图像中的细粒度注意力模式进行探索 [T2, T3]。这两种视图都与单一视图同步，匹配叠加在每个 Query/Key 散点图上的注意力线，以提供流畅的用户体验。

句子视图：当使用 BERT 或 GPT-2 时，用户可以在单一视图中点击一个点，以在左侧边栏打开句子视图，该视图显示受点击标记突出显示的句子级注意力的 BertViz 启发式可视化 [53]（图 5c）。我们还考虑过使用热图可视化（例如，[39]），但似乎双分图方法将为更长句子的模式探索提供更好的可读性和便利性。连接左列中的 Query 标记和右列中的 Key 标记的线条的不透明度表示它们相应的注意力强度。将鼠标悬停在标记上会突出显示特定于标记的注意力线。为了减少来自 BERT 中的分类标记和分隔符或 GPT-2 中的第一个标记的干扰（第 7 节），用户可以隐藏这些特殊标记的注意力线。其他 Query 和 Key 标记也可以开启/关闭，所有注意力线将相应地重新标准化。用户还可以选择查看每个 Attention-head 的聚合注意力模式，以提供另一层比较（图 11a）。

图像视图：对于 ViT 中基于图像的输入，当用户点击一个图像块时，侧面板将显示其对应的原始图像，并用彩色边框突出显示点击的标记（图 7a）。用户还会看到叠加了注意力热图的图像，其中透明度表示点击的图像块与图像其他区域之间的注意力权重（图 7b）。

除了可视化单个标记的注意力外，图像视图还允许用户通过显示不同图像块之间的箭头注意力线来探索图像内的整体注意力模式。我们为可视化注意力箭头提供了两个选项。第一个选项是在原始图像块的顶部叠加箭头，每个箭头代表起始图像块和目的地块之间最强的注意力连接（图 7c）。这为用户创建了一个简化的双分注意力图，以表征特定头中最重要的模式。第二个选项显示所有强注意力连接（即，\(attn(x_i, x_j) > 0.1\)）在原始图像旁边，提供更全面的注意力视图（图 7d）。在这种可视化中，不透明度和线条厚度用于 encoding 注意力连接的强度。我们也尝试可视化 Query 和 Key 之间的所有权重以更紧密地模仿 [53]，但这通常产生过于拥挤、难以理解的结果。

6 系统实现¶

为了处理模型输入并计算注意力信息，我们使用了 Hugging Face Transformers 库和 PyTorch。我们使用了来自 Google 和 OpenAI 的 BERT、GPT-2（小型）和 ViT-16/32 的预训练实现。对于每个 NLP 数据集，我们随机抽样 200 个句子（每个 Attention-head 约 10k 个标记，包括 Query 和 Key）。由于图像注意力数据的计算规模增加，我们为 ViT-32 显示每个头 10 张图片（1000 个标记），对于 ViT-16 是每个头 4 张图片（1576 个标记）。更大的数据集可以输入到 AttentionViz 中，但以增加系统延迟为代价。目前，加载每个 NLP 数据集需要约 6 秒，加载 ViT-16 数据需要约 10 秒。在为每个 Attention-head 提取 Query 和 Key 向量 embedding 后，我们生成相应的 2D/3D t-SNE、UMAP 和 PCA 坐标（第 5.1 节）。为了为 ViT 中的图像块产生语义标签（例如，“狗”或“背景”），我们使用了 DeepLabv3 分割模型 [7]。总的来说，预处理每个 BERT/GPT-2 数据集大约需要 3 小时，在 NVIDIA A100 GPU 上；对于 ViT-32，大约需要 30 分钟。

我们的最终 AttentionViz 原型由一个 Python/Flask 后端组成，该后端与用 Vue 和 Typescript 编写的前端通信。演示系统可在：http://attentionviz.com 访问。由于数据体积大和浏览器内存限制，我们通过后端加载预先计算的注意力/投影信息 JSON 文件。对于 ViT，后端还执行图像处理（例如，贴片突出显示和透明度调整）以在前端显示。我们使用 Deck.gl 来可视化生成的 Query - Key 联合 embedding。AttentionViz 高度可扩展且与模型无关，允许用户向系统添加新的 transformer 和数据集。

图 9：左图：具有不同频率和角度的空间模式数据集。中间：在单视图中，我们观察到 ViT-32 的一个 Attention-head 根据它们的空间模式的频率和角度排列图像标记。右图：图像视图中的注意力热图进一步确认了这些发现——具有相似角度的空间模式更加关注彼此。

7.1 目标：理解机器视觉注意力¶

AttentionViz 在揭示视觉变换器中的注意力新见解方面尤其有帮助，这得益于图像补丁数据的本质上是视觉的 [G1]。视觉注意力中的色调/亮度特化。我们好奇是否有任何视觉 Attention-head 专注于基于颜色和亮度的模式。为了测试这一点，我们向预训练的 ViT-32 模型提供了合成的颜色和亮度梯度图像（图 8），并将结果的 Query 和 Key 标记加载到 AttentionViz 中。在矩阵视图中浏览全局 PCA 模式，我们识别了两个类似于有色和无色视觉的 Attention-head。一个头部根据亮度对黑白图像标记进行对齐，另一个则根据色调对彩色补丁进行对齐。我们的数据集包含所有方向的颜色和亮度梯度图像，我们看到相似的补丁在联合 embedding 空间中聚集在一起，不受其在原始图像中的位置的影响。图像视图中的注意力热图确认了这些发现；标记最多关注与它们颜色或亮度相同的其他标记。E7 对这些结果感到好奇，他之前研究过卷积神经网络（CNN）的颜色潜在空间，并表示有兴趣使用我们的工具进一步探索 CNN 和 ViT 行为之间的差异。

频率过滤和角度检测。频率和角度是图像数据的低级特征。为了调查视觉变换器是否有一个 Attention-head 能够根据这些特征将视觉模式关联起来，我们创建了具有不同频率和方向的正弦信号图像，并使用我们预训练的 ViT-32 模型对其进行处理。在矩阵视图中检查结果的 Query 和 Key embedding，我们识别了一个将图像标记基于其空间模式的频率（x 轴）和角度（y 轴）分开的 Attention-head（图 9）。通过图像视图，我们观察到同心圆的图像中的标记关注于具有相似曲率的其他标记，进一步确认这个 Attention-head 根据它们的角度关联视觉模式。 E7 表示这个结果很有趣，但鉴于我们的色调/亮度发现，并不太令人惊讶，他更好奇的是那些不展现这种“关注相似补丁”行为的头部。他们提出的一个实验是研究注意力调节，例如，如果同一个图像补丁（例如，竖条纹）在两幅图像中的不同上下文中出现（例如，斑马与雨伞），我们是否会看到独特的注意力模式？

模型层间注意力距离的增加。如 [14] 所述，视觉变换器在更深的层中的 self-attention 在图像中的关注范围更广。我们使用 AttentionViz 的交互式、联合可视化 Query 和 Key 标记确认了这一发现。在矩阵视图中，我们通过图像“行”和“列”对补丁进行着色，找到了 ViT-32 第 1 层和第 2 层的四个 Attention-head，它们将标记与其最近的空间邻居分组：左侧、右侧、顶部和底部。在第 3 层和第 4 层，我们看到了类似的位置注意力模式，但图像标记关注于同一行或列中的所有补丁，超出了它们最近的邻居（图 10）。这表明，与使用方形滤波器处理图像的 CNN 不同，变换器中的 self-attention 机制通常逐行和逐列处理图像，类似于一个拉长的滤波器。

图 10：通过按行对图像补丁进行着色，突出显示了 ViT-32 中的位置注意力模式。在第 1 层，同一行和相邻列的标记形成小型群集。图像视图揭示了一种向左看的模式。在第 4 层，基于行位置的标记形成大型群集。使用带箭头的线条，我们可以看到更广泛的双向注意力流动。

7.3 目标：识别异常和意外行为¶

图 11：其他视觉迹象的注意力。左侧：具有小“团块”的头部通常比螺旋形的位置模式更紧密。我们的离散位置 encoding，根据每个标记的位置模 5 着色，突出显示了一个“下一个标记”的注意力趋势。右侧：Query 和 Key 的分层带状只出现在 SuperGLUE AXb 数据 [56] 中，表明对文本开始、结束和中点的强烈关注。

通过与 AttentionViz 中的联合 Query - Key embedding 交互，我们发现了一些不规则的模型行为 [G3]。

normalization 差异和空注意力。在研究 GPT-2 的矩阵视图时，我们观察到在早期模型层中，一些 Query 和 Key 群集即使在 Key 转换后（第 5.1.1 节）也保持良好的分离。通过按 normalization 着色（在 normalization 缩放步骤之前测量），我们看到在许多头部中，Query 和 Key 向量的 normalization 之间存在显著差异（图 12a）。当 Querynormalization 较小（即浅绿色）时，Keynormalization 倾向于较大（即深粉色），反之亦然。计算 GPT-2 与 BERT 中 Query 和 Key 的平均 normalization 差异，我们发现在前者中，平均 Querynormalization - Keynormalization = -4.59，而在后者中，平均差异仅为 0.41。没有任何专家能解释这一发现：“Query 和 Key 的 normalization 为何会有如此大的差异，真是说不通”（E6）。有趣的是，我们做出这一观察后发表的一篇论文 [11] 指出，失控的 Query 和 Keynormalization 可能导致严重的训练不稳定，表明这种现象可能值得进一步研究。这一观察也启发了我们从第 5.1.1 节开始的缩放方法。

我们还注意到，在许多 GPT-2 的头部中，大部分注意力都集中在第一个标记上（图 12b），尤其是在后期层中。[54] 简要提到，当 Attention-head 捕捉的语言属性没有出现在输入文本中时，GPT-2 会将第一个标记视为一个空位置用于接收注意力。然而，这一现象仍然未被充分探索，提出了另一个开放的可解释性问题。E2 和 E6 都单独使用我们的工具注意到了这种异常行为，所有的专家都对这一发现感到惊讶。[55] 显示，修剪变换器中的大部分 Attention-head 可能不会显著影响模型性能，这或许可以部分归因于这种占主导地位的空注意力模式。无论如何，AttentionViz 允许用户过滤掉对第一个标记的注意力，揭示隐藏的 Query - Key 交互。

“自我关注”Attention-head。AttentionViz 也可以揭示视觉变换器中的令人惊讶的注意力模式。在矩阵视图中，我们识别了 ViT-32 早期层中的几个 Key - Query 群集非常分散的头部（图 13a）。观察其中一个 Attention-head（第 0 层头 8），我们发现同一个标记的 Query 和 Key embedding 形成了一个小而密集的群集，每对 Query - Key 之间都与其他群集明显分离（图 13b）。从图像视图中的透明度热图中，我们看到补丁仅关注自身（图 13c）。切换到带箭头的注意力线，我们发现这个图像的整体注意力模式是“自我关注”，在这个头部中图像标记之间没有信息流动。

在识别这种不规则的注意力模式后，我们用相关性测试检查了 Query 和 Key 矩阵的学习参数。我们发现一个强相似性分数（线性相关性 = 0.94），表明这个 ViT 头部的 Query 和 Key 层确实在学习冗余的投影（图 13d）。E3 指出，这一知识可以用于指导模型修剪实验，而 AttentionViz 也可以类似地用于检测潜在的训练失败或其他不规则性。

图 12：GPT-2 中的异常情况。 (a) 在早期模型层中，我们看到许多 Attention-head 的 Query-Key normalization 之间存在显著差异（例如，normalization 缩放前的 L1 H8）。 (b) 后期层中普遍存在的“首先关注”模式的示例。句子视图在隐藏第一个标记后揭示了潜在的注意力行为。

图 13：在 ViT-32 中识别一个“自我关注”的 Attention-head。 (a) 单视图显示 Query 和 Key 在联合 embedding 空间中分布稀疏。 (b) 放大观察，同一图像标记的 Query 和 Key 向量紧密重叠。 (c) 图像视图揭示标记最多关注自己。 (d) 比较 Query 和 Key 投影层的学习参数确认它们学习了冗余的投影。

8 结论与未来工作¶

在本项工作中，我们介绍了一种新技术，用于可视化基于 Query 和 Key 的联合 embedding 空间的变换器 self-attention。应用我们的技术，我们创建了 AttentionViz（演示：http://attentionviz.com），一个交互式可视化工具，并使用它来获得关于语言和视觉变换器中注意力的洞察。例如，我们在 ViT 中发现了新的色调/频率行为，并在 GPT-2 中发现了显著的 Query-Keynormalization 差异。尽管我们的方法专为 self-attention 设计，但它可以通过对向量归一化进行适当的修改，推广到其他注意力机制（例如，交叉注意力）。同样，我们关注语义模式，但 AttentionViz 也可以用来研究句法特征，例如，通过使用 NLP 库添加元数据，如词性标注。

专家反馈还指出了未来工作的几个方向，例如找到管理多重 embedding 可视化的复杂性和将用户聚焦于感兴趣特征的方法。我们的系统目前受到数据预处理时间和内存需求的限制，但我们计划改进 AttentionViz 的可扩展性，并允许用户即时添加新输入。足够大的随机样本（例如，几千个标记）足以表现大规模的全球现象，但更大的数据集和模型可能会揭示额外的语义洞察，正如 [38] 所建议的。尝试其他数据采样方法也可能是有益的（例如，减少来自齐普夫定律 [62] 的偏差）。

未来研究的另一个自然方向是探索如何在每个 Attention-head 中整合来自值向量的信息 [52]。这些值向量是注意力机制的一个重要组成部分，尽管目前还不清楚如何在 Query 和 Key 的背景中可视化它们。找到合适的可视化方法可能会让我们更深入了解 Attention-head 的功能。最后，尽管 AttentionViz 是一个探索性工具，将其适配用于假设测试和/或因果追踪可能为实际模型调试提供支持。