Lighthouse Attention：让长上下文预训练只在"塔尖"上算注意力

长上下文训练的瓶颈，从来不是"模型不够好"，是"算不动"

前沿模型的上下文已经卷到 128K、1M 甚至更长——agentic 多步推理、长文档理解、交错的多模态输入都在往上推。但在这种尺度下，训练才是真正的硬件瓶颈：标准的缩放点积注意力（SDPA）有 $\Theta(N^2)$ 的算力和显存开销。FlashAttention 把常数项压下去了，但没有改变这个二次方的渐近复杂度——当 $N\geq 10^5$ 时，这一项就开始主宰一切。

一条主流的破局思路是用选择（selection）替代稠密注意力：每个 query 只关注一小撮 key。但论文指出，现有这些"稀疏注意力"方法虽然在推理上有效，却带着两个对长上下文预训练很不友好的设计基因：

在"打破二次方注意力"这条赛道上，前人都试过哪些路？

这是理解 Lighthouse 定位的关键一节。论文把破解 $\Theta(N^2)$ 的努力分成两大脉络——压缩与剪枝（Compression & Pruning）和层次化与训练期正确性（Hierarchies & Training-time correctness）。下面把每条路、代表方法、以及它们的"软肋"都摊开讲。

脉络一：压缩与剪枝——三种"少算一点"的思路

面对二次方注意力，第一类回应是直接放弃 softmax、改用一个有界大小的状态；第二类保留 softmax 但在块（block）粒度上剪枝；第三类则在token 粒度上剪枝。论文把它们整理成下面这张谱系表：

脉络二：层次化注意力 & 训练期正确性

多分辨率（multi-resolution）注意力近年又回到了稀疏 LLM 注意力里，主要有两种玩法：

但论文一针见血地指出这两类的共同局限：层次结构只作用在 key 和 value 上，而最终选出来的结果仍然要喂进一个自定义的稀疏注意力核。换句话说，不对称 + 架构纠缠这两个原罪，它们一个都没躲掉。

Lighthouse 和它们到底差在哪？三个轴

把注意力层换成一条"四阶段流水线"——但绝不碰注意力核本身

Lighthouse 用一条四阶段流水线替换掉标准注意力层。关键词是"包住但不修改"：选择发生在注意力之前，FlashAttention 在被挑出来的子序列上原样运行，注意力之后再把结果散射回原位置。选择由一个无参数的打分函数驱动，作用在该层自己 $Q,K,V$ 的多分辨率金字塔上——所以 Lighthouse 不引入任何新的可学参数。

3.1 先回顾：标准注意力为什么贵

设输入 $X\in\mathbb{R}^{N\times d_{\text{model}}}$，投影矩阵 $W_Q,W_K,W_V$，因果掩码 $M$，标准缩放点积注意力为：

3.2 四阶段流水线总览

3.3 金字塔构造（Pyramid）

第 $\ell$ 级是上一级的非重叠窗口池化。第 $\ell$ 级第 $i$ 个窗口定义为 $\mathcal{W}^{(\ell)}_i=[\,ip^\ell,\,(i+1)p^\ell-1\,]$，然后对三个投影对称地做均值池化：

对称带来两个后续阶段会用到的关键性质：池化后的 query $Q^{(\ell)}_i$ 和池化后的 key $K^{(\ell)}_j$ 处在同一表示空间；每个金字塔条目都是一个概括同一段 $p^\ell$ token 跨度的连贯 $(Q,K,V)$ 三元组。金字塔总条目数 $\sum_{\ell=0}^{L-1}N/p^\ell \leq N\cdot p/(p-1)$，所以构造金字塔只花 $\Theta(N)$ 时间和显存。

3.4 打分与选择（Scoring & Selection）

每个金字塔条目拿两个标量分——一个作为 query（QK），一个作为 key（KQ）。第 0 级直接用每头的 $\ell_2$ 范数：

更粗的层级不重算，而是从第 0 级max-pool 上来——让一个粗跨度继承它内部最强 token 的重要性。然后跨所有层级、把 QK 和 KQ 两条流拼起来联合选 top-$K$：

3.5 稠密子序列注意力（Gathered-Sequence Attention）

按索引 $\mathcal{I}$ 把选中的三元组拼成一条连续子序列 $\widetilde{Q},\widetilde{K},\widetilde{V}$，长度为：

代入一个真实数字感受一下：当 $N=10^6,\,L=4,\,p=4,\,k=4096$ 时，$S\approx 6.5\times10^4 \ll N$——一百万 token 的序列，实际进注意力的只剩六万五。然后直接用原版 SDPA / FlashAttention：

由于 gather 是按金字塔坐标拓扑排序过的，因果掩码 $\widetilde{M}$ 退化成一个标准的 $S\times S$ 因果掩码——式中没有任何稀疏索引。这就是 Lighthouse 能直接复用 FlashAttention 的根本原因。

3.6 散射回写（Scatter-Back）

最后一步用一个确定性的"整数原子散射"kernel，把每个条目的输出 $\widetilde{O}_m$ 分发回它所代表的 $p^\ell$ 个基础位置，得到完整的 $O_t$。整条流水线里只有② top-K 和 ④ scatter-back 是自定义 CUDA/Triton kernel，其余全是 torch.compile 能融合的 PyTorch 原语。

为什么它能做到"近线性"，又为什么能直接吃 FlashAttention 的红利

整条流水线唯一超线性于 $N$ 的项，就是那条稠密子序列注意力 $\Theta(S^2 d)$。而 $S=N/p^{L-1}+(L-1)pk$。如果选 $L=\log_p(N/k)$ 来平衡 $S$ 的两项，就得到 $S=\Theta(k\log_p(N/k))$，注意力成本变成：

加上线性的打分和 $\Theta(N\log k)$ 的选择，单层总算力对 $N$ 在固定 $k$ 下是近线性（仅差一个 $\log k$ 因子）。

三个问题：能恢复吗？能扩展吗？四个旋钮怎么调？

实验基座：一个 530M 参数的 Llama-3 风格 decoder（$d_{\text{model}}{=}1024$，30 层，8 头，head dim 128），在 C4 上以序列长度 98,304 训练，AdamW，总预算固定 16,000 步（约 50B tokens）。第 0、1、28、29 层保留稠密 SDPA，其余 26 层用 Lighthouse。两阶段配方：阶段一用 Lighthouse 训练，阶段二载入 checkpoint 用稠密 SDPA 续训一小段。

5.1 可恢复性（Recoverability）——全篇的"命门实验"

固定 16k 步总预算，变化阶段一的长度（10k / 11k / 12k），剩下的步数切回稠密 SDPA 续训，与一个从头稠密训练、架构/数据/token 全部对齐的基线对比。

规律很干净：切回稠密时 loss 会先瞬间尖刺（1.12–1.57，因为模型第一次跑它没训过的注意力），然后约 1–1.5k 步内恢复，并反超基线。三个切换点最终都落在 0.6980–0.7102，全部低于稠密基线的 0.7237——而且 B200 小时还更省。

5.2 扩展性：和稠密注意力的延迟对比

论文 Fig.3(a)：前向延迟 vs 上下文长度（单 B200，$L{=}3,p{=}4$，稀疏度约 1:64）。 绿线 SDPA(cuDNN) 按 $\Theta(N^2 d)$ 飙升——512K 时超 400ms；红线 Lighthouse 按 $\Theta(S^2 d)$ 几乎贴地。512K 上 Lighthouse 快 21.0×；换个说法，SDPA 跑到 ~113K 上下文才达到 Lighthouse 在 512K 的延迟。

论文 Fig.3(b)：前向+反向延迟。 同样的趋势——512K 上 Lighthouse 快 17.3×，相当于 SDPA 在 ~122K 上下文的运行时间。反向比前向加速略低，因为反向要穿过 gather/scatter，但量级优势依旧保持。

全模型训练故事类似但需要小心：530M 架构在单卡 B200 上超过 ~100K 就会因激活/梯度/优化器状态 OOM（与注意力方法无关）。于是论文实现了上下文并行：金字塔池化、打分、top-K 都分片本地跑，gather 出的子序列走原版 ring attention 轮转，无稀疏感知集合通信。CP 引入约 10% 的每卡吞吐开销，但 Lighthouse-vs-SDPA 的乘性加速在同等 CP 几何下完整保留，一路干净地扩展到 1M token / 32 卡。

5.3 设计消融：四个旋钮（scorer / $p$ / $L$ / $k$）

吞吐侧：Lighthouse 阶段一稳定在 84–126k tok/s/GPU，对比稠密 SDPA 的 ~46k，约 2× 每步优势。端到端（10k+6k 配方）总运行时间从 22.5h（179.6 B200-h）到 27.0h，对比稠密从头训练的 37.9h（303.2 B200-h）——即 1.4–1.7× 端到端提速。

5.4 长上下文检索（Needle-in-a-Haystack）

论文 Fig.4：单数字 passkey 检索热力图（深度 0–100%，上下文 4K–96K，随机基线 10%）。 四个 Lighthouse 两阶段 run + 稠密基线（0.72）。绿=高命中、红=低。三个 Lighthouse run 打平或超过稠密基线：$k{=}2048$ dilated 以 0.76 夺冠，$k{=}1536$ dilated 0.73，$k{=}2048$ norm 0.72 持平；只有 $k{=}1536$ norm 掉到 0.65。

一个"训练用稀疏、推理用稠密"的干净配方

Lighthouse 把 $Q,K,V$ 对称池化进多分辨率金字塔、把选择放在注意力核外面，让注意力步退化成稠密子序列上的原版 FlashAttention。它无参数、端到端训练、无辅助 loss、无 straight-through 估计，还能原样继承上游 FlashAttention 的每一次改进。一段短暂的稠密 SDPA 续训后，在等 token 预算下追平或超过从头稠密训练（loss 和长上下文检索都成立），≥100K 上对 cuDNN SDPA 端到端快 1.4–1.7×，并干净扩展到多节点 Blackwell 的 1M token。

这篇论文真正聪明的地方在哪？

站在做长上下文训练和推理（KV-Cache、前缀缓存这些）的人的角度，我觉得这篇有三点值得反复琢磨：