💬 NLP

Long Context Pre-Training with Lighthouse Attention

本文介绍了灯塔注意力，这是一种仅用于训练、无需梯度的分层选择算法，它封装了标准的缩放点积注意力，以支持在极端序列长度下对因果Transformer进行高效的次二次方预训练，并可通过一个简短的恢复阶段无缝移除，从而得到一个具有更快训练速度和更低最终损失的全注意力模型。

原作者： Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

发布于 2026-05-08

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你试图一次性阅读一个庞大的图书馆（即一段非常长的文本）来理解一个故事。在人工智能领域，这被称为“长上下文训练”。

问题在于，标准人工智能模型（Transformer）试图将每一个单词与所有其他单词进行比对，以寻找关联。如果你有 10 万个单词，模型就必须进行 100 亿次比对。这就像试图通过在图书馆里同时向每个单词大喊每一个单词来寻找特定句子一样。这需要耗费漫长时间，产生巨额电费，并且会导致计算机内存耗尽。

本文介绍了一种名为**灯塔注意力（Lighthouse Attention）**的新方法来解决这一问题。以下是其工作原理，使用了简单的类比：

1. 问题：“全知之眼”过于沉重

标准人工智能的注意力机制就像一名保安，坚持要同时盯着体育场里的每一个人，以查看谁在和谁交谈。随着体育场变得越来越大，这名保安会不堪重负。

2. 解决方案：“灯塔”策略

灯塔注意力不像全分辨率地审视一切，而是像灯塔扫描黑暗的海洋。它并非忽略海洋，只是分层扫描，以快速找到最重要的部分。

以下是论文描述的三步流程：

步骤 A：“金字塔”（对人群进行总结）

想象你有一大群人（即文本）。与其观察每一张脸，不如将他们分成小组（如家庭或团队）。

技巧：论文在此处做了一个独特的处理。大多数其他方法仅总结“被谈论的人”（即键和值），而将“正在说话的人”（即查询）保留在高细节中。
灯塔的做法：它平等地总结所有人。它构建了一个总结金字塔：
- 第 0 层：每一个单词。
- 第 1 层：每 4 个单词被总结为一个。
- 第 2 层：每 16 个单词被总结为一个。
- 依此类推。
  这创建了一个从“超级详细”到“宏观概览”的多层级文本地图。

步骤 B：“聚光灯”（挑选最佳部分）

现在，模型需要决定金字塔中的哪些部分重要到值得全细节阅读。

选择：它使用一条简单、免费且无需额外学习的规则来给每个小组打分。它会问：“哪些小组拥有最多的‘能量’或重要性？”
裁剪：它从金字塔的所有层级中挑选出最重要的前 $K$ 个小组。
结果：模型不再需要阅读 10 万个单词，而只需阅读一个微小、密集的、最重要的“块”列表（可能约为 5,000 个单词）。

步骤 C：“闪光”（阅读选定的部分）

一旦模型挑选出前 5,000 个单词，它就会将它们输入标准的、超高速的"FlashAttention"引擎。由于列表现在变短了，这一步极其迅速，并且可以轻松适应计算机内存。

3.“魔法”恢复（两阶段训练）

这是本文最关键的部分。作者曾担心：“如果我们训练人工智能只查看摘要，它以后是否会忘记如何阅读完整句子？”

为了解决这个问题，他们采用了一种两阶段训练配方：

阶段 1（灯塔阶段）：在大部分训练时间里，他们使用灯塔方法训练模型。模型学会变得高效并挑选正确的重点。
阶段 2（恢复阶段）：在最后的少量训练时间里，他们关闭“总结”和“挑选”部分。他们强制模型使用标准方法再次阅读完整文本。

结果：模型从高效训练中“苏醒”过来，并完美地记住了如何使用完整注意力。论文声称，经过这段短暂的恢复后，模型的表现与那些全程都在完整文本上训练的模型一样好（甚至更好），但它到达这一状态的速度更快、成本更低。

为什么这很重要？

速度：论文表明，对于非常长的文本（如 10 万 + 单词），该方法比标准方法快 17 到 21 倍。
无“垃圾”代码：与其他需要构建定制、复杂的计算机芯片（内核）来处理“挑选”过程的方法不同，灯塔注意力在实际阅读时使用标准的、现成的计算机部件。它只是在移交数据之前重新排列了数据。
对称性：它平等地对待句子中的“问题”和“答案”部分，这使得数学计算更清晰、更稳定。

总结

灯塔注意力就像雇佣了一位聪明的研究助理。与其逐字阅读报告的 1,000 页，这位助理会快速扫描整份文件，高亮显示最重要的 50 个段落，然后仅深入详细地阅读这 50 个段落。

论文证明，如果你以这种方式训练人工智能，并在最后给它一个快速阅读整体的“复习课程”，它就会变成一个超级快速的阅读者，且不会损失任何智能。

技术摘要：基于灯塔注意力（Lighthouse Attention）的长上下文预训练

1. 问题陈述

在极端序列长度（128K 到 1M+ 个 token）下训练因果 Transformer，目前受限于缩放点积注意力（SDPA）的二次方时间和内存复杂度（ $\Theta(N^2)$ ）。尽管 FlashAttention 等方法优化了常数项，但并未改变其渐近复杂度。现有的稀疏注意力方法试图通过将密集注意力替换为选择机制（即查询仅关注键的子集）来缓解这一问题。然而，当这些方法应用于预训练时，存在两个主要设计缺陷：

不对称性：大多数方法（如 NSA、HISA）对键（Keys）和值（Values）进行池化，但保持查询（Queries）为全分辨率。这将层级结构视为压缩内存而非多尺度表示，限制了模型学习跨尺度连贯交互的能力。
架构纠缠：选择逻辑通常嵌入在自定义的稀疏注意力内核中。这阻碍了高度优化的密集注意力内核（如 FlashAttention）的复用，并需要自定义的反向传播，从而增加了训练和推理部署的复杂性。

此外，训练时稀疏方法的一个关键关注点是正确性：与仅推理阶段的稀疏性（假设存在一个胜任的密集骨干网络）不同，训练时稀疏方法必须证明，一旦移除稀疏性，所得权重能够作为胜任的密集注意力模型发挥作用。

2. 方法论：灯塔注意力（Lighthouse Attention）

灯塔注意力是一种仅用于训练、对称的、基于选择的分层注意力算法，旨在包裹普通的 SDPA 而无需修改注意力内核本身。它作为一个四阶段流水线运行：

A. 对称金字塔构建

与仅对上下文（K/V）进行池化的先前工作不同，灯塔注意力对查询、键和值（ $Q, K, V$ ）同时应用对称平均池化。

序列被组织为具有池化因子 $p$ 的 $L$ 层金字塔。
第 0 层是完整分辨率的序列。随后的每一层 $\ell$ 对第 $\ell-1$ 层的 $p$ 个连续条目进行汇总。
这在每一层生成了连贯的三元组 $(Q^{(\ell)}, K^{(\ell)}, V^{(\ell)})$ ，确保池化后的查询和键位于相同的表示空间中。

B. 无参数评分与选择

一个分层选择器对金字塔条目进行排名，以确定关注哪些 token。

评分：该方法使用无参数评分函数。对于每个条目，计算投影后的 $Q$ 和 $K$ 向量的 $\ell_2$ 范数。较粗的层级通过最大池化从较细的层级继承分数，确保一个粗粒度跨度保留其最强 token 的重要性。
选择：一个融合的分块双调 Top-K 内核在所有层级上联合选择前 $K$ 个条目。
不可微性：Top-K 选择步骤是离散的且不可微的。不使用直通估计器（straight-through estimator）或辅助损失。梯度仅通过 gather、FlashAttention 和 scatter-back 操作流向投影矩阵（ $W_Q, W_K, W_V$ ）。因此，投影学习的是生成“在被选中时有用”的值，而不是学习如何操纵评分器。

C. 密集子序列注意力

选中的条目被收集成一个连续的、因果一致的子序列，长度为 $S \ll N$ 。

收集（Gather）：选中的三元组经过拓扑排序，形成密集序列。
注意力：在此子序列上执行标准的FlashAttention（或原生 SDPA）。由于输入是密集且连续的，因此不需要自定义的稀疏内核。
回散（Scatter-Back）：一个确定性的散播内核将注意力输出重新分布回完整序列中的原始位置。

D. 两阶段训练方案

为了确保模型保留完整的注意力能力，作者提出了一种两阶段训练方法：

阶段 1：在大部分预训练时间内使用灯塔注意力。
阶段 2：使用相同的检查点，在全密集 SDPA下恢复训练一小段时间。这个“恢复阶段”允许模型将其权重适应到完整注意力机制，而无需从头开始。

3. 主要贡献

对称分层注意力：一种新颖的设计，对称地池化 $Q, K, V$ ，创建真正的多尺度表示而非压缩内存，同时保持从左到右的因果性。
内核解耦：通过将选择置于注意力路径之外，该方法在密集子序列上利用原生 FlashAttention。这允许复用高度优化的密集内核，并简化了反向传播。
训练时正确性：本文引入了一个严格的经验标准：使用灯塔预训练的模型，在短暂恢复密集 SDPA 后，必须达到或超过从头开始训练的密集基线。作者证明，灯塔训练的模型满足这一标准，证明了分层训练不会“掏空”模型使用完整注意力的能力。
高效实现：使用融合的分块双调 Top-K 和自定义的 scatter-back 内核，使得在多级 Blackwell GPU 上针对极端上下文长度（高达 1M 个 token）进行高效训练成为可能。

4. 实验结果

作者使用一个 5.3 亿参数的 Llama-3 风格解码器进行了实验，该模型在 C4 数据集上训练，上下文长度为 98,304 个 token。

可恢复性：使用灯塔预训练并随后在密集 SDPA 下恢复训练 4,000 步的模型，其最终训练损失（0.6980–0.7102）低于在相同 token 预算下从头开始训练的密集基线（0.7237）。这证实了分层训练信号不会降低模型的密集注意力能力。
吞吐量和加速比：
- 在预训练阶段，与密集 SDPA 相比，灯塔实现了2 倍的吞吐量优势（tokens/sec/GPU）。
- 端到端来看，对于相同的 token 预算，两阶段方案相比从头开始的密集训练实现了1.40 倍到 1.69 倍的挂钟时间加速，且最终损失相当或更低。
- 在 512K 的上下文长度下，与密集 SDPA 相比，灯塔的前向传播速度快21 倍，前向加反向传播速度快17.3 倍。
扩展性：该方法成功扩展到100 万 token，使用了 32 张 Blackwell GPU 并结合上下文并行性，同时保持了对密集基线的加速优势。
消融实验：
- 评分器：无参数的 $\ell_2$ 范数评分器表现与更昂贵的膨胀 softmax 评分器相当，但速度快约 9%。
- 超参数：较小的池化因子（ $p=2$ ）、较浅的金字塔（ $L=3$ ）和较小的 Top-K 预算（ $k=1536$ ）通常能产生最佳的损失/性能权衡。
检索（大海捞针）：在恢复密集 SDPA 后，灯塔训练的模型实现了与密集基线相当或更好的检索率（例如，平均检索率 0.76 对比 0.72），验证了长上下文保留能力。

5. 意义与主张

本文声称，灯塔注意力通过将注意力复杂度从二次方降低到次二次方（对于有界的 $K$ ，实际上在 $N$ 中是线性的），同时保持标准 Transformer 的架构完整性，从而解决了长上下文预训练的根本瓶颈。

作者强调，他们的方法是仅用于训练的；选择机制会被移除（或者更准确地说，模型会过渡）为全密集模型以用于推理。这将灯塔与仅推理的稀疏方法区分开来，后者继承自其密集骨干网络的“正确性底线”，但从未证明其稀疏训练过程能产生胜任的密集权重。

作者承认的局限性：

对称池化假设所有查询在一次前向传播中共同出现，这在自回归解码期间被违反。因此，该方法依赖于密集 SDPA 恢复阶段来生成可用于推理的模型。
内部注意力在收集后的子序列上是 $\Theta(S^2 d)$ 。虽然相对于 $N$ 是次二次方的，但如果选择预算 $K$ 必须随 $N$ 扩展，则它并非严格线性。
该方法目前需要两阶段训练方案；尚未集成原生的稀疏推理目标（如 DSA 或 NSA），尽管作者建议将其作为未来的方向。

总之，灯塔注意力提供了一条途径，可以在显著降低计算成本的同时，在百万 token 的上下文中训练模型，并保证所得模型在推理时完全有能力利用密集注意力机制。