150億參數,谷歌開源瞭史上最大視覺模型V-MoE的全部代碼


還記得谷歌大腦團隊去年 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

幹貨推薦

  • 浙大博士導師深度整理:Tensorflow 和 Pytorch 的筆記(包含經典項目實戰)
  • 值得收藏,這份機器學習算法資料著實太香
  • 比 PyTorch 的官方文檔還香啊,吃透PyTorch中文版來瞭
  • 趕快收藏,PyTorch 常用代碼段PDF合輯版來瞭

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 的一些語義聚類出現在網絡的更深層。

_更高的路由決策與圖像類別相關。
_

谷歌大腦相信這隻是計算機視覺大規模條件計算的開始。異構專傢架構和條件可變長度路由也是有潛力的研究方向。稀疏模型尤其有益於數據豐富的領域,例如大規模視頻建模。他們希望開源的代碼和模型能夠吸引更多研究人員關註該領域。

技術交流

目前已開通瞭技術交流群,群友已超過1000人,添加時最好的備註方式為:來源+興趣方向,方便找到志同道合的朋友

  • 方式①、發送如下圖片至微信,長按識別,後臺回復:加群
  • 方式②、微信搜索公眾號:機器學習社區,後臺回復:加群
  • 方式③、可以直接加微信號:mlc2060。加的時候備註一下:研究方向 +學校/公司+CSDN,即可。然後就可以拉你進群瞭。