在數(shù)學(xué)上把稀疏注意力做對！騰訊Hy開源HiLS-Attention: 計(jì)算更少效果更好, 外推512倍

讓大模型 "讀得更長" 一直是 Agent、深度推理和海量資料整合等場景的剛需，但標(biāo)準(zhǔn)全注意力機(jī)制的計(jì)算量隨序列長度呈平方級增長，始終是橫亙在長上下文建模面前的三座大山。

本周，騰訊混元團(tuán)隊(duì)正式開源 HiLS-Attention（分層地標(biāo)稀疏注意力），提出了一種全新的分塊稀疏注意力范式，首次在數(shù)學(xué)層面上同時(shí)解決了 chunk 重要性估計(jì)的 "表達(dá)力不足" 和選擇過程的 "端到端不可導(dǎo)" 兩大根本難題，真正將稀疏注意力做到了 "Done Right"。

在 345M 至 7B 參數(shù)規(guī)模上的系統(tǒng)驗(yàn)證顯示：HiLS-Attention 在短文本場景下語言建模困惑度（PPL）與全注意力幾乎重合，在 8K 訓(xùn)練條件下可實(shí)現(xiàn) 4M 上下文（512 倍）免訓(xùn)外推，512K 上下文下的 prefill 與單步 decode 分別加速 13.5 倍和 15.7 倍。

更重要的是，該方法在部分長上下文檢索任務(wù)上反超了全注意力本身 —— 效率與效果的 "二選一" 困境被首次同時(shí)打破。相關(guān)論文與代碼已在 GitHub 上公開。

• 論文標(biāo)題：Hierarchical Sparse Attention Done Right: Toward Infinite Context Modeling
• 論文鏈接：https://arxiv.org/pdf/2607.02980
• 代碼鏈接：https://github.com/Tencent-Hunyuan/HiLS-Attention

一、長上下文的老大難：

全注意力扛不動，稀疏注意力又 "不準(zhǔn)"

讓大模型讀得更長，幾乎是今天所有 Agent、深度推理、海量資料整合任務(wù)的剛需。

但 "讀得長" 對標(biāo)準(zhǔn)的全注意力（Full Attention）來說，一直是三座大山：

• 計(jì)算量是平方級的：序列翻倍，算力翻四倍；
• 長度外推差：訓(xùn)練 8K，測試 32K 就開始崩；
• KV Cache 隨長度線性膨脹，顯存很快吃不消。

于是大家把目光投向分塊稀疏注意力（Chunk-wise Sparse Attention）：把上下文切成一個(gè)個(gè) chunk，每個(gè) query 只挑最相關(guān)的 Top-K 個(gè) chunk去算注意力，多余的 kv cache 卸載到 cpu memory, 計(jì)算和顯存開銷都被控制在常數(shù)，看起來思路非常美好。

但現(xiàn)實(shí)是：迄今為止，沒有任何一種分塊稀疏注意力能真正追平全注意力。

圖：最直接的證據(jù) —— 即便讓 345M 模型在 RULER 任務(wù)上訓(xùn)練，現(xiàn)有分塊稀疏注意力依然打不滿全注意力。

問題的根子在哪？——chunk 選不準(zhǔn)。

二、現(xiàn)有方法為什么會 "選錯(cuò) chunk"？

mean /max logits 的先天缺陷

要選對 chunk，前提是能準(zhǔn)確估計(jì)每個(gè) chunk 的 "重要性"。

圖：從最樸素的 block sparse attention 出發(fā) —— 一個(gè) chunk 的重要性，就是它內(nèi)部所有 token 注意力質(zhì)量的總和

基于上圖，形式上有:

按 Zc 從大到小取 Top-K，選出來的就是和全注意力完全一致的精確排序。問題在于：想精確算出 Zc，就得把 query 和 chunk 里每一個(gè) token 都點(diǎn)積一遍 —— 這等于又把全注意力算了一遍，稀疏的意義蕩然無存。

現(xiàn)有方法是怎么造的？最主流的是均值池化（mean pooling）：直接把 chunk 內(nèi)所有 key 求平均當(dāng)摘要 key。稍作推導(dǎo)就會發(fā)現(xiàn)，它算出來的分?jǐn)?shù)其實(shí)就是 token logit 的均值

NSA、InfLLM v2、MoBA 等都屬此類；另一類方法（如 MiniMax 稀疏注意力）則改用max logits來近似。

也就是說：

• mean logits 只在 "chunk 內(nèi)注意力均勻分布" 時(shí)才準(zhǔn)；
• max logits 只在 "單個(gè) token 獨(dú)占注意力" 時(shí)才準(zhǔn)。

可真實(shí)場景里，logit 分布隨 query隨 head隨數(shù)據(jù)劇烈變化，根本不會乖乖落在某一種極端。結(jié)果就是：無論用 mean 還是 max，都是在用一個(gè)只在極端情況才成立的代理去硬湊 LogSumExp，從而系統(tǒng)性地錯(cuò)估 chunk 重要性，打亂 chunk 排序，讓真正關(guān)鍵的 chunk 落選

這一點(diǎn)在實(shí)驗(yàn)里看得很清楚：在最簡單的單針大海撈針任務(wù)上，用均值池化的 NSA / DashAttention / InfLLM v2 在 8K 域內(nèi)就已經(jīng)明顯掉點(diǎn)—— 因?yàn)榇蠛漆樓∏∈?"少數(shù)針 token 獨(dú)占注意力" 的高度集中分布，而均值池化會把這種尖峰稀釋掉

三、想用 "參數(shù)化 summary" 救場？

先過端到端反傳這一關(guān)

既然非參數(shù)化的 mean/max 表達(dá)力不夠，那很自然的想法是：給每個(gè) chunk 學(xué)一個(gè)參數(shù)化的 summary，讓它更有表達(dá)力地概括整塊內(nèi)容。

聽起來對，但這里藏著一個(gè)被幾乎所有現(xiàn)有方法忽略的致命斷點(diǎn)

現(xiàn)有方法即便用了參數(shù)化 summary，也只拿它來打分選 Top-K。一旦 Top-K 的 chunk ID 被 "硬選" 出來，summary 和打分就被丟棄了，不再參與后續(xù)的注意力計(jì)算。

這意味著什么？

意味著語言建模（LM）loss 的梯度，根本傳不到 summary 和選擇分?jǐn)?shù)上。

Top-K 選擇是一個(gè)離散不可導(dǎo)的操作。打分→排序→選 ID，這條鏈路梯度無法反傳到 summary，LM loss 無法告訴 summary："你這次把重要的 chunk 排低了，下次該調(diào)高一點(diǎn)。"

于是 summary 的學(xué)習(xí)變成了 "盲訓(xùn)"—— 它學(xué)不會去抑制無關(guān) chunk、抬高對預(yù)測真正有用的 chunk。選擇過程沒有被端到端優(yōu)化，再有表達(dá)力的 summary 也無濟(jì)于事。

這就引出了兩個(gè) research question：

RQ1. 數(shù)學(xué)表達(dá)能力足夠的 chunk 重要性估計(jì).

RQ2. chunk summary 必須能跟著 LM loss 端到端訓(xùn)練.

把這兩點(diǎn)同時(shí)做到，才算把稀疏注意力"做對"

四、思路：HiLS-Attention ——

把 chunk 選擇變成可微分的 "分層 softmax"

騰訊混元提出HiLS-Attention（Hierarchical Landmark Sparse Attention，分層地標(biāo)稀疏注意力）。它的核心，是把上面兩個(gè)訴求拆成兩個(gè)問題逐一攻破。

RQ1：用 "一階泰勒展開" 構(gòu)造表征能力足夠的 chunk 算分函數(shù)

一個(gè)直覺的想法是對 LogSumExp 進(jìn)行一階泰勒展開，觀察數(shù)學(xué)性質(zhì)。研究團(tuán)隊(duì)發(fā)現(xiàn)，chunk 的對數(shù)重要性可以被近似成一個(gè)非常優(yōu)雅的形式：

它由兩部分組成：

RQ2：如何讓 chunk summary 跟著 LM loss 端到端訓(xùn)練？

光有好 summary 還不夠 —— 要解決那個(gè)致命斷點(diǎn)：讓梯度真正流到 summary 上。

HiLS 的做法是把注意力權(quán)重分層因式分解成兩級 softmax：

圖：先按 summary 算每個(gè) chunk 應(yīng)該分?jǐn)偠嗌亠灒儆?chunk 內(nèi) token 繼續(xù)分餅

• chunk 內(nèi)（intra-chunk）softmax：在每個(gè)被選中的 chunk 內(nèi)部，決定 token 之間的相對權(quán)重；

斷點(diǎn)被打通了。chunk 選擇第一次成為了在 LM 目標(biāo)下端到端可學(xué)習(xí)的過程，而且訓(xùn)練和推理全程都是真稀疏（native sparse training）。

一個(gè)反直覺的彩蛋：它不只是 "模仿" 全注意力，還更準(zhǔn)

研究團(tuán)隊(duì)最初是想讓 HiLS 去逼近 "全注意力誘導(dǎo)的 chunk 選擇"（即樸素 BSA）。但實(shí)驗(yàn)給了驚喜：

HiLS 不僅追平了樸素 BSA，還在長上下文檢索上反超了全注意力本身。

原因可能在于壓縮本身能去噪。全注意力有個(gè)固有毛病：只要一個(gè) token 的 logit 不是負(fù)無窮，它就會分到一點(diǎn)點(diǎn)注意力質(zhì)量。上下文越長，這些無關(guān) token 的微小噪聲越積越多，反而污染了檢索信號。而 HiLS 把多個(gè) key 壓縮成一個(gè) summary key 時(shí)，不對齊的噪聲相互抵消共享的語義信號被保留，于是檢索反而更干凈 —— 這正是它在變量追蹤（VT）這類多跳任務(wù)上能比全注意力高出多達(dá) 50% 的根源。

五、實(shí)驗(yàn)：從 345M 到 7B，全面驗(yàn)證

345M → 1.4B → 7B三個(gè)尺度上做了系統(tǒng)驗(yàn)證，結(jié)論高度一致：

• 短文本不掉點(diǎn)：345M 與 1.4B 從零訓(xùn)練時(shí)，HiLS 在各上下文長度、各訓(xùn)練階段的 PPL 與全注意力幾乎重合，8K 處持平甚至略低；
• 超長外推炸裂：僅用 8K 訓(xùn)練，外推到 4M（512 倍）仍保持 90%+ 大海撈針準(zhǔn)確率，遠(yuǎn)超全注意力；
• 低成本改造存量模型：把 OLMo3-7B 這類全注意力模型轉(zhuǎn)成 HiLS，只要續(xù)訓(xùn) 50B token 即可實(shí)現(xiàn)切換。短程任務(wù)不掉點(diǎn)，長序列任務(wù) LongBench 在 in-domain 長度甚至能超越全注意力基線，無縫繼承 HiLS 的外推能力，在 out-of-domain length 顯著碾壓 YaRN 等各類 baseline；
• 推理還更快：512K 上下文下，prefill 快 13.5×、單步 decode 快 15.7×。

稀疏注意力長期以來的 "效率 — 性能" 二選一困境，被第一次同時(shí)打破了。

六、寫在最后

回頭看這條邏輯鏈其實(shí)很清晰：

1. 稀疏注意力的瓶頸是chunk 選錯(cuò)

2. 選錯(cuò)的根源是mean /max logits 系統(tǒng)性失準(zhǔn)

3. 想用參數(shù)化 summary 補(bǔ)救，又卡在端到端反傳斷點(diǎn)上

4.HiLS 用泰勒線性化（提出足夠表達(dá)力的估分函數(shù)）+ 分層 softmax（把代理分?jǐn)?shù)送進(jìn)前向），一舉解決了表達(dá)力和可微分兩個(gè)問題。

HiLS 證明了：稀疏注意力可以同時(shí)提升效率和效果。效果提升的根源或許源自壓縮導(dǎo)致去噪，帶來更純凈的檢索表征。

這，才是把分層稀疏注意力 "Done Right" 的樣子。