Lory：用于自回归语言模型预训练的完全可微分混合专家（MOE）模型

24年5月来自普林斯顿大学和Meta AI的论文“Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training”。

混合专家 (MoE) 模型有助于实现高效扩展；然而，训练路由器网络带来了优化不可微分离散目标的挑战。最近，提出了一种完全可微分的 MoE 架构 SMEAR。Lory，是一个将此类架构扩展到自回归语言模型预训练的方法。Lory 引入了两种关键技术：(1) 因果分段路由策略，可在保留语言模型自回归性质的同时实现专家合并操作的高效率；(2) 基于相似性的数据批处理方法，在训练实例中对相似文档进行聚类来鼓励专家专业化。从头开始在 150B 个 token 上预训练一系列 Lory 模型，最多有 32 位专家和 30B（1.5B 活跃）参数。实验结果表明，在困惑度 (+13.9%) 和各种下游任务 (+1.5%-11.1%) 方面，与参数相匹配的密集模型相比，Lory 模型的性能有显著提升。尽管采用的是分段路由，但与采用 token 级路由的最先进 MoE 模型相比，Lory 模型的性能仍然具有竞争力。

如图所示是Lory的概览图：

基于 Transformer 的 MoE 语言模型通常用稀疏激活的 MoE 层替代前馈网络 (FFN) 层（Shazeer，2017；Fedus，2022；Zoph，2022）。假设 MoE 层由 E 个专家 FFN 组成，每个参数化为 FFN(·; θ1)，…，FFN(·; θE)，其中函数 FFN 定义单个专家模块。对于序列中的每个 token x，MoE 层将隐表征 hx 作为输入，并通过稀疏激活此层中的 k 个专家并通过加权和聚合输出来计算其输出 ox：

第 i 位专家的路由权重 ei 由路由网络或路由器 R 测量，该路由器将 hx 作为输入并计算每个专家的权重。实际上，为了实现稀疏性和计算效率，每个 MoE 层仅激活具有最高路由权重的一个（Fedus，2022）或前 k 个（Lepikhin，2021）专家。其余专家的权重设置为 0（即 ei = 0），从而无需计算 FFN(hx; θi) 并有效地停用第 i 位专家。

训练稀疏激活的 MoE 模型的主要挑战来自于训练离散路由器的难度。一个有希望的方向是设计完全可微分的 MoE 架构，这些架构不依赖于额外的损失公式来进行稳定训练。最近的模型架构 (Muqeeth et al., 2023) 通过计算参数空间中所有专家 FFN 的加权平均值 (Matena & Raffel, 2022; Wortsman et al., 2022) 证明了可行性，从而创建了一个“合并的 FFN”。给定一个输入 x 及其相应的路由权重 ei，合并的 FFN 的输出 ox 计算如下：

然而，将其简单地扩展到自回归语言模型不可行，因为这需要为序列中的每个 token 计算合并的 FFN，因为合并操作的计算成本与专家的数量成线性关系。 SMEAR（Muqeeth，2023）仅在文本分类任务的下游微调中进行了评估，该任务根据整个输入序列的池化表示做出路由决策，即 ei = Softmax(R(∑hxj/L))i。此类操作将破坏语言模型预训练中的 L 自回归属性。

在Lory中，提出因果分段路由，以自回归方式有效地跨段的路由信息。根据前一个段的信息在 MoE 层中合并 FFN，并使用它来处理当前段。具体来说，给定一个由 L 个 tokens（例如 L = 4096）组成的训练实例 X，将训练实例拆分为 N 个段，每个段包含 T（例如 T = 256）个连续 tokens。对于第 k 个段 Sk（当 k > 1 时），计算其前一个段 Sk−1 的隐表示平均值，表示为 hk−1。平均隐表示允许模型在推理过程中适应不同长度的提示。然后用 hk−1 来确定路由权重，从而得到合并的专家 θ：

然后用合并的专家 θ 来处理当前段 Sk 中的所有 tokens，即 ox = FFN(hx; θ)，∀x ∈ Sk。这种方法保证模型做出的路由决策完全基于先前位置的数据。对于第一个段 S1，段本身的表示用于计算其自身 FFN 的合并权重。为了防止信息泄露，在 R(h1) 上实施了“停止-梯度”操作。与密集模型的训练相比，在分段级合并专家产生的开销最小。

推理期间仅提示路由。在推理过程中，从给定的提示开始，并根据提示的平均隐表示对每层做出单个路由决策。此路由决策确定合并的 FFN，并在整个生成过程中一致使用。值得注意的是，此推理过程与密集模型一样简单且高计算效率。

在Lory中，采用了一种受 Shi (2024) 启发的基于相似性的数据批处理技术，该技术按顺序连接相似的文档以构建训练实例。这鼓励了相邻片段之间的高相似性，使专家能够专注于不同的领域或主题。用 Contriever (Izacard，2022) 测量文档相似性，并基于贪婪搜索算法连接相似的文档。虽然采用与 Shi (2024) 类似的数据批处理技术，但动机与他们不同。虽然他们的工作旨在改进语言模型跨文档边界的推理，不过这种技术在鼓励在训练 MoE 模型时专家专业化是有效的。

以下是Lory实验的设置环境等信息。

模型。通过训练仅包含 0.3B 和 1.5B 个活动参数的解码器 Transformer 模型来评估方法。对于 Transformer 模型中的每个 FFN 层，将其替换为具有全相同架构的 E个（取8, 16, 32）专家的 MoE 层。遵循 LLaMA 并使用 SwiGLU 作为 FFN 中的激活函数。我用与 LLaMA 模型相同的token化器。所有模型都使用 4096 个tokens的上下文窗口进行训练。在因果分段路由策略中，将每段长度设置为 T=256。

训练细节。用 AdamW 优化器（Loshchilov & Hutter，2019），β1 = 0.9 和 β2 = 0.95，并用 2e-4 的学习率和余弦学习率调度。所有模型的批量大小为 100 万个tokens。采用数据并行性和 ZeRO 优化进行分布式训练。在训练开始时，训练一个参数匹配的密集模型，并复制 FFN 层作为 MoE 模型的初始化。在实验中，用前 5% 的训练步骤作为热身来初始化 MoE 权重。如果没有热身训练，可能更多专家未被充分利用。还对前 5% 的训练步学习率调度程序应用线性热身。用多达 64 个 A100 GPU 来训练模型。

训练数据集。随机抽取 Commoncrawl 数据集 (Wenzek et al., 2019) 的一个子集作为训练数据。完整的训练数据集总共包含 1500 亿个 tokens。按照 Shi et al. (2024) 的方法，在这个子集上应用基于相似性数据批处理方法构建所有训练实例。

评估数据集。从 arXiv、Books、Wikipedia、C4（Raffel，2020）和 Python 代码（Github 的 Python 子集）中采样的支持评估数据集上，计算训练模型的困惑度，评估语言建模任务上的所有模型。每个评估数据集包含 1K 个样本，每个样本由 4096 个tokens组成。

还使用上下文学习（Brown，2020）评估下游任务中的模型，包括

• 常识推理：BoolQ（Clark，2019）、PIQA（Bisk ，2020）、SIQA（Sap，2019）、HellaSwag（Zellers，2019）、WinoGrand（Sakaguchi，2020）；
• 阅读理解：RACE (Lai et al., 2017)、ARC (Clark et al., 2018))；
• 闭卷问答：Natural Questions (Kwiatkowski et al., 2019)、TriviaQA (Joshi et al., 2017)；
• 文本分类：AGNews (Zhang et al., 2015)、SST-2 Socher et al. (2013)、Amazon 和 Yelp (Zhang et al., 2015)、FEVER (Thorne et al., 2018)、MRPC (Dolan & Brockett, 2005)。

对于文本分类任务，遵循 Min et al. (2022) 的评估设置；对于其余任务，遵循与 Touvron et al. (2023b) 相同的设置。

如下是因果分段路由的算法伪代码：