01
引言
当我们深入研究 LLaMA大型语言模型架构的复杂性时，建立在 Transformer 模型的基础知识之上至关重要。为了确保对 LLaMA 有全面的了解，我建议大家首先回顾一下我在之前的文章中讨论过的 Transformer结构的系列概念。这些文章涵盖了自注意力机制、多头注意力机制、位置编码和 Transformer 整体架构等基本主题。牢固掌握这些概念将为理解 LLaMA 引入的新功能提供宝贵的框架。
闲话少说，我们直接开始吧！
02
LLAMA结构
LLAMA结构与Transformer结构的对比如下所示：
在 LLaMA 中，我们只使用Transformer结构中的解码器部分，因为它是根据上下文预测下一个Token的大型语言模型。因此，我们只需要自注意力机制来预测下一个标记。Transformer结构在输入嵌入后直接使用位置编码，而 LLaMA 则不同，它在这一阶段加入了 RMS 归一化。 此外，在Transformer中，归一化发生在每个block块之后，但在 LLaMA 中，归一化是事先应用的。归一化后，我们会计算 Q、K和 V矩阵，以进行自注意力的计算。此外，LLaMA 的自注意力机制采用了 KV-Cache和分组多查询注意力来提高模型性能和推理效率。在前馈处理层中，使用 SwiGLU 代替 ReLU激活函数。
03
LLaMA1 VS LLaMA2
LLaMA1和LLaMA2模型的对比如下：
LLaMA 1于2023年2月发布，共有四个版本，参数分别为 6.7B、13B、32.5B 和 65.2B。在 LLaMA 2 中，大多数模型参数量都增加了一倍。上下文长度（指序列长度）也有所扩展。GQA 一栏表示哪些模型使用了分组查询注意力（Grouped Query Attention）。
04
Embeddings
与Transformer类似，LLaMA 使用字节对编码（BPE）进行标记化。每个标记都映射到词汇表中的一个唯一 ID。然后，这些输入 ID 会被转换成大小为 4096 的嵌入向量（而Transformer中该维度大小为 512）。这些嵌入向量是可学习的参数。
05
RMS Norm
上图显示了 nn.Linear(in_features=3, out_features=6, bias=True)背后的数学原理。如果我们输入 X (6,3)，输出 O (6,4)。偏置向量将广播到 XWᵀ 矩阵的每一行。之后，我们将对矩阵 O 中的每一项应用一个激活函数。神经元对给定数据项的输出取决于输入数据的特征和神经元的参数。我们可以将神经元的输入视为上一层的输出。如果上一层的权重因梯度下降而更新后，其输出发生了很大变化，那么下一层就会收到明显改变的输入。因此，下一层将被迫在梯度下降的后续步骤中大幅重新调整权重。神经网络中内部节点（神经元）的分布发生变化的现象被称为内部变量偏移（Internal Covariate Shift）。这种偏移是不可取的，因为它会减慢训练过程，因为神经元必须根据前几层输出的显著变化重新调整权重。
上图中的每个item都用其归一化值进行更新，将其转换为均值为 0、方差为 1 的正态分布。参数 gamma 和 beta 是可学习参数，使模型能够根据损失函数的要求调整每个特征的比例和平移。采用批归一化BatchNorm时，我们按列（特征）归一化，采用层归一化LayerNorm时，我们按行（数据项）归一化。
06
旋转位置编码
07
Self Attention
08
KV Cache
为了解决上述问题，作者采用了 KV 缓存技术。KV 缓存通过存储和重复使用之前遇到过的标记的计算来减少推理过程中的计算冗余。
上图说明了自注意力在下一个标记预测任务中的运作方式。Q（查询）矩阵由代表序列中 n 个标记的向量组成。Kᵀ 矩阵是 Q 矩阵的转置。Q 乘以 Kᵀ 得到一个 (n, n) 矩阵，然后再乘以 V 矩阵，计算出注意力分数。然后，这些分数会通过Transformer的线性层。线性层产生对数，然后将其输入 softmax 函数。Softmax 可以帮助从词汇中确定模型预测的下一个标记。让我们一步步深入了解。
通过使用 KV 缓存，我们存储了 K（键）和 V（值）矩阵。当遇到一个新的标记时，我们会将其添加到缓存的K 和 V 矩阵中，而 Q（查询）矩阵只包含上一步的输出。这意味着在第二步推理中，我们不会将第一个标记追加到 Q 矩阵中，而是用下一个标记替换第一个标记。重要的是，我们不会将注意力机制的输出作为下一步的输入。相反，我们将注意力机制映射到的标记作为下一步的输入。这种方法减少了冗余计算，提高了推理效率。
09
Multi-Query Attention
GPU 擅长快速执行并行计算，但在从内存中传输数据时却不那么高效。这意味着，处理速度的瓶颈往往不是我们执行的运算数量，而是这些运算所需的数据传输量。计算中涉及的张量的大小和数量在很大程度上会影响数据传输。因此，我们的目标不仅要优化运算次数，还要尽量减少内存访问和传输。
在《Multi-Query Attention》论文中，作者计算了所执行的算术运算次数和这些运算涉及的总内存。他们发现，所执行的算术运算次数为 O(bnd²)，其中 b 是批量大小，n 是序列长度，d 是嵌入向量的大小。运算所涉及的总内存由计算所涉及的所有张量之和给出，结果为 O(bnd+bn²+d²)。总内存与算术运算次数之比为 O(1/k+1/bn)。由于这个比率远小于 1，表明内存访问次数大大少于算术运算次数。因此，在这种情况下，内存访问并不是瓶颈。
当我们引入 KV 缓存方法时，执行的算术运算次数仍为 O(bnd²)。然而，运算涉及的总内存却变为 O(bn²d+nd²)。如果序列长度接近嵌入向量的大小（n≈d），或批处理大小接近 1（d≈1），这一比率就接近 1。在这种情况下，内存访问成为算法的瓶颈。通常情况下，批量大小大于 1，因此这不是一个重要问题。但是，对于 n/d 项，我们需要减少序列长度来缓解瓶颈。
为了解决这个问题，我们引入了带 KV 缓存的多查询注意力。这种方法去除了 K 和 V 矩阵中的 h 维度，但保留了 Q 的 h 维度，这意味着所有不同的查询头共享相同的键和值。采用这种方法后，执行的算术运算次数仍为 O(bnd²)。但是，操作涉及的总内存变为 O（bnd + bn²k + nd²），比率变为 O（1/d + n/dh + 1/b）。与之前的方法相比，我们将 n/d 项减少了 h 倍。这一优化在显著提高性能的同时，模型的质量也略有下降。
10
Grouped Query Attention
鉴于以上分析，我们有多个查询头，但键和值只有一个查询头。通过分组多头注意力机制，我们将查询分成若干组，每组的 K 和 V 都有一个不同的头。这就减少了所需的计算量，虽然对质量略有影响，但影响很小。
11
前馈层
前馈层中的激活函数用 SwiGLU 激活函数。作者比较了在Transformer架构的前馈层中使用不同激活函数的Transformer模型的性能。在这个修改后的架构中，矩阵由原来的两个变为三个。为确保公平比较，作者缩小了维度，使参数总数保持不变。
12
总结
本文讨论的 LLaMA 模型增强了用于预测下一个标记的Transformer架构。该模型的主要特点包括旋转位置嵌入、KV 缓存和多查询注意力机制，从而优化了效率和性能。重点介绍了归一化、激活函数和 SwiGLU 使用方面的差异，强调了这些功能在提高训练速度和处理长期依赖性方面的重要性。
您学废了嘛？
点击上方小卡片关注我
添加个人微信，进专属粉丝群！
