近年來，大語言模型(LLM)領(lǐng)域出現(xiàn)了一個引人注目的發(fā)展趨勢：研究人員開始轉(zhuǎn)向具有亞二次復(fù)雜度的循環(huán)模型架構(gòu)，如Mamba、RWKV和RecurrentGemma等。這些模型在處理長上下文時展現(xiàn)出顯著的效率優(yōu)勢，相比傳統(tǒng)Transformer架構(gòu)的二次方復(fù)雜度，它們能夠以更低的計算成本處理更長的序列。然而，這些循環(huán)模型存在一個關(guān)鍵的局限性：固定大小的循環(huán)記憶容量。

來自特拉維夫大學、IBM研究院和MIT CSAIL的研究團隊在最新論文《溢出預(yù)防增強了長上下文循環(huán)LLM》中，深入探討了這一問題，并提出了一種簡單而有效的解決方案，顯著提升了循環(huán)LLM在長上下文任務(wù)中的表現(xiàn)。該研究不僅解決了實際問題，還對循環(huán)模型如何處理長距離依賴關(guān)系提出了深刻的理論思考。

記憶溢出問題：循環(huán)模型的隱藏瓶頸

研究團隊首先通過聯(lián)想回憶(Associative Recall)任務(wù)揭示了循環(huán)LLM的記憶溢出現(xiàn)象。在這項任務(wù)中，模型需要從上下文中檢索特定的鍵值對。實驗結(jié)果顯示，即使是擁有較大隱藏狀態(tài)的模型(如Falcon-Mamba-Inst-7B，具有4096個通道，每個通道狀態(tài)大小為16)，當上下文中的信息量增加時，其檢索準確率會顯著下降。

這一現(xiàn)象表明，循環(huán)模型的記憶容量是有限的，當輸入信息超過其容量時，模型會出現(xiàn)"溢出"，導(dǎo)致性能急劇下降。更令人擔憂的是，這種溢出現(xiàn)象在相對較短的上下文(僅1200個標記)中就已經(jīng)出現(xiàn)，這對于長上下文應(yīng)用場景構(gòu)成了嚴峻挑戰(zhàn)。

為了進一步驗證這一發(fā)現(xiàn)，研究人員在受控環(huán)境中訓(xùn)練了多個2層Mamba模型，測試不同通道維度和狀態(tài)維度組合對記憶容量的影響。結(jié)果表明，雖然增加隱藏狀態(tài)大小可以緩解溢出問題，但無法完全解決它。即使是經(jīng)過專門訓(xùn)練以檢索128個事實的模型，其實際容量仍然有限，無法充分利用所有訓(xùn)練數(shù)據(jù)。

OPRM：一種簡單而有效的解決方案

基于對記憶溢出問題的深入理解，研究團隊提出了OPRM(Overflow Prevention for Recurrent Models)方法。這種方法的核心思想非常簡單：將長上下文分割成多個小塊，確保每個塊的信息量不超過模型的記憶容量限制，然后只處理最相關(guān)的塊來生成答案。

OPRM的工作流程分為兩個階段：

1. 推測性預(yù)填充(Speculative Prefill)：將上下文C分割成b個相同長度L的塊，并為每個塊構(gòu)建單獨的提示，保持原始結(jié)構(gòu)：Xi=[P,Ci,S]，其中P是前綴，S是包含查詢Q的后綴。這些提示被并行處理，為每個塊計算輸出分布。
2. 選擇性解碼(Selective Decoding)：基于選擇標準(如熵或概率)，從最具信息量的提示Xj中選擇狀態(tài)Hj和標記Aj,0，然后執(zhí)行自回歸解碼。

為了處理某些塊可能不包含相關(guān)信息的情況，研究者還引入了"IDK過濾器"(I Don't Know Filter)。該過濾器會丟棄那些預(yù)測"Error"標記的塊，確保模型只從包含有用信息的塊中生成回答。

實驗結(jié)果：簡單方法帶來驚人提升

研究團隊在多個基準測試上評估了OPRM方法，結(jié)果令人印象深刻：

1. 聯(lián)想回憶任務(wù)：使用OPRM后，F(xiàn)alcon-Mamba-Instruct-7B的準確率不再隨上下文中信息量的增加而下降，實際上解決了這一任務(wù)。
2. LongBench基準：OPRM顯著提升了多個循環(huán)LLM的性能，包括Falcon3-Mamba-Inst-7B(提升14%)、Falcon-Mamba-Inst-7B(提升28%)、RecurrentGemma-IT-9B(提升50%)和RWKV6-Finch-7B(提升51%)。隨著上下文長度增加，OPRM的優(yōu)勢變得更加明顯。
3. LongBench v2基準：在這個更具挑戰(zhàn)性的基準測試中，F(xiàn)alcon3-Mamba-Inst-7B配合OPRM達到了30.8分的成績，為同等規(guī)模模型設(shè)立了新的最高水平，甚至超過了同等大小的Transformer模型。
4. 上下文擴展：OPRM還能自然地擴展模型的可用上下文長度，無需額外訓(xùn)練。在文檔檢索任務(wù)中，OPRM使Mamba-130M模型能夠處理比訓(xùn)練時長256倍的序列。

這些結(jié)果表明，OPRM不僅能有效緩解記憶溢出問題，還能顯著提升循環(huán)LLM在各種長上下文任務(wù)中的表現(xiàn)。

理論意義：循環(huán)模型真的能捕獲長距離依賴關(guān)系嗎？

除了實際應(yīng)用價值，這項研究還引發(fā)了一個深刻的理論問題：循環(huán)模型是否真正能夠捕獲輸入中相距甚遠部分之間的長距離依賴關(guān)系？

令人驚訝的是，OPRM的單塊策略在各種長上下文任務(wù)中都表現(xiàn)出色，甚至在那些理論上需要跨上下文關(guān)系的任務(wù)中也是如此。這一發(fā)現(xiàn)質(zhì)疑了循環(huán)LLM處理長距離依賴關(guān)系的能力，暗示它們可能主要依賴于局部上下文信息，而非全局關(guān)系。

OPRM的其他優(yōu)勢

除了提升性能外，OPRM還具有多項實用優(yōu)勢：

1. 效率提升：通過并行處理多個較小的塊，OPRM降低了計算復(fù)雜度，從O(Lb·log(Lb))降低到O(b·L·log(L))。
2. 內(nèi)存效率：盡管使用了更多狀態(tài)，OPRM的內(nèi)存使用量僅略高于基線模型，因為單個狀態(tài)占用的內(nèi)存遠小于模型權(quán)重。
3. 靈活的記憶-召回權(quán)衡：塊大小L作為唯一超參數(shù)，提供了一種簡單的方式來平衡記憶約束和召回性能。
4. 與RAG兼容：OPRM自然適合基于檢索增強生成(RAG)的設(shè)置，可作為各種任務(wù)和實際應(yīng)用的即插即用組件。

局限性與未來方向

盡管OPRM取得了顯著成功，但研究團隊也承認其存在一些局限性：

1. 缺乏跨塊處理能力，無法充分利用全局上下文依賴。
2. 作為一種無訓(xùn)練方法，OPRM依賴于預(yù)訓(xùn)練模型的能力，某些組件(如IDK過濾器)可能通過額外微調(diào)獲得更好效果。

未來的研究方向可能包括開發(fā)更復(fù)雜的塊間信息整合機制，以及探索如何通過特定訓(xùn)練進一步增強OPRM的效果。

結(jié)論：簡單而強大的解決方案

這項研究表明，通過理解循環(huán)LLM的記憶溢出現(xiàn)象并采用簡單的基于塊的推理策略，可以顯著提升這些模型在長上下文任務(wù)中的表現(xiàn)。OPRM不僅在實際應(yīng)用中證明了其價值，還為理解循環(huán)模型如何處理長距離依賴關(guān)系提供了新的視角。

最令人驚訝的是，這種簡單的方法能夠使循環(huán)LLM在長上下文任務(wù)上匹配甚至超越同等規(guī)模的Transformer模型，同時保持亞二次復(fù)雜度的效率優(yōu)勢。這一發(fā)現(xiàn)對于開發(fā)更高效、更強大的長上下文語言模型具有重要意義。

GitHub：???https://github.com/assafbk/OPRM??

論文：???https://arxiv.org/abs/2505.07793??

本文轉(zhuǎn)載自?????頓數(shù)AI???，作者：小頓