Transformer

注意力机制#

RNN 中的注意力机制由于串行架构的存在还是计算复杂度过高。因此我们对 RNN 的注意力计算公式进行改进，引入 Query、Key、Value 的概念。其中 Query 为查询信息，Key 为知识库，Value 为知识库中每个知识的释义。形式化表示为：

w_{ij}=\alpha{(q_{i-1},k_j)} \\ c_i=\sum_{j=1}^{t}w_{ij}v_j

即通过评分函数 $\alpha(\cdot)$ 计算 Query 对每一个 Key 的相似度，然后将相似度矩阵与对应的 Value 聚合起来得到上下文向量，表示查询中的哪些部分是对输入是有影响的。

用查字典来举例：我们需要查询的信息为 “中国”，计算完对字典中每个词的相似度之后，得到相似度最大的两个词 “华夏”、“唐宋”。然后将这两个词对应的信息聚合起来，来代表 “中国” 这个词，即这两个词对 “中国” 的影响是最大的。

Transformer#

推荐查看 - https://jalammar.github.io/illustrated-transformer/ ↗

缩放点积注意力#

自注意力(Self-Attention) 指的是计算同一段序列中不同令牌/词元之间的注意力，通过该机制可以对序列中的每一个令牌进行上下文建模。

缩放点积注意力

上图表示了如何计算第一个令牌与其余令牌的注意力，并将其汇总为上下文向量。这里有几点需要注意：

缩放参数 $\small \sqrt{d_k}$ 是防止 softmax 梯度饱和，且 softmax 是在注意力分数矩阵的列维度上进行计算；
$\small Q$ 、 $\small K$ 、 $\small V$ 矩阵是同一段序列通过不同的权重矩阵 $\small W^q$ 、 $\small W^k$ 、 $\small W^v$ 加权得到；
每一个 $q$ 、 $k$ 、 $v$ 、 $y$ 的维度都是 $d$ ，即单个令牌的词向量维度。

多头注意力#

单头无法表征序列的所有信息，通过引入不同的头在同一时刻关注序列中的不同部分，从而学习到不同的上下文信息。类似 CNN 中的通道。

多头注意力

上图表示了如何计算第二个头所对应的注意力、并将其汇总得到第二个头的上下文向量。然后将所有头的上下文向量拼接起来，这里需要注意拼接之后的向量是高维向量，需要压缩使得输入和输出维度保持不变。

实践中需要设计的超参数为头的个数 $h$ ，每个头对应的 $q$ 、 $k$ 、 $v$ 、 $y$ 的维度都是 $\frac{d}{h}$ ；

逐位前馈神经网络#

将注意力层输出的特征进一步进行非线性变换，提升模型的表达能力。逐位指的是每个词元的特征重计算都是独立的。在 Transformer 架构中，这一层的参数量是自注意力层参数量的两倍。

逐位前馈神经网络

上图右侧的计算表示 “深” 这个特征通过权重矩阵加权之后得到另一个更好的特征。这里使用的是两层 FFN，且权重矩阵 $\small W_1$ 和 $\small W_2$ 是共享的，这就类似于 RNN 的思想，因此就能够很好的处理变长序列。

绝对位置编码#

自注意力机制具有置换不变形，无法捕捉序列中的顺序关系。因此引入位置编码对序列信息进行位置建模。

位置编码

原始的 Transformer 使用的是基于正余弦 的绝对位置编码(Sinusoid) - 即根据正余弦函数来生成对应位置的唯一位置嵌入，并将其加到令牌的词向量上。对于维度大小为 $d$ 的词向量，它的编码方式为：

\small \begin{cases} p_{k, 2i} = \sin\left(k/10000^{2i / d}\right) \\ p_{k, 2i+1} = \cos\left(k/10000^{2i / d}\right) \end{cases}

其中 $\small p_{k,2i}$ 和 $\small p_{k,2i+1}$ 是位置 $\small k$ 的词向量的第 $\small 2i$ 和 $\small 2i+1$ 个分量、 $\small d$ 是词向量的维度。将该结果加到词向量对应位置上形成最终的模型输入。

作者也使用了可学习的位置编码，发现使用 Sinusoid 编码可以让模型拥有更好的长度外推能力。

编码器 / 解码器#

原始 Transformer 模型分别堆叠了六层编码器(Encoder)和解码器(Decoder)。且每层输入和输出的维度都是 (batch, time, dimension)，即批量大小、序列长度、词向量维度。

解码器具有自回归属性，也叫 Next-token prediction(NTP) 或 rollout。具体步骤为：

根据输入特征，获得词表上的概率分布，然后采样得到最终输出词元；
将输出词元添加到输入中继续第一步，直到输出 <eos>。

在训练阶段，使用自回归模型会有弊端。因为提供了输入序列和目标序列，使得模型在训练的时候 “作弊”，即参考目标序列直接获得下一个单词。因此提出掩码多头注意力。即通过对注意力分数矩阵进行遮掩，使得模型在训练阶段只能参考历史信息。具体来说通过下述规则进行遮掩(下三角遮掩)：

\small a_{ij}=\begin{cases} 0 & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases}

这样，在使用 softmax 计算上下文向量的时候，由于出现了极大的负值，导致对应位置的指数趋近于 0，从而屏蔽了未来信息，保持了解码器的自回归属性。

掩码多头注意力机制

为了保证输入序列长度一致，会在输入向量中添加 [pad] 来对齐长度。而在计算注意力分数矩阵的时候需要忽略这些信息。

并且引入交叉注意力(Cross Attention) 处理来自编码器的 K/V 和来自解码器的 Q。需要注意传入的 K/V 是编码器最后一层的输出，且传入到解码器的每一层当中。但是解码器的 Q 是当前层计算的结果。

交叉注意力

激活函数#

GeLU 全称为 Gaussian Error Linear Unit，它利用高斯分布的累积分布函数来调节输入信息的激活程度，使得中间区域的输入梯度更为平滑，起到正则化的作用。

\small \begin{aligned} \text{GeLU}(x)&=x\cdot \Phi(x) \\ &= x\cdot \frac{1}{2}\left(1+\text{erf}(\frac{x}{\sqrt{2}})\right) \end{aligned}

其中 $\small \Phi(x)$ 是累积分布函数、 $\small \text{erf}(\cdot)$ 表示误差函数。

SwiGLU 全称为 Swish-Gated Linear Unit，它将 Swish 的平滑与 GLU 的门控机制结合起来使得梯度能够有效传播，目前是 LLM 使用的主流激活函数。

\small \text{SwiGLU}(x)=\text{Swish}(xW_1+b_1)\otimes \sigma(xW_2+b_2)

注意力模块的优化#

自注意力模块是二次复杂度 $\small O(n^{2})$ ，因为它要对序列中的任意两个词向量都要计算相似度，得到一个 $\small n^2$ 大小的注意力分数矩阵。

Sparse Attention#

从注意力分数矩阵的角度来说，就是除了相对距离不超过 $\small k$ 的、相对距离为 $\small k$ 的倍数的注意力分数都设置为 0。这样注意力就具有了局部紧密相关和远程稀疏相关的特性。

该方法的不足之处就是需要人工选择保留的注意力区域，不利于扩展。但是由于将每个词元的注意力压缩在了较小的空间中(每个词元只能看到训练长度的词元)，能一定程度上缓解长度外推问题。

Linear Attention#

线形注意力

制约注意力性能的关键因素是 softmax 函数，如果没有该函数，去掉缩放系数的注意力公式实际上就是三个矩阵连乘，复杂度仅有 $\small O(n)$ 。一个自然的想法就是拿掉该函数，并用一般函数 $\small \text{sim}(\cdot)$ 进行替代：

\small \text{attention}(q_i,k_j,v_j)=\frac{\sum_{j=1}^{n}\text{sim}(q_i,k_j)v_j}{\sum_{j^\prime=1}^{n}\text{sim}(q_i,k_{j^\prime})}

为了保留注意力相似的分布特征，需要保证 $\small \text{sim}(q_i,k_j) \geq0$ 恒成立。

线形注意力就是用核函数代替这个一般函数，即 $\small \text{sim}(q_i, k_j) = \phi(q_i)^{\text{T}}\phi(k_j)$ 。此时注意力公式为：

\small \text{attention}(q_i, k_j, v_j) = \phi(q_i)\cdot \left[\phi(k_j)^{\text{T}}v_j\right]

其中核函数选择 $\small\phi(x)=\text{elu}(x)+1$ 。

KV Cache#

自回归模型中的 Causal Decoder 在 token-by-token 递归生成的时候，每个生成的 Q 都需要与历史时刻的 K/V 进行注意力计算。KV Cache 指的是缓存之前的前缀，仅计算当前时刻的 K/V，然后与之前的结果连接起来即可。

但是 KV Cache 通常保存在 GPU 的 HBM 中，并且序列长度越长，缓存大小呈现线形增长，导致带宽瓶颈。后续的 MQA、GQA、MLA 都是围绕如何减少 KV Cache 的同时实现在更少的设备上推理更长的上下文，或者在相同的上下文长度下让推理的批量大小更大，从而实现更快的推理速度或者更大的吞吐总量。

对比示例

GQA / MQA#

Group-Query Attention(GQA) 提出将Attention Head 分成 $\small g$ 个组，每组共享同一个 K/V，它将 KV Cache 压缩 $\small g/h$ 。当 $\small g=1$ 的时候，就是 Multi-Query Attention(MQA)。

其中 GQA 在 Llama2/3、DeepSeek-V1、ChatGLM2/3 等模型中使用。MQA 在 PaLM、StarCode、Gemini 等模型中使用。

MLA#

工程文章 - https://zhuanlan.zhihu.com/p/714761319 ↗

在推理的过程中，MHA 会为每个令牌缓存 $\small 2dn_hd_h$ 个元素，而在模型部署阶段，大量的 KV Cache 会导致能同时处理的 Batch 变小和支持的最大序列长度变短；

将 Key、Value 进行低秩联合压缩，其中 $\small d$ 为嵌入维度、 $\small n_h$ 为注意力头的数量、 $\small d_h$ 为每个注意力头的维度、 $\small c_t^{KV} \in \mathbb{R}^{d_c}$ 是 KV 的压缩潜在向量、 $\small d_c \ll (d_hn_h)$ 表示 KV 的压缩维度、 $\small W^{DKV} \in \mathbb{R}^{d_c \times d}$ 是下投影矩阵、 $\small W^{UK}$ ， $\small W^{UV} \in \mathbb{R}^{d_hn_h \times d_c}$ 分别是 Key 和 Value 的上投影矩阵。

\small \begin{aligned} c_t^{KV}&=W^{DKV}h_t \\ k_t^{C}&=W^{UK}c_t^{KV} \\ v_t^{C}&=W^{UV}c_t^{KV} \end{aligned}

在推理的过程中，MLA 方法只需要缓存 $\small c_t^{KV}$ ，只有 $\small d_cl$ 个元素，其中 $\small l$ 表示层数。同时 $\small W^{UK}$ 和 $\small W^{UV}$ 可以合并到 $\small W^{Q}$ 和 $\small W^{O}$ 中，因此不需要在注意力计算中真正获得 Key 和 Value。

为了在训练过程中进一步减少激活内存，还可以对 Query 进行低秩压缩：

\small \begin{aligned} c_t^{Q}&=W^{DQ}h_t \\ q_t^{C}&=W^{UQ}c_t^{Q} \end{aligned}

其中 $\small c_t^{Q} \in \mathbb{R}^{d_c^{\prime}}$ 是 Query 的压缩潜在向量、 $\small d_c^{\prime} \ll (d_hn_h)$ 表示 Query 的压缩维度、 $\small W^{DQ} \in \mathbb{R}^{d_c^{\prime} \times d}$ 是下投影矩阵、 $\small W^{UQ} \in \mathbb{R}^{d_hn_h \times d_c^{\prime}}$ 是上投影矩阵。

但是 RoPE 与低秩的 KV Cache 并不兼容，回忆一下 RoPE 的公式：

\small (\mathcal{R}_m q)^{\text{T}}(\mathcal{R}_n k) = q^{\text{T}}\mathcal{R}_m^{\text{T}}\mathcal{R}_n k = q^{\text{T}}\mathcal{R}_{n-m} k

如果 Query 和 Key 中间有一个跟位置相关的矩阵，那么 $W^{UK}\mathcal{R}_{m-n}W^Q$ 就不能合并成一个固定的投影矩阵。MLA 给出的解决办法是用额外的多头 Query 和共享的 Key 去结合 RoPE：

\small \begin{gathered}\left[q_{t,1}^R;q_{t,2}^R;\ldots;q_{t,n_h}^R\right]=q_t^R=\mathrm{RoPE}(W^{QR}c_t^Q)\\k_t^R=\operatorname{RoPE}(W^{KR}h_t)\\q_{t,i}=[q_{t,i}^C;q_{t,i}^R]\\k_{t,i}=[k_{t,i}^C;k_t^R]\\o_{t,i}=\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^\text{T}k_{j,i}}{\sqrt{d_h+d_h^R}}\right)v_{j,i}^C\\u_t=W^O[o_{t,1};o_{t,2};\ldots;o_{t,n_h}]\end{gathered}

其中 $\small d_n^{R}$ 是单个头的维度， $\small W^{QR} \in \mathbb{R}^{d_hn_h \times d^\prime_c}$ ， $\small W^{KR} \in \mathbb{R}^{d_h \times d}$ 是用来生成额外的解耦多头 Query 和共享 Key 的矩阵，推理的时候需要把额外的共享 Key 也缓存了，所以总共需要缓存的元素为 $\small d_c + d_n^{R}l$ 个。

mlaa

Flash Attention#

原始的注意力计算过程中，最大的中间过程结果就是注意力分数矩阵，而最终的上下文向量反而很小。反复读取大尺寸注意力矩阵会导致 HBM 带宽瓶颈。

Flash Attention 通过 矩阵分块 和 算子融合 等方法，将中间计算结果保留在大带宽的 SRAM 中，获得最终结果再写回小带宽的 HBM 中，从而避免了带宽瓶颈。

Flash Attention

上图左侧是原始数值稳定版本的 softmax 实现，右侧是将输入分块后，分别计算 softmax 并融合的实现。具体来说对于 Key/Value 的每一块，计算它们和所有 Query 的注意力矩阵，并保存在 SRAM 中。当计算得到最终上下文向量的时候在将结果写回 HMB 中。

HBM IO 复杂度对比：

Vanilla Attention - $\small O(Nd+N^{2})$ ；
Flash Attention - $\small O(N^2d^2/m) \ll O(N^2) \ll O(Nd+N^2)$ 。

算子融合指的是通过将上图右侧的整个计算过程融合成一个高效的算子，从而减少中间数据的存储。

Page Attention#

Ring Attention#

Ring Attention

混合专家模型(MoE)#

核心思想 - 用来替代 FNN 层，且每个专家在训练的过程中学习不同的信息，然后在推理的过程中只有与特定任务最相关的专家被使用。

混合专家模型

首先通过门控函数 / 路由网络计算出每个专家的权重，该权重表示每个专家对当前词元的关注程度。通常是一个简单的神经网络，将输入 $\mathbf{u}_t^{l}$ 通过 $\text{softmax}$ 映射为概率分布：

s_{i,t}=\text{softmax}_i(\mathbf{u}_t^{l}\mathbf{e}_i^{l})

其中 $\mathbf{e}_i^{l}$ 表示第 $l$ 层第 $i$ 个专家的可学习权重，最后的结果 $s_{i,t}$ 表示第 $i$ 个专家对第 $t$ 个词元的关注程度。

然后从 N 个专家中保留 top-k 个最高权重的专家参与计算，其余权重置零(稀疏激活)：

\small g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}\in\text{Topk}(\{s_{j,t}|1\leqslant j\leqslant N\},K) \\0,&\text{otherwise}&&\end{cases}

最后将路由权重与所有专家的输出进行聚合并加上残差连接作为最终的输出：

\small \mathbf{h}_t^l=\sum_{i=1}^N\left[g_{i,t} \cdot \text{FFN}_i(\mathbf{u}_t^l)\right]+\mathbf{u}_t^l

传统的 MoE 存在以下问题：

单个专家涵盖多种知识，无法充分同时利用；
多个专家涵盖通用的冗余知识，阻碍了 MoE 的理论性能上限。

DeepSeekMoE 通过 划分更细粒度的专家 以及 利用单个专家存储通用知识 来达到单个专家的高度专业化。形式化表达为：

\small \begin{aligned}&\mathbf{h}_t^l=\textcolor{blue}{\sum_{i=1}^{K_s}\text{FFN}_i\left(\mathbf{u}_t^l\right)}+\textcolor{blue}{\sum_{i=K_s+1}^{mN}}\left[g_{i,t}\cdot\text{FFN}_i(\mathbf{u}_t^l)\right]+\mathbf{u}_t^l,\\&g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}\in\text{Topk}(\{s_{j,t}|K_s+1\leqslant j\leqslant mN\},mK-K_s),\\0,&\text{otherwise},&&\end{cases}\end{aligned}

具体来说，将 $\small N$ 个专家划分为 $\small mN$ 个专家，即将隐藏层的维度乘上 $\small 1/m$ ，同时选取 $\small mK-K_s$ 个专家保证参数量与传统 MoE 模型相当。然后隔离出单个专家 $\small K_s$ 来存储通用知识。

归一化#

Post / Pre-Norm#

Post-Norm - 归一化放置在残差计算之后：

\small \text{Post-Norm}(x)=\text{Norm}(x+\text{Sublayer(x)})

训练时间较快、网络容易调优、输出层梯度较大引发的训练不稳定。

Pre-Norm - 归一化放置在子层计算中：

\small \text{Pre-Norm}(x)=x+\text{Sublayer}(\text{Norm}(x))

训练稳定、效果不如 Post-Norm、缓解梯度消失(爆炸)问题，现代 LLM 首选。

RMS / Layer-Norm#

对比图

位置编码#

绝对位置编码#

假设输入向量由词向量 $\small E_{x_i}$ 和位置向量 $\small U_i$ 组成，对注意力分数计算公式 $\small q_i^{\text{T}}k_j$ 进行因式分解得到：

\small A_{i,j}^{\text{abs}}=\underbrace{E_{x_{i}}^{\text{T}}W_q^{\text{T}}\textcolor{green}{W_k}E_{x_j}}_{\text{内容相关项}}+\underbrace{E_{x_{i}}^{\text{T}}W_q^{\text{T}}\textcolor{green}{W_{k}}\textcolor{blue}{U_{j}}}_{\text{内容-位置耦合项}}+\underbrace{\textcolor{red}{U_{i}^{\text{T}}W_q^{\text{T}}}\textcolor{green}{W_k}E_{x_j}}_{\text{位置-内容耦合项}}+\underbrace{\textcolor{red}{U_{i}^{\text{T}}W_{q}^{\text{T}}}\textcolor{green}{W_{k}}\textcolor{blue}{U_{j}}}_{\text{位置相关项}}

可以发现只有最后一项包含 $\small U_i$ 和 $\small U_j$ ，理论上包含相对位置 $\small(j-i)$ 的信息，但因为 $\small W_q$ 和 $\small W_k$ 的非线形变换，导致相对位置信息被破坏了。

相对位置编码#

\small A_{i,j}^{\text{rel}}=E_{x_{i}}^{\text{T}}W_q^{\text{T}}\textcolor{green}{W_{k,E}}E_{x_j}+E_{x_{i}}^{\text{T}}W_q^{\text{T}}\textcolor{green}{W_{k,R}}\textcolor{blue}{R_{i-j}}+\textcolor{red}{u^{\text{T}}}\textcolor{green}{W_{k,E}}E_{x_j}+\textcolor{red}{v^{\text{T}}}\textcolor{green}{W_{k,R}}\textcolor{blue}{R_{i-j}}

其中 $\small R$ 为相对位置编码矩阵、 $\small u$ 和 $\small v$ 都是所有层共享的可学习参数。并且由于 $\small R_{i-j}$ 的编码空间与输入空间 $\small E$ 不一定相同，所以将 $\small W_k$ 矩阵替换为了 $\small W_{k,E}$ 和 $\small W_{k,R}$ 。

T5 位置编码#

如果认为输入信息和位置信息是应该独立/解耦的，那么内容-位置耦合项和位置-内容耦合项应该被移除，并且加上一个可训练的偏置项：

\small A_{i,j}^{\text{rel}}=E_{x_{i}}^{\text{T}}W_{q}^{\text{T}}W_{k,E}E_{x_{j}}+r_{i-j}

其中可学习参数 $r$ 在所有层都是共享的，并且不同的注意力头使用不同的位置编码。

ALiBi#

将注意力分数的计算改为如下格式：

\small q_i^{\text{T}}k_j-\lambda(i-j)

其中 $\small i-j$ 是 Query 和 Key 之间的位置偏移、 $\small \lambda=2^{(-8i/k)}$ 是每个注意力头 $\small k$ 的惩罚系数。当 Query 和 Key 的位置偏移越大，它对注意力分数的惩罚就越大。该方式具有优秀的长度外推能力。

RoPE#

为 Query 和 Key 设置了单独的旋转矩阵 $\small \mathcal{R}$ 。考虑一个二维的情况：

\small \mathcal{R}_mq= \begin{pmatrix} \cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} \quad

由于内积满足线形叠加性，因此任意偶数维的 RoPE，都可以表示为二维形式的拼接：

\small {\underbrace{\begin{pmatrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \cdots & 0 & 0 \\ 0 & 0 & \sin m\theta_1 & \cos m\theta_1 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \cdots & \cos m\theta_{d/2-1} & -\sin m\theta_{d/2-1} \\ 0 & 0 & 0 & 0 & \cdots & \sin m\theta_{d/2-1} & \cos m\theta_{d/2-1} \\ \end{pmatrix}}_{\mathcal{R}_m} \begin{pmatrix}q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1}\end{pmatrix}}

也就是说，给位置 $\small m$ 的向量 $\small q$ 乘上矩阵 $\small \mathcal{R}_m$ 、位置 $\small n$ 的向量 $\small k$ 乘上矩阵 $\small \mathcal{R}_n$ ，并用变换后的序列做注意力，那么就自动包含相对位置信息了，因为如下恒等式成立：

\small (\mathcal{R}_m q)^{\text{T}}(\mathcal{R}_n k) = q^{\text{T}} \mathcal{R}_m^{\text{T}}\mathcal{R}_n k = q^{\text{T}} \mathcal{R}_{n-m} k

由于 $\small \mathcal{R}_m$ 的稀疏性，直接使用矩阵乘法来计算会很浪费算力。因此使用下述方法来实现 RoPE：

\small \begin{pmatrix}q_0 \\ q_1 \\ q_2 \\ q_3 \\ \vdots \\ q_{d-2} \\ q_{d-1} \end{pmatrix}\otimes\begin{pmatrix}\cos m\theta_0 \\ \cos m\theta_0 \\ \cos m\theta_1 \\ \cos m\theta_1 \\ \vdots \\ \cos m\theta_{d/2-1} \\ \cos m\theta_{d/2-1} \end{pmatrix} + \begin{pmatrix}-q_1 \\ q_0 \\ -q_3 \\ q_2 \\ \vdots \\ -q_{d-1} \\ q_{d-2} \end{pmatrix}\otimes\begin{pmatrix}\sin m\theta_0 \\ \sin m\theta_0 \\ \sin m\theta_1 \\ \sin m\theta_1 \\ \vdots \\ \sin m\theta_{d/2-1} \\ \sin m\theta_{d/2-1} \end{pmatrix}

其中 $\otimes$ 是逐位相乘。在 $\small \theta_i$ 的选择上，沿用了和 Sinusoidal 位置编码的方案，即 $\small \theta_i = 10000^{-2i/d}$ 。

RoPE 不带有显式的远程衰减，通过不同频率的三角函数有效区分了长程和短程。并且直接作用于 Q/K，不改变注意力计算的形式，与 Flash Attention 更为契合，容易 Scale Up。

长度外推优化#

长度外推问题指的是模型在推理阶段无法处理比训练阶段更长的输入序列的现象(Train Short, Test Long)，从而导致模型无法捕捉全局的上下文信息，造成信息丢失或模糊的建模结果。

推理的时候用到了没训练过的位置编码(绝对/相对)；
推理的时候注意力机制处理的词元数量远超训练时的数量。

注意力机制理论上可以处理任意长序列，但是越多的词元去平均注意力就会导致注意力分布越均匀，也就不能很好地表征词元之间的关系。

Position Interpolation#

NTK-aware#

YaRN#

解码策略#

贪婪采样的改进策略#

贪心搜索(Greedy Search) - 在每个采样步骤中选择概率最高的词元作为下一个词元。但是该方法容易陷入局部最优，生成重复、不自然的句子；
束搜索(Beam Search) - 在每个采样步骤中选取概率最高的 $k$ 个句子，并最终选取整体概率最高的生成回复。其中 $k$ 称为束宽度。详细可以查看循环神经网络 / 束搜索 ↗ 中的部分。

各类解码参数：

长度惩罚(Length Penalty) - 通过将句子概率除以其长度的指数幂 $\alpha$ ，缓解束搜索倾向于生成较短的句子(每生成一个单词，都会乘以一个小于 1 的概率，使得句子的总体概率逐渐变小)；
出现惩罚(Presence Penalty) - 将生成词元的 logits 减去惩罚项 $\alpha$ 来降低该词元之后出现的概率；
频率惩罚(Frequency Penalty) - 将生成词元的 logits 减去其出现次数乘以惩罚项 $\alpha$ 来降低该词元之后出现的概率。

随机采样的改进策略#

温度采样(Temperature Sampling) - 通过调整 logits 的温度系数，从而保证采样过程的随机性，其越小，概率分布越极端，反之越平坦；
Top-k 采样(Top-k Sampling) - 从概率最高的 $k$ 个词元中进行采样，但是不考虑整体概率分布，无法适应不同的上下文语境；
Top-p 采样(Top-p Sampling) - 从一个符合特定概率条件的最小词元集合中进行采样，要求其中包含的所有词元的累积概率大于或等于预设阈值 $p$ 。

语言模型#

预训练模型确立了预训练-微调的范式。通过大量无标注文本建立模型的基础能力，然后通过有标注数据进行下游任务的微调。

GPT / Decoder#

在大规模语料 $\small u=\{u_1, \ldots, u_n\}$ 上进行无监督预训练，根据前 $\small k$ 个词来预测下一个词，最大化输出词元的似然：

\small L_{u}=\sum\nolimits_{i}\log P(u_i|u_{i-k},\ldots,u_{i-1})

然后取最后一个词元的向量，接一个任务专用的线形层，完成各类任务的微调，并加入预训练的目标函数，避免模型遗忘通用知识：

\small L_s=\sum\nolimits_{x,y}\log P(y|x_1,\ldots,x_m), \hspace{0.2cm} L_{mtl}=L_s+\lambda L_u

GPT-2 发现通过海量数据训练出来的模型在不做微调的情况下，通过在 Prompt 里加入自然语言指令，就能完成下游任务，即零样本(Zero-Shot)的学习能力。而 GPT-3 发现在 Prompt 中给定少量示例就能让模型适应新任务(In-Context Learning)，即少样本(Few-Shot)的学习能力。

BERT / Encoder#

传统的基于自回归架构的 LM 只能单向建模，而 BERT 通过双向建模，使得模型能够聚合左右上下文信息，从而提高下游任务的能力。

模型的输入嵌入由三个部分构成，词元嵌入、位置嵌入和段嵌入(Segment Embedding)，其中句子嵌入是用来区分不同的句子，例如 B 是否是 A 的下文等。并且每个输入序列的第一个标记是 [CLS]，用作当前句子的聚合表征。用 [SEP] 来区分句子，即 [CLS] 句子A [SEP] 句子B [SEP]。

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

encoded_input = tokenizer("我是一句话")

模型的预训练目标由两个部分构成：

Masked LM(MLM) - 随机遮掩输入中 15% 的子词。但是微调的时候不会出现 [MASK]，就会导致训练推理失配问题，可以采用 Scheduled Sampling 解决。即将遮掩词元的 80% 替换为 [MASK]、10% 替换为随机词、10% 保持原词。然后让模型预测遮掩位置的原始词元，损失为所有遮掩位置的交叉熵；
Next Sentence Prediction(NSP) - 随机从语料库中采样句子对，其中 50% 为相邻句子、50% 为不相邻句子。最后取 [CLS] 表示的嵌入向量输入到分类层中，用来表明输入句子对是否相邻。损失为交叉熵。

后续的改进包括 RoBERTa，其去掉了 NSP 任务，引入了动态遮码的机制，即每个回合输入句子的 mask 位置都不同，使得模型适应不同的遮码策略，学习不同的表征。还有 ALBERT，将输入嵌入 V x H 转换为两个较小的矩阵 V x E + E x H，从而大幅减少内存消耗和训练速度。

T5#

将所有的自然语言处理任务都转换成文本-文本的形式，并用一个统一的模型。其输入是带有任务前缀的文本序列，输出是对应任务的结果。

模型结构 - 编码器(解码器) 各 6 层、512 词向量维度、8 个头、32,000 词表大小(SP)、512 tokens 上下文长度；
训练范式 - Span Corruption(随机遮掩连续 span，预测被遮掩的文本)、多任务统一学习(分类、翻译、摘要等任务均转换成文本生成任务)；
实验数据 - 750GB 的 C4 数据集、220M/770M 参数量。

架构分析#

Transformer 架构

还有一种特殊的解码器结构 Prefix Decoder，它允许模型对输入序列中的一部分(“前缀”部分)使用双向注意力，而对其余部分(生成部分)保持传统的单向自注意力。这种设计结合了 Encoder-Decoder 模型的上下文理解能力和 Decoder-Only 模型的高效生成能力。

视觉模型#

ViT#

Transformer 架构

首先将输入图像 $\small[H,W,C]$ 划分成若干 2D 块然后展平得到 $\small[N, (P^2 \times C)]$ 。以 ViT-B/16 为例，将输入图像 (224, 224) 按固定大小 (16, 16) 划分为 (224/16)^2=196 个块。将每个块 (16, 16, 3) 通过线性映射转换为长度为 16x16x3=768 的向量。实际通过卷积核为 16x16，步长为 16 的卷积层实现。此时输入为 (196, 768)。

在嵌入向量的前端添加一个可学习的类别嵌入，在编码器输出时可作为图像的特征表示。预训练阶段，MLP 分类器将该类别嵌入作为输入(预训练阶段 MLP 只有一个隐藏层，而微调阶段 MLP 只有一个线形层)。为了保留图像的位置信息，引入了可学习的 1D 位置编码，直接叠加在嵌入向量上。最终输入为 (197, 768)。

微调阶段使用更高分辨率的图像，且保持每一个块的尺寸不变。但是输入序列的有效长度变长，导致预训练阶段的位置编码不再匹配，因此使用 2D 插值调整位置编码。

Swin Transformer#

Swin Transformer 引入基于卷积神经网络的归纳偏置：

局部性 - 只计算窗口和移动窗口内的注意力，同时跨窗口连接；
层次化 - 通过分层结构，提取不同尺度的特征。

Swin Transformer Locality

该模型的每一个阶段都采用两层移动窗口的设置。第一层划分窗口的自注意无法捕捉全局的特征信息。因此第二层将窗口进行滑动，这样原本的四个区域就变成了九个区域。且新窗口包含之前窗口的边界，能够建立不同窗口的连接。

高效计算方法

为了高效地计算，首先将九个区域调整为四个区域，其中三个区域包含来自不同区域的 patch。当两个来自不同区域的 patch 交互时，在它的位置上增加一个比较大的负值，进行 softmax 时候该位置便会趋于 0。

Swin Transformer Architecture

首先将输入图像划分成不重叠的 patch，每个 patch 的大小为 $4\times 4$ ，可以使用 $4\times4\times3=48$ 维的向量表示。因此网络的输入特征维度时 $\frac{H}{4}\times \frac{W}{4} \times 48$ 。

在第一个阶段，首先使用线形嵌入(Linear Embedding)层将每个 patch 的特征维度映射进行映射。然后使用多组连续的 Swin Transformer 模块处理，如上图 (b) 所示。分别使用了基于窗口和移动窗口的多头自注意力机制。在后三个阶段，每个阶段首先使用图像块合并(Patch Merging)层产生分层表示。通过合并相邻的 $2\times2$ patch 使得特征的维度不断发生变化。

Image GPT#

Image GPT

首先将输入图像进行下采样，将其转换为 1D 序列；然后进行模型预训练，采用两种预训练方法：

Next pixel prediction 是一种自回归的预训练方法，该方法根据前面的像素值预测下一个像素值(采用光栅顺序)，并最终对图像的概率密度进行整体建模。它训练的目标是最小化负对数似然：

L_{\mathrm{AR}}=\mathbb{E}_{x\sim X}[-\log p(x_{\pi_i}|x_{\pi_1},x_{\pi_2},\cdots,x_{\pi_{i-1}},\theta)]

Masked pixel prediction 首先遮掩输入序列若干位置的值，并对这些值进行预测。它训练的目标是最小化遮掩位置元素的负对数似然：

L_{\mathrm{MASK}}=\mathbb{E}_M\sum_{i\in M}[-\log p(x_i|x_{[1,n]\setminus M})]

通过预训练，模型学习到输入序列的分层特征表示。

微调阶段同时优化 $L_{GEN}+L_{CLF}$ ，其中 $L_{GEN}$ 代表 AR 或 MASK 损失、 $L_{CLF}$ 代表分类损失。该方法也被称为带有 辅助训练目标 的微调(Fine-tuning with auxiliary training objective)。

注意力机制#

Transformer#

缩放点积注意力#

多头注意力#

逐位前馈神经网络#

绝对位置编码#

编码器 / 解码器#

激活函数#

注意力模块的优化#

Sparse Attention#

Linear Attention#

KV Cache#

GQA / MQA#

MLA#

Flash Attention#

Page Attention#

Ring Attention#

混合专家模型(MoE)#

归一化#

Post / Pre-Norm#

RMS / Layer-Norm#

位置编码#

绝对位置编码#

相对位置编码#

T5 位置编码#

ALiBi#

RoPE#

长度外推优化#

Position Interpolation#

NTK-aware#

YaRN#

解码策略#

贪婪采样的改进策略#

随机采样的改进策略#

语言模型#

GPT / Decoder#

BERT / Encoder#

T5#

架构分析#

视觉模型#

ViT#

Swin Transformer#

Image GPT#

Credit#