模型架构

到目前为止，我们已经将语言模型定义为对词元序列的概率分布 $p(x_1,\dots ,x_L)$，我们已经看到这种定义非常优雅且强大（通过提示，语言模型原则上可以完成任何任务，正如GPT-3所示）。然而，在实践中，对于专门的任务来说，避免生成整个序列的生成模型可能更高效。

上下文向量表征 (Contextual Embedding): 作为模型处理的先决条件，其关键是将词元序列表示为响应的上下文的向量表征：

$$[the,mouse,ate,the,cheese]\stackrel{\varphi }{\Rightarrow }[\left(\begin{array}{c}1\\ 0.1\end{array}\right),\left(\begin{array}{l}0\\ 1\end{array}\right),\left(\begin{array}{l}1\\ 1\end{array}\right),\left(\begin{array}{c}1\\ -0.1\end{array}\right),\left(\begin{array}{c}0\\ -1\end{array}\right)].$$

正如名称所示，词元的上下文向量表征取决于其上下文（周围的单词）；例如，考虑mouse的向量表示需要关注到周围某个窗口大小的其他单词。

• 符号表示：我们将 $\varphi :V^L\to {ℝ}^{d\times L}$ 定义为嵌入函数（类似于序列的特征映射，映射为对应的向量表示）。

• 对于词元序列 $x1:L=[x_1,\dots ,x_L]$，$\varphi$ 生成上下文向量表征 $\varphi (x_{1:L})$。

编码端（Encoder-Only）架构

编码端架构的著名的模型如BERT、RoBERTa等。这些语言模型生成上下文向量表征，但不能直接用于生成文本。可以表示为， $x_{1:L}\Rightarrow \varphi (x_{1:L})$ 。这些上下文向量表征通常用于分类任务（也被称为自然语言理解任务）。任务形式比较简单，下面以情感分类/自然语言推理任务举例：

$$\mathrm{情}\mathrm{感}\mathrm{分}\mathrm{析}\mathrm{输}\mathrm{入}\mathrm{与}\mathrm{输}\mathrm{出}\mathrm{形}\mathrm{式}：[[CLS],\mathrm{他}\mathrm{们},\mathrm{移}\mathrm{动},\mathrm{而},\mathrm{强}\mathrm{大}]\Rightarrow \mathrm{正}\mathrm{面}\mathrm{情}\mathrm{绪}$$$$\mathrm{自}\mathrm{然}\mathrm{语}\mathrm{言}\mathrm{处}\mathrm{理}\mathrm{输}\mathrm{入}\mathrm{与}\mathrm{输}\mathrm{出}\mathrm{形}\mathrm{式}：[[CLS],\mathrm{所}\mathrm{有},\mathrm{动}\mathrm{物},\mathrm{都},\mathrm{喜}\mathrm{欢},\mathrm{吃},\mathrm{饼}\mathrm{干},\mathrm{哦}]\Rightarrow \mathrm{蕴}\mathrm{涵}$$

该架构的优势是对于文本的上下文信息有更好的理解，因此该模型架构才会多用于理解任务。该架构的优点是对于每个 $x_i$ ，上下文向量表征可以双向地依赖于左侧上下文 $(x_{1:i-1})$ 和右侧上下文 $(x_{i+1:L})$ 。但是缺点在于不能自然地生成完整文本，且需要更多的特定训练目标（如掩码语言建模）。

解码器（Decoder-Only）架构

解码器架构的著名模型就是大名鼎鼎的GPT系列模型。这些是我们常见的自回归语言模型，给定一个提示 $x_{1:i}$ ，它们可以生成上下文向量表征，并对下一个词元 $x_{i+1}$ （以及递归地，整个完成 $x_{i+1:L}$） 生成一个概率分布。 $x_{1:i}\Rightarrow \varphi (x_{1:i}),p(x_{i+1}\mid x_{1:i})$ 。我们以自动补全任务来说，输入与输出的形式为， $[[CLS],\mathrm{他}\mathrm{们},\mathrm{移}\mathrm{动},\mathrm{而}]\Rightarrow \mathrm{强}\mathrm{大}$ 。与编码端架构比，其优点为能够自然地生成完整文本，有简单的训练目标（最大似然）。缺点也很明显，对于每个 $x_i$ ，上下文向量表征只能单向地依赖于左侧上下文 ($x_{1:i-1}$) 。

编码-解码端（Encoder-Decoder）架构

编码-解码端架构就是最初的Transformer模型，其他的还有如BART、T5等模型。这些模型在某种程度上结合了两者的优点：它们可以使用双向上下文向量表征来处理输入 $x_{1:L}$ ，并且可以生成输出 $y_{1:L}$ 。可以公式化为：

$$x1:L\Rightarrow \varphi (x1:L),p(y1:L\mid \varphi (x1:L))。$$

以表格到文本生成任务为例，其输入和输出的可以表示为：

$$[\mathrm{名}\mathrm{称}:,\mathrm{植}\mathrm{物},|,\mathrm{类}\mathrm{型}:,\mathrm{花}\mathrm{卉},\mathrm{商}\mathrm{店}]\Rightarrow [\mathrm{花}\mathrm{卉},\mathrm{是},\mathrm{一},\mathrm{个},\mathrm{商}\mathrm{店}]。$$

该模型具有编码端，解码端两个架构的共同的优点，对于每个 $x_i$ ，上下文向量表征可以双向地依赖于左侧上下文 $x_{1:i-1}$ ) 和右侧上下文 ( $x_{i+1:L}$ )，可以自由的生成文本数据。缺点就说需要更多的特定训练目标。

语言模型理论

下一步，我们会介绍语言模型的模型架构，重点介绍Transformer架构机器延伸的内容。另外我们对于架构还会对于之前RNN网络的核心知识进行阐述，其目的是对于代表性的模型架构进行学习，为未来的内容增加知识储备。

深度学习的美妙之处在于能够创建构建模块，就像我们用函数构建整个程序一样。因此，在下面的模型架构的讲述中，我们能够像下面的函数一样封装，以函数的的方法进行理解:

$$TransformerBlock(x_{1:L})$$

为了简单起见，我们将在函数主体中包含参数，接下来，我们将定义一个构建模块库，直到构建完整的Transformer模型。

基础架构

首先，我们需要将词元序列转换为序列的向量形式。 $EmbedToken$ 函数通过在嵌入矩阵 $E\in {ℝ}^{|v|\times d}$ 中查找每个词元所对应的向量，该向量的具体值这是从数据中学习的参数：

def $EmbedToken(x_{1:L}:V^L)\to {ℝ}^{d\times L}$ ：

• 将序列 $x_{1:L}$ 中的每个词元 $xi$ 转换为向量。
• 返回[Ex1,…,ExL]。

以上的词嵌入是传统的词嵌入，向量内容与上下文无关。这里我们定义一个抽象的 $SequenceModel$ 函数，它接受这些上下文无关的嵌入，并将它们映射为上下文相关的嵌入。

$defSequenceModel(x_{1:L}:{ℝ}^{d\times L})\to {ℝ}^{d\times L}$ ：

• 针对序列 $x_{1:L}$ 中的每个元素xi进行处理，考虑其他元素。
• [抽象实现（例如， $FeedForwardSequenceModel$ ， $SequenceRNN$ ， $TransformerBlock$ ）]

最简单类型的序列模型基于前馈网络（Bengio等人，2003），应用于固定长度的上下文，就像n-gram模型一样，函数的实现如下：

def $FeedForwardSequenceModel(x_{1:L}:{ℝ}^{d\times L})\to {ℝ}^{d\times L}$ ：

• 通过查看最后 $n$ 个元素处理序列 $x_{1:L}$ 中的每个元素 $xi$ 。
• 对于每个 $i=1,\dots ,L$ ：
  • 计算 $h_i$=$FeedForward(x_{i-n+1},\dots ,x_i)$ 。
• 返回[ $h_1,\dots ,h_L$ ]。

递归神经网络

第一个真正的序列模型是递归神经网络（RNN），它是一类模型，包括简单的RNN、LSTM和GRU。基本形式的RNN通过递归地计算一系列隐藏状态来进行计算。

def $SequenceRNN(x:{ℝ}^{d\times L})\to {ℝ}^{d\times L}$ ：

• 从左到右处理序列 $x_1,\dots ,x_L$ ，并递归计算向量 $h_1,\dots ,h_L$ 。
• 对于 $i=1,\dots ,L$ ：
  • 计算 $h_i=RNN(h_{i-1},x_i)$ 。
  • 返回 $[h_1,\dots ,h_L]$ 。

实际完成工作的模块是RNN，类似于有限状态机，它接收当前状态h、新观测值x，并返回更新后的状态：

def $RNN(h:{ℝ}^d,x:{ℝ}^d)\to {ℝ}^d$ ：

• 根据新的观测值x更新隐藏状态h。
• [抽象实现（例如，SimpleRNN，LSTM，GRU）]

有三种方法可以实现RNN。最早的RNN是简单RNN（Elman，1990 (opens new window)），它将$h$和$x$的线性组合通过逐元素非线性函数 $\sigma$ （例如，逻辑函数 $\sigma (z)=(1+e-z)-1$ 或更现代的 $ReLU$ 函数 $\sigma (z)=max(0,z)$ ）进行处理。

def $SimpleRNN(h:ℝd,x:ℝd)\to ℝd$ ：

• 通过简单的线性变换和非线性函数根据新的观测值 $x$ 更新隐藏状态 $h$ 。
• 返回 $\sigma (Uh+Vx+b)$ 。

正如定义的RNN只依赖于过去，但我们可以通过向后运行另一个RNN来使其依赖于未来两个。这些模型被ELMo和ULMFiT使用。

def $BidirectionalSequenceRNN(x_{1:L}:{ℝ}^{d\times L})\to {ℝ}^{2d\times L}$ ：

• 同时从左到右和从右到左处理序列。
• 计算从左到右： $[h{\to }_1,\dots ,h{\to }_L]\leftarrow SequenceRNN(x_1,\dots ,x_L)$ 。
• 计算从右到左： $[h{\leftarrow }_L,\dots ,h{\leftarrow }_1]\leftarrow SequenceRNN(x_L,\dots ,x_1)$ 。
• 返回 $[h{\to }_{1}h{\leftarrow }_1,\dots ,h{\to }_{L}h{\leftarrow }_L]$ 。

注：

• 简单RNN由于梯度消失的问题很难训练。
• 为了解决这个问题，发展了长短期记忆（LSTM）和门控循环单元（GRU）（都属于RNN）。
• 然而，即使嵌入h200可以依赖于任意远的过去（例如，x1），它不太可能以“精确”的方式依赖于它（更多讨论，请参见Khandelwal等人，2018）。
• 从某种意义上说，LSTM真正地将深度学习引入了NLP领域。

Transformer

现在，我们将讨论Transformer（Vaswani等人，2017 (opens new window)），这是真正推动大型语言模型发展的序列模型。正如之前所提到的，Transformer模型将其分解为Decoder-Only（GPT-2，GPT-3）、Encoder-Only（BERT，RoBERTa）和Encoder-Decoder（BART，T5）模型的构建模块。

关于Transformer的学习资源有很多：

• Illustrated Transformer (opens new window)和Illustrated GPT-2 (opens new window)：对Transformer的视觉描述非常好。
• Annotated Transformer (opens new window)：Transformer的Pytorch实现。

强烈建议您阅读这些参考资料。该课程主要依据代码函数和接口进行讲解。

注意力机制

Transformer的关键是注意机制，这个机制早在机器翻译中就被开发出来了（Bahdananu等人，2017）。可以将注意力视为一个“软”查找表，其中有一个查询 $y $，我们希望将其与序列 $x_{1:L}=[x_1,\dots ,x_L]$ 的每个元素进行匹配。我们可以通过线性变换将每个 $x_i$ 视为表示键值对：

$$(W_{key}{x}_i)：(W_{value}{x}_i)$$

并通过另一个线性变换形成查询：

$$W_{query}y$$

可以将键和查询进行比较，得到一个分数：

$$scor{e}_i=x_i^{\top }{W}_{key}^{\top }{W}_{query}y$$

这些分数可以进行指数化和归一化，形成关于词元位置$1,\dots ,L$的概率分布：

$$[{\alpha }_1,\dots ,{\alpha }_L]=softmax([scor{e}_1,\dots ,scor{e}_L])$$

然后最终的输出是基于值的加权组合：

$$\sum _{i=1}^L{\alpha }_i\left(W_{value}{x}_i\right)$$

我们可以用矩阵形式简洁地表示所有这些内容：

def $Attention(x_{1:L}:{ℝ}^{d\times L},y:{ℝ}^d)\to {ℝ}^d$ ：

• 通过将其与每个$x_i$进行比较来处理$y$。
• 返回 $W_{value}{x}_{1:L}\mathrm{softmax}\left(x_{1:L}^{\mathrm{\top }}{W}_{key}^{\mathrm{\top }}{W}_{query}y/\sqrt{d}\right)$

我们可以将注意力看作是具有多个方面（例如，句法、语义）的匹配。为了适应这一点，我们可以同时使用多个注意力头，并简单地组合它们的输出。

def $MultiHeadedAttention(x_{1:L}:{ℝ}^{d\times L},y:{ℝ}^d)\to {ℝ}^d$ :

• 通过将其与每个xi与nheads个方面进行比较，处理y。
• 返回 $W_{output}[\underset{n_{heads}times}{\underbrace{[\mathrm{Attention}(x_{1:L},y),\dots ,\mathrm{Attention}(x_{1:L},y)]}}$

对于自注意层，我们将用$x_i$替换$y$作为查询参数来产生，其本质上就是将自身的$x_i$对句子的其他上下文内容进行 $Attention$ 的运算：

def $SelfAttention(x_{1:L}:{ℝ}_{d\times L})\to {ℝ}_{d\times L})$ ：

• 将每个元素xi与其他元素进行比较。
• 返回 $[Attention(x_{1:L},x_1),\dots ,Attention(x_{1:L},x_L)]$ 。

自注意力使得所有的词元都可以“相互通信”，而前馈层提供进一步的连接：

def $FeedForward(x_{1:L}:{ℝ}^{d\times L})\to {ℝ}^{d\times L}$ ：

• 独立处理每个词元。
• 对于 $i=1,\dots ,L$ ：
  • 计算 $y_i=W_{2}max(W_1{x}_i+b_1,0)+b_2$ 。
• 返回 $[y_1,\dots ,y_L]$ 。

对于Transformer的主要的组件，我们差不多进行介绍。原则上，我们可以只需将 $FeedForward\circ SelfAttention$ 序列模型迭代96次以构建GPT-3，但是那样的网络很难优化（同样受到沿深度方向的梯度消失问题的困扰）。因此，我们必须进行两个手段，以确保网络可训练。

残差连接和归一化

残差连接：计算机视觉中的一个技巧是残差连接（ResNet）。我们不仅应用某个函数f：

$$f(x1:L)，$$

而是添加一个残差（跳跃）连接，以便如果$f$的梯度消失，梯度仍然可以通过 $x_{1:L}$ 进行计算：

$$x_{1:L}+f(x_{1:L})。$$

层归一化:另一个技巧是层归一化，它接收一个向量并确保其元素不会太大：

def $LayerNorm(x_{1:L}:{ℝ}^{d\times L})\to {ℝ}^{d\times L}$ ：

• 使得每个 $x_i$ 既不太大也不太小。

我们首先定义一个适配器函数，该函数接受一个序列模型$f$并使其“鲁棒”：

def $AddNorm(f:(ℝd^{\times L}\to {ℝ}^{d\times L}),x_{1:L}:{ℝ}_{d\times L})\to {ℝ}^{d\times L}$ ：

• 安全地将f应用于 $x_{1:L}$ 。
• 返回 $LayerNorm(x_{1:L}+f(x_{1:L}))$ 。

最后，我们可以简洁地定义Transformer块如下：

def $TransformerBlock(x_{1:L}:{ℝ}^{d\times L})\to {ℝ}^{d\times L}$ ：

• 处理上下文中的每个元素 $x_i$ 。
• 返回 $AddNorm(FeedForward,AddNorm(SelfAttention,x_{1:L}))$ 。

位置嵌入

最后我们对目前语言模型的位置嵌入进行讨论。您可能已经注意到，根据定义，词元的嵌入不依赖于其在序列中的位置，因此两个句子中的𝗆𝗈𝗎𝗌𝖾将具有相同的嵌入，从而在句子位置的角度忽略了上下文的信息，这是不合理的。

[𝗍𝗁𝖾,𝗆𝗈𝗎𝗌𝖾,𝖺𝗍𝖾,𝗍𝗁𝖾,𝖼𝗁𝖾𝖾𝗌𝖾]
[𝗍𝗁𝖾,𝖼𝗁𝖾𝖾𝗌𝖾,𝖺𝗍𝖾,𝗍𝗁𝖾,𝗆𝗈𝗎𝗌𝖾]

为了解决这个问题，我们将位置信息添加到嵌入中：

def $EmbedTokenWithPosition(x_{1:L}:{ℝ}^{d\times L})$ ：

• 添加位置信息。
• 定义位置嵌入：
  • 偶数维度：$P_{i,2j}=sin(i/{10000}^{2j/dmodel})$
  • 奇数维度：$P_{i,2j+1}=cos(i/{10000}^{2j/dmodel})$
• 返回 $[x_1+P_1,\dots ,x_L+P_L]$ 。

上面的函数中， $i$ 表示句子中词元的位置， $j$ 表示该词元的向量表示维度位置。

最后我们来聊一下GPT-3。在所有组件就位后，我们现在可以简要地定义GPT-3架构，只需将Transformer块堆叠96次即可：

$$GPT-3(x_{1:L})=TransformerBloc{k}^{96}(EmbedTokenWithPosition(x_{1:L}))$$

架构的形状（如何分配1750亿个参数）：

• 隐藏状态的维度：dmodel=12288
• 中间前馈层的维度：dff=4dmodel
• 注意头的数量：nheads=96
• 上下文长度：L=2048

这些决策未必是最优的。Levine等人（2020） (opens new window)提供了一些理论上的证明，表明GPT-3的深度太深，这促使了更深但更宽的Jurassic架构的训练。

不同版本的Transformer之间存在重要但详细的差异：

• 层归一化“后归一化”（原始Transformer论文）与“先归一化”（GPT-2），这影响了训练的稳定性（Davis等人，2021 (opens new window)）。
• 应用了丢弃（Dropout）以防止过拟合。
• GPT-3使用了sparse Transformer (opens new window)（稀释 Transformer）来减少参数数量，并与稠密层交错使用。
• 根据Transformer的类型（Encdoer-Only, Decoder-Only, Encdoer-Decoder），使用不同的掩码操作。

Transformer 中的 mask 机制

在 Transformer 模型中，mask 机制是一种用于在 self-attention 中的技术，用以控制不同 token 之间的注意力交互。具体来说，Transformer 中使用两种类型的 mask**：padding mask** 和 sequence mask

Padding mask（填充掩码）

Padding mask（填充掩码）：在自注意力机制中，句子中的所有单词都会参与计算。但是，在实际的句子中，往往会存在填充符（比如-1），用来填充句子长度不够的情况。Padding mask 就是将这些填充符对应的位置标记为 0，以便在计算中将这些位置的单词忽略掉。 例如，假设我们有一个 batch_size 为 3、句子长度为 5 的输入序列：

[
    [1, 2, 3, -1, -1],
    [2, 3, -1, -1, -1],
    [1, -1, -1, -1, -1]
]

其中-1 表示填充符，用于填充长度不到 5 的位置。则 padding mask 为:

[
    [1, 1, 1, 0, 0],
    [1, 1, 0, 0, 0],
    [1, 0, 0, 0, 0]
]

其中 1 表示可以参与计算的位置，0 表示需要被忽略的位置。例如，第一句话前三个位置有具体的词，对应设为 1；而最后两个位置是 padding，对应设为 0。huggingface 的 BertTokenizer 分词所返回的 attention mask 指的就是 padding mask (opens new window)。

Sequence mask（序列掩码）

sequence mask 用于在 Decoder 端的 self-attention 中，以保证在生成序列时不会将未来的信息泄露给当前位置的单词。

例如，假设我们要生成一个长度为 5 的序列。在第 $i$ 个位置上生成的单词，需要将前 $i-1$ 个单词作为输入，但是不能将第 $i$ 个位置以后的单词作为输入。这就需要使用 sequence mask 将第 $i$ 个位置以后的单词掩盖掉。具体而言，sequence mask 会将第 $i$ 个位置以后的所有位置标记为 0，表示在计算中需要忽略这些位置的信息。

本文要比较不同 LLM 架构，其实是在比较 sequence mask。事实上，用户所使用的 LLM 所做的都是 seq2seq 任务，用户输入一个文本 prompt，LLM 对应输出一个文本 completion。为了方便以后的讨论，采用 A Survey of Large Language Models 这篇综述中的例子，用户输入的 prompt 为 A Survey of，期望语言模型的输出为 Large Language Models。

Causal Decoder

Causal LM 是因果语言模型，目前流行的大多数模型都是这种结构，别无他因，因为 GPT 系列模型内部结构就是它，还有开源界的 LLaMa 也是。Causal Decoder 架构的典型代表就是 GPT 系列模型，使用的是单向注意力掩码，以确保每个输入 token 只能注意到过去的 token 和它本身，输入和输出的 token 通过 Decoder 以相同的方式进行处理。在下图中，灰色代表对应的两个 token 互相之间看不到，否则就代表可以看到。例如，”Survery”可以看到前面的“A”，但是看不到后面的“of”。Causal Decoder 的 sequence mask 矩阵是一种典型的下三角矩阵。

Causal LM 只涉及到 Encoder-Decoder 中的 Decoder 部分，采用 Auto Regressive 模式，直白地说，就是根据历史的 token 来预测下一个 token，也是在 Attention Mask 这里做的手脚。

参照着 Prefix LM，可以看下 Causal LM 的 Attention Mask 机制(左)及流转过程(右)。

Prefix Decoder

Prefix LM 的代表模型有 UniLM、T5、GLM

Prefix LM，即前缀语言模型，Prefix Decoder 架构也被称为 non-causal Decoder 架构，该结构是 Google 的 T5 模型论文起的名字，望文知义来说，这个模型的“前缀”有些内容，但继续向前追溯的话，微软的 UniLM 已经提及到了。

Prefix LM 其实是 Encoder-Decoder 模型变体：

1. 在标准的 Encoder-Decoder 模型中，Encoder 和 Decoder 各自使用一个独立的 Transformer

2. 而在 Prefix LM，Encoder 和 Decoder 则共享了同一个 Transformer 结构，在 Transformer 内部通过 Attention Mask 机制来实现。

Prefix LM 是输入部分采用双向注意力，而输出部分采用单向注意力；和 Encoder-Decoder 不同的是，处理输入和输出的模型参数是完全共享的，从这一点看，它又和 causal Decoder 比较接近，都属于 Decoder-Only 架构：

与标准 Encoder-Decoder 类似，Prefix LM 在 Encoder 部分采用 Auto Encoding (AE-自编码)模式，即前缀序列中任意两个 token 都相互可见，而 Decoder 部分采用 Auto Regressive (AR-自回归)模式，即待生成的 token 可以看到 Encoder 侧所有 token(包括上下文)和 Decoder 侧已经生成的 token，但不能看未来尚未产生的 token。

下面的图很形象地解释了 Prefix LM 的 Attention Mask 机制(左)及流转过程(右)。

Encoder Decoder

Transformer 最初被提出来的时候采用的就是 Encoder-Decoder 架构，模型包含两部分 Encoder 和 Decoder，两部分参数独立。其中 Encoder 将输入序列处理为一种中间表示，而 Decoder 则基于中间表示自回归地生成目标序列，典型的 LLM 是 Flan-T5。Encoder 部分采用双向注意力，对应的 prompt 的每个 token 都可以互相看到；而 Decoder 部分仍然采用单向注意力，对应的 completion 仍然保证前面的 token 看不到后面的 token。

混合专家模型

基础知识

混合专家的想法可以追溯到Jacobs et al. (1991) (opens new window)。

为了介绍基本思想，假设我们正在解决一个预测问题：

$$x\in {\mathbb{R}}^d\Rightarrow y\in {\mathbb{R}}^d.$$

让我们从学习前馈（ReLU）神经网络开始：

$$h_{\theta }(x)=W_{2}max(W_{1}x,0),$$

其中参数为 $\theta =(W_1,W_2)$ 。

• 然而，这个函数可能表达能力不足。
• 我们可以使神经网络更宽或更深。

但专家的混合方法是：

• 定义 $E$ 个专家。

• 每个专家 $e=1,\dots ,E$ 都具有自己的嵌入 $w_e\in {\mathbb{R}}^d$ 。

• 将门控函数定义为 $E$ 个专家上的概率分布：

$$g_e(x)=\frac{\exp (w_e\cdot x)}{\sum _{e^{\prime }=1}^E\exp (w_{e^{\prime }}\cdot x)}.$$

• 每个专家 $e=1,\dots ,E$ 都具有自己的参数 ${\theta }^{(e)}=(W_1^{(e)},W_2^{(e)})$ 。

• 根据专家特定参数定义每个专家函数：

$$h_{{\theta }_e}(x)=W_2^{(e)}max(W_1^{(e)}x,0).$$

• 将最终函数定义为专家的混合：

$$f(x)=\sum _{e=1}^E\underset{\text{gating}}{\underbrace{g_e(x)}}\underset{\text{expert}}{\underbrace{h_{{\theta }_e}(x)}}.$$

4.1.1.1 示例

考虑d=2，并且每个专家都是一个线性分类器（来源 (opens new window)）：

4.1.1.2 训练

我们可以通过反向传播来学习混合专家模型。根据链式法则，可以得到：

$$\mathrm{\nabla }f(x)=\sum _{e=1}^E{g}_e(x)(\mathrm{\nabla }(\log g_e(x))h_{{\theta }_e}(x)+\mathrm{\nabla }{h}_{{\theta }_e}(x)).$$

注意到，梯度与 $g_e(x)$ 成比例，并且同时更新门控函数和专家。

4.1.1.3 节约计算

• 注意到，门控函数 $g(x)=[g_1(x),\dots ,g_E(x)]$ 对于每个专家都是非零的。例如：

$$g(x)=[0.04,0.8,0.01,0.15].$$

• 正如公式所言，专家的混合不会节省任何计算，因为前向传播仍然需要评估每个专家，而反向传播也必须接触每个专家。

• 然而，如果我们将门控函数 $g(x)=[g_1(x),\dots ,g_E(x)]$近似为$\tilde{g}(x)=[{\tilde{g}}_1(x),\dots ,{\tilde{g}}_E(x)]$ ，其中大多数专家都是零。因此，在前向和反向传播时，我们只需要使用非零 ${\tilde{g}}_e(x)$ 的专家 $e$ 。

例如，我们可以选取值排名前两位（top 2）的专家，并重新规范化：

$$\tilde{g}(x)=[0,0.84,0,0.16].$$

4.1.1.4 平衡专家

• 只有所有专家都参与进来，混合专家才有效。
• 如果只有一个专家处于活跃状态（例如， $g(x)=[0,1,0,0]$ )，那么这就是浪费。
• 此外，如果我们一直处于这种状态，那么未使用的专家的梯度将为零，因此他们将不会收到任何梯度并得到改善。
• 因此，使用混合专家的主要考虑因素之一是确保所有专家都能被输入使用。

101.1.5 并行

• 混合专家非常有利于并行。
• 每个专家都可以放置在不同的机器上。
• 我们可以在中心节点计算近似门控函数 $\tilde{g}(x)$ 。
• 然后，我们只要求包含激活专家的机器（稀疏）来处理 $x$ 。

4.1.2 Sparsely-gated mixture of experts (Lepikhin et al. 2021 (opens new window))

• 现在我们考虑如何将混合专家思想应用于语言模型。
• 最简单的解决方案是仍然使用96层Transformer，但是
  • 门控函数以某种方式应用于序列；
  • 只在顶层进行专家的结合。
• 因此，我们将混合专家的想法应用于：
  • 每个token
  • 每层Transformer block（或者隔层使用）
• 由于前馈层对于每个token是独立的，因此，我们将每个前馈网络转变为混合专家（MoE）前馈网络：

$$\text{MoETransformerBlock}(x_{1:L})=\text{AddNorm}(\text{MoEFeedForward},\text{AddNorm}(\text{SelfAttention},x_{1:L})).$$

• 隔层使用MoE Transformer block。

我们将top-2专家的近似门控函数定义如下：

• 计算第一个专家： $e_1=\arg \underset{e}{max}{g}_e(x)$ 。

• 计算第二个专家： $e_2=\arg \underset{e\ne e_1}{max}{g}_e(x)$ 。

• 始终保留第一个专家，并随机保留第二个专家：

  • 设 $p=min(2g_{e_2}(x),1)$
  • 在概率为 $p$ 的情况下， ${\tilde{g}}_{e_1}(x)=\frac{g_{e_1}(x)}{g_{e_1}(x)+g_{e_2}(x)},{\tilde{g}}_{e_2}(x)=\frac{g_{e_2}(x)}{g_{e_1}(x)+g_{e_2}(x)}$ 。对于其他专家 $e\notin \{e_1,e_2\}$，${\tilde{g}}_e(x)=0$ 。
  • 在概率为 $1-p$ 的情况下， ${\tilde{g}}_{e_1}(x)=1$。对于$e\ne e_1$，$\tilde g_e(x) = 0 $ 。

4.1.2.1 符号定义

• 设 $B$ 是一个batch中的token数量（在所有序列中）；通常在百万数量级。
• 设 $E$ 是专家数目；通常在千数量级。
• 设 $x_1,\dots ,x_B$ 为一个batch中的token。

4.1.2.2 平衡专家

• 设 $c_e=\sum _{i=1}^B\mathbf{1}[{\tilde{g}}_e(x_i)>0]$ 是专家 $e$ 被选中的次数。
• 注意，处理完一个batch后， $\sum _e{c}_e=B$ 。
• 如果所有专家都是平衡的，那么 $c_e=\frac{B}{E}$ 。
• 溢出：如果 $c_e>2\frac{B}{E}$ ，则设 $f(x)=x$ （带残差的旁路），其中2是容量系数。
• 辅助损失：我们期望 $c=[c_1,\dots ,c_E]$ 接近均匀分布。
• 我们可以惩罚 $\parallel c{\parallel }_2^2=\sum _{e=1}^E{c}_e^2$ ，但这是不可微分的。
• 定义 $m_e=\sum _{i=1}^B{g}_e(x_i)$ （这是 $c_e$ 的软版本)。
• 相反，我们在目标函数中添加 $\text{load-balancing-loss}=\sum _{e=1}^E{m}_e{c}_e$ 。这样，通过 $m_e$ 的梯度将为非零。

$\text{loss}=\text{negative-log-likelihood}+\lambda \text{load-balancing-loss}.$

例如，我们可以取 $\lambda =\frac{0.01}{B}$ 。

4.1.2.3 示例

下面是一个 $B=2$ 个token， $E=4$ 个专家的例子：

$$\begin{gathered} g(x_1)=[0.2,0.6,0.1,0.1]\Rightarrow \tilde{g}(x_1)=[0.25,0.75,0,0] \\ g(x_2)=[0.1,0.6,0.2,0.1]\Rightarrow \tilde{g}(x_2)=[0,0.75,0.25,0] \end{gathered}$$

统计为

$$c=[1,2,1,0]m=[0.3,1.2,0.3,0.2]$$

也就是说，我们会尝试降低专家2的权重，避免其被过度使用。

4.1.3 Switch Transformer (Fedus et al. 2021 (opens new window))

• 定义近似门控函数 $\tilde{g}(x)$ 只有一个专家（获得更多稀疏性）。
• 技巧：
  • 将FP32训练替换成FP16
  • 使用的较小参数进行初始化
  • 专家dropout
  • 专家并行
• 训练了一个1.6万亿参数模型
• 与T5-XXL（110亿参数）相比，训练速度提高了4倍

4.1.4 Balanced Assignment of Sparse Experts (BASE) layers (Lewis et al., 2021 (opens new window))

• BASE将近似门控函数 $\tilde{g}(x)$ 定义为对batch中的所有token进行联合优化的结果。
• 我们将为每个token分配1名专家，但负载平衡是一种约束，而不是软惩罚。
• 我们定义 $a=[a_1,\dots ,a_B]\in \{1,\dots ,E{\}}^B$ 作为联合分配向量。

$$\text{maximize}\sum _{i=1}^B{w}_{a_i}\cdot x_i\text{subject to}\mathrm{\forall }e:\sum _{i=1}^B\mathbf{1}[a_i=e]=\frac{B}{E}.$$

• 这是一个可以有效求解的线性方程。
• 在实践中，我们将线性方程并行化。
• 在测试时，只需选择top 1的专家即可。

4.1.4.1 实验设置

• Sparsely gated MoE (top-2 experts): 52.5B 参数
• Switch Transformer (top-1 expert): 52.5B 参数
• BASE (1 jointly optimized expert): 44.4B 参数 (1.3B shared 参数, 335M x 128 expert 参数)

BASE需要更多的计算来优化 $a$ ，但更稳定。

4.1.4.2 总结和下一步工作

• Switch Transformer（谷歌）使用了top-1专家。
• BASE（Facebook）为每个token分配1名专家，但进行了联合优化。
• 这两个模型的性能都无法与GPT-3可比。虽然谷歌和Facebook都发布了两个最新的高性能MoE语言模型，它们的性能确实与GPT-可比，但有趣的是，它们仍然基于最初简单的top-2专家：
  • 谷歌的GLaM
  • 来自Facebook的“FacebookMoE”

4.1.5 Generalist Language Model (GLaM) (Du et al. 2021 (opens new window))

4.1.5.1 规格

• 1.2万亿个参数（GPT-3有1750亿个参数）
• 64个专家，64层，32K个隐藏单元
• 每个token激活95B（1.2T的8%）的参数

4.1.5.2 其他

• 创建了共有1.6万亿个token的新数据集（GLaM dataset），来源包括网页、论坛、书籍、新闻等。
• 相对位置编码、门控线性单元、GeLU激活函数、RMSNorm（非LayerNorm）
• 如果遇到NaN/Inf，跳过权重更新/回滚到早期检查点。
• “通过仔细实施上述技巧，我们观察到，稀疏激活的模型在各个尺度上的训练都变得相当稳定。”

4.1.5.3 结果

• 与GPT-3相比，训练成本仅为1/3
• 在与GPT-3相同的基准上进行评估（开放域问答、阅读理解、SuperGLUE等）
• 与GPT-3相比，实现了更好的0-shot和1-shot性能（尤其是在知识密集型任务中的性能）
• 注：他们没有在GPT-3更强的few-shot中进行评估

4.1.5.4 WinoGender上的结果

示例：The nurse notified the patient that {her/his,their} shift would be ending in an hour.

GLaM的性别偏见少于GPT-3。

4.1.6 FacebookMoE (Artetxe et al., 2021 (opens new window))

4.1.6.1 实验设置

• 训练了一个1.1T参数的模型
• 512名专家（超过GLaM），32层，4096个隐藏单元
• 使用112 billion token进行训练，来源包括网页、论坛、书籍、新闻等。
• 小模型收益更大，模型越大，收益递减

StereoSet (opens new window)上的结果：

4.1.6.2 示例

The assistant went to work. {She brought her boss coffee., She was valued for her input.}

刻板印象随着模型大小的增加而变得更糟（与GLaM结果相反）。

4.1.7 Decentralized mixture-of-experts (Ryabinin & Gusev, 2020 (opens new window))

4.1.7.1 动机

• 到目前为止，混合专家纯粹是中心机构（如谷歌或Facebook）从扩大大语言模型的角度出发的。
• 然而，混合专家自然地指示了一种更激进的权力下放。
• 为例训练GPT-3，Azure超级计算机集群 (opens new window)耗资2.5亿美元。
• 我们如何利用数以亿计 (opens new window)的消费PC？
• Folding@Home (opens new window)是一个志愿者计算项目，利用世界各地的志愿者捐赠计算机进行分子动力学模拟。
• 2020年4月，Folding@Home有70万人捐赠了产生2.43 exaFLOP（GPT-3需要350千兆FLOP）（文章 (opens new window)）。
• 主要区别在于分子动力学模拟计算量大，不需要网络带宽。

4.1.7.2 主要考虑因素

• 节点众多（ ${10}^3\sim {10}^6$ 异构PC）
• 频繁的节点故障（5-20%的节点每天至少有一次故障）
• 家庭互联网通信带宽（100Mbps；相比之下，Azure超级计算机为400Gbps）

4.1.7.3 分布式哈希表

• $N$ 个节点
• 单个节点需要与其他 $O(\log N)$ 节点通信
• 使用Kademlia DHT协议（被BitTorrent和以太坊使用）

4.1.7.4 论文实验

• 选取top-4的专家（共256名专家）
• 每个专家都是一个Transformer层
• 在4个GPU上训练了一个小型Transformer LM

4.1.8 Diskin et al., 2021 (opens new window)

• 40名志愿者
• 为孟加拉语训练了一个ALBERT的掩码语言模型
• 一起训练Transformer：任何人都可以加入并贡献计算

4.1.9 总结

• 混合专家：起源于将不同专家应用于不同输入的经典理念
• 允许训练更大的语言模型（1.1万亿个参数）
• 与稠密Transformer模型相比，每个输入的效率高得多（FLOP更少）
• 效果难以比较：在相同规模上，直接比较仍然具有挑战性（GPT-3与GLaM与FacebookMoE）
• 对权力下放的重大影响

4.2 基于检索的模型

现在，我们转向另一类语言模型，基于检索的（或检索增强的、记忆增强的模型），它可以帮助我们突破稠密Transformer的缩放上限。

4.2.1 编码器-解码器

让我们首先关注使用编码器-解码器框架的序列到序列任务：

$$\text{input }x\Rightarrow \text{output }y$$

示例（开放问答）：

• 输入$x$：What is the capital of Canada?
• 输出$y$：Ottawa

回想一下，BART (opens new window)和T5 (opens new window)是编码器-解码器模型的代表：

$$p(y\mid x)$$

其使用去噪目标函数进行训练。 例如： 输入 $x$ ：Thank you $<X>$ me to your party $<Y>$ week. 输出 $y$ ： $<X>$ for inviting $<Y>$ last

10.2.2 检索方法

假设我们有一个存储库 $S$ ，它是一组序列（通常是文档或段落）的集合。

$$S=\{\text{Why is the...},\text{Thanks for},...,\text{The quick...},\text{Stanford...}\}.$$

基于检索的模型直观的生成过程：

• 基于输入 $x$ ，检索相关序列 $z$ 。
• 给定检索序列 $z$ 和输入 $x$ ，生成输出 $y$ 。

示例（开放问答）：

• 输入 $x$ ：What is the capital of Canada?
• 检索 $z$ ：Ottawa is the capital city of Canada.
• 输出 $y$ ：Ottawa

最近邻是最常用的一种检索方法：

• $S$ 是训练集。
• 检索 $(x^{\prime },y^{\prime })\in S$ ，使得 $x^{\prime }$ 和 $x $最相似。
• 生成 $y=y^{\prime }$ 。

10.2.3 Retrieval-augmented generation (RAG) (Lewis et al., 2020 (opens new window))

形式上，RAG模型定义如下：

$$(y\mid x)=\sum _{z\in S}\underset{\text{retriever}}{\underbrace{p(z\mid x)}}\underset{\text{generator}}{\underbrace{p(y\mid z,x)}}.$$

在实践中， $\sum _{z\in S}$ 由前k个代替（类似于为混合专家选择前1个或2个专家）。

10.2.3.1 检索器

Dense Passage Retrieval (DPR)** (Karpukhin et al., 2020 (opens new window))

$$p(z\mid x)=\frac{\exp ({\text{BERT}}_{\text{d}}(z)\cdot {\text{BERT}}_{\text{q}}(x))}{\sum _{z^{\prime }\in S}\exp ({\text{BERT}}_{\text{d}}(z^{\prime })\cdot {\text{BERT}}_{\text{q}}(x))}.$$

• 这里以用维基百科文章的标题来检索段落为例
• 使用QA数据集（如NaturalQuestions、TriviQA等）的query、正例、负例 $(q,p^{+},p_1^{-},\dots ,p_n^{-})$ 来训练模型：
  • 负例：随机或者使用BM25检索出的不包含答案的段落
• 推理：使用FAISS (opens new window)（Facebook AI相似性搜索）

10.2.3.2 生成器

$$p(y\mid z,x)=p(y\mid \text{concat}(z,x)).$$

• 使用BART-large（400M参数），其中输入为检索出的段落 $z$ 和输入 $x$
• 回想一下，BART是基于网络、新闻、书籍和故事数据，使用去噪目标函数（例如，掩码）训练得到的

10.2.3.3 训练

• 用BART、DPR（用BERT初始化）初始化
• 训练 $\text{BART}$和${\text{BERT}}_{\text{q}}$

10.2.3.4 实验

• 在Jeopardy问题生成任务上，输入Hemingway的检索结果：

• 实验结果表明，优于非检索方法：

这里引用GPT-3 few-shot的结果进行比较：NaturalQuestions (29.9%), WebQuestions (41.5%), TriviaQA (71.2%)

10.2.4 RETRO (Borgeaud et al., 2021 (opens new window))

• 基于32个token的块进行检索
• 存储库：2 trillion tokens
• 70亿参数（比GPT-3少25倍）
• 使用冻结的BERT进行检索（不更新）
• 在MassiveText上训练（与训练Gopher的数据集相同）

10.2.4.1 实验结果

• 在语言建模方面表现出色
• NaturalQuestions准确率：45.5%（SOTA为54.7%）

10.2.5 讨论

• 基于检索的模型高度适合知识密集型的问答任务。
• 除了可扩展性之外，基于检索的模型还提供了可解释性和更新存储库的能力。
• 目前尚不清楚这些模型是否具有与稠密Transformer相同的通用能力。

10.3 总体总结

• 为了扩大模型规模，需要改进稠密Transformer。
• 混合专家和基于检索的方法相结合更有效。
• 如何设计更好的、可扩展的体系结构仍然是一个悬而未决的问题。

Reference

• LLM 主流框架：Causal Decoder、Prefix Decoder 和 Encoder-Decoder (opens new window)
• 比较 Causal decoder、Prefix decoder 和 encoder-decoder (opens new window)