當(dāng)線性注意力學(xué)會(huì)「寫入前思考」：并行化的多步記憶寫入

該工作已被機(jī)器學(xué)習(xí)領(lǐng)域頂級(jí)會(huì)議 ICML 2026 錄用，論文題目 “PRISM: Parallel Residual Iterative Sequence Model”。

一、背景：從無限背包到有限背包

（一）Transformer 的無限背包與線性注意力的有限背包

背包容量有限，每來一個(gè)新 token，模型必須決定往里寫什么、同時(shí)擦掉什么。這個(gè) "寫與擦" 的規(guī)則，決定了有限背包模型的天花板。但在深入討論 "寫與擦" 之前，我們先要回答一個(gè)更基本的問題。

（二）有限背包本質(zhì)上是 RNN，為何還能并行？

確實(shí)如此，有限背包模型的數(shù)學(xué)形式本質(zhì)上就是 RNN：

關(guān)鍵在于一個(gè)數(shù)學(xué)技巧：Parallel Scan（并行前綴掃描）。

（三）為什么并行這么重要？GPU 的 "搬運(yùn)工" 瓶頸

一個(gè)常見的誤解是將 "串行慢" 歸因于更多的浮點(diǎn)運(yùn)算。實(shí)際上，瓶頸在別處。現(xiàn)代 GPU 的計(jì)算核心（Tensor Core / CUDA Core）算力極為充沛，A100 GPU 每秒能做 312 萬億次浮點(diǎn)運(yùn)算（312 TFLOPS）。真正的瓶頸不是 "算"，而是 "搬"。

GPU 的存儲(chǔ)分為兩層：

• HBM（High Bandwidth Memory，高帶寬顯存）：容量大（40-80 GB），但讀寫速度 "慢"（約 2 TB/s）。模型參數(shù)、state 矩陣 S、中間 activation 都存在這里。
• SRAM（片上緩存）：容量小（每個(gè) SM 約 192 KB），但讀寫速度極快（約 19 TB/s，快 10 倍）。GPU 的計(jì)算核心只能直接訪問 SRAM。

打個(gè)比方：SRAM 像工作臺(tái)（小但觸手可及），HBM 像倉(cāng)庫（大但每次取貨要走一趟）。

所以每一次計(jì)算都要經(jīng)歷一個(gè) "搬運(yùn)" 流程：把數(shù)據(jù)從 HBM 搬進(jìn) SRAM，在 SRAM 里算完，再把結(jié)果搬回 HBM。這個(gè)搬運(yùn)的時(shí)間往往遠(yuǎn)超計(jì)算本身，這就是所謂的 memory-bound（存儲(chǔ)帶寬瓶頸）。

能否適配parallel scan 不僅是算法設(shè)計(jì)上的美學(xué)選擇，更直接決定了 10-100 倍的實(shí)際運(yùn)行速度差異。

（四）Rank-1 寫入的瓶頸

以 GDN （Gated DeltaNet）為代表的線性注意力模型，每個(gè) token 對(duì) S 做的是一次 rank-1 更新：

如果一個(gè) token 攜帶的語義是多維度的（它同時(shí)是某個(gè)句法結(jié)構(gòu)的成分、某個(gè)語義角色的載體、某個(gè) topic 的關(guān)鍵詞），rank-1 的一行寫入無法同時(shí)在這些維度上做精細(xì)調(diào)整。信息在壓縮寫入時(shí)不可避免地丟失。

核心矛盾：背包有限，每次卻只允許寫一行。這是當(dāng)前所有線性復(fù)雜度模型的共有瓶頸。

（五）TTT 的突破與代價(jià)

既然 rank-1 寫入太淺，一個(gè)自然的想法是：讓模型學(xué)會(huì)更深的寫入規(guī)則。

TTT（Test-Time Training）系列工作采取了一種根本性不同的策略：把記憶狀態(tài)從一個(gè) linear 矩陣 S 升級(jí)為一個(gè) MLP 的權(quán)重矩陣。每來一個(gè) token，對(duì) MLP 的權(quán)重做多步梯度下降（multi-step GD），逐步精煉寫入內(nèi)容。這帶來了顯著的質(zhì)量提升。

二、分析：TTT-MLP 為什么效果好，但速度慢？

在設(shè)計(jì) PRISM 之前，我們首先深入分析 TTT-MLP 的梯度結(jié)構(gòu)，弄清楚它的高表達(dá)力到底從何而來。

（一）步長(zhǎng) × 殘差 × 方向 模式的涌現(xiàn)

每步更新具有一個(gè)結(jié)構(gòu)模式：

TTT-MLP 的高表達(dá)力正來自這個(gè) 步長(zhǎng) × 殘差 × 方向 模式：多步殘差遞減提供了優(yōu)化深度（depth），W? 多行提供多個(gè)方向則提供了表達(dá)寬度（width /rank-L）（即同時(shí)修改 S 矩陣的 L 個(gè)獨(dú)立維度）。

（二）高表達(dá)力與串行是同一根因的兩面

具體來說，它造成了兩個(gè)維度的串行瓶頸：

1. Token 間串行（Inter-token Seriality）

2. Step 間串行（Intra-step Seriality）

瓶頸 C（方向與殘差的同步）：在多步 GD 中，第 l+1 步的寫入方向必須等待第 l 步的權(quán)重更新完畢才能確定，殘差也必須等上一步算完才能得到，強(qiáng)制引入一個(gè)無法展開的循環(huán)。

瓶頸 C 是最核心的矛盾：它同時(shí)是 rank-L 表達(dá)力的載體和步間串行的根源。因此消除瓶頸 C 不能簡(jiǎn)單取消迭代，必須在取消同步耦合的同時(shí)保留多方向和殘差遞減帶來的表達(dá)力。

三、方法：PRISM 的設(shè)計(jì)與實(shí)現(xiàn)

基于上述分析，PRISM 的策略非常明確：在兼容 parallel scan 的線性狀態(tài) S 上顯式重建 TTT-MLP 的 步長(zhǎng) × 殘差 × 方向 模式，然后分維度消除串行。

（一）核心迭代形式：步長(zhǎng) × 殘差 × 方向

PRISM 顯式構(gòu)造了 TTT-MLP 的多步迭代模式：

與 TTT-MLP 的對(duì)應(yīng)關(guān)系：

（二）消除 Token 間串行：A/B 分離 + 局部 Anchor 代理

至此，序列級(jí)別的 parallel scan 已完全恢復(fù)。anchor 讓不同 token 的迭代可以同時(shí)啟動(dòng)，但每個(gè) token 內(nèi)部的 L 步之間仍需順序執(zhí)行（瓶頸 C）。

（三）消除 Step 間串行：解耦鏈 + 閉合式預(yù)計(jì)算

解決瓶頸 C。因?yàn)橛辛?anchor，兩條鏈自然解耦：

由此多步迭代推算得到閉合式：

L 步的串行循環(huán)被消解為單步閉合式計(jì)算。整個(gè)多步梯度下降計(jì)算過程可以編譯成一個(gè) fused kernel，數(shù)據(jù)只需要從 HBM 搬進(jìn) SRAM 一次。

（四）架構(gòu)全貌與 GDN 退化

多步梯度下降計(jì)算過程的原始產(chǎn)出是 L 個(gè) rank-1 迭代計(jì)算：

PRISM 可以視為一種多步殘差擬合計(jì)算過程，L=1 時(shí)精確退化為 GDN。 后續(xù)步只是在第一步的基礎(chǔ)上追加非線性修正，且可以使用 low rank 網(wǎng)絡(luò)增量，額外參數(shù)量不超過基礎(chǔ)模型的 10%。

四、實(shí)驗(yàn)結(jié)果

（一）序列推薦

在公開序列推薦基準(zhǔn) Amazon 上，PRISM 表現(xiàn)與 Transformer baseline 效果接近，超過大多數(shù)線性注意力類方法。計(jì)算效率方面，PRISM 與 GDN 同級(jí)，比 TTT-MLP 快 174 倍。

（二）語言建模（基于 SlimPajama 2B 訓(xùn)練，130M 參數(shù)）

在更大規(guī)模的語言建模實(shí)驗(yàn)上（SlimPajama 2B tokens, Mistral tokenizer），PRISM 同樣取得了全面領(lǐng)先：

PRISM 在 WikiText PPL、LAMBADA PPL 和 9 項(xiàng) Zero-Shot 下游任務(wù)平均準(zhǔn)確率上均為最優(yōu)，領(lǐng)先 GDN 3.2 個(gè)百分點(diǎn)。

（三）組件消融

訓(xùn)練 PPL 差異極小，但下游泛化差異巨大。單步 solver (L=1) 的訓(xùn)練 PPL 幾乎等于完整版，但 Avg ACC 下跌 2.9 個(gè)百分點(diǎn) ——rank-L 的真正價(jià)值不在 next-token prediction 上，而在需要精確長(zhǎng)程檢索的下游任務(wù)上。

五、延伸思考

（一）有限背包終究有限，混合架構(gòu)也許是必然

從 PRISM 的視角看，這個(gè)直覺有一個(gè)很好的技術(shù)解釋。PRISM 用短卷積（ShortConv）計(jì)算的局部 anchor 替代全局狀態(tài) S 來近似殘差。由于短卷積窗口通常只覆蓋最近 3-4 個(gè) token，對(duì)于需要跨越數(shù)千步的長(zhǎng)程依賴，近似質(zhì)量必然下降。

如果在 PRISM 層之間穿插少量 Transformer 層，后者就充當(dāng)了一種全局的、非線性的歷史狀態(tài)精確計(jì)算器，能補(bǔ)償 anchor 在長(zhǎng)程上的近似誤差。從這個(gè)角度看，Transformer 本身就是 ShortConv anchor 的 "全局升級(jí)版"：ShortConv 用固定窗口的局部卷積近似歷史狀態(tài)，Transformer 用全局 attention 精確算歷史狀態(tài)。

（二）線性注意力的 LoRA？

PRISM 的最終形式有一個(gè)有趣的結(jié)構(gòu)特征：

這個(gè) "基礎(chǔ)迭代過程 + low rank 旁路" 的形式，跟 LoRA（Low-Rank Adaptation） 非常相似，這啟發(fā)了一個(gè)微調(diào)場(chǎng)景下的有趣思路。

LoRA 的核心思想是：凍結(jié)預(yù)訓(xùn)練好的大模型權(quán)重，只在關(guān)鍵層旁邊加一條 low-rank 旁路來做微調(diào)。受 PRISM 形式的啟發(fā)，我們可以設(shè)想一種面向 Linear Attention / SSM 模型的參數(shù)高效微調(diào)方法：對(duì)已訓(xùn)練好的模型，凍結(jié)基礎(chǔ)迭代過程，只在寫入支路上增加一條 PRISM 風(fēng)格的殘差擬合旁路，此外，這條旁路有閉合式（不增加訓(xùn)練時(shí)間），而且第一步退化為原模型的標(biāo)準(zhǔn)寫入（不破壞預(yù)訓(xùn)練知識(shí)）。這意味著它滿足 LoRA 的兩個(gè)關(guān)鍵要求：參數(shù)高效和不損害原模型能力。

結(jié)語

PRISM 驗(yàn)證了 "寫入前思考" 范式在線性注意力模型中的可行性：通過分析 TTT-MLP 的梯度結(jié)構(gòu)揭示 步長(zhǎng) × 殘差 × 方向 迭代模式，在線性狀態(tài)上顯式重建該模式并通過 anchor 代理和閉合式預(yù)計(jì)算實(shí)現(xiàn)完全并行。最終架構(gòu)極簡(jiǎn) ——GDN + 非線性旁路，訓(xùn)練速度與 GDN 同級(jí)，參數(shù)增量不到 10%。在推薦和語言建模兩個(gè)場(chǎng)景上的驗(yàn)證表明，這是一項(xiàng)通用的線性注意力增強(qiáng)技術(shù)。未來我們將進(jìn)一步探索 PRISM 在更大參數(shù)規(guī)模上的 scaling 行為和推薦系統(tǒng)上的應(yīng)用效果，以及其作為線性注意力模型參數(shù)高效微調(diào)方法的實(shí)際效果。

參考文獻(xiàn)：

[1] Sun et al. “Learning to (Learn at Test Time): RNNs with Expressive Hidden States.” NeurIPS 2024.

[2] Yang et al. “Gated Delta Networks with Pairwise Tokenized Graphs.” NeurIPS 2024.

[3] Katharopoulos et al. “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention.” ICML 2020.