[论文解读] Sparsely Activated Mixture-of-Experts are Robust Multi-Task Learners

Shashank Gupta, Subhabrata Mukherjee+5 more|arXiv (Cornell University)|Apr 16, 2022

Domain Adaptation and Few-Shot Learning被引用 8

一句话总结

本文提出一种针对稀疏激活的专家混合模型（MoE）的任务感知门控机制，以提升自然语言处理中的多任务学习性能。通过将不同任务的 token 路由至特定任务的专家，同时共享部分权重，该模型实现了更高的参数效率与鲁棒性——在低资源任务上的迁移能力更强，对未见任务的样本高效泛化能力更优，并且相比密集模型更能抵抗灾难性遗忘。

ABSTRACT

Traditional multi-task learning (MTL) methods use dense networks that use the same set of shared weights across several different tasks. This often creates interference where two or more tasks compete to pull model parameters in different directions. In this work, we study whether sparsely activated Mixture-of-Experts (MoE) improve multi-task learning by specializing some weights for learning shared representations and using the others for learning task-specific information. To this end, we devise task-aware gating functions to route examples from different tasks to specialized experts which share subsets of network weights conditioned on the task. This results in a sparsely activated multi-task model with a large number of parameters, but with the same computational cost as that of a dense model. We demonstrate such sparse networks to improve multi-task learning along three key dimensions: (i) transfer to low-resource tasks from related tasks in the training mixture; (ii) sample-efficient generalization to tasks not seen during training by making use of task-aware routing from seen related tasks; (iii) robustness to the addition of unrelated tasks by avoiding catastrophic forgetting of existing tasks.

研究动机与目标

• 解决多任务学习中的干扰问题，即不同任务使模型参数向相互冲突的方向更新。
• 克服密集模型的局限性，即所有输入均使用全部参数，导致参数冲突与性能下降。
• 通过设计基于任务感知的门控机制，使稀疏 MoE 模型适用于多任务学习，实现输入到特定任务专家的路由。
• 从三个关键维度评估稀疏 MoE 模型的鲁棒性：对低资源任务的迁移能力、对未见任务的泛化能力，以及对灾难性遗忘的抵抗能力。
• 证明采用任务感知路由的稀疏 MoE 模型在迁移能力、样本效率与任务增加下的稳定性方面，优于密集多任务模型。

提出的方法

• 采用基于 Transformer 的架构，并引入专家混合（MoE）层，其中每层包含多个特定任务的专家以及一个任务感知门控网络。
• 设计一种门控机制，利用特定任务的门控信号，根据任务身份将输入 token 路由至合适的专家，从而减少任务间的干扰。
• 使用预训练的 BERT 权重初始化 MoE 模型，以从强大的基础表征出发，实现高效的多任务微调。
• 采用 top-1 路由并结合负载均衡策略，在保持计算效率的同时实现条件计算与高参数容量。
• 在共享编码器和特定任务的专家头基础上，联合训练包含多样化自然语言理解任务的混合数据集，路由决策基于每个 token 的任务身份。
• 实施一种路由策略，通过解耦各任务的路由决策，避免共享门控的竞争，实现专家专业化且无干扰。

实验结果

研究问题

• RQ1 与共享门控 MoE 或密集模型相比，稀疏 MoE 模型中采用任务感知门控是否能有效降低多任务 NLP 学习中的任务间干扰？
• RQ2 在训练混合中利用相关任务时，稀疏 MoE 模型在多大程度上能将知识迁移至低资源任务？
• RQ3 在仅使用少量样本的情况下，稀疏 MoE 模型如何利用已见任务的路由机制，实现对未见任务的有效泛化？
• RQ4 在添加无关任务时，稀疏 MoE 模型对先前学习任务的性能保持能力如何，尤其在避免灾难性遗忘方面表现如何？
• RQ5 与密集多任务模型相比，MoE 模型中使用任务感知路由是否能在多样化的自然语言理解基准上实现更优的性能与稳定性？

主要发现

• 与单任务基线相比，所提出的 MT-TaG 模型在 GLUE++ 的 Small Tasks 子集上实现了 2.28 个百分点的绝对性能提升（83.56 vs. 81.28）。
• 在完整的 All Tasks 基准上，MT-TaG 模型取得了 86.46 的平均得分，优于单任务基线（85.00），展现出对多样化任务的强大可扩展性。
• 该模型通过利用训练混合中相关任务的知识，在低资源任务（如 RTE、MRPC 和 WiC）上表现出优异的迁移性能。
• MT-TaG 能够高效泛化至未见任务，通过将来自相关已见任务的输入有效路由至合适专家，显著降低对大规模微调数据的需求。
• 即使在添加无关任务后，该模型仍能保持对先前学习任务的性能，展现出对灾难性遗忘的强鲁棒性。
• 实证结果证实，任务感知 MoE 路由在所有三个关键鲁棒性维度（迁移、泛化与遗忘抵抗）上，均显著优于共享门控 MoE 和密集多任务模型。