NFNet：NF-ResNet的延伸，不用BN的4096超大batch size訓練 | 21年論文
論文認為Batch Normalization並不是網絡的必要構造，反而會帶來不少問題，於是開始研究Normalizer-Free網絡，希望既有相當的性能也能支持大規模訓練。論文提出ACG梯度裁剪方法來輔助訓練，能有效防止梯度爆炸，另外還基於NF-ResNet的思想將SE-ResNet改造成NFNet系列，可以使用4096的超大batch size進行訓練，性能超越了Efficient系列
來源：曉飛的算法工程筆記 公眾號
論文: High-Performance Large-Scale Image Recognition Without Normalization
Introduction
目前，計算機視覺的大部分模型都得益於深度殘差網絡和batch normalization，這兩大創新能夠幫助訓練更深的網絡，在訓練集和測試集上達到很高的準確率。特別是batch normalization，不僅能夠平滑損失曲線，使用更大的學習率和batch size進行訓練，還有正則化的作用。然而，batch normalization並不是完美，batch normalization在實踐中有三個缺點：
- 計算消耗大，內存消耗多。
- 在訓練和推理上的用法不一致，並且帶來額外的超參數。
- 打破了訓練集的minibatch的獨立性。
其中，第三個問題最為嚴重，這會引發一系列的負面問題。首先，batch normalization使得模型難以在不同的設備上復現精度，而且分佈式訓練經常出問題。其次，batch normalization不能用於要求每輪訓練樣本獨立的任務中，如GAN和NLP任務。最後，batch normalization對batch size十分敏感，在batch size較低時表現較差，限制了有限設備上的模型大小。
因此，儘管batch normalization有很強大的作用，部分研究者仍在尋找一種簡單的替代方案，不僅需要精度相當，還要能用在廣泛的任務中。目前，大多數的替代方案都着力於抑制殘差分支的權值大小，比如在殘差分支的末尾引入一個初始為零的可學習的標量。但這些方法不是精度不夠，就是無法用於大規模訓練，精度始終不如EfficientNets。
至此，論文主要基於之前替代batch normalization的工作，嘗試解決其中的核心問題，論文的主要貢獻如下：
- 提出Adaptive Gradient Clipping(AGC)，以維度為單位，基於權值範數和梯度範數的比例進行梯度裁剪。將AGC用於訓練Normalizer-Free網絡，使用更大batch size和更強數據增強進行訓練。
- 設計Normalizer-Free ResNets系列，命名為NFNets，在ImageNet上達到SOTA，其中NFNet-F1與EfficientNet-B7精度相當，訓練速度快8.7倍，最大的NFNet可達到86.5%top-1準確率。
- 實驗證明，在3億標籤的私有數據集上預訓練後，再在ImageNet上進行finetune，準確率能比batch normalization網絡要高，最好的模型達到89.2%top-1準確率。
Understanding Batch Normalization
論文探討了batch normalization的幾個優點，這裡簡單說一下：
- downscale the residual branch：batch normalization限制了殘差分支的權值大小，使得信號偏向skip path直接傳輸，有助於訓練超深的網絡。
- eliminate mean-shift：激活函數是非對稱且均值非零的，使得訓練初期激活後的特徵值會變大且均為正數，batch normalization恰好可以消除這一問題。
- regularizing effect：由於batch normalization訓練時用的是minibatch統計信息，相當於為當前batch引入了噪聲，起到正則化的作用，可以防止過擬合，提高準確率。
- allows efficient large-batch training：batch normalization能夠平滑loss曲線，可以使用更大的學習率和bach size進行訓練。
Towards Removing Batch Normalization
這篇論文的研究基於作者之前的Normalizer-Free ResNets(NF-ResNets)進行拓展，NF-ResNets在去掉normalization層後依然可以有相當不錯的訓練和測試準確率。NF-ResNets的核心是採用h_{i+1}=h_i+\alpha f_i(h_i/\beta_i)形式的residual block，h_i為第i個殘差塊的輸入，f_i為第i個residual block的殘差分支。f_i要進行特殊初始化，使其有保持方差不變的功能，即Var(f_i(z))=Var(z)。\alpha=0.2用於控制方差變化幅度，\beta_i=\sqrt{Var(h_i)}為h_i的標準差。經過NF-ResNet的residual block處理後，輸出的方差變為Var(h_{i+1})=Var(h_i)+\alpha^2。
此外，NF-ResNet的另一個核心是Scaled Weight Standardization，用於解決激活層帶來的mean-shift現象，對卷積層進行如下權值重新初始化：
其中，\mu_i=(1/B)\sum_jW_{ij}和\sigma^2_i=(1/N)\sum_j(W_{ij}-\mu_i)^2為對應卷積核的某行(fan-in)的均值和方差。另外，非線性激活函數的輸出需要乘以一個特定的標量\gamma，兩者配合確保方差不變。
之前發佈的文章也有NF-ResNet的詳細解讀，有興趣的可以去看看。
Adaptive Gradient Clipping for Efficient Large-Batch Training
梯度裁剪能夠幫助訓練使用更大的學習率，還能夠加速收斂，特別是在損失曲線不理想或使用大batch size訓練的場景下。因此，論文認為梯度裁剪能幫助NF-ResNet適應大batch size訓練場景。對於梯度向量G=\partial L/\partial\theta，標準的梯度裁剪為：
裁剪閾值\lambda是需要調試的超參數。根據經驗，雖然梯度裁剪可以幫助訓練使用更大的batch size，但模型的效果對閾值\lambda的設定十分敏感，需要根據不同的模型深度、batch size和學習率進行細緻的調試。於是，論文提出了更方便的Adaptive Gradient Clipping(AGC)。
定義W^l\in\mathbb{R}^{N\times M}和G^l\in\mathbb{R}^{N\times M}為l層的權值矩陣和梯度矩陣，\|\cdot\|_F為F-範數，ACG算法通過梯度範數與權值範數之間比值\frac{\|G^l\|_F}{\|W^l\|_F}來進行動態的梯度裁剪。在實踐時，論文發現按卷積核逐行(unit-wise)進行梯度裁剪的效果比整個卷積核進行梯度裁剪要好，最終ACG算法為：
裁剪閾值\lambda為超參數，設定\|W_i\|^{*}_F=max(\|W_i\|_F, \epsilon=10^{-3})，避免零初始化時，參數總是將梯度裁為零。藉助AGC算法，NF-ResNets可以使用更大的batch size(4096)進行訓練，也可以使用更複雜的數據增強。最優的\lambda需考慮優化器、學習率和batch size，通過實踐發現，越大的batch size應該使用越小的\lambda，比如batch size=4096使用\lambda=0.01。
ACG算法跟優化器歸一化有點類似，比如LARS。LARS將權值更新值的範數固定為權值範數的比值\Delta w=\gamma * \eta * \frac{\|w^l\|}{\|\nabla L(w^l)\|} * \nabla L(w^l_t)，從而忽略梯度的量級，只保留梯度方向，能夠緩解梯度爆炸和梯度消失的現象。ACG算法可以認為是優化器歸一化的鬆弛版本，基於權值範數約束最大梯度，但不會約束梯度的下限或忽略梯度量級。論文也嘗試了ACG和LARS一起使用，發現性能反而下降了。
Normalizer-Free Architectures with Improved Accuracy and Training Speed
論文以帶GELU激活的SE-ResNeXt-D模型作為Normalizer-Free網絡的基礎，除訓練加入ACG外，主要進行了以下改進：
- 將$3\times 3$卷積變為分組卷積，每組的維度固定為128，組數由卷積的輸入維度決定。更小的分組維度可以降低理論的計算量，但計算密度的降低導致不能很好地利用設備稠密計算的優勢，實際不會帶來更多加速。
- ResNet的深度擴展(從resnNet50擴展至ResNet200)主要集中在stage2和stage3，而stage1和stage4保持3個block的形式。這樣的做法不是最優的，因為不管低層特徵或高層特徵，都需要足夠的空間去學習。因此，論文先制定最小的F0網絡的各stage的block數為[1,2,6,3]，後續更大網絡都在此基礎上以倍數擴展。
- ResNet的各stage維度為[256,512,1024,2048]，經過測試之後，改為[256,512,1536,1536]，stage3採用更大的容量，因為其足夠深，需要更大的容量去收集特徵，而stage4不增加深度主要是為了保持訓練速度。
-
將NF-ResNet的bottleneck residual block應用到SE-ResNeXt中並進行修改，在原有的基礎上添加了一個$3\times 3$卷積，在計算量上僅有少量的增加。
-
構建一個縮放策略來生產不同計算資源的模型，論文發現寬度擴展對網絡增益不大，於是僅考慮深度和輸入分辨率的縮放。按前面說的，以倍數形式對基礎網絡進行深度擴展，同時縮放分辨率，使其訓練和測試速度能達到上一個量級的一半。
-
當網絡體積增大時，加強正則化強度。通過實驗發現，調整weight decay和stochastic depth rate（訓練過程隨機使某些block的殘差分支失效）都沒有很大的收益，於是通過加大dropout的drop rate來達到正則化的目的。由於網絡缺少BN的顯示正則化，所以這一步是十分重要的，防止過擬合的出現。
根據上述的修改，得出的NFNet系列的各參數如表1所示。這裡網絡的最後有全局池化層，所以訓練和測試的分辨率可以不一樣。
Experiment
對比AGC在不同batch size下的效果，以及\lambda與batch size的關係。
在ImageNet對比不同大小的網絡的性能。
基於ImageNet的10 epoch預訓練權重，進行NF-ResNet改造並Fine-tuning，性能如表4所示。
Conclusion
論文認為Batch Normalization並不是網絡的必要構造，反而會帶來不少問題，於是開始研究Normalizer-Free網絡，希望既有相當的性能也能支持大規模訓練。論文提出ACG梯度裁剪方法來輔助訓練，能有效防止梯度爆炸，另外還基於NF-ResNet的思想將SE-ResNet改造成NFNet系列，可以使用4096的超大batch size進行訓練，性能超越了Efficient系列。
如果本文對你有幫助，麻煩點個贊或在看唄～
更多內容請關注 微信公眾號【曉飛的算法工程筆記】
