学习在专业化专家之间进行路由实现零样本泛化

24年6月来自MIT-IBM实验室、多伦多大学、Vector Institute和UNC的论文“Learning to Route Among Specialized Experts for Zero-Shot Generalization”。

最近，通过参数高效微调（PEFT）专门用于特定任务或领域的“专家”语言模型得到了广泛采用。如何才能从大量专家语言模型提高对未见过任务的零样本泛化能力？这项工作提出事后因果自适应逐token的门控（PHATGOOSE），它可以学习在PEFT生成的专业化模块之间进行路由。与过去学习在专业化模型之间进行路由的方法不同，PHATGOOSE 探索了如果可以为每个token和模型的每一层自适应地选择不同的专家，那么零样本泛化能力将得到改善的可能性。至关重要的是，该方法是事后因果方法 - 它不需要同时访问用于创建专业化模型的数据集，并且在每个专家模型训练完成后只需要少量额外的计算。在涵盖一系列专业化模型和零样本泛化基准的实验中，PHATGOOSE 优于过去的事后因果路由方法，在某些情况下，优于显式多任务训练（需要同时访问数据）。

高性能预训练语言模型的出现，导致大量经过微调的“专家”模型出现，这些模型专门用于特定任务或领域。许多此类专家模型都是通过参数高效微调 (PEFT) 技术创建的（Ding，2022；Lialin，2023；He，2021），该技术添加仅引入或修改少量参数的小“模块”, 例如低秩适配器（Hu，2021）或 (IA)3 向量（Liu，2022）, 生成微调模型。由于规模小，专用的 PEFT 模块可以轻松共享，这导致各种平台上的适配器数量不断增加——例如，目前已有超过 17,000 个基于 PEFT 库 (Mangrulkar，2022) 的适配器上传到 Hugging Face Model Hub。这些 PEFT 模块的可用性，使得一件事变得容易，即以低成本和简单的方式将给定的预训练模型模块化地适配到特定任务或域。

与此同时，大语言模型 (LLM) 现在被视为“通用”AI 系统，可以执行任何任务而无需任何特定于任务的训练或调整。这种方法源于这样的观察：LLM 通常表现出强大的零样本泛化能力，即能够执行未明确训练的新任务。这种零样本泛化能力通常通过大规模多任务微调（也称为“指令调整”）来改进（Sanh，2021；Wei，2021；Mishra，2022）。依赖零样本泛化与上述为每个任务训练专门模型的方法（通过 PEFT 或其他方式）形成鲜明对比。

通用语言模型的吸引力和基于 PEFT 专业化模型的激增，提出了一个自然的问题：能否利用大量专业化模块来改进基础语言模型的零样本泛化？这种方法很有吸引力，原因如下：首先，它将为通用语言模型的去中心化开发提供一条途径，否则需要大量的中心化计算（Kaplan，2020；Hoffmann，2022）；此外，它将提供一种方法来回收已经为创建专业化模型而花费的广泛努力和计算。鉴于大量证据表明多任务训练可以提高零样本泛化能力 (Sanh, 2021; Wei, 2021; Mishra, 2022)，可能希望这种方法能够成功，并且结合专业化模型可以看作是一种不需要同时访问数据的多任务学习形式。

为了解决这个问题，大多数过去的工作 (Jang, 2023; Belofsky, 2023; Durbin, 2024; Maxine, 2023) 将输入查询的嵌入与用于训练每个专家的每个数据集中示例的平均嵌入进行比较，学习事后因果路由策略。这种方法隐式假设有一个非常适合查询的专家，并希望检索算法能够准确识别这个最佳专家。然而，Jang (2023) 表明，这种方法落后于始终为给定查询选择最佳专家的“神谕”路由器。为了探索替代的路由方法，许多 PEFT 方法通常会在模型的许多地方插入小型可训练模块，例如在每个权重矩阵中（Hu，2021）。同时，许多稀疏门控的混合专家模型会针对每个 token 分别做出路由决策（Shazeer，2016；Fedus，2022；Du，2022）。

这项工作的目标是开发一种适用问题的方法——即，它回收来自各个LLM模型贡献者的 PEFT 模块，改进零样本泛化，而无需大量额外工作或同时访问贡献者的数据集。此外，这个工作的目标是开发一种方法，遵循以下假设：学习一个对每个token和每个模块的路由策略，可以实现更好的零样本泛化，这是相比一个从所有tokens选择单个专家模型的路由策略而言。提出的PHATGOOSE，让贡献者训练特定于任务的门控，然后使用每个门控的参数执行离散 top-k 路由来实现这些目标。该过程如图所示，其说明在一个插入模块的一个层上如何学习路由：通常，基于 PEFT 的模型会在整个模型各个层中引入许多这样的模块；左图是在训练一个专门的模块（此处显示为 LoRA模块）后，它会被冻结，并训练一个 S 形门控来选择应该将哪些激活输入到模块中；右图是在推理过程中，根据归一化门控和激活之间的点积计算路由分布（显示为条形图）；然后通过根据这个路由分布，选择模块来执行 Top-k 路由。

在提示数据集上对 T5 (Raffel，2020) 进行显式多任务训练可产生在未见过任务上具有强大零样本性能的模型。这已成为对零样本泛化进行基准测试的常见实验设置，例如 (Chung，2022；Longpre，2023；Jang，2023；Zhou，2022) 。具体来说，实验中，作为基础模型，用 LM 适配的 T5.1.1 XL (Lester，2021)，这是 T5 语言模型 (Raffel，2020) 的 3B 参数变型，使用标准的语言建模目标在 C4 数据集上进行额外的 100K 步训练。

为了创建用于路由的专家模块池，考虑两个数据集。对于第一个（“T0 Held-In”），用与训练 T0（Sanh，2021）相同的 36 个保留提示数据集和任务集。对于第二个（“FLAN”），考虑大型 FLAN 提示数据集集合（Longpre，2023）。FLAN Collection 使用来自 SuperGLUE（Wang，2019）、Super Natural Instructions（Wang，2022b）、对话数据集和思维链（Wei，2022）数据集的零样本和少量样本提示数据集扩展 T0 Held-in。只考虑那些具有零样本提示格式的数据集，从而从 FLAN Collection 中得出总共 166 个专业化的模型。

对于 PEFT 模块，专注于LoRAs，但PHATGOOSE 不需要使用 LoRA，预计它同样适用于其他模块架构（例如 (IA)3、适配器等）。为每个数据集训练一个 PEFT 模块，从而产生两种设置，分别为 T0 Held-In 和 FLAN 提供 36 个或 166 个专家模型。

考虑三个零样本泛化基准进行评估。对于第一个（“T0HO”），用与评估 T0（Sanh，2021）相同的保留数据集。由于FLAN 集合包含来自 T0 的保留数据集，因此在使用 FLAN 专家池时，不会对 T0HO 进行评估。对于第二和第三种情况，考虑 BIG-bench 的两个变型（BIG-bench，2023），这是一个社区策划的数据集，用于衡量模型的不同能力，如推理、创造力、偏见等。具体来说，在 BIG-bench Hard（Suzgun，2022）和 BIG-bench Lite（BIG-bench，2023）上进行评估。BIG-Bench Hard (BBH) 是 23 个数据集的集合，在这些数据集上，最先进的模型的表现明显差于人类。BIG-Bench Lite (BBL) 包含 24 个不同的数据集，旨在作为完整 BIG-Bench 基准的轻量级智体。由于 T5 token化器无法token化 BIG-bench 集合中的某些数据集，因此在评估期间排除在外。在所有情况下，从 Hugging Face Hub 获取数据集。

虽然 PHATGOOSE 不要求不同的贡献者使用相同的超参，但为了简单起见，用 AdamW (Loshchilov & Hutter, 2017) 优化器对每个数据集的 1000 步进行秩 r = 16 LoRA 训练，每个批次包含 1024 个最大长度为 512 的序列，学习率为 5e-3，预热率为 0.06。在验证步以 100 步的粒度执行检查点选择。对于 PHATGOOSE，在训练每个模块后，冻结所有参数并使用相同的超参对门控向量进行额外的 100 步训练。按照过去工作中的标准做法，用 k = 2 进行 top-k 路由。