NFNet (Normalizer-Free ResNets)論文閱讀

Published in

謦伊的閱讀筆記

7 min readMar 31, 2021

NFNet 全名為 Normalizer-Free ResNets，是一種不使用 Batch Normalization、網路架構基於 ResNet 的模型，並且還提出了自適應梯度修剪 (Adaptive Gradient Clipping，AGC)，由 DeepMind 發表。

📝 Paper: https://arxiv.org/abs/2102.06171

🔖 Github: https://github.com/deepmind/deepmind-research/tree/master/nfnets

之前的文章有講過 Batch Normalization (BN) 能夠使模型更加穩定、加速收斂而且有正則化的效果，是當前神經網路的標準配備，但還是存在以下缺點:

在計算平均數和標準差時，需要將其值保存在記憶體中，除了會增加記憶體使用量，也增加了網路評估梯度的時間。
造成模型訓練和推理時的差異，代表 BN 引入了必須調整的隱藏超參數。
破壞了 minibatch 的訓練資料間的獨立性，因此在 minibatch 中選擇哪些樣本變得很重要。

而這三個缺點會帶來一些負面影響

很難準確地複製到不同的硬體上，尤其是在分散式訓練 (distributed training)，因為需要將不同機器上的資料互相傳送，容易導致訊息洩漏的問題。
對於 batch size 非常敏感，若 batch size 太小會導致結果較不穩定。

為了解決這些問題，作者提出了 Adaptive Gradient Clipping (AGC)，能夠使用更大的 batch size 和更強的 data augmentations 進行訓練。

Adaptive Gradient Clipping

Gradient Clipping 通常是基於限制 gradient norm 來實現。在梯度更新前，會採用下列公式進行裁剪梯度，其中 gradient vector G = ∂L/∂θ，L 為 loss, θ 為參數向量, λ 為 clipping threshold，是需要調整的超參數。

Gradient Clipping 能夠使用比以往更大的 batch size 訓練，但訓練穩定性會對 λ 特別敏感。在改變模型深度、batch size、learning rate 時，都要仔細調整 λ 值，因而引入了 Adaptive Gradient Clipping。

AGC 的思想源自於觀察梯度與權重的範數比 (norm ratio)，若訓練時 gradient descent 不使用 momentum，則會有以下關係式

作者認為若關係式左式 |ΔW_l|/|W_l| 很大會導致訓練不穩定，而採用了基於梯度與權重的範數比的 Gradient Clipping 策略，並且在 unit-wise 上的效果比 layer-wise 還要好。unit-wise 是指在第 l 層上 unit i 的梯度矩陣，layer-wise 則是指在第 l 層上每個 unit 的梯度矩陣。

AGC 的公式如下，其中 ε=10^(-3)，為了避免 zero-initialized parameters 總是被裁剪為 0，且在 channel and spatial 維度上都使用 unit-wise norms。

下圖為 Ablation experiment result，由 (a) 可以得知 NFNet 能夠使用 4096 batch size+RandAugment 訓練，而 batch-normalized networks 的結果很差。圖 (b) 為不同的 batch size 在每個 clipping threshold λ 上的表現，藉由經驗，當 batch size 越大應選擇較越小的 λ。

Normalizer-Free Architectures

作者設計了一系列的 NFNet，其架構在 SE-ResNeXt-D 的基礎上應用 Normalizer-Free 和 GELU activations，並修改寬度與深度。關於 ResNeXt 的詳細介紹，可參考我的文章: ResNeXt 論文閱讀。

下圖為 NFNet 的網路架構，作者發現第三個 stage 足夠深，能學習更深層次的特徵且具有較大的感受野，因此選擇在第三個 Stage Widths 增加 capacity。

α 是指在每個 residual block 之後，激活函數的變異數 (variance) 所增加的速率，默認值為 0.2。
βi 則是第 i 個 residual block 輸入的標準差，除了 transition blocks 以外。transition blocks 的架構在文章後面會提到，其 shortcut connection 會在縮減輸入 (即 Scaled Activation，downscaled: hi/βi) 之後進行，通過 transition blocks 後會將預期的變異數設置為 h_(i+1) = 1 + α²。

此外，為了防止隱藏層的激活函數出現 mean-shift 而引入 Scaled Weight Standardization，公式如下

NFNet 中的每個 Stage Depths 為 [1, 2, 6, 3] x N，其中 N 表示 NFNet-F0~F6 所對應的模型，NFNet-F0: N = 1, Stage Depths=[1, 2, 6, 3]、NFNet-F1: N = 2, Stage Depths=[2, 4, 12, 6]，以此類推。