小米給KV Cache減負80%！MiMo團隊推出混合稀疏注意力架構

小米Mimo大模型團隊投稿
量子位 | 公眾號 QbitAI

小米MiMo大模型團隊，加入AI拜年戰(zhàn)場——

推出HySparse，一種面向Agent時代的混合稀疏注意力架構

HySparse創(chuàng)新使用極少的全注意力（Full Attention）層提供“token選擇+KV Cache”，其余稀疏注意力（Sparse Attention）層直接復用這些信息，實現(xiàn)高效精準的長上下文建模。

在總共49層的80B-A3BMoE模型實驗中，僅保留5層Full Attention仍能保持甚至提升模型能力，同時顯著降低KVCache存儲與計算開銷，實現(xiàn)效果與效率的兼顧，展示出混合稀疏注意力在超長上下文建模中的巨大潛力。

HySparse的設計靈感來源于學術界已有研究工作的經(jīng)驗和觀察之上。

一部分是顯著token在相鄰層之間相對穩(wěn)定。

已有工作如TidalDecode等，觀察到連續(xù)層的 “重要 token” 會高度重合，因此可以在某層識別重要token并在后續(xù)層復用。

HySparse將這一觀察提升用于模型結構設計并直接訓練。

還有部分受啟發(fā)于跨層KV Cache共享能顯著省顯存且不顯著傷性能，YOCO、Gemma3n等架構層面工作已經(jīng)證明了跨層共享KV的可行性。

HySparse將共享直接落在“Full Attention層 → 其后Sparse Attention層” 的hybrid block內(nèi)共享上。

背景：Sparse Attention的 “兩朵烏云”

多數(shù)Sparse Attention方法都遵循同一基本范式：先選擇，再計算。

在真正計算注意力之前，先用一個“選擇器” 決定哪些重要token是值得注意的，只在這些位置上進行注意力計算。

這一范式降低了計算量，但始終繞不開兩個根本問題。

一個問題是重要 Token 的選擇依賴代理（Proxy）

Sparse Attention的核心在于 “選哪些重要token”。

但現(xiàn)實中，這一選擇往往依賴proxy信號：固定模式、啟發(fā)式規(guī)則、近似估計，或額外的輕量化選擇模塊。

這些proxy本質(zhì)上都是對真實注意力分布的近似，無法保證能完全精準地識別最重要的token；在長上下文、持續(xù)生成的場景下，這種近似誤差還會不可避免地累積。

即便是可訓練的Sparse Attention方法，往往也只是將人工設計的選擇規(guī)則替換為可學習的輕量化選擇代理模塊，通過訓練來對齊選擇行為。

這在一定程度上緩解了選擇誤差，但同時顯著增加了訓練復雜度，難以從根本上擺脫這一“代理瓶頸”。

第二個問題是計算量降低，但KV Cache存儲未減

目前，主流的動態(tài)Sparse Attention注意力方法主要減少的是計算開銷。

這類做法不再對所有token執(zhí)行注意力計算；但為了避免在生成過程中不可逆地錯誤丟棄KV Cache（因為token的重要性會隨生成動態(tài)變化），推理階段通常仍需保留全量KV Cache。

帶來的結果就是計算量確實下降了，但顯存與帶寬這一主要瓶頸依然存在。

HySparse：把 “選擇” 和 “緩存”交給Full Attention

HySparse做的事是把 “選擇” 和 “緩存” 這兩件對Sparse Attention來說最難的事，交給Full Attention來做

為此，HySparse采用hybrid block結構。

每個hybrid block由1層Full Attention+N層Sparse Attention組成。

Hybrid block內(nèi)部的Sparse Attention層并不再獨立做token選擇和維護全量KV，而是直接復用前置Full Attention層產(chǎn)生的重要token索引和 KV Cache。

這背后的動機是Full Attention在完成自身計算的同時，已經(jīng)生成了KV Cache，并且計算出了最準確的token重要性信息，自然可以供后續(xù)N個Sparse Attention層直接復用。

這一設計表面上非常簡潔，卻巧妙地同時解決了上述 Sparse Attention 的兩大核心問題，

• 選擇不再依賴proxy；
• Sparse層不引入額外KV Cache開銷。

同時，HySparse稀疏層不是單一路徑，稀疏層內(nèi)部也做了一次 “全局稀疏 + 局部窗口” 的混合結構。

HySparse 的每一層 Sparse Attention 包含兩條分支：

• 塊級稀疏注意力分支（全局）：在TopK索引對應的共享KV Cache上進行全局Sparse Attention計算；
• 滑動窗口注意力分支（局部）：維護一個很小的本地窗口KV Cache（默認窗口大小為128）以保證局部建模能力。

兩分支輸出通過輕量門控（sigmoid gate）進行融合。

直觀地看，HySparse并不是用Sparse Attention取代Full Attention ，而是將全局信息通路拆解為 “少量昂貴但可靠的全注意力”+“多次廉價而高效的全局稀疏檢索與局部建模”。

實驗結果：整體更優(yōu)

為了驗證HySparse的有效性，研究團隊將其與兩類主流架構進行對比：

• Full-Attn：所有層均為Full Attention；
• Hybrid SWA：按比例混合Full Attention和Sliding Window Attention；
• HySparse：采用與Hybrid SWA相同的Full Attention比例，但將滑動窗口層升級為“全局稀疏 + 局部窗口”的Sparse Attention。

研究人員分別在7B的Dense模型和80B的MoE模型上進行了實驗。

結果顯示，7B模型共有36層，其中9層是Full Attention；80B MoE模型共有49層，只有5層是Full Attention，這為HySparse帶來了接近10×的KV Cache存儲降低。

在多項通用、數(shù)學、代碼和中文評測中，HySparse在7B Dense和80B MoE兩種規(guī)模均帶來穩(wěn)定提升。

特別是在80B MoE的激進混合比例下，HySparse通過全局稀疏檢索通路有效保持長距離信息訪問，維持了全局推理能力，還在部分任務上超越了全注意力基線，充分體現(xiàn)了全局稀疏通路的重要作用。

RULER長文測試表明，HySparse即便將Full Attention層壓到極少，也能穩(wěn)定保持長距離關鍵信息訪問，充分展現(xiàn)了其混合稀疏結構的優(yōu)勢。

小米MiMo表示，計劃在更大規(guī)模模型上進一步驗證HySparse的極限和潛力，并持續(xù)探索降低Full Attention層數(shù)量的可能性，讓超長上下文更高效。

arXiv鏈接：
https://arxiv.org/abs/2602.03560