拋棄歸一化,深度學習模型準確率卻達到了前所未有的水平
我們知道,在傳遞給機器學習模型的數據中,我們需要對數據進行歸一化(normalization)處理。
在數據歸一化之后,數據被「拍扁」到統一的區間內,輸出范圍被縮小至 0 到 1 之間。人們通常認為經過如此的操作,最優解的尋找過程明顯會變得平緩,模型更容易正確的收斂到最佳水平。
然而這樣的「刻板印象」最近受到了挑戰,DeepMind 的研究人員提出了一種不需要歸一化的深度學習模型 NFNet,其在大型圖像分類任務上卻又實現了業內最佳水平(SOTA)。

該模型(紅色)與其他模型在 ImageNet 分類準確度和訓練時間上的對比。
該論文的第一作者,DeepMind 研究科學家 Andrew Brock 表示:「我們專注于開發可快速訓練的高性能體系架構,已經展示了一種簡單的技術(自適應梯度裁剪,AGC),讓我們可以訓練大批量和大規模數據增強后的訓練,同時達到 SOTA 水平。」
該研究一經提交,便吸引了人們的目光。

- 論文鏈接:https://arxiv.org/abs/2102.06171
- DeepMind 還放出了模型的實現:https://github.com/deepmind/deepmind-research/tree/master/nfnets
NFNet 是不做歸一化的 ResNet 網絡。具體而言,該研究貢獻有以下幾點:
- 提出了自適應梯度修剪(Adaptive Gradient Clipping,AGC)方法,基于梯度范數與參數范數的單位比例來剪切梯度,研究人員證明了 AGC 可以訓練更大批次和大規模數據增強的非歸一化網絡。
- 設計出了被稱為 Normalizer-Free ResNets 的新網絡,該方法在 ImageNet 驗證集上大范圍訓練等待時間上都獲得了最高水平。NFNet-F1 模型達到了與 EfficientNet-B7 相似的準確率,同時訓練速度提高了 8.7 倍,而 NFNet 模型的最大版本則樹立了全新的 SOTA 水平,無需額外數據即達到了 86.5%的 top-1 準確率。
- 如果在對 3 億張帶有標簽的大型私人數據集進行預訓練,隨后針對 ImageNet 進行微調,NFNet 可以比批歸一化的模型獲得更高的 Top-1 準確率:高達 89.2%。
研究方法
在沒有歸一化的情況下,許多研究者試圖通過恢復批歸一化的好處來訓練深度 ResNet 以提升其準確率。這些研究大多數通過引入小常數或可學習的標量來抑制初始化時殘差分支上的激活尺度。
DeepMind 的這項研究采用并建立在「Normalizer-Free ResNet(NF-ResNet)」上,這是一類可以在沒有歸一化層的情況下,被訓練成具有訓練和測試準確率的預激活 ResNet。
NF-ResNet 使用如下形式的殘差塊:

其中,h_i 代表第 i 個殘差塊的輸入,f_i 代表由第 i 個殘差分支計算的函數。
用于高效大批量訓練的自適應梯度裁剪
為了將 NF-ResNet 擴展到更大的批規模,研究者探索了一系列梯度裁剪策略。梯度裁剪通常被用于語言建模中以穩定訓練。近來一些研究表明:與梯度下降相比,梯度裁剪允許以更高的學習率進行訓練,從而加快收斂速度。這對于條件較差的 loss landscape 或大批量訓練尤為重要。因為在這些情況下,最佳學習率受到最大穩定學習率的限制。因此該研究假設梯度裁剪應該有助于將 NF-ResNet 有效地擴展到大批量設置。
借助一種稱為 AGC 的梯度裁剪方法,該研究探索設計了 Normalizer-Free 架構,該架構實現了 SOTA 的準確率和訓練速度。
當前圖像分類任務的 SOTA 大多是 EfficientNet 系列模型 (Tan & Le, 2019)取得的,該系列的模型經過優化以最大化測試準確率,同時最小化參數量和 FLOP 計數,但它們的低理論計算復雜度并沒有轉化為訓練速度的提高。

該研究通過手動搜索設計導向來探索模型設計的空間,這些導向對比設備上的實際訓練延遲,可帶來 ImageNet 上 holdout top-1 的帕累托前沿面的改進。它們對 holdout 準確率的影響如下表 2 所示:

實驗
表 3 展示了六個不同的 NFNets(F0-F5)與其他模型在模型大小、訓練延遲和 ImageNet 驗證準確率方面的對比情況。NFNets-F5 達到了 86.0%的 SOTA top-1 準確率,相比 EfficientNet-B8 有了一定提升;NFNet-F1 的測試準確率與 EfficientNet-B7 相媲美,同時訓練速度提升了 8.7 倍;NFNet-F6+SAM 達到了 86.5%的 top-1 準確率。

NFNets 和其他模型在 ImageNet 數據集上的準確率對比。延遲是指在 TPU 或 GPU(V100)上運行單個完整訓練步驟所需要的毫秒時間。
此外,研究者使用了一個 3 億標注圖像的數據集對 NFNet 的變體進行了預訓練,并針對 ImageNet 進行微調。最終,NFNet-F4 + 在 ImageNet 上獲得了 89.2% 的 top-1 準確率。這是迄今為止通過額外訓練數據達到的第二高的驗證準確率,僅次于目前最強大的半監督學習基線 (Pham et al., 2020) 和通過遷移學習達到的最高準確率。

表 5:使用額外數據進行大規模預訓練后,ImageNet 模型遷移性能對比。
Andrew Brock 表示,雖然我們對于神經網絡信號傳遞、訓練規律的理解還有很多需要探索的方向,但無歸一化的方法已經為人們提供了一個強有力的參考,并證明了發展這種深度理解能力可以有效地在生產環境中提升效率。