模塊化MoE將成為視覺(jué)多任務(wù)學(xué)習(xí)基礎(chǔ)模型

作者：機(jī)器之心 2023-04-03 10:32:56

UMass Amherst 淦創(chuàng)團(tuán)隊(duì)提出了 Mod-Squad 模型，它可以從多任務(wù)大模型中提取針對(duì)單一任務(wù)的相同性能小模型，在 Taskonomy 大數(shù)據(jù)集和 PASCALContext 數(shù)據(jù)集上取得了最佳效果。

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

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

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

模型結(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í)模型，其中多個(gè)專家組成了一個(gè)混合模型。每個(gè)專家都是一個(gè)獨(dú)立的模型，并且每個(gè)模型對(duì)于不同的輸入有不同的貢獻(xiàn)。最后，所有專家的貢獻(xiàn)被加權(quán)并組合在一起以得到最終的輸出。這種方法的優(yōu)勢(shì)在于它可以根據(jù)輸入圖像的內(nèi)容動(dòng)態(tài)地選擇最佳的專家并且控制計(jì)算量。

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

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

本文提出了一個(gè)任務(wù)和專家的聯(lián)合概率模型來(lái)優(yōu)化專家 E 和任務(wù) T 之間的分配。這個(gè)概率模型會(huì)用來(lái)計(jì)算專家和任務(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ù)和專家之間的 mutual information 之后，模型就可以讓專家和任務(wù)擁有稀疏且非常強(qiáng)的依賴關(guān)系，如圖 3 所示。最左邊的就是 Mod-Squad 的任務(wù)使用專家頻率。可以看出，Mod-Squad 的任務(wù)和專家之間有著更稀疏但尖銳的頻率。