Transformer與RNN合體，谷歌打下顯存門檻，解鎖超長(zhǎng)上下文

編輯｜澤南

最近，谷歌跟內(nèi)存干上了。

上個(gè)月，谷歌的 TurboQuant 研究曾經(jīng)引發(fā)過(guò)一場(chǎng)行業(yè)地震，其宣稱能直接把大模型最吃顯存的 KV Cache 壓縮幾倍，讓市場(chǎng)開始擔(dān)憂未來(lái)對(duì)內(nèi)存的需求，引發(fā)了一波內(nèi)存股暴跌，后續(xù)還有不小的學(xué)術(shù)爭(zhēng)議。

這個(gè)星期，又有一篇谷歌論文引發(fā)了 AI 圈的關(guān)注，作者表示他們提出的方法解決了大模型處理長(zhǎng)文本時(shí)的「內(nèi)存瓶頸」（又來(lái)了），但實(shí)施的是完全不同的技術(shù)路線。

他們新開一條道路，通過(guò)對(duì)于大模型架構(gòu)的機(jī)制創(chuàng)新，賦予了 RNN 「可生長(zhǎng)的記憶容量」，找到了一種兼顧 Transformer 與 RNN 優(yōu)勢(shì)的新方法

該方法可以讓 AI 處理更長(zhǎng)的文本，解鎖「超長(zhǎng)上下文」能力，與此同時(shí)極大地降低了推理的資源門檻。

人們紛紛表示歡迎：大模型生產(chǎn)環(huán)境要的就是這個(gè)。

當(dāng)前的大模型幾乎都是基于 Transformer 架構(gòu)打造的，它占據(jù)統(tǒng)治地位，主要?dú)w功于其可增長(zhǎng)的記憶容量（注意力機(jī)制的計(jì)算和空間復(fù)雜度隨上下文長(zhǎng)度呈二次方增長(zhǎng)）。這使得 Transformer 極其擅長(zhǎng)長(zhǎng)上下文的信息召回（Recall）。

然而，這種二次方復(fù)雜度也導(dǎo)致了嚴(yán)重的算力和顯存瓶頸，使得處理超長(zhǎng)文本的成本極高。

為解決這個(gè)問(wèn)題，社區(qū)一直在復(fù)興 RNN、線性注意力模型（Linear Attention）以及狀態(tài)空間模型（SSM，如 Mamba 等）。這些循環(huán)架構(gòu)的優(yōu)勢(shì)在于記憶容量固定（復(fù)雜度為），推理速度快且顯存占用低。然而，它們的致命弱點(diǎn)在于：無(wú)論讀了多少文本，都必須把所有過(guò)去的信息壓縮到一個(gè)「固定大小」的隱藏狀態(tài)（Hidden State）中。

這種「信息漏斗」導(dǎo)致它們?cè)诿芗倩厝蝿?wù)（Recall-intensive tasks，比如從極長(zhǎng)的文檔中精準(zhǔn)提取某個(gè)細(xì)節(jié)）中表現(xiàn)往往遠(yuǎn)不如 Transformer。

為打破僵局，來(lái)自 Google Research 的研究團(tuán)隊(duì)提出了一種名為Memory Caching（MC）的技術(shù)，據(jù)說(shuō)簡(jiǎn)單而有效。

• 論文：《Memory Caching: RNNs with Growing Memory》
• 論文鏈接：https://arxiv.org/abs/2602.24281

在這項(xiàng)研究的視角中，存在一個(gè)架構(gòu)光譜：一端是 Transformer（無(wú)壓縮，Token 級(jí)緩存），另一端是傳統(tǒng)的 RNN（全壓縮，單一記憶）。而「記憶緩存（MC）」則解鎖了介于兩者之間的新形態(tài)：將成組 Token 壓縮并緩存到長(zhǎng)期記憶狀態(tài)中，然后在需要時(shí)進(jìn)行檢索。

Transformer 會(huì)緩存每一個(gè)單獨(dú)的標(biāo)記，而 RNN 則考慮一個(gè)固定大小的記憶，并將上下文中的一切壓縮到記憶的參數(shù)中。那么，如果我們把 RNN 的歷史記憶也緩存下來(lái)會(huì)怎樣？

簡(jiǎn)單來(lái)說(shuō)，與其讓 RNN 只維護(hù)一個(gè)不斷被覆蓋和更新的「當(dāng)前狀態(tài)」，不如定期對(duì)這些隱藏狀態(tài)進(jìn)行「快照打卡」（Caching checkpoints）。這樣，在進(jìn)行信息檢索時(shí)，模型不僅可以查看當(dāng)前的「在線記憶」，還能直接調(diào)取「緩存記憶」中的歷史快照，瞬間找回過(guò)去的相關(guān)信息。

在研究的過(guò)程中，作為概念驗(yàn)證，研究人員提出了三種變體，基于過(guò)去的信息如何組合在一起。

1、門控殘差記憶（Gated Residual Memory）：使用查詢從過(guò)去檢索相關(guān)信息，然后執(zhí)行類似注意力的池化來(lái)組合檢索到的信息。實(shí)際上，RNN 的記憶在增長(zhǎng)，因此解碼成本也在增長(zhǎng)：

2、記憶湯（Memory Soup）：另一種結(jié)合過(guò)去記憶的方式，是直接結(jié)合記憶的權(quán)重，而不是針對(duì)特定查詢的輸出。在這種情況下，我們需要對(duì)過(guò)去記憶的權(quán)重執(zhí)行類似注意力的池化操作，然后對(duì)池化后的記憶執(zhí)行一次檢索。同樣，這種變體相對(duì)于上下文長(zhǎng)度具有不斷增長(zhǎng)的有效記憶，因此解碼成本也在增長(zhǎng)。

3、稀疏選擇性緩存（Sparse Selective Caching，SSC）：到目前為止，似乎沒(méi)有免費(fèi)的午餐，我們需要在不斷增長(zhǎng)的有效記憶和每 token 恒定的解碼成本之間做出選擇。

于是作者提出了 SSC，這是一種類似于 MoBA 的專家混合模型，在序列維度上稀疏地選擇過(guò)去緩存記憶的一個(gè)子集，從而引出一個(gè)模型，其有效記憶在增長(zhǎng)，但其每 token 解碼成本保持相對(duì)恒定：

那么效果如何呢？

該方法可以作為一種通用框架，插入到各種現(xiàn)有的循環(huán)架構(gòu)中，如線性注意力模型，或作者之前提出的深度內(nèi)存模型 Titans 等。

實(shí)驗(yàn)結(jié)果表明了其強(qiáng)大的有效性：

模型在語(yǔ)言建模和常識(shí)推理任務(wù)上的表現(xiàn)。

作者在 13 億參數(shù)的模型上進(jìn)行了語(yǔ)言建模、召回密集型、長(zhǎng)上下文以及 needle-in-a-haystack 等實(shí)驗(yàn)，結(jié)果顯示 MC 相較于基礎(chǔ)模型提供了改進(jìn)。其中包括：

• 長(zhǎng)上下文能力提升：在語(yǔ)言建模和長(zhǎng)上下文理解任務(wù)中，加入了 MC 機(jī)制的循環(huán)模型性能得到了全面提升。
• 縮小與 Transformer 的差距：在最具挑戰(zhàn)性的「上下文內(nèi)召回（in-context recall）」任務(wù)中，加入了 MC 的模型擊敗了目前最先進(jìn)（SOTA）的其他循環(huán)模型。
• 依然存在上限：雖然 MC 極大地彌補(bǔ)了 RNN 的召回短板，并極大縮小了與 Transformer 之間的性能差距，但論文也指出，在純粹的密集召回任務(wù)上限上，Transformer 依然保持著最佳的準(zhǔn)確率。

總的來(lái)說(shuō)，該研究利用一種極其優(yōu)雅的算法直覺(jué)（緩存歷史狀態(tài)快照），解決了一個(gè)長(zhǎng)期存在的理論難題，讓非 Transformer 架構(gòu)在實(shí)用性上又向前邁進(jìn)了一大步。

盡管在極限的密集召回上，它依然尚未徹底超越 Transformer，但新路已經(jīng)鋪就，隨著 RNN、SSM 等架構(gòu)的持續(xù)進(jìn)化，Transformer 一家獨(dú)大的現(xiàn)狀或許要有所改變了？