萬物互聯：測時記憶、注意偏好、信息留存與在線優化之旅

萬物互聯：測時記憶、注意偏好、信息留存與在線優化之旅

It’s All Connected: A Journey Through Test-Time Memorization,Attentional Bias, Retention, and Online Optimization

摘要：

設計高效且有效的架構主干，一直是增強基礎模型能力的核心研究方向。受人類注意力偏向認知現象的啟發——即自然傾向于優先處理某些事件或刺激——我們重新概念化了神經架構，包括 Transformers、Titans 及現代線性循環神經網絡，將它們視為聯想記憶模塊，這些模塊使用內部目標（稱為注意力偏向）來學習鍵與值的映射。令人驚訝的是，我們觀察到大多數現有序列模型要么利用 (1) 點積相似性，要么利用 (2) ?2 回歸目標作為其注意力偏向。超越這些目標，我們提出了一組替代的注意力偏向配置及其有效近似，以穩定其訓練過程。接著，我們將現代深度學習架構中的遺忘機制重新解釋為一種留存正則化的形式，從而為序列模型提供了一套新穎的遺忘門?；谶@些見解，我們提出了 Miras，這是一個基于以下四種選擇來設計深度學習架構的通用框架：(i) 聯想記憶架構，(ii) 注意力偏向目標，(iii) 留存門，以及 (iv) 記憶學習算法。我們提出了三個新穎的序列模型——Moneta、Yaad 和 Memora——它們超越了現有線性 RNN 的能力，同時保持了快速可并行的訓練過程。我們的實驗表明，Miras 中不同的設計選擇會產生具有不同優勢的模型。例如，Miras 的某些實例在特定任務（如語言建模、常識推理和召回密集型任務）中取得了卓越的性能，甚至超越了 Transformers 和其他現代線性循環模型。

1 引言

為序列建模設計高效的架構主干是增強基礎模型在語言、計算機視覺、計算生物學和神經科學等領域能力的關鍵。雖然 Transformers 主要得益于其上下文學習能力和規模化學習能力，已在序列建模中穩固確立為最先進的模型，但其二次時間與空間復雜度限制了其在需要長上下文建模任務中的應用。

近期的努力旨在通過設計高效的循環替代方案來克服 Transformers 在長上下文建模中的局限性。與 Transformers 線性增長的內存不同，這些模型將上下文壓縮到固定大小的內存中，需要改進內存管理以獲得可比的性能。為了設計更有效的架構，研究重點在于通過使用/設計更具表現力的以下方面來提升內存容量及其管理：(1) 學習規則：從赫布規則到 Delta 規則；(2) 遺忘門：從 LSTM 到 Mamba2 再到 Titan 的遺忘門；(3) 更具表現力的內存架構：從 RetNet 和 LRU 的向量值內存，到 Titans 和 TTT 的神經深度內存。

這些進步的核心存在一個關鍵問題：“這些序列模型背后的基本設計框架是什么，以及如何增強這些模型？”。受神經心理學文獻中關于聯想記憶和學習的廣義定義的啟發，多項研究討論了 Transformers 與（線性）循環神經網絡同聯想記憶之間的聯系。然而，這些研究要么 (1) 缺乏能完整闡明底層學習算法的普適性解釋，要么 (2) 局限于特定的聯想記憶定義而缺乏通用性，和/或 (3) 無法描述標準且廣泛使用的組件，如遺忘門。

貢獻：受人類注意力偏向認知現象的啟發——即自然傾向于優先處理某些事件或刺激——我們基于具有注意力偏向的廣義聯想記憶定義，重新概念化了神經架構，包括 Transformers、Titans 和其他現代線性循環神經網絡。我們將注意力偏向定義并形式化為序列模型的內部記憶目標，旨在學習輸入之間的底層映射。我們的表述揭示，幾乎所有現有的序列模型都是利用同一類型注意力偏向的聯想記憶。我們將現代深度學習架構中現有的遺忘機制重新解釋為注意力偏向的一種留存 ?2 正則化形式，然后為序列模型提供了一套新穎的替代留存門，為如何平衡學習新概念與保留已學概念提供了新的見解。

基于我們對記憶和遺忘門的表述，我們提出了 Miras1，這是一個通過以下四種選擇來設計新穎序列建模架構的基礎框架：(1) 注意力偏向（即記憶目標），(2) 留存門，(3) 記憶架構，以及 (4) 記憶學習算法（即優化器）。我們提出并討論了幾種新穎的設計選擇，從而產生了超越現有序列建模架構的新架構。

最后，我們專注于 Miras 的三個新穎變體——Moneta、Yaad 和 Memora——它們基于超越簡單 ?2 回歸目標的注意力偏向，以及比現有機制更穩健的新型留存門控機制。我們進一步在語言建模、常識推理、大海撈針和召回密集型任務上對這三個變體進行了實驗評估。結果表明這些變體性能優越，超越了最先進的序列模型。

路線圖：在第 2 節中，我們回顧了文獻并討論了本文使用的相關概念。在第 3 節中，我們提出并討論了聯想記憶的廣義定義，并正式定義了注意力偏向的概念。接著，我們討論了兩種視角——學習-保留和跟隨正則化領導者——以通過優化透鏡來解讀序列建模，并證明學習-保留相對于 FTRL 的普適性。在第 4 節中，我們介紹了 Miras 框架并討論了它如何統一現代序列模型。在第 5 節中，為了展示 Miras 框架的潛力，我們討論了針對 (1) 注意力偏向和 (2) 留存門（遺忘門）的多種新穎設計選擇。隨后在第 5.3 節，我們介紹了作為 Miras 變體的三個新穎序列模型，并討論了如何以可并行化的方式訓練它們。最后，我們的實驗評估在第 6 節中報告。

2 預備知識與背景

在本節中，我們將回顧全文所涉及的相關研究和背景概念。

為了解決上述限制，近年來，使用矩陣值內存并采用 Delta 學習規則的循環模型變得越來越受歡迎。盡管有顯著優勢，即使是這些基于 Delta 規則的循環模型也面臨理論限制，并且在實踐中性能一般。最近，多項研究旨在通過添加標量或通道級遺忘門機制、使用負特征值以及多步學習來提高此類模型的性能。然而，由于內存架構的表達能力不足，它們在長上下文任務中仍然存在性能下降的問題。

深度記憶模塊：Titans 與測試時訓練 為了克服有限的內存并擴展深度序列模型的有效上下文長度，近期的研究聚焦于具有深度記憶模塊的新一代架構。這些架構建立在元學習視角之上，其中記憶是一個使用梯度下降（可能帶動量）更新的 MLP 架構。Sun 等人 (2024) 進一步提供了一個統一的視角，闡釋了線性和 softmax 注意力如何分別是（核）回歸損失的參數化和非參數化解，但將其他現代線性 RNNs 視為該類模型之外的方法。最近，在我們的一項并行工作中，Wang 等人 (2025) 表明，通過對現代 RNNs（如 RetNet、Mamba）進行額外簡化，它們大致可歸為內部優化回歸損失的同類模型。然而，“這些序列模型背后能夠準確統一現有架構的基本設計框架是什么？”這個問題仍未得到解答。此外，遺忘門的作用及其在現代序列模型中的替代選擇，目前的研究仍顯不足。

3 聯想記憶、注意力偏向與留存

聯想記憶是人類學習中不可或缺的組成部分，一直是文獻中許多人造神經架構的靈感來源。然而，這些研究只定義了聯想記憶概念的特定實例，將架構限制在實體間某種特定類型的相似性度量范圍內。廣義而言，聯想記憶是一種將一組鍵 映射到一組值 的算子。因此，為了學習數據中潛在的映射模式，它需要一個針對某種記憶類型并衡量所學映射質量的目標：

需注意以下幾點：

備注1.當我們用參數 W W參數化記憶時，我們使用 M ( W , k )
。在這種參數化設置下，公式 (4) 中的優化問題應在參數 W W上進行。此外，在參數化設置中，我們可能會使用額外的正則化項 R ( W )
來控制對過去數據的保留。

備注2.學習鍵與值之間的映射（公式 4）是一個元學習問題，其中注意力偏向在內循環中被優化，而神經網絡的所有其他參數（例如線性投影、卷積等）則在外循環中被優化。因此，模型學習如何在測試時將數據存儲到其參數中。

3.1 通過優化的視角學習記憶與留存

定義 3.1 將基于聯想記憶概念的神經架構設計，轉化為通過最小化目標函數 L L來學習鍵與值之間的底層映射。為了優化公式 4，一個簡單的方法是采用梯度下降的思想。具體來說，給定一個新的鍵值對，我們按如下方式更新記憶：

3.2 視角一：在線回歸與跟隨正則化領導者

公式 (5) 可視為對損失函數序列執行一步在線梯度下降：

公式 (7) 使用了損失函數的線性近似和二次正則化。然而，原則上我們可以使用損失函數的其他近似以及其他正則化函數，正如過去在線性優化或一般優化中所使用的那樣。這種改變是開發其他優化算法（如鏡像下降）背后的思想。更具體地說，我們可以將公式 (7) 中的更新規則推廣為以下形式：

3.3 視角二：學習最新令牌同時保留先前信息

我們在附錄 B 中提供了證明。上述命題表明，在一些溫和假設下，（學習-保留視角）也可以解釋由（跟隨正則化領導者視角）得到的方法。因此，（學習-保留視角）可被視為一個更通用的版本。這就是為什么我們在接下來的大部分推導中主要關注這一視角。

備注3：根據上述視角，我們可以看到，即使使用額外的全局正則化，也不存在記憶擦除或遺忘過程（現代架構中的常見術語），但模型可能會決定不保留記憶的過去狀態。有趣的是，這一觀察也與人類的記憶過程相符，即大腦不會擦除記憶，但記憶可能因提取失敗而變得無法訪問。因此，我們隨后將使用 “留存門”來指代這一項，而非遺忘門。

備注4：如我們在第4節討論并總結在表1中的那樣，大多數現有的現代序列模型都在使用梯度下降來優化聯想記憶目標（公式4中的注意力偏向）。因此，為了進一步理解現有序列模型及其在線學習解釋之間的聯系，我們討論了上述兩種限于基于梯度下降更新規則的視角。然而，我們在公式4中對注意力偏向和聯想記憶的初始定義更為寬泛，并且可以通過任何優化算法（例如，甚至牛頓法或非參數解）進行優化。

4 MIRAS：學習用魯棒且富有表達力的記憶進行記憶

基于我們關于聯想記憶、注意力偏差的定義以及之前的觀點，我們提出了 MIRAS 框架，它不僅精確地統一了現有的骨干架構，還為如何設計下一代序列模型提供了見解。如第3節早先討論的，學習聯想記憶可以解釋為一個元學習任務，在此任務中，聯想記憶學習如何在測試時壓縮數據并將其存儲到其參數中。在這類任務中，記憶的架構尤為重要，因為在較長上下文中，記憶結構的表達能力可能會限制其學習潛在模式的能力。因此，設計序列模型的第一個選擇是記憶的結構。給定由一組參數 W W參數化的記憶結構，如前所述，我們旨在通過一個學習算法（例如梯度下降）來最小化帶有記憶保留正則項 Ret ( ? ) 的損失函數 ? ( W ; ? )
。相應地，MIRAS 需要四個設計選擇：

5 超越現有的注意力偏差與保留門機制

5.2 替代的保留門機制

變體 1：通過 -散度在縮放概率單純形上進行記憶化。在學習過程中，一種防止數值不穩定和值爆炸的常見技術是將搜索空間限制在一個有界域內。遵循這一原則，為了避免數值不穩定性，我們可以將變量 約束在（縮放的）概率單純形內。換句話說，我們可以將狀態限制在如下約束集內：

5.3 MIRAS 的變體：MONETA、YAAD 和 MEMORA

在上一節中，我們討論了注意力偏差和保留門的不同潛在選擇，以展示 MIRAS 的通用性和潛力。在本節中，基于我們的框架，我們提出了三種新穎的序列模型，每種模型都基于不同的動機設計，并討論了它們如何利用快速并行訓練。

YAAD。基于我們關于保護記憶免受極端事件（令牌）影響的魯棒性記憶重要性的討論，我們基于 Huber 目標設計了 YAAD。即在 MIRAS 框架中，對于記憶結構的選擇，我們遵循 MONETA，使用與上述相同架構的 MLP；對于注意力偏差的選擇，我們使用 Huber 損失（由公式 16 定義）；對于保留門的選擇，為簡潔起見，我們結合使用局部和全局保留項，表示為

這等價于 Behrouz 等人（2024c）引入的“遺忘門”機制；最后，我們直接使用梯度下降作為記憶學習算法。根據以上選擇，我們可以將得到的記憶學習過程描述如下：

Memora。 最后，在 Memora 中，我們采用了彈性網絡正則化的思想（即硬保留和軟保留）。為此，在 MIRAS 框架中：(1) 在記憶架構的選擇上，與上述變體類似，我們使用 MLP（架構與之前的變體相同）。(2) 在注意力偏差的選擇上，我們使用簡單的 ?? 回歸損失。(3) 在保留門的選擇上，我們使用如公式 21 所示的 KL 散度。(4) 最后，我們使用梯度下降來優化記憶，從而得到以下更新規則：

5.4 架構骨干與快速訓練

架構骨干。對于架構骨干，我們完全遵循近期研究（Behrouz 等人 2024c; Yang 等人 2024a）的做法：在 Llama 的宏觀架構中，使用帶有 SwiGLU() 激活函數的 MLP、旋轉位置編碼（RoPE）（Su 等人 2024）和 RMSNorm（Zhang 等人 2019），將注意力模塊替換為我們 MIRAS 的變體。對于 MIRAS 層模塊，我們遵循近現代線性循環模型（Behrouz 等人 2024c; Yang 等人 2024a），在查詢、鍵和值投影之后分別加入一維深度可分離卷積層（卷積核大小為 4）。為了訓練穩定性，我們同樣對 q 和 k 使用 ? 2
歸一化。MIRAS 層模塊的輸出經過歸一化，并通過一個線性層進行門控（Mehta 等人 2023）。

6 實驗

在我們的實驗評估中，我們旨在回答三個主要問題：(1) 不同的注意力偏差在實踐中是否會導致不同的架構？(2) 不同類型的保留門（即保留門機制）如何影響模型在長上下文中的性能？(3) 與基線模型相比，Memora、Moneta 和 Yaad 在下游任務中的表現如何？

實驗設置。我們使用訓練上下文窗口大小為 4096 來訓練我們的模型，使用的數據集為 FineWeb-Edu (Penedo 等人 2024)（用于語言建模和常識推理任務）或 C4 數據集 (Raffel 等人 2020)（用于擴展規律分析）。我們使用的模型參數量級為 120M、340M、760M 和 1.3B。小模型（120M 和 340M）在從數據集中采樣的 15B tokens 上訓練，中等模型（760M）在 30B tokens 上訓練，大模型在 100B tokens 上訓練。基線結果由 Behrouz 等人 (2024c) 報告。

6.1 語言建模與常識推理

我們遵循近期研究 (Behrouz 等人 2024c; Yang 等人 2024a,c)，首先關注語言建模的困惑度以及常識推理任務。Memora、Yaad、Moneta 以及參數量為 340M、760M 和 1.3B 的基線模型的結果在表 2 中報告。我們所有的變體都優于所有基線，包括 Transformer++、現代線性循環模型和混合方法。相較于混合模型的優越性能尤為重要，因為我們所有的變體都是純循環的（無注意力機制）。在 Miras 的三個變體中，雖然 Moneta 的表現略弱于 Memora 和 Yaad，但其他兩個變體表現接近，并且根據任務和模型規模，最佳模型會有所不同。

6.2 擴展規律

為了評估模型的擴展規律并與基線進行比較，本節中，我們繪制了它們在不同模型規模和上下文窗口下的性能表現。

上下文長度。我們首先將訓練上下文長度從 2K 變化到 32K，對規模為 340M 和 760M 的兩個模型版本進行評估。結果報告在圖 3（中和右）中。當增加上下文長度時，Miras 的所有三個變體都比最先進的基線模型擴展得更好。我們將這種優越性能歸因于：(1) 富有表達力的記憶架構。與使用向量值和矩陣值記憶的 Mamba2 和 GSA 等基線不同，我們的變體使用具有更強表達能力的 2 層 MLP 來從更長的序列中學習。(2) 保留門和注意力偏差的選擇：我們所有的三個變體都超越了標準的注意力偏差和保留門。這些選擇可以幫助記憶更好地管理其固定大小的容量。

模型規模。我們還在圖 3（左）中報告了我們的模型和基線在 FLOPs 與困惑度方面的表現。在幾乎相同的 FLOPs 預算下，所有三個變體都優于所有基線。這些結果再次支持了強大記憶設計的重要性。

6.3 大海撈針

為了評估我們的模型和基線的有效上下文窗口，我們使用“大海撈針”任務。在該任務中，我們評估模型從長干擾文本（即“干草堆”）中檢索特定信息（即“針”）的能力。我們專注于 RULER 基準測試 (Hsieh 等人 2024) 中的單針任務 (S-NIAH)，并在長度為 1K、2K、4K 和 8K 的序列上評估我們的模型和基線。結果報告在表 3 中。我們所有的變體都以相當大的優勢優于所有基線。有趣的是，當數據是合成噪聲（S-NIAH-PK）時，Moneta 表現出比其他模型更好的性能。這一觀察結果驗證了 -范數目標和保留門的有效性，因為它們對噪聲更加魯棒。

6.4 消融實驗

在本節中，我們進行消融實驗，以驗證本文中討論的不同設計選擇是否對實現更好的結果有積極貢獻。

p值對性能的影響。我們首先評估p值對Moneta性能的影響。我們將p值在{1, 1.5, 2, 2.8, 3, 3.2, 4}范圍內變化，并將上下文窗口從2K調整至16K。結果如圖4所示。有趣的是，隨著p值增加，性能并未呈現單調變化趨勢，當p=3時達到最佳性能，而p=4時性能最差。此外，盡管不同p值會導致具有不同性能的記憶模塊，但隨著上下文長度增加，擴展規律幾乎保持一致。

q值對性能的影響。類似地，我們通過在{2, 3, 4, 5}范圍內變化q值來評估其影響。有趣的是，與p值不同，q值能夠改變上下文長度增加時的擴展規律。這一現象的主要原因是q值決定了保留門的特性，一個強大的保留門可以改善記憶管理，從而帶來更好的性能表現。

設計選擇的影響。為了評估架構設計選擇，我們對Yaad進行了消融實驗。結果見表4。第一行報告了Yaad的原始性能，而（1）第二行移除了保留機制（即設β=1），（2）第三行使δ參數獨立于輸入，（3）第三行從Huber損失中移除?2損失項，（4）第四行移除?1條件約束，（5）最后一行用線性層替換MLP。這些結果表明所有設計選擇都對模型性能具有貢獻。

7 結論

本文提出了一種通用框架 Miras，用于解釋在線優化與測試時記憶化之間的關聯。Miras 框架能夠闡釋文獻中多種標準架構選擇（例如遺忘門）的作用，并有助于設計能夠更有效管理記憶的新一代架構。基于我們的框架，我們提出了三種新穎的序列模型，每種模型都具有其各自的優勢與局限性。實驗評估表明，在各種下游任務中，所有這些變體均優于 Transformer 和線性 RNN。本研究通過 Miras 展示了一系列多樣化的變體。未來，探索這些替代架構在不同下游任務中的應用是一個值得關注的研究方向。

原文鏈接：https://arxiv.org/pdf/2504.13173