删掉多余的注意力层——最不像手术的手术
这一章讲三件事: 怎么单独测量「块里的注意力那一半」贡献了多少(钩子要挂在哪两个点、为什么必须连残差——把输入抄近道加回输出的那条线——一起量);删掉之后怎么让模型还能保存、还能重新加载(架构改了,救援代码也要跟着改);以及为什么基准成绩好看不代表模型没受伤。 本章是原书第 8 章的下半场:上一章不动架构,本章永久改变架构。
1. 这一章讲什么
「Transformer 不能没有注意力」近乎信仰。本章的立足点是 2024 年的一篇论文:《What Matters in Transformers? Not All Attention is Needed》——块内冗余分布不均,注意力层往往比 MLP 层更可弃1。于是第 04 章那套「钩子+余弦相似度」的瞄准镜,从「整块」推近到「块里的注意力模块」:量出哪些层的注意力几乎什么都没干,把它们的注意力阶段整个旁通,MLP 保留2。
它在全书链条里的位置: 它把第 04 章的测量学、第 08 章上半场的内存账、第 06 章的蒸馏善后接成一条线,是全书「测量→下刀→善后→诚实复核」流程最完整的一次示范。
2. 顶层全景
测量:钩子挂在 input_layernorm(进)与 post_attention_layernorm(出)
→ 对注意力阶段(含残差)算 1 − 余弦相似度 = 重要性分
选层:升序排,取最低的几层(3B 上是 21/22/23 层)
手术:delattr 删 self_attn 与 input_layernorm
→ type() 造子类、类级替换 forward(跳过注意力阶段)
善后:config 里登记删了谁 + 随模型发一个自定义加载类
复核:选择题基准几乎不动,但 LAMBADA 和长文生成暴露暗伤
→ 蒸馏 10 分钟能修
3. 核心原理
3.1 更细的探针:钩子挂哪、量什么
第 04 章的钩子挂在整块(DecoderLayer)上,量「注意力+MLP 合计改了多少」。要单独审注意力,挂载点就得挪到块内的两个归一化层上——它们的妙处在于恰好框住注意力阶段,而且所有主流 decoder 家族(Llama、Mistral、Qwen、Gemma)都标准化地暴露这两个名字3:
- input_layernorm:截获的是即将进入注意力计算的隐状态,记作 X_A;
- post_attention_layernorm:截获的是注意力输出加残差连接之后、即将进入 MLP 的张量,记作 Y_A4。
「残差连接」这里必须补课:它是把块的输入原样抄一条近道、在出口加上去——块学的其实是「输出与输入的差」,这让深层网络的训练稳定、梯度畅通。所以 Y_A 里已经装着 X_A:它的成分是「X_A + 注意力(归一化后的 X_A)」5。重要性分还是老公式:对 X_A 与 Y_A 逐 token 算余弦相似度、剔除 padding、取平均、用 1 减掉6。
技术细节一处:这次用 register_forward_pre_hook(前向前钩子)而不是普通钩子——它只收模块的输入参数,在这两个点位上正好是要的那一口干净数据7。
为什么「含残差」是生死线? 这是论文专门的消融结论:只量注意力模块自身的输出、不加残差,得到的估分会误导——删几层就明显掉成绩;把残差算进来,度量才可靠。本书的钩子锚点正是照这个要求选的8。
3.2 冗余地图:低洼带在 60-80% 深度处
校准数据这次不用领域数据(目标不是特化是通用提速),改用 Cosmopedia 六个子集按权重混装 400 条——配比是照着评测基准的能力面倒推的:故事和网页文本喂给常识题,指南文对物理题,教材内容对科学题9。
Llama-3.2-3B 的分数形状(28 层):第 0、1 层远高于其余;2-17 层中等;真正的低洼带在 18-25 层;26、27 略回升——尾部在输出前整合表示,职责还在10。换更大的 Llama-3.1-8B 用同一份代码重跑,结构完全复现:头尾活跃、低洼居中,最低分都落在全网 60%-80% 深度的地方;变化的只是比例——阈值(划的一条分数线)以下,3B 有 18% 的层(5/28),8B 有 34%(11/32)。冗余空间随模型长大而变宽,这正是大模型更扛注意力剪枝的原因11。
3.3 主走查:删层手术,从下刀到重新加载
走查对象:Llama-3.2-3B,删第 21、22、23 层的注意力。 只删三层的理由本章自己给了:3B 的冗余空间窄,从「可控退化」到「生成崩塌」的过渡来得很快12。
第一步 下刀
delattr(层, "self_attn") # 删注意力模块
delattr(层, "input_layernorm") # 连它前面的归一化一起删
——只删 self_attn 不删 layernorm 毫无意义:没有输入的归一化
是一步空转的计算,白烧周期[^13]。
第二步 改执行路径
原来的块 forward 会依次调 input_layernorm→self_attn→……
部件删了,老 forward 下一次执行就会当场报错(找不到部件)。
解法:写一个跳过注意力阶段的新 forward——
残差保留,张量直接走 post_attention_layernorm→MLP→加回残差;
再用 type() 造一个子类、在【类级别】替换 forward
——实例级替换会被 transformers 的内部机制在生成时覆盖回去[^14]。
第三步 教模型「自报家门」
直接保存再加载会出大事:保存只存权重和配置,不存你的新 forward;
重新加载时按原类重建——被删层的权重在文件里根本不存在,
框架就随机初始化补齐,模型「加载成功」,生成出来的是垃圾[^15]。
解法三件套:
① config 里登记 model.config.dropped_attn_layers = [21,22,23]
② config.auto_map 指向一个自定义加载类(随模型一起发文件)
③ 加载时 trust_remote_code=True
第四步 验证闭环
自定义类的加载顺序是精髓:先按原结构搭好全部 28 块的骨架 →
读 config、对登记过的层执行同样的删除与 forward 替换 →
然后才载入权重文件。
「checkpoint 里的每个键都在补丁骨架里找到归宿,
补丁骨架里的每个部件也都在 checkpoint 里有权重」——零缺失键[^16]。
3.4 代价核算:基准说没事,生成说有事
显存(最直接的收益): 每层注意力约 2500 万参数,删三层共 7500 万(FP16 下 150 MB);动态显存增量从基线 109.57 MB 降到 96.06(-12.3%),删 6 层 -28.2%,删 15 层 -46.5%——每个被删的注意力层,意味着每个新 token 少记一份 KV 账13。
基准(看似没事): 删 3 层, 五项基准里四项波动在 ±1.4% 以内;删 6 层依然几乎不动。但 LAMBADA 一枝独敏感:3 层 -9.2%,6 层 -19.1%,15 层 -93.3% 全场崩盘14。
为什么选择题基准这么钝?它们的机制是「给上下文、比各选项的概率」——考的是局部判别力;而 LAMBADA 要生成出段落末词,答案被设计成不看全文绝无可能猜到——考的是长程整合。开放词表、五万选一、没有干扰项可排除。注意力层被删后,第一个受伤的就是这种能力15。
生成(暗伤现形): 短提示下,6 层版开始复读(「2300 万游客…1680 万游客…」);长文生成下,连 3 层版都会冒出细密的事实不一致;15 层版短提示即崩16。作者点破:论文里那套「免训练即可删」的结论,建立在只考一次静态前向的基准上;自回归生成暴露的退化,静态基准设计不出来17。
善后: 跑一小段蒸馏就能修——A100 上 10 分钟、一万条 Cosmopedia,复读循环消失,巴黎测试恢复正常。定性:退化不是不可逆的能力损失,是架构与权重之间的暂时错位,恢复校正它18。
4. 作者的判断与证据
给了证据的: 冗余分数、显存表、基准表全部可复现;论文侧数字更硬——四种删法同台比「性价比」:SDR(性能下降百分点 ÷ 加速百分点)越低越好,删注意力的 SDR 恒定比删整块、删 MLP 低 10-20 倍;Llama-2-13B 删 8 个注意力层,平均分 68.1 对基线 68.2,几乎纹丝不动;Llama-2-70B 删 40%(32/80 层),SDR=0.00(测不出退化)+ 1.35 倍加速,删到一半也只 0.05 + 1.48 倍19。本书 3B 上的对应数字诚实得多:删 11% 无可见加速(25.0 对 25.4 token/秒),删 54% 换约 6% 吞吐但生成当场崩——规模就是本钱20。
作者的两个方法论判断:
-
基准要配任务。 论文的基准全是选择题或二分类(二选一)任务,作者特意补了 LAMBADA 和自由生成,把论文看不见的暗伤照了出来——结论不是推翻技巧,而是「免训练删除」只对选择/排序类任务成立,要生成就必须配恢复21;
-
结论挂在任务上,而任务定义度量。 同是马里兰团队,后续一篇论文把同样的 探针搬到了另一类任务上——检索,也就是按内容相似度从资料库里找答案的那类任务——得出完全反转的结论22。
那篇论文做的是稠密(全量数值参与查库)检索:那个场景里注意力层变得关键,反而是 MLP 高冗余、可大量删。差异不来自谁测错了,来自度量各自对准了任务的本质要求。
与上一章的接缝(论文给出的叠加许可): 论文验证了删注意力与 4-bit 量化(AWQ)完全兼容——加了量化,各项基准的差异不足 1%。也就是说第 08 章的 KV cache 量化和本章的删注意力层不是二选一,是可以叠加在同一模型上的两笔独立收益23。
5. 边界与局限
- 3B 的结论外推要打折:小模型冗余窄,能删的比例小;「几乎零损失」的漂亮数字全部来自 13B/70B。
- 本章的速度收益在小模型上不成立(25.0 vs 25.4);它的实际卖点在 3B 上是动态显存,不是速度。
- 保存/重载方案绑定 Hugging Face 的加载机制(auto_map + trust_remote_code);换推理框架要另想办法,而且第三方引擎未必执行你的自定义类。
- 「恢复 10 分钟就好」基于 3B + 特定数据量;更大模型、更长恢复的成本没有给出。
- 「稠密检索」——把文本压成数值去查库——场景的反转提醒:本章的全部「删注意力」结论,只在生成/判别场景成立;做 embedding 模型前先看下一节那张反转牌。
6. 可带走的
- 块内冗余不均,注意力先挨刀:同一把余弦软尺,挂载点从整块挪到两个归一化层,就能单独审注意力;
- 残差必须算进度量:不含残差的重要性分会误导——这是论文消融实验(逐个拆掉部件看掉多少分)的结论;
- 低洼带在 60-80% 深度处,且模型越大低洼越宽;小模型能删的注意力层很少;
- 删部件必须连着它的配套删(注意力与其前置归一化),否则留一步空转;
- 改了架构就要管两件事:类级替换 forward 防覆盖;config 登记+自定义加载类+trust_remote_code 保证能重新加载;
- 选择题基准测不出自回归暗伤:LAMBADA 与自由生成长文才是照妖镜;
- 暗伤可修:短期蒸馏即可校正架构-权重错位;
- SDR 是挑手术的好尺子:性能下降÷加速,删注意力的性价比是删整块的 10-20 倍(大模型上);
- 把机器改成查资料的取数模型(embedding)之前,结论要反转:那个场景里该删的是 MLP,不是注意力。
7. 原文地图
| 主题 | 原书章 | 原文位置 |
|---|---|---|
| 从整块到注意力层的探针 | 8.3 Measuring attention redundancy | text/33-ch08-03-8-3-measuring-attention-redundancy.txt:7(搜「attention layer of the transformer」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:13(搜「keep the MLP layer」) |
| 论文与冗余不均 | 8.3 Measuring attention redundancy | text/33-ch08-03-8-3-measuring-attention-redundancy.txt:19(搜「Not All Attention is Needed」) |
| 两个锚点与标准化 | 8.3 Measuring attention redundancy | text/33-ch08-03-8-3-measuring-attention-redundancy.txt:55(搜「input_layernorm」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:109(搜「decoder-only models」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:280(搜「register_forward_pre_hook」) |
| X_A/Y_A 与残差补课 | 8.3 Measuring attention redundancy | text/33-ch08-03-8-3-measuring-attention-redundancy.txt:266(搜「Y_A」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:275(搜「residual connection」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:293(搜「stabilizes training」) |
| 校准集配比 | 8.3 Measuring attention redundancy | text/33-ch08-03-8-3-measuring-attention-redundancy.txt:147(搜「0.300」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:182(搜「maps naturally」) |
| 冗余地图与规模化 | 8.3 Measuring attention redundancy | text/33-ch08-03-8-3-measuring-attention-redundancy.txt:396(搜「18 to 25」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:402(搜「consolidating representations」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:444(搜「0.009053」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:509(搜「60% and 80%」) · text/33-ch08-03-8-3-measuring-attention-redundancy.txt:515(搜「34%」) |
| 选层与手术 | 8.4 Removing attention modules | text/34-ch08-04-8-4-removing-attention-modules.txt:29(搜「narrower than in larger architectures」) · text/34-ch08-04-8-4-removing-attention-modules.txt:66(搜「burns cycles」) · text/34-ch08-04-8-4-removing-attention-modules.txt:113(搜「class level, not the instance level」) · text/34-ch08-04-8-4-removing-attention-modules.txt:123(搜「AttributeError」) |
| 保存/重载难题与解法 | 8.4 Removing attention modules | text/34-ch08-04-8-4-removing-attention-modules.txt:228(搜「missing keys」) · text/34-ch08-04-8-4-removing-attention-modules.txt:254(搜「dropped_attn_layers」) · text/34-ch08-04-8-4-removing-attention-modules.txt:266(搜「PrunedLlamaForCausalLM」) · text/34-ch08-04-8-4-removing-attention-modules.txt:402(搜「patched skeleton」) |
| 显存与参数账 | 8.4 Removing attention modules | text/34-ch08-04-8-4-removing-attention-modules.txt:414(搜「150 MB」) · text/34-ch08-04-8-4-removing-attention-modules.txt:485(搜「109.57」) · text/34-ch08-04-8-4-removing-attention-modules.txt:529(搜「fewer KV cache entries」) |
| 基准分化与 Lambada | 8.4 Removing attention modules | text/34-ch08-04-8-4-removing-attention-modules.txt:439(搜「-9.2%」) · text/34-ch08-04-8-4-removing-attention-modules.txt:445(搜「local discrimination」) · text/34-ch08-04-8-4-removing-attention-modules.txt:451(搜「open-vocabulary generation task」) |
| 生成暗伤 | 8.4 Removing attention modules | text/34-ch08-04-8-4-removing-attention-modules.txt:553(搜「23 million tourists」) · text/34-ch08-04-8-4-removing-attention-modules.txt:603(搜「factual inconsistencies」) · text/34-ch08-04-8-4-removing-attention-modules.txt:603(搜「static benchmarks are not designed to detect」) |
| 蒸馏善后 | 8.4 Removing attention modules | text/34-ch08-04-8-4-removing-attention-modules.txt:622(搜「distill_model」) · text/34-ch08-04-8-4-removing-attention-modules.txt:640(搜「ten minutes」) · text/34-ch08-04-8-4-removing-attention-modules.txt:662(搜「temporary misalignment」) |
| 论文:SDR 与规模化 | 8.5 From paper to practice | text/35-ch08-05-8-5-from-paper-to-practice.txt:25(搜「misleading estimate」) · text/35-ch08-05-8-5-from-paper-to-practice.txt:67(搜「ten to twenty times lower」) · text/35-ch08-05-8-5-from-paper-to-practice.txt:80(搜「1.35 speedup」) · text/35-ch08-05-8-5-from-paper-to-practice.txt:86(搜「25.0 vs 25.4」) |
| 任务适配与叠加 | 8.5 From paper to practice | text/35-ch08-05-8-5-from-paper-to-practice.txt:104(搜「not used in the paper」) · text/35-ch08-05-8-5-from-paper-to-practice.txt:116(搜「perfectly viable」) · text/35-ch08-05-8-5-from-paper-to-practice.txt:128(搜「less than 1%」) |
| 稠密检索器反转 | 8.5 From paper to practice | text/35-ch08-05-8-5-from-paper-to-practice.txt:134(搜「completely reversed」) · text/35-ch08-05-8-5-from-paper-to-practice.txt:140(搜「metrics used」) |
| 动手练习(无新机制):反向删层验证度量、与上一章技巧叠加 | 8.6 Hands-on lab | text/36-ch08-06-8-6-hands-on-lab.txt:13(搜「inverse pruning」) |