一文读懂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的推理过程可以归纳为下图: