大家好，我是肆〇柒。今天要和大家分享的是由中國人民大學統計學院、上海人工智能實驗室和上海財經大學聯合研究團隊在ICLR 2026會議投稿的突破性工作——MemMamba。這項研究首次系統揭示了Mamba模型的記憶衰減機制，通過受人類"做筆記"行為啟發的創新設計，成功解決了長序列建模中的關鍵記憶瓶頸問題，在400k tokens的極端長序列任務中仍能保持90%的準確率，為高效且強大的序列建模開辟了全新道路。

不同上下文長度下的Passkey檢索準確率

想象一下：當序列長度達到40萬token時，MemMamba仍能準確記住關鍵信息（90%準確率），而原始Mamba完全失效（0%準確率）。這一驚人差距揭示了長序列建模領域的一個根本性突破。在當今數據爆炸性增長的時代，長序列建模已成為自然語言處理和生物信息學等任務中的關鍵挑戰。長序列數據通常指跨越數千至數百萬時間步或 token 的連續序列，廣泛存在于現代機器學習應用中：從 NLP 中的書籍長度文檔建模（PG19 數據集包含約 1 億 token，由 1919 年左右出版的 Project Gutenberg 英文小說組成，平均長度達 69k tokens）、生物信息學中的 DNA 序列分析，到復雜多模態醫療記錄的處理。傳統架構在處理這類數據時面臨效率與記憶保留能力的根本性權衡：RNN 系列模型易受梯度消失/爆炸困擾，難以穩定處理長依賴；Transformer 雖能建模全局依賴，但其 O(n2) 復雜度使其在真正長上下文中效率低下。MemMamba 的突破性在于首次系統揭示了 Mamba 的記憶衰減機制，提出了一種可線性擴展的"做筆記"式記憶增強架構，實現了從單純追求"更長"到關注"如何記住關鍵信息"的范式轉變。

現有架構為何難以兼顧效率與記憶？

不同上下文長度下模型的困惑度(PPL)比較

傳統序列建模架構在處理超長序列時表現出顯著局限。循環神經網絡(RNN)及其變體(LSTM, GRU)本質上是順序處理的，受到梯度消失或爆炸問題的困擾，使其難以穩定處理長距離依賴。有研究顯示，這些模型在長序列上的表現受限于其內在的不穩定性。

Transformer 通過自注意力機制和全局 context 建模引入了范式轉變，但其序列長度的二次復雜度使其在真正長上下文中效率低下。對于序列長度 n、層數 L 和特征維度 d，Transformer 自注意力的時間復雜度為 TC = O(L·n2·d)。當序列長度達到 10? 時，n2 項導致約 101? 次操作，遠超當前硬件能力。實踐中常采用滑動窗口注意力(窗口大小 w=512)或稀疏注意力等近似方法，但這些截斷不可避免地會丟棄窗口外的信息，導致有效建模長度(EML)受限：EML ≤ w ? n，無法捕捉真正長距離依賴。

選擇性狀態空間模型(SSMs)，特別是 Mamba 架構，提供了一個有前景的替代方案。通過解耦序列長度與計算，Mamba 實現了線性時間復雜度 O(n) 和常數時間遞歸推理 O(1)，定位為長序列建模的潛在基礎。然而，盡管在計算效率上取得了飛躍，其記憶保真度在規模擴大時迅速下降。隨著序列長度增長，Mamba 及其后續版本在需要強記憶保留的任務中表現出急劇下降，如 5-shot MMLU 或長距離鍵值檢索任務。在 400k tokens 的 Passkey Retrieval 任務中，原始 Mamba 的準確率降至 0%，而 MemMamba 仍能保持 90% 的準確率，這一差距直觀展示了傳統架構在長距離記憶保留上的根本局限。

核心發現：Mamba 的"遺忘"機制數學解析

不同Mamba變體的ETMF和ECLMF比較

實證分析表明，Mamba 在兩項指標上均顯著劣于理想模型，揭示了其在超長序列中"忘記"關鍵信息的本質原因。隨著序列長度增加，Mamba 的記憶保真度急劇下降，而 MemMamba 通過創新機制有效緩解了這一問題。在 ETMF 指標上，MemMamba 的表現明顯優于原始 Mamba 和 DeciMamba，即使在長距離依賴情況下也能保持較高的語義保真度。論文指出，ETMF 和 ECLMF 兩個指標提供了互補視角：ETMF 反映了遞歸傳播后長距離 token 語義是否保持忠實，而 ECLMF 量化了跨層信息傳遞過程中的退化程度，共同揭示了 Mamba 在記憶衰減和外推能力方面的雙重挑戰。

解法創新：MemMamba 的"做筆記"記憶增強機制

MemMamba整體工作流程

如上圖，MemMamba 受人類閱讀長文檔時做筆記行為的啟發，將狀態空間模型重新構想為結構化記憶系統。想象你正在閱讀一本600頁的小說（約69k tokens）。隨著閱讀深入，你會自然地在關鍵情節處做筆記，幫助記憶重要線索。同樣，MemMamba 通過 Note Block 在關鍵 token 處"做筆記"，將摘要存入狀態池。當模型"忘記"某些信息時，它會像你翻閱筆記一樣，從狀態池中檢索相關信息進行恢復。

MemMamba Block Layer工作流程

這種雙閾值和稀疏跨層機制確保跨 token 補充在每一層發生，而跨層記憶交互則稀疏激活，在記憶保留和計算效率之間取得平衡。每個狀態摘要向量被壓縮到 64 維，狀態池大小固定為 50，這些設計選擇在廣泛參數范圍內表現穩定，證明了架構的魯棒性。

核心機制消融實驗結果

如上，消融實驗結果清晰展示了 MemMamba 各個組件的必要性。在 PG19 語言建模任務中，完整 MemMamba 的困惑度(PPL)為 17.35，而移除 Note Block 后 PPL 升至 19.99，移除跨 token 注意力后 PPL 升至 18.63，移除跨層注意力后 PPL 升至 18.19。這表明所有三個組件都對 MemMamba 的性能提升有重要貢獻，特別是 Note Block 機制在長序列建模中起到關鍵作用。實驗還表明，最大池化(max pooling)在保真度-效率平衡上表現最佳，優于均值池化(mean)、T-Max-Avg 和 S3Pool 等替代方案。

如何兼顧線性復雜度與強記憶能力？

MemMamba 在引入狀態摘要和跨層注意力的同時，仍保持線性時間和空間復雜度。具體而言，其計算成本隨序列長度 n 和隱藏維度 d 的縮放為 O(n·d)，與 Transformer 的 O(n2d) 形成鮮明對比。這通過將狀態維度 d_s 和注意力池大小 k 限制為常數來實現。

實驗驗證：MemMamba 的記憶能力突破

MemMamba 在多個長序列基準測試中實現了顯著改進。

不同上下文長度下模型的困惑度(PPL)比較

如上表，在 PG19 語言建模任務中，MemMamba 在 60k tokens 時的困惑度(PPL)僅為 17.35，而參數規模相似的 Mamba 和 DeciMamba 已完全崩潰(>100)。在 200M 參數規模下，MemMamba 的表現超越了 400M 參數的 Compressive Transformer，證明了其參數效率優勢。隨著上下文長度增加，MemMamba 保持低且穩定的 PPL：在 1.5k tokens 時 PPL 為 19.35，在 60k tokens 時 PPL 在 17.33–18.25 之間波動。

這相當于在閱讀整本小說的過程中，MemMamba 能夠像人類一樣"記住"關鍵情節，而其他模型在閱讀到一半時就已經"忘記"了故事主線。在醫療領域，患者可能有長達數萬token的電子健康記錄，包含多年診療歷史。MemMamba 能夠記住關鍵病史信息（如藥物過敏史），而不會在長序列處理中"忘記"，這直接關系到診斷準確性和患者安全。

不同上下文長度下的Passkey檢索準確率

如上，在 Passkey Retrieval 任務中，MemMamba 即使在輸入長度達幾十萬 token 時仍保持高檢索精度。當目標 token 位于預測點 200k tokens 之外時，MemMamba 仍能準確檢索關鍵信息，而 Mamba 和 Pythia 在此類長度下完全失效。在 400k tokens 時，MemMamba 保持 90% 的檢索準確率，而 DeciMamba 降至 60%，Mamba 完全失效。這一結果驗證了 MemMamba 在處理極端長距離依賴任務中的優勢，特別是在需要精確記憶特定位置信息的場景下。

不同噪聲文檔數量下各模型的性能

如上，在文檔檢索基準測試中，MemMamba 在簡單和詳細檢索設置下均取得領先性能。隨著噪聲文檔數量增加，Mamba 的性能急劇下降，DeciMamba 顯示部分改進但仍不穩定。相比之下，MemMamba 在高噪聲條件下持續保持較高分數，在 200 個噪聲文檔干擾下得分為 0.24，而 DeciMamba 僅為 0.12，Mamba 接近于 0，突顯其在跨文檔和跨領域推理任務中的優勢。這相當于在200本無關書籍的干擾下，MemMamba 仍能準確找到目標信息，而其他模型早已迷失方向。

效率方面，在相同硬件條件下，MemMamba 的端到端延遲僅為 Transformer 的 0.52 倍(即 48% 的速度提升)。盡管引入了增強建模能力的額外計算，MemMamba 通過緊湊表示和跨層/跨 token 注意力優化信息流，從而在超長序列上保持高計算效率。ETMF 和 ECLMF 評分顯示，MemMamba 在兩項指標上均顯著優于原始 Mamba 和 DeciMamba，驗證了其記憶保真度的提升。

MemMamba 的魯棒性

MemMamba 在關鍵超參數方面表現出強大魯棒性。

狀態池大小和窗口大小對PPL的影響

如上圖，在廣泛范圍內，窗口大小和狀態池大小對性能影響甚微，表明架構穩定性強。例如，當狀態池大小從 10 變化到 100 時，PPL 僅在 17.33-17.45 之間波動；當窗口大小從 16 變化到 128 時，PPL 僅在 17.33-17.42 之間波動。這一結果證明 MemMamba 的設計選擇具有廣泛的適用性，無需針對特定任務進行精細調參。

不同池化函數對建模質量的影響

如上圖，在池化函數選擇上，簡單的最大池化(max pooling)在保真度-效率平衡上表現最佳，優于均值池化(mean)、T-Max-Avg 和 S3Pool 等替代方案。在 60k tokens 長度下，最大池化的 PPL 為 17.35，均值池化為 17.61，T-Max-Avg 為 17.52，S3Pool 為 17.48。這一結果表明，最大池化能更有效地保留關鍵信息，同時保持計算效率。論文指出，這是因為最大池化能夠捕獲最顯著的特征，而這些特征往往對應于對任務最為關鍵的信息。

不同融合方法在上下文長度上的PPL比較

如上，融合方法的對比實驗表明，在短序列上各種方法差異較小，但在長上下文中，殘差融合和加權融合表現更優。在 60k tokens 時，加權融合的 PPL 為 17.36，殘差融合為 17.49，元素乘積為 17.49，1D 卷積為 17.41，門控融合為 18.01。這些結果驗證了 MemMamba 設計選擇的合理性，特別是加權融合方法在超長序列上的穩定性優勢。加權融合通過學習可調節的權重，能夠更靈活地平衡原始輸入和注意力補充之間的貢獻，從而在不同長度的序列上保持穩定性能。

ETMF 和 ECLMF 評分進一步驗證了 MemMamba 的記憶增強效果。MemMamba 在兩項指標上均顯著優于所有 Mamba 變體，盡管 DeciMamba 在極長距離跨層傳輸中顯示微弱優勢，但其不穩定性構成重大缺陷。這些量化指標為 MemMamba 的記憶增強機制提供了有力證據，證實了狀態摘要和跨層注意力在緩解長距離遺忘方面的有效性。

不同上下文長度下模型的困惑度(PPL)比較

如上，參數規模對比實驗顯示，MemMamba 在較小參數規模下就能達到與大規模模型相當的性能。在 1.5k tokens 時，200M 參數的 MemMamba PPL 為 19.35，而 400M 參數的 Compressive Transformer PPL 為 33.09，表明 MemMamba 具有顯著的參數效率優勢。在 60k tokens 時，MemMamba 仍保持 PPL 17.35，而 Compressive Transformer 因序列長度限制無法評估。這種參數效率使得 MemMamba 能夠在資源受限的環境中部署，同時保持高性能。

總結：MemMamba 的啟示

MemMamba 代表了長序列建模領域的一個重要突破，它通過將狀態空間模型增強為結構化記憶系統，彌合了可擴展性與長距離依賴建模之間的長期差距。通過將動態狀態摘要與輕量級跨層和跨 token 注意力相結合，MemMamba 為現有 SSM 中限制性的記憶衰減問題提供了原則性解決方案。MemMamba 的成功表明，長序列建模不應只追求"更長"，而應關注"如何記住關鍵信息"。將人類"做筆記"行為形式化為可學習的記憶機制，為高效記憶建模開辟了新范式。

MemMamba 的架構設計表明，通過合理設計的機制，可以有效平衡計算效率與記憶保留能力。其"做筆記"機制不僅解決了 Mamba 的記憶衰減問題，還為其他序列模型提供了可借鑒的思路。狀態池的固定容量設計和雙閾值觸發機制確保了模型在保持線性復雜度的同時，能夠有針對性地保留和檢索關鍵信息，避免了傳統注意力機制的二次復雜度問題。

在法律領域，合同文檔往往長達數萬詞，MemMamba 能準確記住關鍵條款的位置和內容，避免因"遺忘"而導致的法律風險。在生物信息學中，DNA序列分析需要處理超長序列，MemMamba 能夠保留關鍵基因信息，為疾病研究提供更準確的分析基礎。

盡管 MemMamba 在 400k tokens 上表現優異，但在極端長序列(>1M tokens)上的表現仍需進一步驗證。MemMamba 的出現為超長序列建模提供了新思路，有望在自然語言處理、生物信息學和多模態分析等需要處理超長序列的領域產生廣泛影響。未來工作將探索擴展至多模態設置、與檢索增強系統的集成，以及將 MemMamba 作為高效基礎模型，支持復雜現實世界任務中的高保真記憶。這些方向將進一步拓展 MemMamba 的應用范圍，推動長序列建模技術的發展。

MemMamba 的核心價值在于其對"記憶模式"的重新思考，它不僅是對 Mamba 的改進，更是對長序列建模范式的重新定義。通過將記憶視為一種需要主動管理和增強的能力，而非被動依賴模型規模的屬性，MemMamba 為高效且強大的序列建模提供的新思路。