Nous Research 本周发表了 Lighthouse Attention —— 一种仅训练阶段使用的层级注意力机制,把 query、key、value 对称地池化进一个多层金字塔,把 top-K 选择放在 kernel 外面跑,让 FlashAttention 在一段较小的稠密子序列上运行。报告的 wall-clock 预训练加速:在 530M Llama-3 风格 decoder 上,相对 cuDNN-backed SDPA,端到端 1.40-1.69×,512K 上下文单卡测试,1M token 32 卡 context parallelism 测试。512K 上下文下 kernel 级别加速更陡:forward 21×,forward+backward 17.3×。作者:Peng、Ghosh、Quesnelle。arXiv 2605.06554,代码在 github.com/ighoshsubho/lighthouse-attention,以 torchtitan 上的一个 patch 加两个新文件的形式提供。
把 Lighthouse 跟之前 NSA、HISA 工作分开来的架构选择,是 Q/K/V 对称池化,而不是只池化 K/V。之前的 selection-based attention 方法把 query 留在全分辨率上,只把 K/V 那一侧池化掉;Lighthouse 把三者都池化进金字塔,然后在它们之上跑一个 ℓ₂-norm 的分块双调 top-K 选择。开销从 O(N·S·d) 变成 O(S²·d)。四阶段 pipeline —— 平均池化到 L 层、打分加 top-K、gather 被选中的项、在 gather 出来的稠密块上跑原版 FlashAttention、用一个确定性 kernel 把输出 scatter 回去 —— 让内层 attention kernel 跟在稠密序列上完全一模一样。这是为什么 FlashAttention 的加速能跟 Lighthouse 的选择叠加,而不是互相打架的实际原因。
训练-only 这个定位很关键。Lighthouse 在推理时被移除:两阶段训练配方 —— stage 1 带选择训练,stage 2 切到稠密 SDPA 上 resume。最终训练 loss 0.6980-0.7102,对照从头训稠密 baseline 0.7237 —— 略好一点 —— wall-clock 22.5-27.0 小时,对照稠密 baseline 37.9 小时,在同一个模型、同样 token 预算下(~50.3B token、16,000 steps)。所以胜负在训练-compute 轴上,不是推理-compute 轴上:用 Lighthouse 训出来的模型,在部署时跟一个正常稠密模型一样。这跟稀疏-attention-at-inference 的工作(StreamingLLM、KV cache 压缩)是不同的问题陈述,也跟那些把稀疏 attention 直接送进生产的架构级方案不一样。Lighthouse 是设计空间里「预训练更便宜、部署稠密」这个点。
周一上手:如果你在 commodity 训练设施上做长上下文模型预训练,Lighthouse 离上你的训练 run 做消融只差一个 torchtitan patch 加两个文件。530M 规模的结果是提示性的,不是承重的 —— 1.4-1.7× 在 7B、70B、或者 405B 上还成立不成立,是悬而未决的问题。选择开销(gather/scatter、top-K)不是随模型大小线性 scale 的,所以加速可能压缩或放大。盯紧 Nous 自己有没有在更大规模上复现,盯紧下一代 Llama、Qwen、或 DeepSeek 的预训练有没有采纳对称金字塔池化这个 trick,盯紧 GitHub repo 上有没有还没发的 cuDNN 级别 fused kernel —— 那才是生产级采用的瓶颈所在。