還記得谷歌大腦團(tuán)隊(duì)去年 6 月份發(fā)布的 43 頁(yè)論文《Scaling Vision with Sparse Mixture of Experts》嗎？他們推出了史上最大規(guī)模的視覺(jué)模型 V-MoE，實(shí)現(xiàn)了接近 SOTA 的 Top-1 準(zhǔn)確率。如今，谷歌大腦開源了訓(xùn)練和微調(diào)模型的全部代碼。

　　在過(guò)去幾十年里，深度學(xué)習(xí)的進(jìn)步是由幾個(gè)關(guān)鍵因素推動(dòng)的：少量簡(jiǎn)單而靈活的機(jī)制、大型數(shù)據(jù)集、更專業(yè)的硬件配置，這些技術(shù)的進(jìn)步使得神經(jīng)網(wǎng)絡(luò)在圖像分類、機(jī)器翻譯、蛋白質(zhì)預(yù)測(cè)等任務(wù)中取得令人印象深刻的結(jié)果。

　　然而，大模型以及數(shù)據(jù)集的使用是以大量計(jì)算需求為代價(jià)的。最近的研究表明，增強(qiáng)模型的泛化能力以及穩(wěn)健性離不開大模型的支持，因此，在訓(xùn)練大模型的同時(shí)協(xié)調(diào)好與訓(xùn)練資源的限制是非常重要的。一種可行的方法是利用條件計(jì)算，該方法不是為單個(gè)輸入激活整個(gè)網(wǎng)絡(luò)，而是根據(jù)不同的輸入激活模型的不同部分。這一范式已經(jīng)在谷歌提出的 pathway（一種全新的 AI 解決思路，它可以克服現(xiàn)有系統(tǒng)的許多缺點(diǎn)，同時(shí)又能強(qiáng)化其優(yōu)勢(shì)）愿景和最近的大型語(yǔ)言模型研究中得到了重視，但在計(jì)算機(jī)視覺(jué)中還沒(méi)有得到很好的探索。

　　稀疏門控混合專家網(wǎng)絡(luò) (MoE) 在自然語(yǔ)言處理中展示了出色的可擴(kuò)展性。然而，在計(jì)算機(jī)視覺(jué)中，幾乎所有的高性能網(wǎng)絡(luò)都是密集的，也就是說(shuō)，每個(gè)輸入都會(huì)轉(zhuǎn)化為參數(shù)進(jìn)行處理。

　　去年 6 月，來(lái)自谷歌大腦的研究者提出了 V-MoE（Vision MoE ），這是一種基于專家稀疏混合的新視覺(jué)架構(gòu)。當(dāng)應(yīng)用于圖像識(shí)別時(shí)，V-MoE 在推理時(shí)只需要一半的計(jì)算量，就能達(dá)到先進(jìn)網(wǎng)絡(luò)性能。此外，該研究還提出了對(duì)路由算法的擴(kuò)展，該算法可以在整個(gè) batch 中對(duì)每個(gè)輸入的子集進(jìn)行優(yōu)先級(jí)排序，從而實(shí)現(xiàn)自適應(yīng)圖像計(jì)算。這允許 V-MoE 在測(cè)試時(shí)能夠權(quán)衡性能和平滑計(jì)算。最后，該研究展示了 V-MoE 擴(kuò)展視覺(jué)模型的潛力，并訓(xùn)練了一個(gè)在 ImageNet 上達(dá)到 90.35% 的 150 億參數(shù)模型。

　　論文地址：https://arxiv.org/pdf/2106.05974.pdf

　　代碼地址：https://github.com/google-research/vmoe

　　V-MoE

　　谷歌大腦在 ViT 的不同變體上構(gòu)建 V-MoE：ViT-S(mall)、ViT-B(ase)、ViT-L(arge) 和 ViTH(uge)，其超參數(shù)如下：

　　ViT 已被證明在遷移學(xué)習(xí)設(shè)置中具有良好的擴(kuò)展性，在較少的預(yù)訓(xùn)練計(jì)算下，比 CNN 獲得更高的準(zhǔn)確率。ViT 將圖像處理為一系列 patch，輸入圖像首先被分成大小相等的 patch，這些 patch 被線性投影到 Transformer 的隱藏層，在位置嵌入后，patch 嵌入（token）由 Transformer 進(jìn)行處理，該 Transformer 主要由交替的自注意力和 MLP 層組成。MLP 有兩個(gè)層和一個(gè) GeLU 非線性。對(duì)于 Vision MoE，該研究用 MoE 層替換其中的一個(gè)子集，其中每個(gè)專家都是一個(gè) MLP，如下圖所示：

　　為了大規(guī)模擴(kuò)展視覺(jué)模型，該研究將 ViT 架構(gòu)中的一些密集前饋層 (FFN) 替換為獨(dú)立 FFN 的稀疏混合（稱之為專家）?？蓪W(xué)習(xí)的路由層為每個(gè)獨(dú)立的 token 選擇對(duì)應(yīng)的專家。也就是說(shuō)，來(lái)自同一圖像的不同 token 可能會(huì)被路由到不同的專家。在總共 E 位專家（E 通常為 32）中，每個(gè) token 最多只能路由到 K（通常為 1 或 2）位專家。這允許擴(kuò)展模型的大小，同時(shí)保持每個(gè) token 計(jì)算的恒定。下圖更詳細(xì)地顯示了 V-MoE 編碼器塊的結(jié)構(gòu)。

　　V-MoE Transformer 編碼器塊

　　實(shí)驗(yàn)結(jié)果

　　谷歌大腦首先在大型圖像數(shù)據(jù)集 JFT-300M 上對(duì)模型進(jìn)行一次預(yù)訓(xùn)練。

　　下圖左展示了模型在所有大小（從 small s/32 到 huge H/14）時(shí)的預(yù)訓(xùn)練結(jié)果。然后，使用一個(gè)新的 head（一個(gè)模型中的最后一層）將模型遷移至新的下游任務(wù)（如 ImageNet）。他們探索了兩種遷移設(shè)置：在所有可用的新任務(wù)示例上微調(diào)整個(gè)模型或者凍結(jié)預(yù)訓(xùn)練網(wǎng)絡(luò)并使用少量示例僅對(duì)新 head 調(diào)整（即所謂的小樣本遷移）。

　　下圖右總結(jié)了模型遷移至 ImageNet 的效果，其中每個(gè)圖像類別僅在 5 張圖像上訓(xùn)練（叫做 5-shot transfer）。

　　左為 JFT-300M 數(shù)據(jù)集上的 Precision@1 曲線圖；右為 ImageNet 5-shot 的準(zhǔn)確率曲線圖。

　　對(duì)于這兩種情況，谷歌大腦發(fā)現(xiàn)，在給定訓(xùn)練計(jì)算量時(shí)，稀疏模型顯著優(yōu)于密集模型或者更快地實(shí)現(xiàn)相似性能。為了探索視覺(jué)模型的極限，他們?cè)?JFT-300M 擴(kuò)展數(shù)據(jù)集上訓(xùn)練了一個(gè)具有 150 億參數(shù)、24 個(gè) MoE 層（出自 48 個(gè)塊）的模型。這個(gè)迄今為止最大的視覺(jué)模型在 ImageNet 上實(shí)現(xiàn)了 90.35 的 Top-1 準(zhǔn)確率。

　　優(yōu)先路由

　　在實(shí)踐中，由于硬件限制，使用動(dòng)態(tài)大小的緩沖區(qū)（buffer）效率不高，因此模型通常為每個(gè)專家使用預(yù)定義的緩沖區(qū)容量。一旦專家變「滿」，超出此容量的分配 token 將被丟棄并不會(huì)被處理。因此，更高的容量會(huì)產(chǎn)生更高的準(zhǔn)確性，但它們的計(jì)算成本也更高。

　　谷歌大腦利用這種實(shí)現(xiàn)約束來(lái)使 V-MoE 在推理時(shí)更快。通過(guò)將總組合緩沖區(qū)容量降低到要處理的 token 數(shù)量以下，網(wǎng)絡(luò)被迫跳過(guò)處理專家層中的一些 token。該模型不是以某種任意方式選擇要跳過(guò)的 token（就像以前的工作那樣），而是學(xué)習(xí)根據(jù)重要性分?jǐn)?shù)對(duì)它們進(jìn)行排序。這樣可以保持高質(zhì)量的預(yù)測(cè)，同時(shí)節(jié)省大量計(jì)算。他們將這種方法稱為批量?jī)?yōu)先級(jí)路由（Batch Priority Routing, BPR），動(dòng)態(tài)示意圖如下所示：

　　在高容量下，Vanilla 和優(yōu)先路由都可以很好地處理所有 patch。但是，當(dāng)減小緩沖區(qū)大小以節(jié)省計(jì)算時(shí)，Vanilla 路由選擇處理任意 patch，通常導(dǎo)致預(yù)測(cè)不佳；BPR 智能地優(yōu)先選擇處理重要 patch，使得以更低的計(jì)算成本獲得更佳的預(yù)測(cè)。

　　事實(shí)證明，適當(dāng)?shù)貏h除 token 對(duì)于提供高質(zhì)量和更有效的推理預(yù)測(cè)至關(guān)重要。當(dāng)專家容量減少時(shí)，Vanilla 路由機(jī)制的性能會(huì)迅速下降。相反，BPR 對(duì)低容量更為穩(wěn)健。

　　總體而言，谷歌大腦觀察發(fā)現(xiàn)，V-MoE 在推理時(shí)非常靈活：例如，可以減少每個(gè) token 選擇的專家數(shù)量以節(jié)省時(shí)間和計(jì)算，而無(wú)需對(duì)模型權(quán)重進(jìn)行任何進(jìn)一步的訓(xùn)練。

　　探索 V-MoE

　　由于關(guān)于稀疏網(wǎng)絡(luò)的內(nèi)部工作原理還有很多待發(fā)現(xiàn)，谷歌大腦還探索了 V-MoE 的路由模式。一種假設(shè)是，路由器會(huì)根據(jù)某些語(yǔ)義背景（如「汽車」專家、「動(dòng)物」專家等）學(xué)會(huì)區(qū)分并分配 token 給專家。

　　為了測(cè)試這一點(diǎn)，他們?cè)谙旅嬲故玖藘蓚€(gè)不同 MoE 層的圖，一個(gè)非常早期（very early-on），另一個(gè)更靠近 head。x 軸對(duì)應(yīng) 32 個(gè)專家中的每一個(gè)，y 軸顯示圖像類別的 ID（從 1 到 1000）。圖中每個(gè)條目都顯示了為與特定圖像類對(duì)應(yīng)的 token 選擇專家的頻率，顏色越深表示頻率越高。

　　結(jié)果顯示，雖然在早期層幾乎沒(méi)有相關(guān)性，但在網(wǎng)絡(luò)后期，每個(gè)專家只接收和處理來(lái)自少數(shù)幾個(gè)類別的 token。因此，可以得出結(jié)論，patch 的一些語(yǔ)義聚類出現(xiàn)在網(wǎng)絡(luò)的更深層。

　　更高的路由決策與圖像類別相關(guān)。

　　谷歌大腦相信這只是計(jì)算機(jī)視覺(jué)大規(guī)模條件計(jì)算的開始。異構(gòu)專家架構(gòu)和條件可變長(zhǎng)度路由也是有潛力的研究方向。稀疏模型尤其有益于數(shù)據(jù)豐富的領(lǐng)域，例如大規(guī)模視頻建模。他們希望開源的代碼和模型能夠吸引更多研究人員關(guān)注該領(lǐng)域。