多任務(wù)學(xué)習(xí)（MTL）存在很多挑戰(zhàn)，因?yàn)椴煌蝿?wù)之間的梯度可能矛盾。為了利用任務(wù)之間的關(guān)聯(lián)，作者引入了 Mod-Squad 模型，它是多個專家組成的模塊化模型。模型可以靈活優(yōu)化任務(wù)和專家的匹配，針對任務(wù)選擇部分專家。模型讓每一個專家只對應(yīng)部分任務(wù)，每一個任務(wù)只對應(yīng)部分專家，以此最大化利用任務(wù)之間的正向聯(lián)系。Mod-Squad 整合了 Mixture of Expert (MoE) 層到 Vision Transformer 模型中，并引入了新的損失函數(shù)鼓勵專家和任務(wù)之間的稀疏但強(qiáng)烈的依賴關(guān)系。此外，對于每個任務(wù)，模型都可以只保留小部分專家網(wǎng)絡(luò)，并且性能與原來的大模型相同。模型在 13 個視覺任務(wù)的 Taskonomy 大數(shù)據(jù)集和 PASCALContext 數(shù)據(jù)集上取得了最佳效果。

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

項(xiàng)目地址：https://vis-www.cs.umass.edu/mod-squad/

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

多任務(wù)學(xué)習(xí)（MTL）的目的是建模任務(wù)之間的關(guān)系，并為多種任務(wù)構(gòu)建統(tǒng)一的模型。如圖 1 所示，Mod-Squad 的主要動機(jī)就是要讓專家只被一些任務(wù)更新而不是所有任務(wù)，且每一個任務(wù)只更新部分專家。這樣可以利用模型的全部容量的同時避免任務(wù)間的互相干擾。

圖 1.Mod-Squad: 專家和任務(wù)互相選擇。MoE ViT: 所有專家都被所有任務(wù)使用。

下面簡單介紹下該文章。

模型結(jié)構(gòu)

圖 2.Mod-Squad: 將專家組 (mixture-of-expert) 插入到 Vision Transformer.

如圖 2 所示， Mod-Squad 的結(jié)構(gòu)就是將 Mixture-of-expert (MoE) 引入 Vision Transformer (ViT)。MoE 是一種機(jī)器學(xué)習(xí)模型，其中多個專家組成了一個混合模型。每個專家都是一個獨(dú)立的模型，并且每個模型對于不同的輸入有不同的貢獻(xiàn)。最后，所有專家的貢獻(xiàn)被加權(quán)并組合在一起以得到最終的輸出。這種方法的優(yōu)勢在于它可以根據(jù)輸入圖像的內(nèi)容動態(tài)地選擇最佳的專家并且控制計(jì)算量。

之前的 MoE 模型收斂后，可以根據(jù)不同圖片使用不同的專家，但是針對某個任務(wù)，模型會收斂到傾向于使用全部專家。Mod-Squad 可以做到讓模型針對圖片來使用不同的專家，并且模型可以在收斂后，達(dá)到一個任務(wù)只使用一部分專家的狀態(tài)。接下來介紹這是怎么實(shí)現(xiàn)的。

最大化專家和任務(wù)之間的 mutual information

本文提出了一個任務(wù)和專家的聯(lián)合概率模型來優(yōu)化專家 E 和任務(wù) T 之間的分配。這個概率模型會用來計(jì)算專家和任務(wù)之間的 mutual information，并作為額外的損失函數(shù)來優(yōu)化 MoE 里的權(quán)重網(wǎng)絡(luò)。Mutual information 公式如下，E 和 T 的概率可以由 MoE 里的權(quán)重網(wǎng)絡(luò)得到，具體可以參見論文。

最大化任務(wù)和專家之間的 mutual information 之后，模型就可以讓專家和任務(wù)擁有稀疏且非常強(qiáng)的依賴關(guān)系，如圖 3 所示。最左邊的就是 Mod-Squad 的任務(wù)使用專家頻率?？梢钥闯觯琈od-Squad 的任務(wù)和專家之間有著更稀疏但尖銳的頻率。

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

這個任務(wù)和專家之間稀疏且非常強(qiáng)依賴關(guān)系的好處就是：

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

2. 專家傾向于被一組正相關(guān)的任務(wù)使用；

3. 模型的容量被全部使用，但每個任務(wù)只使用部分容量，可以根據(jù)任務(wù)調(diào)整使用容量；

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

根據(jù)任務(wù)之間分享專家的頻率，模型還可以算出任務(wù)之間的相似性，如下圖所示?？梢钥闯?，偏 3D 的任務(wù)之間更傾向于使用相同專家，因此更加相似。

實(shí)驗(yàn)部分

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

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

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

具體細(xì)節(jié)可以參照原文。