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

2023-04-02 04:04

机器之心专栏

机器之心编辑部

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 的主要动机就是要让专家只被一些任务更新而不是所有任务，且每一个任务只更新部分专家。这样可以利用模型的全部容量的同时避免任务间的互相干扰。

图 1.Mod-Squad: 专家和任务互相选择。MoE ViT: 所有专家都被所有任务使用。

下面简单介绍下该文章。

模型结构

图 2.Mod-Squad: 将专家组 (mixture-of-expert) 插入到 Vision Transformer.

如图 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 里的权重网络得到，具体可以参见论文。