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 —— 那才是生產級採用的瓶頸所在。