目录文章目录目录Transformer 架构Input EmbeddingPositional EncodingEncoder软对齐注意力的思想词向量相似度的计算方法Scaled Dot-Product AttentionSelf-AttentionMulti-Head AttentionSelf-AttentionFeed Forward Neural NetworkResidual ConnectionLayer NormalizationDecoder自回归生成模型Multi-Head AttentionCross-AttentionMasked Multi-Head AttentionSelf-AttentionLinear SoftmaxLayerAttention 的优化方向Sparse Attention稀疏注意力多头注意力优化Flash AttentionPage AttentionRadix Attention参考文档Transformer 架构Transformer 架构的重点可以总结为 3 句话Transformer 延续了 Encoder-Decoder 架构所以它可以轻松处理长度不一致的输入序列和输出序列Transformer 的 Encoder-Decoder 完全由 Attention 机制构成所以它没有 RNN 的长距离依赖问题以及串行计算问题可以轻松处理超长文本和使用多机多卡并行计算Transformer 的 Encoder 采用了 Self-Attention 机制Decoder 采用了 Self-Attention 和 Cross-Attention所以它既实现了输入序列的自关联特征捕捉Self-Attention又实现了两个序列之间的联合对齐Cross-Attention是一个通用 Seq2Seq 框架不仅仅适配 MT 任务场景。下面我们逐一展开 Transformer 架构的细节。Input Embedding在 NLP 任务中往往需要将自然语言的输入转化为机器可以处理的向量。在深度学习中承担这个任务的组件就是 Embedding 层。Encoder 的 Inputs 是一个文本序列首先需要进行 Embedding词嵌入操作根据词汇表将本文中的每个词转化为词向量表示。以下是 Embedding 阶段的关键概念document训练场景中可以将数据集文件例如 jsonl 或 bin中的一行/一句理解为一个 document包含多个 token每个 document 的长度可能不一样有长有短。token文本的最小单位可以是字、词、子词比如 “开源软件” 可能被切分为 [“开”, “源”, “软”, “件”]。tokenizetoken 化将文本转换为一个个 token。token IDtoken ID 化将每个 token 转化为一个数字 ID称为 token ID。embedding词嵌入/词向量化根据词汇表将每个 token ID 转化为词向量。vocab_size词汇表的大小往往高达数万/数十万。embedding_dim词向量的维度/大小/长度。hidden_size是词向量的大小/长度和 embedding_dim 相同代码中惯用 hidden_size 变量名。batch_size一次性读取的 document 的数量比如多少句话用于同一个 Step 的训练。seq_len表示 document 的合法输入长度进行向量化之前模型会对输入的文本进行裁剪将段落等裁剪成短文本。以此来约束输入的长度。head_size多头注意力的头数在 MHA 中每个 head 计算的 d_k 等于 hidden_size / head_size。sample训练场景中表示一个数据样本。在采样过程中每条 document 的长短各不相同为了按照 seq_len 大小来取出出相同大小的 sample可以将各个 document 拼在一起形成一条很长的 token 序列然后按照 seq_len 截取得到长度一致的 sample。这样可以避免数据空间浪费。在 Input Embedding 阶段会执行以下关键步骤对输入的文本进行裁剪将长文本裁剪成长度为 seq_len 的短文本。对文本进行分词tokenize采用 BPE 等 tokenizer分词器将连续的文本转换成一个个 token。包括规范化Normalization对文本进行必要的清理工作例如清理特殊符号、格式转换、过滤停用词、Unicode 标准化等。预分词Pre-tokenization将输入拆分为 token。后处理Post-processing添加分词器的特殊 token、生成注意力掩码等等。Token Index 索引化通过词汇表将 token 转换成计算机更容易处理的数字信息。每个 token 都会去词会表中查询得到该 token 在词汇表中的 ID 序号和输入矩阵形状为 [seq_len, batch_size, token_index]。Word Embedding输入矩阵经过 Embedding 权重矩阵形状为 [vocab_size, hidden_size]。然后通过 index 的匹配最终得到 Input Tensor形状为 [seq_len, batch_size, hidden_size]再进入到后续张量计算的环节。如下图所示。Positional Encoding序列模型处理序列数据的前提是 “时序位置”在 RNN 中输入序列会沿着语句本身的顺序被依次递归处理因此输入序列的顺序提供了极其重要的信息这也和自然语言的本身特性非常吻合。但正如前文提到过Transformer 没有 RNN 结构Self-Attention 又本身不感知词的时序顺序所以需要结合 Positional Encoding位置编码来提供时序能力即捕捉序列中词汇的位置信息使得模型能够理解单词之间的相对顺序。Positional Encoding即根据序列中 token 的相对位置对其进行编码再将位置编码加入词向量编码中。具体而言Positional Encoding 会为 Input Sequence 中不同位置的词分配一个唯一的位置向量并将这个位置向量与词的嵌入向量相加使得词向量的特征既包含语义信息又包含了位置信息。位置编码的方式有很多绝对位置编码思路是为序列中每个位置生成一个唯一的编码代表该位置在整个序列中的固定 “绝对” 位置。相对位置编码没有完整建模每个输入的位置信息而是在算 Attention 的时候考虑当前位置与被 Attention 的位置的相对距离由于自然语言一般更依赖于相对位置所以相对位置编码通常也有着优秀的表现。旋转位置编码RoPE是目前 LLM 的主流编码RoPE 的目标是让 Transformer 模型在计算 Query 和 Key 的内积时自然地包含相对位置信息而无需引入额外的相对位置偏置。它通过对 Query 和 Key 向量进行位置相关的旋转操作使得不同位置的向量在进行点积时其结果能够反映出 token 之间的相对距离。Transformer 采用的是正余弦函数sin/cos绝对位置编码公式如下pos 为 token 在句子中的位置2i 和 2i1 则指示了词向量 index 是奇数位置还是偶数位置从上式中可以看出对于奇数 index 和偶数 index 的元素Transformer 采用了不同的函数sin/cos进行编码。如下图所示用 X_d*n 表示d 为 token 数量n 为词向量的维度。得到位置编码矩阵之后再和输入矩阵进行相加就最终得到了包含位置信息的 Input Tensor 了。Encoder软对齐注意力的思想公式可知Attention 的 3 个核心参数是 Query查询、Key键、Value值它们的命名借鉴了 IT 信息检索系统的设计思想。因为 Attention 计算得出了每个词和其他词的上下文语义关系之后再反过来查找某个词的关联关系时本质上就是一次 “信息检索” 过程 —— 用 Q 发起检索用 K 匹配关联对象用 V 获取关联的实际语义信息。一个通俗的例子来说明假设你在淘宝上进搜索 “李宁鞋” 时Query在搜索栏输入的查询内容 “李宁鞋”。Key模糊搜索后展示出来的商品条目只包括标题、图片等简略信息。Value点击某个商品进入查看详细的信息。软对齐注意力机制的本质就是这个 “信息检索” 过程Q 和 K 计算相似度。得到 QK 相似度/权重列表。根据相似度列表得到商品 ID。根据商品 ID 来获取实际的 V。再举一个字典查询的例子假设一个字典Key 是词条Value 是数值。{apple:10,banana:5,chair:2......}此时我通过字典查询 apple 的含义会经历 3 个步骤发起 Query查输入 apple 作为一个 Query。匹配 Key键字典里有很多 Keysapple、banana、chair 等等词条此时用 apple 和 Keys 进行匹配。获取 Value值匹配到 apple 的 Key 后得到对应的 Value 数值 10。这是最终我想要获取的实际信息。不过上述是一个简单的一对一精确匹配的理解说明 Q、K、V 的命名含义。但实际上MT 场景中常见的是一对多的场景例如多义词、纲目类属等等这时 Query 要匹配多个 Keys例如我想要查询 fruit那么此时应该 apple 和 banana 都能够匹配到但不能匹配到 chair。因此我们需要引入 Soft Attention软注意力机制其具有注意力权重的特性。当 Query 为 fruit 时三个 Key 的注意力权重分别为{apple:0.6,banana:0.4,chair:0}然后Query 为 fruit 的 Value 计算公式为 “加权求和” 计算公式Value10*0.6 5*0.4 2*08上面通过简单的例子帮助我们理解了 QKV 的含义和工作流程下面我们再具体解释 NLP 领域中 Attention 计算的实际情况。在 NLP 领域Attention 机制就做 3 件事情QK^T计算相似度对 2 个文本序列的元素之间依次进行相似度计算寻找出一个序列的每个元素对另一个序列的每个元素的相关度。Softmax(QK^T)计算注意力权重根据相似度的大小计算词与词之间的注意力权重数值即分配注意力。Softmax(QK^T)V计算加权的特征融合如下图特征矩阵 A 包含了词与词之间的注意力权重和 V 相乘后得到最终的 “搜索” 结果通过加权求和的方式计算。注特征是一个数学概念不能使用物理世界概念来进行理解。特征使用数学数字和方法来抽象的表示了某种状态或关系。这里的特征矩阵表示了词与词之间的关注度大小。词向量相似度的计算方法在 NLP 领域中注意力权重的计算是由 “文本向量化表示法” 中的向量空间距离决定的。例如 Word2Vec 通过合理的训练拟合词向量能够表征语义信息从而让语义相近的词在向量空间中距离更近语义较远的词在向量空间中距离更远。通常的可以用向量点积法Dot-Product来衡量词向量之间的相似性根据词向量的定义语义相似的两个词对应的词向量的点积应该大于0而语义不相似的词向量点积应该小于0。如下为 2 个向量的点积计算公式。注点积Dot Product运算表征一个向量在另一个向量上的投影。投影值越大说明两个向量的相关性越高如果两个向量的夹角为 90 度则这两个向量线性无关完全没有相关性。基于效率和建模能力的考虑Transformer 选择了点积的方式来计算词向量之间的相似度。Scaled Dot-Product AttentionSelf-Attention、Cross-Attention 和 Masked-Attention 都是一种 Attention 思想而 Scaled Dot-Product Attention缩放点积注意力则是 Attention 的具体算法其数学公式如下图所示。​Cross-Attention用于 MT 场景Q 往往来自于一个源语言序列K 与 V 来自于另一个目标语言序列通过学习来拟合这 2 个序列之间的关系。Self-Attention用于更通用的场景Q、K、V 都来自同一个文本序列通过学习来拟合词与词之间的相似度。Masked-AttentionQ、K、V 都来自同一个文本序列但是使用掩码掩盖部分信息时。假设序列的长度都为 seq_len、词向量的维度为 d_k那么经过 Embedding 之后 Q、K、V 矩阵的形状如下Q[seq_len, d_k]K[seq_len, d_k]V[seq_len, d_k]它们之间的 Attention 计算的步骤入如下求 Q 和 K 的语义相似度 Q 矩阵和 K 矩阵的转置相乘即多个词向量和键向量之间的向量点积计算结果得到 “相似度矩阵 X”形状为 [seq_len, seq_len]其中元素 X_ij 表示 Q 中第 i 个词和 K 中第 j 个词的相似程度。值越大关联度越高。缩放防止梯度消失还需要注意的是在上一个公式中如果 Q 和 K 对应的 d_k 维度比较大那么 QK^T 的数值会被拉得很大。因为向量点积结果与 d_k 维度成正比维度越大乘积累加的数值就越大或越小-。这会导致下一步的 softmax 归一化后的权重值会极端化少数正相关接近 1大部分负相关接近 0此时 softmax 的梯度会趋近于 0梯度消失导致模型训练无法收敛低 Loss。因此我们还需要对 Q 和 K 乘积的结果做一次缩放Q、K 点积后再除以根号 d_k使得 softmax 结果的方差稳定在 1 附近避免数值极端化保证 Softmax 有正常的梯度。这样也就得到了最终的 Attention 计算公式。将相似度转为注意力权重相似度矩阵 X 再通过 Softmax归一化得到注意力权重矩阵 A形状为 [seq_len, seq_len]其中元素 A_ij 表示 Q 中第 i 个词和 K 中第 j 个词的注意力权重取值范围在 0~1 之间。且对于 Q 中第 i 个词而言每一行的权重和为 1相似度越高占比权重也越大。加权融合计算 Attention 输出注意力权重矩阵 A 和 V 矩阵相乘得到最终 Attention 输出矩阵 O形状为 [seq_len, d_k]其中向量 O_i 表示 Q 中第 i 个词 “加权融合” 了全序列上下文的动态特征是 Query 检索中获取匹配到的 Key 的实际 Value 信息。根据上述过程就得到了完整的注意力计算公式了下面再给出具体的代码实现defattention(query,key,value,dropoutNone): query: 查询矩阵 key: 匹配键矩阵 value: 真值矩阵 # d_khidden_size/head_sized_kquery.size(-1)# key.transpose(-2, -1) 将 K 进行转置# torch.matmul(query, key.transpose(-2, -1)) 求 Q、K 的点积# math.sqrt(d_k) 根号 d_k# scores 是 Q、K 的点积除以根号 d_k 的结果scorestorch.matmul(query,key.transpose(-2,-1))/math.sqrt(d_k)# Softmax 归一化p_attnscores.softmax(dim-1)# 采样ifdropoutisnotNone:p_attndropout(p_attn)# 加权融合得到最终的语义单元returntorch.matmul(p_attn,value),p_attnSelf-AttentionSelf-Attention 是 Transformer Encoder 的核心负责处理 Input Sequence输入序列目的是要理解序列中每个词的时序和语义。所谓 “语义” 就是找对词与词之间的上下文关系并通过权重值来表示关系的深浅。下面以句子 “我吃了苹果它很甜” 为例。找关系Self-Attention 并行扫描每个词和其他词之间的关系比如 “苹果” 和 “它” 语义最相关和 “甜” 次之和 “很” 几乎没有关系。这些 “关系” 的远近是由词汇表中的高维词向量空间模型距离计算公式得到的。加权重为每个词和其他词之间的关系远近分配一个 “注意力权重”比如 “苹果” 和 “它” 语义最相关权重值最大比如 0.8和 “很” 几乎没关系权重值最小比如 0.00001。注意力权重是软注意力的特性之一。融信息把每个词和其他词的信息进行 “加权融合”比如把 “苹果” 的原始信息和 “它” 的信息进行融合其他词权重太低几乎不融信息最终得到一个 “语义单元”。如此的Self-Attention 令每句话中的每个词都成为了带着全句上下文信息的语义单元。后续将这些 “语义单元” 传递到 Decoder 之后Decoder 就可以根据每个 “语义单元” 进行对齐翻译输出序列了。例如句子 The animal didn’t cross the street because it was tired其中 it 和 animal 的关联必然最大其权重也最大it 的语音单元必然包含 animal 的信息如此的 Decoder 就知道了代词 it 指代的是名称 animal 而不是其他词翻译为 “动物”。并且因为 Self-Attention 处理的 Input Tensor 中已经包含了 Positional Encoding不再依赖 RNN所以 Self-Attention 没有 RNN 的缺陷不依赖定长向量 C 来传递上下文语义所以不管序列有多长Self-Attention 都能建立序列中所有词相互之间的关系。由于 Self-Attention 计算序列中每一个词的过程都是相互独立的所以可以应用多机多卡并行计算。值得注意的是​经典的 Attention 机制中如 Cross-AttentionQ 往往来自于一个序列K 与 V 来自于另一个序列。例如在 Transformer Decoder 中Q 来自于 Decoder 的输入K 与 V 来自于 Encoder 的输出从而拟合了编码信息与历史信息之间的关系便于综合这两种信息实现未来的预测。但是 Transformer Encoder 中的 Self-Attention 即是计算本身序列中每个元素对其他元素的注意力分布。也就是说Self-Attention 的 Q、K、V 都由同一个输入序列通过不同的参数矩阵计算得到。从而拟合输入语句中每一个词对其他所有词的关系。具体而言经过 Embedding 和 Positional Encoding 之后的 Input Tensor 进入 Encoder 时首先会通过 3 个独立的 Linear 线性层它们分别对应将 Input Tensor 映射到 Q、K、V 空间的权重矩阵 Wk、Wq、Wv在训练的初始化阶段随机产生。此处 Linear线性层的作用是对 Input Tensor 做纯线性变换输入矩阵 × 权重矩阵 可选偏置。如下公式是 Self-Attention 中 Linear 进行的线性变换计算公式。假设 Input Sequence 的长度为 seq_len那么Input Tensor 形如 [batch_size, seq_len, d_k]其中 batch_size1一句话d_khidden_size/head_size词向量维度除以多头注意力的头数。这句话的 Q 就是一个形如 [seq_len, d_k] 的矩阵每行表示一个 q 词向量这句话的 K 也是一个形如 [seq_len, d_k] 的矩阵每行表示一个 k 键向量。这句话的 V 还是一个形如 [seq_len, d_k] 的矩阵每行表示一个 v 值向量。注意QK 匹配后找到的 V 同样来自于输入本身。通过 Self-Attention 计算之后最终输出矩阵 OHidden States。其中O_i 表示第 i 个词的输出 Q 中第 i 个词对 K 中所有词的注意力权重 x V_j 值向量然后再求和。例如Q_它 0.9×V_苹果 0.05×V_我 0.05×V_甜让 “它” 的特征里 90% 都是 “苹果” 的语义这样 Encoder 就知道 “它” 指代 “苹果” 了。可见O_i 蕴含着该元素与序列中所有元素之间的关系即把上下文的信息融入当前元素之中借助 “上下文” 来理解这个元素从而捕捉序列内部的依赖关系。所以被称之为 “加权融合” 计算。通过 Self-Attention 机制我们可以找到一段文本中每一个词与其他所有词的相关度大小从而建模文本之间的依赖关系。​在代码中的实现self-attention 机制其实是通过给 Q、K、V 的输入传入同一个参数实现的attention(x, x, x)值得一提是的 Self-Attention 的并行计算特性如下图并行化处理序列数据避免了诸如 RNN 的时序依赖。Multi-Head AttentionSelf-AttentionSelf-Attention 也称为 “单头自注意力”其存在一个局限性 —— 一次只能拟合一种相关关系例如只能捕捉语法层面的关联或只能捕捉语义层面的关联但却无法一次拟合序列中可能存在的多种相关关系。因此Transformer 中使用的其实是 Multi-Head AttentionMHA多头自注意力Self-Attention 的一种变体。MHA 通过将 Q/K/V 矩阵切分为多个 Head 来解决了单头自注意力的局限性。例如在 NLP 任务中不同的 Attention Head注意力头能够拟合不同的相关关系有的头关注语法依赖比如 I love China 中love 与 I主语、China宾语的语法关联有的头关注语义指代比如前面的 it 与 animal 的指代关联有的头关注词序 / 标点比如逗号、连词前后的词关联。在 Transformer 论文中作者也通过实验证实了 MHA 每个不同的注意力头能够拟合语句中的不同信息。如下图所示自注意力红色和多头自注意绿色相比较显然后者更加关注全局信息。通过多个注意力头的同时计算能够更全面地拟合语句关系。不同 Head 对关注点不同。简而言之MHA 让模型能同时捕捉序列的多维度特征比单头自注意力的特征表达能力强得多。目前主流 LLM 的多头数量通常为 12/16/32如 GPT-3 是 96 头是大模型能力的重要指标。在具体的算法实现上MHA 的处理流程如下假设采用 h 个 Attention Head将原始的 Q/K/V 矩阵通过 3 个 Linear特征投影 拆分之后被拆分为 h 组独立的 Sub-Q/K/V。所以Sub-Q/K/V 的 d_k 维度等于 hidden_size / h。h 组 Sub-Q/K/V 分别做独立的 Self-Attention 计算缩放点积注意力得到 h 个独立的自注意力特征输出将 h 组的特征输出 Concat 拼接起来。最后再通过一个 Linear多头信息融合 维度还原做融合得到最终的多头注意力特征输出。这个 Linear 可以学习如何最好地组合来自不同 Head 的信息。Feed Forward Neural NetworkTransformer 由 2 大组件构成一块是 MHA另一块就是 FNNFeed Forward Neural Network前馈神经网络它们之间是互为补充的关系。MHA 关注 “全局元素关联建模”而 FFN 关注 “单一元素特征精修”。如果将 Transformer 比喻成拼图MHA 就是要找到许多快拼图之间的关系进行拼接但也只能做到简单的关联拼接。而 FFN 则是对单一块拼图进行 “放大、修图、加细节”让图片变得更清晰、更丰富让人们可以理解拼接起来之后的图画的实际含义。具体来说前面提到的 MHA多头自注意力、Layer Norm层归一化、Add残差连接都是纯线性变换。但是线性变换的任意组合依然是线性的而语言的语义是高度非线性的例如一词多义、复合概念、句法逻辑等等这些语义无法用简单的线性叠加来表示。所以 Transformer 架构中需要引入一种 “非线性变换”这就是 FFN 的意义 —— FFN 是整个 Transformer 架构中唯一提供非线性特征变换的核心组件。注语义是语言字、词、句所表达的内在含义是语言背后对应的客观事物、概念、逻辑关系等也是人们使用语言时想要传递的核心信息。语义关注的是 “想表达的是什么”而非单纯的文字形式或语言结构规则。简单说语义就是语言的 “实际内容”是脱离表面文字后能被人理解的核心信息。例如“苹果” 的语义在 “我吃了苹果” 中是一种可食用的水果在 “苹果手机” 中是一个数码品牌。词的语义真实含义由上下文决定。线性变换的数学本质是 yaxb特点是可加性例如 112。但这个现实世界的复杂规律是非线性的在语言领域这种非线性体现为以下几点复合概念“苹果 手机”MHA 能够线性的拼接为 “苹果手机”而 FFN 能够非线性的理解语义为 “一个品牌的手机”。句法逻辑“我 吃 苹果”MHA 能够线性的拼接为 “我吃苹果”而 FFN 能够非线性的理解语义为 “主谓宾的动作逻辑”。一词多义MHA 能够线性的拼接为 “吃苹果” 和 “苹果手机”而 FFN 能够非线性的理解 “苹果” 的语义为 “食物” 和 “手机品牌”。可见语言的语义是非线性的词的语义由上下文决定不同上下文的线性叠加会使得词出现非线性的变化。线性模型的表达能力远不足以描述真实的语义、物理等规律而 FFN 非线性变换就是 Transformer 能建模现实世界复杂规律的唯一方法突破了 MHA 纯线性表达的根本瓶颈能够适配语言的非线性语义本质。深度学习之所以能超越传统机器学习关键不是因为层数多而是因为 “多层线性变换 非线性激活” 的组合非线性变换是让多层网络变得有意义的唯一前提。反之如果没有非线性函数无论多少层全连接层本质上都可以化简成一层 MLP。1986 年Hinton 第一次在 MLP多层感知机中采用 Sigmoid 非线性激活函数并且结合 BP 算法进行多层网络反向传播。开启了深度学习的序幕。1989 年乔治·西本科提出万能近似定理Universal Approximation Theorem从理论证明了一个包含足够多隐含层神经元的、具有非线性变换激活函数的多层 FFN前馈网络能以任意精度逼近任意复杂度的函数。简而言之理论上神经网络可以解决世界上任何问题其表达力与图灵机等价。万能近似定理从根本上消除了对神经网络表达力的质疑。奠定了深度学习成功的理论基础。如上图所示Transformer 中的 FFN 由下列 3 部分组成。MHA 输入到 FFN 的 Input Tensor 形状为 [batch_size, seq_len, hidden_size] 且记 hidden_size d_model。如果 Linear 的升维系数 d_ffn 为 4Transformer 论文的经验值那么升维线性变换层LinearFFN W1 矩阵形状为 [hidden_size , 4*hidden_size]将每个 token 的 hidden_size 特征维度进行升维。在高维空间做更精细的非线性特征变换高维空间能表达更复杂的语义。非线性变换激活函数ReLU输入、输出矩阵均为 [batch_size, seq_len, 4*hidden_size]对每个 token 的特征向量进行非线性运算。降维线性变化层LinearW2 矩阵形状为 [4*hidden_size , hidden_size]将每个 token 的 hidden_size 特征维度进行降维浓缩特征。为了适配 Transformer 的残差连接必须保证输入输出维度一致。此外还可选的加入了一个 Dropout 层来防止过拟合。具体的代码实现如下classMLP(nn.Module):前馈神经网络def__init__(self,dim:int,hidden_dim:int,dropout:float):super().__init__()# 定义第一层线性变换从输入维度到隐藏维度self.w1nn.Linear(dim,hidden_dim,biasFalse)# 定义第二层线性变换从隐藏维度到输入维度self.w2nn.Linear(hidden_dim,dim,biasFalse)# 定义 dropout 层用于防止过拟合self.dropoutnn.Dropout(dropout)defforward(self,x):# 前向传播函数# 首先输入 x 通过第一层线性变换和 RELU 激活函数# 最后通过第二层线性变换和 dropout 层returnself.dropout(self.w2(F.relu(self.w1(x))))假设 Input Tensor 为 [3, 4, 8]那么W1 升维运算 [3, 4, 8] * [8, 84] [3, 4, 84]W2 降维运算 [3, 4, 84] * [84, 8] [3, 4, 8]W1 升维运算后每个 token 的特征向量被映射到更高的维度。最理想的情况当然是在无限维空间去展开所有的特征在无限维空间中任何复杂的非线性关系都可以被完美地展开、整理和分离。升维之后特征之间的关系可以被更充分地展开和重组让更加不可分的细微的特征结构也能暴露在 ReLU 面前。然后 token 向量中的每个元素都会经过 ReLU 激活函数的非线性变换比如负数置 0。ReLU 函数会将一部分元素置 0相当于做了一次筛选去掉不相关的信息。经过 ReLU 的筛选再次通过一个全连接网络降维回到之前的特征维度相当于把之前过滤的重要特征组合通过重组压缩回原来的维度。经过升维、非线性计算、降维之后token 向量的特征空间的局部表达能力就得到了很大的提升。ReLU 是 2017 年深度学习领域最主流的激活函数在效果、效率、工程实现上达到了最优平衡。后续 BERT、GPT 等 Transformer 变体将 ReLU 替换为 GELU它针对 NLP 任务的效果进行了优化。ReLU 的数学公式如下本质就是 “负数置 0”极其简单的计算使得其具有 3 个关键优势低损耗仅做了一次 “比较 取值”没有任何复杂的指数、对数、三角函数计算GPU 硬件可极致优化计算耗时几乎可以忽略。相较于 sigmoid、tanh 等需要 e^z 指数运算它们的计算量是 ReLU 的数倍且指数运算易产生数值精度问题如浮点数下溢 / 上溢。适用高维场景高维特征下ReLU 的低计算量优势会被放大若用 sigmoid、tanh高维逐元素的指数运算会让 FFN 计算效率大幅下降。缓解梯度消失ReLU 相比 sigmoid、tanh从激活函数层面大幅缓解了梯度消失问题配合 Transformer 的残差连接 层归一化让深层 Transformer 可以稳定训练。Residual Connection由于 Transformer 模型结构较复杂、层数较深​为了避免模型退化Transformer 采用了 Residual Connection残差连接的思想来连接每一个 Layer —— 下一层的输入不仅是上一层的输出还包括上一层的输入。每个 Layer 的输出与输入相加形成残差连接这有助于解决深层网络中的梯度消失问题。​例如在 Encoder 的第一个子层 MHA 中残差连接层即接受 MHA 的输出还接受 MHA 的原输入然后将 MHA 的输出和原输入相加再进行标准化。所以在 Transformer 架构中使用 “Add 和 2 条连线” 表示残差连接。Layer Normalization归一化核心是为了让不同层输入的取值范围或者分布能够比较一致。由于深度神经网络中每一层的输入都是上一层的输出因此多层传递下对网络中较高的层之前的所有神经层的参数变化会导致其输入的分布发生较大的改变。也就是说随着神经网络参数的更新各层的输出分布是不相同的且差异会随着网络深度的增大而增大。但是需要预测的条件分布始终是相同的从而也就造成了预测的误差。因此在深度神经网络中往往需要归一化操作将每一层的输入都归一化成标准正态分布。神经网络主流的归一化一般有两种批归一化Batch Norm层归一化Layer NormTransformer 架构采用的是效果更好的 Layer Norm相较于 Batch Norm 在每一层统计所有样本的均值和方差Layer Norm 在每个样本上计算其所有层的均值和方差从而使每个样本的分布达到稳定。Transformer 在每个 LayerSelf-Attention FFN之后都会进行层归一化操作这有助于模型的训练稳定性。classLayerNorm(nn.Module): Layer Norm 层def__init__(self,features,eps1e-6):super().__init__()# 线性矩阵做映射self.a_2nn.Parameter(torch.ones(features))self.b_2nn.Parameter(torch.zeros(features))self.epsepsdefforward(self,x):# 在统计每个样本所有维度的值求均值和方差meanx.mean(-1,keepdimTrue)# mean: [bsz, max_len, 1]stdx.std(-1,keepdimTrue)# std: [bsz, max_len, 1]# 注意这里也在最后一个维度发生了广播returnself.a_2*(x-mean)/(stdself.eps)self.b_2Decoder自回归生成模型Decoder 采用自回归Autoregressive的方式逐一生成任意长度输出序列的元素每一步的生成结果都送给下一步作为输入如下图所示。通过输入开始标识符Begin来启动 Decoder 输出。进入 MMHA 模块进行 Self-Attention自注意力计算输入序列为 Output Sequence。进入 MHA 模块进行 Cross-Attention交叉注意力计算把 Decoder 的输入序列Q和 Encoder 的输出结果KV进行融合转换。输出一个概率分布Output Probabilities表示词汇表中每个单词作为下一个输出单词的概率并依据某种策略输出一个最可能的单词。假设生成第一个单词为 I进入自回归把生成的第一个单词 “I” 作为 Decoder 的输入Outputs再次进入 Decoder 解码生成出第二个单词 ate。重复上述步骤直到遇见停止符End。Multi-Head AttentionCross-AttentionDecoder 作为序列的生成器继承了 Bahdanau Attention 的 “源-目标词语义对齐” 思想。使得 Decoder 生成每个目标词时都可以精准对齐源语言的相关词。MHA 使用 Encoder 的输出作为 Key 和 Value接受源语言特征。MHA 使用 MMHA 的输出作为 Query接受目标语言特征。Masked Multi-Head AttentionSelf-AttentionDecoder 使用 Masked Multi-Head AttentionMMHA掩码多头自注意力的目的是为了符合语言生成的时序性只能从左到右生成不能提前看到后面的词。因为 Transformer 采用的是 CLM 预训练任务根据上文生成下一个词所以 Decoder MMHA 对目标语言序列的自注意力做掩码处理屏蔽当前词之后的所有词让模型只能基于前面的词生成下一个词。掩码Mask的作用是遮蔽一些特定位置的词。具体而言掩码会被应用到点积的结果上通常是将某些位置的值设置为一个非常大的负数如负无穷这样在应用 Softmax 时这些位置的权重会接近于零从而忽略对某些位置的学习。例如待学习的文本序列是 the cake was sour进行 Masked 处理之后在每一行输入中Decoder 仍然是只看到前面的词并预测下一个词。观察上述的例子可以发现 Masked Attention 其实就是一个和输入序列等长的上三角矩阵。可以通过创建一个与输入序列等长的上三角矩阵作为注意力掩码再使用掩码来遮蔽掉输入序列的矩阵即可。也就是说当 Input Tensor 维度为 [batch_size, seq_len, hidden_size] 时Mask 矩阵的维度一般为 [1, seq_len, seq_len]。多层注意力的运算效果随着 layer 层数叠加可以看到注意力运算愈发集中都某几个词黄色表示注意力得分高相似度高。这是一个 Mask 的注意力得分因此上三角为得分为 0。一个词对自己的关注度是最高的因此注意力得分的黄色集中在对角线上。一个词和周围上下文的关注度也很大因此多层运算中会慢慢出现 Localized Attention 的三角格式。经过多层的 Attention 运算之后注意力得分集中在相似度最高的词之间。Linear Softmax最终输出 token 之前还需要进行一次概率分布从词汇表中选出最终真正输出的 token。步骤如下线性层Linear将 Output token embedding 向量乘以词汇表矩阵得到一个长度为词汇表大小的概率分布向量。表示每一个 token 作为真正输出的概率大小。例如[1, embed_size] * [embed_sizevocab_size ] [1, vocab_size]。Softmax将概率分布向量进行归一化计算得到最终的 % 概率值Output Probabilities。LayerTransformer 由 Encoder 和 Decoder 组成而 Encoder 和 Decoder 都会由多个相同的 Layer 堆叠而成。Input Sequence 进入 Encoder 进行编码到 Encoder Layer 的最顶层再将编码结果K、V输出给 Decoder Layer 的每一层通过 Decoder 解码后就可以得到 Output Sequence。Layer 结构使得每层 MHA 计算得到的中间结果还可以在下一层继续做 MHA这样一层层做下去每前进一层实质都是在更大的范围内交融序列中的各个元素、解析其间的关系这样的结构使得 Transformer 能够捕获多种不同的相关关系。最终就能够得到表示着整个序列的总体特征的向量了。Attention 的优化方向模型算法层面针对 Attention 的优化主要有 2 个方向一方面利用注意力得分矩阵的稀疏性选择部分 token 进行注意力运算从而减少运算次数另一方面从压缩 KV 矩阵的方向出发共享 KV 矩阵或者压缩 KV 矩阵的维度。工程层面主要是面对训练和推理的场景一方面为了根据 GPU 的存储架构做的软件适配比如 FlashAttention、Paged Attention。另一方面是根据模型结构做的框架层的优化比如 KV Cache、Radix Attention。Sparse Attention稀疏注意力假设 seq_len 是 n那么 Self-Attention 的 QK 计算就会产生一个 形状为 [n, n] 的注意力得分矩阵相似度矩阵。所以从理论上来讲Self Attention 的计算时间和显存占用量都是 O(n^2)。也就是说围绕该矩阵的计算量和显存占用量Scaled、Masked、Softmax、乘 V 加权融合计算等等会随 n 呈平方级增长例如如果 seq_len 变成原来的 2 倍显存占用量就是原来的 4 倍计算时间也是原来的 4 倍。但实际上Child 等人2019的研究发现在训练好的 Transformer 模型中注意力矩阵往往是稀疏的这意味着并不是每个 token 都需要关注其他所有 token每个 token 只关注非常有限个其他 token。有些 token 之间的相互作用可能对最终的输出贡献不大可以被忽略。稀疏注意力机制的核心思想是在自注意力计算中引入稀疏性即不是让序列中的每个位置都与其他所有位置进行注意力计算而是仅选择部分位置进行计算。所以稀疏注意力具有以下优势减少计算量通过减少参与注意力计算的位置数稀疏注意力显著降低了计算复杂度使得模型能够处理更长的序列。减少显存占用量稀疏操作减少了需要存储的注意力权重的数量从而降低了模型的内存需求。提高长距离依赖学习能力某些稀疏模式如分层或跳跃连接可以帮助模型更有效地学习序列中的长距离依赖关系。下图是 Self-Attention 的一个注意力矩阵。左边显示了注意力矩阵右边显示了关联性这表明每个元素都跟序列内所有元素有关联。Atrous Self Attention空洞注意力启发于 “膨胀卷积Atrous Convolution”它对相关性进行了约束强行要求每个元素只跟它相对距离为 k, 2k, 3k 的元素关联其中 k1 是超参数。如此的运行效率和显存占用都变成了 O(n^2/k) 也就是说能直接降低到原来的 1/k。Local Self Attention局部自注意力约束每个元素只与前后 k 个元素以及自身有关联。保留了一个 2k1 大小的窗口每个元素只跟 2k1 个元素算相关性这样一来理想情况下运行效率和显存占用都变成了 O(kn)也就是说随着 n 而线性增长非指数增长。这是一个很理想的性质当然也直接牺牲了长程关联性。Sparse Self Attention稀疏自注意力将 Atrous Self Attention 和 Local Self Attention 合并为一个除了相对距离不超过 k 的、相对距离为 k,2k,3k,… 的注意力都设为 0这样一来 Attention 就具有了 “局部紧密相关和远程稀疏相关” 的特性。多头注意力优化GQA、MQA 和 MLAFlash AttentionStanford DAWN Lab 实验室提出https://arxiv.org/pdf/2205.14135。将部分内容存储到 SRAM 中以减少访问全局内存 HBM 的频率。Page AttentionPaged AttentionvLLM 提出https://arxiv.org/pdf/2309.06180。利用分页技术减少显存碎片化问题。Radix AttentionRadix AttentionSGLang 提出https://arxiv.org/pdf/2312.07104。提高缓存命中率。参考文档https://medium.com/vmirly/tutorial-on-scaled-dot-product-attention-with-pytorch-implementation-from-scratch-66ed898bf817https://poloclub.github.io/transformer-explainer/