RWKV 架构及历史

以下是 RWKV 架构相关的知识,包括 RWKV 架构名称的由来、RWKV 架构特性、架构历史、以及每个架构的模型发布等内容。

RWKV 是 RNN 的一个变体。因此,有必要先介绍:什么是 RNN 架构,以及 RNN 架构的隐藏状态是什么?

RNN 架构和隐藏状态

循环神经网络(RNN)是一种广泛应用于深度学习领域的神经网络模型。

RNN 网络在运行过程中会保持一个隐藏状态(state),隐藏状态可看作 RNN 模型的“心理状态”。就像人类在思考时,脑海中会保留与当前事件关联度最高的“关键信息”。随着思考内容的变化,我们脑海中的“关键信息”会不断更新。同样的,RNN 网络也会通过特定的函数不断更新其隐藏状态

经典 RNN 图解

如图所示,RNN 网络依次处理每一个输入的 token,并根据“当前隐藏状态”来预测下一个可能出现的 token(如果有需求)。每处理一个 token,RNN 会将结果反馈给网络自身,从而“更新其隐藏状态”,再用“更新后的状态”去预测下一个 token。如此循环,直到达到“完成任务”的状态。

作为 RNN 的一种变体,RWKV 支持对隐藏状态进行微调(state tuning (opens in a new tab))。通过调整模型的“心理状态”,可使 RWKV 模型在特定任务上表现更佳。

RWKV 架构名称的由来

RWKV 架构的名称来源于时间混合和通道混合块中使用的四个主要模型参数,分别如下:

  • RR:Receptance,作为过去信息的接受程度的接受向量
  • WW:Weight,位置权重衰减向量,可训练的模型参数
  • KK:键(Key)是类似于传统注意力中 KK 的向量
  • VV:值(Value)是类似于传统注意力中 VV 的向量

RWKV 架构的发展历程

2020 年,BlinkDL 开始研究 Transformer ,立刻发现其有两个明显的改进方向:引入显式 decay 和 Token-shift(或者说短卷积)。在 https://github.com/BlinkDL/minGPT-tuned (opens in a new tab) 上测试后,发现这些技巧对于 Transformer 的性能有显著提升。

随后,他注意到 Apple 的 Attention Free Transformer (opens in a new tab) (AFT)论文并对其进行测试,发现这两种技巧也为 AFT 带来了显著的性能提升。

RWKV-V1

2021 年 8 月,RWKV 架构的初版:RWKV-V1 被提交到 RWKV-LM 仓库 (opens in a new tab) 中。RWKV-V1 首次 commit (opens in a new tab) 于 2021 年 8 月 9 日。

RWKV-V1 使用长卷积代替 Attention 机制,其架构由交替的 Time-mix 和 Channel-mix 组成。Channel-mix 是 Transformer 的 GeGLU 层的变种。Time-mix 则是对于 AFT 的显著改进

 Time-mix :TMt,c=sigmoid(Rt,c)uWt,u,csoftmaxt(Ku,c)Vu,c\text { Time-mix }: \mathbf{TM}_{t, c}=\operatorname{sigmoid}\left(R_{t, c}\right) \cdot \sum_{u} W_{t, u, c} \cdot \operatorname{softmax}_{t}\left(K_{u, c}\right) \cdot V_{u, c}  Channel-mix CMt,c=sigmoid(Rt,c)dWc,dgelu(Kt,d)Vt,d\text { Channel-mix } \mathbf{CM}_{t, c}=\operatorname{sigmoid}\left(R_{t, c}\right) \cdot \sum_{d} W_{c, d} \cdot \operatorname{gelu}\left(K_{t, d}\right) \cdot V_{t, d}

其中 RRKKVV 由输入线性变换生成,WW 则是长卷积中的卷积核。

ℹ️

Time-mix 和 Channel-mix 的结构设计均基于 RRWWKKVV 四个主要参数,这便是 RWKV 这个名称的由来。

RWKV-V1 更接近 Linear Transformer 而非 RNN ,因为其每一个时间步都依赖于历史时刻所有的输入。

RWKV-V2-RNN

RWKV-V2 版本首次为 RWKV 实现了 RNN 模式,伪代码图如下:

RWKV-v2-RNN-Architecture

代码图的补充说明:

  • a 和 b 是 kv 和 k 的 EMA(exponential moving average)
  • c 和 d 是 a 和 b 加上纯 self-attention 效应(原地的自己对自己的注意力)
  • c / d 是记忆机制,因为如果某个字在某个通道的 k 很强,且 W 接近 1,那么这个字就会被后文记住
  • TT, KK, VV, RR, WW, XX, PP 都是可以训练的参数
  • WW 使用预计算的值进行初始化,与 Alibi(Attention with Linear Biases)有类似之处,但都是可训练的非固定值,且在每个通道都可以不同,表达能力显著提高。
  • 发明了一种 headQK 机制,让模型可以直接从前文复制或避免生成某些字,这对于 ICL(in context learning) 很重要。几年后,其它研究者也发现了这一现象。但为了测试纯 RNN 的能力极限,RWKV 主线模型并未加入这一功能。
q = self.head_q(x)[:,:T,:] 
k = self.head_k(x)[:,:T,:] 
c = (q @ k.transpose(-2, -1)) * (1.0 / 256)
c = c.masked_fill(self.copy_mask[:T,:T] == 0, 0)
c = c @ F.one_hot(idx, num_classes = self.config.vocab_size).float()       
x = self.head(x) + c

通过指数衰减(exponential decay),RWKV-V2 实现了 RNN 形式的推理:模型每次生成都依赖于上一个时刻的输入和当前时刻的输入,同时拥有固定大小的隐藏状态(在 RWKV-V2 里是 a 和 b)。

RWKV-V2 的自注意力层简化表达:

xt+1=σ(Rxt)exp(Kxt)(Vxt)+exp(W)atexp(Kxt)+exp(W)btx_{t+1} = \sigma(Rx_t) \cdot \frac{ \exp(Kx_t) \cdot (V_xt) + \exp(W) \cdot a_t }{ \exp(Kx_t) + \exp(W) \cdot b_t }

RWKV-V3

RWKV-V3 是一个短期的过渡版本,其相对 RWKV-V2 使用更全面的 token-shift(对 SA 和 FF 层中的 R / K / V 分别使用不同的可训练 TimeMix 因子) :

xx = self.time_shift(x)
xk = x * self.time_mix_k + xx * (1 - self.time_mix_k)
xv = x * self.time_mix_v + xx * (1 - self.time_mix_v)
xr = x * self.time_mix_r + xx * (1 - self.time_mix_r)

RWKV-V3 对比 V2 的 token-shift 改进图示:

此外,这个版本使用 preLN 代替 postLN(更稳定,且收敛更快):

if self.layer_id == 0:
	x = self.ln0(x)
x = x + self.att(self.ln1(x))
x = x + self.ffn(self.ln2(x))

RWKV-V4

RWKV-V4 是 RWKV 架构的首个正式版本,版本代号 "Dove(鸽)" 。RWKV-V4 解决了 RWKV-v3 架构的数值稳定性问题,同时也发布了 RWKV 项目的首篇论文 《RWKV: Reinventing RNNs for the Transformer Era》 (opens in a new tab)

RWKV-V4 架构的论文由 RWKV 作者彭博和 RWKV 社区共同完成,初次发表于 2023 年 5 月 22 日。同年 10 月,RWKV-V4 架构论文被顶级会议 EMNLP 2023 (opens in a new tab) 收录。

RWKV-V4 论文中的模型架构图

ℹ️

上图是 RWKV-V4 论文中的模型架构概览,其中:

  • 左侧:RWKV-V4 的 time-mixing 和 channel-mixing 模块
  • 右侧:RWKV-V4 的语言建模流程

Token Shift 概念

RWKV-V4 论文中正式提出 “Token Shift” 概念:将 RWKV 接收的每个 token 和前一个 token 做混合,类似于大小 = 2 的一维因果卷积。Token Shift 让模型可以控制每个时间步将多少新信息与旧信息分配给每个头的接收、键、值和门向量(即 rrkkvvgg)。

RWKV-V4 语言建模的架构图,演示了 RWKV-V4 的 Token Shift 过程,以及 RWKV 的 state 更新过程:

RWKV-V4 语言建模的架构

RWKV-V4 版本的一大特点是 state(RNN 的隐藏状态)非常小,这是因为彭博专注于在“最小的 state 下实现最大的性能”。

DD 为模型主维数, LL 为模型层数。RWKV-4 模型的 state 大小为 5DL5DL,包括 2DL2DL 的 mixing 和 3DL3DL 的 WKV。

RWKV-V4 的模型发布

RWKV-V4 的研究(架构迭代和模型训练等)横跨了 2022 和 2023 年,主要发布了以下 4 类模型:

  • RWKV-4-Pile:在 331B tokens 的 Pile 数据集上进行预训练,包含 169m、430m、1B5、3B、7B、14B 六种参数的模型
  • RWKV-4-Raven:RWKV-4-Pile 的指令微调模型,使用 Alpaca、CodeAlpaca、Guanaco、GPT4All、ShareGPT 等开源数据集进行微调,包含 1B5、3B、7B、14B 四种参数的模型。
  • RWKV-4-World:在 RWKV World v1 数据集上训练的多语言模型(100 多种世界语言,训练数据包含 Pile),包含 169m、430m、1B5、3B、7B 五种参数的 Base 模型,有一些中文微调模型。
  • RWKV-4-Music:有 MIDI 和 ABC 两种格式的作曲模型,分别使用 bread-midi-dataset 和 irishman 乐谱数据进行训练。其中 ABC 模型参数为 82m,MIDI 模型参数分别为 120m 和 560m 。

RWKV-V4 模型对比表格:

模型名称描述训练数据(数据量)词表(词表大小)
RWKV-4-Pile (opens in a new tab)基于 Pile 数据集的预训练语言模型EleutherAI/pile (opens in a new tab) (331B tokens)20B_tokenizer (opens in a new tab) (50,277)
RWKV-4-Raven (opens in a new tab)RWKV-4-Pile 的指令微调语言模型Alpaca、CodeAlpaca、Guanaco、GPT4All、ShareGPT 等20B_tokenizer (50,277)
RWKV-4-World (opens in a new tab)基于 100 多种全球语言数据集的预训练语言模型RWKV World v1 (opens in a new tab) (590B tokens)rwkv_vocab_v20230424 (opens in a new tab) (65,536)
RWKV-4-Music(MIDI) (opens in a new tab)基于 MIDI 乐谱数据训练的音乐(作曲)模型bread-midi-dataset (opens in a new tab)MIDI-LLM-tokenizer (opens in a new tab)(20,095)
RWKV-4-Music(ABC) (opens in a new tab)基于 ABC 乐谱数据训练的音乐(作曲)模型irishman (opens in a new tab)-

此外,RWKV-V4 模型首次训练了 14B 参数的模型,并从 RWKV-V4-World 起正式启用了包含 100 多种世界语言的 World 系列数据集,以及对应的多语言词表 rwkv_vocab_v20230424

ℹ️

RWKV-V5 和 RWKV-V6 等更新的 RWKV 架构沿用了 World 系列数据集和 rwkv_vocab_v20230424 词表。

World 数据集会持续补充训练数据,并更新版本号:

  • World v1 ≈ 0.59T tokens
  • World v2 ≈ 1.1T tokens
  • World v2.1 ≈ 1.42T tokens,v2.1 模型的总训练数据 ≈ 2.5T tokens
  • World v3 ≈ 3T tokens,v3 模型的总训练数据 ≈ 5.5T tokens

RWKV-V5

RWKV-V5 是 RWKV-V4 架构的改良版本,版本代号“Eagle(鹰)”。

RWKV-V5 和 RWKV-V6 这两个架构在同一篇论文 《Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence》 (opens in a new tab)中发布。

该论文由 RWKV 作者彭博和 RWKV 社区共同完成,初次发表于 2024 年 4 月 9 日。同年 10 月,RWKV-5/6 架构论文被 LLM 领域顶级会议 COLM 2024 (opens in a new tab) 收录。

RWKV 5/6 架构概览

ℹ️

上图是论文中 RWKV 5/6 的架构概览,其中:

  • 左侧:RWKV 的 time-mixing 和 channel-mixing 模块
  • 右上角:作为 RNN 单元时的 RWKV time-mixing 模块
  • 底部的中间:前向传播模式下 RWKV-V5 time-mixing 的 token-shift 模块
  • 右下角:前向传播模式下 RWKV-V6 time-mixing 的 token-shift 模块
  • 虚线箭头(左侧,右上角):表示 RWKV-V6 架构的连接,但在 RWKV-V5 中不存在

RWKV-V5 架构优化细节

相对 RWKV-V4, RWKV-V5 的最重点改动在于引入了多头的基于矩阵值的状态(state),即论文中的 “multi-headed matrix-valued states”。

在 RWKV-V4 架构的 time mixing 计算中,uuwwkkvv 参数都是维度为 DD 的向量,而 head size 是 1 :

ttRWKV-4 u,w,kt,vtRDu, w, k_t, v_t \in \mathbb{R}^D, head size 1
0σ(r0)uk0v0uk0\sigma(r_0) \odot \frac{u \odot k_0 \odot v_0}{u \odot k_0}
1σ(r1)uk1v1+k0v0uk1+k0\sigma(r_1) \odot \frac{u \odot k_1 \odot v_1 + k_0 \odot v_0}{u \odot k_1 + k_0}
2σ(r2)uk2v2+k1v1+wk0v0uk2+k1+wk0\sigma(r_2) \odot \frac{u \odot k_2 \odot v_2 + k_1 \odot v_1 + w \odot k_0 \odot v_0}{u \odot k_2 + k_1 + w \odot k_0}
3σ(r3)uk3v3+k2v2+wk1v1+w2k0v0uk3+k2+wk1+w2k0\sigma(r_3) \odot \frac{u \odot k_3 \odot v_3 + k_2 \odot v_2 + w \odot k_1 \odot v_1 + w^2 \odot k_0 \odot v_0}{u \odot k_3 + k_2 + w \odot k_1 + w^2 \odot k_0}

而 RWKV-V5 则将 uuww参数分别进行对角化,而 kkvv 从维度为 DD向量转化为维度为 64 * 64 的矩阵, head size 大小改为固定的 64 。

RWKV-V5 的 time-mixing 时间步:

ttEagle (RWKV-5) diag(u)\text{diag}(u), diag(w)\text{diag}(w), ktk_t, vtR64×64v_t \in \mathbb{R}^{64 \times 64} for each head, head size 64
0r0(diag(u)k0v0)r_0 \cdot (\text{diag}(u) \cdot k_0^\top \cdot v_0)
1r1(diag(u)k1v1+k0v0)r_1 \cdot (\text{diag}(u) \cdot k_1^\top \cdot v_1 + k_0^\top \cdot v_0)
2r2(diag(u)k2v2+k1v1+diag(w)k0v0)r_2 \cdot (\text{diag}(u) \cdot k_2^\top \cdot v_2 + k_1^\top \cdot v_1 + \text{diag}(w) \cdot k_0^\top \cdot v_0)
3r3(diag(u)k3v3+k2v2+diag(w)k1v1+diag(w2)k0v0)r_3 \cdot (\text{diag}(u) \cdot k_3^\top \cdot v_3 + k_2^\top \cdot v_2 + \text{diag}(w) \cdot k_1^\top \cdot v_1 + \text{diag}(w^2) \cdot k_0^\top \cdot v_0)

RWKV-V5 前向传播(推理过程)的 time-mixing 计算公式:

t=lerp(xt,xt1)W,{r,k,v,g}\square_t = \text{lerp}_{\square}(x_t, x_{t-1}) W_{\square}, \quad \square \in \{ r, k, v, g \} w=exp(exp(ω))w = \exp(-\exp(\omega)) wkvt=diag(u)ktvt+i=1t1diag(w)t1ikiviR(D/h)×(D/h)wkv_t = \text{diag}(u) \cdot k_t^\top \cdot v_t + \sum_{i=1}^{t-1} \text{diag}(w)^{t-1-i} \cdot k_i^\top \cdot v_i \in \mathbb{R}^{(D/h) \times (D/h)} ot=concat(SiLU(gt)LayerNorm(rtwkvt))WoRDo_t = \text{concat}(\text{SiLU}(g_t) \odot \text{LayerNorm}(r_t \cdot wkv_t)) W_o \in \mathbb{R}^D

通过将 RWKV-V4 的向量变成矩阵,RWKV-V5 的 state 计算从“基于向量”变成了“基于 64×64 的矩阵值”,即 “matrix-valued states” 。假设当前 RWKV 模型的维度是 512 ,则可以说有 512/64 = 8 个头 (八头×64 维),这就是 RWKV-V5 的“多头-multi-headed” 概念。

因此,我们可以把 RWKV-V5 的优化细节总结为:RWKV-V5 消除了归一化项(RWKV-V4 公式中的分母),并引入了矩阵值状态代替以往的向量值状态。

通过这种方式,RWKV-V5 巧妙地扩大了 state 的规模,使得 RWKV 模型有更好的记忆力和模型容量。

ℹ️

DD 为模型主维数, LL 为模型层数。RWKV-5/6 模型的 state 大小为 66DL66DL,包括 2DL2DL 的 mixing 和 64DL64DL 的 WKV。

RWKV-V5 迭代过程

实际上,RWKV-V5 架构的研究并非一蹴而就。从 2023 年 7 月提出到 2023 年 9 月定型,RWKV-V5 的优化过程可分为三个子版本:RWKV-5.0、5.1 和 5.2。

RWKV-5.0

RWKV-5.0 重新设计了 RWKV-V4 架构的 wkvwkv 模块,其中 kkvv 两个参数从维度为 DD 的向量转化为维度为 64 * 64 的矩阵,因此 RWKV-5.0 的 wkvwkv 状态比 RWKV-V4 大 32 倍。

RWKV-5.0 架构仅发布了 0.1B 模型:RWKV-5-World-0.1B-v1-20230803-ctx4096 ,此模型仍然基于 RWKV-V4 时期的 World-v1 数据集训练。

RWKV-5.1

RWKV-5.1 在 RWKV-5.0 的基础上引入了 Time-mixing 门控机制,即在 time-mixing 模块中添加额外的矩阵 WgW_g 和门控激活函数 SiLU 实现。

RWKV-5.1 架构仅发布了音乐模型:RWKV-5-ABC-82M-v1-20230901-ctx1024 和 RWKV-5-MIDI-560M-v1-20230902-ctx4096。

RWKV-5.2

RWKV-5.2 在 RWKV-5.1 的基础上引入了对角衰减矩阵,也就是将 uuww 向量参数分别进行对角化。

RWKV-5.2 架构发布了 0.4B、1B5、3B、7B,以及 3B(ctx16k)六类模型,这些模型基于 World-v2 数据集训练。

RWKV-V5 子版本对比表格

架构版本优化细节发布模型数据集
RWKV-5.0在 RWKV-V4 的基础上重新设计了 wkvwkv 模块,其中 kkvv 两个参数从维度为 DD 的向量转化为维度为 64 * 64 的矩阵,因此 RWKV-5.0 的 wkvwkv 状态比 RWKV-V4 大 32 倍RWKV-5-World-0.1B-v1 (opens in a new tab)World-v1
RWKV-5.1在 RWKV-5.0 的基础上引入了 Time-mixing 门控机制,即在 time-mixing 模块中添加额外的矩阵 WgW_g 和门控激活函数 SiLU 实现RWKV-5-music (opens in a new tab) 系列,含 ABC-82M 和 MIDI-560Mirishman (opens in a new tab)bread-midi-dataset (opens in a new tab)
RWKV-5.2在 RWKV-5.1 的基础上引入了对角衰减矩阵,也就是将 uuww 向量参数分别进行对角化RWKV-5-World-V2.1 (opens in a new tab) 系列,参数包含 0.4B、1B5、3B、7B,以及 3B(ctx16k)World-v2
️🚫

RWKV-V5 全系列模型均已过时,建议使用 RWKV-V6 模型。

RWKV-V6

RWKV-V6 的版本代号是“Finch(雀)”,这个版本自 2023 年 10 月开发,是当前(2024 年 11 月)的稳定架构。

RWKV-V6 完整架构图

RWKV-V6 架构优化细节

RWKV-V6 在 RWKV-V5 的基础上引入了基于 LoRA 的动态递归机制,进一步优化了 Token Shift 和 time-mixing 过程。

RWKV-V5 的 Token Shift 和 RWKV-V4 类似,是一个非常简单的线性插值(linear interpolation - lerp),且这个线性插值是数据无关的(data-independent),只由参数 μ\mu 决定当前 token 和前一个 token 混合到模型输入的占比。

RWKV-V5 的 Token Shift 中,xxx1x-1 之间的线性插值公式如下:

lerp(a,b)=a+(ba)μx\text{lerp}_{\Box}(a, b) = a + (b - a) \odot \mu_x

RWKV-V6 则借鉴了低秩适应(LoRA)的技术,将静态的参数 μ\mu 替换成了动态的 LoRA : μx\mu_x 和每个 λ\lambda_\square 引入了维度为 DD 的可训练向量,并且每个 ARD×32A_\square \in \mathbb{R}^{D \times 32}BR32×DB_\square \in \mathbb{R}^{32 \times D} 引入了新的可训练权重矩阵。

RWKV-V6 的 Token Shift 中,xxx1x-1 之间的线性插值公式如下:

lora(x)=λ+tanh(xA)B\text{lora}_{\Box}(x) = \lambda_{\Box} + \tanh(x A_{\Box}) B_{\Box} ddlerp(a,b)=a+(ba)lora(a+(ba)μx)\text{ddlerp}_{\Box}(a, b) = a + (b - a) \odot \text{lora}_{\Box}(a + (b - a) \odot \mu_x)

相对 RWKV-V4/RWKV-V5 ,RWKV-V6 这种增强了数据依赖性的新型 Token Shift 有效地扩展模型的能力,每个通道分配的新旧数据量取决于当前和之前时间步骤的输入。

ℹ️

通俗地理解,这种动态递归机制使“重要的信息”可以有效地标记自身,以待在后续任务使用;而“不重要的信息”也可以标记自身,以减少或完全避免进入后续的数据流,从而为更重要的现有数据保留空间。

此外,如果某些信息对于特定任务没有用,那么动态递归机制可以允许这些信息预先被过滤掉。

RWKV-V6 的动态 Time Mixing 计算公式:

t=ddlerp(xt,xt1)W,{r,k,v,g}\square_{t} = \mathrm{ddlerp}_{\square}(x_t, x_{t-1}){W}_{\square}, \quad \square\in \{r,k,v,g\} dt=lorad(ddlerpd(xt,xt1))d_t = \mathrm{lora}_d( \mathrm{ddlerp}_d ( x_t, x_{t-1} ) ) wt=exp(exp(dt))w_t = \exp(-\exp(d_t)) wkvt=diag(u)ktTvt+i=1t1diag(j=1i1wj)kiTviR(D/h)×(D/h){wkv}_{t} = \mathrm{diag}(u)\cdot k_{t}^\mathrm{T} \cdot v_{t} + \sum_{i=1}^{t-1} \mathrm{diag}\left(\bigodot_{j=1}^{i-1}w_{j}\right) \cdot k_{i}^\mathrm{T} \cdot v_{i} \in \mathbb{R}^{(D/h) \times (D/h)} ot=concat(SiLU(gt)LayerNorm(rtwkvt))WoRDo_t = \mathrm{concat}\left(\mathrm{SiLU}(g_t) \odot \mathrm{LayerNorm}(r_{t} \cdot {wkv}_{t})\right){W}_o \in \mathbb{R}^{D} wkv=s+diag(u)kTv{wkv}' = {s} + \mathrm{diag}(u) \cdot k^\mathrm{T} \cdot v s=diag(w)s+kTv{s}' = \mathrm{diag}(w) \cdot {s} + k^\mathrm{T} \cdot v

wkvt{wkv}_{t} 的注意力计算可以用递归方式写成:

wkv=s+diag(u)kTv{wkv}' = {s} + \mathrm{diag}(u) \cdot k^\mathrm{T} \cdot v s=diag(w)s+kTv{s}' = \mathrm{diag}(w) \cdot {s} + k^\mathrm{T} \cdot v

与 RWKV-V5 不同,RWKV-V6 的 wtw_t 在整个序列中不是静态的。这是 RWKV-V6 衰减的核心变化:wtw_t 的每个通道可以根据数据依赖性独立变化,而以前是一个固定的学习向量。

上述新的 LoRA 机制用于获取混合向量。注意,LoRA 过程本身使用 Eagle(RWKV-V5) 风格的 Token Shift 值作为输入(RWKV-V6 完整架构图左下方μx\mu_x)而不仅仅是最新的 token。新的时间变化衰减 wtw_t 也需要再次应用 LoRA。

直观上,这是 Token Shift 的二阶变体,允许 wtw_t 的每个通道根据当前和先前 token 的混合来变化,而混合本身由两个 token 的嵌入决定。

RWKV-V6 的模型发布

RWKV-V6 架构本身并没有子版本,但发布了不同类型的模型,包括不同训练集的预训练基底模型,以及中文小说和日文微调模型。以下是基于 RWKV-V6 架构的模型概览:

模型类别模型状态模型描述训练数据
RWKV-6-World-v2过时基于 World-v2 数据集的多语言预训练模型World-v2
RWKV-6-World-v2.1稳定基于 World-v2.1 数据集的多语言预训练模型World-v2.1
RWKV-6-World-v3训练中基于 World-v3 数据集的多语言预训练模型World-v3
RWKV-6-ChnNovel稳定中文小说微调模型,在 RWKV-6-World-v2.1 的基础上,使用中文小说+指令数据微调World-v2.1,中文小说数据
RWKV-6-Jpn稳定日文微调模型,在 RWKV-6-World-v2.1 的基础上,使用日文+指令数据微调World-v2

RWKV-V6 的 state tuning

ℹ️

除了完整的模型权重,社区在 RWKV-V6 架构迭代期间研发出了 state tuning。这是一种新颖的微调方法,微调 RWKV 的初始 state 相当于最彻底的 prompt tuning,甚至可以用于 alignment,因为迁移能力很强。有关 state tuning 的详细方法,请查看 RWKV 微调- state tuning (opens in a new tab)

微调得来的 state 文件可以合并到基底模型中,也可以作为“RWKV 模型的增强插件”使用:即在加载基底 RWKV 模型之前初始化模型的 state ,以影响模型的回答文风、回答格式等效果。

有关 RWKV State 文件的加载方法,请查看 RWKV state 文件用法 (opens in a new tab)

RWKV 官方发布了不同类型的 state 文件

state 类型state 描述适用 RWKV 模型
chn-single-round中文单轮对话增强,更符合人类语言习惯,带丰富的 Emoji 表情通用 RWKV 模型
eng-single-round英文单轮对话增强,更符合人类语言习惯,带丰富的 Emoji 表情通用 RWKV 模型
chn-小说扩写-single-round中文单轮对话,会根据用户输入进行小说扩写(建议使用 RWKV-V6-ChnNovel 模型和 state 作为替代)通用 RWKV 模型
chn-打油诗-single-round中文单轮对话,会根据用户输入创作打油诗通用 RWKV 模型
chn-文言文-single-round中文单轮对话,回答的风格会偏向文言文通用 RWKV 模型
chn-文言文和古典名著-single-round中文单轮对话,回答的风格会偏向文言文和古典名著通用 RWKV 模型
OnlyForChnNovel_小说扩写用于扩写中文小说,适用于同尺寸的 RWKV-V6-ChnNovel 模型RWKV-V6-ChnNovel 中文小说

RWKV-7

RWKV-7 的架构代号是 “Goose(雁)” 。RWKV-7 超越了 attention / linear attention 范式,它的状态演化很灵活,可以解决在相同算力消耗下 attention 无法解决的问题。与此同时,RWKV-7 超越了 TC0 约束。

RWKV-7 的研究始于 2024 年 9 月,其预览版 RWKV-7 "Goose" x070.rc2-2409-2r7a-b0b4a 的训练代码首次提交于 RWKV-LM 仓库的 commit (opens in a new tab) 中。

RWKV-7 "Goose" x070.rc2-2409-2r7a-b0b4a 的时间步递归公式:

St=St1(diag(wt)αtβt)+vtk~tS_t = S_{t-1} \left( \text{diag}(w_t) - \alpha_t^\top \beta_t \right) + v_t^\top \tilde{k}_t αt=normalize(kt+bt)\alpha_t = \text{normalize}(k_t + b_t) βt=αtat\beta_t = \alpha_t \cdot a_t
ℹ️

其中,ata_t 可看作 in-context learning rate。RWKV-7 拥有很强的 ICL 能力。

RWKV -7 与历史版本(RWKV5/6)的时间步公式和状态更新机制对比:

模型版本时间步公式状态更新机制
RKWV-5 EagleSt=(w000wn)St1S_t = \begin{pmatrix} w_0 & \cdots & 0 \\ \cdots & \cdots & \cdots \\ 0 & \cdots & w_n \end{pmatrix} S_{t-1}Trainable State Decay
RKWV-6 FinchSt=(wt,000wt,n)St1S_t = \begin{pmatrix} w_{t,0} & \cdots & 0 \\ \cdots & \cdots & \cdots \\ 0 & \cdots & w_{t,n} \end{pmatrix} S_{t-1}Dynamic State Decay
RKWV-7 GooseSt=(wt,0,0wt,0,nwt,n,0wt,n,n)St1S_t = \begin{pmatrix} w_{t,0,0} & \cdots & w_{t,0,n} \\ \cdots & \cdots & \cdots \\ w_{t,n,0} & \cdots & w_{t,n,n} \end{pmatrix} S_{t-1}Dynamic State Evolution

核心机制代码

def ref_fwd(r, w, k, v, a, b):
    r = r.view(B, T, H, N)
    k = k.view(B, T, H, N)
    v = v.view(B, T, H, N)
    a = a.view(B, T, H, N)
    b = b.view(B, T, H, N)
    w = torch.exp(-torch.exp(w.view(B, T, H, N)))
    out = torch.zeros((B, T, H, N), device=DEVICE)
    state = torch.zeros((B, H, N, N), device=DEVICE)
 
    for t in range(T):
        kk = k[:, t, :]
        rr = r[:, t, :]
        vv = v[:, t, :]
        aa = a[:, t, :]
        bb = b[:, t, :]
        sab = torch.einsum('bhik,bhk,bhj->bhij', state, aa, bb)
        state = state * w[: , t, :, None, :] + sab + torch.einsum('bhj,bhi->bhij', kk, vv)
        out[:, t, :] = torch.einsum('bhj,bhij->bhi', rr, state)
 
    return out.view((B, T, C))
 
⚠️

RWKV-7 处于早期预览版本,其论文仍然由 RWKV 作者彭博和社区成员共同撰写。截至 2024 年 11 月,论文仍处于撰写阶段,暂无具体发布计划。

RWKV 架构特性

RWKV 大模型架构的特点包括:

  • 高效且稳定的推理速度
  • 低且固定的显存占用量(支持在CPU 上运行)
  • 能够处理无限上下文,非常适合长文本处理和多轮对话等应用
  • 对硬件友好,仅执行矩阵与矢量的乘法操作,无需 KV 缓存

RWKV 架构由一系列堆叠的残差块组成,每个残差块由具有循环结构的时间混合(timemixing)和通道混合(channel-mixing)子块组成,该循环通过将当前输入和上一时间步的输入之间进行线性插值来实现(在 RWKV-4 架构论文中,这一过程被称为 token shift)。RWKV 6 通过借鉴 LoRA 技术优化了 token shift 过程,使其 RWKV4/5 的简单线性插值(lerp)变成了数据依赖的、动态的线性插值(ddlerp)。

根据对不同模型的推理复杂度进行比较,Transformer 的时间复杂度为:O (T^2)、空间复杂度为:O (T^2),因此推理速度将会越来越慢,也越来越耗内存。而 RWKV 的时间复杂度为:O(T)、空间复杂度为 O(1)。RWKV 大模型通过对计算流程的优化,实现了恒定的推理速度,极大地减少了推理过程中的时间消耗。

此外,RWKV 架构设计显著降低了显存占用,使得模型即使在标准配置的 CPU 或非专业级的 GPU上也能高效运行,无需依赖昂贵或高端的计算硬件。这一突破性的进展使得大规模深度学习模型不再受限于特定的硬件平台,拓宽了应用范围。

RWKV 架构参考资料

有关 RWKV 架构的核心概念,可以查看 RWKV 150 行代码 (opens in a new tab)

有关 RWKV-6 架构的代码解读,可以参考这篇博客:介绍 RWKV-6 的模型设计,代码带注释 (opens in a new tab)

或者,你可以通过阅读 RWKV 论文学习:

如果你对基础有了掌握,可以开始在RWKV 主仓库 (opens in a new tab)中研究 RWKV 的训练和 CUDA 代码。

RWKV 架构发展的综述论文:RWKV 架构发展综述 | arXiv(2411.02795) (opens in a new tab)