RWKV——一种具有Transformer级别LLM性能的RNN

山行AI

发布于 2023-08-10 13:49:29

7520

发布于 2023-08-10 13:49:29

文章被收录于专栏：山行AI山行AI

前言

RWKV是一种具有Transformer级别LLM性能的RNN，也可以像GPT Transformer一样直接进行训练（可并行化）。它是100%无注意力的。您只需要在位置t处的隐藏状态来计算位置t+1处的状态。您可以使用“GPT”模式快速计算“RNN”模式的隐藏状态。

因此，它结合了RNN和Transformer的优点-出色的性能、快速推理、节省VRAM、快速训练、“无限”ctx_len和免费的句子嵌入（使用最终的隐藏状态）。

体验感受：一个字，快；两个字，很快；三个字，真的快。

Raven 14B 体验示例（亲测：生成速度很快）：

体验地址：https://huggingface.co/spaces/BlinkDL/ChatRWKV-gradio，PS：中文生成效果也尚可。

RWKV-4-World-7B 体验效果：

体验地址：Raven RWKV 7B - a Hugging Face Space by BlinkDL[1]

RWKV语言模型（以及我的LM技巧）

RWKV：具有Transformer级别LLM性能的并行化RNN（发音为“RwaKuv”，由4个主要参数R、W、K、V组成）

因此，它结合了RNN和Transformer的优点-出色的性能、快速推理、节省VRAM、快速训练、“无限”ctx_len和免费的句子嵌入（使用最终的隐藏状态）。

Raven 14B（在Alpaca+ShareGPT+...上微调）演示：https://huggingface.co/spaces/BlinkDL/ChatRWKV-gradio

World 7B（支持100多种世界语言）演示：https://huggingface.co/spaces/BlinkDL/RWKV-World-7B

ChatRWKV： 具有“stream”和“split”策略以及INT8。3G VRAM足以运行RWKV 14B :) https://github.com/BlinkDL/ChatRWKV

下载RWKV-4 0.1/0.4/1.5/3/7/14B权重：https://huggingface.co/BlinkDL

RWKV-4-World是最佳模型：支持100多种世界语言的生成、聊天和代码，还具有最佳的英语零-shot和上下文学习能力。

RWKV pip包: 链接[2]

os.environ["RWKV_JIT_ON"] = '1'
os.environ["RWKV_CUDA_ON"] = '0'  # 如果为'1'，则使用CUDA内核进行序列模式（速度更快）
from rwkv.model import RWKV  # pip install rwkv
model = RWKV(model='/fsx/BlinkDL/HF-MODEL/rwkv-4-pile-1b5/RWKV-4-Pile-1B5-20220903-8040', strategy='cuda fp16')

out, state = model.forward([187, 510, 1563, 310, 247], None)  # 使用20B_tokenizer.json
print(out.detach().cpu().numpy())  # 获取logits
out, state = model.forward([187, 510], None)
out, state = model.forward([1563], state)  # RNN具有状态（如果要克隆，请使用深拷贝）
out, state = model.forward([310, 247], state)
print(out.detach().cpu().numpy())  # 与上述结果相同

Cool Community RWKV Projects (check them!):

•saharNooby/rwkv.cpp[3] 使用ggml[4]进行快速i4 i8 fp16 fp32 CPU推理•harrisonvanderbyl/rwkv-cpp-cuda[5] 快速的Windows/Linux和CUDA/ROCM/Vulkan GPU推理（无需Python和PyTorch）•Blealtan/RWKV-LM-LoRA[6] LoRA微调•josStorer/RWKV-Runner[7] 酷炫的GUI

更多RWKV项目：链接[8]

加入我们的Discord：链接[9]（有很多开发者）

Twitter: 链接[10]

RWKV in 150 lines（模型、推理、文本生成）：链接[11]

RWKV预印本：链接[12]

RWKVpaper

RWKV介绍，以及使用100行numpy代码实现：链接[13] 链接[14]

使用RWKV的一篇酷炫论文（脉冲神经网络）：链接[15]

欢迎加入RWKV Discord 链接[16] 一起进行进一步的构建。我们现在有大量的计算资源（A100 40Gs）（感谢Stability和EleutherAI），所以如果你有有趣的想法，我可以运行它们。

RWKVeval2

RWKV在Pile的10000个ctx4k+文档中的[损失 vs. 令牌位置]。RWKV 1B5-4k在ctx1500之后基本保持平稳，但3B-4k、7B-4k和14B-4k有一些斜率，并且它们正在变得更好。这推翻了RNN无法建模长ctx长度的旧观点。我们可以预测RWKV 100B将会很棒，而RWKV 1T可能是你所需要的 :)

RWKVctxlen

使用RWKV 14B ctx8192的ChatRWKV：

RWKVchat

我相信RNN是基础模型的更好选择，因为：(1) 它对ASIC更友好（没有kv缓存）。(2) 它对RL更友好。(3) 当我们写作时，我们的大脑更类似于RNN。(4) 宇宙也像一个RNN（因为局部性）。Transformers是非局部模型。

RWKV-3 1.5B在A40（tf32）上=每个令牌始终为0.015秒，使用简单的pytorch代码进行测试（无CUDA），GPU利用率为45%，VRAM 7823M

GPT2-XL 1.3B在A40（tf32）上=每个令牌为0.032秒（对于ctxlen 1000），使用HF进行测试，GPU利用率也为45%（有趣），VRAM 9655M

训练速度：（新的训练代码）RWKV-4 14B BF16 ctxlen4096 = 在8x8 A100 80G（ZERO2+CP）上每秒处理114K个令牌。（旧的训练代码）RWKV-4 1.5B BF16 ctxlen1024 = 在8xA100 40G上每秒处理106K个令牌。

我也在进行图像实验（例如：链接[17]），并且RWKV将能够进行txt2img扩散 :) 我的想法是：256x256的RGB图像 -> 32x32x13位潜变量 -> 应用RWKV计算每个32x32网格的过渡概率 -> 假设网格是独立的，并使用这些概率进行"扩散"。

平滑训练-没有损失峰值！（lr和bsz在处理150亿个令牌时发生变化）

RWKVloss

RWKVeval

所有训练的模型将开源。即使在CPU上，推理速度也非常快（只有矩阵-向量乘法，没有矩阵-矩阵乘法），因此您甚至可以在手机上运行LLM。

工作原理：RWKV将信息收集到多个通道中，随着移动到下一个令牌，这些通道也以不同的速度衰减。一旦理解，它就非常简单。

RWKV是可并行化的，因为每个通道的时间衰减是数据独立的（且可训练的）。例如，在通常的RNN中，您可以将一个通道的时间衰减从0.8调整为0.5（这些称为“门控”），而在RWKV中，您只需将信息从W-0.8通道移动到W-0.5通道即可实现相同的效果。此外，如果您想要额外的性能，您可以将RWKV微调为不可并行化的RNN（然后可以使用上一个令牌的后几层的输出）。

RWKVformula

以下是我的一些待办事项。让我们一起合作 :)

•HuggingFace集成（请查看链接[18]），以及优化的CPU和iOS和Android和WASM和WebGL推理。RWKV是一个RNN，非常适合边缘设备。让我们使在手机上运行LLM成为可能。•在双向和MLM任务以及图像和音频和视频令牌上进行测试。我认为RWKV可以通过以下方式支持编码器-解码器：对于每个解码器令牌，使用一个学习到的混合物，包括[解码器先前隐藏状态]和[编码器最终隐藏状态]。因此，所有解码器令牌都可以访问编码器的输出。•现在正在使用一个额外的微小注意力（与RWKV-4相比，只需添加几行代码）来训练RWKV-4a，以进一步改善一些困难的零样本任务（如LAMBADA）的较小模型。请参阅链接[19]用户反馈：

到目前为止，我在我们相对较小的预训练数据集（约10GB的文本）上玩弄了基于字符的模型，结果非常好-与训练时间长得多的模型相似的ppl。
天哪，rwkv速度真快。我在从头开始训练它后切换到另一个标签页，当我回来时，它正在产生可信的英文和毛利语词汇，我离开去微波一些咖啡，当我回来时，它正在生成完全符合语法的句子。

Sepp Hochreiter的推文（感谢！）：链接[20]您还可以在EleutherAI Discord中找到我（BlinkDL）：链接[21]

RWKVdemo

快速开始

重要提示：使用 deepspeed==0.7.0 pytorch-lightning==1.9.2 torch 1.13.1+cu117

使用链接[22]（最新代码，与v4兼容）。

这是一个用于测试LLM的问答的绝佳提示。适用于任何模型：（通过最小化ChatGPT的ppl来找到的RWKV 1.5B）

prompt = f'\nQ & A\n\nQuestion:\n{qq}\n\nDetailed Expert Answer:\n' # 在此之后让模型生成

推理

运行 RWKV-4 Pile 模型： 从 https://huggingface.co/BlinkDL 下载模型。在 run.py 中设置 TOKEN_MODE = 'pile'，然后运行它。即使在 CPU 上也可以快速运行（默认模式）。

RWKV-4 Pile 1.5B 的 Colab： https://colab.research.google.com/drive/1F7tZoPZaWJf1fsCmZ5tjw6sYHiFOYVWM

在浏览器中运行 RWKV-4 Pile 模型（以及 ONNX 版本）：参见此问题 https://github.com/BlinkDL/RWKV-LM/issues/7

RWKV-4 Web 演示：https://josephrocca.github.io/rwkv-v4-web/demo/（注意：目前仅支持贪婪采样）[23]

对于旧版 RWKV-2：在此发布中查看 enwik8 上的 27M 参数模型和 0.72 BPC（开发）的运行方法 https://github.com/BlinkDL/RWKV-LM/tree/main/RWKV-v2-RNN。甚至可以在浏览器中运行：https://github.com/BlinkDL/AI-Writer/tree/main/docs/eng[24] https://blinkdl.github.io/AI-Writer/eng/（这是使用[25] tf.js WASM 单线程模式）。

训练/微调

pip install deepspeed==0.7.0 // pip install pytorch-lightning==1.9.5 // torch 1.13.1+cu117

注意：在少量数据上进行训练时，添加权重衰减（0.1 或 0.01）和丢弃（0.1 或 0.01）。尝试 x=x+dropout(att(x))、x=x+dropout(ffn(x))、x=dropout(x+att(x))、x=dropout(x+ffn(x)) 等。

从头开始训练 RWKV-4： 运行 train.py，默认使用 enwik8 数据集（解压缩 https://data.deepai.org/enwik8.zip）。[26]

你将训练 "GPT" 版本，因为它可以并行化，训练速度更快。RWKV-4 能够推断，因此使用 ctxLen 1024 可以适用于 ctxLen 为 2500+ 的情况。你可以使用更长的 ctxLen 对模型进行微调，它能够快速适应更长的 ctxLen。

微调 RWKV-4 Pile 模型： 使用 https://github.com/BlinkDL/RWKV-v2-RNN-Pile/tree/main/RWKV-v3 中的 'prepare-data.py' 将 .txt 文件标记为 train.npy 数据。然后使用 https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v4neo/train.py 进行训练。

阅读 src/model.py 中的推理代码，并尝试使用最终隐藏状态（.xx .aa .bb）作为其他任务的可靠句子嵌入。可能应该从 .xx 和 .aa/.bb（.aa 除以 .bb）开始。

用于 RWKV-4 Pile 模型微调的 Colab：Google Colab[27]

大型语料库： 使用 GitHub - Abel2076/json2binidx_tool[28] 将 .jsonl 转换为 .bin 和 .idx 格式。

以下是 jsonl 格式的示例（每个文档一行）：

{"text": "This is the first document."}
{"text": "Hello\nWorld"}
{"text": "1+1=2\n1+2=3\n2+2=4"}

通过类似以下代码生成：

ss = json.dumps({"text": text}, ensure_ascii=False)
out.write(ss + "\n")

无限 ctxlen 训练（进行中）： https://github.com/Blealtan/RWKV-LM-LoRA/tree/dev-infctx

如何使用 RWKV 隐藏状态作为文本嵌入

以 RWKV 14B 为例。该状态具有 200 个向量，即每个块有 5 个向量：fp16（xx）、fp32（aa）、fp32（bb）、fp32（pp）、fp16（xx）。

不要进行平均池化，因为状态中的不同向量（xx、aa、bb、pp、xx）具有非常不同的含义和范围。你可能可以去除 pp。

建议首先收集每个向量的每个通道的均值+标准差统计信息，并对它们进行归一化（注意：归一化应该是与数据无关的，并且应该从各种文本中收集）。然后训练一个线性分类器。

迈向 RWKV-5（仅用于记录一些新想法）

一些想法

1.现在的时间衰减是类似于 0.999^T（其中 0.999 是可学习的）。将其改为类似于（0.999^T + 0.1），其中 0.1 也是可学习的。0.1 部分将永远保留下来。或者，A^T + B^T + C = 快速衰减 + 慢速衰减 + 常量。甚至可以使用不同的公式（例如，对于衰减分量，使用 K^2 而不是 e^K，或者不进行归一化）。2.在某些通道中使用复数衰减（因此，使用旋转而不是衰减）。3.注入一些可训练且可外推的位置编码？4.除了二维旋转，我们可以尝试其他李群，如三维旋转（SO(3)）。非阿贝尔 RWKV，哈哈。5.RWKV 在模拟设备上可能效果很好（搜索模拟矩阵-向量乘法和光子矩阵-向量乘法）。RNN 模式非常适合硬件（在内存中进行处理）。也可以是 SNN（https://github.com/ridgerchu/SpikeGPT）。我想知道它是否可以针对量子计算进行优化。[29]6.可训练的初始隐藏状态（xx aa bb pp xx）。7.逐层（甚至逐行/逐列、逐元素）学习率，测试 Lion 优化器。

视觉任务

1.我发现添加一个二维位置编码效果很好：

  self.pos_emb_x = nn.Parameter(torch.zeros((1,args.my_pos_emb,args.n_embd)))
self.pos_emb_y = nn.Parameter(torch.zeros((args.my_pos_emb,1,args.n_embd)))
...
x = x + pos_emb_x + pos_emb_y

2. 在 BPE 语言模型中，最好使用 [tokenShift 偏移 1 个标记]（在字符级英文模型中可以混合更多标记）。然而，如果图像尺寸为 N x N，可以尝试使用 [tokenShift 偏移 N（或 N-1）（或 N+1）个标记]，因为这样会将 [当前位置上方的标记（或将要预测的位置上方的标记）] 与 [当前标记] 混合在一起。你可以为 "ATT" 和 "FFN" 尝试不同的 tokenShift 风格，或者混合不同的 tokenShift 风格 - 例如将 [标记 A] 与 [标记 A-1] 和 [标记 A-(N-1)] 等混合在一起。

其他

也许我们可以通过简单地重复上下文来改进记忆能力（我猜重复两次就足够了）。示例：参考 -> 参考（再次）-> 问题 -> 答案

想法：字节感知嵌入

这个想法是确保词汇表中的每个标记都了解其长度和原始的 UTF-8 字节。

令 a = max(len(token))，其中 token 是词汇表中的所有标记。定义 AA：float[a][d_emb]

令 b = max(len_in_utf8_bytes(token))，其中 token 是词汇表中的所有标记。定义 BB：float[b][256][d_emb]

对于词汇表中的每个标记 X，假设 [x0, x1, ..., xn] 是它的原始 UTF-8 字节。我们将在其嵌入 EMB(X) 中添加一些额外的值：

EMB(X) += AA[len(X)] + BB[0][x0] + BB[1][x1] + ... + BB[n][xn]（注意：AA 和 BB 是可学习的权重）

•我们也可以对最终的 Linear(d_emb, n_vocab) 投影进行类似操作。•我们可以使用一些小型网络来生成 AA 和 BB，以提供额外的正则化（例如，BB[m][xi] 和 BB[n][xi] 应该是相关的）。

旧想法

我有一个改进标记化的想法。我们可以硬编码一些通道来表示特定含义。例如：

通道 0 = "空格"

通道 1 = "首字母大写"

通道 2 = "所有字母大写"

因此：

"abc" 的嵌入：[0, 0, 0, x0, x1, x2, ...]

" abc" 的嵌入：[1, 0, 0, x0, x1, x2, ...]

" Abc" 的嵌入：[1, 1, 0, x0, x1, x2, ...]

"ABC" 的嵌入：[0, 0, 1, x0, x1, x2, ...]

......

它们将共享大部分嵌入。我们可以快速计算所有变体的 "abc" 的输出概率。

注意：上述方法假设对于任何 "xyz"，p(" xyz") / p("xyz") 都是相同的，但这可能是错误的。

更好的方法是：将 emb_space、emb_capitalize_first、emb_capitalize_all 定义为 emb 的函数。

也许最好的方法是：让 'abc'、' abc' 等共享它们嵌入的最后 90% 部分。

当前，我们的所有分词器都花费了很多项来表示 'abc'、' abc'、' Abc' 等的所有变体。此外，如果数据集中某些变体很少出现，模型无法发现它们实际上是相似的。这种方法可以改善这一点。我计划在 RWKV 的新版本中测试这个想法。

想法：更好的初始状态

示例（单轮问答）：

1.生成所有维基文档的最终状态。2.对于任何用户的问题 Q，找到最佳的维基文档，并将其最终状态作为初始状态。3.训练一个模型，直接为任何用户的问题 Q 生成最佳的初始状态。

然而，对于多轮问答，这可能会稍微复杂一些 :)

工作原理

RWKV 受到 Apple 的 AFT (https://arxiv.org/abs/2105.14103) 的启发。

此外，它使用了一些我自己的技巧，例如：

•SmallInitEmb: https://github.com/BlinkDL/SmallInitEmb（适用于所有[30] Transformer），它有助于提高嵌入质量，并稳定了后置层归一化（我正在使用该方法）。•Token-shift: https://github.com/BlinkDL/RWKV-LM#token-shift-time-shift-mixing（适用于所有[31] Transformer），对字符级模型尤其有帮助。•Head-QK: https://github.com/BlinkDL/RWKV-LM#the-head-qk-trick-learning-to-copy-and-avoid-tokens（适用于所有[32] Transformer）。注意：它很有用，但我在 Pile 模型中禁用了它，以保持 100% RNN。•FFN 中的额外 R-gate（适用于所有 Transformer）。我还使用了来自 Primer 的 reluSquared。•更好的初始化：我将大多数矩阵初始化为零（参见 https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v2-RNN/src/model.py 中的 RWKV_Init）。•你可以将一些参数从小模型传输到大模型（注意：我也对其进行了排序和平滑处理），以加快和改善收敛速度（参见 https://www.reddit.com/r/MachineLearning/comments/umq908/r_rwkvv2rnn_a_parallelizable_rnn_with/）。[33]•我的 CUDA 内核：https://github.com/BlinkDL/RWKV-CUDA，用于加速训练。[34]

伪代码（从上到下执行）：

RWKVv2RNN

a、b、c、d 因子共同构建了一个时间衰减曲线：[X, 1, W, W^2, W^3, ...]。

写出“位置为 2 的标记”和“位置为 3 的标记”的公式，你就能理解：

•a 和 b：kv 和 k 的 EMA（指数移动平均）。•c 和 d：这些是 a 和 b 与“自注意力”结合的结果。

kv / k 是记忆机制。如果通道中的 W 接近 1，高 k 值的标记可以被长时间记住。

R-gate 对于性能很重要。k 是该标记的信息强度（传递给后续标记），r 表示是否将信息应用于该标记。

RWKV-3 的改进

在 SA（Self-Attention）和 FF（Feed-Forward）层中，使用不同的可训练 TimeMix 因子来处理 R / K / V。示例：

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)

使用 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-3 GPT 模型的代码

GPT 模型 - 概述

RWKV-3 GPT 模型的构建模块与通常的 preLN GPT 相似。

唯一的区别是嵌入之后有一个额外的 LN（层归一化）。请注意，在训练结束后，可以将该 LN 吸收到嵌入中。

x = self.emb(idx)  # input: idx = token indices
x = self.ln_emb(x) # extra LN after embedding
x = x + self.att_0(self.ln_att_0(x)) # preLN
x = x + self.ffn_0(self.ln_ffn_0(x))
...
x = x + self.att_n(self.ln_att_n(x))
x = x + self.ffn_n(self.ln_ffn_n(x))
x = self.ln_head(x) # final LN before projection
x = self.head(x)    # output: x = logits

重要的是将 emb 初始化为极小的值，例如 nn.init.uniform_(a=-1e-4, b=1e-4)，以利用我的技巧 https://github.com/BlinkDL/SmallInitEmb。[35]

对于 1.5B 的 RWKV-3，我在 8 个 A100 40G 上使用 Adam 优化器（无权重衰减，无 dropout）。

batchSz = 32 * 896，ctxLen = 896。我使用的是 tf32，因此批量大小稍小。

对于前 150B 个标记，学习率（LR）固定为 3e-4，beta=(0.9, 0.99)。

然后，我将 beta 设置为 (0.9, 0.999)，并对 LR 进行指数衰减，达到 332B 标记时的学习率为 1e-5。

在 GPT 模式下的 RWKV-3 中，通常意义上的注意力是不存在的，但我们仍然将这个块称为 ATT（注意力）。

B, T, C = x.size() # x = (Batch, Time, Channel)

# 使用前一个时间步骤的 x 混合生成 xk、xv、xr
xx = self.time_shift(x) # self.time_shift = nn.ZeroPad2d((0,0,1,-1))
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)

# 使用 xk、xv、xr 生成 k、v、r
k = self.key(xk).transpose(-1, -2)
v = self.value(xv).transpose(-1, -2)
r = self.receptance(xr)
k = torch.clamp(k, max=60) # 将 k 截断以避免溢出
k = torch.exp(k)
kv = k * v

# 计算 W-curve = [e^(-n * e^time_decay), e^(-(n-1) * e^time_decay), ..., 1, e^(time_first)]
self.time_w = torch.cat([torch.exp(self.time_decay) * self.time_curve.to(x.device), self.time_first], dim=-1)
w = torch.exp(self.time_w)

# 使用 W 分别对 kv 和 k 进行混合。添加 K_EPS 到 wk 以避免除以零
if RUN_DEVICE == 'cuda':
    wkv = TimeX.apply(w, kv, B,C,T, 0)
    wk = TimeX.apply(w, k, B,C,T, K_EPS)
else:
    w = w[:,-T:].unsqueeze(1)
    wkv = F.conv1d(nn.ZeroPad2d((T-1, 0, 0, 0))(kv), w, groups=C)
    wk = F.conv1d(nn.ZeroPad2d((T-1, 0, 0, 0))(k), w, groups=C) + K_EPS

# RWKV 公式
rwkv = torch.sigmoid(r) * (wkv / wk).transpose(-1, -2)
rwkv = self.output(rwkv) # 最终输出的投影

self.key、self.receptance、self.output 矩阵都被初始化为零。

time_mix、time_decay、time_first 向量是从一个较小的训练模型中传输过来的（注意：我也对其进行了排序和平滑处理）。

GPT 模式下的 FFN 块

与通常的 GPT 相比，FFN 块具有三个技巧：

1.我的 time_mix 技巧。2.来自 Primer 论文的 sqReLU。3.额外的 receptance-gate（类似于 ATT 块中的 receptance-gate）。

# 使用前一个时间步骤的 x 混合生成 xk、xr
xx = self.time_shift(x)
xk = x * self.time_mix_k + xx * (1 - self.time_mix_k)
xr = x * self.time_mix_r + xx * (1 - self.time_mix_r)

# 常规的 FFN 操作
k = self.key(xk)
k = torch.square(torch.relu(k)) # 来自 Primer 论文
kv = self.value(k)

# 对 kv 应用额外的 receptance-gate
rkv = torch.sigmoid(self.receptance(xr)) * kv
return rkv

self.value、self.receptance 矩阵都被初始化为零。

RWKV-4 的改进

RWKVv3plan

从GPT到RWKV（公式）

设F[t]为时刻t的系统状态。

设x[t]为时刻t的新外部输入。

在GPT中，预测F[t+1]需要考虑F[0]，F[1]，...，F[t]。因此，生成长度为T的序列的时间复杂度为O(T^2)。

GPT的简化公式为：

从理论上讲，这非常强大，然而这并不意味着我们可以通过常规的优化器充分利用其能力。我怀疑损失函数的优化空间对于我们当前的方法来说过于困难。

与RWKV（并行模式，类似于Apple的AFT）的简化公式进行比较：

其中，R、K、V为可训练矩阵，W为可训练向量（每个通道的时间衰减因子）。

在GPT中，F[i]对F[t+1]的贡献由

加权。

在RWKV-2中，F[i]对F[t+1]的贡献由

加权。•这里的𝛔是一个非线性函数，我们可以使用sigmoid函数。

•注意𝛔(Rx[t])不在分母中，我将R称为"receptance"（响应度）。

• exp(W*(t-i))是时间衰减因子。我在2020年8月提出了相同的思路（通过距离对注意力进行缩放），我将其称为"时间加权"（time-weighting）（请查看https://github.com/BlinkDL/minGPT-tuned的提交历史记录）。[36]

这些公式描述了GPT到RWKV的转变，其中RWKV使用不同的加权方案来生成系统在每个时刻的状态。这些公式提供了更高效的计算方式，并且可以通过训练可学习的矩阵和向量来调整权重，以适应不同的任务和数据。

以下是重要结论：我们可以将其重写为一个RNN（递归公式）。请注意：

其中，A[t]和B[t]分别是上一步的分子和分母。

我相信RWKV的性能良好，因为W类似于重复应用对角矩阵。请注意(P^{-1} D P)^n = P^{-1} D^n P，因此它类似于重复应用一般的可对角化矩阵。

此外，我们还可以将其转化为连续的ODE（类似于状态空间模型）。稍后我会写一篇关于这个的文章。

多模态的想法

我有一个关于[text --> 32x32 RGB image]的想法，使用LM（transformer，RWKV等）将进行测试。

首先，使用LM损失（而不是L2损失），这样图像就不会模糊。

其次，颜色量化。例如，只允许R/G/B有8个级别。然后，图像词汇表的大小为8x8x8 = 512（每个像素），而不是2^24。因此，一个32x32的RGB图像可以表示为一个长度为1024的词汇表大小为512的序列（图像标记），这对于通常的LM来说是一个典型的输入。（以后我们可以使用扩散模型进行上采样和生成RGB888图像。我们可能也可以使用LM来做这个。）

第三，使用二维位置嵌入，使模型容易理解。例如，在前64个（=32+32）通道中添加X和Y坐标的独热编码。例如，如果像素位于x=8，y=20，则在通道8和通道52（=32+20）中添加1。此外，我们可能还可以将浮点型X和Y坐标（归一化为0~1范围）添加到另外2个通道中。其他周期性的位置编码可能也有帮助（将进行测试）。

最后，在DataLoader中进行颜色量化时使用RandRound。例如，如果浮点级别为4.578，则有57.8％的机会使用5，（1-57.8％）的机会使用4。我们可以允许预测结果为4和5，但是如果预测结果是4，则损失将更高。

多任务训练也可能有帮助。我将尝试以下数据集格式：[TxtFirst] [Img描述（文本标记）] [Img] [图像标记] 有时也可以是 [ImgFirst] [图像标记] [Txt] [Img描述（文本标记）] ...在DataLoader中，图片的顺序应该是随机的，[TxtFirst] [ImgFirst] [Img] [Txt]是特殊标记。对完整数据集进行随机采样。因此，模型有时将先看到图像标记，然后是相应的文本标记，这是一个[img -> txt]任务。模型将看到一些部分图像和部分文本。我认为字符级LM可能有助于模型在图像上编写正确的文本。

如何对大型数据集进行采样（用于训练）

我使用了一个技巧来确定性地但足够随机地对Pile进行采样。

假设pile有x个块（一个块=ctx_len个标记）。

选择一个小于x的素数p，并确保p = 2（mod 3）。

使用(step * step * step) mod p进行采样。为step添加一些偏差以增加额外的随机性。

top-p-x采样方法（用于推断）

我们提出了一种名为top-p-x的新采样方法：

它类似于top-p，唯一的区别是您还保留所有概率大于x的标记。

首先尝试x = 0.01。

通过损失曲线的变分方法获得更好的学习率调度

我提出了一种简单的新方法来找到更好的学习率调度。该方法对于大型语言模型来说成本高效且实用。关键是我们可以对损失曲线动态（现象学）建模，从中直接使用变分方法计算出良好的学习率曲线。此外，我们可以以合理的准确性预测最终的损失。

更新：在"结论1"中，使用最适合的区间（忽略我们的近似失效的初始步骤）来拟合参数。

尝试以下方法：固定的学习率1小时，然后指数衰减到0.2 * lr，持续12小时，并选择t = [1小时，13小时]段。

在最后三个图中，黑色=新学习率调度的预测损失曲线，蓝色=原始（未优化）真实损失曲线，橙色=新学习率调度。

betterlrschedule

RWKV v1

我们提出了RWKV语言模型，其中包含交替的时间混合（Time-mix）和通道混合（Channel-mix）层：

•时间混合（Time-mix）:

TM_{t,c} = sigmoid(R_{t,c}) * sum_{u} (W_{t,u,c} * softmax_t(K_{u,c}) * V_{u,c})

•通道混合（Channel-mix）:

CM_{t,c} = sigmoid(R_{t,c}) * sum_d (W_{c,d} * gelu(K_{t,d}) * V_{t,d})

其中，R、K、V是输入的线性变换生成的，而W是参数。RWKV的思想是将注意力分解为R（目标）* W（源，目标）* K（源）。因此，我们可以将R称为"接收度"（receptance），而sigmoid函数将其限制在0~1范围内。

时间混合与AFT（https://arxiv.org/abs/2105.14103）相似，但有两个区别：[37]

1.我们更改了归一化（分母）的定义。对于遮蔽语言模型，我们定义：

softmax_t(K_{u,c}) = exp(K_{u,c}) / sum_{v <= t} exp(K_{v,c})

(更新：在v2中，我们使用了AFT的原始归一化方法)

2.我们将W_{t,u,c}分解，并引入多头W（其中h是c对应的头）：

W_{t,u,c} = f_h(t-u) * α_h(u) * β_h(t)

此外，我们通过γ(t)乘以时间混合层的最终输出。α、β、γ因子的原因是当t较小时，上下文大小较小，可以使用α、β、γ因子进行补偿。

为了快速稳定地收敛，初始化K和R矩阵（以及输出投影矩阵）为零。

注意：上述公式中的softmax、sigmoid、gelu函数分别表示softmax函数、sigmoid函数和gelu函数。

RWKV v2

在 RWKV v2 中，我们对架构进行了一些修改。以下是更新的内容：

•在 v2-RNN 中，我们删除了 α、β、γ 因子，并将 W 限制为简单形式，使其可以被重写为 RNN。•Channel-mix 层类似于 GeGLU（https://arxiv.org/abs/2002.05202），但添加了额外的[38] R 因子。初始化 R 和 W 矩阵为零，以实现快速且稳定的收敛。•最后，我们添加了额外的 token-shift 或 time-shift 混合，与 (https://github.com/BlinkDL/minGPT-tuned) 中的方法类似。

Token-shift（时间偏移混合）

Token-shift 明确使用当前令牌的一半通道和上一个令牌的一半通道来生成所有向量（QKV、RWKV 等）。

self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))

x = torch.cat([self.time_shift(x[:, :, :C//2]), x[:, :, C//2:]], dim=-1)

将通道数除以 2 并将其向左偏移 1 对于字符级英语和字符级中文的语言模型效果很好。

然而，对于基于 BPE 的英语语言模型而言，只有在嵌入大小足够大时才有效（至少为 1024）。通常的小型 L12-D768 模型可能不足够。

关于 token-shift 的有效性，我有以下理论：

在训练 GPT 时，令牌的隐藏表示必须实现两个不同的目标：

1.预测下一个令牌。有时这很容易（下一个令牌很明显）。2.收集所有先前的上下文信息，以供后续的令牌使用。这始终是一项困难的任务。

信息传播与令牌偏移

通过通道的偏移，我们可以更好地传播信息，使其更专注于（2）这个目标。这就像一种残差连接，或者在变换器中嵌入了一个小型 RNN。

你也可以在常规的 QKV 自注意力中使用令牌偏移。我查看了权重，并发现 V 很喜欢偏移后的通道，而 Q 则不太喜欢。如果仔细思考，这是有道理的。我还发现你可能希望在较高层次中使用较少的混合。

附注：这个仓库中有一个名为 MHA_pro 的模型，具有强大的性能。可以试试看：）

头部 QK 技巧：学习复制和避免令牌

在通常的 Transformer 中，小型模型在复制上下文中的令牌（例如人名）时会遇到困难。我们在最终输出中添加了额外的 Q 和 K，使模型能够直接复制（或避免）上下文中的令牌。在观察学习到的权重时，模型将自动学习命名实体识别（NER）。

q = self.head_q(x)[:,:T,:] # 投影到 256 维
k = self.head_k(x)[:,:T,:] # 投影到 256 维
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

注意：当一个令牌在上下文中出现多次时，使用 max(prob) 而不是 sum(prob) 可能会更好。

top-a 抽样方法

我们还提出了一种称为 top-a 的新的抽样方法（如 src/utils.py 中所示）：

（1）找到 softmax 之后的最大概率 p_max。

（2）删除所有概率低于 0.2 * pow(p_max, 2) 的条目。因此它是自适应的，因此称为 "top-a"。

（3）可以根据需要调整 0.2 和 2 的因子。首先调整 0.2。

top-a 的思想：

1.如果 max_prob=0.9，则删除概率小于 0.162 的所有令牌（因此删除所有备选项）。2.如果 max_prob=0.5，则删除概率小于 0.05 的所有令牌（因此允许更多选择）。3.如果 max_prob=0.1，则删除概率小于 0.002 的所有令牌（因此允许许多可能性）。

probs = F.softmax(logits, dim=-1)

limit = torch.pow(torch.max(probs), 2) * 0.02
logits[probs < limit] = -float('Inf')

性能

基于 simplebooks-92 数据集的字符级损失情况：simplebooks.zip[39]

RWKVvsMHA

灰色：常规的 MHA+Rotary+GeGLU - 性能不太好。参数数量为 17.2M。

红色：RWKV（"linear" attention）- 对 VRAM 友好 - 当上下文窗口较长时速度较快 - 性能良好。参数数量为 16.6M。

绿色：MHA+Rotary+GeGLU+Token_shift。参数数量为 17.2M。

蓝色：MHA_pro（带有各种改进和 RWKV 类型 FFN 的 MHA）- 速度较慢 - 需要更多 VRAM - 性能良好。参数数量为 16.6M。

@software{peng_bo_2021_5196578,
  author       = {PENG Bo},
  title        = {BlinkDL/RWKV-LM: 0.01},
  month        = aug,
  year         = 2021,
  publisher    = {Zenodo},
  version      = {0.01},
  doi          = {10.5281/zenodo.5196577},
  url          = {https://doi.org/10.5281/zenodo.5196577}
}

初始化

为了实现快速收敛，我们对 RWKV 进行了谨慎的初始化 - 使用正交矩阵进行适当缩放，并使用特殊的 time_w 曲线。详细信息请查看 model.py。

一些学习到的 time_w 示例：

RWKVtimew

中文模型

中文使用教程：https://zhuanlan.zhihu.com/p/618011122 https://zhuanlan.zhihu.com/p/616351661

推荐UI：GitHub - wenda-LLM/wenda: 闻达：一个LLM调用平台。目标为针对特定环境的高效内容生成，同时考虑个人和中小企业的计算资源局限性，以及知识安全和私密性问题[40]

声明

本文翻译整理自：

•https://github.com/BlinkDL/RWKV-LM•GitHub - BlinkDL/ChatRWKV: ChatRWKV is like ChatGPT but powered by RWKV (100% RNN) language model, and open source.[41]

感兴趣的同学请点赞、收藏。

References

[1] Raven RWKV 7B - a Hugging Face Space by BlinkDL: https://huggingface.co/spaces/BlinkDL/RWKV-World-7B [2] 链接: https://pypi.org/project/rwkv/ [3] saharNooby/rwkv.cpp: https://github.com/saharNooby/rwkv.cpp [4] ggml: https://github.com/ggerganov/ggml [5] harrisonvanderbyl/rwkv-cpp-cuda: https://github.com/harrisonvanderbyl/rwkv-cpp-cuda [6] Blealtan/RWKV-LM-LoRA: https://github.com/Blealtan/RWKV-LM-LoRA [7] josStorer/RWKV-Runner: https://github.com/josStorer/RWKV-Runner [8] 链接: https://github.com/search?o=desc&q=rwkv&s=updated&type=Repositories [9] 链接: https://discord.gg/bDSBUMeFpc [10] 链接: https://twitter.com/BlinkDL_AI [11] 链接: https://github.com/BlinkDL/ChatRWKV/blob/main/RWKV_in_150_lines.py [12] 链接: https://arxiv.org/abs/2305.13048 [13] 链接: https://johanwind.github.io/2023/03/23/rwkv_overview.html [14] 链接: https://johanwind.github.io/2023/03/23/rwkv_details.html [15] 链接: https://github.com/ridgerchu/SpikeGPT [16] 链接: https://discord.gg/bDSBUMeFpc [17] 链接: https://huggingface.co/BlinkDL/clip-guided-binary-autoencoder [18] 链接: https://github.com/huggingface/transformers/issues/17230 [19] 链接: https://github.com/BlinkDL/RWKV-LM/commit/a268cd2e40351ee31c30c5f8a5d1266d35b41829 [20] 链接: https://twitter.com/HochreiterSepp/status/1524270961314484227 [21] 链接: https://www.eleuther.ai/get-involved/ [22] 链接: https://github.com/BlinkDL/RWKV-LM/tree/main/RWKV-v4neo [23] https://josephrocca.github.io/rwkv-v4-web/demo/（注意：目前仅支持贪婪采样）: https://josephrocca.github.io/rwkv-v4-web/demo/%EF%BC%88%E6%B3%A8%E6%84%8F%EF%BC%9A%E7%9B%AE%E5%89%8D%E4%BB%85%E6%94%AF%E6%8C%81%E8%B4%AA%E5%A9%AA%E9%87%87%E6%A0%B7%EF%BC%89 [24] https://github.com/BlinkDL/RWKV-LM/tree/main/RWKV-v2-RNN。甚至可以在浏览器中运行：https://github.com/BlinkDL/AI-Writer/tree/main/docs/eng: https://github.com/BlinkDL/RWKV-LM/tree/main/RWKV-v2-RNN%E3%80%82%E7%94%9A%E8%87%B3%E5%8F%AF%E4%BB%A5%E5%9C%A8%E6%B5%8F%E8%A7%88%E5%99%A8%E4%B8%AD%E8%BF%90%E8%A1%8C%EF%BC%9Ahttps://github.com/BlinkDL/AI-Writer/tree/main/docs/eng [25] https://blinkdl.github.io/AI-Writer/eng/（这是使用: https://blinkdl.github.io/AI-Writer/eng/%EF%BC%88%E8%BF%99%E6%98%AF%E4%BD%BF%E7%94%A8 [26] https://data.deepai.org/enwik8.zip）。: https://data.deepai.org/enwik8.zip%EF%BC%89%E3%80%82 [27] Google Colab: https://colab.research.google.com/github/resloved/RWKV-notebooks/blob/master/RWKV_v4_RNN_Pile_Fine_Tuning.ipynb [28] GitHub - Abel2076/json2binidx_tool: https://github.com/Abel2076/json2binidx_tool [29] https://github.com/ridgerchu/SpikeGPT）。我想知道它是否可以针对量子计算进行优化。: https://github.com/ridgerchu/SpikeGPT%EF%BC%89%E3%80%82%E6%88%91%E6%83%B3%E7%9F%A5%E9%81%93%E5%AE%83%E6%98%AF%E5%90%A6%E5%8F%AF%E4%BB%A5%E9%92%88%E5%AF%B9%E9%87%8F%E5%AD%90%E8%AE%A1%E7%AE%97%E8%BF%9B%E8%A1%8C%E4%BC%98%E5%8C%96%E3%80%82 [30] https://github.com/BlinkDL/SmallInitEmb（适用于所有: https://github.com/BlinkDL/SmallInitEmb%EF%BC%88%E9%80%82%E7%94%A8%E4%BA%8E%E6%89%80%E6%9C%89 [31] https://github.com/BlinkDL/RWKV-LM#token-shift-time-shift-mixing（适用于所有: https://github.com/BlinkDL/RWKV-LM#token-shift-time-shift-mixing%EF%BC%88%E9%80%82%E7%94%A8%E4%BA%8E%E6%89%80%E6%9C%89 [32] https://github.com/BlinkDL/RWKV-LM#the-head-qk-trick-learning-to-copy-and-avoid-tokens（适用于所有: https://github.com/BlinkDL/RWKV-LM#the-head-qk-trick-learning-to-copy-and-avoid-tokens%EF%BC%88%E9%80%82%E7%94%A8%E4%BA%8E%E6%89%80%E6%9C%89 [33] https://www.reddit.com/r/MachineLearning/comments/umq908/r_rwkvv2rnn_a_parallelizable_rnn_with/）。: https://www.reddit.com/r/MachineLearning/comments/umq908/r_rwkvv2rnn_a_parallelizable_rnn_with/%EF%BC%89%E3%80%82 [34] https://github.com/BlinkDL/RWKV-CUDA，用于加速训练。: https://github.com/BlinkDL/RWKV-CUDA%EF%BC%8C%E7%94%A8%E4%BA%8E%E5%8A%A0%E9%80%9F%E8%AE%AD%E7%BB%83%E3%80%82 [35] https://github.com/BlinkDL/SmallInitEmb。: https://github.com/BlinkDL/SmallInitEmb%E3%80%82 [36] https://github.com/BlinkDL/minGPT-tuned的提交历史记录）。: https://github.com/BlinkDL/minGPT-tuned%E7%9A%84%E6%8F%90%E4%BA%A4%E5%8E%86%E5%8F%B2%E8%AE%B0%E5%BD%95%EF%BC%89%E3%80%82 [37] https://arxiv.org/abs/2105.14103）相似，但有两个区别：: https://arxiv.org/abs/2105.14103%EF%BC%89%E7%9B%B8%E4%BC%BC%EF%BC%8C%E4%BD%86%E6%9C%89%E4%B8%A4%E4%B8%AA%E5%8C%BA%E5%88%AB%EF%BC%9A [38] https://arxiv.org/abs/2002.05202），但添加了额外的: https://arxiv.org/abs/2002.05202%EF%BC%89%EF%BC%8C%E4%BD%86%E6%B7%BB%E5%8A%A0%E4%BA%86%E9%A2%9D%E5%A4%96%E7%9A%84 [39] simplebooks.zip: https://dldata-public.s3.us-east-2.amazonaws.com/simplebooks.zip [40] GitHub - wenda-LLM/wenda: 闻达：一个LLM调用平台。目标为针对特定环境的高效内容生成，同时考虑个人和中小企业的计算资源局限性，以及知识安全和私密性问题: https://github.com/l15y/wenda [41] GitHub - BlinkDL/ChatRWKV: ChatRWKV is like ChatGPT but powered by RWKV (100% RNN) language model, and open source.: https://github.com/BlinkDL/ChatRWKV

本文参与腾讯云自媒体分享计划，分享自微信公众号。

原始发表：2023-07-13，如有侵权请联系 cloudcommunity@tencent.com 删除

https

本文分享自山行AI 微信公众号，前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体分享计划，欢迎热爱写作的你一起参与！