機器之心報道

編輯：陳陳、笑寒

線性循環模型（如 Mamba）和線性注意力機制都具備這樣一個顯著優勢：它們能夠處理極長的序列，這一能力對長上下文推理任務至關重要。

事實上，這正是它們相較于 Transformer 的關鍵優勢 —— 后者受限于有限的上下文窗口，且在序列長度上的計算復雜度是二次的，成為性能瓶頸。

過去，循環模型面臨的主要問題是性能不足：在處理短序列時，它們的表現往往不如 Transformer。然而，隨著架構上的一系列突破，循環模型的性能已經顯著提升，在許多任務中已能與 Transformer 媲美，甚至已經被應用于多個工業場景中，如音頻建模和代碼補全等。

但近期的多項研究發現，循環模型仍存在一個關鍵短板：盡管它們在訓練長度范圍內表現良好，但在處理超出訓練長度的序列時，往往難以泛化，表現明顯下降。

事實也確實如此，舉例來說，下圖為 Mamba-2 檢查點在不同序列位置上的性能變化（以困惑度 Perplexity 衡量，數值越低代表性能越好）?？梢悦黠@看到，一旦序列位置超出了訓練上下文范圍，這些模型幾乎就變得毫無作用：它們無法實現長度泛化。

這就帶來一個問題：現有的循環模型在長序列上表現較差，而在短序列上相比 Transformer 也沒有明顯的效率優勢；換句話說，它們在兩個維度上似乎都顯得不夠理想。

那這是否意味著循環模型就毫無用處了呢？

當然不是！

本文，來自 CMU、 Cartesia AI 的研究者證明了通過簡單的訓練干預，循環模型是完全可以實現長度泛化的。只需 500 步后訓練（約占預訓練預算的 0.1%），就能讓模型在高達 256k 長度的序列上實現泛化！

因此，循環模型并不是存在根本性缺陷，而是擁有尚未被充分釋放的潛力。

值得一提的是，Mamba 作者之一 Albert Gu 也參與了這項研究。2023 年他和 Karan Goel、Chris Ré、Arjun Desai、Brandon Yang 共同創立了 Cartesia 。公司的使命就是「構建具有長記憶能力的實時智能，無論你身在何處都能運行」，也和這篇文章的主題不謀而合。

• 論文地址：https://arxiv.org/pdf/2507.02782
• 博客地址：https://goombalab.github.io/blog/2025/improving-length-generalization/
• 論文標題： Understanding and Improving Length Generalization in Recurrent Models

為什么循環模型不能進行長度泛化？

函數 f 和 g 不依賴于位置 t，因此理論上循環模型可以自然地處理任何序列長度。但是，當 t 很大時，它們為什么會失敗呢？

本文展示了狀態 h_t 的分布會隨時間的變化。因此，即使 g 和 f 在某個 T 之前工作正常，其他 t>T 的 h_t 可能會有顯著不同，從而導致模型無法產生正確的輸出。實際上，下圖展示了 Mamba-2 的狀態范數隨時間顯著增加：

這就解釋了為什么循環模型無法實現長度泛化：當處理超過訓練長度的序列時，模型會遇到在訓練過程中從未接觸過的狀態 h_t，因此模型并沒有學會如何處理這些狀態。

基于這一觀察，本文提出了一個新的解釋框架 —— 未探索狀態假說（unexplored states hypothesis），用于說明循環模型在長度泛化上失敗的根本原因。

未探索狀態假說（Unexplored States Hypothesis）

當循環模型只在所有可能狀態分布的一個子集上進行訓練時，它們就難以實現長度泛化 —— 也就是說，模型只學習了在有限訓練步數內可能出現的狀態，而未曾接觸那些在無限時間展開狀態遞推后可能出現的狀態分布。
當訓練時間足夠長時，模型會過擬合于這一有限狀態子集，從而在處理更長序列時表現不佳，因為它們會遭遇未被探索過的狀態分布，從而導致性能下降。

訓練干預，使長度泛化

未探索狀態假說指出：要實現長度泛化，并不需要改變模型的架構或機制，而是應該讓模型在訓練階段接觸到更加多樣的狀態分布 —— 尤其是那些在長序列狀態遞推過程中自然產生的分布。

為此，最直接的方法是讓模型直接在更長的序列上進行訓練，但這在實際中往往不可行，原因包括：

• GPU 顯存限制；
• 缺乏足夠長的訓練數據。

因此，我們需要尋找更高效的訓練方法來覆蓋這些狀態分布，從而提升模型的長度泛化能力。

實現長度泛化的方法是：對初始狀態進行干預

一般而言，現代模型架構假設初始狀態為 h_(-1)=0，本文考慮了對初始狀態 h_(-1) 的四種簡單干預。這四種訓練干預措施可以看作是從四種不同的分布中采樣初始狀態 h_(-1) ：

• 隨機噪聲（Random Noise）：將模型狀態初始化為獨立同分布（IID）高斯噪聲，均值為 0，標準差為常數。所有層和注意力頭使用相同的均值和標準差。
• 擬合噪聲（Fitted Noise）：在訓練過程中，記錄所有層和注意力頭上，序列最終狀態的均值和標準差。然后使用與這些統計量相匹配的 IID 高斯分布來初始化狀態，即為每一層和每一個頭分別設置不同的均值和標準差。
• 狀態傳遞（State Passing，SP）：使用先前（不相關）序列的最終狀態作為初始狀態。這些最終狀態是通過在給定序列上應用狀態遞歸獲得的，得到 h_T 并將其用作另一個序列的 h_(-1)。這類似于驗證過程中發生的情況：模型不會在 T 停止，而是繼續滾動狀態并從 h_T 生成輸出。
• TBTT（Truncated Backpropagation Through Time）： 將一條長序列劃分為多個較小的片段，并將每個片段的最終狀態作為下一個片段的初始狀態。這等價于處理整條完整序列，但在片段之間停止梯度的反向傳播。

下圖展示了 Mamba-2 模型經過 500 步（約預訓練總步數的 0.1%）后訓練后，在不同干預措施下的結果：

核心發現 1：SP 與 TBTT 機制可實現長度泛化

SP 和 TBTT 這兩種干預方法，能使模型在遠超訓練序列長度的情況下實現泛化。由此可見：長度泛化有望通過簡單的訓練干預在循環模型中容易實現。

請注意，結果只達到了原始預訓練預算的 0.02% ！

核心發現 2：循環模型狀態的性質，可以通過觀察干預的表現來推斷循環模型狀態分布的性質

在 370M 參數規模的模型中，隨機噪聲干預未能實現長度泛化，而擬合噪聲則有效。這表明，對于 370M 模型來說，模型可達狀態的分布無法通過具有固定方差的高斯分布來近似，但可以通過在每一層和每個注意力頭中使用擬合方差的 IID 高斯分布來近似。

然而，擬合噪聲在 1.3B 模型中未能成功實現長度泛化，這說明大模型的狀態在其元素之間可能存在更復雜的依賴關系，因此無法再用簡單的 IID 分布來建模。

此外，這些干預方法還能解決此前展示的狀態范數隨時間增長的問題，使模型在所有時間步輸出的狀態保持相近的范數，從而提升整體穩定性。

長上下文任務的表現

本文觀察到，這些干預措施能夠實現長度魯棒性（即在訓練上下文長度 T 后，性能不會下降），但尚不清楚它們是否能實現長度泛化（即解決需要利用距離超過 T 個位置的 token 之間關系的任務）。

可能會有疑問，干預措施是否只是簡單地通過阻止模型在訓練上下文長度之外進行推理來實現長度魯棒性 ？

這類似于滑動窗口注意力，無法推理超過滑動窗口的 token，模型在所有評估 t > T 的上下文中會保持恒定性能，但無法解決需要長上下文推理的任務。

在本文的工作中，通過在三個長上下文任務上的實驗結果，展示了這些干預措施確實能夠實現長度泛化。

BABILong

BABILong 是一個具有挑戰性的基準測試，它考察了模型的常識理解能力以及捕捉文本中長程依賴關系的能力。

從下圖可以觀察到，狀態傳遞（State Passing）在少樣本和微調設置下都增強了模型的長度泛化能力（模型是在長度為 2048 的序列上進行訓練和微調的）。

因此，狀態傳遞不僅有助于解決已建立語言模型的困惑度發散問題，還能增強它們解決長上下文推理任務的能力。

密碼檢索任務

密碼檢索任務要求模型在長上下文中的給定深度處檢索一個 5 位數字的密碼。

下圖展示了 Mamba-2 370M 和 780M 官方檢查點在三種設置下的表現：零樣本、常規微調和使用擬合噪聲進行微調。經過擬合噪聲微調的模型能夠利用超過 2048 個位置（訓練上下文長度）之間的 token 關系。特別地，780M 模型能夠完美地解決長度為 256k 的序列中的密碼檢索任務。

合成復制任務

合成復制任務要求模型復制一個任意的 token 序列。

下表展示了在訓練過程中使用狀態傳遞顯著提高了模型在長度超過三倍的序列中的驗證表現。因此，狀態傳遞幫助模型實現長度泛化，解決了比訓練過程中遇到的任務更為復雜的長上下文任務。

深入探討循環模型如何處理上下文

本文已經展示了對初始狀態的干預能夠實現長度魯棒性，并使模型能夠解決長上下文任務。在這些發現的基礎上，本文提出一個度量標準，幫助我們深入了解序列模型是如何處理上下文的。

理想情況下，在文本建模中，希望模型能夠關注最近的上下文，而不是過多地關注那些距離過遠的 token。那么，該如何量化這種行為呢？

盡管每個 token 都會對模型的輸出產生影響，但從直覺上我們會期望最近的 token 應該具有顯著更強的影響力。

然而，注意到 EffRem 曲線在立即上升后逐漸平緩下去。這種情況顯然是有問題的：在 T=8192 時，下一 token 的預測不應該因為模型是只看到最近的 token 還是完整的序列而發生劇烈變化。

在自然語言中，模型應該主要依賴于最近的上下文，而早期的 token 不應該完全改變預測，尤其不應該改變到兩個輸出概率分布之間的整體變差接近 1 的程度。這意味著模型在序列開頭的 token 上受到了不成比例的影響。

狀態傳遞修正了有效記憶

經過狀態傳遞的后訓練，EffRem 曲線顯示出逐漸上升，表明模型對遠程 token 的權重最小，并逐漸增加對最近 token 的權重。特別是，緊鄰上下文中的 token（例如句子中的前一個詞）對下一 token 的預測具有重要影響，這正是文本建模中期望的行為。

簡而言之，通過有效記憶，我們可以確認狀態傳遞幫助模型優先考慮最近的上下文，而不會被遠處的過往 token 不必要地干擾。