宅中地 - 每日更新
宅中地 - 每日更新

贊助商廣告

X

南加州大學與芝加哥大學聯手破解AI記憶瓶頸:讓推理模型邊思考邊「忘記」,還不出錯

2026年06月09日 首頁 » 熱門科技

這項由南加州大學與芝加哥大學聯合開展的研究,於2026年6月發表於機器學習領域的預印本平台arXiv,論文編號為arXiv:2606.03928。研究提出了一種名為VASE(Value-aware Stochastic KV Cache Eviction,價值感知隨機鍵值緩存淘汰)的方法,專門針對當前最前沿的推理型大語言模型在長思考過程中面臨的內存瓶頸問題。

當你讓一個聰明的AI幫你解一道複雜的數學題時,它並不是一步就給出答案的。它會先在腦子裡默默地推導,一步一步地想,把中間過程全部記錄下來,最後才告訴你結果。這種"先想再答"的方式讓AI變得更準確,但也帶來了一個麻煩:它記錄的中間內容越來越多,占用的內存也越來越大,到最後甚至會把整個系統拖垮。

研究團隊發現了兩個此前幾乎被忽視的關鍵規律,並據此設計出一套既省內存又不影響準確率的解決方案,在六項不同推理任務上的表現甚至超過了當時最頂尖的同類方法。

一、AI的"工作記憶"為什麼會爆滿

要理解這個研究解決的問題,可以把推理型AI的工作過程想像成一個偵探在辦案。偵探在解案過程中會把所有的線索、推斷、筆記全部寫在一塊巨大的白板上,這樣他回頭翻看時就能找到之前想到的每一條推斷。AI也是如此——每生成一個新詞,它就需要把所有之前生成的內容"記"在一種叫做"KV緩存"(Key-Value Cache)的儲存結構里,每次回答問題時都要翻閱這塊"白板"。

問題在於,這塊白板的大小是有限的,而推理型AI思考的鏈條又極其漫長。研究團隊發現,即使是一道簡單的算術題,模型也可能洋洋灑灑地寫下超過900個詞才給出答案。對於難度更高的競賽數學題,模型平均要生成將近18000個詞。這塊白板很快就會被寫滿,內存和計算資源都會承受巨大壓力。

現有的解決思路分為兩大類。第一類叫做"選擇性注意",相當於白板上的內容全都保留著,但偵探每次只精選其中一部分來參考,不用全部看一遍——這樣可以節省時間,但白板本身還是越來越大,內存問題沒有從根本上解決。第二類叫做"淘汰機制",意思是當白板寫滿時,把偵探認為"不重要"的筆記擦掉,給新內容騰出地方——這樣白板永遠不會超出預設大小,內存占用是固定的,但問題是一旦擦錯了關鍵線索,案子就破不了了。

研究團隊聚焦的正是第二類方法,也就是淘汰機制。他們的核心問題是:為什麼現有的淘汰機制總是擦錯線索,導致AI推理能力大打折扣?

二、偵探最不該丟的兩類線索

研究團隊提取了AI在推理過程中產生的所有"記憶條目",仔細分析它們的特徵,從中找到了讓擦除操作出錯的根本原因。

KV緩存里儲存的內容可以拆分成兩個部分:一部分叫做"鍵"(Key),相當於偵探白板上每條筆記的"標題索引";另一部分叫做"值"(Value),相當於筆記的"正文內容"。現有的淘汰方法基本上只看標題索引,也就是只依據"這條筆記被翻查過多少次"來決定要不要擦掉它。但研究團隊發現,正文內容本身的重要性完全被忽視了。

團隊統計了大量"值"向量(即那些正文內容)的數值範圍,發現它們的分布極度不均勻:絕大多數值的數值範圍很窄,像一批普通的筆記,內容平平;但有極少數的值,數值範圍異常地大,像幾頁密密麻麻、充滿關鍵推斷的筆記。這種分布就像一群身高差不多的普通人里突然混進了幾個兩米五的巨人,極為醒目。

為了驗證這些"大範圍值"(large-magnitude value states)是否真的至關重要,研究團隊做了一個特意破壞實驗:故意把這些"巨人筆記"率先從緩存中擦掉,看看AI的表現會怎樣。結果令人震驚——AI在一個叫做GSM8K的數學問答測試集上的準確率從88%直線崩潰至14.3%,連隨機刪除筆記的效果(53.2%)都遠不如。

更直觀的是AI生成的內容:在這些關鍵條目被刪除之後,模型開始陷入無限循環,反覆重新審視題目,不斷自我懷疑,像一個偵探把所有關鍵推斷筆記都弄丟了,只剩下索引卻找不到對應內容,於是開始反覆問自己"這道題到底要求什麼?等等,我再想想……等等,我剛才說的對嗎?",卻永遠無法得出結論。

這個現象揭示了一個深層機制:AI注意力的輸出是所有"值"的加權求和,數值範圍越大的值,對輸出結果的影響就越大,因為無論權重怎麼分配,它們的絕對貢獻量都會顯著高於那些"普通筆記"。一旦這些高影響力的條目被刪除,AI的推理進展就會卡住,陷入一種自我強化的循環困境。

三、隨機性為什麼反而幫了大忙

除了保護關鍵筆記之外,研究團隊還發現了第二個被忽視的問題:現有的淘汰方法太"固執"了。

現有方法在決定保留哪些筆記時,總是按照固定的規則選擇得分最高的那些,把其餘的一律丟棄。這種方式就像偵探每次只保留自己最近最頻繁翻閱的那幾頁筆記,久而久之,留下來的內容越來越同質化——全是偵探反覆強調的那幾條推斷,而那些雖然不常被翻閱、但在某個關鍵時刻可能至關重要的線索,就這樣被永久丟失了。

研究團隊引入了"隨機採樣"的思路:保留筆記時不再總是按分數排名嚴格選取前幾名,而是根據每條筆記的重要性分數作為權重進行隨機抽籤。這就像偵探在挑選保留筆記時,給每條筆記貼上不同大小的標籤,重要的標籤大、抽中概率高,但並非百分之百;不太重要的標籤小,但也有一定機會被保留。

這種方式帶來的好處是,保留下來的筆記更加多樣化,覆蓋的資訊面更廣。在長達數萬詞的推理過程中,任何一條"看起來不那麼重要"的筆記,在之後的某個推導節點都可能突然變得關鍵。隨機性相當於給緩存裝上了"多樣性緩衝器",防止AI的記憶越來越片面。

實驗數據非常直觀:對於一種叫做SnapKV的基礎淘汰方法,僅僅把它的"嚴格選最高分"改成"按分數權重隨機抽籤",準確率就從64.3%提升到了70.9%,而什麼都沒改動,只是引入了隨機性。對於另一種叫做CurDKV的方法,通過在每次淘汰時重新隨機生成評分矩陣,準確率從78.6%躍升至87.6%,提升幅度高達9個百分點。

四、VASE是怎麼把兩個發現合在一起的

有了這兩個發現,研究團隊設計了VASE方案,把"保護大範圍值"和"引入隨機性"結合成一套完整的操作流程。

VASE的整體框架採用了一種叫做"周期性淘汰"的工作方式。白板(緩存)有一個固定容量上限,旁邊還留著一個小區域作為"緩衝區",專門放最近剛寫下的筆記。每當緩衝區寫滿,就觸發一次淘汰操作:從整個白板上選出要清除的內容,然後清空緩衝區,重新開始。

具體來說,團隊提出了兩種VASE變體。第一種叫做VASE-AttnV,它在整個緩存預算里專門劃出一塊固定區域——類似於白板上永遠不允許擦除的"重要檔案區"——專門保存那些"值的範圍最大"的筆記。這塊區域大小約占總預算的四分之一。剩餘的空間則用隨機加權抽籤的方式決定保留哪些普通筆記,抽籤的權重來自於每條筆記被最近幾次推理步驟查閱的頻率。這樣一來,高價值筆記穩穩噹噹地待在檔案區,其餘位置則通過隨機性保證覆蓋面足夠廣。

第二種叫做VASE-DKV,它的出發點是另一種基礎方法CurDKV。CurDKV已經考慮到了"值"的資訊,它通過一種叫做隨機投影的數學技巧來評估每條筆記的重要性——大致上,範圍更大的值會在投影后得到更高的分數,因此CurDKV天然就有一定的"值感知"能力。但它的問題是每次淘汰時都用同一個投影矩陣評分,某些筆記一旦第一次被判定為低分,就會在所有後續輪次中永遠得低分,最終被淘汰出局。VASE-DKV的改進非常簡潔:每次觸發淘汰時,重新隨機生成一個新的投影矩陣。這樣每條筆記每次都面臨獨立的"重新評判",不會因為某次偶然的低分而被永久放棄。

五、實驗結果:新方法到底強在哪裡

研究團隊在六項推理任務上進行了全面測試,包括頂級競賽數學題(AIME25、AIME26、HMMT25)、研究生水平的科學問答(GPQA-Diamond)、數學題庫(MATH)以及代碼生成(LiveCodeBench-v6),分別在Qwen3-4B和Qwen3-14B兩個規模的模型上運行,測試時將緩存壓縮到原來的四分之一大小。

VASE-AttnV在4B模型上的平均準確率達到59.09%,在14B模型上達到65.81%。與之相比,之前最強的淘汰方法R-KV在4B和14B模型上分別只有54.69%和60.90%,差距分別達到4.4和4.9個百分點。VASE-DKV的表現同樣突出,在4B模型上比CurDKV提高了7.7個百分點,在14B模型上提高了9.2個百分點。

更值得關注的是,VASE方法達到的平均準確率甚至超越了一種叫做SeerAttention-R的"選擇性注意"方法——後者的內存占用會隨著推理時間的增長而線性增加,VASE則始終維持固定的內存上限。換句話說,VASE用更省內存的方案,取得了比"不省內存"方案還要好的效果。

在極度壓縮的情況下(將緩存壓縮到原來約七分之一或八分之一),VASE相對於基線方法的優勢最為明顯;隨著緩存預算增大,各方法的差距逐漸縮小,最終都趨近於完整模型的準確率。這說明VASE在資源最匱乏的場景下發揮作用最為顯著,正是最需要幫助的地方,幫助最大。

代碼生成任務上有一個有趣的反轉:SeerAttention-R在LiveCodeBench上只有45.3%,遠低於R-KV的62.6%和VASE-AttnV的63.5%。研究團隊分析認為,SeerAttention-R的注意力預測門控需要經過專門訓練,面對代碼生成這種與訓練數據分布差異較大的場景,泛化能力弱;而無需訓練的淘汰方法則不存在這個問題。

六、這個發現還能用在哪裡

研究團隊的發現並不局限於"淘汰哪些緩存條目"這一個問題,它實際上揭示了"值的範圍"(Range)這個數值特徵的深層意義,而這個意義同樣適用於另一種完全不同的壓縮技術——量化(Quantization)。

量化是什麼?簡單來說,就是把一個精度很高的數字(比如小數點後很多位)壓縮成一個精度更低的整數來儲存,以此節省內存。就像把一張高清照片壓縮成低解析度版本——大多數情況下還認得出來,但如果原圖里有很精細的細節,壓縮後就會丟失。

研究團隊發現,"值的範圍"越大,量化時損失的資訊就越多。原因是量化的步長(也就是數字分級的精細程度)直接由範圍決定:範圍小,步長小,量化精細,誤差小;範圍大,步長大,量化粗糙,誤差大。團隊用實際數據驗證了這個關係,發現在不同的壓縮精度和分組方式下,"值的範圍"與量化誤差的相關係數都在0.83到0.92之間,高度正相關。

這意味著,凡是"值範圍大"的緩存條目,無論是在淘汰問題中還是在量化問題中,都應該被優先保護、優先分配更高的精度資源。研究團隊提出了一個展望:未來可以設計一種混合精度方案,把大範圍值存放在高精度緩存區,其餘條目壓縮到低精度,從而在同樣的內存預算下進一步提升保真度。

七、實際運行的速度和內存表現

方法再好,如果實際跑起來速度很慢或者占內存很多,也沒有實用價值。研究團隊專門在一塊A100 80GB顯卡上對Qwen3-14B進行了實際吞吐量和內存占用的測試,以確保理論上的優勢能落地成真實的工程收益。

測試結果顯示,在生成1.6萬詞並使用2048條目緩存預算的條件下,VASE-DKV的生成速度達到每秒411個詞,而完整模型(不壓縮緩存)只有133個詞每秒,速度提升了3.1倍。在生成3.2萬詞時,完整模型因內存不足直接崩潰(OOM),而所有淘汰方法都能正常運行。

內存方面,以使用4096條目緩存預算為例,完整模型占用68.7GB顯存,R-KV占用42.4GB,VASE-AttnV占用39.6GB,VASE-DKV只需38.9GB。如果刨去模型權重本身占用的約30GB不變內存,VASE-DKV的KV緩存實際壓縮比與理論預期的4倍非常接近,說明方法的內存控制非常精確。

此外,VASE-DKV比VASE-AttnV更快、更省內存,原因在於VASE-DKV在淘汰時不需要計算注意力分數,而VASE-AttnV需要;VASE-AttnV又比R-KV更快,因為R-KV每次淘汰時需要額外計算所有筆記之間的相似度來衡量"冗餘度",這一步計算量不小。

歸根結底,這項研究做的事情,就是找到了兩把被忽視已久的鑰匙——"保護高價值記憶"和"用隨機性帶來多樣性"——然後把它們插入了推理AI的記憶管理系統,讓AI在思考極長鏈條時,既不會因為記不住關鍵內容而陷入死循環,也不會因為記的內容太片面而偏離推理軌道。在壓縮到四分之一內存的條件下,這套方案的平均推理準確率不僅超過了其他所有淘汰方法,甚至與那些完整保留全部記憶(但內存持續增長)的方法打了個平手。

這對普通用戶意味著什麼?當你在手機或個人電腦上運行一個推理型AI助手時,有限的內存一直是制約模型能力的瓶頸。這項研究提供的方案,有望讓同等硬體條件下的AI能處理更長、更複雜的問題,而不會因內存不足而截斷思考或降低質量。有興趣深入了解技術細節的讀者,可以通過arXiv編號2606.03928查閱完整論文。

Q&A

Q1:KV緩存淘汰方法和KV緩存選擇方法有什麼區別?

A:KV緩存選擇方法把所有歷史記憶都保留在內存中,每次推理時只挑部分來用,內存占用會隨生成長度持續增長;KV緩存淘汰方法則是把"不重要"的記憶條目永久刪除,將內存控制在一個固定上限內,適合長時間連續推理,但刪錯了就會造成不可挽回的資訊損失。

Q2:VASE方法裡說的"值的範圍"(Range)到底是怎麼算的?

A:對於緩存中儲存的每一條記憶向量,"值的範圍"就是這個向量里數值最大的那個維度減去數值最小的那個維度得到的差值。範圍越大,說明這條記憶對最終輸出的影響越大,也越容易在量化壓縮時產生較大誤差,因此VASE會優先將這類條目保護起來,不允許被淘汰。

Q3:VASE方法為什麼在代碼生成任務上比SeerAttention-R表現更好?

A:SeerAttention-R依賴一個經過專門訓練的"注意力門控"模組來預測哪些緩存條目值得關注,這個模組對訓練數據分布以外的任務泛化能力有限;而VASE是完全無需訓練的方法,純粹基於當前數據的統計特徵做決策,因此在代碼生成這類與數學推理分布差異較大的任務上,反而展現出更強的適應性。

宅中地 - Facebook 分享 宅中地 - Twitter 分享 宅中地 - Whatsapp 分享 宅中地 - Line 分享
相關內容
Copyright ©2026 | 服務條款 | DMCA | 聯絡我們
宅中地 - 每日更新