多头注意力的计算冗余分析：哪些头可以被安全剪掉而不损失精度

多头注意力的计算冗余分析：哪些头可以被安全剪掉而不损失精度

一、多头注意力并非"头越多越好"

Transformer 论文引入多头注意力的初衷是让不同头关注不同的表示子空间——可以类比为 CNN 中不同卷积核学习不同特征。然而后续研究（Voita et al., 2019; Michel et al., 2019）发现了一个被忽视的事实：在训练收敛后，很多注意力头的行为高度冗余。具体表现为：某些头可以被完整移除而模型精度几乎不受影响；某些头在不同输入上的注意力分布几乎完全相同（可合并）；还有些头在任何输入上都输出近乎均匀的注意力分布（干脆不工作）。

本文聚焦注意力头冗余的量化检测方法和安全剪枝策略。核心问题是：如何在训练后识别哪些头是冗余的，以及剪枝后如何在不重新训练的前提下恢复精度。

flowchart TD A[已训练的Transformer模型] --> B[头重要性评估] B --> C1[基于权重的评估] B --> C2[基于激活的评估] B --> C3[基于梯度的评估] C1 --> C1a[L1范数: 权重绝对值求和] C2 --> C2a[注意力熵: 分布的集中度] C2 --> C2b[头间余弦相似度: 检测冗余] C3 --> C3a[Taylor展开: 对loss的影响] C1a & C2a & C2b & C3a --> D[重要性排序] D --> E[剪枝策略选择] E --> F1[逐层均匀: 每层剪相同比例] E --> F2[全局排序: 跨层竞争] E --> F3[层敏感: 底层保留更多头] F1 & F2 & F3 --> G[剪枝 + 微调恢复] G --> H[精度验证]

二、三种头重要性评估方法的优劣

方法一：基于权重的L1范数。计算每个头对应的 Q、K、V、O 投影权重的 L1 范数之和。优点是计算成本极低（不需要任何前向传播），缺点是完全忽略了输入数据的影响——某个头的权重大，可能只是因为初始化比别的头大，而非它学到了更有意义的表示。

方法二：基于注意力分布的熵。对一批代表性输入，收集每个头的注意力分布，计算其熵。高熵意味着分布接近均匀——这个头没有在"选择性关注"，而是在"平均地看所有位置"。这种头的信息贡献很小，通常是首选的剪枝目标。缺点是计算成本随数据量和序列长度线性增长。

方法三：基于Taylor展开的损失敏感度。计算移除某个头对训练损失的近似影响（一阶Taylor展开），即"权重 × 梯度"的乘积。这是理论最严谨的方法，因为它直接估计剪枝对模型性能的影响，而非间接度量"头的重要性"。缺点是需要在训练数据上做一次完整的 forward+backward pass。

import torch import torch.nn as nn from typing import Dict, List, Tuple import numpy as np def compute_head_importance_by_entropy( model: nn.Module, dataloader, num_samples: int = 500, device: str = "cuda" ) -> Dict[str, List[float]]: """通过注意力分布熵评估每个头的信息集中度。 原理： 一个有效的注意力头应该对少数关键位置赋予高注意力权重 （低熵）。如果某个头对所有位置均匀分配注意力（高熵）， 说明它没有学到有意义的位置选择性——可以安全剪掉。 为什么用熵而非方差： 方差在大值和大量小值混合时也会很高，但高方差不一定 意味着信息集中。熵直接度量分布的"峰度"。 """ model.eval() head_entropies = {} with torch.no_grad(): sample_count = 0 for batch in dataloader: if sample_count >= num_samples: break batch = {k: v.to(device) for k, v in batch.items()} outputs = model(**batch, output_attentions=True) # outputs.attentions: tuple of [batch_size, num_heads, seq_len, seq_len] for layer_idx, layer_attn in enumerate(outputs.attentions): # 计算每个头的平均熵 bsz, n_heads, seq_len, _ = layer_attn.shape # 熵 = -sum(p * log(p)), 对最后一维（目标位置） # 加极小值避免 log(0) eps = 1e-9 entropies = -torch.sum( layer_attn * torch.log(layer_attn + eps), dim=-1 ).mean(dim=(0, 2)) # 在 batch 和 query 位置维度上平均 if layer_idx not in head_entropies: head_entropies[layer_idx] = [] head_entropies[layer_idx].append(entropies.cpu()) sample_count += batch[list(batch.keys())[0]].size(0) # 聚合所有样本的熵（取均值） aggregated = {} for layer_idx, entropies_list in head_entropies.items(): stacked = torch.stack(entropies_list) # [num_batches, num_heads] aggregated[f"layer_{layer_idx}"] = stacked.mean(dim=0).tolist() return aggregated def compute_head_cosine_similarity( model: nn.Module, dataloader, num_samples: int = 200, device: str = "cuda" ) -> Dict[str, np.ndarray]: """检测同层内注意力头之间的行为相似度。 原理： 如果两个头在所有输入上的注意力分布高度相似 （余弦相似度 > 0.9），则它们捕获的信息几乎重叠。 可以保留其中一个，剪掉另一个。 """ model.eval() # 存储每个头在所有样本上的注意力分布（展平） head_patterns = {} with torch.no_grad(): sample_count = 0 for batch in dataloader: if sample_count >= num_samples: break batch = {k: v.to(device) for k, v in batch.items()} outputs = model(**batch, output_attentions=True) for layer_idx, layer_attn in enumerate(outputs.attentions): bsz, n_heads, seq_len, _ = layer_attn.shape # 展平为 [n_heads, bsz * seq_len * seq_len] flat_attn = layer_attn.permute(1, 0, 2, 3).reshape(n_heads, -1) if layer_idx not in head_patterns: head_patterns[layer_idx] = [] head_patterns[layer_idx].append(flat_attn.cpu()) sample_count += batch[list(batch.keys())[0]].size(0) # 计算同层内头间余弦相似度 similarity_matrices = {} for layer_idx, patterns_list in head_patterns.items(): # 拼接所有样本: [n_heads, total_elements] all_patterns = torch.cat(patterns_list, dim=1) # 行归一化后计算余弦相似度矩阵: [n_heads, n_heads] normed = all_patterns / (all_patterns.norm(dim=1, keepdim=True) + 1e-9) sim_matrix = (normed @ normed.T).numpy() similarity_matrices[f"layer_{layer_idx}"] = sim_matrix return similarity_matrices

三、安全剪枝策略与恢复微调

识别出冗余头后的操作顺序直接影响最终模型精度：

1. 全局排序 > 逐层均匀：不同层的注意力头冗余程度不同。通常底层（靠近输入）的头冗余度更高，因为底层更多在做语法层面的模式匹配，多头的互补性较弱。应允许底层剪掉更多头，高层保留更多头。

2. 每次剪枝后微调：一次性剪掉所有冗余头会导致精度急剧下降。渐进式剪枝（每次剪10-20%的头，然后微调100-500步）的最终精度远高于一次性剪枝。

3. 微调时冻结已剪枝的结构：对剩余头使用较小的学习率（原学习率的1/10），避免剧烈改变已学到的表示。

四、头剪枝的失效场景

1. 跨语言模型：在多语言BERT中，不同头负责不同语言的现象已被多项工作证实。剪枝可能导致某些低资源语言的性能崩溃，而英文评测指标完全看不出来。

2. 长序列任务：对于长文档摘要或长文本分类，某些头专门负责长距离依赖。这些头的注意力分布可能看起来"稀疏"（对大多数位置权重近似0）但恰好在关键远程位置上有高权重。基于熵的评估会错误地将它们标记为低效头。

3. 对抗鲁棒性：有证据表明冗余头在对抗攻击和分布外泛化中起到"缓冲"作用——剪掉它们后模型的对抗鲁棒性可能下降，即使干净样本上的精度不变。

五、总结

注意力头的冗余检测需要多维度交叉验证：

1. 注意力熵检测"不工作"的头（均匀分布），余弦相似度检测"重复工作"的头。
2. 逐层均匀剪枝是次优策略，应使用全局重要性排序。
3. 渐进式剪枝 + 低学习率微调的恢复效果远优于一次性剪枝。
4. 多语言和长序列场景下，冗余头的定义需要重新审视。