CVPR 2023 | 模块化MoE将成为视觉多任务学习基础模型

UMass Amherst 淦创团队提出了 Mod-Squad 模型，它可以从多任务大模型中提取针对单一任务的相同性能小模型，在 Taskonomy 大数据集和 PASCALContext 数据集上取得了最佳效果。

多任务学习（MTL）存在很多挑战，因为不同任务之间的梯度可能矛盾。为了利用任务之间的关联，作者引入了 Mod-Squad 模型，它是多个专家组成的模块化模型。模型可以灵活优化任务和专家的匹配，针对任务选择部分专家。模型让每一个专家只对应部分任务，每一个任务只对应部分专家，以此最大化利用任务之间的正向联系。Mod-Squad 整合了 Mixture of Expert (MoE) 层到 Vision Transformer 模型中，并引入了新的损失函数鼓励专家和任务之间的稀疏但强烈的依赖关系。此外，对于每个任务，模型都可以只保留小部分专家网络，并且性能与原来的大模型相同。模型在 13 个视觉任务的 Taskonomy 大数据集和 PASCALContext 数据集上取得了最佳效果。

论文地址：https://arxiv.org/abs/2212.08066

项目地址：https://vis-www.cs.umass.edu/mod-squad/

Github地址：https://github.com/UMass-Foundation-Model/Mod-Squad

多任务学习（MTL）的目的是建模任务之间的关系，并为多种任务构建统一的模型。如图 1 所示，Mod-Squad 的主要动机就是要让专家只被一些任务更新而不是所有任务，且每一个任务只更新部分专家。这样可以利用模型的全部容量的同时避免任务间的互相干扰。

下面简单介绍下该文章。

模型结构

如图 2 所示， Mod-Squad 的结构就是将 Mixture-of-expert (MoE) 引入 Vision Transformer (ViT)。MoE 是一种机器学习模型，其中多个专家组成了一个混合模型。每个专家都是一个独立的模型，并且每个模型对于不同的输入有不同的贡献。最后，所有专家的贡献被加权并组合在一起以得到最终的输出。这种方法的优势在于它可以根据输入图像的内容动态地选择最佳的专家并且控制计算量。

之前的 MoE 模型收敛后，可以根据不同图片使用不同的专家，但是针对某个任务，模型会收敛到倾向于使用全部专家。Mod-Squad 可以做到让模型针对图片来使用不同的专家，并且模型可以在收敛后，达到一个任务只使用一部分专家的状态。接下来介绍这是怎么实现的。

最大化专家和任务之间的 mutual information

本文提出了一个任务和专家的联合概率模型来优化专家 E 和任务 T 之间的分配。这个概率模型会用来计算专家和任务之间的 mutual information，并作为额外的损失函数来优化 MoE 里的权重网络。Mutual information 公式如下，E 和 T 的概率可以由 MoE 里的权重网络得到，具体可以参见论文。

最大化任务和专家之间的 mutual information 之后，模型就可以让专家和任务拥有稀疏且非常强的依赖关系，如图 3 所示。最左边的就是 Mod-Squad 的任务使用专家频率。可以看出，Mod-Squad 的任务和专家之间有着更稀疏但尖锐的频率。

这个任务和专家之间稀疏且非常强依赖关系的好处就是：

1. 相近的任务倾向于使用同一个专家；

2. 专家倾向于被一组正相关的任务使用；

3. 模型的容量被全部使用，但每个任务只使用部分容量，可以根据任务调整使用容量；

4. 可以针对特定任务从多任务大模型中提取出单任务小模型，并具有和大模型一样的性能。这个特性能用于从超大多任务模型中提取出单任务小模型。

根据任务之间分享专家的频率，模型还可以算出任务之间的相似性，如下图所示。可以看出，偏 3D 的任务之间更倾向于使用相同专家，因此更加相似。

实验部分

Mod-Squad 可以在不损失精度的情况下针对单一任务进行剪枝，下图纵轴是性能，横轴是参数量。

在大数据集 Taskonomy 上也有很大的提升，可以看到，Mod-Squad 比单纯的 MTL 平均高了 2.8 个点，并且在剪枝以后保持着一样的性能。

在 PASCAL-Context 上跟其他方法的对比，Mod-Squad 比其他 MoE 方法平均高出了接近两个点。

具体细节可以参照原文。