DC娱乐网

NousResearch:长上下文预训练加速新方案

在 LLM 朝着 128K、1M+ 上下文狂飙的今天,标准 Attention 的 O(N²) 复杂度成了最硬的硬件墙。FlashAttention 优化了常数,但没解决根本问题。

今天给大家分享 Nous Research 的最新论文《Long Context Pre-Training with Lighthouse Attention》,论文核心提出了一种新的注意力机制——Lighthouse Attention 。它像一个「包装器」包裹在普通 FlashAttention/SDPA 外面:
- 大部分预训练步骤用 Lighthouse 加速;
- 训练末期短暂切换回全注意力(resume)恢复,得到一个全注意力兼容的模型。

🔆 核心流程
对称金字塔池化(Q/K/V 同时池化) → 参数-free 的 Top-K 选择 → 聚集成长序列用标准 FlashAttention 计算 → Scatter-back 重构。选择逻辑完全在注意力内核之外,完美复用现有高度优化的内核。

🧠 创新点
1.对称压缩:不像多数稀疏方法只压缩 KV,Lighthouse 对 Query 也进行多尺度池化,构建真正的多分辨率表示,同时保持因果性。
2.解耦选择与计算:Top-K 和 gather/scatter 是独立内核,主体用 stock FlashAttention,无需自定义稀疏 attention 内核。
3.梯度-free 选择 + Resume 策略:Top-K 不可微但梯度正常流动;通过短暂全 SDPA 恢复,确保最终模型在 loss 和长上下文检索(NIAH)上匹配甚至优于从头全注意力 baseline。
4.高效内核:支持极长上下文和 context parallelism。

这篇论工作为解决 LLM 「长文本预训练贵、训练难」 的问题提供了一个高效的工程化方案。
Lighthouse Attention 通过巧妙的稀疏策略,平衡了「计算效率」与「长距离建模能力」,对于未来开发具备超长上下文理解能力的基座模型具有重要的参考意义。

代码已开源,对于做大模型预训练工程师(尤其做 100K+ 上下文的团队)和关注高效训练、稀疏注意力、层次化表示 的 ML 研究人员来讲,这是一篇值得关注的论文,推荐一读。