注意力机制#
RNN 中的注意力机制由于串行架构的存在还是计算复杂度过高。因此我们对 RNN 的注意力计算公式进行改进,引入 Query、Key、Value 的概念。其中 Query 为查询信息,Key 为知识库,Value 为知识库中每个知识的释义。形式化表示为:
即通过评分函数 计算 Query 对每一个 Key 的相似度,然后将相似度矩阵与对应的 Value 聚合起来得到上下文向量,表示查询中的哪些部分是对输入是有影响的。
用查字典来举例:我们需要查询的信息为 “中国”,计算完对字典中每个词的相似度之后,得到相似度最大的两个词 “华夏”、“唐宋”。然后将这两个词对应的信息聚合起来,来代表 “中国” 这个词,即这两个词对 “中国” 的影响是最大的。
Transformer#
缩放点积注意力#
自注意力(Self-Attention) 指的是 计算同一段序列中不同令牌之间的注意力,通过该机制可以对序列中的每一个令牌进行上下文建模。
上图表示了如何计算第一个令牌与其余令牌的注意力,并将其汇总为上下文向量。这里有几点需要注意:
- 缩放参数 防止 softmax 梯度饱和。且 softmax 函数是在注意力权重矩阵的 列维度 上进行计算;
- 、、 矩阵是同一段序列通过不同的权重矩阵 、、 加权得到;
- 每一个 、、、 的维度都是 ,即单个令牌的词向量维度。
多头注意力#
单头无法表征序列的所有信息,通过引入不同的头在 同一时刻关注序列中的不同部分,从而学习到不同的上下文信息。
上图表示了如何计算第二个头所对应的注意力、并将其汇总得到第二个头的上下文向量。然后将所有头的上下文向量拼接起来,这里需要注意拼接之后的向量是 高维向量,需要压缩使得输入和输出维度保持不变。
- 实践中需要设计的超参数为头的个数 ,每个头对应的 、、、 的维度都是 ;
- 不同头之间的计算都是并行的,但是十分占用显存。
逐位前馈神经网络#
将注意力层输出的特征进一步进行非线性变换,提升模型的表达能力。逐位指的是 每个词元的特征重计算都是独立的。
上图右侧的计算表示 “深” 这个特征通过权重矩阵加权之后得到另一个更好的特征。这里使用的是两层 FFN,且权重矩阵 和 是共享的,这就类似于 RNN 的思想,因此就能够很好的处理变长序列。
绝对位置编码#
自注意力机制具有 置换不变形,无法捕捉序列中的顺序关系。因此引入位置编码对序列信息进行位置建模。
原始的 Transformer 使用的是基于 正余弦 的绝对位置编码(Sinusoid) - 即根据正余弦函数来生成 对应位置的唯一位置嵌入,并将其加到令牌的词向量上。对于维度大小为 的词向量,它的编码方式为:
其中 和 是位置 的词向量的第 和 个分量、 是词向量的维度。将该结果加到词向量对应位置上形成最终的模型输入。
作者也使用了可学习的位置编码,但是发现使用 Sinusoid 编码可以让模型拥有更好的长度外推能力。
编码器 / 解码器#
原始 Transformer 模型分别堆叠了六层编码器(Encoder)和解码器(Decoder)。且每层输入和输出的维度都是 (batch, time, dimension),即批量大小、序列长度、词向量维度。
解码器是 自回归模型,也叫 Next-token prediction。在训练阶段,使用自回归模型会有弊端。因为提供了输入序列和目标序列,使得模型在训练的时候 “作弊”,即参考目标序列直接获得下一个单词。
因此提出 掩码多头注意力。即通过对注意力矩阵进行遮掩,使得模型在训练阶段 只能参考历史信息。具体来说,对注意力矩阵中的元素通过下述规则进行遮掩(下三角遮掩):
这样,在使用 softmax 计算上下文向量的时候,由于出现了极大的负值,导致对应位置的指数趋近于 0,从而屏蔽了未来信息,保持了解码器的自回归属性。
为了保证输入序列长度一致,会在输入向量中添加 [pad] 来对齐长度。而在计算注意力矩阵的时候需要忽略这些信息,实践中通过获取 q/k 的非填充重叠部分作为掩码添加到注意力矩阵中即可(与下三角掩码原理一致)。Causal Attention 中需要用到两个掩码,其余只需要用到对齐掩码。
并且引入 交叉注意力(Cross Attention) 处理来自编码器的 K/V 和来自解码器的 Q。需要注意传入的 K/V 是编码器最后一层的输出,且传入到解码器的每一层当中。但是解码器的 Q 是当前层计算的结果。
激活函数#
GeLU 全称为 Gaussian Error Linear Unit,它利用高斯分布的累积分布函数来调节输入信息的激活程度,使得中间区域的输入梯度更为平滑,起到正则化的作用。
其中 是累积分布函数、 表示误差函数。
SwiGLU 全称为 Swish-Gated Linear Unit,它将 Swish 的平滑与 GLU 的门控机制结合起来使得梯度能够有有效传播,适用于深度模型。
在 Llama3.1 中的代码实现为:
self.w2(F.silu(self.w1(x)) * self.w3(x)) # 没有添加 bias
注意力模块的优化#
自注意力模块是二次复杂度 ,因为它要对序列中的任意两个词向量都要计算相似度,得到一个 大小的注意力矩阵。
Sparse Attention#
从注意力矩阵的角度来说,就是除了相对距离不超过 的、相对距离为 的倍数的注意力都设置为 0。这样注意力就具有了 局部紧密相关和远程稀疏相关 的特性。
该方法的不足之处就是需要人工选择保留的注意力区域,不利于扩展。但是由于将每个词元的注意力压缩在了较小的空间中(每个词元只能看到训练长度的词元),能一定程度上缓解长度外推问题。
Linear Attention#
制约注意力性能的关键因素是 softmax 函数。如果没有该函数,去掉缩放系数的注意力公式实际上就是三个矩阵连乘,复杂度仅有 。一个自然的想法就是拿掉 softmax 函数,并用一般函数 进行替代。此时注意力公式为:
为了保留注意力相似的分布特征,需要保证 恒成立。
线形注意力就是用 核函数 代替这个一般函数,即 。此时注意力公式为:
其中核函数选择 。
KV Cache#
自回归模型中的 Causal Decoder 在 token-by-token 递归生成的时候,每个时刻都需要重复计算历史时刻的 K/V。
KV Cache 指的是 缓存之前的结果,仅计算当前时刻的 K/V,然后与之前的结果连接起来即可。
但是反复读取显存、导致带宽瓶颈。后续的 MQA、GQA、MLA 都是围绕如何减少 KV Cache 的同时尽可能地保证效果。目的就是要实现在更少的设备上推理更长的上下文,或者在相同的上下文长度下让推理的批量大小更大,从而实现更快的推理速度或者更大的吞吐总量。
Multi-Query Attention(MQA)#
当 KV Cache 中存储的历史信息越来越多的时候,反复进行读取会导致带宽瓶颈。
MQA 的思想是让 所有的 Attention Head 都共享同一个 K/V,它将 KV Cache 压缩 。
该技术在 PaLM、StarCode、Gemini 等模型中使用。
Group-Query Attention(GQA)#
MQA 对 KV Cache 的压缩太过严重,会影响模型的学习效率以及最终效果。
GQA 的思想是将 Attention Head 分成 个组,每组共享同一个 K/V,它将 KV Cache 压缩 。通常设置 保证单机八卡中的每张卡负责计算一组 K/V。
该技术在 Llama2/3、DeepSeek-V1、ChatGLM2/3 等模型中使用。
Multi-Head Latent Attention(MLA)#
Flash Attention / 单机多卡#
原始的注意力计算过程中,最大的中间过程结果就是注意力矩阵,而最终的上下文向量反而很小。反复读取大尺寸注意力矩阵会导致 HBM 带宽瓶颈。
Flash Attention 通过 矩阵分块 和 算子融合 等方法,将中间计算结果保留在大带宽的 SRAM 中,获得最终结果再写回小带宽的 HBM 中,从而避免了带宽瓶颈。
上图左侧是原始数值稳定版本的 softmax 实现,右侧是将输入分块后,分别计算 softmax 并融合的实现。具体来说对于 Key/Value 的每一块,计算它们和所有 Query 的注意力矩阵,并保存在 SRAM 中。当计算得到最终上下文向量的时候在将结果写回 HMB 中。
HBM IO 复杂度对比:
- Vanilla Attention - ;
- Flash Attention - 。
算子融合指的是通过将上图右侧的整个计算过程融合成一个高效的算子,从而减少中间数据的存储。
Ring Attention / 多机多卡#
混合专家模型(MoE)#
首先通过 门控函数 / 路由网络 计算出每个专家的权重,该权重表示每个专家对当前词元的关注程度。通常是一个简单的神经网络,将输入 通过 映射为概率分布:
然后根据路由权重选择对应的专家进行计算,通常是稀疏激活,即 个专家里只有 个专家参与计算:
最后得到所有专家的输出结果之后,与路由权重进行加权,并进行残差连接:
其中 表示第 个专家的输出、 表示第 个词元与第 个专家的路由权重。
上图 表示传统的 MoE 模型架构,而图 和 则是改进的 DeepSeekMoE 架构。因为传统的 MoE 存在问题:
- 单个专家涵盖多种知识,无法充分同时利用;
- 多个专家涵盖通用的冗余知识,阻碍了 MoE 的理论性能上限。
DeepSeekMoE 通过划分更细粒度的专家以及利用单个专家存储通用知识来达到单个专家的高度专业化。形式化表达为:
具体来说,将 个专家划分为 个专家,即将隐藏层的维度乘上 ,同时选取 个专家保证参数量与传统 MoE 模型相当。然后隔离出单个专家 来存储通用知识。
归一化#
Post / Pre - Norm#
Post-Norm - 归一化放置在残差计算之后
训练时间较快、网络容易调优、输出层梯度较大引发的训练不稳定。
Pre-Norm - 归一化放置在子层计算中
训练稳定、效果不如 Post-Norm、缓解梯度消失(爆炸)问题。
RMSNorm#
RMSNorm 沿着输入特征的通道维度进行归一化操作,相较于层归一化而言,它仅保留缩放参数。形式化表示为:
其中 为输入特征、 为缩放参数。并且计算的是输入特征平方的均值。
位置编码#
相对位置编码#
注意力分数计算公式 进行因式分解得到:
其中 、 表示词向量和位置向量。可以发现只有最后一项与相对位置有关。假设去除 、 那么原始的 Transformer 是包含相对位置关系的。因为加入了非线性变换,导致相对位置信息被抹除了。
Transformer-XL 位置编码#
注意力分数计算公式修改为下述形式:
其中 为绝对位置编码矩阵、 和 都是可学习参数。
T5 位置编码#
注意力分数中加入可学习的标量:
可学习参数在所有层之间都是共享的,且每个头使用不同的位置嵌入。
ALiBi#
在 softmax 之前,将注意力分数的计算改为如下格式:
其中 是 Query 和 Key 之间的位置偏移、 是每个注意力头 的惩罚系数。该编码方式具备优秀的长度外推能力(16000 序列长度)。
长度外推问题指的是模型在推理阶段无法处理比训练阶段更长的输入序列的现象。
RoPE#
为 Query 和 Key 设置了单独的旋转矩阵 。考虑一个二维的情况:
由于内积满足线形叠加性,因此任意偶数维的 RoPE,都可以表示为二维形式的拼接:
也就是说,给位置 的向量 乘上矩阵 、位置 的向量 乘上矩阵 ,用变换后的 、 序列做注意力,那么注意力就自动包含相对位置信息了,因为如下恒等式成立:
由于 的稀疏性,直接使用矩阵乘法来计算会很浪费算力。因此使用下述方法来实现 RoPE:
其中 是逐位相乘。在 的选择上,沿用了和 Sinusoidal 位置编码的方案,即 。
RoPE 不带有显式的远程衰减,通过不同频率的三角函数有效区分了长程和短程。并且直接作用于 、,不改变注意力计算的形式,与 Flash Attention 更为契合。容易 Scale Up。
长度外推优化#
长度外推问题指的是模型在推理阶段无法处理比训练阶段更长的输入序列的现象(Train Short, Test Long)。具体来说:
- 推理的时候用到了没训练过的位置编码(绝对/相对);
- 推理的时候注意力机制处理的词元数量远超训练时的数量。
注意力机制理论上可以处理任意长序列,但是越多的词元去平均注意力就会导致注意力分布越均匀,也就不能很好地表征词元之间的关系。
Position Interpolation#
NTK-aware#
YaRN#
解码策略#
大模型的输出本质上是一个概率采样的过程,可以分为以下步骤:
- 将输入序列传入模型并生成下一个词元的概率分布 / ;
- 从分布中根据不同的采样/解码方法选择下一个词元 / ;
- 将选择的词元添加到输入序列中,重复步骤 1-2 / 。
贪婪采样的改进策略#
贪心搜索(Greedy Search) - 在每个采样步骤中选择概率最高的词元作为下一个词元。但是该方法容易陷入局部最优,生成重复、不自然的句子。
束搜索(Beam Search) - 在每个采样步骤中选取概率最高的 个句子,并最终选取整体概率最高的生成回复。其中 称为束宽度。
详细可以查看 循环神经网络 / 束搜索 ↗ 中的部分。
长度惩罚(Length Penalty) - 通过将句子概率除以其长度的指数幂 ,缓解束搜索倾向于生成较短的句子(每生成一个单词,都会乘以一个小于 1 的概率,使得句子的总体概率逐渐变小)。
出现惩罚(Presence Penalty) - 将已经生成词元的 logits 减去惩罚项 来降低该词元之后出现的概率。
频率惩罚(Frequency Penalty) - 将已经生成词元的 logits 减去其出现次数乘以惩罚项 来降低该词元之后出现的概率。
随机采样的改进策略#
温度采样(Temperature Sampling) - 通过调整 logits 的温度系数,从而保证采样过程的随机性。
Top-k 采样(Top-k Sampling) - 从概率最高的 个词元中进行采样。但是不考虑整体概率分布,无法适应不同的上下文语境。
Top-p 采样(Top-p Sampling) - 从一个符合特定概率条件的最小词元集合中进行采样,要求其中包含的所有词元的累积概率大于或等于预设阈值 。
语言模型#
预训练模型确立了 预训练-微调 的范式。通过大量无标注文本建立模型的基础能力,然后通过有标注数据进行下游任务的微调。
GPT / Decoder#
- 模型结构 - 12 层、768 词向量维度、12 个头、40,000 词表大小(B-BPE)、512 tokens 上下文长度;
- 训练范式 - 标准语言模型训练、SFT 特定任务微调(多任务学习损失);
- 实验数据 - 4.6GB 数据集、117M 参数量。
GPT-2
- 模型结构 - 48 层、1600 词向量维度、25 个头、50,257 词表大小、1024 tokens 上下文长度;
- 训练范式 - 多任务学习、通过 Prompt 实现 Zero-Shot 微调(使用无监督预训练做有监督任务);
- 实验数据 - 40GB 数据集、1.5B 参数量。
GPT-3
- 模型结构 - 96 层、12,288 词向量维度、96 个头、50,257 词表大小、2048 tokens 上下文长度;
- 训练范式 - 大规模上下文学习、Few-Shot 微调;
- 实验数据 - 570GB 数据集、175B 参数量。
BERT / Encoder#
下面的模型结构均是 base 设置。
-
模型结构
- 12 层、768 词向量维度、12 个头、30,522 词表大小(WP)、512 tokens 上下文长度;
- 在输入维度上加入 Segment Embedding,用来区分不同的句子、双向 Transformer。
假设输入为
[CLS] 句子A [SEP] 句子B [SEP]
,SE 为句子 A 的所有词元分配一个固定的嵌入向量(0),句子 B 分配另一个嵌入向量(1),从而区分两个句子。
-
训练范式
- MLM - 训练过程将输入遮掩 15%。然后预测遮掩的部分;
- NSP - 判断两个句子是否是上下文关系(50% 正/负样本)。
微调的时候不会出现 [MASK],会出现训练推理失配问题,可以采用 Scheduled Sampling 解决。即将确定要遮掩的词元,80% 替换成 [MASK]、10% 替换成随机 token、10% 保持不变。
-
实验数据 - 102 种语言、16GB 数据集、110M/340M 参数量。
RoBERTa
- 模型结构 - 12 层、768 词向量维度、12 个头、50,265 词表大小(BPE)、512 tokens 上下文长度;
- 训练范式 - 去掉 NSP 任务、引入 动态掩码 的机制,即每个回合输入句子的 mask 位置都不同;
- 实验数据 - 160G 数据集、125M/355M 参数量。
ALBERT
- 模型结构
- 12 层(参数共享)、128 词向量维度、12 个头、30,000 词表大小(SP)、512 tokens 上下文长度;
- 将 输入向量分解成两个较小的矩阵,维度分别为 和 ,其中 ;
- 训练范式 - 使用 SOP 任务预测句子顺序是否正确;
- 实验数据 - 16G 数据集、12M/18M 参数量。
T5#
将所有的自然语言处理任务都转换成 文本-文本 的形式,并用一个统一的模型。其输入是带有任务前缀的文本序列,输出是对应任务的结果。
- 模型结构 - 编码器(解码器) 各 6 层、512 词向量维度、8 个头、32,000 词表大小(SP)、512 tokens 上下文长度;
- 训练范式 - Span Corruption(随机遮掩连续 span,预测被遮掩的文本)、多任务统一学习(分类、翻译、摘要等任务均转换成文本生成任务);
- 实验数据 - 750GB 的 C4 数据集、220M/770M 参数量。
架构分析#
视觉模型#
ViT#
首先将输入图像 划分成若干 2D 块然后展平得到 ,序列有效长度为 。
然后通过线性映射转化为维度为 的嵌入向量(patch embedding),并在嵌入向量的前端添加一个可学习的类别嵌入,在编码器输出时可作为图像的特征表示。预训练阶段,MLP 分类器将该类别嵌入作为输入(预训练阶段 MLP 只有一个隐藏层,而微调阶段 MLP 只有一个线形层)。
微调阶段使用更高分辨率的图像,且保持每一个块的尺寸不变。但是输入序列的有效长度变长,导致预训练阶段的位置编码不再匹配,因此使用 插值调整位置编码。
Swin Transformer#
Swin Transformer 引入基于卷积神经网络的归纳偏置:
- 局部性 - 只计算窗口和移动窗口内的注意力,同时跨窗口连接;
- 层次化 - 通过分层结构,提取不同尺度的特征。
该模型的每一个阶段都采用两层移动窗口的设置。第一层划分窗口的自注意无法捕捉全局的特征信息。因此第二层将窗口进行滑动,这样原本的四个区域就变成了九个区域。且新窗口包含之前窗口的边界,能够建立不同窗口的连接。
为了高效地计算,首先将九个区域调整为四个区域,其中三个区域包含来自不同区域的 patch。当两个来自不同区域的 patch 交互时,在它的位置上增加一个比较大的负值,进行 softmax 时候该位置便会趋于 0。
首先将输入图像划分成不重叠的 patch,每个 patch 的大小为 ,可以使用 维的向量表示。因此网络的输入特征维度时 。
在第一个阶段,首先使用线形嵌入(Linear Embedding)层将每个 patch 的特征维度映射进行映射 。然后使用多组连续的 Swin Transformer 模块处理,如上图 (b) 所示。分别使用了基于窗口和移动窗口的多头自注意力机制。在后三个阶段,每个阶段首先使用图像块合并(Patch Merging)层产生分层表示。通过合并相邻的 patch 使得特征的维度不断发生变化。
Image GPT#
首先将输入图像进行下采样,将其转换为 序列;然后进行模型预训练,采用两种预训练方法:
Next pixel prediction 是一种自回归的预训练方法,该方法根据前面的像素值预测下一个像素值(采用光栅顺序),并最终对图像的概率密度进行整体建模。它训练的目标是最小化负对数似然:
Masked pixel prediction 首先遮掩输入序列若干位置的值,并对这些值进行预测。它训练的目标是最小化遮掩位置元素的负对数似然:
通过预训练,模型学习到输入序列的分层特征表示。
微调阶段同时优化 ,其中 代表 AR 或 MASK 损失、 代表分类损失。该方法也被称为带有 辅助训练目标 的微调(Fine-tuning with auxiliary training objective)。