跳到主内容
快讯直播
AI智模界
AI 词典

RWKV 架构:训练像 Transformer,推理像 RNN

RWKV(Receptance Weighted Key Value)是一种把 Transformer 的并行训练能力和 RNN 的线性推理能力结合起来的序列模型架构。

先看两个极端。Transformer 像一场圆桌会议:每个词都能直接看到所有词,训练时 GPU 可以并行计算,效率很高;但生成文本时,每多写一个字,都要把之前所有字的键值缓存(AI 词典:KV Cache">KV Cache)重新过一遍,显存和计算量随上下文长度增长很快。传统 RNN(Recurrent Neural Network,循环神经网络)像传话游戏:每个时刻只接收上一个时刻的状态和当前输入,推理时一步一个脚印,计算量随长度线性增长,状态大小固定;但训练时后一步依赖前一步,没法并行,长序列还容易忘事。

RWKV 想两全。它设计了一种带衰减的加权求和,把注意力里的 softmax 换成可学习的指数衰减,使同一套计算既能写成递推形式(RNN 模式),也能写成并行形式(Transformer 模式)。训练时用并行模式,像 Transformer 一样吃满 GPU;推理时用递推模式,只需维护一个固定大小的状态,每来一个 token 更新一次,不需要让 KV Cache 随长度膨胀。可以比作“可切换的流水线”:训练时所有工位同时开工,推理时切成传送带,每个工位只拿上一站的结果和当前原料。

维度Transformer传统 RNN/LSTMRWKV
训练并行可以不行可以
推理计算随长度平方增长随长度线性随长度线性
推理状态KV Cache 随长度增长固定大小隐藏状态固定大小状态
长程记忆直接注意力,较强易衰减用衰减机制平衡

和相邻概念的区别:RWKV 常被归入线性注意力或高效序列模型一类,和 Mamba(基于状态空间模型 SSM)目标类似,都想打破 Transformer 的平方复杂度。区别在于 RWKV 的递推更像带门控的线性注意力,名字里的 Weighted Key Value 就体现了它保留 key/value 加权求和的思路。具体实现细节各版本有差异,以官方页面为准。

对从业者,RWKV 的价值在于长文本、流式输出和端侧部署:显存占用更可控,推理成本随长度线性增长。对普通人,它可能让本地跑模型、长对话变得更便宜。但它的生态和工具链不如 Transformer 成熟,很多现成优化和加速库未必直接支持,选型时还得看具体场景。

AI 生成本文由 AI 基于公开信息自动生成,仅供参考。