注意力权重的可视化与可解释性从热力图到归因分析的工具链Transformer模型的可解释性研究日益受到关注注意力权重作为自然的信息瓶颈提供了观察模型内部决策过程的窗口。本文系统梳理了注意力可视化的四层工具链原始注意力热力图、头重要性剪枝、输入归因分析和电路发现分析每层方法的信息粒度和适用场景并给出基于bertviz、captum和transformer_lens的具体实现。一、注意力权重的信息瓶颈特性Transformer的每一层自注意力计算可以表述为$$ \text{Attention}(Q,K,V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$注意力权重矩阵 $A \in \mathbb{R}^{n \times n}$ 是输入token之间信息流动的显式表示$A_{ij}$ 量化了token $i$ 从token $j$ 聚合信息的程度。这一矩阵构成了模型内部唯一的结构化、可人类阅读的信息表示——残差流的加性混合使得其他中间表示难以直接解读。然而原始注意力权重存在三个被广泛讨论的局限性其一注意力权重的高稀疏性长尾分布多数token间的注意力接近零使得热力图的可读性随序列长度增加而急剧下降其二多层多头的组合使单一头的注意力模式无法独立解释模型行为其三注意力权重反映的是信息混合比例而非因果重要性——高注意力并不等同于高重要性。二、第一层注意力热力图与模式识别注意力热力图是最直观的可视化手段。对于编码器模型如BERT可以通过热力图观察不同层和不同头捕捉的语言学模式。典型的可识别模式包括对角注意力低层常见每个token主要关注自身或相邻token反映局部句法依赖垂直注意力中层常见特定token如[CLS]或标点获得广泛关注全局注意力高层常见注意力分布趋于均匀化反映高层语义聚合# 使用 bertviz 进行多头注意力可视化 from transformers import AutoModel, AutoTokenizer import torch def visualize_attention_patterns( model_name: str bert-base-uncased, text: str The cat sat on the mat because it was comfortable. ): 提取并分析 BERT 各层各头的注意力矩阵。 bertviz 的核心接口model(**inputs, output_attentionsTrue) 返回的 attentions 元组形状为 (num_layers, batch, num_heads, seq_len, seq_len) tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModel.from_pretrained( model_name, output_attentionsTrue # 关键参数强制返回所有注意力矩阵 ) inputs tokenizer(text, return_tensorspt) with torch.no_grad(): outputs model(**inputs) # attentions 是一个 tuple长度等于 num_layers # 每个元素的形状: (batch_size, num_heads, seq_len, seq_len) attentions outputs.attentions tokens tokenizer.convert_ids_to_tokens(inputs[input_ids][0]) # 分析每一层的注意力熵——衡量注意力分布的集中程度 for layer_idx, layer_attn in enumerate(attentions): # layer_attn shape: (1, 12, seq_len, seq_len) avg_entropy 0.0 for head_idx in range(layer_attn.shape[1]): attn_matrix layer_attn[0, head_idx] # 计算每个 query token 的注意力熵 # 低熵 注意力集中在一个或少数 token 上 # 高熵 注意力均匀分布在多个 token 上 entropy -torch.sum( attn_matrix * torch.log(attn_matrix 1e-12), dim-1 ).mean() avg_entropy entropy.item() avg_entropy / layer_attn.shape[1] print(fLayer {layer_idx}: avg entropy {avg_entropy:.3f}) return tokens, attentions热力图的问题在于当序列长度超过50时单个头的热力图已经难以用人眼解析出有意义的模式。此时需要转向更高层级的抽象方法。三、第二层头重要性的量化分析头重要性分析回答的问题是如果移除这个注意力头模型性能会下降多少这是一个典型的消融实验ablation study设计。常用的头重要性度量包括基于梯度的头重要性计算损失函数对头输出掩码的梯度梯度越大说明头对当前预测越重要。基于消融的头重要性将头的输出替换为零zero ablation或均值mean ablation测量任务指标的变化幅度。def compute_head_importance_by_ablation( model, tokenizer, eval_dataset, task_metric_fn ): 通过消融实验计算每个注意力头的重要性得分。 方法逐头进行 zero ablation记录任务指标的下降幅度。 指标下降越多该头越重要。 import copy import numpy as np # 基线性能 baseline_score task_metric_fn(model, eval_dataset) head_importance {} for layer_idx in range(model.config.num_hidden_layers): for head_idx in range(model.config.num_attention_heads): # 创建带特定头掩码的模型副本 masked_model copy.deepcopy(model) # 注册 hook 以在注意力计算后将指定头的输出置零 def make_head_mask_hook(target_layer, target_head): def hook(module, input, output): # output shape: (batch, seq, hidden_size) head_dim output.shape[-1] // module.num_attention_heads start target_head * head_dim end start head_dim output[:, :, start:end] 0.0 return output return hook attn_layer masked_model.encoder.layer[layer_idx].attention.self handle attn_layer.register_forward_hook( make_head_mask_hook(attn_layer, head_idx) ) # 评估消融后的性能 ablated_score task_metric_fn(masked_model, eval_dataset) handle.remove() importance baseline_score - ablated_score head_importance[(layer_idx, head_idx)] importance return head_importance一系列实验表明BERT中约60%的注意力头可以在不显著影响下游性能的情况下被剪枝这与过参数化假说一致。重要的头通常集中在中间层6-9层这些层被认为是句法-语义转换的关键阶段。四、第三层和第四层归因分析与电路发现输入归因分析的目标是将模型预测分解为各输入token的贡献。在注意力机制的场景下Attention Rollout和Attention Flow是两种经典方法。Attention Rollout基于注意力矩阵的连乘$\tilde{A}^{(l)} A^{(l)} \cdot \tilde{A}^{(l-1)}$通过逐层累积注意力流来估计输入token对最终表示的影响。其数学直觉是如果第l层token i关注token j而第l-1层token j关注token k那么信息从k到i的流动路径可以被建模为两次注意力的乘积。**电路发现Circuit Discovery**是当前可解释性研究的前沿方向。其核心假设是模型中的特定功能如间接对象识别、大写字母检测可以由一小部分attention头和MLP神经元组成的电路来解释。ACDCAutomatic Circuit Discovery算法通过迭代式消融来自动发现这些功能性子网络相比于穷举搜索将计算复杂度从$O(2^N)$降至$O(N^2)$。五、总结本文构建了注意力可解释性的四层工具栈热力图可视化适合初步探索和模式发现头重要性消融实验提供了量化的贡献度评估Attention Rollout等归因方法将模型预测与输入token建立因果链接电路发现算法在功能层面揭示了模型内部的模块化子网络。每一层方法服务于不同的分析目的——从定性的模型在看什么到定量的哪些部分在驱动决策再到因果性的这个功能是如何实现的。将这些方法组合使用构成了当前Transformer可解释性分析的主流范式。