記憶體快取:具備成長型記憶力的循環神經網路

記憶體快取:具備成長型記憶力的循環神經網路

作者:Ali Behrouz1,2,†, Zeman Li1,3, Yuan Deng1, Peilin Zhong1, Meisam Razaviyayn1,3, Vahab Mirrokni1
Google Research 標誌
康乃爾大學標誌
南加州大學標誌
† 通訊作者:alibehrouz@google.com

摘要

Transformer 模型已成為近期序列建模進展的骨幹,這主要歸功於其隨上下文長度增長的記憶容量。雖然這種機制有利於檢索任務,但也導致了二次方的計算複雜度,促使近期研究致力於探索可行的次二次方(subquadratic)循環替代方案。儘管這些循環架構在多項領域展現出初步成果,但在需要高度回憶能力的任務中,其表現仍不及 Transformer,這通常被歸因於其固定大小的記憶體限制。在本文中,我們引入了「記憶體快取」(Memory Caching, MC),這是一項簡單卻有效的技術,透過快取記憶體狀態(又稱隱藏狀態)的檢查點來增強循環模型。MC 使得循環神經網路(RNN)的有效記憶容量能隨序列長度增長,提供了一種靈活的權衡方案,能在 RNN 的固定記憶(即 O(L) 複雜度)與 Transformer 的成長型記憶(即 O(L2) 複雜度)之間進行插值。我們提出了四種 MC 變體,包括閘道聚合與稀疏選擇機制,並討論了它們在線性及深度記憶模組上的影響。我們在語言建模與長上下文理解任務上的實驗結果顯示,MC 顯著提升了循環模型的效能,證實了其有效性。在上下文回憶任務的結果表明,雖然 Transformer 仍保持最高準確率,但我們的 MC 變體展現出極具競爭力的表現,成功縮小了與 Transformer 的差距,並優於現有的最先進循環模型。

1. 引言

Transformer 模型 (Vaswani et al., 2017) 是近期機器學習在各領域取得進展的基石 (Jumper et al., 2021); (Dosovitskiy et al., 2021); (Comanici et al., 2025)。這種成功通常歸因於其大規模學習能力 (Kaplan et al., 2020) 以及上下文學習(in-context learning)能力 (Brown et al., 2020),這兩者皆源自其核心構件——注意力模組(attention module),該模組作為一種具有成長容量的關聯記憶運作 (Ramsauer et al., 2021); (Bietti et al., 2024); (Behrouz et al., 2026)。雖然這種成長型記憶對許多檢索任務非常有效 (Arora et al., 2024b),但它帶來了二次方的計算複雜度以及推論時高昂的記憶體使用量(KV-caching)。這激發了對次二次方架構的開發,旨在提高效率同時維持效能 (Dai et al., 2019); (Child et al., 2019); (Poli et al., 2023)

特別是,循環神經網路(RNNs)旨在將過去數據壓縮至其記憶狀態中,在整個輸入序列期間維持固定大小,近年來重新受到關注 (Katharopoulos et al., 2020); (Irie et al., 2021); (Sun et al., 2023); (Behrouz et al., 2025c)。儘管在多項短上下文語言建模任務 (Irie et al., 2022) 及其他序列建模任務(如影片數據)(Park et al., 2025) 中展現出有希望的成果,但此類循環架構的固定記憶狀態仍是限制其發揮潛力的瓶頸。這些架構的基礎在於循環與數據壓縮,若設計得當,可產生高效且具表現力的學習演算法 (Merrill et al., 2024); (Huang et al., 2024)。然而,其壓縮不斷增長序列的固定容量迫使它們遺忘過去的資訊,這在需要高度回憶能力及長上下文的任務中構成了關鍵瓶頸 (Arora et al., 2024b); (Kuratov et al., 2024)

記憶體快取方法總覽圖:每個 token 除了關注其在線記憶外,還關注一組來自過去的快取記憶。
圖 1:記憶體快取方法總覽。每個 token 不僅關注其在線記憶(online memory),還關注一組來自過去的快取記憶(cached memories)。

貢獻。 我們引入了「記憶體快取」(Memory Caching, MC),這是一項通用技術,透過快取記憶狀態的檢查點,使循環模型的有效記憶力能隨序列長度增長(參見圖 1)。MC 提供了一種介於標準循環與注意力機制之間的靈活中間地帶,提供了可控制的 O(NL) 複雜度。這使得我們能在 RNN 的 O(L) 複雜度與 Transformer 的 O(L2) 複雜度之間進行靈活插值。我們的主要貢獻如下:

  • MC 框架: 我們提出將序列分段,並快取每個分段的壓縮記憶狀態,使模型能直接存取整個歷史的壓縮資訊。
  • 新穎的聚合策略: 我們引入了四種利用這些快取記憶的方法:(i, ii) (閘道)殘差記憶,利用殘差連接與新穎的上下文感知閘道機制;(iii) 記憶湯(Memory Soup),靈感來自權重湯(weight souping),對快取記憶模組的參數進行平均(適用於非線性記憶);以及 (iv) 稀疏選擇性快取(SSC),使用類似混合專家(Mixture-of-Experts)的路由器,僅選擇與上下文最相關的快取記憶進行高效聚合。
  • 實證驗證: 作為概念驗證,我們在三种架構上展示了 MC 的有效性:線性注意力(Linear Attention, LA)(Katharopoulos et al., 2020)、深度記憶模組 Titans (Behrouz et al., 2025c)、滑動視窗線性注意力(SWLA)及深度線性注意力(DLA)(Behrouz et al., 2025a)。橫跨語言建模、長上下文及檢索任務的實驗結果顯示,MC 提升了效能並擴展了 RNN 的有效上下文長度。

2. 預備知識與背景

本節回顧必要的背景知識並建立符號系統。特別是,我們將回顧注意力及其線性變體的概念,接著討論參數化上下文學習與嵌套學習範式(nested learning paradigm)(Behrouz et al., 2026); (2025b),這是我們構建記憶體快取的基礎。

符號。 我們使用粗體小寫(分別為大寫)字母表示向量(分別為矩陣),並使用下標 t 指代時間 t 時的實體狀態。文中,設 xL×din 為輸入,t 為時間 t 時的記憶狀態 ()𝐊 為鍵(keys),𝐕 為值(values),𝐐 為查詢矩陣(query matrices),L 表示序列長度。我們專注於基於多層感知器(MLP)的記憶架構,其層數 1。值得注意的是,當 =1 時,此公式包含了線性矩陣值記憶模組。如有需要,我們將記憶模組 () 參數化為 θ:={W1,,W,},其中至少包含 MLP 中線性層的參數。

注意力機制。 注意力 (Vaswani et al., 2017) 是 Transformer 的主要構件,充當其關聯記憶 (Bietti et al., 2023); (Behrouz et al., 2026); (Wang et al., 2025)。給定輸入 xL×din,因果注意力(causal attention)根據輸入依賴的鍵、值和查詢矩陣 𝐐=x𝐖𝐐, 𝐊=x𝐖𝐊, 及 𝐕=x𝐖𝐕 計算輸出 𝐲L×din 如下:

𝐲i=t=1iexp(𝐪i𝐤t)𝐯t=1iexp(𝐪i𝐤)=1Zit=1iexp(𝐪i𝐤t)𝐯t公式 (1)

其中 𝐖𝐐, 𝐖𝐊, 和 𝐖𝐕din×din 為可學習參數,Zi 為歸一化項。由於需要存取所有過去的 token,注意力機制需要 O(L2) 的運算量。

線性注意力。 線性注意力 (Katharopoulos et al., 2020) 及其變體 (Schlag et al., 2021); (Peng et al., 2023); (Yang et al., 2024b) 透過將公式 (1) 中的 exp() 運算子替換為可分離核函數 ϕ() 來提高效率,從而產生高效的循環公式:

𝐲i=t=1iϕ(𝐪i))ϕ(𝐤t)𝐯t=1iϕ(𝐪i))ϕ(𝐤)=1Ziiϕ(𝐪i)公式 (2)

其中 t=t1+𝐯tϕ(𝐤t) 充當固定大小的記憶 (Katharopoulos et al., 2020)

測試時記憶化與嵌套學習觀點。 近期的一個統一框架將序列模型(包括注意力和現代 RNN)的更新規則解釋為具有不同目標的動態上下文學習/記憶化過程 (Behrouz et al., 2026); (2025b)。在此觀點下,模型作為關聯記憶,主動學習輸入 token(鍵與值)之間的映射。這種記憶化是通過優化內部目標實現的,通常形式化為 L2 回歸問題 (Wang et al., 2025) 或具有「注意力偏見」(attentional bias)的一般目標 (Behrouz et al., 2026); (2025b)。此觀點將記憶狀態視為在前向傳播過程中優化的動態實體。特別是在 Miras 框架的最簡單形式中 (Behrouz et al., 2026),關聯記憶 () 旨在基於某個目標(稱為「注意力偏見」)學習鍵 {𝒌t}t=1L 與值 {𝒗t}t=1L 之間的映射:

t+1=argmin((𝒌t);𝒗t)+Ret(;t)公式 (3)

其中目標函數 () 衡量映射的品質,而 Ret(;t) 則確保新解與上一個記憶狀態保持接近。針對特定的注意力偏見選擇,可以還原出已知的架構:例如,令 ((𝒌t);𝒗t)=(𝒌t),𝒗t()d×d,即可還原出未歸一化的線性注意力架構 (Katharopoulos et al., 2020)。我們利用此觀點引入記憶體快取,其中快取狀態作為此優化過程的檢查點,增強了模型在長序列中檢索資訊的能力。

3. 具備記憶體快取的循環神經網路

RNN 維持固定大小的記憶來壓縮輸入序列。隨著序列變長,這會導致記憶溢位與效能下降。相對地,注意力機制快取所有過去的 token,導致記憶增長但帶來二次方成本。我們提出記憶體快取(MC)來快取中間記憶狀態,提供一個中間地帶,使模型的記憶能隨規模任意增長。這使得計算成本能在 O(L)(類似 RNN)與 O(L2)(類似 Transformer)之間進行插值。為此,給定一序列 token xL×din,我們將序列分割為片段 S(1),,S(N),大小為 L(1),,L(N),並使用記憶 (1),,(N) 來壓縮這些片段。對應於第 s 個片段的記憶更新規則或循環如下:

𝒌t=xtW𝒌,𝒗t=xtW𝒗,𝒒t=xtW𝒒公式 (4a)
t(s)=f(t1(s);𝒌t,𝒗t),其中1tL(s)公式 (4b)

其中 f() 是學習更新規則(例如,對於線性注意力 (Katharopoulos et al., 2020)f(t1(s);𝒌t,𝒗t)=t1(s)+𝒗t𝒌t)。使用上述公式,在更新記憶後,我們快取每個片段的最後狀態(即 {L(s)(s)}s=1T,其中 T 是當前片段的索引,xtS(T))。標準 RNN 僅使用當前記憶狀態計算輸出:𝐲t=t(𝒒t)。相較之下,我們的公式利用所有快取記憶以及當前記憶(在線記憶)來計算查詢 𝒒t 的輸出。給定任意聚合函數 Agg(;;),輸出為:

𝐲t=Agg({L(1)(1)(),,L(s1)(s1)()};t(s)();𝐪t)公式 (5)

其中 s 是當前片段的索引。請注意,對於 1is,項 L(i)(i)(𝐪t) 提供了第 i 個片段中對應於查詢 𝒒t 的資訊。在接下來的章節中,我們將介紹 Agg(;;) 函數的不同有效選擇,以將過去資訊納入當前輸出的計算,並增加模型的有效記憶容量。

3.1 殘差記憶 (Residual Memory)

我們從最簡單的 Agg(;;) 運算子開始:作為記憶狀態間的殘差連接的總和。在這種情況下,給定鍵、值和查詢(參見公式 4)以及片段 S(1),,S(N),我們將時間 t(位於片段 s 中)的記憶更新與輸出計算定義為:

t(s)=f(t1(s);𝒌t,𝒗t),其中1tL(s)公式 (6)
𝐲t=t(s)(𝐪t)在線記憶+i=1s1L(i)(i)(𝐪t)快取記憶公式 (7)

記憶體快取的關鍵變革在於輸出的計算方式。事實上,為了檢索記憶,模型對當前記憶(稱為在線記憶)和快取記憶都進行了針對輸入查詢 𝒒t 的前向傳播。

閘道殘差記憶 (Gated Residual Memory, GRM)。 當記憶模組嚴格為線性時(即 為矩陣),殘差記憶公式(公式 7)在數學上會坍縮為標準的固定大小記憶,因為快取記憶可以預先加總(參見下文公式 13)。然而,實驗結果顯示,即使是這種簡單的公式也能增強循環模型的能力(參見第 5 節)。主要原因是簡單的殘差記憶充當了保留算子,增強了對長久過去資訊的存取。殘差方法的另一個限制是它同等對待所有快取記憶,忽略了它們與查詢 𝒒t 的相關性。為了實現選擇性檢索,我們引入了輸入依賴的閘道機制。給定片段 s 中的輸入 xt,我們定義參數 0γt(1),,γt(s)1 為輸入依賴參數,並將輸出重新表述為:

t(s)=f(t1(s);𝒌t,𝒗t),for1tL(s)公式 (8)
𝐲t=γt(s)t(s)(𝐪t)+i=1s1γt(i)L(i)(i)(𝐪t)公式 (9)

此處,參數 γt(i) 調節每個片段對輸出的貢獻。當 γt(i)1(分別為 0)時,第 i 個片段對輸出的貢獻較大(分別為較小)。由於這些輸入依賴參數的存在,上述公式無法在此 token 之前預先計算,也無法複用於下一個 token/片段。因此,與先前的變體不同,它不會坍縮為固定大小的記憶情況(即使是在線性記憶的情況下),因此需要為每個 token 重新計算並快取記憶狀態。對於 γt(i) 的一個簡單參數化選擇是將其定義為輸入 xt 的線性投影(類似於鍵、值和查詢的投影)。然而,在這種參數化下,γt(i) 充當基於位置的過濾/聚焦,意味著 xt 的上下文僅決定第 i 個片段的記憶(基於位置)貢獻多少,無論其上下文為何。為了克服這個問題,我們建議使 γt(i) 成為 xt 和第 i 個片段 S(i) 的函數,結合兩者的上下文及其相似度。為此,我們引入連接器參數 𝒖t 作為輸入的線性投影,並將 γt(i) 定義為 𝒖t 與第 i 個片段 S(i) 的相似度:

γt(i)=𝒖t,MeanPooling(S(i))where𝒖t=xtW𝒖公式 (10)

此處,MeanPooling() 提供了片段上下文的簡單表示,即所有 token 的平均值。然而,它可以被任何其他池化過程替換。在實踐中,我們還使用 softmax(·)γt(i) 進行歸一化。作為替代參數化,我們可以使用 𝒖t=𝒒t。當 γt(i) 為常數時,GRM 等價於殘差記憶變體。

範例。 為了更好地說明上述公式,我們舉一個說明性範例。設 f(t1(s);𝒌t,𝒗t)=t1(s)t1(s)(𝒌t),𝒗t,其中記憶 () 是任意前饋層(例如 MLP 或閘道 MLP 層)。這種通用形式等價於深度線性注意力(DLA)(Behrouz et al., 2025a),而當記憶為矩陣時(即單層 MLP),它等價於線性注意力 (Katharopoulos et al., 2020)。對 DLA 使用殘差記憶快取會產生具有以下更新和檢索規則的模型:

滑動視窗線性注意力(Sliding Window Linear Attention, SWLA)。近期,Behrouz 等人(2025a)提出了滑動視窗線性注意力(SWLA),其記憶權重是根據一組 c1 個過去 token 進行更新(這與線上 RNN 不同,後者僅根據最後一個 token 更新記憶)。更具體地說,給定一個記憶模組 () 以及鍵、值和查詢 {(𝒌t,𝒗t,𝒒t)}t=1L,其更新與檢索規則定義如下:

t=αtt1+i=tc+1tβi(t)𝒗i𝒌i,(26)
𝐲t=t𝒒t,(27)

c=1 時,此設計會退化為簡單的線性注意力(線上線性 RNN)及其閘道變體(Katharopoulos 等人,2020;Sun 等人,2023;Li 等人,2025)。作為概念驗證,我們對 c=2 的 SWLA 使用記憶快取,產生的遞迴與檢索公式為:

t(s)=αtt1(s)+(βt𝒗t1𝒌t1+λt𝒗t𝒌t),(28)
𝐲t=γt(s)t(s)𝒒t+i=1s1γt(i)L(i)(i)𝒒t.(29)

請注意,如前所述,由於 SWLA 是線性記憶模組,因此 GRM 和 Memory Soup 變體最終會導出相同的公式。

表 1:模型在語言建模與常識推理任務上的表現。
模型Wiki.
ppl
LMB.
ppl
LMB.
acc
PIQA
acc
Hella.
acc_n
Wino.
acc
ARC-e
acc
ARC-c
acc_n
SIQA
acc
BoolQ
acc
Avg.
7.6 億參數 / 300 億 token
Transformer++24.1824.2736.367.241.852.065.633.439.161.749.64
Samba21.0722.8539.268.947.853.165.834.938.963.151.46
RetNet25.7724.1934.566.841.251.963.632.538.856.248.19
DeltaNet24.5224.3836.867.344.551.864.232.739.660.149.63
RWKV-723.7523.0837.167.347.652.264.734.239.461.950.55
Miras (Memora)22.2822.3138.267.849.353.363.636.140.963.051.53
SWLA23.8322.7436.566.944.154.964.234.139.660.150.05
+ Log-Linear++23.3722.1936.967.344.755.064.934.639.460.450.40
+ GRM (= Soup)22.8121.5037.868.345.855.065.436.240.661.051.26
+ SSC23.0622.3937.267.945.254.965.235.539.860.650.79
DLA23.1222.0936.168.047.952.765.834.639.159.650.48
+ Log-Linear++23.0821.1536.868.147.753.065.635.139.259.350.60
+ GRM22.9120.1037.569.248.752.866.136.840.359.951.41
+ Memory Soup22.7820.4937.269.648.353.465.836.539.660.251.33
+ SSC23.1420.8637.068.447.752.766.035.239.760.150.85
Titans (LMM)20.0421.9637.469.348.552.366.335.840.162.851.56
+ Log-Linear++19.7920.6237.870.148.052.566.835.640.362.851.74
+ GRM19.1420.2138.370.648.454.067.536.441.763.552.55
+ Memory Soup19.5220.3838.071.448.653.767.135.441.363.152.33
+ SSC19.3920.4637.770.948.753.566.936.341.263.152.29
13 億參數 / 1000 億 token
Transformer++17.9217.7342.671.451.354.169.936.041.858.453.19
Samba16.1513.2145.271.553.855.869.136.740.663.054.46
RetNet18.9117.0441.271.349.155.267.534.141.461.052.60
DeltaNet18.6217.1041.670.149.452.767.635.239.754.851.39
Miras (Memora)15.9012.0448.773.156.057.471.537.940.261.355.76
SWLA18.4716.2339.470.948.856.567.335.841.560.252.55
+ Log-Linear++18.6716.0939.971.249.356.668.136.341.460.452.90
+ GRM (= Soup)18.5115.9540.672.650.557.869.540.842.862.254.60
+ SSC18.6116.0140.471.950.057.168.938.642.261.253.79
DLA16.3112.2944.570.653.954.269.636.040.860.253.72
+ Log-Linear++16.2212.2544.971.154.554.870.036.641.360.754.24
+ GRM16.0812.1045.872.555.955.871.541.242.862.255.96
+ Memory Soup16.1612.1745.671.955.455.670.937.742.061.555.08
+ SSC16.2012.1945.371.754.855.370.437.141.461.154.64
Titans (LMM)15.6011.4149.173.156.359.872.440.842.161.056.82
+ Log-Linear++15.4911.3849.473.656.560.372.841.142.561.357.19
+ GRM15.3711.2950.474.557.461.573.842.643.962.558.33
+ Memory Soup15.4211.3149.974.257.360.873.542.243.462.057.91
+ SSC15.4411.3549.673.857.060.673.141.942.861.857.58
代表注意力機制與線性 RNN 的混合架構 (Ren 等人,2024)。

深度線性注意力(Deep Linear Attention, DLA)。DLA 使用與線性注意力相同的更新規則(即赫布式規則),但採用深度記憶模組。也就是說,給定一個記憶模組 () 以及鍵、值和查詢 {(𝒌t,𝒗t,𝒒t)}t=1L,其更新與檢索規則定義如下:

t=t1ηt(t1;𝒌t,𝒗t),(30)
𝐲t=t(𝒒t),(31)

其中注意力偏置目標函數定義為 (t1;𝒌t,𝒗t)=t1(𝒌t),𝒗t。使用記憶快取(GRM 變體)時,DLA 的更新與檢索過程定義為:

t(s)=t1(s)ηt(t1(s);𝒌t,𝒗t),for 1tL(s),(32)
𝐲t=γt(s)t(s)(𝐪t)+i=1s1γt(i)L(i)(i)(𝐪t).(33)

同樣地,我們可以將公式 14公式 17替換為 DLA 的更新規則(類似於公式 32),以推導出其他記憶快取變體。請注意,當記憶模組 () 為矩陣時,上述公式即等價於線性注意力(Katharopoulos 等人,2020)。

Titans。在 Titans 中,與 DLA 相比,其注意力偏置目標函數以及內部優化器皆有所不同。更具體地說,給定一個記憶模組 () 以及鍵、值和查詢 {(𝒌t,𝒗t,𝒒t)}t=1L,Titans 的更新與檢索規則定義為:

t=αtt1𝒮t,(34)

其中,注意力偏置目標函數定義為 (t1;𝒌t,𝒗t)=t1(𝒌t)𝒗t22\mathcal{L}\left(\mathcal{M}_{t-1};{\bm{k}}_{t},{\bm{v}}_{t}\right)=\|\mathcal{M}_{t-1}({\bm{k}}_{t})-{\bm{v}}_{t}\|^{2}_{2}。在使用記憶快取(MC)時,雖然每個區段的記憶更新運算與 公式 34公式 35 相同,但 Titans 搭配記憶快取的檢索過程則定義為與 公式 33 相同。

Log-Linear++ 變體。 最近,Guo 等人 (2025) 利用線性 RNN 的結構化矩陣形式,設計了 Log-linear Attention。這是一種基於 Fenwick 樹結構 (Fenwick, 1994) 的分層演算法,允許隱藏狀態以對數方式增長。我們旨在將 Log-linear Attention 作為實驗中的基準,以展示分段(Segmentation)對效率與檢索效能的影響。然而,其公式存在我們先在 第 3.1 節 討論過的位置偏誤(Positional Bias)以及檢索過程缺乏上下文依賴性(Context-dependency)等問題。為了公平起見,我們改進了 Log-linear 公式(在實驗中標記為 Log-Linear++),將其重新表述為一種具有 GRM 與對數規模區段集合的記憶快取變體。其分段過程與 第 4.2 節 所述過程相同,其他組成部分則保持不變,與我們的記憶快取變體一致。

作為訓練後處理的記憶快取。 記憶快取亦可應用於模型預訓練之後;在推論階段,我們快取每個區段(例如:訓練序列長度)結束後的記憶狀態。在解碼時,我們使用過去快取記憶的移動平均,且不涉及可學習權重。在我們的實驗結果中觀察到,即便是這種簡單的技術,也能顯著增強循環模型(Recurrent Models)的長度外推(Length Extrapolation)能力。

表 2:NIAH(大海撈針)實驗,包含三種難度等級:單針任務——S-NIAH-1(密鑰檢索)、S-NIAH-2(數字針)、以及 S-NIAH-3(UUID 針)。
S-NIAH-1S-NIAH-2S-NIAH-3
(pass-key retrieval)(number in haystack)(uuid in haystack)
模型4K8K16K4K8K16K4K8K16K
Transformer88.676.479.810098.894.278.069.240.8
DLA96.471.244.079.642.628.218.28.84.0
+ Log-Linear++10096.270.487.670.418.028.820.46.0
+ GRM10010082.494.682.854.848.234.418.2
+ Memory Soup10010078.291.877.240.443.032.814.8
+ SSC10098.276.889.274.837.634.028.611.2
Titans (LMM)10010010099.684.675.474.242.821.2
+ Log-Linear++10010010095.688.474.876.048.424.2
+ GRM10010010099.896.688.289.469.032.2
+ Memory Soup10010010098.892.283.084.261.828.6
+ SSC10010010098.690.479.681.054.227.0

5 實驗

接下來,我們評估記憶快藏在提升模型於語言建模、常識推理、大海撈針(Needle-in-a-Haystack)以及上下文召回(In-context Recall)等任務上的成效。

實驗設定。 在我們的實驗評估中,主要遵循 Guo 等人 (2025) 的方法。我們使用 FineWeb 數據集 (Penedo 等人,2024) 與 Long-Data-Collections (Together AI, 2024) 的混合數據訓練模型,訓練上下文視窗大小為 {2K, 4K, 8K, 16K, 32K},區段長度範圍為 {16, 32, 64, 128, 256, 512} 個 token。在語言建模與常識推理任務(表 1)中,預設模型使用 4K 上下文長度與 256 區段長度進行訓練。我們使用了 760M 與 1.3B 參數量的模型,並在從 FineWeb 數據集中採樣的 300 億與 1000 億 token 上進行訓練。困惑度(Perplexity)是在保留的驗證數據上測量。至於下游任務,我們在 Wikitext (Merity 等人,2017)、LMB (Paperno 等人,2016)、PIQA (Bisk 等人,2020)、HellaSwag (Zellers 等人,2019)、WinoGrande (Sakaguchi 等人,2021)、ARC-easy (ARC-e) 與 ARC-challenge (ARC-c) (Clark 等人,2018)、SIQA (Sap 等人,2019) 以及 BoolQ (Clark 等人,2019) 上評估訓練好的模型。在其他下游任務如大海撈針、上下文檢索與 LongBench 中,我們使用 16K 上下文長度訓練模型,以便更好地區分模型在短上下文與長上下文上的表現。有關實驗設定與其他使用數據集的更多細節請參閱 附錄 B

5.1 語言建模

我們從常見的學術規模語言建模開始。SWLA、DLA 與 Titans 在有无記憶快藏下的結果列於 表 1。觀察結果如下:(1) 比較 DLA、Titans 與 SWLA 及其加入記憶快藏的增強版本,我們發現所有記憶快藏變體在不同下游任務上均提供一致的提升,且平均表現優於基準模型。這顯示了記憶快藏對於進一步增強受限記憶體模型的重要性。(2) 如前所述,記憶快藏可視為(稀疏)注意力機制與循環模型的混合體。比較記憶快藏增強模型與基於注意力的模型(即混合架構與 Transformers),記憶快藏為循環模型的記憶體限制問題提供了更強大的解決方案。特別是,Titans + MC 與 DLA + MC 相比 Titans 基線提升了 0.8% 的效能。(3) 比較 MC 的定長分段與 Log-Linear++ 方法,我們觀察到定長分段變體提供了更好的結果。此外,在我們提供的方法中,GRM 接著是 SSC 表現最佳。我們將此效能提升歸因於 MC 為模型提供了更大的有效記憶容量。

表 3:輸入截斷至不同長度時的檢索任務準確率。
SWDESQuADFDA
模型5121024204816k5121024204816k5121024204816k
Transformer46.243.744.444.033.133.333.633.471.069.571.671.0
Titans (MAL)51.948.648.348.528.329.229.128.871.173.972.171.7
DLA44.539.932.732.523.824.023.824.155.640.225.923.3
+ Log-Linear++43.737.730.430.627.827.827.928.355.139.622.318.9
+ GRM52.448.948.748.529.530.730.730.163.351.648.941.5
+ Memory Soup49.545.038.037.728.428.628.529.160.548.437.234.6
+ SSC47.042.535.535.326.028.127.128.858.046.028.829.4
Titans (LMM)43.234.429.229.725.726.226.325.659.345.535.432.5
+ Log-Linear++48.041.437.237.027.227.327.227.167.055.541.232.4
+ GRM52.649.349.550.129.730.431.532.072.968.761.152.6
+ Memory Soup50.346.744.845.429.229.729.830.370.363.855.745.8
+ SSC48.644.241.041.428.328.828.528.868.259.447.638.9
TriviaQADropNQAvg.
模型5121024204816k5121024204816k51210242048
Transformer47.548.547.447.621.822.021.521.423.623.123.741.00
Titans (MAL)44.845.144.644.820.620.520.820.922.122.422.540.46
DLA43.344.243.543.220.119.920.620.019.718.418.530.51
+ Log-Linear++43.744.843.643.820.320.220.820.219.918.821.030.75
+ GRM50.147.344.850.021.921.822.021.723.523.323.438.03
+ Memory Soup48.046.444.248.721.521.321.721.222.822.422.535.05
+ SSC45.845.543.946.120.920.721.220.621.420.621.833.09
Titans (LMM)44.244.743.944.520.220.120.320.620.119.519.131.75
+ Log-Linear++44.544.944.144.720.420.420.520.721.519.820.434.37
+ GRM50.247.545.350.921.721.821.921.523.723.423.340.50
+ Memory Soup48.346.644.849.421.321.421.721.122.922.222.538.43
+ SSC46.145.744.346.920.820.721.220.921.920.421.536.27

5.2 大海撈針(Needle-In-A-Haystack)任務

我們使用大海撈針(NIAH)任務(表 2)評估 MC 對長上下文檢索的影響。經 MC 增強的 DLA 與 Titans 持續超越基準模型。此外,MC 變體在較長上下文中表現優於 Log-Linear 方法。Log-Linear 表現掙扎的原因在於它強迫單一記憶壓縮非常大的初始區段(例如在 16K 序列中的前 8K token),而 MC 則能更有效地分攤壓縮負載。

5.3 上下文檢索(In-context Retrieval)任務

上下文召回任務是循環神經網絡最具挑戰性的基準之一。在本節中,我們遵循 Arora 等人 (2024b) 的方法,在 SWDE (Lockard 等人,2019)、NQ (Kwiatkowski 等人,2019)、DROP (Dua 等人,2019)、FDA (Arora 等人,2023)、SQUAD (Rajpurkar 等人,2016) 與 TQA (Kembhavi 等人,2017) 上進行實驗,以評估並比較 MC 增強變體與基準模型及 Transformers 的表現。結果列於 表 3。雖然 Transformer 在上下文召回任務中仍取得最佳結果,但我們的 MC 變體展現了具競爭力的效能,縮小了與 Transformer 的差距,且表現優於最先進的循環模型。我們再次將此效能提升歸因於隨序列長度擴展的更大記憶容量。

表 4:LongBench 任務準確率 (Bai 等人,2024):包含 NarrativeQA, QasperQA, MultiFieldQA, HotpotQA, 2WikiMultiQA, Musique, GovReport, QMSum, MultiNews, TREC, TriviaQA, SamSum, LCC, 與 RepoBench-P。
Single-Doc QAMulti-Doc QASummarizationFew-shotCode
模型NQAQQAMFQHQA2WMMusGvRQMSMNsTRCTQASSMLCCRBP
Transformer11.59.619.121.528.96.513.09.23.127.227.915.122.929.1
DLA9.417.512.111.822.34.89.57.45.14.823.59.738.434.9
+ Log-Linear++10.110.217.112.423.35.56.612.75.818.624.716.231.631.0
+ GRM11.610.319.818.226.96.413.514.16.925.728.218.332.733.9

表格 4(續):LongBench 任務準確率 (Bai 等人,2024):包含 NarrativeQA, QasperQA, MultiFieldQA, HotpotQA, 2WikiMultiQA, Musique, GovReport, QMSum, MultiNews, TREC, TriviaQA, SamSum, LCC, 與 RepoBench-P。

模型Single-Doc QAMulti-Doc QASummarizationFew-shotCode
NQAQQAMFQHQA2WMMusGvRQMSMNsTRCTQASSMLCCRBP
+ Memory Soup11.210.319.516.725.16.311.213.86.222.526.917.732.333.5
+ SSC10.710.218.814.224.85.98.412.96.120.525.716.831.932.6
Titans (LMM)8.712.518.415.626.16.710.512.611.837.126.224.531.331.4
+ Log-Linear++9.68.919.318.726.96.86.712.92.811.242.725.029.529.7
+ GRM11.89.419.921.429.17.28.413.33.114.849.725.531.032.8
+ Memory Soup10.79.219.620.228.27.17.813.13.013.747.125.330.831.4
+ SSC9.99.119.419.827.56.97.113.02.812.544.825.229.930.8

5.4 長上下文理解(Long Context Understanding)任務

我們使用 LongBench (Bai 等人,2024) 進行長上下文理解任務的實驗。結果列於表 4。所有經 MC 增強的變體相較於其基礎 RNN 皆提供了效能提升,這再次歸因於其增加的記憶容量。

記憶體快取變體與基準模型的訓練吞吐量比較 變體速度比較
圖 4:記憶體快取變體與基準模型的訓練吞吐量比較。
MQAR 平均準確率圖表
圖 5:MQAR 在 5 個種子上的平均準確率。
表 5:MC 的消融研究。MC 的所有設計選擇皆對其有效性有正面貢獻。
模型語言建模 (ppl ↓)C.S. 推理 (acc ↑)檢索 (acc ↑)
Titans (GRM)13.358.340.5
- Context-dependent13.457.433.0
- Gating13.556.932.4
- Linear Memory13.756.334.5
- Shared u and q00.000.000.0
Titans (SSC)13.457.636.3
- Context-dependent13.457.132.6
- Gating13.556.831.9
- Linear Memory13.856.833.4
- Shared u and q00.000.000.0

5.5 多查詢關聯召回(MQAR)

在本節中,我們評估 MC 增強變體在多查詢關聯召回(MQAR)任務 (Arora 等人,2024a) 中的表現。結果列於圖 5。我們的模型相較於其基礎 RNN 以及最先進的循環模型皆展現了良好的效能,與 Atlas (Behrouz 等人,2025a) 等最先進模型相比,在每維度数值的效能表現上達到最佳。

5.6 消融研究(Ablation Studies)

接下來,我們評估 MC 框架中設計選擇的影響。第一個選擇是 γ 應該僅是輸入的函數,還是也應包含區塊上下文的函數。結果列於表 5。此設計選擇顯示出顯著的平均改善。第二個設計是移除閘道(gating)。請注意,若無閘道,設計將退化為殘差記憶。結果顯示,即使是這種簡單的設計也能提升模型效能。最後,在第三個設計中,我們使用線性記憶模組。令人驚訝的是,使用記憶體快取能帶來相對於記憶架構和表達性更強健的效能表現。

5.7 效率(Efficiency)

最後,我們評估我們的變體與基準模型的訓練吞吐量。結果列於圖 4。我們的 MC 變體提供了 Transformer 與 RNN 之間的中間地帶,並且在增加上下文長度時,相較於 Transformer 變得極為高效。這些結果表明,我們的 SSC 變體兼具兩者優點,在我們之前討論的多樣化下游任務中,其表現與其他變體相當或更佳,同時相較於其原始基礎 RNN 變體僅增加極小的開銷。此外,它們在較長序列中顯示出顯著更好的效率。

6 結論

在本文中,我們提出了記憶體快取(Memory Caching, MC),這是一種適用於所有循環神經網絡的簡單技術,它快取一部分記憶狀態,允許後續 token 直接關注其過去相關的 token。我們的實驗顯示,相較於一組基準模型,本方法有所改進。本文中的許多選擇都是為了使最終模型盡可能簡單,以更好地展示記憶體快取概念的效果。然而,在未來的工作中,可以使用更具表達性的池化或路由機制來進一步提升效能。

參考文獻

  • [1] Z. Allen-Zhu. Physics of language models: part 4.1, architecture design and the magic of canon layers. 收錄於 The Thirty-ninth Annual Conference on Neural Information Processing Systems. 引用於:附錄 A
  • S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Re (2024a). Zoology: measuring and improving recall in efficient language models. 收錄於 The Twelfth International Conference on Learning Representations. 外部連結:Link. 引用於:§5.5
  • S. Arora, S. Eyuboglu, M. Zhang, A. Timalsina, S. Alberti, J. Zou, A. Rudra, and C. Re (2024b). Simple linear attention language models balance the recall-throughput tradeoff. 收錄於 Forty-first International Conference on Machine Learning. 外部連結:Link. 引用於:附錄 A, §1, §1, §5.3
  • S. Arora, B. Yang, S. Eyuboglu, A. Narayan, A. Hojel, I. Trummer, and C. Ré (2023). Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv preprint arXiv:2304.09433. 引用於:§5.3
  • Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, Y. Dong, J. Tang, and J. Li (2024). LongBench: a bilingual, multitask benchmark for long context understanding. 收錄於 ACL (1), pp. 3119–3137. 外部連結:Link. 引用於:§5.4, 表 4
  • A. Behrouz, Z. Li, P. Kacham, M. Daliri, Y. Deng, P. Zhong, M. Razaviyayn, and V. Mirrokni (2025a). Atlas: learning to optimally memorize the context at test time. arXiv preprint arXiv:2505.23735. 引用於:附錄 A, 附錄 A, 第 3 項, §3.1, §4.3, §4.3, §5.5
  • A. Behrouz, M. Razaviyayn, P. Zhong, and V. Mirrokni (2025b). Nested learning: the illusion of deep learning architectures. 收錄於 The Thirty-ninth Annual Conference on Neural Information Processing Systems. 外部連結:Link. 引用於:附錄 A, §2, §2, 註腳 1
  • A. Behrouz, M. Razaviyayn, P. Zhong, and V. Mirrokni (2026). It's all connected: a journey through test-time memorization, attentional bias, retention, and online optimization. 收錄於 The Fourteenth International Conference on Learning Representations. 外部連結:Link. 引用於:附錄 A, 附錄 A, 附錄 A, 附錄 B, §1, §2, §2, §2, 註腳 1
  • A. Behrouz, P. Zhong, and V. Mirrokni (2025c). Titans: learning to memorize at test time. 收錄於 The Thirty-ninth Annual Conference on Neural Information Processing Systems. 外部連結:Link. 引用於:附錄 A, 附錄 B, 第 3 項, §1, §4.3
  • A. Bietti, V. Cabannes, D. Bouchacourt, H. Jegou, and L. Bottou (2023). Birth of a transformer: a memory viewpoint. Advances in Neural Information Processing Systems 36, pp. 1560–1588. 引用於:§2
  • A. Bietti, V. Cabannes, D. Bouchacourt, H. Jegou, and L. Bottou (2024). Birth of a transformer: a memory viewpoint. Advances in Neural Information Processing Systems 36. 引用於:§1
  • Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. (2020). Piqa: reasoning about physical commonsense in natural language. 收錄於 Proceedings of the AAAI conference on artificial intelligence, Vol. 34, pp. 7432–7439. 引用於:附錄 B, §5
  • T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. (2020). Language models are few-shot learners. Advances in neural information processing systems 33, pp. 1877–1901. 引用於:§1
  • R. Child, S. Gray, A. Radford, and I. Sutskever (2019). Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509. 引用於:§1
  • C. Clark, K. Lee, M. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova (2019). BoolQ: exploring the surprising difficulty of natural yes/no questions. 收錄於 Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), J. Burstein, C. Doran, and T. Solorio (Eds.), Minneapolis, Minnesota, pp. 2924–2936. 外部連結:Link, Document. 引用於:附錄 B, §5
  • P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord (2018). Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457. 引用於:附錄 B, §5
  • G. Comanici, E. Bieber, M. Schaekermann, I. Pasupat, N. Sachdeva, I. Dhillon, M. Blistein, O. Ram, D. Zhang, E. Rosen, et al. (2025). Gemini 2.5: pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06261. 引用於:§1
  • R. Csordás, C. Potts, C. D. Manning, and A. Geiger (2024). Recurrent neural networks learn to store and generate sequences using non-linear representations. 收錄於 Proceedings of the 7th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, pp. 248–262. 引用於:附錄 A
  • Z. Dai, Z. Yang, Y. Yang, J. G. Carbonell, Q. V. Le, and R. Salakhutdinov (2019). Transformer-xl: attentive language models beyond a fixed-length context. 收錄於 ACL (1), A. Korhonen, D. R. Traum, and L. Màrquez (Eds.), pp. 2978–2988. 外部連結:ISBN 978-1-950737-48-2. 引用於:§1
  • T. Dao, B. Chen, N. S. Sohoni, A. Desai, M. Poli, J. Grogan, A. Liu, A. Rao, A. Rudra, and C. Ré (2022). Monarch: expressive structured matrices for efficient and accurate training. 收錄於 International Conference on Machine Learning, pp. 4690–4721. 引用於:附錄 A
  • T. Dao, A. Gu, M. Eichhorn, A. Rudra, and C. Ré (2019). Learning fast algorithms for linear transforms using butterfly factorizations. 收錄於 International conference on machine learning, pp. 1517–1527. 引用於:附錄 A
  • A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby (2021). An image is worth 16x16 words: transformers for image recognition at scale. 收錄於 International Conference on Learning Representations. 外部連結:Link. 引用於:§1
  • D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner (2019). DROP: a reading comprehension benchmark requiring discrete reasoning over paragraphs. arXiv preprint arXiv:1903.00161. 引用於:§5.3
  • P. M. Fenwick (1994). A new data structure for cumulative frequency tables. Software: Practice and experience 24 (3), pp. 327–336. 引用於:§4.3
  • X. Gonzalez, A. Warrington, J. Smith, and S. Linderman (2024). Towards scalable and stable parallelization of nonlinear rnns. Advances in Neural Information Processing Systems 37, pp. 5817–5849. 引用於:附錄 A
  • H. Guo, S. Yang, T. Goel, E. P. Xing, T. Dao, and Y. Kim (2025). Log-linear attention. arXiv preprint arXiv:2506.04761. 引用於:附錄 A, 附錄 B, §4.3, §5
  • D. O. Hebb (2005). The organization of behavior: a neuropsychological theory. Psychology press. 引用於:附錄 A
  • D. Hendrycks and K. Gimpel (2016). Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415. 引用於:附錄 B
  • J. J. Hopfield (1982). Neural networks and physical systems with emergent collective computational abilities.. Proceedings of the national academy of sciences 79 (8), pp. 2554–2558. 引用於:附錄 A, 附錄 A
  • J. Y. Hu, D. Wu, and H. Liu (2024). Provably optimal memory capacity for modern hopfield models: transformer-compatible dense associative memories as spherical codes. arXiv preprint arXiv:2410.23126. 引用於:附錄 A
  • J. Hu, Y. Pan, J. Du, D. Lan, X. Tang, Q. Wen, Y. Liang, and W. Sun (2025). Improving bilinear RNN with closed-loop control. 收錄於 The Thirty-ninth Annual Conference on Neural Information Processing Systems. 外部連結:Link. 引用於:附錄 A
  • Y. Huang, J. Zhang, Z. Shan, and J. He (2024). Compression represents intelligence linearly. 收錄於 First Conference on Language Modeling. 外部連結:Link. 引用於:§1
  • K. Irie, I. Schlag, R. Csordas, and J. Schmidhuber (2021). Going beyond linear transformers with recurrent fast weight programmers. Advances in neural information processing systems 34, pp. 7703–7717. 引用於:附錄 A, §1
  • K. Irie, I. Schlag, R. Csordás, and J. Schmidhuber (2022). A modern self-referential weight matrix that learns to modify itself. 收錄於 International Conference on Machine Learning, pp. 9660–9677. 引用於:§1
  • K. Jordan, Y. Jin, V. Boza, Y. Jiacheng, F. Cecista, L. Newhouse, and J. Bernstein (2024). Muon: an optimizer for hidden layers in neural networks, 2024b. URL https://kellerjordan.github.io/posts/muon. 引用於:附錄 A
  • J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, et al. (2021). Highly accurate protein structure prediction with alphafold. nature 596 (7873), pp. 583–589. 引用於:§1
  • Y. Kang, G. Tran, and H. De Sterck (2023). Fast multipole attention: a divide-and-conquer attention mechanism for long sequences. arXiv preprint arXiv:2310.11960. 引用於:附錄 A
  • J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. 引用於:§1
  • M. Karami and V. Mirrokni (2025). Lattice: learning to efficiently compress the memory. 引用於:附錄 A
  • A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret (2020). Transformers are rnns: fast autoregressive transformers with linear attention. 收錄於 International conference on machine learning, pp. 5156–5165. 引用於:附錄 A, 第 3 項, §1, §2, §2, §2, §3.1, §3, §4.1, §4.3, §4.3, §4.3
  • A. Kembhavi, M. Seo, D. Schwenk, J. Choi, A. Farhadi, and H. Hajishirzi (2017). Are you smarter than a sixth grader? textbook question answering for multimodal machine comprehension. 收錄於 Proceedings of the IEEE Conference on Computer Vision and Pattern recognition, pp. 4999–5007. 引用於:§5.3
  • N. Kitaev, Ł. Kaiser, and A. Levskaya (2020). Reformer: the efficient transformer. arXiv preprint arXiv:2001.04451. 引用於:附錄 A
  • D. Krotov and J. J. Hopfield (2016). Dense associative memory for pattern recognition. Advances in neural information processing systems 29. 引用於:附錄 A
  • D. Krotov (2021). Hierarchical associative memory. arXiv preprint arXiv:2107.06446. 引用於:附錄 A
  • Y. Kuratov, A. Bulatov, P. Anokhin, I. Rodkin, D. I. Sorokin, A. Sorokin, and M. Burtsev (2024). BABILong: testing the limits of LLMs with long context reasoning-in-a-haystack. 收錄於 The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track. 外部連結:Link. 引用於:§1
  • T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, J. Devlin, K. Lee, et al. (2019). Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics 7, pp. 453–466. 引用於:§5.3
  • A. Li, B. Gong, B. Yang, B. Shan, C. Liu, C. Zhu, C. Zhang, C. Guo, D. Chen, D. Li, et al. (2025). Minimax-01: scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313. 引用於:附錄 A, §4.3
  • S. Li, X. Jin, Y. Xuan, X. Zhou, W. Chen, Y. Wang, and X. Yan (2019). Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. Advances in neural information processing systems 32. 引用於:附錄 A
  • X. Li, Y. Li, Y. Liang, Z. Shi, and Z. Song (2024). On the expressive power of modern hopfield networks. arXiv preprint arXiv:2412.05562. 引用於:附錄 A
  • Y. H. Lim, Q. Zhu, J. Selfridge, and M. F. Kasim (2024). Parallelizing non-linear sequential models over the sequence length. 收錄於 The Twelfth International Conference on Learning Representations. 外部連結:Link. 引用於:附錄 A
  • B. Liu, R. Wang, L. Wu, Y. Feng, P. Stone, and Q. Liu (2024). Longhorn: state space models are amortized online learners. arXiv preprint arXiv:2407.14207. 引用於:附錄 A
  • C. Lockard, P. Shiralkar, and X. L. Dong (2019). Openceres: when open information extraction meets the semi-structured web. 收錄於 Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 3047–3056. 引用於:§5.3
  • E. Lu, Z. Jiang, J. Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang, Z. Huang, H. Yuan, S. Xu, X. Xu, G. Lai, Y. Chen, H. Zheng, J. Yan, J. Su, Y. Wu, Y. Zhang, Z. Yang, X. Zhou, M. Zhang, and J. Qiu (2025). MoBA: mixture of block attention for long-context LLMs. 收錄於 The Thirty-ninth Annual Conference on Neural Information Processing Systems. 外部連結:Link. 引用於:附錄 A
  • C. Lucibello and M. Mézard (2024). Exponential capacity of dense associative memories. Physical Review Letters 132 (7), pp. 077301. 引用於:附錄 A
  • S. Merity, C. Xiong, J. Bradbury, and R. Socher (2017). Pointer sentinel mixture models. 收錄於 International Conference on Learning Representations. 外部連結:Link. 引用於:附錄 B, §5
  • W. Merrill, J. Petty, and A. Sabharwal (2024). The illusion of state in state-space models. 收錄於 Forty-first International Conference on Machine Learning. 外部連結:Link. 引用於:附錄 A, §1
  • T. Munkhdalai, M. Faruqui, and S. Gopal (2024). Leave no context behind: efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143. 引用於:附錄 A
  • T. Munkhdalai, A. Sordoni, T. Wang, and A. Trischler (2019). Metalearned neural memory. Advances in Neural Information Processing Systems 32. 引用於:附錄 A
  • T. Munkhdalai and H. Yu (2017). Neural semantic encoders. 收錄於 Proceedings of the conference. Association for Computational Linguistics. Meeting, Vol. 1, pp. 397. 引用於:附錄 A
  • T. Nguyen, V. Suliafu, S. Osher, L. Chen, and B. Wang (2021). Fmmformer: efficient and flexible transformer via decomposed near-field and far-field attention. Advances in neural information processing systems 34, pp. 29449–29463. 引用於:附錄 A
  • D. Paperno, G. Kruszewski, A. Lazaridou, N. Q. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernandez (2016). The LAMBADA dataset: word prediction requiring a broad discourse context. 收錄於 Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), K. Erk and N. A. Smith (Eds.), Berlin, Germany, pp. 1525–1534. 外部連結:Link, Document. 引用於:附錄 B, §5
  • Y. Park, M. Seo, and H. Jeon (2025). VideoTitans: scalable video prediction with integrated short- and long-term memory. 收錄於 The Thirty-ninth Annual Conference on Neural Information Processing Systems. 外部連結:Link. 引用於:§1
  • G. Penedo, H. Kydlíček, A. Lozhkov, M. Mitchell, C. A. Raffel, L. Von Werra, T. Wolf, et al. (2024). The fineweb datasets: decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems 37, pp. 30811–30849. 引用於:§5
  • B. Peng, E. Alcaide, Q. G. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. N. Chung, L. Derczynski, X. Du, M. Grella, K. K. GV, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, J. Lin, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. S. Wind, S. Wozniak, Z. Zhang, Q. Zhou, J. Zhu, and R. Zhu (2023). RWKV: reinventing RNNs for the transformer era. 收錄於 The 2023 Conference on Empirical Methods in Natural Language Processing. 外部連結:Link. 引用於:附錄 A, §2

本段摘要

本段深入探討記憶體快取(MC)在長上下文理解、多查詢關聯召回(MQAR)等任務中的實驗結果,並通過消融研究驗證了各項設計選擇(如上下文依賴性、閘道機制、線性記憶模組)的有效性。研究顯示,MC 變體在保持高效率的同時,顯著提升了模型在長序列任務中的表現,特別是在處理長上下文時優於傳統 Transformer 模型。此外,本段亦提供了完整的實驗數據表格與相關文獻引用,佐證 MC 技術的實用性與擴展性。

B. Peng, D. Goldstein, Q. Anthony, A. Albalak, E. Alcaide, S. Biderman, E. Cheah, X. Du, T. Ferdinan, H. Hou, 等人 (2024). Eagle and finch: rwkv with matrix-valued states and dynamic recurrence. arXiv preprint arXiv:2404.05892. 引用於:附錄 A.

M. Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré (2023). Hyena hierarchy: towards larger convolutional language models. 收錄於 International Conference on Machine Learning, pp. 28043–28078. 引用於:§1.

D. Prados and S. Kak (1989). Neural network capacity using delta rule. Electronics Letters 25 (3), pp. 197–199. 引用於:附錄 A.

S. Qiu, A. Potapczynski, M. Finzi, M. Goldblum, and A. G. Wilson (2024). Compute better spent: replacing dense layers with structured matrices. arXiv preprint arXiv:2406.06248. 引用於:附錄 A.

P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang (2016). Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250. 引用於:§5.3.

H. Ramsauer, B. Schäfl, J. Lehner, P. Seidl, M. Widrich, L. Gruber, M. Holzleitner, T. Adler, D. Kreil, M. K. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter (2021). Hopfield networks is all you need. 收錄於 International Conference on Learning Representations. 外部連結:Link. 引用於:附錄 A, §1.

L. Ren, Y. Liu, Y. Lu, Y. Shen, C. Liang, and W. Chen (2024). Samba: simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522. 引用於:表 1.

K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi (2021). Winogrande: an adversarial winograd schema challenge at scale. Communications of the ACM 64 (9), pp. 99–106. 引用於:附錄 B, §5.

M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi (2019). Social IQa: commonsense reasoning about social interactions. 收錄於 Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), K. Inui, J. Jiang, V. Ng, and X. Wan (Eds.), Hong Kong, China, pp. 4463–4473. 外部連結:Link, Document. 引用於:附錄 B, §5.

I. Schlag, K. Irie, and J. Schmidhuber (2021). Linear transformers are secretly fast weight programmers. 收錄於 International Conference on Machine Learning, pp. 9355–9366. 引用於:附錄 A, 附錄 A, §2.

J. Schmidhuber (1992). Learning to control fast-weight memories: an alternative to recurrent nets. accepted for publication in. Neural Computation. 引用於:附錄 A, 附錄 A.

J. Schmidhuber (1993). Reducing the ratio between learning complexity and number of time varying variables in fully recurrent nets. 收錄於 ICANN'93: Proceedings of the International Conference on Artificial Neural Networks Amsterdam, The Netherlands 13–16 September 1993 3, pp. 460–463. 引用於:附錄 A.

M. Schöne, B. Rahmani, H. Kremer, F. Falck, H. Ballani, and J. Gladrow (2025). Implicit language models are rnns: balancing parallelization and expressivity. arXiv preprint arXiv:2502.07827. 引用於:附錄 A.

N. Shazeer, *. Mirhoseini, *. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean (2017). Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. 收錄於 International Conference on Learning Representations. 外部連結:Link. 引用於:§3.3.

J. Siems, T. Carstensen, A. Zela, F. Hutter, M. Pontil, and R. Grazzi (2025). DeltaProduct: increasing the expressivity of deltanet through products of householders. arXiv preprint arXiv:2502.10297. 引用於:附錄 A.

J. T.H. Smith, A. Warrington, and S. Linderman (2023). Simplified state space layers for sequence modeling. 收錄於 The Eleventh International Conference on Learning Representations. 外部連結:Link. 引用於:附錄 A.

Y. Sun, X. Li, K. Dalal, J. Xu, A. Vikram, G. Zhang, Y. Dubois, X. Chen, X. Wang, S. Koyejo, 等人 (2024). Learning to (learn at test time): rnns with expressive hidden states. arXiv preprint arXiv:2407.04620. 引用於:附錄 A, 附錄 A.

Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei (2023). Retentive network: a successor to transformer for large language models. arXiv preprint arXiv:2307.08621. 引用於:附錄 A, §1, §4.3.

M. Tiezzi, M. Casoni, A. Betti, T. Guidi, M. Gori, and S. Melacci (2024). On the resurgence of recurrent models for long sequences: survey and research opportunities in the transformer era. arXiv preprint arXiv:2402.08132. 引用於:附錄 A.

Together AI (2024). Long data collections. 外部連結:Link. 引用於:§5.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017). Attention is all you need. 收錄於 Advances in Neural Information Processing Systems, I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Eds.), Vol. 30. 外部連結:Link. 引用於:§1, §2.

J. Von Oswald, M. Schlegel, A. Meulemans, S. Kobayashi, E. Niklasson, N. Zucchet, N. Scherrer, N. Miller, M. Sandler, M. Vladymyrov, 等人 (2023). Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858. 引用於:附錄 A.

K. A. Wang, J. Shi, and E. B. Fox (2025). Test-time regression: a unifying framework for designing sequence models with associative memory. arXiv preprint arXiv:2501.12352. 引用於:附錄 A, §2, §2.

M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S. Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, 等人 (2022). Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. 收錄於 International conference on machine learning, pp. 23965–23998. 引用於:§3.2.

S. Yang, J. Kautz, and A. Hatamizadeh (2024a). Gated delta networks: improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464. 引用於:附錄 A, 附錄 A, 附錄 B.

S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim (2024b). Gated linear attention transformers with hardware-efficient training. 收錄於 Forty-first International Conference on Machine Learning. 外部連結:Link. 引用於:附錄 A, §2.

S. Yang, B. Wang, Y. Zhang, Y. Shen, and Y. Kim (2024c). Parallelizing linear transformers with the delta rule over sequence length. Advances in Neural Information Processing Systems 37, pp. 115491–115522. 引用於:附錄 A, 附錄 A.

R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi (2019). HellaSwag: can a machine really finish your sentence?. 收錄於 Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, A. Korhonen, D. Traum, and L. Marquez (Eds.), Florence, Italy, pp. 4791–4800. 外部連結:Link, Document. 引用於:附錄 B, §5.

Z. Zeng, S. Pal, J. Kline, G. M. Fung, and V. Singh (2022). Multi resolution analysis (mra) for approximate self-attention. 收錄於 International conference on machine learning, pp. 25955–25972. 引用於:附錄 A.

T. Zhang, S. Bi, Y. Hong, K. Zhang, F. Luan, S. Yang, K. Sunkavalli, W. Freeman, and H. Tan (2025). Test-time training done right. arXiv preprint arXiv:2505.23884. 引用於:附錄 A, 附錄 B.

H. Zhou, S. Zhang, J. Peng, S. Zhang, J. Li, H. Xiong, and W. Zhang (2021). Informer: beyond efficient transformer for long sequence time-series forecasting. 收錄於 Proceedings of the AAAI conference on artificial intelligence, Vol. 35, pp. 11106–11115. 引用於:附錄 A.

附錄 A 相關文獻

線性記憶模組。 近期的研究致力於緩解 Transformer 在解決複雜問題時面臨的二次方複雜度、上下文長度限制以及表達能力不足等問題,這促使了高效能循環替代方案的發展,這些方案能提供更快的推論與訓練速度 (Tiezzi 等人,2024)。更具體地說,Katharopoulos 等人 (2020) 指出,在計算注意力結果時,若以可分離的核心函數(separable kernel)替換 softmax,即可導出線性注意力的表述,使其能進行遞迴計算。基於此見解,多項研究專注於提升線性注意力的效能,縮小其與二次方 Transformer 之間的差距。為此,RetNet (Sun 等人,2023)、RWKV (Peng 等人,2023)、Lightning Attention (Li 等人,2025) 以及 S5 (Smith 等人,2023) 在線性注意力的公式中引入了遺忘閘道(forget gate)機制。隨後,其他研究進一步調整這些公式,透過使線性注意力架構中既有的遺忘閘道依賴於輸入,以適應需要更具選擇性遺忘的任務 (Yang 等人,2024b; Peng 等人,2024)。與此同時,Schlag 等人 (2021) 提出了 DeltaNet,這是一種基於 Delta 法則(Delta-rule)的循環神經網絡遞迴學習更新替代方案,旨在改善線性注意力模型的記憶管理。後來,多項研究設計了不同的演算法來訓練 Delta 更新規則 (Yang 等人,2024c; Sun 等人,2024; Liu 等人,2024)。此外,基於這些現有技術——涵蓋遺忘閘道、學習演算法及訓練演算法設計——並加以整合,近年來已設計出多種線性注意力模組的變體 (Yang 等人,2024a; c; a; Allen-Zhu; Liu 等人,2024)。最近,Siems 等人 (2025) 透過對每個 token 應用多次更新,增強了 Delta 法則模型,從而提升了狀態追蹤的表達能力。除了線性循環模型外,還有幾項研究探討了具有非線性遞迴但採用線性矩陣值記憶體的 RNN (Csordás 等人,2024; Merrill 等人,2024; Lim 等人,2024; Behrouz 等人,2026; 2025a; Schöne 等人,2025; Karami and Mirrokni, 2025; Von Oswald 等人,2023; Gonzalez 等人,2024; Hu 等人,2025),並著重於加速其訓練過程 (Gonzalez 等人,2024; Lim 等人,2024; Schöne 等人,2025)。

深層記憶模組。 另一條研究路徑則專注於增強記憶模組的容量並改進其學習更新規則。Sun 等人 (2024) 提出了 TTT 層,這是一種快速權重程式(fast-weight program)(Schmidhuber, 1992),其權重更新基於L2-迴歸損失。Sun 等人 (2024) 討論了注意力機制與簡單線性注意力如何作為 TTT 層的特例,但將其他循環神經網絡排除在 TTT 層之外,主要原因是它們無法透過內部 L2-迴歸損失準確還原,而這正是 TTT的定義。Titans (Behrouz 等人,2025c) 建議採用更複雜的優化演算法,並以之取代梯度下降法。作為概念驗證,Titans 使用帶有動量與權重衰減的梯度下降法來優化內部 L2-迴歸損失。基於 TTT 層的公式(即優化內部 L2-迴歸損失),Wang 等人 (2025) 展示了如何近似L2-迴歸損失,進而近似還原其他現代循環神經網絡。基於此見解,Wang 等人 (2025) 提出了一種高階注意力變體,其表達能力優於標準 softmax 注意力。與此同時,Behrouz 等人 (2026) 提出了「測試時記憶化(test-time memorization, TTM)」框架,該框架能根據任意目標在內部學習映射的關聯記憶概念,準確地還原出各種架構。事實上,與限制內部模型為 L2-迴歸損失的 TTT 層 (Sun 等人,2024) 相反,TTM 建議基於關聯記憶概念設計架構,並包含四項設計選擇:(1) 記憶體的架構;(2) 內部目標函數;(3) 內部保留閘道;以及 (4) 內部優化演算法。

遵循此方向,鑑於為內部迴路選擇新的目標函數與優化演算法能促成更具表達力的架構發展 (Behrouz 等人,2026),近期已出現新一代透過改變內部目標函數而產生的架構。Moneta 和 Yaad 分別將 L2-迴歸損失替換為 Lp 和 Huber 損失。Atlas (Behrouz 等人,2025a) 引入了 Omega 學習法則,它不是針對最後一個 token 更新記憶,而是針對過去數據的局部上下文進行更新。該研究進一步建議使用 Muon (Jordan 等人,2024) 作為內部優化器。Zhang 等人 (2025) 將 L2-迴歸損失替換為點積相似度,並建議使用較大的區塊大小以提升訓練效率。最近,為了進一步提升機器學習模型的長期記憶能力,Behrouz 等人 (2025b) 提出了連續記憶系統(Continuum Memory System, CMS),該系統建議不需替換注意力區塊,而是將 Transformer 中單一的靜態 MLP 區塊替換為多個 MLP 區塊,每個區塊根據當前任務需求,以端到端方式按各自頻率進行更新(更新方式與 MLP 區塊相同)。這種由注意力機制與多個動態 MLP 區塊組成的架構稱為 Hope-attention,其展現出比傳統 Transformer 更優異的長上下文理解能力。

快速權重程式與元學習。 將線性層視為鍵值關聯記憶系統的觀點可追溯至 Hopfield 網絡 (Hopfield, 1982)。此概念後來透過快速權重程式(fast weight programmers)的發展得到擴展,即將動態快速程式整合至循環神經網絡中,作為可寫入的記憶儲存庫 (Schlag 等人,2021; Schmidhuber, 1992; 1993)。在此類系統的學習範式中,Hebbian 學習 (Hebb, 2005) 與 Delta 法則 (Prados and Kak, 1989) 最為突出。這兩項法則在文獻中已得到廣泛研究 (Munkhdalai and Yu, 2017; Schmidhuber, 1992; Munkhdalai 等人,2019; Schlag 等人,2021; Irie 等人,2021; Yang 等人,2024c; a)。

Hopfield 網絡。 我們的公式建立在關聯記憶的廣泛概念之上,其目標在於學習鍵值之間的映射。Hopfield (1982) 的開創性工作引入了 Hopfield 網絡,這是最早明確基於關聯記憶的神經架構之一,其形式化是透過最小化能量函數來儲存鍵值對。儘管經典 Hopfield 網絡因向量值記憶容量及其能量函數結構的限制而應用減少,但近期的研究試圖透過多種方法提升其容量 (Krotov, 2021; Li 等人,2024; Krotov and Hopfield, 2016)。特別是,已有研究探索了使用指數核心函數擴展其能量函數 (Krotov and Hopfield, 2016; Lucibello and Mézard, 2024)。此外,現代 Hopfield 網絡與 Transformer 架構之間的關聯性也成為積極探討的議題 (Ramsauer 等人,2021; Hu 等人,2024)。

高效能注意力機制。 除了循環架構外,近期工作亦提出使用結構化矩陣來提升 token 混合與通道混合層的效率。例如,Butterfly 矩陣 (Dao 等人,2019)、Monarch 矩陣 (Dao 等人,2022) 以及 Block Tensor-Train 矩陣 (Qiu 等人,2024) 提供了緊湊且具表達力的參數化方式,減輕了密集投影的計算負擔。其他方法則設計稀疏或混合注意力機制,例如滑動視窗注意力,或結合局部遞迴與選擇性長程連接的模型 (Nguyen 等人,2021; Arora 等人,2024b; Munkhdalai 等人,2024)。另一類方法將注意力的二次方複雜度降低至接近對數線性時間。經典案例包括 Reformer (Kitaev 等人,2020),其使用局部敏感哈希(locality-sensitive hashing)將查詢與鍵進行分群;以及 LogSparse Transformer (Li 等人,2019) 和 Informer (Zhou 等人,2021),它們依賴結構化稀疏模式以提升長序列與時間序列任務的效率。後續研究引入了更精細的設計,例如多解析度注意力 (Zeng 等人,2022),其從粗到細逐步優化注意力分數;以及 Fast Multipole Attention (Kang 等人,2023),其調整快速多極子方法以實現可擴展的長程互動。另一組研究則專注於區塊或 token 級別的稀疏注意力模組。具體而言,Lu 等人 (2025) 提出了 MoBA,建議將序列分塊並在序列維度上執行 MoE。此設計不僅基於注意力模組,且與我們的 MoE 有根本差異:其注意力計算是針對每個區塊和 token 臨時進行的。在此,記憶狀態是預先計算的,無需臨時計算。最近,Guo 等人 (2025) 引入了 Log-Linear Attention,這是一個透過 Fenwick 樹分區組織、具有對數增長隱藏狀態集的線性注意力增強框架。此設計實現了 O(L log L) 的訓練複雜度與 O(log L) 的解碼記憶體需求,同時保留了硬體友善的平行化能力。

表 6: 架構細節。
模型區塊數維度頭數峰值學習率Token 數
760M241536161.25e-330B
1.3B18204887e-4100B

附錄 B 實驗細節

在我們的實驗設置中,我們遵循近期關於循環模型的研究 (Yang 等人,2024a; Behrouz 等人,2025c; 2026; Zhang 等人,2025; Guo 等人,2025),使用了 Wikitext (Merity 等人,2017)、LMB (Paperno 等人,2016)、PIQA (Bisk 等人,2020)、HellaSwag (Zellers 等人,2019)、WinoGrande (Sakaguchi 等人,2021)、ARC-easy (ARC-e) 與 ARC-challenge (ARC-c) (Clark 等人,2018)、SIQA (Sap 等人,2019) 以及 BoolQ (Clark 等人,2019) 等數據集。在訓練方面,我們使用 32K 的詞彙量,訓練長度為 4K 至 32K 個 token。我們採用 AdamW 優化器,學習率為 4e-4,並搭配餘弦退火調度(cosine annealing schedule),批次大小為 0.5M tokens,權重衰減為 0.1。對於記憶架構,除非另有說明,我們使用具有 2 層、擴展係數為 4 且激活函數為 GELU 的 MLP (Hendrycks and Gimpel, 2016)。我們亦在每個區塊的末尾使用殘差連接與層歸一化(layer norm):M(x) = x + W₁σ(W₂x)。

相關文章推薦

分享網址
AINews·AI 新聞聚合平台
© 2026 AINews. All rights reserved.