本文提出灯塔注意力（Lighthouse Attention），一种训练专用的对称分层选择注意力算法，通过降低计算复杂度实现因果Transformer在超长序列上的高效训练。该方法包含自适应压缩/解压缩、对称池化保持因果性、两阶段训练（预训练+恢复）等创新，在小型LLM预训练实验中达到更快的总训练时间和更低的最终损失。
核心要点
- 提出次二次复杂度的分层预处理/后处理步骤，实现序列的自适应压缩与解压缩
- 对称压缩策略同时池化查询、键和值，保持从左到右的因果性，大幅提升并行性
- 两阶段训练：大部分时间使用灯塔注意力预训练，最后通过短训练恢复全注意力模型
- 在小型LLM预训练实验中，与全注意力训练相比，总训练时间更快且恢复阶段后最终损失更低
- 算法为训练专用，可轻松移除，且分层选择无梯度，避免复杂反向传播
正文
原始摘要
在极端序列长度上训练因果Transformer受到缩放点积注意力（SDPA）二次时间和内存的瓶颈限制。本文提出灯塔注意力，一种训练专用的对称选择分层注意力算法，它包裹普通SDPA，并可在训练结束时轻松移除。我们的分层选择也是无梯度的，这使我们免于处理复杂且可能低效的反向传播核。我们的贡献有三方面：(i) 一个次二次复杂度的分层预处理和后处理步骤，对序列进行自适应压缩和解压缩。(ii) 一种对称压缩策略，同时池化查询、键和值，同时保持从左到右的因果性，大大提高了并行性。(iii) 一种两阶段训练方法：大部分时间使用灯塔注意力进行预训练，最后通过短训练恢复全注意力模型。我们进行了初步的小规模LLM预训练实验，在所有其他设置匹配的情况下，与全注意力训练相比，我们的方法实现了更快的总训练时间和恢复阶段后更低的最终损失。完整代码见：https://github.com/ighoshsubho/lighthouse-attention
关联概念
- 缩放点积注意力 (SDPA)
- 因果Transformer
- 分层注意力
- 次二次复杂度
- 两阶段训练
可操作项
可尝试在小型LLM预训练任务中实现灯塔注意力，使用提供的代码库（https://github.com/ighoshsubho/lighthouse-attention）进行实验，对比全注意力训练的训练时间和最终损失。
原文: Long Context Pre-Training with Lighthouse Attention
自动加工于 2026-05-19 05:53