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

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

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

現(xiàn)有架構為何難以兼顧效率與記憶？

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

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

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

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

核心發(fā)現(xiàn)：Mamba 的"遺忘"機制數(shù)學解析

不同Mamba變體的ETMF和ECLMF比較

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

解法創(chuàng)新：MemMamba 的"做筆記"記憶增強機制

MemMamba整體工作流程

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

MemMamba Block Layer工作流程

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

核心機制消融實驗結果

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

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

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

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

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

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

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

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

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

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

不同噪聲文檔數(shù)量下各模型的性能

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

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

MemMamba 的魯棒性

MemMamba 在關鍵超參數(shù)方面表現(xiàn)出強大魯棒性。

狀態(tài)池大小和窗口大小對PPL的影響

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

不同池化函數(shù)對建模質(zhì)量的影響

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

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

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

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

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

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

總結：MemMamba 的啟示

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

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

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

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

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