20億參數,大型視覺Transformer來了,刷新ImageNet Top1

編輯:杜偉、陳萍

2020 年 10 月,谷歌大腦團隊提出將標準 Transformer 應用於圖像,提出了視覺 Transformer(ViT)模型,並在多個圖像識別基準上實現了接近甚至優於當時 SOTA 方法的性能。近日,原 ViT 團隊的幾位成員又嘗試將 ViT 模型進行擴展,使用到了包含 30 億圖像的 JFT-3B 資料集,並提出了參數量高達 20 億參數的 ViT 變體模型 ViT G/14,在 ImageNet 圖像資料集上實現了新的 SOTA Top-1 準確率。

基於注意力機制的 Transformer 架構已經席捲了 CV 領域,併成為研究和實踐中日益流行的選擇。此前,Transformer 被廣泛用於 NLP 領域。有研究者仔細研究了自然語言處理中 Transformer 最優擴展,主要結論是大型模型不僅性能更好,而且更有效地使用了大量計算預算。

然而,目前尚不清楚這些發現在多大程度上能夠遷移到視覺領域。例如,視覺中最成功的預訓練方案是有監督的,而 NLP 領域是無監督預訓練。

在今天介紹的這篇論文中,原 ViT 團隊成員、谷歌大腦的幾位研究者集中研究了預訓練 ViT 模型用於圖像分類任務的遷移性能的擴展規則(scaling law)。特別是,研究者試驗了從 500 萬到 20 億個參數不等的模型、從 3000 萬到 30 億個訓練圖像不等的資料集以及從低於 1 個 TPUv3 核每天(core-day)到超過 10000 個核每天的計算預算。其主要貢獻是描述 ViT 模型的性能計算邊界。

論文連結:https://arxiv.org/pdf/2106.04560.pdf

在這個過程中,研究者創建了一個改進的大規模訓練方案,探索了訓練超參數以及發現微妙的選擇,大幅改善小樣本遷移性能。具體來說,研究者發現非常強的 L2 正則化,僅應用於最終的線性預測層,導致學習到的視覺表徵具有很強的小樣本學習能力。

例如,在 ImageNet 資料集(有 1000 個類)上,每個類只有一個示例,該研究的最佳模型達到 69.52% 的準確率;如果每個類有 10 個示例,準確率達到了 84.86%。此外,該研究大大減少了 [11] 中提出的原始 ViT 模型的記憶體佔用,通過特定於硬體的體系架構更改和不同的最佳化器來實現這一點。結果表明,該研究訓練了一個具有 20 億個參數的模型,在 ImageNet 資料集上達到了新的 SOTA 性能 90.45% 的準確率。可以看到,在 SOTA 基準排行榜上,ViT-G/14 模型的 Top-1 準確率已經超越了谷歌之前提出的 Meta Pseduo Labels 模型。

如何做到的?

如何做到的?

研究者展示了對 ViT 模型及訓練的改進,這些改進大多數易於實現,並顯著提升了記憶體使用率和模型質量。如此一來,研究者可以單獨使用資料並行化訓練 Vit-G/14 模型,並在單個 TPUv3 上實現整個模型擬合。

擴展資料

該研究使用專有的 JFT-3B 資料集,它是 JFT-300M 資料集的更大規模版本,在之前的許多大型計算機視覺模型工作中使用過 [31, 18, 11]。該資料集由近 30 億張圖像組成,標註為包含 30k 類別的層級類別,並且由於採用半自動標註,所以標註是有噪音的。所有的測試資料也均從 JFT-3B 中移除。

下圖 5 展示了資料集從 JFT-300M 到 JFT-3B 過程中對模型性能的影響。可以觀察到,更大規模的 JFT-3B 資料集可以得到更好的模型,所以 JFT-300M 資料集的過擬合併不是實現性能提升的唯一原因。

「head」 的解耦權重衰減

「head」 的解耦權重衰減

權重衰減對低資料情況下的模型自適應具有重大影響。研究者在中等規模程度上研究了這一現象,並發現可以從模型中最終線性層(「head」)和剩餘權重(「body」)的權重衰減強度解耦中獲益。

下圖 4 展示了這一效果。研究者在 JFT-300M 上訓練了一個 ViT-B/32 模型,每個單元格對應不同 head/body 權重衰減值的性能。他們觀察到的有趣的一點是:儘管提升了遷移性能,但 head 中高權重衰減卻降低了預訓練(上游)任務的性能。

通過移除 [class] token 節省記憶體

對於 ViT 模型,當前的 TPU 硬體將 toek 維數填充為 128 的倍數,這可能導致高度 50% 的記憶體開銷。為了解決這一問題,研究者嘗試探索「使用額外[class] token」的替代方法。具體地,他們對全局平均池化( GAP)和多頭注意力池化(MAP)進行評估以聚合來自所有 patch token 的表示,並將 MAP 中 head 的數量與模型其他部分中注意力 head 的數量設為相同。為了進一步簡化 head 設計,研究者原始 ViT 論文中出現的、最終預測層之前的最終非線性對映。

為了選擇最佳 head,研究者對[class] token 和 GAP/MAP head 進行了並排比較,結果如上圖 4(右)所示。他們發現,所有 head 的表現類似,同時 GAP 和 MAP 由於進行了填充(padding)考慮,因而具備更高的記憶體效率。此外,非線性對映還可以安全地進行移除。因此,研究者選擇了 MAP head,這是因為其表現力最強,並且能夠生成最統一的架構。

模型尺寸

在本文中,研究者選擇了 ViT 模型,模型參數從 500 萬到 2 億,訓練資料量從 3000 萬到 30 億。下表 2 為具體的模型架構細節:

實驗結果

實驗結果

在接下來的實驗中,研究者在多達 30 億個弱標記圖像上訓練了幾個 ViT 模型。研究者改變架構大小、訓練圖像的數量和訓練持續時間。所有模型都在 TPUv3 上訓練,因此總計算量是以 TPUv3 核每天(core-days)數來衡量的。

為了評估由模型學習的表徵的質量,研究者進行了一下測量(i)通過在固定權值上訓練線性分類器來進行小樣本遷移;(ii)通過對所有資料的整個模型進行微調,將其傳遞給多個基準測試任務。

將計算、模型和資料一起擴展

圖 1 顯示了在 ImageNet 上的 10-shot 線性評估和微調評估,並進行了高層級的觀測。首先,將計算、模型和資料一起擴展可以提高表徵質量。第二,模型尺寸會影響表徵質量。第三,大型模型受益於額外的資料,甚至超過 1B 圖像。

圖 2 顯示了在預訓練期間「可見」的圖像總數(批大小乘以 step 數)的表徵質量。除了在公共驗證集上進行 ImageNet 微調和線性 10-shot 結果外,研究者還報告了在 ImageNet-v2 測試集 [27] 上的 ImageNet 微調模型的結果,作為魯棒泛化的指標。下圖中展示了對 30 億張圖像進行預訓練的三個 ViT 模型。

實驗可得,有足夠的資料,以較少的步驟訓練一個更大的模型是可取的。

ViT-G/14 結果

該研究訓練了一個大型的視覺 Transformer,ViT-G/14,它包含近 20 億個參數。實驗評估了 ViT-G/14 模型在一系列下游任務中的應用,並將其與 SOTA 結果進行了比較。研究者在 ImaegNet 上進行微調,並報告 ImageNet[28]、ImageNet-v2[27]、ReaL[3]和 ObjectNet[1]的準確率。此外,該研究還報告了在 VTAB-1k 基準上的遷移學習效果,該基準包括 19 個不同的任務[43]。

圖 3 顯示了在 ImageNet 上進行小樣本學習(few-shot learning)結果。由結果可得,ViT-G/14 比以前最好的 ViT-H/14 模型 [11] 表現優異(超過 5%),達到 84.86% 的準確率,每類 10 個例子。

下表 1 展示了其他基準上的結果。ViT-G/14 在 ImageNet 資料集上實現了 90.45% 的 Top-1 準確率,成為新的 SOTA。此外,在 ImageNet-v2 上,ViT-G/14 比基於 EfficientNet-L2 的 Noisy Student 模型提升了 3%。在 ReaL 資料集上,ViT-G/14 略微優於 ViT-H 和 BiT-L,再次表明 ImageNet 分類任務性能可能達到了飽和點。在 ObjectNet 資料集上,ViT-G/14 大幅度優於 BiT-L,較 Noisy Student 模型提升 2%,比 CLIP 落後約 2%。

相關文章