小米Mimo大模型團(tuán)隊(duì)投稿
量子位 | 公眾號(hào) QbitAI

小米MiMo大模型團(tuán)隊(duì)，加入AI拜年戰(zhàn)場(chǎng)——

推出HySparse，一種面向Agent時(shí)代的混合稀疏注意力架構(gòu)

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

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

HySparse的設(shè)計(jì)靈感來源于學(xué)術(shù)界已有研究工作的經(jīng)驗(yàn)和觀察之上。

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

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

HySparse將這一觀察提升用于模型結(jié)構(gòu)設(shè)計(jì)并直接訓(xùn)練。

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

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

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

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

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

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

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

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

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

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

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

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

第二個(gè)問題是計(jì)算量降低，但KV Cache存儲(chǔ)未減

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

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

帶來的結(jié)果就是計(jì)算量確實(shí)下降了，但顯存與帶寬這一主要瓶頸依然存在。

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

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

為此，HySparse采用hybrid block結(jié)構(gòu)。

每個(gè)hybrid block由1層Full Attention+N層Sparse Attention組成。

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

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

這一設(shè)計(jì)表面上非常簡(jiǎn)潔，卻巧妙地同時(shí)解決了上述 Sparse Attention 的兩大核心問題，

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

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

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

• 塊級(jí)稀疏注意力分支（全局）：在TopK索引對(duì)應(yīng)的共享KV Cache上進(jìn)行全局Sparse Attention計(jì)算；
• 滑動(dòng)窗口注意力分支（局部）：維護(hù)一個(gè)很小的本地窗口KV Cache（默認(rèn)窗口大小為128）以保證局部建模能力。

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

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

實(shí)驗(yàn)結(jié)果：整體更優(yōu)

為了驗(yàn)證HySparse的有效性，研究團(tuán)隊(duì)將其與兩類主流架構(gòu)進(jìn)行對(duì)比：

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

研究人員分別在7B的Dense模型和80B的MoE模型上進(jìn)行了實(shí)驗(yàn)。

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

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

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

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

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

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