RWKV架构原理：RNN的Transformer式复兴

可视化演示

展示RWKV时间混合的指数衰减记忆机制，以及并行训练与线性推理的效率对比。

架构缘起

让我们回到那间教室。Transformer老师每次写板书都要“翻全班旧作文”，当班级人数从20变成40，他要做的工作不是翻一倍，而是翻四倍——因为每个学生都要和所有其他学生两两比较，这就是二次方增长的代价。另一边，RNN老师只拿小本子按顺序记，读起来很快，但记到后面就把前面忘了，而且全班只能一个一个排队交作业，训练慢得让人着急。

RWKV的诞生，正是想同时拿到两个老师的好处：像RNN一样推理时只读一块黑板，又像Transformer一样训练时全班并行。它没有发明全新的架构魔法，而是把Transformer里最昂贵、最消耗内存的自注意力模块，替换成两个更“朴素”的模块：时间混合和通道混合。

传统Transformer的自注意力可以写成下面的公式：

RWKV老师不翻旧账，他把所有历史信息浓缩成一块“黑板状态”。黑板虽然小，但更新规则很聪明，既能保留顺序感，又能在训练时算得快。这，就是“RNN的Transformer式复兴”的第一块基石。

系统全景

黑板教学法有一套固定流程：新同学先把手上的便签和上一张旧便签“叠影”一下，感受时间先后；然后把要点写进黑板，让全班共享；最后按科目分栏整理，方便不同同学各取所需。RWKV的架构也按这个流程组装起来，形成四个核心部件。

这四个部件分别是：输入嵌入、时间混合、通道混合和输出头。输入嵌入把词变成向量；时间混合负责跨时间步传递信息；通道混合负责在同一时刻做特征变换；输出头把最终向量映射成下一个词的概率。前两个部件组合起来替代了Transformer自注意力，后两个部件替代了Transformer的前馈网络（FFN）。

整个前向过程可以写作下面的流程：

RWKV把“自注意力”这块昂贵的拼图拆掉，换成了两块更轻巧的积木：TimeMix负责班级内部的信息共享，ChannelMix负责把知识整理成不同科目。整条流水线既保留了Transformer堆叠层的熟悉结构，又悄悄把内存陷阱填平了。

WKV状态

那块黑板并不只是一块白板，而是写满了“加权总分”。越久远的笔记，字迹会随时间指数褪色；而今天刚写的内容，老师会用红笔额外加粗。RWKV的核心状态就叫WKV（Weighted Key Value），它维护着一个由历史Key和Value加权得到的动态记忆。

WKV可以被理解为一种“线性注意力”。传统注意力用softmax让所有历史token争抢权重，WKV则给每个历史位置一个固定的指数衰减权重，距离当前越近权重越大。这个衰减速度由一个可学习的向量w控制，而且每个通道都有自己独立的衰减快慢，就好像黑板上的不同科目褪色速度不一样。

WKV的经典计算公式如下：

如果把Transformer比作“全班重读所有作文”，WKV就是在黑板上保留一份“滚动总结”：每来一条新笔记，旧内容按时间褪色，新内容被红笔加粗，最后归一化得到当前最合理的知识状态。它不是没有回忆，而是把回忆压缩进了一块固定大小的黑板。

时间混合

新同学想理解“今天讲的微分方程”，至少得先看一眼昨天讲的“极限”。如果完全没有历史信息，每个词都像失忆一样孤立。RWKV的时间混合模块用一个叫Token Shift的小技巧，让每个token先和上一个token做一次“叠影”，再进入后续计算。

具体做法是：把x_t和x_{t-1}按一定比例线性插值，然后将结果分别投影成r和k。r叫Receptance，可以理解为“黑板是否打开门接收新知识”；k叫Key，是要写入黑板的关键词。通过这种“叠影”，模型不需要额外绝对位置编码，就能感知谁先谁后。

公式可以写成两路并行计算：

时间混合像是给黑板加了一道“时间标签”：每一笔新字都带着上一笔的淡淡墨迹，因此顺序信息自然渗入。正因为这个简单的“叠影”，RWKV不再需要Transformer那种复杂的位置编码，也能知道谁先谁后。

通道混合

黑板上的信息不能只按时间堆叠，还要按科目分栏。数学公式不要混到语文古诗里，这一步叫通道混合。它负责在同一个时间位置上，把向量的不同特征通道进行非线性重组，相当于Transformer里的前馈网络。

通道混合的做法和传统FFN类似，但激活函数更有特点：它使用平方ReLU代替普通ReLU，并用sigmoid门控控制信息流量。平方ReLU会把正特征放大得更剧烈，同时把负特征压成零，让模型学会“哪些科目要重点记，哪些科目直接擦掉”。

通道混合的公式如下：

如果把时间混合比作“按时间顺序更新黑板”，通道混合就是“按科目维度整理黑板”。两者各司其职，共同替代了自注意力与FFN，让RWKV在每一层都能高效地“读黑板、做笔记、分栏存放”。

并行训练

黑板更新看似是串行的——新内容必须等旧内容写完才能写。但RWKV发现，黑板的状态本质上是一个“累加器”：分子S不断累加历史Key-Value对的贡献，分母C不断累加历史权重。累加器这类运算，恰恰可以用前缀和算法让所有位置同时计算。

这意味着训练时，RWKV不需要像普通RNN那样一个时间步一个时间步地跑。它可以把整个序列的K、V一次性算出来，再用前缀和（或分块扫描）得到所有位置的S和C，最后一起得到输出。这就是“Transformer式并行训练”的核心秘密。

状态递推可以写成：

RWKV把“必须排队写黑板”改成了“全班先用公式算好自己那份，再合并累加结果”。这种改写，让训练时GPU可以被充分榨干，而推理时又恢复到只需更新两个小状态变量的RNN形态。

推理优势

推理阶段，黑板教学法的优势终于彻底爆发：从第一个词到最后第十万个词，老师只需要同一块黑板和一支粉笔。不需要把全班所有历史作文堆在桌上，更不需要为每一个新词重新翻一遍旧账。

RWKV推理时，内存只取决于模型自身维度，和输入长度无关。不管生成100个token还是一万个token，固定状态S和C的大小不变。同时每生成一个新token，只需要做常数次矩阵运算，总时间随序列长度线性增加，而不是平方增加。

用复杂度公式表达就是：

传统Transformer的KV缓存会随对话轮数越积越大，就像老师把全班旧作文越堆越高。RWKV则永远只带一块黑板，无论聊多久，桌面依然整洁。这也是它常被用于机器人、手机助手等资源受限场景的原因。

权衡取舍

黑板字迹会随时间褪色，这是RWKV最明显的代价。WKV中的衰减因子e^{-dw}意味着，当两个token的距离d很大时，历史贡献趋近于零。因此对于需要精确“翻几十万年前旧账”的超长程任务，RWKV不如完整的全局注意力。

不过，这个“遗忘”并非缺陷，而是一种工程取舍。相比Transformer的精确全量检索，RWKV更像人类记忆：近事清晰，远事模糊。它用可控的衰减换来了恒定内存和线性推理，在大多数实际场景中，这种平衡点是可以接受的。

数学上这个衰减极限可以写成：

RWKV的启示不是“替代Transformer”，而是提供一条“效率优先”的路线。如今许多混合架构开始把RWKV等线性注意力与局部注意力结合，既要黑板的省事，又要档案室的精确。未来的大模型，可能既不是纯Transformer，也不是纯RNN，而是一块会分区、会褪色、也会被红笔加粗的智能黑板。

AI 应用场景

Python 演示

import numpy as np def sigmoid(x):    return 1.0 / (1.0 + np.exp(-x)) class TinyRWKV:    def __init__(self, dim):        self.dim = dim        self.wk = np.eye(dim) * 0.5        self.wv = np.eye(dim) * 0.5        self.wr = np.eye(dim) * 0.1        self.u = np.ones(dim) * 1.0        self.state_s = np.zeros(dim)        self.state_c = np.zeros(dim)     def reset(self):        self.state_s = np.zeros(self.dim)        self.state_c = np.zeros(self.dim)     def step(self, x_t, x_prev):        shift = 0.7 * x_t + 0.3 * x_prev        k = self.wk @ shift        v = self.wv @ shift        r = sigmoid(self.wr @ shift)        a = np.exp(k)        u_a = a * np.exp(self.u)        wkv = (self.state_s + u_a * v) / (self.state_c + u_a)        self.state_s = self.state_s + a * v        self.state_c = self.state_c + a        return r * wkv     def forward_parallel(self, xs):        shift = np.empty_like(xs)        shift[0] = xs[0]        shift[1:] = 0.7 * xs[1:] + 0.3 * xs[:-1]        k = shift @ self.wk        v = shift @ self.wv        r = sigmoid(shift @ self.wr)        a = np.exp(k)        u_a = a * np.exp(self.u)        cum_s = np.cumsum(a * v, axis=0)        cum_c = np.cumsum(a, axis=0)        s_prev = np.vstack([np.zeros((1, self.dim)), cum_s[:-1]])        c_prev = np.vstack([np.zeros((1, self.dim)), cum_c[:-1]])        wkv = (s_prev + u_a * v) / (c_prev + u_a)        return r * wkv     def check_consistency(self, xs):        self.reset()        loop_out = []        prev = np.zeros(self.dim)        for x in xs:            loop_out.append(self.step(x, prev))            prev = x        self.reset()        return np.stack(loop_out), self.forward_parallel(xs) if __name__ == "__main__":    rng = np.random.default_rng(0)    model = TinyRWKV(8)    xs = rng.normal(size=(5, 8))    loop_out, parallel_out = model.check_consistency(xs)    print("loop output shape:", loop_out.shape)    print("parallel output shape:", parallel_out.shape)    print("outputs consistent:", np.allclose(loop_out, parallel_out, atol=1e-6))

1. shift = 0.7 * x_t + 0.3 * x_prev模拟Token Shift，把当前词与上一个词混合，注入顺序信息。
2. a = np.exp(k)把Key变成正数权重，作为写入黑板的强度。
3. wkv = (self.state_s + u_a * v) / (self.state_c + u_a)用分子和分母的比值得到加权平均后的黑板内容。
4. self.state_s = self.state_s + a * v更新分子状态，累加历史Key-Value贡献。
5. self.state_c = self.state_c + a更新分母状态，累加历史权重，保证归一化。
6. shift[1:] = 0.7 * xs[1:] + 0.3 * xs[:-1]并行版本一次性算出所有位置的Token Shift。
7. np.cumsum(a * v, axis=0)用前缀和一次算出所有时刻的分子状态，这就是RWKV能并行训练的根基。
8. np.allclose(loop_out, parallel_out)验证顺序推理和并行计算的输出完全一致。

常见误区

面试题