線性循環(huán)模型（如 Mamba）和線性注意力機(jī)制都具備這樣一個(gè)顯著優(yōu)勢(shì)：它們能夠處理極長(zhǎng)的序列，這一能力對(duì)長(zhǎng)上下文推理任務(wù)至關(guān)重要。

事實(shí)上，這正是它們相較于 Transformer 的關(guān)鍵優(yōu)勢(shì) —— 后者受限于有限的上下文窗口，且在序列長(zhǎng)度上的計(jì)算復(fù)雜度是二次的，成為性能瓶頸。 

過(guò)去，循環(huán)模型面臨的主要問(wèn)題是性能不足：在處理短序列時(shí)，它們的表現(xiàn)往往不如 Transformer。然而，隨著架構(gòu)上的一系列突破，循環(huán)模型的性能已經(jīng)顯著提升，在許多任務(wù)中已能與 Transformer 媲美，甚至已經(jīng)被應(yīng)用于多個(gè)工業(yè)場(chǎng)景中，如音頻建模和代碼補(bǔ)全等。

但近期的多項(xiàng)研究發(fā)現(xiàn)，循環(huán)模型仍存在一個(gè)關(guān)鍵短板：盡管它們?cè)谟?xùn)練長(zhǎng)度范圍內(nèi)表現(xiàn)良好，但在處理超出訓(xùn)練長(zhǎng)度的序列時(shí)，往往難以泛化，表現(xiàn)明顯下降。

事實(shí)也確實(shí)如此，舉例來(lái)說(shuō)，下圖為 Mamba-2 檢查點(diǎn)在不同序列位置上的性能變化（以困惑度 Perplexity 衡量，數(shù)值越低代表性能越好）。可以明顯看到，一旦序列位置超出了訓(xùn)練上下文范圍，這些模型幾乎就變得毫無(wú)作用：它們無(wú)法實(shí)現(xiàn)長(zhǎng)度泛化。 

這就帶來(lái)一個(gè)問(wèn)題：現(xiàn)有的循環(huán)模型在長(zhǎng)序列上表現(xiàn)較差，而在短序列上相比 Transformer 也沒(méi)有明顯的效率優(yōu)勢(shì)；換句話說(shuō)，它們?cè)趦蓚€(gè)維度上似乎都顯得不夠理想。

那這是否意味著循環(huán)模型就毫無(wú)用處了呢？

當(dāng)然不是！

本文，來(lái)自 CMU、 Cartesia AI 的研究者證明了通過(guò)簡(jiǎn)單的訓(xùn)練干預(yù)，循環(huán)模型是完全可以實(shí)現(xiàn)長(zhǎng)度泛化的。只需 500 步后訓(xùn)練（約占預(yù)訓(xùn)練預(yù)算的 0.1%），就能讓模型在高達(dá) 256k 長(zhǎng)度的序列上實(shí)現(xiàn)泛化！

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

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

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

為什么循環(huán)模型不能進(jìn)行長(zhǎng)度泛化？

對(duì)于一個(gè)包含 t 個(gè)元素的輸入序列，循環(huán)模型將輸入上下文壓縮成一個(gè)固定大小的循環(huán)狀態(tài)，在時(shí)間 t=0 時(shí)，狀態(tài)用某個(gè)值 h_(?1) 初始化，然后在每個(gè)時(shí)間步 t 通過(guò)更新函數(shù) f 進(jìn)行更新：。同樣地，時(shí)間 t 時(shí)的輸出僅依賴(lài)于狀態(tài) h_t 和當(dāng)前輸入 x_t，即對(duì)于某個(gè)其他函數(shù) g，輸出 y_t 可以表示為：。

函數(shù) f 和 g 不依賴(lài)于位置 t，因此理論上循環(huán)模型可以自然地處理任何序列長(zhǎng)度。但是，當(dāng) t 很大時(shí)，它們?yōu)槭裁磿?huì)失敗呢？

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

這就解釋了為什么循環(huán)模型無(wú)法實(shí)現(xiàn)長(zhǎng)度泛化：當(dāng)處理超過(guò)訓(xùn)練長(zhǎng)度的序列時(shí)，模型會(huì)遇到在訓(xùn)練過(guò)程中從未接觸過(guò)的狀態(tài) h_t，因此模型并沒(méi)有學(xué)會(huì)如何處理這些狀態(tài)。

基于這一觀察，本文提出了一個(gè)新的解釋框架 —— 未探索狀態(tài)假說(shuō)（unexplored states hypothesis），用于說(shuō)明循環(huán)模型在長(zhǎng)度泛化上失敗的根本原因。

未探索狀態(tài)假說(shuō)（Unexplored States Hypothesis）

當(dāng)循環(huán)模型只在所有可能狀態(tài)分布的一個(gè)子集上進(jìn)行訓(xùn)練時(shí)，它們就難以實(shí)現(xiàn)長(zhǎng)度泛化 —— 也就是說(shuō)，模型只學(xué)習(xí)了在有限訓(xùn)練步數(shù)內(nèi)可能出現(xiàn)的狀態(tài)，而未曾接觸那些在無(wú)限時(shí)間展開(kāi)狀態(tài)遞推后可能出現(xiàn)的狀態(tài)分布。

當(dāng)訓(xùn)練時(shí)間足夠長(zhǎng)時(shí)，模型會(huì)過(guò)擬合于這一有限狀態(tài)子集，從而在處理更長(zhǎng)序列時(shí)表現(xiàn)不佳，因?yàn)樗鼈儠?huì)遭遇未被探索過(guò)的狀態(tài)分布，從而導(dǎo)致性能下降。

訓(xùn)練干預(yù)，使長(zhǎng)度泛化

未探索狀態(tài)假說(shuō)指出：要實(shí)現(xiàn)長(zhǎng)度泛化，并不需要改變模型的架構(gòu)或機(jī)制，而是應(yīng)該讓模型在訓(xùn)練階段接觸到更加多樣的狀態(tài)分布 —— 尤其是那些在長(zhǎng)序列狀態(tài)遞推過(guò)程中自然產(chǎn)生的分布。

為此，最直接的方法是讓模型直接在更長(zhǎng)的序列上進(jìn)行訓(xùn)練，但這在實(shí)際中往往不可行，原因包括：

• GPU 顯存限制；
• 缺乏足夠長(zhǎng)的訓(xùn)練數(shù)據(jù)。

因此，我們需要尋找更高效的訓(xùn)練方法來(lái)覆蓋這些狀態(tài)分布，從而提升模型的長(zhǎng)度泛化能力。

實(shí)現(xiàn)長(zhǎng)度泛化的方法是：對(duì)初始狀態(tài)進(jìn)行干預(yù)。

一般而言，現(xiàn)代模型架構(gòu)假設(shè)初始狀態(tài)為 h_(-1)=0，本文考慮了對(duì)初始狀態(tài) h_(-1) 的四種簡(jiǎn)單干預(yù)。這四種訓(xùn)練干預(yù)措施可以看作是從四種不同的分布中采樣初始狀態(tài) h_(-1) ：

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

下圖展示了 Mamba-2 模型經(jīng)過(guò) 500 步（約預(yù)訓(xùn)練總步數(shù)的 0.1%）后訓(xùn)練后，在不同干預(yù)措施下的結(jié)果：

核心發(fā)現(xiàn) 1：SP 與 TBTT 機(jī)制可實(shí)現(xiàn)長(zhǎng)度泛化。

SP 和 TBTT 這兩種干預(yù)方法，能使模型在遠(yuǎn)超訓(xùn)練序列長(zhǎng)度的情況下實(shí)現(xiàn)泛化。由此可見(jiàn)：長(zhǎng)度泛化有望通過(guò)簡(jiǎn)單的訓(xùn)練干預(yù)在循環(huán)模型中容易實(shí)現(xiàn)。

請(qǐng)注意，結(jié)果只達(dá)到了原始預(yù)訓(xùn)練預(yù)算的 0.02% ！

核心發(fā)現(xiàn) 2：循環(huán)模型狀態(tài)的性質(zhì)，可以通過(guò)觀察干預(yù)的表現(xiàn)來(lái)推斷循環(huán)模型狀態(tài)分布的性質(zhì)。

在 370M 參數(shù)規(guī)模的模型中，隨機(jī)噪聲干預(yù)未能實(shí)現(xiàn)長(zhǎng)度泛化，而擬合噪聲則有效。這表明，對(duì)于 370M 模型來(lái)說(shuō)，模型可達(dá)狀態(tài)的分布無(wú)法通過(guò)具有固定方差的高斯分布來(lái)近似，但可以通過(guò)在每一層和每個(gè)注意力頭中使用擬合方差的 IID 高斯分布來(lái)近似。

然而，擬合噪聲在 1.3B 模型中未能成功實(shí)現(xiàn)長(zhǎng)度泛化，這說(shuō)明大模型的狀態(tài)在其元素之間可能存在更復(fù)雜的依賴(lài)關(guān)系，因此無(wú)法再用簡(jiǎn)單的 IID 分布來(lái)建模。

此外，這些干預(yù)方法還能解決此前展示的狀態(tài)范數(shù)隨時(shí)間增長(zhǎng)的問(wèn)題，使模型在所有時(shí)間步輸出的狀態(tài)保持相近的范數(shù)，從而提升整體穩(wěn)定性。

長(zhǎng)上下文任務(wù)的表現(xiàn)

本文觀察到，這些干預(yù)措施能夠?qū)崿F(xiàn)長(zhǎng)度魯棒性（即在訓(xùn)練上下文長(zhǎng)度 T 后，性能不會(huì)下降），但尚不清楚它們是否能實(shí)現(xiàn)長(zhǎng)度泛化（即解決需要利用距離超過(guò) T 個(gè)位置的 token 之間關(guān)系的任務(wù)）。

可能會(huì)有疑問(wèn)，干預(yù)措施是否只是簡(jiǎn)單地通過(guò)阻止模型在訓(xùn)練上下文長(zhǎng)度之外進(jìn)行推理來(lái)實(shí)現(xiàn)長(zhǎng)度魯棒性 ？

這類(lèi)似于滑動(dòng)窗口注意力，無(wú)法推理超過(guò)滑動(dòng)窗口的 token，模型在所有評(píng)估 t > T 的上下文中會(huì)保持恒定性能，但無(wú)法解決需要長(zhǎng)上下文推理的任務(wù)。

在本文的工作中，通過(guò)在三個(gè)長(zhǎng)上下文任務(wù)上的實(shí)驗(yàn)結(jié)果，展示了這些干預(yù)措施確實(shí)能夠?qū)崿F(xiàn)長(zhǎng)度泛化。

BABILong 

BABILong 是一個(gè)具有挑戰(zhàn)性的基準(zhǔn)測(cè)試，它考察了模型的常識(shí)理解能力以及捕捉文本中長(zhǎng)程依賴(lài)關(guān)系的能力。

從下圖可以觀察到，狀態(tài)傳遞（State Passing）在少樣本和微調(diào)設(shè)置下都增強(qiáng)了模型的長(zhǎng)度泛化能力（模型是在長(zhǎng)度為 2048 的序列上進(jìn)行訓(xùn)練和微調(diào)的）。

因此，狀態(tài)傳遞不僅有助于解決已建立語(yǔ)言模型的困惑度發(fā)散問(wèn)題，還能增強(qiáng)它們解決長(zhǎng)上下文推理任務(wù)的能力。

密碼檢索任務(wù)

密碼檢索任務(wù)要求模型在長(zhǎng)上下文中的給定深度處檢索一個(gè) 5 位數(shù)字的密碼。

下圖展示了 Mamba-2 370M 和 780M 官方檢查點(diǎn)在三種設(shè)置下的表現(xiàn)：零樣本、常規(guī)微調(diào)和使用擬合噪聲進(jìn)行微調(diào)。經(jīng)過(guò)擬合噪聲微調(diào)的模型能夠利用超過(guò) 2048 個(gè)位置（訓(xùn)練上下文長(zhǎng)度）之間的 token 關(guān)系。特別地，780M 模型能夠完美地解決長(zhǎng)度為 256k 的序列中的密碼檢索任務(wù)。

合成復(fù)制任務(wù)

 合成復(fù)制任務(wù)要求模型復(fù)制一個(gè)任意的 token 序列。

下表展示了在訓(xùn)練過(guò)程中使用狀態(tài)傳遞顯著提高了模型在長(zhǎng)度超過(guò)三倍的序列中的驗(yàn)證表現(xiàn)。因此，狀態(tài)傳遞幫助模型實(shí)現(xiàn)長(zhǎng)度泛化，解決了比訓(xùn)練過(guò)程中遇到的任務(wù)更為復(fù)雜的長(zhǎng)上下文任務(wù)。

深入探討循環(huán)模型如何處理上下文

本文已經(jīng)展示了對(duì)初始狀態(tài)的干預(yù)能夠?qū)崿F(xiàn)長(zhǎng)度魯棒性，并使模型能夠解決長(zhǎng)上下文任務(wù)。在這些發(fā)現(xiàn)的基礎(chǔ)上，本文提出一個(gè)度量標(biāo)準(zhǔn)，幫助我們深入了解序列模型是如何處理上下文的。

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

本文引入了「有效記憶（Effective Remembrance）」來(lái)衡量一個(gè)自回歸模型在多大程度上有效地記住了先前的 token。用表示在給定上下文的情況下，自回歸序列模型為下一個(gè) token 輸出的概率。然后，定義，其中是概率分布之間的距離度量。該度量大致衡量了模型在時(shí)間 T 時(shí)刻有效記住過(guò)去的 token x [0:t?1]。如果，這意味著使用 x [0:T] 和 x [t:T] 進(jìn)行預(yù)測(cè)的結(jié)果相同，即模型沒(méi)有有效記住任何過(guò)去的 token。

相反，如果較高，則表示模型受到過(guò)去 token 的顯著影響，因?yàn)閺纳舷挛闹幸瞥鼈儠?huì)顯著改變預(yù)測(cè)結(jié)果。

下圖展示了對(duì)于兩個(gè)官方的 Mamba-2 檢查點(diǎn)（它們無(wú)法進(jìn)行長(zhǎng)度泛化），在不同的 t 和 T=8192（訓(xùn)練上下文的四倍）下的：

盡管每個(gè) token 都會(huì)對(duì)模型的輸出產(chǎn)生影響，但從直覺(jué)上我們會(huì)期望最近的 token 應(yīng)該具有顯著更強(qiáng)的影響力。

然而，注意到 EffRem 曲線在立即上升后逐漸平緩下去。這種情況顯然是有問(wèn)題的：在 T=8192 時(shí)，下一 token 的預(yù)測(cè)不應(yīng)該因?yàn)槟Ｐ褪侵豢吹阶罱?token 還是完整的序列而發(fā)生劇烈變化。

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

狀態(tài)傳遞修正了有效記憶

 經(jīng)過(guò)狀態(tài)傳遞的后訓(xùn)練，EffRem 曲線顯示出逐漸上升，表明模型對(duì)遠(yuǎn)程 token 的權(quán)重最小，并逐漸增加對(duì)最近 token 的權(quán)重。特別是，緊鄰上下文中的 token（例如句子中的前一個(gè)詞）對(duì)下一 token 的預(yù)測(cè)具有重要影響，這正是文本建模中期望的行為。

簡(jiǎn)而言之，通過(guò)有效記憶，我們可以確認(rèn)狀態(tài)傳遞幫助模型優(yōu)先考慮最近的上下文，而不會(huì)被遠(yuǎn)處的過(guò)往 token 不必要地干擾。