注意力权重可视化基础

咱们直接进入正题。注意力权重可视化，说白了就是把GNN模型"看"哪里给亮出来。我刚开始做这个的时候，也觉得不就是画个热力图吗？后来发现，权重提取这一步就够你折腾半天的。

3.1 权重矩阵提取

权重矩阵藏在哪儿？对于图注意力网络（GAT），注意力系数就在 forward 过程中的 alpha 变量里。我个人的习惯是注册一个 forward hook，在模型跑前向时把注意力系数"截获"下来。

来，直接看代码。下面这段是从 PyTorch Geometric 的 GAT 层里提取权重的典型写法：

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv

# 定义带钩子的包装类
class GATWithAttn(torch.nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.conv = GATConv(in_dim, out_dim, heads=4, concat=False)
        self.attn_weights = None

    def forward(self, x, edge_index):
        # 注册钩子
        self.conv.attn = None  # 重置
        out = self.conv(x, edge_index)
        # 注意力系数存在 self.conv 的 _alpha 属性中
        # 需要手动从 GATConv 内部取出
        # 这里为了演示，我们用 hook 的方式
        return out

def get_attn_hook(module, inputs, outputs):
    # 在 GATConv 的 forward 中，注意力系数被保存在 module._alpha
    if hasattr(module, '_alpha') and module._alpha is not None:
        module._cached_attn = module._alpha.detach().cpu()

model = GATWithAttn(16, 8)
handler = model.conv.register_forward_hook(get_attn_hook)

# 跑一次前向
x = torch.randn(10, 16)
edge_index = torch.randint(0, 10, (2, 40))
out = model(x, edge_index)

# 提取注意力权重
attn_weights = model.conv._cached_attn
print(f"注意力权重形状: {attn_weights.shape}")
# 输出: (边数, 头数) 或 (头数, 边数) 取决于版本

拿到原始权重后，通常需要 reshape 成 (头数, 节点数, 节点数) 的稠密矩阵，才能方便地画热力图。为什么会这样？因为稀疏格式的边列表虽然省内存，但可视化的时候你得把缺失的位置补成 0。

import torch
from torch_geometric.utils import to_dense_adj

# 假设 attn_weights 形状为 (边数, 头数)
# 需要转成稠密邻接矩阵形式

edge_index = edge_index  # (2, 边数)
num_nodes = x.shape[0]
num_heads = attn_weights.shape[1]

# 为每个头构造稠密注意力矩阵
attn_dense = []
for h in range(num_heads):
    # 取出第 h 头的注意力值
    attn_h = attn_weights[:, h]
    # 转成稠密矩阵 (节点数 x 节点数)
    dense_h = to_dense_adj(edge_index, max_nodes=num_nodes, edge_attr=attn_h)
    attn_dense.append(dense_h.squeeze(0))

attn_dense = torch.stack(attn_dense, dim=0)
print(f"稠密注意力矩阵形状: {attn_dense.shape}")
# 输出: (头数, 节点数, 节点数)

嗯，这一步要注意：to_dense_adj 会把没有边的位置填成 0，所以你的注意力矩阵里会有大量的 0 值区域。这是正常的。

3.2 热力图绘制基础

权重提出来了，下一步就是画图。我个人偏爱 seaborn 的 heatmap，配色干净，参数直观。不过 matplotlib 的 imshow 也能干，就是调颜色麻烦点。

来一个直接能跑的例子：

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# 取第 0 个头的注意力矩阵
attn_map = attn_dense[0].numpy()  # (节点数, 节点数)

plt.figure(figsize=(8, 6))
sns.heatmap(
    attn_map,
    cmap='YlOrRd',      # 黄-橙-红渐变，高亮明显
    annot=False,         # 数据点太多时不要标数字
    cbar=True,           # 显示颜色条
    square=True,         # 格子保持正方形
    linewidths=0.2,     # 格子间距线宽
    linecolor='white'
)
plt.title('Head 0 注意力权重热力图', fontsize=14)
plt.xlabel('目标节点')
plt.ylabel('源节点')
plt.tight_layout()
plt.show()

我在项目中遇到过一个问题：多个头的热力图怎么排布比较直观？我的做法是画一个 2x2 或 3x3 的子图矩阵，每个子图对应一个头，统一颜色范围。

fig, axes = plt.subplots(2, 2, figsize=(12, 10))
axes = axes.flatten()

# 统一颜色范围
vmin = attn_dense.min().item()
vmax = attn_dense.max().item()

for h in range(4):
    sns.heatmap(
        attn_dense[h].numpy(),
        ax=axes[h],
        cmap='YlOrRd',
        vmin=vmin,
        vmax=vmax,
        cbar=True,
        square=True
    )
    axes[h].set_title(f'头 {h}')

plt.suptitle('GAT 4头注意力热力图对比', fontsize=16)
plt.tight_layout()
plt.show()

这样一眼就能看出哪个头关注的是全局，哪个头只盯着局部。我记得有次分析一个分子图分类任务，头 2 和头 3 的模式完全不一样——一个抓官能团内部连接，一个抓分子间长程作用。这种差异从热力图上看特别明显。

3.3 注意力模式初探

热力图画出来了，怎么看？我总结了几种常见的注意力模式，用表格列出来，你对照着看就行：

你想想看，如果热力图上全是均匀模糊型，那模型基本没学到什么东西。我遇到过一个失败的案例——训练了 200 轮的 GAT，热力图全是均匀的黄色，后来发现是学习率太大，注意力参数全崩了。

下面这张图概括了从模型到模式分析的完整流程，我把它画成了 SVG，方便你理解各个环节的关系：

最后说一个我踩过的坑。有一次我把注意力矩阵画出来，发现所有头几乎一模一样，当时百思不得其解。后来排查了半天，原来是共享了同一个 hook 对象的引用，所有头的权重都指向同一份内存。嗯，这种 bug 光看代码很难发现，热力图一出来就现原形了。

这一节的内容就这些。你拿到手里的应该是权重的提取方法、热力图的绘制套路，以及五种常见的注意力模式。下回看到一张注意力热力图，至少能说出"哦，这是局部集中型，模型在聚合邻居信息"这样的话。