[论文] 控制对注意力 logits 的更改

Source: arXiv - 2511.21377v1

概览

训练大规模 transformer 模型可能出奇地脆弱：注意力机制中的查询（query）和键（key）矩阵常常在幅度上爆炸，导致训练不稳定和性能下降。Anson 与 Aitchison 的这篇论文指出，根本原因是 注意力 logits（softmax 之前的原始分数）出现了失控的漂移。通过对查询和键权重使用一种简单的、依赖参数的学习率方案来限制这些 logits 的变化幅度，作者在不使用笨重的 QK norm 技巧的情况下恢复了稳定性，即使在无法使用该技巧的情形下亦是如此。

关键贡献

• 识别 logits 漂移 为 transformer 训练不稳定的主要来源，尤其是针对 query/key 权重。
• 提出轻量级干预：对每个参数的学习率进行缩放，直接限制 logits 更新的幅度。
• 展示与 Multi‑Latent Attention (MLA) 的兼容性，该记忆高效的注意力变体无法使用完整的 query/key 归一化。
• 展示实证提升：该方法允许使用更高的基础学习率，在 MLA 中优于现有的稳定器，并在标准多头注意力（MHA）中匹配 QK norm 的性能。
• 提供实用配方，只需几行代码且无需额外的前向传播计算。

方法论

1. 问题定义 – 作者从注意力 logits (L = QK^\top / \sqrt{d}) 在训练步骤之间可能剧烈变化的观察出发，指出这会导致 softmax 分布过于尖锐或平坦。
2. 参数依赖的学习率 – 与对所有权重使用统一学习率 (\eta) 不同，他们为查询和键矩阵分配一个 缩放 学习率 (\eta_{Q,K} = \alpha \cdot \eta)，其中 (\alpha) 为一个小常数（例如 0.1）。这直接限制了 logits 在单次更新中的移动幅度。
3. 实现细节 – 缩放在优化器层面实现（例如在 Adam 中通过自定义参数组）。不需要额外的前向或反向传播，且该方法可与任何支持每参数学习率的优化器配合使用。
4. 实验设置 – 他们在两个方面进行评估：
   • (a) 标准多头注意力（MHA）在语言建模基准上的表现，
   • (b) Multi‑Latent Attention（MLA），该方法在推理时避免显式生成完整的 query/key 张量。
     基线包括普通训练、QK norm 以及其他近期的稳定器。

结果与发现

• 参数化学习率方案 保持注意力 logits 的方差低，防止 softmax 饱和。
• 在 QK norm 无法使用的 MLA 中，新方法 优于所有先前的稳定器，并加速收敛。
• 在两种设置下，该方法 不降低最终模型质量；它仅使优化器能够安全地使用更大的学习率。

实际意义

• 更快的训练周期 – 开发者可以将学习率提升 2–3 倍而不担心发散，从而缩短大规模 transformer 预训练的实际时间。
• 内存高效的注意力 – 对依赖 MLA 或其他低内存注意力技巧（如流式或端侧推理）的模型而言，该方法提供了 QK norm 无法实现的稳定性修复。
• 即插即用 – 由于该技术仅是学习率的微调，可在现有代码库（PyTorch、TensorFlow、JAX）中几乎不做改动即可加入。
• 更好的超参数鲁棒性 – 该方法降低了对学习率调度细致调参的需求，这在训练成本高昂的生产流水线中尤为重要。

局限性与未来工作

• 该方法 依赖手动选择的缩放因子 (\alpha)；尽管作者报告了相对稳健的默认值，但最佳取值仍可能随任务或模型规模而异。
• 它 未解决其他不稳定来源，例如前馈层的梯度爆炸或 layer‑norm 的缩放问题。
• 论文主要聚焦于 语言建模和视觉分类；将分析扩展到多模态或强化学习 transformer 仍是开放问题。
• 未来工作可以探索 自适应方案，根据观测到的 logits 漂移自动调整 (\alpha)，或将此方法与其他归一化技巧结合，以获得更强的鲁棒性。

作者

• Ben Anson
• Laurence Aitchison

论文信息

• arXiv ID: 2511.21377v1
• 分类: cs.LG
• 发布日期: 2025 年 11 月 26 日
• PDF: Download PDF