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

來源：機(jī)器之心

編輯：澤南

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

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

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

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

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

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

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

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

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

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

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

• 論文：《Memory Caching: RNNs with Growing Memory》

• 論文鏈接：https://arxiv.org/abs/2602.24281

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

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

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

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

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

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

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

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

那么效果如何呢？

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

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

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

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

• 長上下文能力提升：在語言建模和長上下文理解任務(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)確率。

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

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

閱讀最新前沿科技趨勢(shì)報(bào)告，請(qǐng)?jiān)L問21世紀(jì)關(guān)鍵技術(shù)研究院的“未來知識(shí)庫”

未來知識(shí)庫是 “21世紀(jì)關(guān)鍵技術(shù)研究院”建 立的在線知識(shí)庫平臺(tái)，收藏的資料范圍包括人工智能、腦科學(xué)、互聯(lián)網(wǎng)、超級(jí)智能，數(shù)智大腦、能源、軍事、經(jīng)濟(jì)、人類風(fēng)險(xiǎn)等等領(lǐng)域的前沿進(jìn)展與未來趨勢(shì)。目前擁有超過8000篇重要資料。每周更新不少于100篇世界范圍最新研究資料。 歡迎掃描二維碼或訪問https://wx.zsxq.com/group/454854145828進(jìn)入。

截止到2月28日 ”未來知識(shí)庫”精選的百部前沿科技趨勢(shì)報(bào)告

（加入未來知識(shí)庫，全部資料免費(fèi)閱讀和下載）