外行也能看懂的大语言模型结构对比！

本文作者：白牛

AI 算法论文力求公正，通常通过客观指标如精度、召回和困惑度来评估模型的优劣，而这些结果都会受到权重数值的影响。如果我们将部署和产品化也纳入考虑范围，则 AI 构成了一个完整系统。

在此基础上，除了分析数值（如《GPTQ-for-LLaMa：量化分析与优化》），我们还需要考虑更多问题，例如：在特定的应用场景下，是否 Transformer 结构是最优选择？换句话说，在该使用场景中，模型结构需要满足哪些特性？

1. 定义场景和特点

在日常工作中，我们经常要回答同事或客户的各种问题。然而受限于时间，微信或 issue 无法实现即时响应。

因此，我们从日常需求出发，开发了一个专业知识问答助手，并将其命名为“茴香豆”。它是一只可爱的布偶猫，悄悄潜入了 MMDeploy 2 群中，随时准备解答大家的疑问。

我们最初定了 3 个目标：
首先，严格遵守法律规定，避免任何有风险的闲聊。目前，“茴香豆”已完全实现这一目标，感谢 ncnn 的热心测试。

其次，应能回答部署专业问题，例如 “遇到 undefined libprotobuf.so 怎么办”、“OpenMMLab 2.0 是否需要 MMCV 2.x 版本”。“茴香豆”在这方面已展现较好的回答能力。

最后，需确保模型训练、初始化和推理的速度，以保证版本迭代效率。

经过一些训练和调优测试，我们目前采用特征检索方案。这个方案类似人脸识别流程，分 3 个阶段。

1. 底库初始化：我们从 OpenMMLab 算法仓库，筛选出 1700 份使用文档。然后把每个文档切分成近似大小的 bread，提取文本特征，并存储到数据库。如果数据较杂，也可以用 LDA（Linear Discriminant Analysis，线性判断分析）等抽取主题词做 bread；

2. 查询目标段落：提取用户输入的文本特征，和底库对比。如果特征距离接近，取目标连续的 bread 作为目标段落；如果距离过远，就认为是闲聊；

3. 获取回答：把用户输入作为 prompt，目标段落作 input，输入给 LLM，返回响应。

特征检索方案满足最初的设计目标，经过实测，确实有能力缓解响应不及时的情况，然而也存在一些不足：

1. 上下文过长导致程序崩溃。然而，在处理专业问题时，往往需要较长的上下文才能更好地理解用户输入。以下是一个真实的例子。

用户的问题含有 “mmdeploy”、“vulkan”、“ncnn” 等多个主题，按照最初的设置，需要 9100 字上下文，这导致了程序崩溃和参数调整。

2. 回答失败。虽然用特征检索方案能有效降低法律风险，但它容易导致无法匹配不到正确的底库，从而无法正确回答问题。这是一个特征匹配失败的例子。

针对 MMDEPLOY_SHARED_LIBS=ON，助手会匹配到 hrnet 相关文档，从而导致结果错误。

3. 整体资源利用率不高。我们为计算特征，额外部署了模型，并且 LLM 推理的 GPU 利用率也不高。所以能否直接使用 LLM 的中间 feature 做检索。

面对这些问题，我们研究了近期流行的开源模型结构，并尝试从中找到答案。

2. 对比方法

我们本次对比 2 种开源模型，均支持中文，详见表格：

抽象地讲，这些 LLM 从输入 “今天天气很好，所以” 开始，都需要执行这些步骤：

1. encode 和 decode。由于中文不能直接参与计算，所以要先分词并编码成 tensor 形式。计算出结果后，同样要把 tensor 解码回中文；

2. embedding 和 head。将分词得到的 tensor 转到词向量空间表示，从而捕捉单词之间的相似性和关联性。head 模块会把词向量空间特征翻译回 tensor 形式；

3. backbone（主干网络）。根据词向量空间特征，该部分负责捕捉文本中的深层次语义关系和上下文信息、理解句子中的长距离依赖关系，完成标注、翻译、续写等具体 NLP 任务。

我们将根据微信助手场景的需求，以白盒的方式，对比相同阶段不同实现。相关论文和技术细节也将一并介绍。

3. LLaMA

3.1 词向量空间

Facebook 开源的 LLaMa 模型在训练时使用的中文语料较少，其 tokenizer 的词表仅包含 32,000 个词条，因此并不能很好地表达中文。

为了解决这一问题，Chinese-LLaMA-Alpaca 增加了大量中文数据，并将词表扩展到了 49,954 个词条。以下是扩充前后分词能力的直观对比：

# 扩词前
>>> model.encode_as_pieces('今天天气怎么样')
['▁', '今', '天', '天', '<0xE6>', '<0xB0>', '<0x94>', '<0xE6>', '<0x80>', '<0x8E>', '么', '样']

# 扩充后
>>> model.encode_as_pieces('今天天气怎么样')
['▁今天', '天气', '怎么样']
>>> model.tokenize('今天天气怎么样')
[36345, 34230, 34164]

在原版 LLaMa 中，“气”字需要 6 个数才能表达。而在词表扩充后，只需 3 个数值即可表示“今天天气怎么样”。

需要注意的是，tokenize 的目的是让中文变得能处理，分词效果好坏和模型效果没有强关联。

在 LLaMa 推理过程中，embedding（词向量空间嵌入）是一个 Gather 操作。使用 llama.onnx 可视化，其表示如下。

例如，我们已经把“今天”表示成 36345，从已经训练好的特征中，取第 36345 个，作为“今天”的词向量空间表示。

可以看到，这种方法得到的词向量空间，在段落匹配方面表现不佳，因此并不适合作为“茴香豆”的知识特征。主要在于实际应用中的句子，长度和表达各不相同，而该方法未能捕捉上下文关联信息。

3.2 attention

LLaMa 使用 Transformer Attention + RoPE 作为主干结构。

Transformer 来自 2017 年发表的《Attention is all you need》，最初用来解决 NLP 领域的问题。针对循环神经网络（RNN）和长短时记忆网络（LSTM）的计算效率不足问题，论文提出了自注意力机制（Attention）和位置编码方法。

Attention 的计算过程如下图所示：

假设有一个输入序列，我们先为每个词创建三个向量：Query、Key 和 Value。接下来：

1. 将每个词的 Query 向量与其他所有词的 Key 向量进行点积，以计算得分；

2. 对得分使用 softmax 函数进行归一化。归一化得分代表输入序列中各词的关注度；

3. 将归一化得分与 Value 向量相乘，使关注度较高的词的 Value 向量保留更多信息；

4. 将加权后的 Value 向量相加，得到输出向量。

RoPe（Rotary Position Embedding，旋转式位置编码）是一种配合注意力机制的位置编码方法，由知乎用户 @苏剑林 提出。这种设计能够实现绝对位置编码的效果，同时避免了传统绝对位置编码的一些缺点。

完整的介绍可以参考以下链接：https://zhuanlan.zhihu.com/p/359502624

在推理过程中，为了保留上下文信息，LLaMa 需要存储 attention 计算得到的历史位置编码。实现里采用 concat 方法来存储这些编码；处理一个长度为 N 的序列，还需要计算一个 NxN 的注意力矩阵（attention map）。这会导致注意力矩阵所需的内存随输入呈二次方增长。

这两种方法都需要额外的内存，因此当 LLaMa 处理大篇幅问题时，容易出现内存不足（OOM）和程序崩溃。以下是触发异常的位置：

        if past_key_value is not None:
            # reuse k, v, self_attention
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)
            
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

即便内存问题能被解决，当前结构仍无法处理长文本，更多信息请阅读 RoPE 作者的 HWFT（Hybird Window-Full Attention）：https://kexue.fm/archives/9603

3.3 FFN

Attention 的实现主要包括诸如矩阵乘法、点积和 Softmax 等线性变换。这些线性变换关注全局信息，并通过计算各个位置间的权重来构建上下文表示。因此，在骨干网络中使用 FFN（Feed-Forward Neural Network，前馈神经网络）来实现非线性变换，可以更专注于处理每个位置的局部信息。

LLaMa 模型采用 SiLU（Sigmoid Linear Units）作为 FFN 的激活函数，以确保特征计算具有非线性特性。

以下是关于 SiLU 的完整解释：https://arxiv.org/abs/1702.03118v3

4. RWKV

RWKV 是一种模型，目前作者并未发表相关论文，解释权归 GitHub 用户 BlinkDL 所有。我们目前理解的 RWKV 含义如下：

• Receptions：新单词对旧单词的信息接收意愿，通过输入线性变换得到。

• Weight：旧字和新字之间的关联程度，也就是二者之间的权重。

• Key/Value：与 Transformer 类似，通过输入变换得到。

RWKV 模型得名于其统一的格式，即 time-mix 和 channel-mix 都遵循相同的 RWKV 结构设计。

实际上，RWKV 的设计灵感来源于 Transformer。从代码本身无法看出其所宣称的“RNN 结构”，需要阅读作者的历史文章才能理解。

若将模型写成 RNN 形式并恢复时间对称性，可以消除 alpha、beta 和 gamma 等因子。

从部署角度来看，RWKV 应该作为一个独立的算子实现。这样可以确保模型自成一体，便于不同场景下应用和优化。

4.1 词向量空间

RWKV 模型的词表长度为 50,254。编码过程中使用的是 huggingface/tokenizers，它基于 Rust 实现，比 SentencePiece 更快。

与 LLaMa 的嵌入层仅用 Gather 不同，RWKV 在嵌入层中增加了一次 LayerNorm（层归一化）。在 NLP 领域，LayerNorm 相对于 Batch Normalization（BN）更常用于实现归一化。主要是因为文本长度不确定，以 batch 为单位计算归一化参数并没有意义。

关于 LayerNorm 如何加速训练和增强泛化能力的深入理解，可以参考这篇文章：https://arxiv.org/pdf/1911.07013.pdf

4.2 time-mix

RWKV-4 的原型使用了 AFT（An Attention Free Transformer），AFT 用下面的操作替换了原本 transformer 中的 attention：

1. 合并 key/value。首先使用一组可学习的位置偏置（position biases）将 key 和 value (上下文信息) 合并在一起。捕捉输入序列的局部相关性；

2. 将 query 与简化的上下文结合。然后用点乘，将 query 与已合并的上下文信息相结合。这保证了 input 任意两个位置互相交互。

新的公式如下图：

其中 QKV 均来自 affine(input) ，K 的首个元素实质是 input 第一行加权累加得到的，因此，可以说 input 中任意两个位置之间存在交互；wt 就是前面提到的“可学习的位置偏置”。

这个过程既保留了原本 attention 机制的特性，同时避免计算 NxN 尺寸的 attention map，因此 AFT 模型支持更长的输入。以下是 AFT 时间、空间复杂度的对比表格，表中的 T、d 分别表示文本序列长度和特征尺寸。

需要注意的是，RWKV 调整了 AFT 实现，使用一个固定大小的 cache 累加上下文信息。因为空间复杂度和 T 不再相关，所以内存角度上支持无限长度。训练角度暂有争议。

time-shift 是作者的局部 attention 设计。举一个粗糙的例子：

1. 用户输入“今天天气怎么样”，tokenizer 编码得到尺寸 [1,3] 的 tensor，经过 embedding 处理，特征尺寸为 [1,3,1024]。这是 attention 的输入；

2. 然后强制相邻的两个特征相互结合。即“今天”和“天气”结合、“天气”和“怎么样”结合。请运行这段代码体会效果。

import torch
x = torch.rand(1, 3, 1024)

time_shift = torch.nn.ZeroPad2d((0,0,1,0))
torch.cat([time_shift(x)[:, :-1, :1024//2], x[:, :, 1024//2:]], dim = -1)

这也是作者要表达的“RNN”语意。

通过这种方法，模型能够更好地关注局部信息、捕获相邻词汇之间的潜在关联。

4.3 channel-mix

RWKV 模型中使用的 FFN 受到了两篇论文 GLU（Gated Linear Units，门控线性单元）和 SquareReLU 的启发。

GLU：是 google 提出的一种激活函数。相较于 LLaMa 使用的公式 FFN.SiLU(x, W1, W2) = SiLU(xW1)W2，GLU 增加了参数，使表达形式更加丰富。新公式为：FFN.SiLUGLU(x, W1, V, W2) = (SiLU(xW1) ⊗ xV)W2。

通过引入门控机制，GLU 能够使模型在处理序列数据时，具有更强的表示能力。

SquareReLU：是通过 Primer 搜索得到的网络结构的一部分。SquareReLU 实际上是在 ReLU 激活函数之后接一个开平方根运算。

Primer 论文 ：https://arxiv.org/pdf/2109.08668v2.pdf

RWKV 的作者将本可以发表为论文的内容，直接发布在知乎上，这客观上加快了我们对其观点的理解。我们完全理解撰写严谨论文是非常耗时的，但希望有一页基本定义，以方便更多人学习这个模型。

5. 总结

我们基于特征检索方案实现了 LLM 专业知识问答助手，并将其部署到微信 MMDeploy 2 群。在实际运行中，我们总结了各种场景问题，带着这些问题比较了 LLaMa 和 RWKV 两种模型结构的差异。在此过程中，我们还介绍了许多基础论文。

以下是对比汇总表格：

关于 GPU 利用率，我们发现这受实现影响很大，例如 ChatGLM-6B 合并了 qkv_weight，这显然能提高效率。难以从模型结构角度直接对比 GPU 情况。