多任務(wù)學(xué)習(xí)（MTL）存在很多挑戰(zhàn)，因?yàn)椴煌蝿?wù)之間的梯度可能矛盾。為了利用任務(wù)之間的關(guān)聯(lián)，作者引入了 Mod-Squad 模型，它是多個(gè)專(zhuān)家組成的模塊化模型。模型可以靈活優(yōu)化任務(wù)和專(zhuān)家的匹配，針對(duì)任務(wù)選擇部分專(zhuān)家。模型讓每一個(gè)專(zhuān)家只對(duì)應(yīng)部分任務(wù)，每一個(gè)任務(wù)只對(duì)應(yīng)部分專(zhuān)家，以此最大化利用任務(wù)之間的正向聯(lián)系。Mod-Squad 整合了 Mixture of Expert (MoE) 層到 Vision Transformer 模型中，并引入了新的損失函數(shù)鼓勵(lì)專(zhuān)家和任務(wù)之間的稀疏但強(qiáng)烈的依賴(lài)關(guān)系。此外，對(duì)于每個(gè)任務(wù)，模型都可以只保留小部分專(zhuān)家網(wǎng)絡(luò)，并且性能與原來(lái)的大模型相同。模型在 13 個(gè)視覺(jué)任務(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 的主要?jiǎng)訖C(jī)就是要讓專(zhuān)家只被一些任務(wù)更新而不是所有任務(wù)，且每一個(gè)任務(wù)只更新部分專(zhuān)家。這樣可以利用模型的全部容量的同時(shí)避免任務(wù)間的互相干擾。

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

下面簡(jiǎn)單介紹下該文章。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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