多任務學習（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 里的權重網絡得到，具體可以參見論文。

最大化任務和專家之間的 mutual information 之后，模型就可以讓專家和任務擁有稀疏且非常強的依賴關系，如圖 3 所示。最左邊的就是 Mod-Squad 的任務使用專家頻率。可以看出，Mod-Squad 的任務和專家之間有著更稀疏但尖銳的頻率。

圖 3. 任務使用不同專家的頻率圖對比。橫軸是不同的專家，縱軸是不同的 task，顏色深代表更高的使用頻率。Mod-Squad 的頻率圖更加稀疏且尖銳。

這個任務和專家之間稀疏且非常強依賴關系的好處就是：

1. 相近的任務傾向于使用同一個專家；

2. 專家傾向于被一組正相關的任務使用；

3. 模型的容量被全部使用，但每個任務只使用部分容量，可以根據任務調整使用容量；

4. 可以針對特定任務從多任務大模型中提取出單任務小模型，并具有和大模型一樣的性能。這個特性能用于從超大多任務模型中提取出單任務小模型。

根據任務之間分享專家的頻率，模型還可以算出任務之間的相似性，如下圖所示。可以看出，偏 3D 的任務之間更傾向于使用相同專家，因此更加相似。

實驗部分

Mod-Squad 可以在不損失精度的情況下針對單一任務進行剪枝，下圖縱軸是性能，橫軸是參數量。

在大數據集 Taskonomy 上也有很大的提升，可以看到，Mod-Squad 比單純的 MTL 平均高了 2.8 個點，并且在剪枝以后保持著一樣的性能。

在 PASCAL-Context 上跟其他方法的對比，Mod-Squad 比其他 MoE 方法平均高出了接近兩個點。

具體細節可以參照原文。