循環LLM的突破：基于塊的推理如何解決記憶溢出問題
近年來，大語言模型(LLM)領域出現了一個引人注目的發展趨勢：研究人員開始轉向具有亞二次復雜度的循環模型架構，如Mamba、RWKV和RecurrentGemma等。這些模型在處理長上下文時展現出顯著的效率優勢，相比傳統Transformer架構的二次方復雜度，它們能夠以更低的計算成本處理更長的序列。然而，這些循環模型存在一個關鍵的局限性：固定大小的循環記憶容量。
來自特拉維夫大學、IBM研究院和MIT CSAIL的研究團隊在最新論文《溢出預防增強了長上下文循環LLM》中，深入探討了這一問題，并提出了一種簡單而有效的解決方案，顯著提升了循環LLM在長上下文任務中的表現。該研究不僅解決了實際問題，還對循環模型如何處理長距離依賴關系提出了深刻的理論思考。
記憶溢出問題：循環模型的隱藏瓶頸
研究團隊首先通過聯想回憶(Associative Recall)任務揭示了循環LLM的記憶溢出現象。在這項任務中，模型需要從上下文中檢索特定的鍵值對。實驗結果顯示，即使是擁有較大隱藏狀態的模型(如Falcon-Mamba-Inst-7B，具有4096個通道，每個通道狀態大小為16)，當上下文中的信息量增加時，其檢索準確率會顯著下降。
這一現象表明，循環模型的記憶容量是有限的，當輸入信息超過其容量時，模型會出現"溢出"，導致性能急劇下降。更令人擔憂的是，這種溢出現象在相對較短的上下文(僅1200個標記)中就已經出現，這對于長上下文應用場景構成了嚴峻挑戰。
為了進一步驗證這一發現，研究人員在受控環境中訓練了多個2層Mamba模型，測試不同通道維度和狀態維度組合對記憶容量的影響。結果表明，雖然增加隱藏狀態大小可以緩解溢出問題，但無法完全解決它。即使是經過專門訓練以檢索128個事實的模型，其實際容量仍然有限，無法充分利用所有訓練數據。
OPRM：一種簡單而有效的解決方案
基于對記憶溢出問題的深入理解，研究團隊提出了OPRM(Overflow Prevention for Recurrent Models)方法。這種方法的核心思想非常簡單：將長上下文分割成多個小塊，確保每個塊的信息量不超過模型的記憶容量限制，然后只處理最相關的塊來生成答案。
OPRM的工作流程分為兩個階段：
- 推測性預填充(Speculative Prefill)：將上下文C分割成b個相同長度L的塊，并為每個塊構建單獨的提示，保持原始結構：Xi=[P,Ci,S]，其中P是前綴，S是包含查詢Q的后綴。這些提示被并行處理，為每個塊計算輸出分布。
- 選擇性解碼(Selective Decoding)：基于選擇標準(如熵或概率)，從最具信息量的提示Xj中選擇狀態Hj和標記Aj,0，然后執行自回歸解碼。
為了處理某些塊可能不包含相關信息的情況，研究者還引入了"IDK過濾器"(I Don't Know Filter)。該過濾器會丟棄那些預測"Error"標記的塊，確保模型只從包含有用信息的塊中生成回答。
實驗結果：簡單方法帶來驚人提升
研究團隊在多個基準測試上評估了OPRM方法，結果令人印象深刻：
- 聯想回憶任務：使用OPRM后，Falcon-Mamba-Instruct-7B的準確率不再隨上下文中信息量的增加而下降，實際上解決了這一任務。
- LongBench基準：OPRM顯著提升了多個循環LLM的性能，包括Falcon3-Mamba-Inst-7B(提升14%)、Falcon-Mamba-Inst-7B(提升28%)、RecurrentGemma-IT-9B(提升50%)和RWKV6-Finch-7B(提升51%)。隨著上下文長度增加，OPRM的優勢變得更加明顯。
- LongBench v2基準：在這個更具挑戰性的基準測試中，Falcon3-Mamba-Inst-7B配合OPRM達到了30.8分的成績，為同等規模模型設立了新的最高水平，甚至超過了同等大小的Transformer模型。
- 上下文擴展：OPRM還能自然地擴展模型的可用上下文長度，無需額外訓練。在文檔檢索任務中，OPRM使Mamba-130M模型能夠處理比訓練時長256倍的序列。
這些結果表明，OPRM不僅能有效緩解記憶溢出問題，還能顯著提升循環LLM在各種長上下文任務中的表現。
理論意義：循環模型真的能捕獲長距離依賴關系嗎？
除了實際應用價值，這項研究還引發了一個深刻的理論問題：循環模型是否真正能夠捕獲輸入中相距甚遠部分之間的長距離依賴關系？
令人驚訝的是，OPRM的單塊策略在各種長上下文任務中都表現出色，甚至在那些理論上需要跨上下文關系的任務中也是如此。這一發現質疑了循環LLM處理長距離依賴關系的能力，暗示它們可能主要依賴于局部上下文信息，而非全局關系。
OPRM的其他優勢
除了提升性能外，OPRM還具有多項實用優勢：
- 效率提升：通過并行處理多個較小的塊，OPRM降低了計算復雜度，從O(Lb·log(Lb))降低到O(b·L·log(L))。
- 內存效率：盡管使用了更多狀態，OPRM的內存使用量僅略高于基線模型，因為單個狀態占用的內存遠小于模型權重。
- 靈活的記憶-召回權衡：塊大小L作為唯一超參數，提供了一種簡單的方式來平衡記憶約束和召回性能。
- 與RAG兼容：OPRM自然適合基于檢索增強生成(RAG)的設置，可作為各種任務和實際應用的即插即用組件。
局限性與未來方向
盡管OPRM取得了顯著成功，但研究團隊也承認其存在一些局限性：
- 缺乏跨塊處理能力，無法充分利用全局上下文依賴。
- 作為一種無訓練方法，OPRM依賴于預訓練模型的能力，某些組件(如IDK過濾器)可能通過額外微調獲得更好效果。
未來的研究方向可能包括開發更復雜的塊間信息整合機制，以及探索如何通過特定訓練進一步增強OPRM的效果。
結論：簡單而強大的解決方案
這項研究表明，通過理解循環LLM的記憶溢出現象并采用簡單的基于塊的推理策略，可以顯著提升這些模型在長上下文任務中的表現。OPRM不僅在實際應用中證明了其價值，還為理解循環模型如何處理長距離依賴關系提供了新的視角。
最令人驚訝的是，這種簡單的方法能夠使循環LLM在長上下文任務上匹配甚至超越同等規模的Transformer模型，同時保持亞二次復雜度的效率優勢。這一發現對于開發更高效、更強大的長上下文語言模型具有重要意義。
GitHub：???https://github.com/assafbk/OPRM??
論文：???https://arxiv.org/abs/2505.07793??
本文轉載自?????頓數AI???，作者：小頓
