選自 Google AI
機器之心編譯
編輯：杜偉、陳萍
還記得谷歌大腦團隊去年 6 月份釋出的 43 頁論文《Scaling Vision with Sparse Mixture of Experts》嗎？他們推出了史上最大規模的視覺模型 V-MoE，實現了接近 SOTA 的 Top-1 準確率。如今，谷歌大腦開源了訓練和微調模型的全部程式碼。
在過去幾十年裡，深度學習的進步是由幾個關鍵因素推動的：少量簡單而靈活的機制、大型資料集、更專業的硬體配置，這些技術的進步使得神經網路在影象分類、機器翻譯、蛋白質預測等任務中取得令人印象深刻的結果。
然而，大模型以及資料集的使用是以大量計算需求為代價的。最近的研究表明，增強模型的泛化能力以及穩健性離不開大模型的支援，因此，在訓練大模型的同時協調好與訓練資源的限制是非常重要的。一種可行的方法是利用條件計算，該方法不是為單個輸入啟用整個網路，而是根據不同的輸入啟用模型的不同部分。這一正規化已經在谷歌提出的 pathway（一種全新的 AI 解決思路，它可以克服現有系統的許多缺點，同時又能強化其優勢）願景和最近的大型語言模型研究中得到了重視，但在計算機視覺中還沒有得到很好的探索。
稀疏門控混合專家網路 (MoE) 在自然語言處理中展示了出色的可擴充套件性。然而，在計算機視覺中，幾乎所有的高效能網路都是密集的，也就是說，每個輸入都會轉化為引數進行處理。
去年 6 月，來自谷歌大腦的研究者提出了 V-MoE（Vision MoE ），這是一種基於專家稀疏混合的新視覺架構。當應用於影象識別時，V-MoE 在推理時只需要一半的計算量，就能達到先進網路效能。此外，該研究還提出了對路由演算法的擴充套件，該演算法可以在整個 batch 中對每個輸入的子集進行優先順序排序，從而實現自適應影象計算。這允許 V-MoE 在測試時能夠權衡效能和平滑計算。最後，該研究展示了 V-MoE 擴充套件視覺模型的潛力，並訓練了一個在 ImageNet 上達到 90.35% 的 150 億引數模型。
論文地址：https://arxiv.org/pdf/2106.05974.pdf
程式碼地址：https://github.com/google-research/vmoe
V-MoE
谷歌大腦在 ViT 的不同變體上構建 V-MoE：ViT-S(mall)、ViT-B(ase)、ViT-L(arge) 和 ViTH(uge)，其超引數如下：
ViT 已被證明在遷移學習設定中具有良好的擴充套件性，在較少的預訓練計算下，比 CNN 獲得更高的準確率。ViT 將影象處理為一系列 patch，輸入影象首先被分成大小相等的 patch，這些 patch 被線性投影到 Transformer 的隱藏層，在位置嵌入後，patch 嵌入（token）由 Transformer 進行處理，該 Transformer 主要由交替的自注意力和 MLP 層組成。MLP 有兩個層和一個 GeLU 非線性。對於 Vision MoE，該研究用 MoE 層替換其中的一個子集，其中每個專家都是一個 MLP，如下圖所示：
為了大規模擴充套件視覺模型，該研究將 ViT 架構中的一些密集前饋層 (FFN) 替換為獨立 FFN 的稀疏混合（稱之為專家）。可學習的路由層為每個獨立的 token 選擇對應的專家。也就是說，來自同一影象的不同 token 可能會被路由到不同的專家。在總共 E 位專家（E 通常為 32）中，每個 token 最多隻能路由到 K（通常為 1 或 2）位專家。這允許擴充套件模型的大小，同時保持每個 token 計算的恆定。下圖更詳細地顯示了 V-MoE 編碼器塊的結構。
V-MoE Transformer 編碼器塊
實驗結果
谷歌大腦首先在大型影象資料集 JFT-300M 上對模型進行一次預訓練。
下圖左展示了模型在所有大小（從 small s/32 到 huge H/14）時的預訓練結果。然後，使用一個新的 head（一個模型中的最後一層）將模型遷移至新的下游任務（如 ImageNet）。他們探索了兩種遷移設定：在所有可用的新任務示例上微調整個模型或者凍結預訓練網路並使用少量示例僅對新 head 調整（即所謂的小樣本遷移）。
下圖右總結了模型遷移至 ImageNet 的效果，其中每個影象類別僅在 5 張影象上訓練（叫做 5-shot transfer）。
左為 JFT-300M 資料集上的 Precision@1 曲線圖；右為 ImageNet 5-shot 的準確率曲線圖。
對於這兩種情況，谷歌大腦發現，在給定訓練計算量時，稀疏模型顯著優於密集模型或者更快地實現相似效能。為了探索視覺模型的極限，他們在 JFT-300M 擴充套件資料集上訓練了一個具有 150 億引數、24 個 MoE 層（出自 48 個塊）的模型。這個迄今為止最大的視覺模型在 ImageNet 上實現了 90.35 的 Top-1 準確率。
優先路由
在實踐中，由於硬體限制，使用動態大小的緩衝區（buffer）效率不高，因此模型通常為每個專家使用預定義的緩衝區容量。一旦專家變「滿」，超出此容量的分配 token 將被丟棄並不會被處理。因此，更高的容量會產生更高的準確性，但它們的計算成本也更高。
谷歌大腦利用這種實現約束來使 V-MoE 在推理時更快。透過將總組合緩衝區容量降低到要處理的 token 數量以下，網路被迫跳過處理專家層中的一些 token。該模型不是以某種任意方式選擇要跳過的 token（就像以前的工作那樣），而是學習根據重要性分數對它們進行排序。這樣可以保持高質量的預測，同時節省大量計算。他們將這種方法稱為批次優先順序路由（Batch Priority Routing, BPR） ，動態示意圖如下所示：
在高容量下，Vanilla 和優先路由都可以很好地處理所有 patch。但是，當減小緩衝區大小以節省計算時，Vanilla 路由選擇處理任意 patch，通常導致預測不佳；BPR 智慧地優先選擇處理重要 patch，使得以更低的計算成本獲得更佳的預測。
事實證明，適當地刪除 token 對於提供高質量和更有效的推理預測至關重要。當專家容量減少時，Vanilla 路由機制的效能會迅速下降。相反，BPR 對低容量更為穩健。
總體而言，谷歌大腦觀察發現，V-MoE 在推理時非常靈活：例如，可以減少每個 token 選擇的專家數量以節省時間和計算，而無需對模型權重進行任何進一步的訓練。
探索 V-MoE
由於關於稀疏網路的內部工作原理還有很多待發現，谷歌大腦還探索了 V-MoE 的路由模式。一種假設是，路由器會根據某些語義背景（如「汽車」專家、「動物」專家等）學會區分並分配 token 給專家。
為了測試這一點，他們在下面展示了兩個不同 MoE 層的圖，一個非常早期（very early-on），另一個更靠近 head。x 軸對應 32 個專家中的每一個，y 軸顯示影象類別的 ID（從 1 到 1000）。圖中每個條目都顯示了為與特定影象類對應的 token 選擇專家的頻率，顏色越深表示頻率越高。
結果顯示，雖然在早期層幾乎沒有相關性，但在網路後期，每個專家只接收和處理來自少數幾個類別的 token。因此，可以得出結論，patch 的一些語義聚類出現在網路的更深層。
更高的路由決策與影象類別相關。
谷歌大腦相信這只是計算機視覺大規模條件計算的開始。異構專家架構和條件可變長度路由也是有潛力的研究方向。稀疏模型尤其有益於資料豐富的領域，例如大規模影片建模。他們希望開源的程式碼和模型能夠吸引更多研究人員關注該領域。
原文連結：
https://ai.googleblog.com/2022/01/scaling-vision-with-sparse-mixture-of.html?continueFlag=b96fa8ed72dfc82b777e51b7e954c7dc
