一句话定义:注意力头剪枝(Attention Head Pruning)是把 Transformer 模型里的一部分注意力头(attention head)整个删掉,连它的参数和计算一起,再看模型还能不能正常运行。这是一种常见的模型压缩方法。
一个头在管什么
Transformer 的每一层里都有一组并行的注意力头,每个头都有自己的 Q/K/V 投影,各自学一种“该看哪里”的模式:有的盯着紧挨着的前一个词,有的把注意力固定在句首,有的擅长把主语和谓语联系起来,有的专门帮助代词找到它指代的对象。它们既各管一摊,又在相当程度上互相冗余。这正是“可以删”的前提。
打个比方:项目组里十几个信息员,各盯一条线,有人盯预算,有人盯进度,有人盯甲方情绪。砍掉一两个,活还是有人接。剪枝的做法是给每个信息员打分,看关掉他之后项目损失多大,再按分数从低到高依次砍掉;砍完让全组重新磨合几天,也就是微调(fine-tuning)。
打分方式常见的有:关掉后损失涨多少、这个头输出的幅值大小、梯度大小。名字不必记,本质都是“先试着拿掉,看影响大不大”。
和相邻概念的区别
| 方法 | 删掉的是什么 | 能否真省算力 | 典型风险 |
|---|---|---|---|
| 注意力头剪枝 | 整个头,即一条完整计算路径 | 能 | 伤特定能力 |
| 非结构化剪枝(unstructured pruning) | 单个权重,也就是把权重置零 | 在普通硬件上常常不能 | 稀疏了但跑不快 |
| 量化(quantization) | 数值精度 | 能 | 精度下降、数值溢出 |
| 蒸馏(knowledge distillation) | 用大模型教一个小模型 | 能 | 基本要重训 |
其中最关键的区分是:剪头属于结构化剪枝(structured pruning),删除的是完整的计算路径,所以显存和时延能真实下降;把权重逐个置零属于非结构化剪枝,稀疏率看着漂亮,落到普通 GPU 上未必更快。
最容易踩的坑:平均分陷阱
很多头确实可以删掉而平均指标几乎不动,但“平均”会掩盖局部。某个头可能专门负责中文量词搭配、日文助词、代码缩进层级、长文档里隔了几百个 token 的指代,这类样本在通用评测里占比很小,删掉它总分看不出变化,实际使用却明显退步。更隐蔽的是头之间有补偿:删掉 A 后 B 顶上,平均分恢复了,但冗余余量也消耗掉了,一遇到分布外(out-of-distribution)输入就露出问题。
所以实操上别只看一个总分:分层、分语言、分任务各测一遍,并留一组“探针”任务:长文本、小语种、结构化抽取、代码;剪枝比例保守一些,剪完一定微调。
对从业者的意义
在端侧部署、推理成本敏感的场景里,剪头是最直接的结构瘦身手段之一;对普通用户来说,这也解释了为什么同一个模型在更小的设备或更便宜的档位上会“变笨”,而且往往先笨在小语种、长文档和专业术语上。剪多少合适没有通用答案,具体工具的支持情况以官方页面为准。
