大型语言模型（LLMs）的性能非常强大，但是现有的模型训练和部署成本都很高。而且在不忘记先前知识的前提，扩展它们去学习新的知识也很困难。也很难针对特定的任务去提取出轻量化的模型。

(资料图片)

最近，来自 MIT-IBM Waston AI Lab、清华大学、Mila 的研究人员联合提出了一种新的神经网络架构ModuleFormer，利用模块化来大幅提高大型语言模型的效率和灵活性。

ModuleFormer 是一种基于稀疏专家混合( SMoE ) 的模块化架构，包括两种不同类型的模块，即新的 stick-breaking 注意力专家模块和传统的 MLP 专家模块。在训练和推理过程中，根据输入的向量，不同的模块会被稀疏地激活。

与之前基于 SMoE 的模块化语言模型不同，ModuleFormer 可以通过其新的负载平衡（load balance）和负载集中（load concentration）损失函数从未经筛选的数据中诱导出模块化。

在实验中，团队发现模块化架构使得大型预训练语言模型具备了三个重要的能力：

1）效率

因为 ModuleFormer 只对每个输入只激活一小部分模块，因此可以以两倍以上的吞吐量达到与常规语言模型相同的性能。

2）可扩展性

实验表明，由于微调阶段 ModuleFormer 只需要更新一部分的模块，因此比常规语言模型更不容易发生灾难性遗忘，并且可以轻松通过新的模块扩展以学习训练数据中不包含的新知识。

3）模块特异化和筛选

在微调阶段 ModuleFormer，新提出的负载集中损失函数可以自动筛选一部分模块，让它们专注于目标任务，而与任务无关的模块可以被直接抛弃掉以实现轻量化部署。

导言

尽管现代大型语言模型（LLM）在某些任务上取得了显著的成果，甚至超过了人类的表现，但其效率和灵活性仍然不高。

大多数 LLM（例如 Llama，Falcon）在推理和训练过程中都使用了它们的全部参数，我们称这些模型为密集模型。

然而，先前的研究已经表明，在执行任何特定任务时，神经模型中的大部分参数都可以在不影响模型性能的前提下被剪枝掉。

此外，一旦训练完成，LLM 就会 " 定格在某个时间点 "，但许多实际应用情况要求 LLM 具有最新的知识。因此模型进行 continue leanring 的扩展能力也十分重要。

随着模型规模的增长，为了领域适应或持续学习而对整个模型进行微调变得代价高昂且计算资源受限，这使得那些计算预算较小的用户无法实施。同时，更新所有参数也使得模型容易遭受灾难性遗忘（catastrophic forgetting）。

为此，像 LoRA 这样仅更新一小部分原始参数修正量的轻量级适应方法正在变得流行。

然而实验表明，这种方法仍然可能遭受灾难性遗忘影响，并且 LoRA 并不擅长需要模型学习大量新知识的场景，比如让模型学习一种新的语言。

文章作者认为模块化是解决前述问题的一个好方法。模块化模型具有以下几个优点：

模型可以在输入或任务上激活一组模块条件，从而比密集激活整个模型需要更少的计算量；

在给定领域或任务的情况下，可以组装一组与领域 / 任务相关的模块，形成一个新的轻量级模型；

模型可以轻松添加新的模块进行领域适应或持续学习；

模型可能更不容易发生灾难性遗忘，因为只有与输入相关的模块在模型微调期间进行更新。

这篇论文提出了一种新的模块化架构，ModuleFormer（图 1a），以及在其中进行模块操作的方法。ModuleFormer 每层包含一个组前馈神经网络（MLP）专家模块和一组新提出的 Stickbreaking 注意力专家模块。

为了平衡在训练过程中不同模块的负载，作者提出了一种新的互信息损失函数。此外，文章还展示了如何在 ModuleFormer 中插入新模块（图 1b）和进行模块修剪（图 1c）。

为了实现模块修建，文章引入了一种新的负载集中损失函数，用于在微调的过程中自动选择最适合执行给定任务的模块，同时通过微调进一步增强了这些模块执行该任务的能力。

△图 1

实验结果显示 ModuleFormer 相对于密集模型在以下方面有显著的提升：

由于稀疏的模块激活机制（图 1a）它在更低的延迟（50％）和更小的内存占用下实现了与密集 LLM 相同的性能，因此 ModuleFormer 可以实现密集模型两倍的吞吐量。

在对新领域进行微调后，它收到灾难性遗忘的影响较低，并且也可以轻松地通过添加新模块来学习新语言和知识。（图 1b）。

它可以在下游任务上进行微调，将一部分模块特异化成为处理该任务专用的模块，而未使用的模块可以被扔掉而不会牺牲模型在该任务上的性能（图 1c）。

详解 ModuleFormer 模型 Sparse Mixture of Experts（SMoE）的基本结构

SMoE 最早由 Shazeer 在《Outrageously large neural networks: The sparsely-gated mixture-of-experts layer》一文中提出。

每个 SMoE 层包含一组用于处理输入产生输出的专家模块 m_1, m_2, … , m_n，以及一个用于挑选专家的路由函数 g。在本篇文章中，路由函数由一个一层的神经网络来建模：

在给定一个输入向量 x 之后，路由函数 g 会计算一个专家模块的概率分布 g ( m|x ) ，然后模型会自动选择分布中前 top k 的专家模块来处理输入 x。SMoE 层的输出就是专家模块的输出通过对应的路由概率家和得到：其中不再 topk 内的专家模块不会被计算，同时对应的概率 g 也会被用 0 替代。

在 ModuleFormer 中，作者使用了两种不同的专家模块来分别构建传统 transformer 解码器中的自注意力层（self attention）和前馈层（MLP）。其中前馈层的专家模块和常见的 SMoE 专家模块一致，是一个单一隐藏层的全连通神经网络。而自注意力层的专家模块则是一个新提出的 stick-breaking 注意力模块。

Stick-Breaking 注意力模块

Stick-Breaking 自注意力是设计用于取代 Transformer 解码器的自注意力层，使用狄利克雷过程中的 Stick-Breaking（折棍子）过程对每个字符 xt 与之前的字符 x

Stick-Breaking 自注意力的优势在于它会自动关注最近的相关字符，而不需要引入额外的位置信息，比如 position embedding 和 relative position bias。

给定一个包含 t 个时间步的输入向量序列 x1, x2, … , xt，每个输入被投影到一系列 key 向量 k1, k2, … , kt 和一系列 value 向量 v1, v2, … , vt。为了计算时间 t 的注意力，输入 x_t 被投影到一个查询向量 q_t = W_q x_t，其中 W_q 是查询投影矩阵。对于所有之前的步骤和当前步骤 i ≤ t，计算时间步 i 的键与时间步 t 的查询匹配的概率 :

需要注意的是，这个查询概率使用了 sigmoid 激活函数，所以没有归一化。接下来通过 stick-breaking 过程来对查询概率进去归一化：这样，注意力就会自动分配给离 t 时刻最近，且具有较大查询概率的时刻。使得自注意力机制在没有额外的位置信息的情况下，也能对于相对位置进行有效的建模。最终，自注意力模块的输出是由注意力权重对历史的 value 向量进行加和并且投影得到：ModuleFormer 中的模块控制预训练中的负载均衡

为了避免 SMoE 反复使用相同的模块并浪费其他模块的额外容量，一般采用负载平衡损失函数来调节每个专家的使用频率。与之前的 SMoE 模型 不同，团队希望最大化输入字符和模块之间的互信息（MI）：

为了简化起见，假设在批次 X 中的令牌分布是均匀的，因此 p ( x ) = 1/X。在去除所有常数成分后，可以简化互信息损失（公式 6）为 p ( m ) 的熵与 p ( m | x ) 的条件熵之间的差异。在上述内容中，p ( m ) = sum_x ( g ( m|x ) p ( x ) ) ，其中 p ( x ) 是批处理中每个字符的概率，H ( m ) 是模块分布的边际熵，H ( m | x ) 是模块在给定输入字符 x 的条件下的熵，|X | 是输入字符的数量。对于长度为 T 的 batch 大小为 B 的小批量，字符的数量是 |X | = BT，字符的概率是 p ( x ) = 1/|X |。

直观地说，互信息损失最大化了模块的概率分布的边际熵，并最小化了给定输入 x 的模块条件分布的商。它平衡了整个 batch 中每个专家的负载（最大化 H ( m ) ），同时也鼓励每个输入 x 将其路由概率集中在较少的模块上（最小化 H ( m | x ) ）。

微调中的负载集中

尽管团队希望在预训练期间最大限度地利用每个专家的能力，但在微调期间希望将少量的模块专注于下游任务。这样可以移除未使用的模块并减少微调后模型的参数数量。为了将负载集中在较少的模块上，团队引入了一个新的负载集中损失函数来最小化模块的边际熵：

这样可以鼓励模型使用更少的模块来处理下游任务。在微调后，可以计算在训练或验证集上使用的模块频率 f_m。f_m 代表了模块 m 对于这个任务的重要性，可以通过移除 f_m 小于某个特定阈值的专家来轻松实现模型剪枝。

用新的模块来学习新的知识

对于模块化模型来说，插入新模块是一种直接且参数高效的方法，可以在不对整个模型进行微调的情况下学习新知识。当向每一层插入 N_new 个随机初始化的模块时，还需要扩展路由器（方程 2 中的 A）中的模块嵌入层 A，使其包含一个形状为（N_new，D_rtr）的新矩阵 A ’。因此，

新的路由函数可以写成：

由于在微调期间其他的模块参数被冻结，因此使用新模块进行持续学习可以在很大程度上避免灾难性遗忘问题。

然而，灾难性遗忘仍然可能影响路由函数。当新模块在一个新领域进行训练时，如果路由函数错误地将来自旧领域的输入路由到新专，模型可能会遭受灾难性遗忘。

为了避免这种情况，团队对路由函数进行了正则化以避免灾难性遗忘，并提出了两种训练策略：

1）全面微调路由，公式 9 中 A 和 B 使用预训练参数进行初始化，而 A ’则是随机初始化的。这个策略是为了训练数据中同时包含新旧数据的情况设计。

2）只训练 A ’，这个策略是为了连续学习（lifelong learning）的情况而设计的，不使用以前训练过的数据。由于这种情况可能导致新的模块使用频率过高，从而带来灾难性遗忘。团队引入了正则项来限制 A ’的范数：

与被指出存在缺陷的传统连续学习正则化方法（如衰减或 L2 损失）不同，路由正则化不限制专家的能力，而只限制对新专家的使用趋势。

评估

基于 ModuleFormer，研究者在 Pile 数据集上预训练了三个不同体积和计算量的 ModuleFormer Language Model（MoLM）语言模型：

基础性能评估

团队使用 Language Model Evaluation Harness 来评估零样本、少样本和语言建模任务中的语言模型。

对于零样本和少样本任务，目标是在给定上下文的基础上从一组给定选项中选择最合适的完成部分。最终选择在给定上下文下具有最高可能性的完成部分。

对于语言建模，在 Wikitext 数据集上进行测试。目标是最小化下一个标记预测的困惑度。

对于代码生成，在 HumanEval 数据集上评估模型。HumanEval 包含 164 个手写的 Python 编程问题。模型需要根据任务描述提示完成一个函数，以便能够通过所有提供的测试案例。

表 2 和表 3 显示了 MoLM 和基准语言模型在常识推理、闭卷问答和代码生成基准上的性能。

总体而言，MoLM-4B-K2 模型的性能与大约 13 亿参数的稠密模型相当，MoLM-4B-K4 和 MoLM-8B-K2 模型的性能与大约 27 亿参数的稠密模型相当。

由于其稀疏计算结构，MoLM 处理每个字符的激活参数仅（等同于计算量）相当于同等性能稠密模型的约 25%。因此，它减少了 50% 的延迟，同时具有较低的内存使用峰值，并在 GPU 内存完全占用时将吞吐量提高了 2 倍。

通过增加模块学习新语言

在本节中，我们测试了模型学习新语言的能力。主要研究两种实验设置：连续联合预训练（continual joint pre-training）和连续终身预训练（continual lifelong pre-training）。

它们的区别在于是否有英文文本的存在。对于这两种设置，我们通过在 CC-100 语料库上进行语言模型任务，不断地对 ModuleFormer 和 GPT-Neo 进行预训练。为了评估质量，我们采用了由 XGLM 和 mGPT 引入的 0-shot 方法的 mLAMA 基准测试。

持续联合预训练：在这部分中，我们对联合训练的模型进行持续预训练。具体而言，我们混合了英语和一种新语言来构建一个新的训练语料库，并保持嵌入层可训练。联合训练 [ Caruana, 1997 ] 是一种众所周知的多任务学习方法，展示了对旧任务和新任务的熟练掌握。然而，它经常在不同任务之间产生负面干扰。

表 4 显示了持续训练模型获得的结果。表格揭示了以下发现：

1）我们观察到稀疏模型在 Fully Tuned 的情况下经历较少干扰，最终得到了最好的的性能；

2）ModuleFormer 通过增加模块（Insert New Expert）的能力，比之前的 LoRA 方法展示出了更好的少量参数（Parameter Efficient）调优的能力。这些结果表明，稀疏架构带来了更强的抗遗忘能力。

持续终身预训练：对于这个实验设定，模型仅在新语言文本上进行训练。Abraham 和 Robins [ 2005 ] 提出了稳定性 - 可塑性困境，这解释了模型面临的一个困难挑战：

1）模型应具有较高的可塑性以学习新语言，

2）模型必须具有出色的稳定性，考虑到在众多的训练迭代中不会接触到任何英语标记。

表 5 显示了 LoRA 基准和我们的方法在不同的路由正则化损失权重下的结果。我们的 ModuleFormer 借助路由正则化损失表现出了强大的平衡稳定性和可塑性的能力。

当我们通过增加损失权重来限制新专家的使用时，模型获得了稳定性，但可塑性下降。相比之下，使用 LoRA 对 GPT-Neo 进行微调在稳定性和可塑性方面都落后。

相比于 1.33 亿可训练参数的高秩 LoRA，低秩 LoRA（减少训练参数到 2400 万）和基本正则化都无法改善稳定性。

微调和压缩模型

在本节中，我们展示了 ModuleFormer 中的模块可以被快速移除，以创建一个在尺寸上更小但性能不受损的任务专用模型。

我们首先从 GitHub-code-clean 数据集中创建了一个包含 150 亿个字符的子集，该子集只包含 Python 代码。然后，我们使用负载集中损失函数（权重为 0.001）对 MoLM-4B-K2 模型在该数据集上进行精调。

在精调之后，我们在从精调数据集中随机抽样的小型评估集上，计算每个专家的激活频率，然后通过将每层除以层内最大频率来进行归一化。之后，我们设定一个阈值 τ，并修剪了所有归一化频率低于该阈值的模块。

我们在 HumanEval 数据集上测试了我们修剪后的 MoLM-4B-K2 模型。

图 2a 说明了 pass@k 指标与剩余参数比例之间的相关性。图 2b 展示了剩余参数比例与阈值之间的关联。我们观察到：

1）修剪不必要的模块对结果影响不大。我们可以修剪 40% 至 50% 的参数而不牺牲性能。相反，适当的修剪（33%）使精调后的模型在任务上表现更好。

2）模块分布存在显著差异，大约有一半的模块的激活频率低于最常使用的专家的 0.3%。这个结果显示了负载集中损失函数的有效性。

总结

在这篇论文中，我们提出了一种新的模块化架构 ModuleFormer，以及与之相关的模块操作方法。

ModuleFormer 包括几个新组件：新的 Stickbreaking 注意力机制、新的互信息负载平衡损失函数用于预训练，以及新的负载集中损失函数用于微调。

基于 ModuleFormer，我们预训练了一个新的语言模型 MoLM。我们的实验结果显示了 MoLM 的相对于稠密 LLM 展现出了一些新的能力：

1）它在更低的延迟（50%）和更小的内存占用下实现了与密集 LLM 相同的性能；从而提高了吞吐量超过 2 倍；

2）在对整个模型进行微调以适应新领域后，它对灾难性遗忘的鲁棒性较强，并且也可以轻松扩展以学习新的语言和知识；

3）它可以在下游任务上进行微调，以使一部分模块专注于任务，并且未被任务使用的模块可以被修剪而不影响性能。

论文地址：

https://arxiv.org/abs/2306.04640