Transformer 架构及配置网上有很多成熟的教程,本篇笔记不下钻细节,主要过一下几个术语。相关资料主要来自中国人民大学高瓴人工智能学院赵鑫教授团队出品的 《大语言模型》,感谢 DataWhale 组织的组队学习活动。

大语言模型与Transformer

当前主流的大语言模型都基于 Transformer 模型进行设计的。Transformer 是由多层的多头自注意力(Multi-head Self-attention)模块堆叠而成的神经网络模型。原始的 Transformer 模型由编码器和解码器两个部分构成,而这两个部分实际上可以独立使用,例如基于编码器架构的BERT模型和解码器架构的GPT模型。与 BERT 等早期的预训练语言模型相比,大语言模型的特点是使用了更长的向量维度、更深的层数,进而包含了更大规模的模型参数,并主要使用解码器架构,对于 Transformer 本身的结构与配置改变并不大。

输入编码

在 Transformer 模型中,输入的词元序列 (𝒖 = [𝑢1, 𝑢2, . . . , 𝑢𝑇 ]) 首先经过一个输入嵌入模块(Input Embedding Module)转化成词向量序列。具体来说,为了捕获词汇本身的语义信息,每个词元在输入嵌入模块中被映射成为一个可学习的、具有固定维度的词向量𝒗𝑡 ∈R𝐻。由于Transformer的编码器结构本身无法识别序列中元素的顺序,位置编码(Position Embedding, PE)被引入来表示序列中的位置信息。给定一个词元 𝑢𝑡 ,位置编码根据其在输入中的绝对位置分配一个固定长度的嵌入向量 𝒑𝑡 ∈ R𝐻。然后,每个词元对应的词向量和位置向量将直接相加,生成了最终的输入嵌入序列 𝑿 = [𝒙1, . . . , 𝒙𝑇 ],并且被传入到后续层中:𝒙𝑡 = 𝒗𝑡 + 𝒑𝑡.

多头自注意力机制

多头自注意力是 Transformer 模型的核心创新技术。相比于循环神经网络(Re- current Neural Network, RNN)和卷积神经网络(Convolutional Neural Network, CNN) 等传统神经网络,多头自注意力机制能够直接建模任意距离的词元之间的交互关 系。作为对比,循环神经网络迭代地利用前一个时刻的状态更新当前时刻的状态, 因此在处理较长序列的时候,常常会出现梯度爆炸或者梯度消失的问题。而在卷积神经网络中,只有位于同一个卷积核的窗口中的词元可以直接进行交互,通过 堆叠层数来实现远距离词元间信息的交换。

多头自注意力机制通常由多个自注意力模块组成。在每个自注意力模块中,对于输入的词元序列,将其映射为相应的查询(Query, 𝑸)、键(Key, 𝑲)和值(Value, 𝑽 )三个矩阵。然后,对于每个查询,将和所有没有被掩盖的键之间计算点积。这些点积值进一步除以 √𝐷 进行缩放(𝐷 是键对应的向量维度),被传入到 softmax 函数中用于权重的计算。

自注意力机制能够直接建模序列中任意两个位置之间的关系,进而有效捕获长程依赖关系,具有更强的序列建模能力。另一个主要的优势是,自注意力的计算过程对于基于硬件的并行优化(如 GPU、TPU 等)非常友好,因此能够支持大规模参数的高效优化。

前馈网络层

为了学习复杂的函数关系和特征,Transformer 模型引入了一个前馈网络层 (Feed Forward Netwok, FFN),对于每个位置的隐藏状态进行非线性变换和特征 提取。具体来说,给定输入 𝒙,Transformer 中的前馈神经网络由两个线性变换和一个非线性激活函数组成:

FFN(𝑿) = 𝜎(𝑿𝑾𝑈 + 𝒃1)𝑾𝐷 + 𝒃2,

其中 𝑾𝑈 ∈ R𝐻×𝐻′ 和 𝑾𝐷 ∈ R𝐻′×𝐻 分别是第一层和第二层的线性变换权重矩阵, 𝒃1 ∈ R𝐻′ 和 𝒃2 ∈ R𝐻 是偏置项,𝜎 是激活函数(在原始的 Transformer 中,采用 ReLU 作为激活函数)。前馈网络层通过激活函数引入了非线性映射变换,提升了 模型的表达能力,从而更好地捕获复杂的交互关系。

编码器

在 Transformer 模型中,编码器(Encoder)的作用是将每个输入词元都编码成一个上下文语义相关的表示向量。编码器结构由多个相同的层堆叠而成,其中每一层都包含多头自注意力模块和前馈网络模块。在注意力和前馈网络后,模型使用层归一化和残差连接来加强模型的训练稳定度。其中,残差连接(Residual Connection)将输入与该层的输出相加,实现了信息在不同层的跳跃传递,从而缓解梯度爆炸和消失的问题。而 LayerNorm 则对数据进行重新放缩,提升模型的训练稳定性。

编码器接受经过位置编码层的词嵌入序列 𝑿 作为输入,通过多个堆叠的编码器层来建模上下文信息,进而对于整个输入序列进行编码表示。由于输入数据是完全可见的,编码器中的自注意力模块通常采用双向注意力,每个位置的词元表示能够有效融合上下文的语义关系。在编码器-解码器架构中,编码器的输出将作为解码器(Decoder)的输入,进行后续计算。

解码器

Transformer 架构中的解码器基于来自编码器编码后的最后一层的输出表示以及已经由模型生成的词元序列,执行后续的序列生成任务。与编码器不同,解码器需要引入掩码自注意力(Masked Self-attention)模块,用来在计算注意力分数的时候掩盖当前位置之后的词,以保证生成目标序列时不依赖于未来的信息。

除了建模目标序列的内部关系,解码器还引入了与编码器相关联的多头注意力层,从而关注编码器输出的上下文信息 𝑿𝐿。同编码器类似,在每个模块之后,Transformer 解码器也采用了层归一化和残差连接。在经过解码器之后,模型会通过一个全连接层将输出映射到大小为 𝑉 的目标词汇表的概率分布,并基于某种解码策略生成对应的词元。在训练过程中,解码器可以通过一次前向传播,让每个词元的输出用于预测下一个词元。而在解码过程,解码器需要经过一个逐步的生成过程,将自回归地生成完整的目标序列。

注意力机制

注意力机制是Transformer架构中的核心技术,它能够针对序列中的词元对构建交互关系,聚合来自于不同位置的语义信息。下面介绍四种常见的注意力机制 的设计方法。

• 完整自注意力机制. 在原始的 Transformer 模型中,注意力机制通过成对的方式进行序列数据的语义建模,充分考虑了序列中所有词元之间的相互关系。其中,每个词元在注意力计算中都需要对于其前序的所有词元的键值对予以访问,因此对于序列长度为 𝑇 的序列需要 𝑂(𝑇2) 的计算复杂度。此外,Transformer 还引入了多头注意力机制,将查询、键和值在不同的语义空间进行线性投影,然后将每个头的输出进行聚合形成最终的输出。

• 稀疏注意力机制. 尽管完整自注意力机制具有较强的建模能力,但是它需要平方级的计算复杂性。在处理长序列时较为显著,带来了较大的计算和存储开销。为了降低注意力机制的计算复杂度,研究人员提出了多种高效的注意力变种。其中,滑动窗口注意力机制(Sliding Window Attention, SWA)是大语言模型中使用最多的一种稀疏注意力机制。不同于完整的注意力机制,滑动窗口注意力根据词元位置,仅仅将位置索引上距离该词元一定范围内的词元考虑到注意力的计算中。具体来说,滑动窗口注意力设置了一个大小为 𝑤 的窗口,对每个词元 𝑢𝑡 ,只对窗口内的词元 [𝑢𝑡−𝑤+1,…,𝑢𝑡] 进行注意力计算,从而将复杂度降低到𝑂(𝑤𝑇)。进一步,通过信息的逐层传递,模型实现了随着层数线性增长的感受野,从而获取远处词元的信息。

• 多查询/分组查询注意力. 为了提升注意力机制的效率,多查询注意力(Multi- Query Attention, MQA)提出针对不同的头共享相同的键和值变换矩阵 。这种 方法减少了访存量,提高了计算强度,从而实现了更快的解码速度(具体可以参考第 9.2.1 节),并且对于模型性能产生的影响也比较小。一些代表性的大语言模 型,如 PaLM 和 StarCoder ,已经使用了多查询注意力机制。为了结合多 查询注意力机制的效率与多头注意力机制的性能,研究人员进一步提出了分组查询注意力机制(Grouped-Query Attention, GQA)。GQA 将全部的头划分为若 干组,并且针对同一组内的头共享相同的变换矩阵。

• 硬件优化的注意力机制. 除了在算法层面上提升注意力机制的计算效率,还可以进一步利用硬件设施来优化注意力模块的速度和内存消耗。其中,两个具有代表性的工作是 FlashAttention 与 PagedAttention 。相比于传统的注意力实现方式,FlashAttention 通过矩阵分块计算以及减少内存读写次数的方式,提高注意力分数的计算效率;PagedAttention 则针对增量解码阶段,对于 KV 缓存进行分块存储,并优化了计算方式,增大了并行计算度,从而提高了计算效率。

混合专家模型

随着模型参数规模的扩大,计算成本也随之增加。为了解决这一问题,研究人员 在大语言模型中引入了基于稀疏激活的混合专家架构(Mixture-of-Experts, MoE), 旨在不显著提升计算成本的同时实现对于模型参数的拓展。

在混合专家架构中,每个混合专家层包含 𝐾 个专家组件,记为 [𝐸1, 𝐸2, . . . , 𝐸𝐾 ], 其中每个专家组件 𝐸𝑖 都是一个前馈神经网络。对于输入的每个词元表示 𝒙𝑡,模型通过一个路由网络(或称为门控函数)𝐺 来计算该词元对应于各个专家的权重。 在路由函数中,首先通过线性层 𝑾𝐺 ∈ R𝐻×𝐾 映射为 𝐾 个专家的得分,并基于此选择出概率最高的 𝑘 个专家进行激活。随后,这 𝑘 个专家的得分将被送入 softmax 函数计算出它们的权重 𝐺 (𝒙𝑡 ) = [𝐺 (𝒙𝑡 )1, . . . , 𝐺 (𝒙𝑡 )𝑘 ],没有被选择的专家权重将 被置为 0。上述路由网络的计算过程如下式所示:

𝐺(𝒙𝑡) = softmax(topk(𝒙𝑡 · 𝑊𝐺)).

之后,每个被选择的词元的输出的加权和将作为该混合专家网络层的最终输出 𝒐𝑡:

∑𝐾 𝑖=1

目前具有代表性的混合专家模型是 Mixtral (8×7B),该模型在 Mistral (7B) 的 基础上,使用了混合专家模块。具体来说,Mixtral 每一层都配备了 8 个专家(7B), 并对每个词元选择 2 个专家进行后续计算。在每次计算被激活的参数仅仅有 13B 的情况下,其性能超越了更熟规模更大的 LLaMA-2 (70B),进一步证明了混合专家架构的有效性。

Q.E.D.