Attention 即 RNN：重寫長文本推理的成本邊界

一篇名為《Attention as an RNN》的新研究，提出將 Attention 機制重新表述為一種遞歸神經網路（RNN）。這個看似純粹的數學轉換，實則可能根本性地改變我們對大型語言模型處理長上下文的成本假設。它不僅是技術上的優化，更是在挑戰 Transformer 架構的邊界，為實現真正高效、低記憶體佔用的無限長度推理，開啟了一扇新的大門。我認為，這項工作的重要性不在於又一個新模型的誕生，而在於它迫使我們重新思考 Attention 的本質與其在系統中的角色。

為什麼 Transformer 的長上下文推理如此昂貴？

要理解這項研究的突破性，我們必須先回到 Transformer 模型的核心瓶頸：KV Cache。在標準的 Transformer 架構中，當模型生成每一個新的 token 時，它需要「關注」（attend to）先前序列中的所有 token。為了避免重複計算，系統會將先前每個 token 的「鍵（Key）」和「值（Value）」向量儲存在 GPU 記憶體中，這就是所謂的 KV Cache。

這種機制的直接後果是，推理時的記憶體佔用量與上下文長度 L 呈線性關係，也就是 O(L) 的複雜度。當上下文從幾千個 token 擴展到數百萬個 token 時，KV Cache 會輕易佔用數百 GB 甚至 TB 等級的 VRAM，這使得長上下文推理變得極其昂貴，且對硬體的要求極高。即便透過量化或分頁（例如 PagedAttention）等技術進行優化，記憶體與序列長度之間的線性增長關係，依然是個難以撼動的物理限制。

如何將 Attention 理解為一種 RNN？

這篇論文的洞見在於，它證明了 Attention 的計算過程，可以被數學上等價地重寫為一個線性 RNN 的狀態更新過程。傳統觀點認為，Attention 是一種全局操作，一次性計算所有 token 之間的關聯性。但研究者發現，在自回歸推理的情境下，Attention 的輸出可以被表達為前一個時間步的隱藏狀態（hidden state）與當前時間步輸入的函數。

這意味著，在推理時，我們不再需要儲存完整的 KV Cache。模型只需要維持一個固定大小的隱藏狀態，並在每一步生成新 token 時對其進行更新。如此一來，記憶體複雜度便從 O(L) 驟降至 O(1)，與序列長度完全脫鉤。

這個轉換並非沒有代價。傳統 RNN 的一個主要缺點是其循序（sequential）特性，導致訓練時無法像 Transformer 那樣進行大規模平行化，因而效率低下。為了解決這個問題，研究團隊引入了「平行掃描（parallel scan）」演算法，這是一種在平行計算領域廣為人知的技術。它允許模型在訓練時，以 O(L log L) 的時間複雜度一次性計算出所有時間步的 RNN 狀態，從而保留了類似 Transformer 的高效平行訓練能力。

Aaren：如何實現高效推理與平行訓練的雙贏？

基於上述理論，團隊實作出一個名為 Aaren（Attention as a Recurrent Neural Network）的模組。Aaren 的設計目標是成為一個可以無縫替換標準 Multi-Head Attention 的模組，它巧妙地融合了兩種架構的優點：

• 訓練階段： 表現得像一個 Transformer。利用平行掃描演算法，它可以高效地在 GPU 上進行平行化訓練，充分利用現代硬體的算力。
• 推理階段： 表現得像一個 RNN。它以循序方式更新一個固定大小的狀態，實現了常數級別的記憶體佔用與每個 token O(1) 的計算時間。

這種雙重特性，使其在概念上與近年來備受關注的狀態空間模型（State Space Models, SSMs）如 Mamba 有異曲同工之妙。兩者都試圖在維持 Transformer 等級性能的同時，引入 RNN 的高效推理模式。但 Aaren 的獨特之處在於，它並非提出一個全新的架構，而是從根本上重新詮釋了我們早已熟悉的 Attention 機制。

為何說這不只是一個新模型，而是對成本假設的重估？

我認為，將 Attention 重塑為 RNN，其意義遠不止於性能優化。它挑戰了我們對 LLM 架構的一些基本假設，並可能改變未來的發展路徑。

首先，它徹底改變了長上下文推理的成本結構。如果推理記憶體不再是瓶頸，那麼運行一個能處理百萬、甚至千萬級 token 上下文的模型，將不再是少數頂級實驗室的專利。這將極大地推動需要長期記憶的應用，例如個人化 AI 助理、複雜的程式碼分析、或對長篇法律或科學文獻的深度理解。

其次，它模糊了 Transformer 與 RNN/SSM 之間的界線。過去，我們常將兩者視為不同的技術路線。但 Aaren 的存在證明，兩者在數學上可能比我們想像的更為接近。這或許意味著，未來的模型架構將不再是非此即彼的選擇，而是根據任務需求，在一個更廣闊的設計空間中，靈活地組合不同的計算基元（computational primitives）。

最後，這項工作為我們提供了一個強有力的啟示：許多看似根深蒂固的架構瓶頸，有時可以透過回到底層數學，用一個全新的視角來重新審視和解決。這不是用更強大的硬體去暴力破解問題，而是透過更深刻的理解來繞過問題。在 AI 系統的設計與建構中，這種來自第一性原理的思考，往往是通往真正創新的路徑。

延伸閱讀

• Attention as an RNN (arXiv)
• Attention Is All You Need (arXiv)
• Mamba: A New State Space Model Architecture (NVIDIA Developer Blog)

我是江中喬，一位具有 TPM 與產品管理背景的 AI 系統建構者，目前專注於 AI 認知增強系統與多 Agent 協作架構的設計與實踐。