用 Lighthouse Attention 实现长上下文预训练：对称池化 + 训练期稀疏化 Long Context Pre-Training with Lighthouse Attention
Lighthouse：训练期对称Q/K/V金字塔+Top-K，续训恢复全注意力
前置知识
Scaled Dot-Product Attention (SDPA)
Transformer 核心算子，公式为 softmax(QK^T/√d)V，时间和显存代价都是 N² 级别的二次方。FlashAttention 通过分块 softmax 降低常数但不改渐近，所以序列长度到 10⁵ 时就会成为瓶颈。
本文所有动机都建立在打破 SDPA 二次方瓶颈上；理解 SDPA 才能体会为什么需要稀疏化、子序列化和重计算。
稀疏/选择式注意力（Sparse / Selection-based Attention）
用某个评分函数（学习或无参数）从全部 $N$ 个 key 中选出 top-K 个，让每个 query 只看一小部分。代表工作有 MoBA（块级）、NSA（块级+学习索引）、DSA（token 级+学习索引）、HISA（分层索引）。
Lighthouse 属于这一家族，但反对把选择耦合进 attention 算子、并首次对称池化 Q/K/V；只有了解前作才能体会它为何自称解决了两个设计缺陷。
多尺度金字塔（Multi-Resolution Pyramid）
把长度 N 的序列按池化因子 p 逐级下采样：第 0 级是原序列，每升一级用平均池化把 [ip^ℓ, (i+1)p^ℓ) 窗口压成一项；总条目数上限约 N·p/(p-1)，构造和遍历代价 Θ(N)。
金字塔是 Lighthouse 的核心数据结构；用 L=3、p=4、k=4096 在 N=10^6 时把子序列压到 S≈6.5×10^4，代价从二次方降到 polylog。
Top-K 选择与 Bitonic Sort
对每个候选打分后选最大的 K 个。Bitonic sort 是一种可在 GPU 上并行、O(N log²N) 的排序网络，chunked-bitonic 把流式数据切成 CTA 块在寄存器内维护一个 top-m 缓冲，避免共享内存爆炸。
Lighthouse 用自己的 chunked-bitonic kernel 替代 torch.topk，论文强调它在 K=4096 时仍高效，并防止 "span collapse"（跨度塌缩）。
Context Parallelism (CP) 与 Ring Attention
把序列维度切到多卡，每卡负责一段子序列，注意力通过环形 all-gather 旋转 K/V 块实现。代价是引入集合通信开销但保持 $O(N^2)$ 计算按卡数线性下降。
本文在 256K-1M 上下文上必须用 CP（单卡 B200 在 100K 处 OOM），Lighthouse 因为选择 decouple from attention，所以可以无修改地套在 stock ring attention 上。
研究动机
当序列长度迈向 128K、1M 甚至更长（agent 多步推理、长文档、交织多模态）时，标准 SDPA 的 $\Theta(N^2)$ 计算和显存成为训练的主要硬件瓶颈——FlashAttention 只降低常数，不改变渐近复杂度。现有稀疏化方案分成三族：线性注意力/SSM 牺牲长程召回；块级稀疏（MoBA、NSA、MInference 等）只在 K/V 侧做块粒度剪枝、且要求一个块全留或全弃；token 级稀疏（DSA、H2O、SnapKV 等）则在推理时砍 KV cache，或在训练时把选择算子焊进 attention kernel，使 dense kernel 无法复用，更糟的是这些方法几乎都是不对称的——query 保持全分辨率，金字塔只充当"压缩可寻址内存"而非真正的多尺度表示。一旦在训练期用了稀疏算子，模型最终是否仍是合格的 dense 模型是个悬而未决的问题：推理期稀疏方法有 dense backbone 当正确性下限，但训练期稀疏方法必须回答"训完后能否被还原回全注意力"。
本文的目标是提出 Lighthouse Attention——一种"训练专用"的对称选择式分层注意力，能包裹普通 SDPA 跑而不修改它、训练末尾可被完全移除。具体三大目标：(i) 设计一个子二次方的分层前后处理，对 Q/K/V 做自适应压缩解压；(ii) 设计对称压缩策略，同步池化 Q、K、V 同时严格保留自左向右的因果性以最大化并行；(iii) 设计两阶段训练流程——绝大部分步数用 Lighthouse 训练，最后一小段恢复成全 SDPA 微调以保证推理期可部署。最终目标是在匹配的 token 预算下达到与 dense-from-scratch 相同或更低的训练损失，并把总训练时间显著缩短。
与已有工作不同的是，切入角度与已有工作有三个本质差别。第一，对称性：把 query 和 key/value 一样池化进同一金字塔，得到相干三元组 $(Q^{(\ell)}, K^{(\ell)}, V^{(\ell)})$，让 query 和 key 共处同一表征空间且每个金字塔项都是同一个 $p^\ell$ 长跨度的摘要，从而在选择阶段能用 $\ell_2$ 范数这种无参数打分器做无差别排序。第二，位置正确：把选择完全搬出 attention kernel——top-K 之后 attention 就是 stock FlashAttention 在一段连续 dense 子序列上跑，前向反向与 dense Transformer bit-for-bit 一致；选择和 attention 解耦也意味着能直接套用 stock ring attention 跑 context parallelism，不需任何 sparse-aware 集合通信。第三，正确性标准：明确把"在 Lighthouse 预训练后做短暂 SDPA 续训能否匹敌 dense-from-scratch"作为核心实证判据，这是第一个以此为硬指标的训练期分层方法。
核心方法
Lighthouse 把一层普通 attention 替换成一条"四阶段流水线"，包裹但不修改 SDPA 算子。直觉上：与其在所有 $N$ 个 token 上做 $N \times N$ 的 softmax 注意力，不如先用一个无参数打分器在多层金字塔上给每段 span 打分、选出 top-K 个最值得注意的 span，把它们 gather 成一段连续、严格因果的子序列 $S \ll N$，让 stock FlashAttention 在 $S$ 上跑出 $\Theta(S^2 d)$ 的代价，最后把结果 scatter 回原 $N$ 个位置。整条流水线不引入新可学习参数——打分器是 Q/K 自己的 $\ell_2$ 范数、Top-K 是离散不可微的算子、金字塔只是固定平均池化；梯度只能从 loss 经过 scatter、FlashAttention、gather 流回投影矩阵 $W_Q, W_K, W_V$，迫使投影学会"在被选中时更有用"，而不是"会打高自己的分"。这样既绕开了学习式 scorer 的优化脆弱性，又让上层的 FlashAttention 升级（如 FA-3、FA-4、ring attention）能直接继承。
三个本质区别于 NSA/HISA/DSA/MoBA 的核心创新点。其一，**对称 Q/K/V 池化**：在 $L$ 层、池化因子 $p$ 的金字塔上对 Q、K、V 同时做 mean-pooling，得到相干三元组；而 NSA/HISA/InfLLM-v2 只池化 K/V、保持 query 全分辨率，因此 dense kernel 仍然吃不下 query 端。其二，**选择外置**：top-K 输出一组整数索引，gather 把对应三元组重排成连续因果子序列，attention 是普通 SDPA/FlashAttention；不存在 NSA 那样的 sparse matmul、也不存在 DSA 那种 per-query gather，dense kernel 库零修改就能用。其三，**无参数无 ST 的打分**：用 $\|Q^{(\ell)}_i\|_2$ 和 $\|K^{(\ell)}_i\|_2$ 当 scorer、coarser 级从 level-0 max-pool 继承范数；top-K 离散不可微，不上 straight-through、不上 Gumbel、不加辅助 loss——这意味着任何正结果都是"更强 scorer 能做到的下界"。
方法步骤详情
算法 1 给出单头单层的完整流程（共 7 步）。**Step 1 投影**：输入 $X \in \mathbb{R}^{N \times d_{\mathrm{model}}}$ 与 $W_Q, W_K, W_V \in \mathbb{R}^{d_{\mathrm{model}} \times d}$ 相乘得到 $Q, K, V \in \mathbb{R}^{N \times d}$（一次融合 GEMM）。**Step 2 金字塔**：用平均池化把 Q、K、V 同步构造成 $L$ 级金字塔，第 0 级是原序列，第 $\ell$ 级窗口 $W^{(\ell)}_i = [i p^\ell, (i+1)p^\ell)$ 摘要为一项，构造代价 $\Theta(N)$。**Step 3 打分**：用 $\|Q_j\|_2$ 和 $\|K_j\|_2$ 做 per-head $\ell_2$ 范数，coarser 级从 level-0 max-pool 继承；这让粗 span 自动继承最强 token 的重要度。**Step 4 Top-K**：把 QK 和 KQ 两路分数拼起来跑 chunked-bitonic top-K，强制保留最粗级全部 $N/p^{L-1}$ 项，剩余预算分给更细级；输出整数索引集 $I$。**Step 5 Gather**：用 $I$ 把 $(Q^{(\ell)}, K^{(\ell)}, V^{(\ell)})$ 三元组按 $\ell, i$ 拓扑排序重排为连续子序列 $\tilde Q, \tilde K, \tilde V \in \mathbb{R}^{S \times d}$，其中 $S = N/p^{L-1} + (L-1) p k$（$p$ 因子来自因果边界处理），并在 $N=10^6, L=4, p=4, k=4096$ 时把 $S$ 压到约 $6.5 \times 10^4 \ll N$。**Step 6 Dense Sub-sequence Attention**：调用 stock FlashAttention 算 $\tilde O = \mathrm{Attn}(\tilde Q, \tilde K, \tilde V; \tilde M)$，$\tilde M$ 由金字塔坐标导出的因果 mask（dense $S \times S$ 即可，无 sparse 索引）。**Step 7 Scatter-Back**：用 deterministic integer-atomic scatter 把 $\tilde O$ 写回原 $N$ 个位置；每个选中项 $(\ell, i)$ 写到偏移 $p^\ell - 1$ 之后的 $p^{\ell-1}$ 个连续位置以保证因果（一个 base 位置不会收到含自己未来的摘要），跨级贡献在每个位置累加，per-position fan-in $\le L$。
技术新颖性
技术新颖性集中在四点。**第一**，**对称 Q/K/V 池化**——和 NSA/HISA/InfLLM-v2 那种"只池 K/V、query 全分辨率"的设计分道扬镳，量化收益是把训练期 dense kernel 调用从 $O(N S d)$ 降到 $O(S^2 d)$。**第二**，**参数自由且无 ST 的 scorer**——用 $\ell_2$ 范数天然比任何 attention-style scorer 弱，因此正结果给"更强 scorer"留出明确提升空间；dilated-softmax 替代 scorer 的对照实验显示两者差距 $\sim 0.01$ loss 但 $\ell_2$ 范数版在 B200-hours 上便宜约 9%。**第三**，**选择和 attention 解耦**——前向反向与 dense Transformer bit-for-bit 一致，能完全复用 FA-3/FA-4 优化、也能零修改套 stock ring attention 跑 context parallelism（论文在 32 卡 B200 上做到 1M token 训练）。**第四**，**明确以 dense-SDPA 续训当正确性准则**——以前没有训练期分层方法把"续训后能否匹敌 dense-from-scratch"当硬指标提出来；本文的实验在 16k 步固定预算下用三种切换点（10k+6k / 11k+5k / 12k+4k）证明都能跨过 dense baseline 的最终 loss。
实验结果
实验沿三个轴评估。**(1) 可恢复性（SDPA Recoverability）**：固定总预算 16k 步（$\approx 50.3$B tokens），把 stage-1 Lighthouse 长度在 10k/11k/12k 间切换、剩余用 dense SDPA 续训。切换瞬间训练损失会有 1.12–1.57 的尖峰（模型第一次跑没训过的算子），但 1–1.5k SDPA 步内就跌回并穿过 dense baseline；最终 loss 分别达到 0.6980 / 0.7001 / 0.7102，均**低于** dense-from-scratch 的 0.7237，最佳点在 stage-1 最短的 10k+6k（dense 续训越长越好）。**(2) Scaling vs Dense**：单层 B200 测注意力延迟，从 8K 到 512K 全面胜出，$N=512$K 时 forward 快 21×、forward+backward 快 17.3×——即 SDPA 跑到 $\sim 113$K（fwd）/$\sim 122$K（fwd+bwd）才与 Lighthouse 在 512K 的耗时打平。**(3) 设计消融**：表 1 的全网格显示 (a) 所有 Lighthouse 配置最终 loss 都 $\le$ 0.7237，可恢复性不是某个超参巧合；(b) projection-norm scorer 与 dilated-softmax scorer 互有胜负但差距 $\sim 0.01$、norm 版本便宜约 9%；(c) 网格中最低 loss 出现在 L=3, p=2, k=1536（dilated scorer）= 0.6825；同时 throughput 在 Lighthouse stage-1 维持 84–126k tok/s/GPU，**vs SDPA 的 45.6k tok/s/GPU**，大约 2× 每步优势；端到端 10k+6k 配方的总墙钟时间在 22.5h（norm、k=1536、p=4）到 27.0h（k=4096、p=2、L=3）之间，**对比 dense-from-scratch 的 37.9h（303.2 B200-h）**，实现 **1.40× 到 1.69× 端到端加速**且最终 loss 持平或更低。**Context Parallelism 块**：在 256K 上下文、CP=8、DP=1 下，最终 loss 降到 0.6721（与 dense-from-scratch 同样的 16k 步但 1.07T tokens，1M 级），证明 CP 下 Lighthouse-vs-SDPA 的乘性优势不变。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| SDPA 可恢复性（10k LH + 6k SDPA） | 530M 模型最终训练 loss | 0.6980（dilated scorer） | 0.7237（dense-from-scratch 16k 步） | loss 降低 0.0257（≈3.5% 相对） |
| SDPA 可恢复性（11k LH + 5k SDPA） | 530M 模型最终训练 loss | 0.7001 | 0.7237 | loss 降低 0.0236（≈3.3% 相对） |
| SDPA 可恢复性（12k LH + 4k SDPA） | 530M 模型最终训练 loss | 0.7102 | 0.7237 | loss 降低 0.0135（≈1.9% 相对） |
| 单层注意力延迟（N=512K, 单 B200） | Forward / Forward+Backward 加速比 | Forward 21×、Fwd+Bwd 17.3× | cuDNN SDPA 在 512K 的耗时 | 等价于 SDPA 需跑到 ~113K（fwd）/ ~122K（fwd+bwd）才能追平 |
| Stage-1 训练吞吐（530M, 98K ctx, 8×B200） | tokens/s/GPU | 84k–126k tok/s/GPU（k 越小越快，norm scorer L=3 p=4 k=1536 顶到 126k） | 45.6k tok/s/GPU（dense SDPA） | 约 2× 每步 throughput 优势 |
| 端到端 16k 步训练墙钟（10k+6k 配方） | B200-hours | 179.6–215.7 B200-h（超参范围内） | 303.2 B200-h（dense-from-scratch） | 1.40×–1.69× 端到端加速，loss 持平或更低 |
| CP 长上下文训练（256K, CP=8 DP=1, 1.07T tokens） | 最终 loss / B200-h | loss 0.6721、1300.3 B200-h | （dense SDPA 在 100K 上下文已 OOM） | 在 dense 跑不到的规模上 Lighthouse 仍可训练并继续降 loss |
| 超参网格最低 loss（消融） | Lighthouse stage-1 + 6k SDPA 续训的 loss | 0.6825（L=3, p=2, k=1536, dilated scorer） | 0.7237（dense-from-scratch） | loss 降低 0.0412（≈5.7% 相对），且是该网格 Pareto-best |
局限与改进
作者明确承认两条硬约束。**第一**，对称 Q/K/V 池化要求所有 query 在同一次前向里出现，自回归解码（一次只生成一个 token）天然违反——因此本文最终可部署的 checkpoint 必须经过 dense-SDPA 续训；论文里所有"下游评估"都是在续训后的模型上跑、而不是直接在分层前向上跑。**第二**，内层 attention 仍是 $\Theta(S^2 d)$ 的子序列注意力——在固定 $k$ 下对 $N$ 是 polylog、但**不是严格线性**，所以"必须让 $k$ 随 $N$ 增长"的更激进场景还没刻画清楚。**我的独立观察**：(a) 整个实验在单模型规模（530M）和单数据集（C4）上，**未在更大规模（≥1B、≥7B）或 Pile/FineWeb 等异构数据上验证**，scaling law 是否仍成立未知；(b) 最终 loss 的提升是 0.025 量级，**没看到困惑度/下游任务（LongBench、Needle-in-a-Haystack、HumanEval）的实测**，对实际部署的"长文档理解/检索"能力证据不足；(c) 1M token 训练的结果只在 1.07T tokens 那条 CP 线上展示（loss 0.6721），但缺少对应 dense 基线作对照，"CP 下 Lighthouse 仍优"的结论是间接成立的；(d) 论文没仔细分析当 $k$ 相对 $N$ 极小时（比如 $k=1536$ 在 98K 上）会不会丢关键 attention sink 现象；(e) 自定义 CUDA/Triton kernel 增加了实现门槛，给非 Nous 团队复现带来挑战。
独立分析的弱点
**弱点一：必须做 SDPA 续训才能部署**——Lighthouse 是"训练专用"，推理时直接用分层前向会因自回归解码破坏对称池化假设，因此实际产出是"Lighthouse pretrain + dense resume"的两阶段 checkpoint。**改进方向**：把 dense resume 换成某种不对称稀疏目标（DSA/NSA/HISA/MoBA），或研究"前 K-1 层用 dense，最后几层保留 Lighthouse 也能解码"的混合方案，让 checkpoint 原生可服务。**弱点二：超参 $\{L, p, k\}$ 是全局固定**——不同层不同 head 适合不同的稀疏度（早期层偏局部、晚期层需要长程），但当前 530M 模型中只有 layer $\{0, 1, 28, 29\}$ 保留 dense、其他 26 层统一用同一组 $(L, p, k)$。**改进方向**：做 per-layer 或 per-head 自适应 $k$、让 L/p 也可学习，或引入 routing 网络根据输入动态决定预算。**弱点三：scorer 仍是 query 独立**——$\ell_2$ 范数不依赖 query-key 交互，对一些 query-specific 重要位置可能不敏感。**改进方向**：保留参数自由的属性但加一个 query-conditioned 评分（如 $s_{i}^{(\ell)} = \|Q^{(\ell)}_i\|_2 \cdot \|K^{(\ell)}_i\|_2 + \alpha \langle Q^{(\ell)}_i, K^{(\ell)}_i\rangle$），既不引入可学习参数又捕到交互信号。**弱点四：评估维度单薄**——只有训练 loss 一个核心指标，缺下游长上下文 benchmark（Needle-in-a-Haystack、LongBench、RULER、InfiniteBench 等）和困惑度曲线。**改进方向**：补齐长文档 QA / 检索 / 代码补全的评估，并发现在哪些任务上"Lighthouse pretrain + dense resume"的 loss 优势能传导到下游能力。**弱点五：只跑了 530M 规模**——能否在 7B/13B/70B 上保持端到端 1.4–1.7× 加速未知。**改进方向**：沿 scaling law 做 $\{350M, 1B, 3B, 7B\}$ 的扫描，给出"Lighthouse 加速比随模型规模如何变化"的经验曲线。
未来方向
**作者明确点出的方向**：(a) 用不对称稀疏目标（DSA/NSA/HISA/MoBA）替代 dense-SDPA 续训，得到原生可服务 checkpoint；(b) per-layer / per-head 自适应 $k$ 超越固定预算；(c) 多尺度金字塔自然扩展到 vision、audio、video 等多模态长序列；(d) serving 集成（continuous batching、speculative decoding、KV-cache 管理）把训练加速翻译成部署加速。**基于成果的延伸方向**：(e) 探索把 $\ell_2$ 范数 scorer 升级为 query-key 交互型 scorer、但仍保持参数自由，把"正结果下界"兑现成实际增益；(f) 研究把 Lighthouse 用在 RLHF / reasoning trace 训练中（这些场景的 context 长度经常突破 100K 且对长程依赖极敏感）；(g) 与 linear attention / SSM / hybrid 架构对比，明确在哪些 N/d_model 区间里各自占优；(h) 把 chunked-bitonic top-K 抽象成一个通用 GPU 算子库给社区复用，因为它的"流式 + 寄存器 top-m"模式对 streaming TopK 场景有普适价值；(i) 探索"Lighthouse pretrain + LoRA/DoRA fine-tune"的轻量化两阶段流程，让中等规模团队也能负担长上下文模型训练。
复现评估
**开源情况**：完整代码已开源在 https://github.com/ighoshsubho/lighthouse-attention，Algorithm 1 把单层单头流程写得相当清楚（投影 → 金字塔 → 打分 → top-K → gather → FlashAttention → scatter-back），理论上任何熟悉 PyTorch + CUDA/Triton 的工程师都能复现。**数据**：训练用 C4（公开数据集），无需特殊协议。**算力门槛**：这是**最大的复现门槛**——(a) 单 8×B200 节点跑 98K 上下文是 baseline 实验的最低配置，B200 是当前 NVIDIA 旗舰，单节点采购价 $\sim \$300k+$；(b) 256K-1M 上下文实验需要多节点 + 32 卡 B200 + context parallelism，几乎只有大厂或一线实验室能复现；(c) 16k 步 + 50B–1T tokens 的训练量对单卡团队来说时间也不可控。**难度评级**：方法本身实现难度中等（自定义 kernel 是 CUDA + Triton 混合，但算法描述清楚），跑通实验的难度高（需要 B200 集群、长上下文训练栈、context parallelism 工程能力）。**作者还提到 LLM pretraining 之外**的"可恢复性"准则对整个训练期稀疏方法家族都有参考价值。
论文图表
