基于 Lighthouse Attention 的长上下文预训练
计算与语言
2026-05-08 v1
摘要
在极长序列长度下训练因果 Transformer 受到缩放点积注意力（SDPA）的二次时间和内存瓶颈的限制。在本工作中，我们提出了 Lighthouse Attention，一种仅用于训练的、基于对称选择的分层注意力算法，它包裹在普通 SDPA 之外，并可在训练接近尾声时轻松移除。我们的分层选择也是无梯度的，这使我们免于处理复杂且可能低效的反向传播内核。我们的贡献包括三个方面： 一个亚二次分层的预处理和后处理步骤，对序列进行自适应压缩和解压缩。 一种对称压缩策略，同时池化查询、键和值，同时保持从左到右的因果性，这极大地提高了并行度。 一种两阶段训练方法，我们在大部分预训练时间里使用 Lighthouse Attention，并在最后通过短时训练恢复为完整的注意力模型。我们进行了初步的小规模 LLM 预训练实验，结果表明在所有其他设置匹配的情况下，与全注意力训练相比我们的方法具有有效性，在恢复阶段后实现了更快的总训练时间和更低的最终损失。完整代码可在以下地址获取：https://github.com/ighoshsubho/lighthouse-attention
引用
@article{arxiv.2605.06554,
title = {Long Context Pre-Training with Lighthouse Attention},
author = {Bowen Peng and Subho Ghosh and Jeffrey Quesnelle},
journal= {arXiv preprint arXiv:2605.06554},
year = {2026}
}
备注
18 pages, 4 figures, 4 tables