NFNet解读
針對BN做的一個工作,這段時間這樣的工作還是不少的,當一個領域的技術發展得比較成熟時,我們往往就會考慮一些我們習以為常的東西的優化和改進,如之前的RepVGG、又如現在的NFNet。
簡介
最近有不少文章介紹了NFNet,但是沒怎么看到針對論文較為詳細的解讀,所以這邊就結合論文談談個人的見解。NFNet(Normalizer-Free ResNets)是DeepMind提出了一種不需要Batch Normalization的基于ResNet的網絡結構,其核心為一種AGC(adaptive gradient clipping technique,自適應梯度裁剪)技術。如下圖所示,最小的NFNet版本達到了EfficientNet-B7的準確率,并且訓練速度快了8.7倍,最大版本的模型實現了新的SOTA效果。
-
論文標題
High-Performance Large-Scale Image Recognition Without Normalization
-
論文地址
http://arxiv.org/abs/2102.06171
-
論文源碼
https://github.com/deepmind/deepmind-research/tree/master/nfnets
介紹
目前計算機視覺中很多網絡網絡都是基于ResNet的變種,使用Batch Normalization(下文簡稱BN)進行訓練。BN和殘差結構的組合已經被業界證明十分有效,可以很容易地訓練深層網絡。BN的存在可以平滑loss使得更大學習率更大batch size的訓練穩定進行,BN也有一定的正則化效果。不可否認,BN是非常有效的,但是它存在問題,論文中總結了三個非常典型的缺點如下。
- BN是帶來了額外的不小的計算開銷。 計算均值等需要在內存中保存臨時變量,導致了內存開銷的增大,并且,在某些網絡中增加了梯度評估的時間。
- BN造成了模型訓練和推理時的行為差異。 在Pytorch中實現上時model.train()和model.eval()的差異,也就是說BN帶來了需要調整的隱藏超參數。
- BN打破了小批量樣本之間的獨立性。 這就是說,其實選擇哪些樣本其實挺重要的。
這三個特性導致了一系列不好的后果。一方面,具體而言,研究者已經發現,使用BN的網絡通常難以在不同的硬件設備之間精確復制,BN往往就是這種細微錯誤的原因,特別是分布式訓練時。分布式訓練時,如果數據并行,在不同的機器上都有BN層,那么就需要將信號發送到BN層,然后在BN層之間傳遞均值等統計信息,不這樣的話這個批次就沒有均值和方差,這使得網絡可以欺騙某些損失函數,造成信息泄露等問題。另一方面,BN層對batch size是非常敏感的,bs過小BN網絡就會效果很差,這是因為bs很小的時候樣本很少,均值其實是噪聲的近似。
因此,盡管BN推動了深度學習的發展,但是從長遠來看,它其實阻礙了深度網絡的進步。其實如上圖所示,業內已經出現了Layer Norm、Group Norm等BN的替代品,它們獲得了更好的精度表現但也帶來了不小的計算量。幸運的是,這些年也出現了一些具有代表性的無BN的網絡,這些工作的核心思路就是通過抑制殘差分支上隱藏激活層的尺度從而訓練非常深的無BN的ResNet網絡。最簡單的實現方法就是在每個殘差分支的末尾引入一個初始值為0的可學習標量,不過這種技巧的精度表現并不好。另一些研究表明,ReLU這個激活函數會帶來均值偏移現象,這導致不同樣本的隱藏激活值隨著網絡的深度增加越來越相關。此前已經有工作提出了Normalize-Free ResNets,它在初始化的時候抑制殘差分支并且使用Scaled Weight Standardization來消除均值偏移現象,通過額外的正則化,這種網絡在ImageNet上獲得了和有BN網絡相媲美的效果,但是它在大的bs時訓練并不穩定并且其實距離目前的SOTA也就是EfficientNet還有一定距離。因此,這篇論文提出了解決它劣勢的新方法,稱為NFNet,論文的主要貢獻如下。
- 提出Adaptive Gradient Clipping (AGC)模塊,該方法基于逐單元梯度范數與參數范數的單位比例來裁剪梯度,實驗表明,AGC允許NFNet以大的bs和強數據增強條件進行訓練。
- 設計了一系列的Normalize-Free ResNets,稱為NFNets,最簡單版本的NFNet達到了EfficientNet-B7的精度,但訓練速度是8.7倍。最優版本的NFNet在不適用額外數據的情況下實現了新的SOTA。
- 在3億張帶有標簽的大型私有數據集進行預訓練后,對ImageNet進行微調時,NFNet與批歸一化網絡相比,其驗證準確率要高得多。最佳模型經過微調后可達到89.2%的top-1 accuracy。
這篇文章后面兩節主要是敘述BN的效果,以及前人如何在去除了BN之后保留這些優勢所做的工作,這里感興趣的可以查看文章的第二節和第三節,我這里就直接來將這篇論文的方法論了。
AGC(自適應梯度裁剪模塊)
梯度裁剪技術常用于語言模型來穩定訓練,最近的研究表明,與梯度下降相比,它允許以更大的學習率進行訓練從而加速收斂。這對于條件較差的損失或大批量訓練尤為重要,因為在這些設置中,最佳學習率往往會受到最大學習率的限制。因此作者假定梯度裁剪有利于NFNet的大批尺寸訓練。梯度裁剪往往是對梯度的范數進行約束來實現的,對梯度向量G=?L/?θG=\partial L / \partial \thetaG=?L/?θ而言,LLL表示損失值,θ\thetaθ則表示模型所有參數向量,標準的裁剪算法會在更新θ\thetaθ之前以如下的公式裁剪梯度。
G→{λG∥G∥if?∥G∥>λGotherwise?G \rightarrow\left\{\begin{array}{ll} \lambda \frac{G}{\|G\|} & \text { if }\|G\|>\lambda \\ G & \text { otherwise } \end{array}\right. G→{λ∥G∥G?G??if?∥G∥>λ?otherwise??
上式的λ\lambdaλ是必須調整的超參數,根據經驗,作者發現雖然這個裁剪算法能夠以比以前更高的批尺寸進行訓練,但訓練穩定性對裁剪閾值λ\lambdaλ的選擇極為敏感,在改變模型深度、批尺寸或學習率時都需要精調閾值。
為了解決這個問題,作者引入了自適應梯度裁剪算法(AGC),下面詳細敘述這個算法。記W?∈RN×MW^{\ell} \in \mathbb{R}^{N \times M}W?∈RN×M為第?\ell?層的權重矩陣,G?∈RN×MG^{\ell} \in \mathbb{R}^{N \times M}G?∈RN×M為對應于W?W^{\ell}W?的梯度矩陣,∥?∥F\|\cdot\|_{F}∥?∥F?表示FFF范數,即有∥W?∥F=∑iN∑jM(Wi,j?)2\left\|W^{\ell}\right\|_{F}=\sqrt{\sum_{i}^{N} \sum_{j}^{M}\left(W_{i, j}^{\ell}\right)^{2}}∥∥?W?∥∥?F?=∑iN?∑jM?(Wi,j??)2?。AGC算法的動機源于觀察到梯度與權重的范數比∥G?∥F∥W?∥F\frac{\left\|G^{\ell}\right\|_{F}}{\left\|W^{\ell}\right\|_{F}}∥W?∥F?∥G?∥F??,這其實是一個單次梯度下降對原始權重影響的簡單度量。舉個例子,如果使用無動量的梯度下降算法,有∥ΔW?∥∥W?∥=h∥G?ˉ∥F∥W?∥F\frac{\left\|\Delta W^{\ell}\right\|}{\left\|W^{\ell}\right\|}=h \frac{\left\|G^{\bar{\ell}}\right\|_{F}}{\left\|W^{\ell}\right\|_{F}}∥W?∥∥ΔW?∥?=h∥W?∥F?∥G?ˉ∥F??,那么第?\ell?層的參數更新公式為ΔW?=?hG?\Delta W^{\ell}=-h G^{\ell}ΔW?=?hG?,其中hhh表示學習率。直觀上,我們認為如果∥ΔW?∥/∥W?∥\left\|\Delta W^{\ell}\right\| /\left\|W^{\ell}\right\|∥∥?ΔW?∥∥?/∥∥?W?∥∥?很大那么訓練就會變得不穩定,這就啟發了一種基于∥G?∥F∥W?∥F\frac{\left\|G^{\ell}\right\|_{F}}{\left\|W^{\ell}\right\|_{F}}∥W?∥F?∥G?∥F??的梯度裁剪策略,然而實際上,逐單元的梯度范數和參數范數比會比逐層的效果好,因此定義第?\ell?層上第iii個單元的梯度矩陣Gi?G_{i}^{\ell}Gi??(表示G?G_{\ell}G??的第iii行)的裁剪公式如下,其中λ\lambdaλ是一個標量超參數,定義∥Wi∥F?=max?(∥Wi∥F,?)\left\|W_{i}\right\|_{F}^{\star}=\max \left(\left\|W_{i}\right\|_{F}, \epsilon\right)∥Wi?∥F??=max(∥Wi?∥F?,?)(其中默認?=10?3\epsilon=10^{-3}?=10?3),這避免0初始化參數總是裁剪為0。對于卷積濾波器中的參數,我們在扇入范圍(包括通道和空間維度)上評估逐單元范數。
Gi?→{λ∥Wi?∥F?∥Gi?∥FGi?if?∥Gi?∥F∥Wi?∥F?>λGi?otherwise.?G_{i}^{\ell} \rightarrow\left\{\begin{array}{ll} \lambda \frac{\left\|W_{i}^{\ell}\right\|_{F}^{\star}}{\left\|G_{i}^{\ell}\right\|_{F}} G_{i}^{\ell} & \text { if } \frac{\left\|G_{i}^{\ell}\right\|_{F}}{\left\|W_{i}^{\ell}\right\|_{F}^{\star}}>\lambda \\ G_{i}^{\ell} & \text { otherwise. } \end{array}\right. Gi??→????λ∥Gi??∥F?∥Wi??∥F???Gi??Gi????if?∥Wi??∥F??∥Gi??∥F??>λ?otherwise.??
使用上述的AGC模塊,NFNet能夠以高達4096的批尺寸訓練同時使用RandAugment這樣的數據增強策略,不適用AGC模塊,NFNet是無法這樣訓練的。注意,最優裁剪參數λ\lambdaλ可能取決于優化器的選擇,學習率和批大尺寸。經驗上,batch size越大,λ\lambdaλ應該越小。
上圖是論文針對AGC做的兩個消融實驗,左圖a表示使用BN的REsNet以及使用和不使用AGC的NFNet之間的對比,實驗表明AGC使得NFNet有著媲美BN網絡的效果,而且批尺寸越小,AGC收益越低。右圖b則表示不同批尺寸不同λ\lambdaλ選擇的效果,結果表明,當批尺寸較大的時候,應該選擇較小的λ\lambdaλ以穩定訓練。
后續作者也對AGC的作用層進行了消融實驗,得到一些結論,比如對最終的線性層裁剪是不需要的等。
NFNet
上一節提到了可以讓網絡以較大批尺寸和較強數據增強方法進行訓練的梯度裁剪手段,同時為了配置這個AGC模塊,論文也對模型結構進行了探索。EfficientNet是當前的分類SOTA網絡,它基于NAS技術搜索而得,擁有很低的理論計算復雜度,但是實際硬件上的表現并不是很好,所以作者這邊選擇了手工探索模型空間以獲得較好的表現,在SE-ResNeXt-D模型的基礎上,對其先是應用了前人的Normalizer-Free配置,修改了寬度和深度模式,以及第二個空間卷積(下圖表示這種配置上的修改,更具體的可以查看論文附加材料,給出的圖如下)。接著,應用AGC到除了最后的線性層上的每一層,得到最終的NfNet配置。
得到的NfNets如下表所示,構成一個網絡系列。
實驗
下表是NFNet使用各種數據增強在ImageNet上和其他方法的對比,當之無愧的SOTA。
也進行了預訓練的實驗,證明了Normalize-Free ResNets在遷移學習上效果也是能夠很強的。
總結
關于BN這個結構我之前其實也有一些思考,雖然不否認它為深度學習做出的巨大貢獻(我自己進行實驗過程中深感BN結構的有效性),但是它也確實存在一些問題,DeepMind的這篇文章在保留BN優勢克服BN劣勢的基礎上,實現了一個非常成功的Normalize-Free ResNets,是很值得關注的工作。
總結
- 上一篇: 新的一年2021
- 下一篇: MMDetection-数据准备