注意力图可视化

了解 Transformer 模型内部运作方式最直接的途径之一是检查其注意力机制 (attention mechanism)。自注意力 (self-attention)作为 Transformer 架构的主要组成部分，让模型在计算特定标记 (token)的表示时，能够衡量输入序列中不同标记的重要性。这些注意力权重 (weight)，为每一层中的每个头计算，构成图谱，显示信息如何在模型中传递。可视化这些图谱可以提供模型在标记之间学习到的联系的线索。

自注意力的核心思想涉及计算一个标记的查询向量 (vector)($Q$)与序列中所有标记（包括其自身）的键向量($K$)之间的分数。这些分数经过缩放，并使用 softmax 归一化 (normalization)，然后用于计算值向量($V$)的加权和。注意力权重是缩放点积后应用 softmax 的结果：

$$\text{权重} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$$

这里，$d_k$是键向量的维度。这些权重表示每个查询标记对所有键标记的注意力分布。权重越高表示模型认为对应的键标记在生成查询标记的表示时更重要。

提取注意力权重 (weight)

大多数现代深度学习 (deep learning)框架，包括 PyTorch，都提供机制来在前向传播过程中访问这些注意力权重。当使用 PyTorch 的 nn.MultiheadAttention 层时，你可以在前向调用中指定need_weights=True。这个参数 (parameter)指示层除了输出之外，还返回所有头的平均注意力权重。对于更细致的、特定于头的权重，你可能需要稍微修改层的实现，或使用钩子在权重平均之前捕获它们。

这里是一个简化示例，展示如何在 PyTorch 中从 nn.MultiheadAttention 层获取注意力权重：

import torch
import torch.nn as nn

# 示例设置
seq_len = 5
embed_dim = 8
num_heads = 2
batch_size = 1

# 确保 embed_dim 可以被 num_heads 整除
assert embed_dim % num_heads == 0

mha_layer = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)

# 虚拟输入 (批量大小, 序列长度, 嵌入维度)
query = torch.randn(batch_size, seq_len, embed_dim)
key = torch.randn(batch_size, seq_len, embed_dim)
value = torch.randn(batch_size, seq_len, embed_dim)

# 请求注意力权重的前向传播
# attn_output: (批量大小, 序列长度, 嵌入维度)
# attn_output_weights: (批量大小, 序列长度, 序列长度) -> 各头平均
attn_output, attn_output_weights = mha_layer(
    query, value,
    need_weights=True,
    average_attn_weights=True
) # 对于每个头的权重设置为 False (如果层支持/已修改)

print("平均注意力权重的形状:", attn_output_weights.shape)
# 如果 average_attn_weights=False (且层已修改/钩子):
# 形状将是 (批量大小, 头数, 序列长度, 序列长度)

# 示例: 访问第一个批量项目的权重
first_batch_weights = attn_output_weights[0] # 形状: (序列长度, 序列长度)
# first_batch_weights[i, j] 是查询标记 i 对标记 j 的注意力

# 获取每个头的权重 (通常需要修改或钩子)
# _, attn_output_weights_per_head = mha_layer(
#     query, value,
#     need_weights=True,
#     average_attn_weights=False
# )
# print("每个头的注意力权重形状:",
#       attn_output_weights_per_head.shape)
# first_batch_head_0_weights = attn_output_weights_per_head[0, 0] # 第0个头

请注意，标准的 nn.MultiheadAttention 在 average_attn_weights 为 True（如果 need_weights 为 True 时的默认值）时，返回在各头之间平均后的权重。访问单个头的权重通常需要修改前向方法，或者更简洁地，在注意力机制 (attention mechanism)的内部 softmax 或矩阵乘法操作上注册一个前向钩子，以便在平均之前捕获权重。

可视化方法

一旦提取，注意力权重 (weight)（通常是每个头/层的 (序列长度, 序列长度) 大小的矩阵）可以通过几种方式可视化：

1. 热力图： 这是最常用的方法。热力图显示注意力矩阵，其中行表示查询标记 (token)（输出位置），列表示键标记（输入位置）。单元格 (i, j) 的颜色强度表示标记 i 对标记 j 的注意力权重。颜色越浅通常表示注意力越高。分析这些热力图可以显示出一些模式，例如强对角线（标记关注自身）、对前置标记的注意力，或者特定标记（如标点符号或特殊标记）充当信息接收器或源。

单个头的注意力权重。请注意表示自注意力 (self-attention)的强对角线，以及“sat”如何强烈关注“cat”。特殊标记 [CLS] 主要关注自身，而 [SEP] 也显示出高自注意力。

2. 多头可视化： 由于每个层包含多个注意力头，将它们全部可视化很重要。常用方法包括：

   • 小倍数图： 显示一个热力图网格，每个头一个。这允许直接比较不同头学习到的模式。
   • 平均热力图： 显示一个代表某层中所有头平均权重的单个热力图。这提供了一个概览，但可能掩盖特定于头的行为。

3. 基于图的可视化： 注意力权重可以表示为有向图，其中标记是节点，如果注意力权重 $w_{ij}$ 超过某个阈值，则存在从标记 i 到标记 j 的有向边。边的粗细或颜色可以表示权重的量级。这对于可视化较短序列中的连接或突出特定的强联系可能有效。

简化图，显示强注意力链接。“sat”强烈关注“cat”，而“cat”对“the”也有明显关注。

解释和理解

分析注意力模式有时可以显示出语言学上合理的行为：

• 句法依赖： 头可能学习到关注句法相关的词语，比如动词关注它们的主语或宾语，或形容词关注它们修饰的名词。
• 指代消解： 注意力可能将代词与其在文本中较早提及的名词联系起来。
• 位置信息： 有些头经常集中关注紧邻的前一个或后一个标记 (token)，或相对位置偏移。
• 特殊标记： 像 [CLS] 或 [SEP] 这样的标记可能汇总整个序列的信息，这表现为源自或指向它们的广泛注意力模式。
• 层级进展： 注意力模式在不同层之间常有差异。早期层可能关注局部、句法关系，而更靠后的层可能捕获更复杂、长距离或语义上的连接。

局限性和注意事项

虽然注意力可视化是一种有用的工具，但了解其局限性很重要：

• 注意力不等于解释： 高注意力权重 (weight)不一定意味着某个标记 (token)是特定输出的主要原因。注意力表示在构建下一层表示时，哪些标记的表示被赋予了高权重，但前馈网络内部和跨层间的复杂转换模糊了直接的因果关系。研究表明，注意力权重可能不总是与梯度等其他特征重要性指标强相关。
• 平均化问题： 在各头之间平均权重（某些框架实现中的默认做法）可能隐藏单个头学习到的多样化甚至矛盾的模式。有些头可能学习到专门的功能，而另一些则显得噪声或冗余。
• Softmax 饱和/扩散： Softmax 函数强制权重总和为1。如果没有哪个标记明确重要，注意力可能会分散到许多标记上，使得解释变得困难。反之，如果某个标记高度相关，其权重可能接近1，从而抑制其他潜在相关标记的可见权重。
• 深度模型的复杂性： 在非常深的 Transformer 中，层之间传递的表示变得越来越抽象。后面层中的注意力模式作用于这些复杂的表示，使得它们直接映射回原始输入标记的解释变得不那么直接。

注意力图可视化提供了一个窗口（尽管有时模糊），用于了解 Transformer 内部的信息流。这是一种有用的诊断技术，用于对模型行为生成假设并找出潜在的关注点，但应谨慎得出结论，最好与本章后面讨论的其他分析方法相互印证，例如探查内部表示或分析神经元激活。