56|状态空间模型 讲的是用可更新状态取代 full attention 的一条路径:Mamba/S4。这一篇讲另外两条并行的路径——直接改写 softmax attention 的数学形式(线性注意力),以及训练走 Transformer 的并行路、推理走 RNN 的循环路(RWKV、RetNet)。三条路径目标一致:把 \(O(n^2)\) 降到 \(O(n)\),同时尽量不丢 Transformer 的训练效率。
常见的误区是把”\(O(n)\)“当成免费午餐:只要复杂度公式好看,就默认它在质量和吞吐上都不吃亏。这个误区有具体的反例。Zoology(Arora et al., ICLR 2024)证明了标准 softmax attention 能用与序列长度无关的维度解出关联检索任务,而线性注意力、RWKV 等循环状态模型做不到,除非状态维度随序列长度增长;Gated Linear Attention(Yang et al., ICML 2024)的论文摘要直接写道,”现有的线性注意力实现缺乏 I/O 感知,因此比高度优化的 softmax attention 实现更慢”——理论复杂度和真实吞吐是两件事。
本篇要交代清楚三件事:
- 线性注意力具体怎么把 softmax 拆成可累积的循环状态,这个拆解在理论上到底丢了什么;
- RWKV、RetNet 各自怎样把”训练并行、推理线性”工程化,它们的失败模式分别是什么,有没有实测数据支撑;
- 为什么 \(O(n)\) 在论文里成立,换到真实 GPU 上却经常打不过精心优化的 FlashAttention,除非专门写 I/O-aware 的 kernel。
一、线性注意力:把 softmax 拆成可累积的形式
标准 causal self-attention 对位置 \(i\) 的输出是
\[ o_i = \sum_{j=1}^{i} \frac{\exp(q_i^\top k_j/\sqrt{d})}{\sum_{l=1}^{i}\exp(q_i^\top k_l/\sqrt{d})}\, v_j \]
分母里的归一化常数必须扫过全部 \(j\le i\),这是 \(O(n^2)\) 的根源:每个 query 都要和所有历史 key 单独算一次相似度。
1.1 用核函数替代 softmax
Katharopoulos et al.(ICML 2020)的做法是把相似度函数从 \(\exp(q^\top k/\sqrt d)\) 换成一个可分解的核:\(\text{sim}(q,k)=\phi(q)^\top\phi(k)\),其中 \(\phi\) 是某个特征映射(论文用 \(\text{elu}(x)+1\))。代入之后,利用矩阵乘法的结合律,输出可以写成
\[ o_i = \frac{\phi(q_i)^\top \sum_{j=1}^i \phi(k_j)v_j^\top}{\phi(q_i)^\top \sum_{j=1}^i \phi(k_j)} \]
定义累积量 \(S_i=\sum_{j\le i}\phi(k_j)v_j^\top\in\mathbb R^{d_k\times d_v}\)、\(z_i=\sum_{j\le i}\phi(k_j)\in\mathbb R^{d_k}\),则
\[ S_i = S_{i-1}+\phi(k_i)v_i^\top,\qquad z_i=z_{i-1}+\phi(k_i),\qquad o_i=\frac{\phi(q_i)^\top S_i}{\phi(q_i)^\top z_i} \]
这正是论文标题”Transformers are RNNs”的字面含义:\(S_i,z_i\) 就是循环状态,形状是固定的 \(d_k\times d_v\) 矩阵和 \(d_k\) 向量,不随位置 \(i\) 增长。这是它和标准 attention 的 KV Cache 的本质区别——KV Cache 的大小是 \(O(i\cdot d)\),随生成长度线性增长;线性注意力的状态大小是常数。论文报告在自回归图像生成和语音识别任务上,这种重写让推理速度提升最多三个数量级。
Performer(Choromanski et al., ICLR 2021)走的是另一条子路径:不是任选一个核函数替代 softmax,而是用随机特征(FAVOR+)去无偏估计 softmax 本身对应的核,理论上更接近原始 attention 的行为,但同样依赖特征维度的选择来控制近似误差。这两条子路径(换核 vs. 近似原核)在后续工作里常被一起归为”线性注意力”,但它们对表达力的取舍并不完全一样。
1.2 表达力差距:不是”看场景”,是有下界的
线性注意力把 attention 矩阵换成了一个固定维度的状态,这个状态能不能撑住”精确检索”是可以被形式化检验的。Zoology(Arora et al., ICLR 2024)提出了 multi-query associative recall(MQAR)任务:在一段文本里多次、在不同位置、从大词表中检索之前出现过的键值对(例如同时记住”Hakuna Matata 意味着 no worries”和别的几组配对)。论文的核心结论:
- 标准 softmax attention 能用与序列长度 \(N\) 无关的模型维度解出 MQAR;
- 门控卷积、朴素线性注意力这类把历史压进固定状态的架构,除非模型维度 \(d\ge N\),否则解不出来;
- 在真实语言建模数据(the Pile)上,这类架构的困惑度差距里 82% 可以用关联检索能力的差距解释。
Based(Arora et al., arXiv:2402.18668, 2024,预印本,未见正式同行评审)在同一套 the Pile 评测(GPT-2 BPE 分词、10B token 预训练、约 360M 参数规模)上给出了具体数字,把整体测试集切成”关联检索(AR)token”和”其他 token”分别算困惑度:
| 架构 | 参数量 | 预填充吞吐 (tok/ms) | 生成吞吐 (tok/ms) | 整体 Pile ppl | AR 切片 ppl |
|---|---|---|---|---|---|
| Transformer++ | 360M | 207.77 | 23.82 | 8.39 | 1.87 |
| Based(线性注意力+局部滑窗混合) | 363M | 514.57 | 47.23 | 8.65 | 2.07 |
| Mamba | 358M | 267.09 | 39.95 | 8.64 | 2.21 |
| 纯线性注意力(无滑窗、无卷积) | 362M | — | — | 9.49 | 2.29 |
| GLA | 362M | — | — | 9.12 | 2.36 |
| RWKV-v5 | 362M | — | — | 9.79 | 2.40 |
数据来源:Arora et al. (2024) Table 1、Table 4,环境为 10 亿参数以下模型在 the Pile 上训练 10B token 的受控对比,口径一致(同一分词器、同一训练 token 数)。
在这个口径下,纯线性注意力和 RWKV-v5 在关联检索切片上比 Transformer++ 差 0.4–0.5 ppl,而 Based 靠”在小窗口里插入精确 softmax attention + 全局线性注意力”把差距压到 0.2 ppl。这说明线性注意力的短板不是”训练不够”,而是用一个不随序列长度增长的状态去存所有历史键值对,在需要多次、精确检索远处不同位置信息时天生吃亏——除非引入局部精确 attention 或数据依赖的门控来补。第二节、第三节要看的是 RWKV 和 RetNet 分别怎么应对这个取舍。
二、RWKV:训练像 Transformer、推理像 RNN 的取舍
RWKV(Receptance Weighted Key Value,Peng et al., Findings of EMNLP 2023)的目标和线性注意力是同一件事的另一种工程化:让模型训练时可以像 Transformer 一样并行,推理时可以像 RNN 一样维护固定大小状态。它的核心算子 WKV 用逐通道(channel-wise)指数衰减替代了线性注意力里的核函数:
\[ wkv_t=\frac{\displaystyle\sum_{i=1}^{t-1}e^{-(t-1-i)w+k_i}v_i+e^{u+k_t}v_t}{\displaystyle\sum_{i=1}^{t-1}e^{-(t-1-i)w+k_i}+e^{u+k_t}} \]
其中 \(w\) 是可学习的逐通道衰减向量(越久远的 token 权重按 \(e^{-w}\) 逐步衰减),\(u\) 是单独给当前 token 的”加成”向量,避免衰减机制过度压低最新信息。输出经过一个 sigmoid 门控:
\[ o_t = W_o\cdot(\sigma(r_t)\odot wkv_t) \]
\(r\) 是 receptance 向量,\(\sigma(r_t)\) 起到类似 GRU/LSTM 输出门的作用。对照第一节的记号,\(wkv_t\) 本质上和线性注意力的 \(o_i=\phi(q_i)^\top S_i/\phi(q_i)^\top z_i\) 是同一个”累积再归一化”结构,区别在于:线性注意力靠特征映射 \(\phi\) 决定 query/key 的相似度权重,RWKV 直接用一个可学习的逐通道衰减 \(w\) 替代了相似度计算本身,再叠加一个门控。这让状态更新在训练时可以按元素并行做 scan(复杂度 \(O(BTd)\)),推理时是严格的逐步 RNN(复杂度 \(O(Td)\),常数内存 \(O(d)\)),而 Transformer 推理的复杂度是 \(O(T^2d)\) 时间、\(O(Td)\) 内存(KV Cache 随长度线性增长)——这组复杂度对比来自论文 Table 1。
2.1 论文自己写的失败模式
RWKV 论文第 9 节”Limitations”直接承认了两个问题,值得原文引用:
线性注意力……可能会限制模型在需要从很长上下文中回忆细枝末节信息的任务上的表现,这是因为信息被压缩进了一个单一向量表示,而不像标准 Transformer 的二次 attention 保留完整信息。换句话说,模型的循环结构从根本上限制了它”回看”之前 token 的能力……另一个局限是 prompt engineering 的重要性显著提高。
论文附录 L 用一个具体实验支撑这个说法:在 RTE(文本蕴含)任务上,把指令和输入的顺序对调(让 RNN 不必”回头”去看已经被状态压缩掉的指令),RWKV-4-Raven-14B 的 F1 分数可以显著提升——这说明 RWKV 对信息在 prompt 里出现的顺序比标准 Transformer 敏感得多,因为循环状态一旦把某段信息”滚过去”,后面就很难再精确取回。
2.2 长上下文基准上的实测差距
RULER(Hsieh et al., COLM 2024,同 arXiv:2404.06654)用比 needle-in-a-haystack 更全面的多任务长上下文基准(含多跳追踪、聚合等任务类型)测试了 17 个长上下文模型,其中专门对比了 RWKV-v5 和 Mamba-2.8B 与 Llama-2-7B:两者在上下文扩展到 8K 时都出现明显性能下降,在 4K 长度以内已经大幅落后 Llama-2-7B 基线。这和论文自己第 9 节的”funneling”说法方向一致:固定大小状态在长上下文精确检索上的短板,不是小模型量级的实验假象,在更大规模、更系统的基准上依然存在。
三、RetNet:retention 与三种计算模式
RetNet(Sun et al., arXiv:2307.08621, 2023,未见正式同行评审版本)提出的 retention 机制,目标同样是”训练并行、推理 \(O(1)\)“,但走的是第三条子路径:用一个显式的、按 head 区分速率的指数衰减,替代 softmax,同时保留矩阵形式的高维状态(而不是 RWKV 的向量状态)。论文用三种数学上等价的表示描述同一个算子。
3.1 parallel:训练时当 attention 用
\[ Q=(XW_Q)\odot\Theta,\qquad K=(XW_K)\odot\overline\Theta,\qquad V=XW_V \]
\[ \Theta_n=e^{in\theta},\qquad D_{nm}=\begin{cases}\gamma^{n-m}, & n\ge m\\[2pt]0, & n<m\end{cases} \]
\[ \text{Retention}(X)=(QK^\top\odot D)V \]
\(\Theta\) 是类似 xPos/RoPE 的旋转位置编码,作用在 \(Q,K\) 上;\(D\) 是把因果 mask 和指数衰减合并成的一个矩阵,\(\gamma\in(0,1)\) 按 head 取不同值(multi-scale decay,不同 head 关注不同的”记忆半衰期”)。这一步和标准 attention 的计算图几乎一样,可以用矩阵乘法在 GPU 上并行训练。
3.2 recurrent:推理时当 RNN 用
\[ S_n=\gamma S_{n-1}+K_n^\top V_n,\qquad \text{Retention}(X_n)=Q_nS_n \]
\(S_n\) 是每个 head 一个的 \(d_k\times d_v\) 矩阵状态——这一点和线性注意力的 \(S_i\) 形状一致,区别只是衰减方式从”隐式核函数”变成了”显式标量 \(\gamma\)“。推理时每步只需要 \(O(1)\) 的状态更新和读出。
3.3 chunkwise recurrent:训练长序列时的折中
把序列切成长度 \(B\) 的 chunk,chunk 内部按 parallel 表示计算(矩阵乘法,吃满 GPU 算力),chunk 之间按 recurrent 表示传递状态(只传一个 \(d_k\times d_v\) 矩阵,不重新展开整段历史)。这既避免了长序列下 parallel 表示 \(O(n^2)\) 的显存开销,又避免了纯 recurrent 表示逐 token 串行的低效率。下图是三种表示的关系:
flowchart LR
A["Parallel representation\n(QK^T ⊙ D)V, full matmul\nused for training, short context"] <--> B["Recurrent representation\nS_n = γS_(n-1) + K_n^T V_n\nO(1) state, used for inference"]
B <--> C["Chunkwise recurrent\nparallel inside each chunk +\nrecurrent state across chunks\nused for long-sequence training"]
A <--> C
图:retention 的三种表示是同一个线性递归的三种展开方式,不是三个不同模型。
这个”chunk 内并行、chunk 间递归”的思路后面会在第四节再次出现——FlashLinearAttention 用几乎一样的技巧让线性注意力变得硬件友好,说明这不是 RetNet 独有的技巧,而是这一整类模型共同的工程解法。
3.4 衰减和门控是不是真的有用:两组消融
RetNet 论文自己的消融(Table 6,同一评测口径)显示,去掉 \(\gamma\) 衰减(等价于 \(\gamma=1\),退化成无衰减的线性递归)或去掉 multi-scale decay(所有 head 用同一个 \(\gamma=127/128\)),in-domain 困惑度都会从 26.05 恶化到 27.86 和 27.02;去掉 swish 门控或 GroupNorm 同样会退化。这说明衰减机制和门控都在实际起作用,不是可以随意砍掉的装饰。
但衰减本身也有上限。Gated Linear Attention(Yang et al., ICML 2024)在自己的受控消融(340M 参数模型,训练 7B token)里对比了三种衰减策略的训练困惑度:无门控的朴素线性注意力 23.21,数据无关的标量衰减(论文明确写”即 RetNet 的方式”)16.55,数据依赖的标量门控 15.56,完整的 GLA 门控 14.77。也就是说,RetNet 式的固定衰减比完全没有衰减好得多,但比”衰减率随输入内容变化”的门控差一截——衰减本身是有用的归纳偏置,但把衰减做成数据依赖,比单纯引入固定衰减更接近解决表达力短板。同一篇论文在三个真实关联检索任务(FDA、SWDE、SQUAD)上的对比(Table 3,100B token 训练)也支持这个结论:RetNet 在 FDA 上是 14.3,GLA 是 19.9;SWDE 上 RetNet 42.8,GLA 50.6;标准 softmax attention 在这三项上接近满分(论文因此略去了具体数字)。
有意思的是,RetNet 论文的 Related Work 部分自己对比过 RWKV:“RWKV 用指数衰减替代了 AFT 的位置编码,训练和推理都跑循环形式;相比之下,retention 保留了高维状态来编码序列信息,这有助于表达能力和最终性能”——即 RetNet 认为自己相对 RWKV 的优势正是”矩阵状态 \(S_n\in\mathbb R^{d_k\times d_v}\) 比 RWKV 的向量状态信息容量更大”。RetNet 论文用同一套训练配置(200M 参数、16 层、隐藏维度 1024、10k 步、0.5M token batch)直接比较了 Linear Transformer、RWKV、H3、Hyena 和 RetNet(Table 5):
| 方法 | In-Domain ppl | PG22 | QMSum | GovReport | SummScreen |
|---|---|---|---|---|---|
| Linear Transformer | 40.24 | 63.86 | 28.45 | 25.33 | 32.02 |
| RWKV | 30.92 | 51.41 | 28.17 | 19.80 | 25.78 |
| H3 | 29.97 | 49.17 | 24.29 | 19.19 | 25.11 |
| Hyena | 32.08 | 52.75 | 28.18 | 20.55 | 26.51 |
| RetNet | 26.05 | 45.27 | 21.33 | 16.52 | 22.48 |
在这个统一口径下 RetNet 全面领先,但这是论文自己报告的结果,没有独立第三方复现——读者应该知道这一点,再引用具体数字。训练成本方面,论文另一组自测(Table 4)显示 1.3B 规模下 RetNet 的训练吞吐(73344.8 词/秒)略高于 Transformer+FlashAttention(63965.2 词/秒),显存占用(34.5GB)也更低;这组数字同样是作者自报告,没有第三方复现。
四、理论 \(O(n)\) 不等于真实吞吐:kernel 和硬件生态的距离
第一到第三节的复杂度分析都是”理想情况”:一次矩阵乘法算作一次操作,扫描算作线性时间。真实 GPU 上决定速度的是显存带宽、片上 SRAM 大小、tensor core 的利用率,这些跟”大 O 记号”不是一回事。
Gated Linear Attention 论文(Yang et al., ICML 2024)在摘要里直接点出这个问题:“现有的线性注意力实现缺乏 I/O 感知,因此比高度优化的 softmax attention 实现更慢”。线性注意力的递归形式如果直接按逐 token 展开写成 CUDA kernel,会产生大量小矩阵操作和频繁的显存读写,吃不满 GPU 的算力峰值;而 FlashAttention 系列恰恰是把 softmax attention 的 \(O(n^2)\) 计算拆成 SRAM 里能装下的 tile,用分块 + 在线 softmax 把显存搬运降到最低——这是经过好几代 CUDA/Triton 迭代才做到的工程成果,不是”\(O(n^2)\) 天生就慢”。
这篇论文提出的 FlashLinearAttention 用的是和 RetNet 的 chunkwise recurrent 几乎一样的思路:chunk 内做并行矩阵乘法、chunk 间传递状态,同时显式设计 tiling 策略减少 HBM 读写次数。结果是:作为独立层,FlashLinearAttention 即便在 1K 这种短序列长度下也比 FlashAttention-2 快——但这个速度优势是专门为线性注意力重新写一遍 I/O-aware kernel 换来的,用同一篇论文的原话说,“当前的线性注意力实现”(即没有做这层工程)反而更慢。换句话说,\(O(n)\) 只给了”有可能更快”的理论空间,真正兑现要靠专门的 kernel 工程,这个工程量和当年 FlashAttention 对 softmax attention 做的事是同一个量级。
同一篇论文还报告了训练吞吐的横向对比:在超过 4096 长度的训练序列上,Mamba 的训练吞吐落后于 Transformer++ 和 GLA。这说明”线性时间模型”内部也存在明显的 kernel 成熟度差异——不是所有号称 \(O(n)\) 的架构都能在当下的软件栈里跑出线性时间对应的实际速度优势,谁的 kernel 被打磨得更细,谁就更快,这和模型本身的数学形式是两个独立的变量。
Mamba 那一侧的进展也指向同一个方向:Dao 和 Gu 在 Transformers are SSMs(ICML 2024,见 56|状态空间模型 参考文献)里证明了 SSM 的选择性扫描和(某种形式的)线性 attention 在结构化矩阵乘法的意义下是同一件事,这也解释了为什么 Mamba 和 GLA 之类的架构能共享相近的硬件优化思路——hardware-aware 设计已经变成决定这几条路线能不能落地的共同瓶颈,而不是各自独立的细节。
五、和 Mamba/SSM 的对照:同一个目标,不同的机制分叉
四条路径(线性注意力、RWKV、RetNet、SSM/Mamba)都在回答同一个问题:怎么用一个不随序列长度增长的状态取代显式的 pairwise 历史访问。分叉点在于用什么函数替代相似度计算,以及衰减/更新是不是数据依赖的:
| 机制 | 用什么替代相似度/衰减 | 状态形状 | 衰减是否数据依赖 | 训练并行方式 | 文献报告的已知短板 |
|---|---|---|---|---|---|
| 线性注意力 | 核特征映射 \(\phi(q)^\top\phi(k)\) | 矩阵 \(d_k\times d_v\)(累积量 \(S_i\)) | 否(\(\phi\) 固定) | cumsum / chunkwise scan | MQAR 需要 \(d\ge N\)(Zoology, ICLR 2024) |
| RWKV | 逐通道指数衰减 \(w\) + receptance 门控 \(\sigma(r)\) | 向量 \(O(d)\) | 部分(门控依赖输入,衰减率本身是固定参数) | time-parallel WKV + 逐通道 scan | 论文自述”funneling”,RULER 上落后 Llama-2-7B(COLM 2024) |
| RetNet | 逐 head 固定标量衰减 \(\gamma\)(multi-scale)+ xPos 旋转 | 矩阵 \(d_k\times d_v\)(每 head) | 否(\(\gamma\) 固定,但按 head 分级) | parallel(含 decay mask)/ chunkwise | 固定衰减弱于数据依赖门控(GLA 消融, ICML 2024) |
| Mamba/SSM | 输入依赖的选择性参数 \(\Delta,A,B,C\) | 由 SSM 状态维度决定(见 56) | 是(selective mechanism) | parallel scan | 信息瓶颈,见 56 节讨论 |
这张表最重要的一格是”衰减是否数据依赖”:Mamba 的 selective mechanism 从一开始就让衰减/更新参数依赖当前 token;RWKV 的衰减 \(w\) 是训练后固定的参数,只有门控 \(\sigma(r)\) 依赖输入;RetNet 的 \(\gamma\) 完全固定,靠多个 head 分摊不同衰减速率;朴素线性注意力则完全没有衰减这个概念,只有一个固定核。这条”数据依赖程度”的梯度,恰好对应第三节 GLA 消融里困惑度从 23.21(无门控)到 16.55(RetNet 式固定衰减)到 14.77(完整数据依赖门控)的排序——不是所有”线性时间模型”都一样,衰减机制越接近输入依赖,在同等参数量下往往越接近 full attention 的质量,这是有具体消融数据支撑的结论,不是”各有优劣看场景”式的空话。
六、争论与开放问题
争论点:线性注意力/RWKV/RetNet 的表达力短板,是理论下界,还是当前工程没做够?
A 派的证据是理论性的:Zoology(Arora et al., ICLR 2024)用 MQAR 证明了朴素的、无输入依赖门控的循环状态模型需要维度 \(d\ge N\) 才能匹配 softmax attention 的关联检索能力,这是一个可以形式化证明的下界,不会随着”训练更久、调得更细”而消失。
B 派的证据是工程性的:GLA 的消融显示,仅仅把固定衰减换成数据依赖的门控,就能把训练困惑度从 23.21 拉到 14.77(第三节数据);Based 通过在线性注意力旁边加一小段精确 softmax attention(滑窗 64–128 token),就把关联检索切片的困惑度从纯线性注意力的 2.29 拉到 2.07,逼近 Transformer++ 的 1.87。这说明”朴素线性注意力”和”加了门控/混合了精确 attention 的现代变体”不该被当成同一件事一并否定。
两派并不矛盾:A 派证明的是”无门控、纯状态压缩”这一类架构的下界,B 派证明的是”引入数据依赖或局部精确检索”能大幅缓解但未必消除这个下界——Based 在 SQUAD 等任务上依然没有追平 Transformer++(第一节表格),差距缩小了,没有归零。
留下的开放问题:
- 门控能不能在保持 \(O(n)\) 训练、\(O(1)\) 推理的前提下,把 MQAR 式检索的维度需求降到与序列长度无关,而不靠局部精确 attention 兜底? 目前的证据(Based、GLA)都还依赖某种形式的局部精确检索或强门控,纯粹的”输入依赖衰减”能不能单独解决这个问题还没有定论。可读入口:Zoology 和 Based 的后续讨论(Hazy Research 实验室的系列工作)。
- RULER 一类长上下文基准上的差距,会不会随模型规模和训练数据增大而消失? 目前只有 7B 及以下规模的系统对比(RULER, COLM 2024),没有覆盖更大规模、专门为长上下文优化过的 RWKV/RetNet 变体的公开可信评测,这条 scaling 曲线是否收敛尚不清楚。
- 线性注意力/RetNet 的 kernel 生态需要多久才能接近 FlashAttention 系列的成熟度? FlashLinearAttention 已经证明可以在短序列上超过 FlashAttention-2,但这类 kernel 经过多硬件代际、多种序列长度和 batch 配置验证的时间还很短,工程追赶速度是个开放的、可以随时间检验的问题,而不是理论问题。
七、关键概念回顾
- 线性注意力:用核特征映射 \(\phi(q)^\top\phi(k)\) 替代 \(\exp(q^\top k/\sqrt d)\),把 attention 重写成可累积的矩阵状态 \(S_i,z_i\),状态大小不随序列长度增长。
- MQAR(multi-query associative recall):Zoology 提出的形式化任务,用来量化”循环状态模型”和”full attention”在关联检索上的表达力差距。
- WKV:RWKV 的核心算子,用逐通道指数衰减 \(w\) 加上当前 token 的加成项 \(u\),替代 softmax 相似度。
- retention:RetNet 的核心机制,用按 head 分级的固定衰减 \(\gamma\)(multi-scale decay)配合旋转位置编码,同时支持 parallel、recurrent、chunkwise 三种数学等价的计算表示。
- 数据依赖门控(data-dependent gating):衰减/更新参数是否随输入内容变化,是区分”朴素线性递归”和”选择性/门控”架构表达力的关键变量。
- I/O-aware kernel:把算法在 HBM 和片上 SRAM 之间的搬运次数降到最少的实现方式,决定理论复杂度能否兑现成真实吞吐。
八、常见误解
8.1 “所有线性时间模型都一样”
不一样。线性注意力、RWKV、RetNet、Mamba 在”用什么替代相似度/衰减、衰减是否数据依赖”上有明确差异,这组差异直接对应到 GLA 消融里 23.21→16.55→14.77 的困惑度梯度(第三、五节),不是文字游戏。
8.2 “复杂度低就一定更快”
不对。Gated Linear Attention 论文自己承认”现有的线性注意力实现缺乏 I/O 感知,因此更慢”,只有专门写 I/O-aware kernel(FlashLinearAttention)之后才反超 FlashAttention-2;同样,Mamba 的训练吞吐在长序列上落后于 Transformer++ 和 GLA(第四节)。\(O(n)\) 只是理论上界,不是实测保证。
8.3 “RWKV/RetNet 已经证明表达力不输 full attention”
没有。RWKV 论文自己在 Limitations 一节承认长上下文精确回忆的短板,RULER 的系统评测也证实了这一点;RetNet 在 FDA/SWDE 等关联检索任务上依然弱于标准 softmax attention(Table 3,GLA 论文)。它们证明的是”在特定规模和任务上可以接近或匹配”,不是”表达力等价”。
九、下一步
各种后 Transformer 路线已经展开。最后一篇回到大局:未来是某个架构替代 Transformer,还是 attention、SSM、MoE、检索和工具共同组成混合系统?
十、参考文献
核心论文
- Katharopoulos, A., Vyas, A., Pappas, N., Fleuret, F. “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention.” ICML 2020(同 arXiv:2006.16236)。线性注意力的原始定义、“as RNN”重写与状态形状分析。
- Choromanski, K. et al. “Rethinking Attention with Performers.” ICLR 2021。FAVOR+ 随机特征,softmax kernel 的无偏近似路线,与 Katharopoulos 等人”换核”路线的分野。
- Peng, B. et al. “RWKV: Reinventing RNNs for the Transformer Era.” Findings of the Association for Computational Linguistics: EMNLP 2023。WKV 算子、receptance 门控,以及论文第 9 节自述的长上下文与 prompt 顺序敏感性限制。
- Sun, Y., Dong, L., Huang, S. et al. “Retentive Network: A Successor to Transformer for Large Language Models.” arXiv:2307.08621, 2023。未见正式同行评审版本;retention 的 parallel/recurrent/chunkwise 三种表示、消融实验(Table 5、6)与自测训练成本(Table 4)。
- Yang, S., Wang, B., Shen, Y., Panda, R., Kim, Y. “Gated Linear Attention Transformers with Hardware-Efficient Training.” ICML 2024(同 arXiv:2312.06635)。FlashLinearAttention 的 I/O-aware 实现,门控消融(Table 4)与关联检索任务对比(Table 3)。
- Arora, S., Eyuboglu, S., Timalsina, A. et al. “Zoology: Measuring and Improving Recall in Efficient Language Models.” ICLR 2024(同 arXiv:2312.04927)。MQAR 任务定义与循环状态模型的关联检索维度下界。
- Arora, S. et al. “Simple linear attention language models balance the recall-throughput tradeoff.” arXiv:2402.18668, 2024。未见正式同行评审版本;Based 架构与线性注意力/RWKV-v5/Mamba 在关联检索切片上的量化对比(Table 1、4)。
- Hsieh, C.-P., Sun, S., Kriman, S. et al. “RULER: What’s the Real Context Size of Your Long-Context Language Models?” COLM 2024(同 arXiv:2404.06654)。长上下文多任务基准,RWKV-v5、Mamba 落后 Transformer 基线的实测数据。
- Dao, T., Gu, A. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.” ICML 2024。SSM 选择性扫描与线性注意力在结构化矩阵乘法意义下的统一,串联本文与 56|状态空间模型。
← 上一篇:56|状态空间模型 | 下一篇:58|后 Transformer 时代 →
同主题继续阅读
把当前热点继续串成多页阅读,而不是停在单篇消费。
【Transformer 与注意力机制】58|后 Transformer 时代:架构会消失还是会进化
后 Transformer 时代的判断不该是「哪个架构赢」,而是系统边界问题。本文用 Jamba、Griffin、NVIDIA 的受控对照实验说明纯 SSM 为何缺 in-context learning、混合架构为何更可能,用 RETRO/RAG/ReAct/MemGPT 说明模型与系统边界如何模糊,并结合 Hardware Lottery 与 SWE-bench/GAIA 等交互式评测讨论硬件和评测如何反向塑造架构。
【Transformer 与注意力机制】56|状态空间模型:Mamba、S4 的线性复杂度路径
Transformer 用 KV Cache 保存全部历史,SSM 用固定状态压缩历史。本文讲清 S4 为何长程却难训、Mamba 的 selective SSM 与 parallel scan 如何让递归既线性又能并行训练,比较推理状态与 KV Cache 的胜负边界,并给出 SSM 能否取代通用 LLM 的争论。
【Transformer 与注意力机制】42|FlashAttention:注意力计算的硬件级重写
FlashAttention 的关键不是近似注意力,也不是把公式改掉,而是重新安排标准 attention 在 GPU 内存层级里的计算路径。本文解释为什么标准 attention 的瓶颈常常是 HBM 读写,FlashAttention 如何用 tiling 和 online softmax 避免物化完整注意力矩阵,以及它为什么省显存、提吞吐,却没有消除 O(n²) 的根本复杂度。
【Transformer 与注意力机制】49|KV Cache:推理为什么是 O(n) 不是 O(n²)
自回归推理和训练不是同一种程序。本文解释 KV Cache 为什么成立:历史 token 的 Key/Value 一旦算出,在后续 decode 中不会改变;缓存它们可以避免反复重算前缀。文章同时讲清 prefill 与 decode 的差异、cache 显存公式、长上下文为什么受限,以及 PagedAttention、MQA/GQA、cache 量化等方向各自在解决什么。