想象你是一座大型图书馆的唯一管理员。有一位读者每次提问,你都必须把所有书从书架上一本本抽出来,两两配对翻看它们的相似之处——书越多,你要做的比对就成倍增长,这就是传统 Transformer 面对长文本时的困境。而 RWKV 带来了一位新管理员:他边走边读,只在一本固定大小的笔记本上记录要点,下次回答问题时只需翻开笔记本,再瞥一眼当前那页。这本书就是 WKV 状态,这个动作就是线性 RNN 的递推更新。这套方法让长文本处理从“平方级翻书”变成“线性级阅读”,推理时背包重量也不再增长。

岗位关联

理解 RWKV 的线性复杂度思路,是设计高效长序列模型、探索 Transformer 替代架构的必备视角,也是大模型推理优化与流式应用的重要技能点。

目录
  1. 可视化演示
  2. 瓶颈剖析(对应思维导图一级分支 1)
  3. RWKV定位(对应思维导图一级分支 2)
  4. WKV机制(对应思维导图一级分支 3)
  5. Token移位(对应思维导图一级分支 4)
  6. 复杂度对比(对应思维导图一级分支 5)
  7. 流式推理(对应思维导图一级分支 6)
  8. 工程实战(对应思维导图一级分支 7)
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示TCP连接建立的三次握手过程:SYN、SYN-ACK、ACK报文交互,客户端与服务端状态变迁与序列号追踪。

本节配有交互动效(进入视口后自动加载,约 3KB)。

TCP三次握手 客户端发送SYN报文请求建立连接,进入SYN_SENT状态 SYN=1, seq=x 三次握手确保双方收发能力正常,建立可靠连接

瓶颈剖析(对应思维导图一级分支 1)

老管理员的工作方式看似严谨,却有一个致命问题:每当有人提问,他都要把图书馆里所有书两两翻开,比对内容相似度。假如馆里有 100 本书,他要做 100×100 次比对;有 1000 本,就要做 100 万次。读的书越多,他花在“翻书配对”上的时间就呈平方级膨胀。

Transformer 的原理正是如此。它靠自注意力让序列中的每个 token 都与所有其他 token 建立联系,先产出一个形状为“序列长度×序列长度”的注意力矩阵。这个矩阵越大,计算量的增长速度就越惊人,于是长文本训练变得昂贵,推理时保存 KV 缓存的内存也成了负担。

用复杂度语言描述,自注意力层的计算量大致是:

其中 T 是序列长度,d 是每个 token 的向量维度。

图书馆里的书越多,管理员要做的比对就越多,最终把时间花在了“反复翻旧书”上。这就是 Transformer 处理长文本时的核心瓶颈:平方级计算和不断增长的缓存。

O(T2d)\displaystyle O\left(T^2 d\right)
  1. 两两计算序列中任意两个 token 之间都要计算一次相关性分数,所以组合数是 T 的平方。
    T×TT \times T
  2. 每次乘法计算一次相关性需要对 d 维向量做点积,包含 d 次乘法运算。
    O(d)O(d)
  3. 总复杂度组合数乘以单次开销,得到关于 T 的平方级复杂度。
    O(T2d)O(T^2 d)

RWKV定位(对应思维导图一级分支 2)

新管理员的策略完全不同:他不追求每次把全部书都摊开,而是随身带一本固定大小的笔记本,每读完一页就记下最关键的信息,然后合上书继续往前走。这个“记录—更新—带走”的循环,就是 RNN 式的状态传递。

RWKV 就是一条介于 RNN 与 Transformer 之间的路线。它本质上是一个线性 RNN:信息被压缩进一个固定大小的隐藏状态,新 token 到来时,只对状态做一次线性更新。妙的是,这种递推关系在训练时可以被展开成并行计算,不必像传统 RNN 那样必须串行地一个词一个词跑。

训练阶段的总计算量约为:

新管理员不再需要把整座图书馆背在背上,他只需要一本会自己更新的笔记。RWKV 用“记录摘要”替代“反复翻阅”,从架构定位上绕开了平方级复杂度。

O(Td)训练;O(d)推理每步\displaystyle O(Td) \qquad \text{训练};\qquad O(d) \qquad \text{推理每步}
  1. 训练时展开因为每一步的线性更新可以写成前缀和形式,所以整段序列可以一次性并行处理,开销与序列长度 T 成线性关系。
    O(Td)O(Td)
  2. 推理时递推生成下一个 token 时,只需把旧状态乘以衰减系数,再加上新 token 的信息,单步开销只与维度 d 有关。
    O(d)O(d)
  3. 常数状态无论读了多少本书,笔记本的页数不变,因此总内存固定。
    O(d)O(d)

WKV机制(对应思维导图一级分支 3)

笔记本上到底记了什么?新管理员的记录不是流水账,而是一份“加权摘要”:越久远的信息颜色越淡,印象越浅;刚刚读到的当前页会被额外重点标注。这份摘要既要保留旧书里的重要信息,又要给当前内容足够的分量。

WKV 正是这个加权摘要的数学化。它把历史 token 乘上指数衰减权重后累加,再除以总权重做归一化。当前 token 还会额外加上一个“提醒值” u,确保模型不会被遥远的过去淹没,也不会忽视眼前的信息。

RWKV 的核心公式可以写成:

看到公式不必害怕,它本质上就是“按距离打折的加权平均”。WKV 用两个累计向量代替了 T×T 的注意力矩阵,既保留了“谁重要谁就权重高”的注意力直觉,又把复杂度降到了线性。

wkvt=∑i=1t−1e−(t−1−i)w+kivi+eu+ktvt∑i=1t−1e−(t−1−i)w+ki+eu+kt\displaystyle \text{wkv}_t = \frac{\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i} v_i + e^{u + k_t} v_t}{\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i} + e^{u + k_t}}
  1. 历史加权每个历史位置 i 的向量 v_i 会乘上一个指数权重,权重由它与当前时刻的距离和衰减参数 w 决定。
    e−(t−1−i)wkie^{-(t-1-i)w} k_i
  2. 当前加成当前 token 使用额外的 u 作为“高亮标记”,避免模型只关注遥远历史而忽视当下。
    eu+ktvte^{u + k_t} v_t
  3. 归一化分子是所有加权值的总和,分母是所有权重之和,两者相除得到最终摘要。
    加权和权重和\frac{\text{加权和}}{\text{权重和}}

Token移位(对应思维导图一级分支 4)

新管理员还有一个习惯:翻到新一页时,会用余光扫一眼上一页的最后一行。这样即使笔记本只记录了摘要,他也能记得上一页的结尾语气,让整个故事读起来更连贯。这个“扫一眼上一行”的动作,就是 Token Shift。

RWKV 在每一层输入前,会把当前 token 的表示和上一个 token 的表示按一定比例混合,再把混合结果送入后续计算。这个操作很简单,却给线性递推补充了重要的局部信息——因为纯 RNN 状态容易丢细节,而相邻 token 的混合能让模型感知到“刚刚发生了什么”。

Token Shift 的数学表达非常简洁:

这个动作就像说话时自然地带出上一句的尾音,让局部语境更连贯。Token Shift 不需要额外增加模型参数,却能让线性 RNN 在逐词递推时保留一阶的“邻近感”,这是 RWKV 在长序列上依然能稳住质量的重要细节。

xt′=μxt+(1−μ)xt−1\displaystyle x'_t = \mu x_t + (1-\mu) x_{t-1}
  1. 当前表示取当前时刻 t 的输入向量 x_t 作为新信息。
    xtx_t
  2. 上一时刻取 t-1 时刻的输入向量 x_{t-1},它带着紧邻的上文。
    xt−1x_{t-1}
  3. 按比例混合用系数 μ 控制当前词和上一个词各占多少比例,形成新的输入。
    μxt+(1−μ)xt−1\mu x_t + (1-\mu) x_{t-1}

复杂度对比(对应思维导图一级分支 5)

让我们把新旧两位管理员放到同一条赛道上。老管理员每处理一个新问题,都要重新翻一遍所有旧书,工作量与书的数量平方上升;新管理员每遇到一个新词,只更新一行笔记,工作量基本不变。两者在短途时看不出差距,一旦书变多,差距立刻拉开。

这种差距可以用两个关键数字概括:RWKV 的 WKV 递推,每一步只做固定次数的加法和乘法,旧状态被“打折”后与当前信息合并。因为状态只有两个固定向量,所以每步计算的复杂度与历史长度无关。把多个步骤叠加起来,训练总开销也只有线性级别。

状态更新和归一化可以写成:

赛道的长度越夸张,新管理员的优势越明显。Transformer 需要保存所有 KV 缓存,RWKV 只需要这两个向量;当序列长度从千级涨到万级、十万级,线性复杂度的价值就从“省一点”变成了“能不能跑得动”的根本差别。

\displaystyle a_{t+1} = e^{-w} a_t + e^{k_t} v_t, \qquadb_{t+1} = e^{-w} b_t + e^{k_t}
  1. 旧信息打折扣上一步的累计向量 a_t 乘上 e^{-w},表示历史信息随时间衰减,而不是被无限累积。
    e−wate^{-w} a_t
  2. 写入新信息当前 token 的 k_t 和 v_t 组成一项新的贡献,被加进状态里。
    ektvte^{k_t} v_t
  3. 同步归一化b 也在做同样的递推,它记录的是权重总和,最后用 a/b 得出归一化的摘要。
    at+1bt+1\frac{a_{t+1}}{b_{t+1}}

流式推理(对应思维导图一级分支 6)

想象一位导游带领游客穿过长长的展馆,他边走边讲,口袋里只有一本小册子,不需要把整座展馆的展品说明都背下来。每走到一个新展品前,他只在小册子上更新几行,然后继续走。这正是 RWKV 在流式生成时的状态:逐词吐出答案,内存占用却纹丝不动。

传统 Transformer 在做流式生成时,必须把已经生成的所有 token 的 KV 缓存都留着,下一轮继续拼接计算。序列越长,缓存的负担越重。RWKV 则是每生成一个 token 后,只更新自己那两个向量,然后旧历史就可以“翻篇”了。于是它可以真正做到边生成边丢书,内存始终恒定。

每一步推理所需内存可以写成:

导游不会因为展馆变长而背包变重,RWKV 也不会因为生成文本变长而内存膨胀。这种恒定内存特性,让流式问答、实时对话和长文本续写变得像散步一样轻快,也正是它受部署工程师青睐的原因。

M=O(d)\displaystyle M = O(d)
  1. 固定状态RWKV 的隐藏状态只包含 a 和 b 两个 d 维向量。
    2×O(d)2 \times O(d)
  2. 与历史无关无论已经生成了多少 token,状态维度都不增长,所以内存中不出现 T。
    M≠O(T)M \neq O(T)
  3. 恒定背包因此整个生成过程的内存占用是一条水平线,而不是不断爬升的楼梯。
    M=O(d)M = O(d)

工程实战(对应思维导图一级分支 7)

理论说得再好,最终要把笔记本变成真正可用的工具。新管理员要上岗,还需要几个实操细节:笔记本怎么初始化、记录时如何避免数值溢出、哪些工作适合交给他。

RWKV 的工程实现有几个关键点。首先,衰减参数 w 通常按通道初始化成不同数值,让不同维度拥有不同的记忆长度,类似多头注意力关注不同粒度。其次,直接使用 e^{k_t} 可能导致指数爆炸,实际代码会对 k 做缩放或减最大值,保证数值稳定。最后,训练时可以通过并行扫描把递推展开,而推理时切换到逐词更新,两端共用同一套 WKV 参数。

从性能对比看,RWKV 在长文本任务上能显著降低训练显存和推理延迟,但在需要高强度“跨位置精确寻址”的任务上,比如要求逐字引用原文的极细节检索,它的表现可能不如全注意力精确。它更适合长文档摘要、对话生成、内容续写这类“抓住大方向”的场景。

工程上的选择从来不是非黑即白。新管理员并非要取代所有管理员,而是在“书太多、路太长”的场景里发挥独特优势。理解 RWKV 的线性 RNN 精神,就是学会在合适的场合放下重读全馆的执念,带上一本更聪明的笔记。

AI 应用场景

长文档摘要与问答

阶段:推理优化与长序列建模

解决的问题:传统 Transformer 对整本书做摘要时,注意力矩阵随页数平方增长,显存可能直接爆掉。RWKV 把全文压缩进固定状态,用线性扫描读完一本书,边读边记,最终生成摘要或回答问题。

具体案例:某知识库系统需要处理 50 万字的企业制度文档,过去必须截断成多个片段分别处理,再拼接结果;改用 RWKV 后,可以一次性读完整个文档,状态内存从与文档长度相关变成恒定,回答时还能结合全局上下文。

聊天机器人的流式对话

阶段:流式推理与部署

解决的问题:对话机器人每次回复都要携带全部历史消息,KV 缓存随轮数增长,延迟越来越高。RWKV 把历史压缩成固定状态,每一轮只更新一次,保证长对话延迟平稳。

具体案例:一款客服机器人上线后,用户在连续提问 100 轮后响应速度仍保持稳定;测试数据显示,使用 RWKV 方案后,第 100 轮的首 token 延迟比 Transformer 方案降低约 40%。

边缘设备上的长文本生成

阶段:模型压缩与端侧部署

解决的问题:手机、嵌入式设备内存有限,跑不动不断增长的 KV 缓存。RWKV 的恒定状态让它特别适合内存受限场景,可以把长文本生成能力带上端侧。

具体案例:某离线写作助手在低端平板上运行,用户输入一篇 8000 字文章,要求续写 3000 字;RWKV 版本在整过程中占用内存始终低于 500MB,而相同参数量的 Transformer 在长文续写中后期内存超过 2GB。

Python 演示

python49 行
import numpy as np def transformer_attention(q, k, v):    # 标准注意力:计算所有两两 token 的分数    T = q.shape[0]    d = q.shape[-1]    scores = np.matmul(q, k.T) / np.sqrt(d)    attn = np.exp(scores - scores.max(axis=-1, keepdims=True))    attn = attn / attn.sum(axis=-1, keepdims=True)    return np.matmul(attn, v) class SimpleRWKV:    # 把 WKV 状态压缩成两个固定大小的向量 a 和 b    def __init__(self, d, w=0.5, u=1.0):        self.d = d        self.w = w        self.u = u        self.a = np.zeros(d)        self.b = np.zeros(d)     def reset(self):        self.a = np.zeros(self.d)        self.b = np.zeros(self.d)     def step(self, k, v):        k_exp = np.exp(k)        self.a = np.exp(-self.w) * self.a + k_exp * v        self.b = np.exp(-self.w) * self.b + k_exp        return self.a / self.b def rwkv_forward(k, v):    # 按时间步递推,模拟 RWKV 的线性 WKV 更新    model = SimpleRWKV(k.shape[-1])    out = []    for t in range(len(k)):        out.append(model.step(k[t], v[t]))    return np.stack(out) if __name__ == '__main__':    rng = np.random.default_rng(0)    T, d = 64, 16    q = rng.normal(size=(T, d))    k = rng.normal(size=(T, d))    v = rng.normal(size=(T, d))    attn_out = transformer_attention(q, k, v)    rwkv_out = rwkv_forward(k, v)    print('attention shape:', attn_out.shape)    print('rwkv shape:', rwkv_out.shape)    print('first row sim:', np.dot(attn_out[0], rwkv_out[0]))
目的:用 NumPy 实现一个简化版 RWKV 的 WKV 递推,并与标准 Transformer 注意力对比输出结构,帮助理解线性状态更新与平方注意力在计算路径上的区别。
  1. scores = np.matmul(q, k.T) / np.sqrt(d)一次矩阵乘法生成 T×T 的注意力分数,复杂度随 T 平方增长。
  2. attn = attn / attn.sum(axis=-1, keepdims=True)对每一行做归一化,让权重之和为 1,这是注意力的“加权平均”。
  3. self.a = np.zeros(d)a 是固定大小的分子状态,负责累计加权后的历史信息。
  4. def step(self, k, v):每一步只更新两个向量,不保留任何历史 token 的 KV 缓存。
  5. self.a = np.exp(-self.w) * self.a + k_exp * v旧状态乘衰减系数后,加上当前 token 的贡献,这是 WKV 递推的核心。
  6. return self.a / self.b用累计权重 b 归一化,得到当前时刻的 WKV 摘要。
  7. for t in range(len(k)):按时间步串行递推,模拟推理时的逐个 token 消费。
  8. rwkv_out = rwkv_forward(k, v)输出与 attention 同样形状,但内部计算路径完全不同。

常见误区

误区听说 RWKV 是线性 RNN,就觉得它和注意力一点关系都没有,注意力能做到的它都做不到。

正解RWKV 的 WKV 机制本质上是一种“线性化的注意力”。它用指数衰减权重模拟距离远近,用当前 token 的 u 加成模拟注意力中的突出项,把注意力矩阵压缩成两个向量的递推。它没有注意力矩阵,却在做相似的事——按重要性加权聚合信息。

为何常见:大家习惯把架构分成“Transformer 派”和“RNN 派”,看到线性 RNN 标签就容易二极管式地认为它与注意力彻底对立。

误区觉得 Transformer 的平方复杂度是“强大”的代价,RWKV 线性复杂度省了算力,效果肯定大打折扣。

正解复杂度只描述计算规模,不直接决定模型能力。RWKV 通过 WKV 和 Token Shift 保留了关键的长程依赖和局部信息,在很多生成任务上能达到接近 Transformer 的效果,同时还能处理更长的序列。当然,在需要精确逐字检索的极端场景下,压缩状态确实会丢细节。

为何常见:人们容易把“计算更强”等同于“效果更好”,忽略了信息压缩和归纳偏差带来的效率收益。

误区认为只要是 Transformer 变体,推理时就必须把 KV 缓存都存下来,RWKV 只是训练快一点。

正解RWKV 推理时只需要维护 a 和 b 两个固定大小的向量,内存是 O(d),与历史长度无关。正因如此,它才能支持超长上下文和流式生成,而不必担心缓存不断膨胀。

为何常见:很多人的经验来自 Transformer 的 KV 缓存机制,容易把“自回归生成必须保留历史”的直觉套用到所有架构上。

误区看到 Token Shift 混合相邻 token,以为它替代了位置编码,模型从此有了完整位置信息。

正解Token Shift 只混合当前与前一个 token 的输入,提供的是“局部突兀感”的一阶线索,不是完整的绝对或相对位置编码。RWKV 真正的位置感知来自 WKV 里的指数衰减权重 w,它让距离越远的 token 权重越低,从而隐式建模相对距离。

为何常见:Token Shift 在输入层制造了“顺序感”,让人误以为它是位置编码的替代品,而忽略了衰减权重才是真正的距离控制器。

误区既然 RWKV 又快又省内存,那就应该在所有场景全面替换 Transformer。

正解RWKV 在线性复杂度和恒定内存上有明显优势,但全注意力的精确寻址能力在需要逐字引用、复杂推理和高质量对齐的任务中仍然重要。实际工程中两者可以互补,甚至可以做混合架构,而不是非此即彼。

为何常见:新架构出现时,人们容易陷入“革命叙事”,把相对优势夸大成绝对优势。

面试题

Q1. Transformer 自注意力的计算复杂度主要受哪个量的平方影响?

单选题 · easy

  • A.序列长度← 正确答案
  • B.隐藏层维度
  • C.训练轮数
  • D.批处理大小

答案:A。自注意力需要计算序列中所有 token 两两之间的分数,因此计算量随序列长度 T 平方增长,即 O(T²d)。隐藏层维度 d 只做线性贡献,训练轮数和批大小与单步复杂度无关。

Q2. RWKV 推理时保存历史信息所用的状态大小是?

单选题 · easy

  • A.固定的向量← 正确答案
  • B.随长度增长
  • C.随平方增长
  • D.完全无状态

答案:A。RWKV 把历史压缩进 WKV 的 a 和 b 两个固定向量中,大小只与维度 d 有关。它不像 Transformer 那样保留全部 KV 缓存,所以内存恒定,不随生成长度增加。

Q3. WKV 公式中参数 w 的主要作用是什么?

单选题 · easy

  • A.控制衰减速度← 正确答案
  • B.控制学习率
  • C.增加模型层数
  • D.替换激活函数

答案:A。w 出现在指数项 e^{-w} 中,决定历史信息随距离衰减的快慢,类似“记忆的半衰期”。它不参与优化器设置,也不改变网络深度或激活函数。

Q4. Token Shift 操作在 RWKV 中的作用是?

单选题 · mid

  • A.混合相邻时间步信息← 正确答案
  • B.随机丢弃部分 token
  • C.交换位置编码顺序
  • D.把句子整体反转

答案:A。Token Shift 把当前 token 与上一个 token 按比例混合,为线性递推补充局部上下文线索。它不是数据增强手段,也不是位置编码或序列变换。

Q5. RWKV 在训练阶段为什么可以并行处理长序列?

单选题 · mid

  • A.递推关系可展开成并行扫描← 正确答案
  • B.因为它完全不需要注意力
  • C.因为序列被截断成小块
  • D.因为它使用更大学习率

答案:A。RWKV 的 WKV 更新本质是线性递推,可以写成前缀和形式,因此训练时能利用并行扫描一次性算出所有时间步,和 Transformer 一样享受 GPU 并行。其他选项都不是并行化的原因。

Q6. RWKV 每步推理的内存复杂度是?

单选题 · mid

  • A.O(d)← 正确答案
  • B.O(T)
  • C.O(T²)
  • D.O(Td)

答案:A。RWKV 推理只维护 a、b 两个 d 维状态向量,单步内存是 O(d)。它不需要像 Transformer 那样保存长度为 T 的 KV 缓存,因此内存与序列长度解耦。

Q7. WKV 公式中的 u 项起到什么作用?

单选题 · hard

  • A.突出当前信息← 正确答案
  • B.计算归一化分母
  • C.替代位置编码
  • D.控制梯度裁剪

答案:A。u 是当前 token 的“加成项”,让模型在聚合历史摘要时不会完全被过去淹没。它只是权重公式中的额外偏置,不负责归一化,也不是位置编码或优化器设置。

Q8. 下列哪种任务更适合优先使用 RWKV 架构?

单选题 · hard

  • A.超长文档流式续写← 正确答案
  • B.短句情感分类
  • C.图像像素级分割
  • D.结构化 SQL 精确查询

答案:A。RWKV 在超长文本生成时线性复杂度和恒定内存优势最明显。短句任务体现不出优势,图像分割不是其目标领域,结构化精确查询需要强逐字寻址能力,更适合全注意力。

Q9. RWKV 相比 Transformer 的优势有哪些?

多选题 · mid

  • A.训练复杂度线性← 正确答案
  • B.推理内存恒定← 正确答案
  • C.完全不需要数据
  • D.长序列更友好← 正确答案

答案:ABD。RWKV 训练 O(Td)、推理 O(d),内存固定,因此长序列更友好。它仍然需要数据训练,“不需要数据”显然错误。

Q10. 关于 RWKV 的状态递推,下列说法正确的有?

多选题 · hard

  • A.只需维护 a 和 b 两个向量← 正确答案
  • B.可以用并行扫描加速训练← 正确答案
  • C.历史信息完全不会衰减
  • D.每一步都要重算全部历史

答案:AB。RWKV 把历史压缩进 a、b 两个向量,递推可展开成并行扫描,所以训练能并行。历史信息会按 e^{-w} 衰减,而不是完全保留;每步也只做固定更新,不需要重算全部历史。