一文读懂Llama2的架构和推理过程

Meta公司发布了一众备受瞩目的Llama系列大模型，其中的佼佼者无疑是优化后的Llama 2系列。该模型训练数据集2万亿token，上下文长度是由llama的2048扩展到4096，在各种基准集的测试上表现突出。深入了解Llama 2，可用于研究和商业用途。我们今天主要来了解一下该模型的整体结构和推理过程。

Llama 2简介

为了更深入了解Llama 2，我们从Meta官网探寻了一些基本信息：

· Llama 2，作为Llama的下一代版本，推出了三种尺寸供选择：7B、13B和70B。其中，7B和13B沿用了Llama 1的经典架构，而70B模型则采用了创新的分组查询注意力（GQA）架构。

· 在预训练数据方面，Llama 2的表现颇为出色。相较于Llama 1，Llama 2的预训练语料增加了40%，这使得模型能够在更丰富的语境中学习。更令人惊艳的是，预训练模型在2万亿个标记上进行了训练，且上下文长度是Llama 1的两倍。而经过微调的模型，也在超过100万条人工标注的数据上进行了精进训练。

进一步观察了Llama 2模型的评估性能，发现它在许多外部基准测试中都优于其他开源语言模型。这些测试涵盖了推理、编码、熟练度和知识等多个方面。

优化后的Llama 2展现出的卓越表现让人对其架构产生了浓厚兴趣。下面，我们将深入探索Llama 2模型的架构，以揭示其强大的能力背后的秘密。

Llama 2模型的整体结构

在深入探讨Llama 2的架构之前，我们先来对比一下传统的Transformer架构和Llama 2的架构，以便更好地理解其创新之处。

与传统的Transformer架构相比，Llama 2架构具有以下独特的特点：

1、取消Encoder，仅保留Decoder：这一设计使得模型结构更为简洁，专注于生成和解码任务。

2、采用RMSNorm并将Norm前置：这种归一化方式有助于提高模型的训练稳定性和收敛速度。

3、在Q和K上使用RoPE旋转式位置编码：这种位置编码方式能够更好地捕捉序列中的位置信息，提高模型的表达能力。

4、使用causal mask确保每个位置只能看到前面的tokens：这保证了模型在生成文本时，每个位置只依赖于前面的tokens，符合语言生成的因果性。

5、更早地将K、V拼接到当前K、V前面：这使得模型能够利用更多的上下文信息，提高文本生成的一致性和连贯性。

6、Llama 2使用了Group query attention来节省cache：这一优化减少了模型的计算量和内存占用，提高了模型的效率。

下面我们详细解释下Llama 2架构中各部分功能：

2.1 Tokenizer

Llama 2架构中的Tokenizer组件负责将输入的文本语句进行分词和编码。

分词是指将连续的文本拆分成一个个独立的词汇单元，而编码则是将这些词汇单元转换为模型能够处理的整数表示。Llama 2使用了专门的分词表（tokenizer.json）来定义词汇与整数之间的映射关系。

具体来说，当我们输入一串语句给Llama 2模型进行推理时，Tokenizer首先会根据分词表将语句中的词拆分成对应的整数序列。这些整数序列将作为模型的输入，经过模型的计算后，输出的也是整数序列。推理结束后，需要再次使用分词表将这些输出整数转换回原始的词汇，形成最终的推理结果句子。

举个例子来说，当我们输入语句“南京市长江大桥”，Tokenizer会将其分词为“南京市”、“长江”和“大桥”，并将这些词汇转换为对应的整数。模型经过计算后，输出的也是一组整数。推理结束后，再次使用分词表将这些整数转换回词汇，得到最终的推理句子。这样，我们就能够理解Llama 2如何处理输入文本，以及如何将模型的输出转换回可读的句子形式。

2.2 Token Embedding

Token Embedding负责将输入的整数序列转换为高维的特征向量表示。

在完成分词和编码后，每个词汇都被表示为一个整数。然而，模型需要将这些整数转换为更具表达能力的向量形式，以便能够捕获词汇之间的语义关系和特征信息，Token Embedding通过查找一个预定义的嵌入表（Embedding table）来实现这一转换过程。该嵌入表将每个整数映射到一个高维（dim维）的特征向量。这样，每个输入的词汇都被映射到一个稠密的向量空间中，这个向量空间能够更好地捕捉词汇的语义信息和上下文关系。

通过Token Embedding，Llama 2能够将输入的整数序列转换为模型可以处理的向量表示，从而为后续的计算提供了更加丰富和表达能力的输入特征。

2.3 RMSNorm

在完成Token Embedding之后，RMSNorm组件负责对输入的向量数据进行归一化处理。

归一化是一种常见的数据预处理技术，它可以将输入的向量数据进行缩放，使其期望值为0，标准差为1。这样做有助于加快模型的收敛速度，提高模型的性能。

RMSNorm是一种特定的归一化方法，它采用了均方根（Root Mean Square）的方式来计算缩放因子。通过对输入向量进行RMSNorm处理，可以使得不同尺度的特征具有相同的权重，从而提高了模型的稳定性和泛化能力。

经过RMSNorm处理后的向量X具有相同的尺度和分布，这有助于模型在训练过程中更快地收敛，同时提高了模型的性能表现。

Norm处理过程可以简化视为：

RMSNorm对应的计算公式为：

2.4 Self-attention

2.4.1 计算Query、Key和Value矩阵

首先，通过线性变换将输入的向量转换为Query、Key和Value矩阵。这些矩阵分别代表了查询、键和值的信息，用于计算注意力权重。

计算公式如下：

其中，WQ、WK和WV是可学习的权重矩阵，用于将输入向量映射到相应的查询、键和值空间。

2.4.2 对Q、K向量进行RoPE (旋转式位置编码)

在进行自注意计算之前，Llama 2会对Query和Key矩阵应用旋转式位置编码（RoPE）。这种位置编码方式能够更好地捕捉序列中的位置信息，提高模型的表达能力。RoPE的原理图如下：

RoPE计算过程为：

2.4.3 进行Softmax(Q,K)以及计算自我注意层的输出

在进行Softmax(Q,K)以及计算自我注意层输出之前，需要先对经过RoPE运算的Q、K向量进行矩阵运算。具体地，这个矩阵运算是将Query矩阵和Key矩阵进行点积，得到注意力得分。注意力得分衡量了每个位置对当前位置的重要性。

然后，对这些注意力得分进行Softmax运算，将它们转换成注意力权重。Softmax函数确保了这些权重的总和为1，且每个权重值都在0到1之间。

再将计算得到的注意力权重与Value矩阵进行矩阵运算，即将注意力权重应用于Value矩阵。这样，每个位置都会根据注意力权重聚合相关的值向量，形成自我注意层的输出。

对应的结构为：

Llama 2使用了Group Query Attention，整个自注意力输出的计算过程，可以归纳为：

2.5 FFN与Final RMSN

在自注意计算输出后，Llama 2会将其与输入进行残差连接，然后进行FFN（前馈神经网络）计算输出y。这个计算过程是为了进一步捕获非线性特征，并增强模型的表达能力。

在Llama 2中，FFN由两个线性层和一个激活函数组成。首先，输入向量会经过一个线性变换，然后经过激活函数（如SiLU），最后再进行另一个线性变换，得到输出向量y。

具体计算过程为：

然后将y进行Final RMSNorm归一化处理，会输出一个概率数组logits，logits里面每个元素代表了出现对应token的概率，然后可以根据策略选择某一个token作为输出。

写在最后

我们以输入“南京市长江大桥”为例，通过上述各个步骤的协同作用，整个Llama 2的推理过程可以归纳为下图：