Command Palette
Search for a command to run...
ALiBi Transformer中塌陷注意力头的修复
ALiBi Transformer中塌陷注意力头的修复
Palmer Schallon
具有线性偏差的注意力(ALiBi)
摘要
我们在BLOOM系列的Transformer语言模型中发现了一种系统性的注意力塌陷病理现象,其中ALiBi位置编码导致31%至44%的注意力头几乎完全关注序列起始标记。这种塌陷在四个模型规模(5.6亿至71亿参数)下遵循可预测的模式,集中在ALiBi斜率调度施加最陡峭距离惩罚的头部索引处。我们引入了手术式重新初始化方法:对Q/K/V进行针对性重新初始化,将输出投影置零,并对所有非手术参数进行梯度掩码冻结。将该技术应用于单消费级GPU上的BLOOM-1b7模型,经过两轮迭代即可恢复98.7%的有效头部容量(从242个恢复到384个中的379个)。与C4训练数据的对照比较证实,驱动恢复的是重新初始化而非语料内容,并揭示了两种不同的术后现象:改善模型的早期全局功能重分布,以及在嘈杂训练信号下累积的晚期局部退化。一项扩展实验在重新初始化大部分健康头部的同时重新初始化塌陷头部,产生了一个在训练困惑度上短暂优于原始BLOOM-1b7 25%的模型(12.70对比16.99),这表明预训练的注意力配置是次优的局部极小值。代码、检查点和诊断工具已作为开源软件发布。
一句话总结
为修复基于 ALiBi 的 BLOOM 转换器中出现的系统性注意力坍塌问题,本研究引入了手术式重初始化(surgical reinitialization),即针对 Q/K/V 进行重初始化,将输出投影置零,并通过梯度掩码冻结非手术参数。该方法在单张消费级 GPU 上运行的 BLOOM-1b7 模型中成功恢复了 98.7% 的有效头容量(384 个头中的 242 个恢复至 379 个),将训练困惑度从 16.99 降至 12.70。通过受控的 C4 数据集对比实验表明,该技术本身即可驱动恢复过程,同时暗示预训练配置可能处于次优局部极小值。
核心贡献
- 分析揭示了 BLOOM 系列转换器模型中存在的系统性注意力坍塌病理现象。ALiBi 位置编码迫使 31% 至 44% 的注意力头不成比例地关注序列起始 token。该坍塌现象在 5.6 亿至 71 亿参数的模型规模中呈现可预测的分布规律,主要集中在 ALiBi 斜率计划施加最陡峭距离惩罚的头索引位置。
- 提出了一种手术式重初始化流程,选择性更新查询、键和值矩阵,同时对所有非手术参数应用输出投影置零与梯度掩码冻结。这种定向更新策略无需修改模型架构即可唤醒休眠的注意力容量。
- 在单张消费级 GPU 上对 BLOOM-1b7 模型应用该技术,仅需两次训练迭代即可恢复 98.7% 的有效注意力头。扩展实验表明,相较于原始模型,训练困惑度降低了 25%,表明标准预训练注意力配置通常位于次优局部极小值中。
引言
注意力头剪枝已成为转换器语言模型的标准压缩技术,通过移除过度关注序列起始 token 的组件来实现更快的推理速度。既往研究将这些坍塌的头视为功能冗余,基于其移除可安全降低计算开销且不损害性能的假设。本研究对此前提提出挑战,通过论证 BOS-sink 坍塌并非闲置容量的症状,而是由 ALiBi 位置编码引发的系统性伪影。研究并未直接丢弃这些组件,而是证明坍塌的头仍处于休眠状态且可被成功重初始化。该方法能够恢复功能性注意力容量,在可量化的层面提升模型表现,同时避免不可逆的结构损伤。
数据集
- 数据集构成与来源:提供的节选未明确说明数据集的构成或来源。
- 各子集的关键细节:缺乏关于子集规模、来源及过滤规则的信息。
- 数据使用与划分:未描述训练集划分、混合比例或处理工作流。
- 处理与元数据:未提及裁剪策略、元数据构建步骤或其他预处理细节。 所提供的文本仅包含论文标题与作者署名。
方法
本研究采用定向手术式重初始化方法修复 BLOOM 系列转换器模型中的注意力坍塌问题,旨在解决显著比例的注意力头不成比例地关注序列起始(BOS)token 的系统性病理现象。该坍塌现象在四种模型规模中均有观察,与 ALiBi 位置编码方案密切相关。该方案对较高头索引施加更陡峭的距离惩罚,导致坍塌头在特定层和头位置呈现可预测的集中分布。
整体框架始于诊断阶段,用于识别坍塌的注意力头。针对每个头计算两项指标:BOS 质量(BOS mass),用于衡量所有查询位置分配给 BOS token 的平均注意力权重;以及熵(entropy),用于量化跨位置注意力分布的广度。这些指标来源于序列长度 T 内第 ℓ 层和第 h 头的注意力矩阵 Aℓ,h。诊断提示词因其语义密度与多样的依赖结构被选中,确保观察到的头行为双峰分布反映的是内在权重特性,而非输入特定效应。头的分类基于阈值设定:高 BOS 质量(高于 0.50)与低熵(低于 0.50),并通过 DEAD 分数阈值 0.95 进行额外验证。
诊断完成后,应用手术式重初始化流程以恢复功能容量。针对每个已识别的坍塌头执行四项操作:使用 Xavier 正态初始化重新初始化 Q、K 和 V 投影权重,以跳出 BOS-sink 局部极小值;将密集输出投影置零,确保该头在初始恢复阶段对残差流不产生任何贡献,从而防止系统失稳;通过梯度掩码冻结所有非手术参数,在反向传播期间将梯度置零;仅在训练语料上训练手术参数。该方法基于实证证据,即仅依赖梯度的微调无法恢复坍塌头,因为 BOS-sink 状态构成一个尖锐的局部极小值,梯度下降法无法在不重置的情况下逃脱。
该技术以两轮迭代方式应用于 BLOOM-1b7。第一轮迭代针对第 5 至 22 层中 H9–H15 频带的 108 个头,同时冻结五个已健康的频带内头以维持稳定性。该轮迭代恢复了大量丢失的容量,但引入了罕见的医源性效应,导致第 23 层中两个原本健康的头发生坍塌。第二轮迭代处理 H9–H15 频带外剩余的 39 个坍塌头,并以第一轮迭代的最佳检查点作为起始状态。这种两轮策略反映了注意力生态机制:优先稳定主频带可使模型注意力拓扑在处理异常值之前完成重组,从而实现 98.7% 的有效头容量恢复(384 个头中的 242 个恢复至 379 个)。
扩展实验表明,将相同的手术技术应用于主要健康的头时,生成的模型在训练困惑度上可短暂超越原始模型 25%。这暗示预训练注意力配置可能代表次优局部极小值,而定向重初始化能够触及标准训练动态无法到达的更优配置。恢复过程由手术干预本身驱动,而非训练语料,这一点已通过与 C4 数据的受控对比得到证实。该对比揭示了两种截然不同的术后现象:提升模型性能的早期全局功能重组,以及在噪声训练信号下累积的晚期局部退化。
实验
评估采用针对特定注意力头的定向重初始化,随后在精选或通用语料上进行短暂微调。实验验证了头的恢复情况、全局拓扑变化及生成质量。定性分析表明,手术干预成功唤醒了坍塌头,并通过共享的残差流触发了注意力网络中广泛的功能性重组。此外,实验证明预训练注意力模式经常占据次优局部极小值,因为对坍塌头和名义上健康的头进行重初始化均能发现更优配置。最终,研究结果确立转换器注意力作为一种相互关联的生态机制运行,定向修改既能实现治疗性修复,也能完成结构重优化,尽管持续的性能提升需要足够大规模的语料以防止过拟合。
该表格比较了不同训练条件对 BLOOM 模型注意力头重分配的影响。结果显示,与 C4 语料相比,精选语料导致更广泛的全局注意力头重分配,表现为漂移头数量更多且 BOS 质量平均变化更大。C4 语料,尤其是在延长训练后,在特定头上表现出更高的最坏情况漂移,表明存在局部退化。这些发现支持训练数据会影响手术干预后注意力重分配质量与范围的结论。精选语料引发的全局注意力头重分配程度高于 C4 语料,以漂移头数量更多和 BOS 质量平均变化更大为衡量标准。在 C4 语料上进行延长训练会导致特定头的最坏情况漂移增加,暗示精选条件下未观察到的局部退化。在匹配的训练轮次下,精选语料实现了更有效的重分配,相较于 C4 基线产生了更低的训练困惑度。
作者比较了 BLOOM-1b7 模型在不同手术干预下的训练困惑度,表明仅针对频带的手术与包含健康头的扩展手术均能降低相较于原始模型的训练困惑度。在最优子轮次点评估时,扩展手术取得最佳性能,表明重初始化非坍塌头可改善注意力模式。这些提升是暂时的,因为过拟合开始后困惑度会上升。包含健康头的扩展手术获得的训练困惑度低于仅针对频带的手术。最佳训练性能出现在过拟合开始前的中间子轮次点。重初始化非坍塌头能够改善注意力模式,暗示预训练模型原本处于局部极小值中。
该表格比较了不同训练条件对注意力头行为的影响。结果显示,在相同训练轮次下,精选语料导致的漂移头数量更多且 BOS 质量平均变化更高,而 C4 语料在较晚轮次则表现出更多的漂移头与更高的平均变化。这表明精选语料在训练初期驱动了更显著的全局注意力模式重分配,而 C4 语料随时间推移导致更局部的漂移。精选语料在相同训练轮次下引发的注意力头漂移与 BOS 质量平均变化均高于 C4 语料。C4 语料在较晚轮次相较于较早轮次显示出头漂移和平均变化增加,表明存在随时间累积的局部退化。精选语料在训练早期实现更广泛的重分配,而 C4 语料随训练推进表现出更局部的漂移。
该表格展示了训练期间不同头列的冻结头漂移数据。结果显示 H15 表现出扩散型漂移,而 H9、H10 和 H12 保持相对稳定。漂移量通过第 3 轮至第 15 轮之间 BOS 质量的平均变化进行衡量,H15 的数值显著更高,表明注意力模式变化呈渐进式增加。这暗示某些头索引在时间推移中更容易受到重分配影响,尤其是在持续使用噪声数据进行训练时。H15 表现出扩散型漂移且平均增量值较高,表明注意力模式发生渐进式变化;H9、H10 和 H12 呈现稳定漂移,表明其对重分配具有抵抗力。各列之间漂移模式的差异突显了架构对级联效应的脆弱性。
该表格展示了 BLOOM-1b7 模型中选定头在手术重初始化前后的 BOS 质量,表明健康头与坍塌头的 BOS 质量均显著降低。结果证明,该手术技术不仅恢复了坍塌头,还改善了健康头的注意力模式,从而达成更优配置。这支持了预训练注意力拓扑处于局部极小值的结论,且重初始化能够解锁更优配置。手术重初始化降低了健康头与坍塌头的 BOS 质量,表明注意力模式得到改善。该技术实现了坍塌头的恢复与健康头的增强,暗示其作用不仅是修复,更是重优化。结果支持预训练注意力模式属于局部极小值的观点,且通过重初始化可访问更优配置。
实验评估了 BLOOM-1b7 模型在不同手术干预与训练语料下的注意力头重分配及训练动态。跨数据源的对比验证了精选语料在训练初期驱动更广泛的全局注意力重分配,而标准数据集在长期训练中往往导致局部退化。手术干预测试证实,同时对坍塌头与健康头进行重初始化可成功跳出局部极小值,相较于部分干预方法更有效地重优化注意力模式并降低困惑度。最后,头漂移分析揭示某些架构位置本质上更易受渐进式重分配影响,强调了整体重训练策略对模型稳定性的重要性。