【NLP】RWKV框架入门

Deeplearning 评论

全文从以下几个方面展开：

RWKV是什么

RWKV架构是一种针对分布式存储系统的架构设计，主要关注于提高数据的可靠性、可用性和一致性，同时优化性能和扩展性。

RWKV代表的是Read-Write-Key-Value，这个架构通过分离读（R）和写（W）操作，以及将键（Key）和值（Value）分离，来实现其设计目标。

架构组成

  1. 读写分离：RWKV架构中，读操作 和 写操作 是分开处理的。这意味着系统可以针对读操作进行优化，比如通过缓存机制来提高读取速度，而写操作则可以专注于数据的持久化和一致性保证。
  2. 键值存储：在这种架构中，键（Key）和值（Value）是分开存储的。键通常包含了定位数据所需的全部信息，而值则是实际的数据内容。这种分离使得系统可以更加灵活地管理数据，比如可以实现键的索引，从而加快查找速度。
  3. 数据分区：为了提高系统的扩展性，RWKV架构通常会将数据进行分区。每个分区可以独立地进行读写操作，这样可以在不影响整个系统的情况下，对单个分区进行扩展或维护。

特殊性分析

  1. 高性能：通过读写分离，RWKV架构能够针对不同的操作类型进行专门的优化。例如，读操作可以通过增加缓存层来提高响应速度，而写操作则可以通过批量处理和日志优化来提高吞吐量。
  2. 高可用性：RWKV架构通过数据分区和副本机制来保证数据的可用性。即使某个分区发生故障，系统也可以通过其他分区或副本来继续提供服务。
  3. 一致性保证：在分布式系统中，数据一致性是一个重要的问题。RWKV架构通过写操作的严格控制和一致性协议，比如两阶段提交（2PC）或Paxos、Raft等算法，来确保数据在多个副本之间的一致性。
  4. 扩展性：由于数据分区的设计，RWKV架构可以很好地支持水平扩展。当系统需要处理更多的数据或请求时，可以通过增加更多的分区来实现。
  5. 灵活性：键值分离的设计使得RWKV架构在数据管理上具有很高的灵活性。可以根据业务需求，对键进行定制化的索引和查询，从而更好地支持复杂的应用场景。

为什么选型RWKV做大语言基底模型的预研

因为：便宜好用（这里的便宜是指算力的便宜）

到底有多便宜呢？

上图是RWKV架构模型对比Transformer类架构模型在完成推理任务时的时间复杂度和空间复杂度对比。

通过复杂度可以看到，RWKV架构模型推理的耗时和内存占用随Tokens(可以简单理解为字数)长度呈线性增长关系。而Transformer类模型则呈幂级数增长关系，如果表现在游戏对话任务中，就是和NPC说的话越多，它回复你的时间就会越久，这也很贴合Transformer模型中的重要结构Attention机制的中文翻译——“注意力”。

前些日子看到有文章说华为内部开始研究RWKV架构去服务于他们的大语言模型，于是去其官网上一探究竟，发现 stability.ai（开源AI绘画架构的开发公司）也对RWKV进行了赞助，说明RWKV在架构一定有它的特殊性和优越性，于是决定研究一下。

早在2023年5月RWKV架构问世就已经，在一众Transformer模型架构的LLM领域下，这是一个使用RNN架构进行模型推理的模型。并且在竞争激烈的大语言模型评测中排到了第八名的成绩。

长话短说，如果使用RWKV，最希望它能够解决的问题是：费用问题

在某上线项目的用户信息采集中，我们发现：玩家其实并不那么care游戏中的AI智不智能(更多的是来看看这个AI是不是一个噱头)，也并没有对于它不能回答自己的奇怪提问而提出不满（注意：这里得到的结果并不具有泛化性，有许多影响这一结果的因素没有考虑，包括但不限于：游戏玩法和用户年龄段画像等）。那么在学术界上的细扣模型智能面板上提升零点几个百分点的性能表现大战就几乎不需要考虑了，与其牺牲更多的开发成本和模型加载成本得到一个智力刚及格的AI，还要被玩家拷打“爸爸的爸爸叫什么~”，不如创造一个弱智吧吧友级别的AI更有节目效果。

RWKV和Transformer的差别

RWKV相较于Transformer：

优点（便宜）

  • 在训练和推理中更少的资源占用（VRAM，CPU， GPU，etc）
  • 对比长上下文的Transformer架构大语言模型，计算要求降低了10到100倍
  • 可以线性缩放到任意长度的上下文（在时间复杂度上Transformer架构的时间复杂度是O(n) ）
  • RWKV通常在这几类语言（例如中文、日语等）上训练的更好，然后才是其他开源的LLM模型

缺点（对提示词的要求高）

  • 对提示词的要求更高，需要用户提出详细描述的问题
  • 在需要回顾的任务上表现较弱，因此需要相应地重新排序提示词的顺序
    • （正例：“阅读下面的文档，帮我执行XX任务。文档为……” √）
    • （反例(回顾任务)：“对于上面一段对话中的内容，帮我执行XX任务。” ×）

以上优缺点对比来自RWKV官网：https://wiki.rwkv.com/#tldr-vs-existing-transformer-models

从缺点上看，RWKV表现出非常经典的“RNN特色”，即：长序列任务训练导致的梯度消失和梯度爆炸问题（根源于反向传播过程中的梯度计算问题）。但是其实没有特别想清楚为什么训练阶段的问题会推广至推理阶段。

看到这里，方便给对于AI领域的专业名词不那么了解的同学讲清楚这个性能侧重问题，我假设衡量大语言模型的性能有以下四个维度：适用领域——多、推理速度——快、表现效果——好、内存占用——省，可绘制的表格如下：

适用领域 推理速度 表现效果 内存占用
Transformer 多（Attention is ALL you need） 慢(O(n2)) 好（或许应该听过Sora是Transformer架构结合Diffusion实现的） 慢(O(n2))
RWKV 较多(RWKV-Runner中可以看到AI绘画和音乐生成等工作) 快(O(n)) 一般（源自于自身结构的脆弱性，在长序列下的表现有时像一个健忘的老人，类似于下面这种表现：input:马冬梅 → output:马冬什么/马什么梅/什么冬梅） 快(O(n))

这样一个简单的性能矩阵关系，我们应该可以大概定义Transformer和RWKV的应用场景：

  • Transformer侧重于整体性能的表现效果，如果在高并发的场景下使用将非常非常吃服务器性能和配置。如果使用Transformer作为底层模型架构启动AI相关服务，要么提升显卡算力性能，要么对模型进行剪枝蒸馏等一些列模型压缩工作。
  • RWKV侧重高效和满足基准线，“量大管饱”，在追求“都用上”的场景下，轻量级的高效模型可以很好的弥补Transformer在推理并发上的缺陷，更适合ToC的场景，制作一个弱智吧吧友应该绰绰有余（理论上）。

结构建模

Self-Attention

一个简单的自注意力机制的实现。

其中涉及到以下几个步骤：

  1. 初始化参数：定义需要的查询（Q）、键（K）和值（V）矩阵。
  2. 计算注意力权重：通过查询和键的点积来计算注意力得分，使用Softmax激活函数使得分数归一化。
  3. 计算输出：使用得到的注意力权重对值进行加权求和。

假设输入是一个小批量的嵌入向量，每个向量维度一致。

下面仅用 Python 语言和 NumPy 库来实现一个基本的自注意力结构。

首先，导入必要的库，其次定义一些参数和初始化权重。接下来是关键的一步，创建一个自注意力层，其中输入、查询、键、值的维度都是相同的。

假设批量大小为一个简单的例子，并处理几个向量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# 导入库函数 NumPy
import numpy as np

# 利用 NumPy 来写 Softmax 激活函数(也可以使用torch中现成的)
def softmax(x):
    e_x = np.exp(x - np.max(x, axis = -1, keepdims = True)) # Numerical stability
    return e_x / e_x.sum(axis = -1, keepdims = True)

class SelfAttention:
    # 权重随机初始化
    def __init__(self, embed_size):
        self.embed_size = embed_size
        # 初始化权重矩阵。简单起见，我们使用**随机初始化**
        self.query_weights = np.random.rand(embed_size, embed_size)
        self.key_weights = np.random.rand(embed_size, embed_size)
        self.value_weights = np.random.rand(embed_size, embed_size)
    
    # 前向传播(Forward Propagation)
    def forward(self, x):
        # x 的形状是(batch_size, seq_length, embed_size)
        batch_size, seq_length, embed_size = x.shape
        
        # 计算Q,K,V
        queries = np.dot(x, self.query_weights)
        keys = np.dot(x, self.key_weights)
        values = np.dot(x, self.valeu_weights)
        
        # 计算**注意力得分**
        attention_scores = np.dot(queries, key.transpose((0, 2, 1))) / np.sqrt(self.embed_size)
        attention_probs = softmax(attention_scores) # 将得分结果进行归一化
        
        # 加权求和得到输出
        output = np.dot(attention_probs, values)
        return output

# 示例使用
embed_size = 4 # 嵌入维度是4
seq_length = 3 # 序列长度是3
batch_size = 1 # 批大小是1

# 创建一个自注意力对象
self_attention = SelfAttention(embed_size)

# 创建一个随机的输入批量
x = np.random.rand(batch_size, seq_length, embed_size)

# 运行前向传播
output = self_attention.forward(x)
print("Output of self-attention:\n", output)

Multi-Head Attention

多头注意力结构的目的：显著提高模型的表达能力和注意力的细致度。

多头注意力机制，顾名思义就是多个头的注意力机制，是自注意力机制的扩展。

多头注意力允许模型在不同的表示子控件中并行地捕获信息。

接下来，我们就在上面已经实现的自注意力代码基础上进一步实现这个多头注意力机制。

实现多头注意力机制的几个步骤：

  1. 初始化参数：为每个头粉笔初始化查询（Q）、键（K）和值（V）矩阵。
  2. 分割输入：将输入向量分割成多个头（Head）。
  3. 并行计算：对每个头执行自注意力计算。
  4. 合并头：将所有头的输出合并回一个单一的输出。
  5. 输出层：可能需要一个额外的线性层（Linear）来整合多头计算的结果。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
# 导入库
import numpy as np

# 手写Softmax
def softmax(x):
    e_x = np.exp(x - np.max(x, axis = -1, keepdims = True)) # Numerical stability
    return e_x / e_x.sum(axis = -1, keepdims = Ture)

class MultiHeadAttention:
    # 多头的初始化
    def __init__(self, embed_size, num_heads):
        self.embed_size = embed_size
        self.num_heads = num_heads
        self.head_dim = embed_size / num_heads
        
        assert self.head_dim * num_heads = embed_size, "embed_size nust be devisible by num_heads"
        
        # 初始化权重矩阵，每个头都有一组
        self.query_weights = np.random.rand(num_heads, self.head_dim, self.head_dim)
        self.key_weights = np.random.rand(num_heads, self.head_dim, self.head_dim)
        self.value_weights = np.random.rand(num_heads, self.head_dim, self.head_dim)
        self.out_weights = np.random.rand(num_heads * self.head_dim, embed_size)
        
    def forward(self, x):
        batch_size, seq_length, embed_size = x.shape
        
        # 分割输入
        x = x.reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        
        # 分别对每个头计算Q、K、V
        queries = np.dot(x, self.query_weights.transpose((1, 0, 2)))
        keys = np.dot(x, self.key_weights.transpose((1, 0, 2)))
        values = np.dot(x, self.value_weights.transpose((1, 0, 2)))
        
        # 计算注意力得分
        attention_scores = np.einsum('bnhd, bmhd -> bhnm', queries, keys) / np.sqrt(self.head_dim)
        attention_probs = softmax(attention_scores)
        
        # 加权求和得到每个头的输出
        out = np.einsum('bhnm, bmhd -> bnhd', attention_probs, values)
        
        # 合并所有头的输出
        out = out.reshape(batch_size, seq_length, self.num_heads * self.head_dim)
        
        # 应用输出层权重
        out = np.dot(out, self.out_weights)
        return out
    
# 示例使用
embed_size = 8 # 嵌入维度
seq_length = 3 # 序列长度
batch_size = 1 # 批量大小
num_heads = 2  # 头的数量

# 创建一个多头注意力对象
multi_head_attention = MultiHeadAttention(embed_size, num_heads)

# 创建一个随机的输入批量
x = np.random.rand(batch_size, seq_length, embed_size)

# 运行前向传播
output = multi_head_attention.forward(x)
print("Output of multi-head attention:\n", output)

通过分割输入并在不同的头上并行运算，然后合并结果，它允许模型从多个子空间同时学习信息。

每个头独立地学习输入的不同方面，最后一个输出层将所有头的信息合并起来，形成最终的输出。

RWKV结构

下面是用于实现RWKV的类。

组成它的部分是初始化、层归一化、通道融合、时序融合、前向传播。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
class RWKV_RNN(torch.jit.ScriptModule):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.eval() # 切换模型至评估模式。运行模型就不会使用Dropout和BatchNormalization
        
        w = torch.load(args.MODEL_NAME + '.pth', map_location = 'cpu')
        for k in w.keys():
            if '.time_' in k:
                w[k] = w[k].squeeze()
            if '.time_decay' in k:
                w[k] = -torch.exp(w[k].float)
            else:
                w[k] = w[k].float
  
        self.w = tpyes.SimpleNamespace()
        self.w.blocks = {}
        for k in w.keys():
            parts =k.split('.')
            last = parts.pop()
            here = welf.w
            for p in parts:
                if p.isdigit():
                    p = int(p)
                    if p not in here:
                        here[p] = types.SimpleNamespace()
                        here = here[p]
                    else:
                        if not hasattr(here, p):
                            setattr(here, p, types.SimpleNamespace())
                            here = getattr(here, p)
                setattr(here, last, w[k])
                
        def layer_norm(self, x, w):
            return F.layer_norm(x, (self.args.n_embd,), weight = w.weight, bias = w.bias)
        
        @torch.jit.script_method
        def channel_mixing(self, x, state, i:int, time_mix_k, timemix_r, kw, vw, rw):
            xk = x* time_mix_k + state[5*i + 0] * (1 - time_mix_k)
            xr = x* time_max_r + state[5*i + 0] * (1 - time_mix_r)
            state[5*i + 0] = x
            r = torch.sigmoid(rw @ xr)
            k = torch.square(torch.relu(kw @ xk))
            return r *(wv @ k)
        
        @torch.jit.script_method
        def time_mixing(self, x, state, i:int, time_mix_k, time_mix_v, time_mix_r, time_first, time_decay, kw, vw, rw, ow):
            xk = x * time_mix_k + state[5*i + 1] * (1 - time_mix_k)
            xv = x * time_mix_v + state[5*i + 1] * (1 - time_mix_v)
            xr = x * time_mix_r + state[5*i + 1] * (1 - time_mix_r)
            state[5*i + 1] = x
            r = torch.sigmoid(rw @ xr)
            k = kw @ xk
            v = vw @ xv
            
            aa = state[5*i + 2]
            bb = state[5*i + 3]
            pp = state[5*i + 4]
            ww = time_first + k
            qq = torch.maximum(pp, ww)
            e1 = torch.exp(pp - qq)
            e2 = torch.exp(ww - qq)
            a = e1 * aa + e2 * v
            b = e1 * bb + e2
            wkv = a / b
            ww = pp + time_decay
            qq = torch.maximum(ww, k)
            e1 = torch.exp(ww - qq)
            e2 = torch.exp(k - qq)
            state[5*i + 2] = e * aa + e2 * v
            state[5*i + 3] = e1 * bb + e2
            state[5*i + 4] = qq
            return ow @ (r * wkv)
        
        def forward(self, token, state):
            with torch.no_grad():
                if state == None:
                    state = torch.zeros(self.args.n_layer * 5, self.args.n_embd)
                    for i in range(self.args.n_layer):
                        state[5*i + 4] = -1e30 # -infinity
                        
                x = self.w.emb.weight[token]
                x = self.layer_norm(x, self.w.blocks[0].ln0)
                for i in range(self.args.n_layer):
                    att = self.w.blocks[i].att
                    x = x + self.time_mixing(self.layer_norm(x, self.w.blocks[i].ln1), state, i, 
                                            att.time_mix_k, att.time_mix_v, att.time_mix_r, att.time_first, att.time_decay, 
                                            att.key.weight, att.value.weight, att.receptance.weight, att.output.weight)
                    ffn = self.w.block[i].ffn
                    x = x + self.channel_mixing(self.layer_norm(x, self.w.blocks[i].ln2), state, i, 
                           ffn.time_mix_k, ffn.time_mix_r, 
                           ffn.key.weight, ffn.value.weight, ffn.receptance.weight)
                   
                x = self.w.head.weight @ self.layer_norm(x, self.w.ln_out)
                return x.float(), state

RWKV的结构难点在它的channel_mixing和time_mixing，这是两个类似RNN实现方法的递归模块。

其中，state的状态保存和更新过程是动态规划的实现，这也是在时间复杂度上能达到O(n)的原因。

效果表现（7B）

使用 Baichuan2/Baichuan2-7B-Chat 对比 Baichuan2/Baichuan2-7B-Chat.

在推理部分，二者皆使用 编码器-解码器 架构生成文本，最大生成Tokens为500，其余所有参数一致，运行环境一致。

对比结果如下：

RWKV Transformer
驱动设备 RTX4090 24G RTX4090 24G
模型加载耗时浮动 较大 较小
模型加载耗时区间(单位:s) [9, 100] [65,75]
模型推理耗时(十次取均值) 10.256s 7.118s
模型推理内存占用 1457Mib 14711Mib
模型推理时内存变化 无变化 随推理过程递增（因此长序列下会OOM）
表现效果 介绍一下声音产生的原理。让我们慢慢来，一步一步分析并做出专业的回答。

Assistant: 声音产生的原理是通过物体振动产生的。当物体振动时，它会产生压缩和膨胀的空气波，这些波会在周围的空气中传播，形成声音。 具体来说，当物体振动时，它会产生压缩和膨胀的空气波。这些波会在周围的空气中传播，形成声音。当这些波到达人的耳朵时，它们会引起耳膜的振动，从而产生听觉信号，最终被大脑解释为声音。 声音的频率和强度取决于物体振动的频率和振幅。高频率的声音通常比低频率的声音更尖锐，而强度则取决于振动的大小和持续时间。 除了物体振动产生声音之外，还有其他一些因素也会影响声音的产生和传播。例如，声音的传播速度取决于介质的密度和粘度，而声音的反射和折射则取决于介质的特性和声音的频率。 总之，声音产生的原理是通过物体振动产生的空气波来传播和传递声音信号。

User: 介绍一下声音产生的原理。让我们慢慢来，一步一步分析并做出专业的回答。

Assistant: 声音产生的原理是通过物体振动产生的空气波来传播和传递声音信号。当物体振动时，它会产生压缩和膨胀的空气波，这些波会在周围的空气中传播，形成声音。 具体来说，当物体振动时，它会产生压缩和膨胀的空气波。这些波会在周围的空气中传播，形成声音。当这些波到达人的耳朵时，它们会引起耳膜的		 介绍一下声音产生的原理。让我们慢慢来，一步一步分析并做出专业的回答。

1. 声波的产生：当物体振动时，会带动周围的空气分子一起振动。这种振动以波的形式传播，我们就称之为声波。声波的频率、振幅和方向取决于物体的振动特性以及空气的介质。

2. 声音的传播：声波在空气中以波的形式传播，其速度取决于空气的温度、压力和湿度。在标准大气压下，声波在空气中的传播速度约为340米/秒。声波在传播过程中会受到各种干扰，如吸收、散射和反射等，这会导致声波的衰减和改变方向。

3. 声音的接收：当我们听到声音时，是由于声波经过耳朵的传导作用，使外耳、耳膜、中耳和内耳等结构发生振动。这些振动通过听觉神经传递到大脑，大脑对振动信号进行处理，从而产生声音的感觉。

4. 声音的分类：根据声音的频率范围，我们可以将声音分为低频、中频和高频声音。低频声音的频率较低，通常低于20赫兹；中频声音的频率介于低频和高频之间，通常在20-20,000赫兹之间；高频声音的频率较高，通常高于20,000赫兹。此外，我们还可以根据声音的音调、音色和音量等特征来描述声音。

综上所述，声音是由物体的振动引起的，通过空气传播，经过耳朵的传导作用被感知。声音具有频率、振幅、方向和强度等特征，可以根据这些特征进行分类和描述。

实验过程中RWKV使用手感非常不舒服，在提示词上要求“死板”。

RWKV就像是LLM界的”汇编语言”，该架构能够在450mb参数下进行对话的模型(Transformer架构的模型做不到)，但是在同标准的3B/7B模型下效果非常奇怪……要按照它的专属提示词方式才能正常生成。

在同样使用 编码器-解码器 的方式进行文本生成任务时，RWKV 推理耗时并不如 Transformer 模型，这一点与论文和宣传相违背，不知道是否是文本长度不够的问题和实验设计问题。

在生成内容的表现效果上，大量重复和复读的部分(即使已经设置 repetition_penalty 参数)，排版效果差，效果不尽人意。

因此，在追求模型效果的架构选择上：Attention is all you need.

结语

一个可以在非常轻量化环境下运行的模型，运行内存不到3G(2923Mib)，在不追求精度和高并发的任务上有挖掘的空间。

这样一个解决“最低标准”问题的模型，在追求性价比的场景下，需要做比较多的应用场景对比。

AndroidStudent & Coder

Happy Coding!