High-Performance Large-Scale Image Recognition Without Normalization

Batch NormalizationなどのNormalizationを一切含まない画像分類モデルNFNetを提案。ImageNetのみの訓練で86.0%のtop-1精度を達成。NFNetの小さめの構成であるNFNet-F1 は、EfficientNet-B7と同程度の精度でありながら、8.7倍も高速に訓練が可能。

新規性としては、勾配のノルムとパラメーターのノルムの比率に基づいて勾配をクリップする、適応勾配クリッピング（AGC）を提案している。

ここ最近BatchNormalizationに関心があって調べていたら、勾配の極端な偏りを減らせれば解決するのではという気はしていたので、やはりそうか、という印象。

https://arxiv.org/abs/2102.06171

takoroy

Adaptive Gradient Clipping

一般的な勾配クリッピングは、以下のような式で表される。勾配ベクトル $G$ のノルムが閾値 $\lambda$ を越えるとき、ノルムが $\lambda$ になるように正規化される。勾配クリッピングは、ハイパーパラメータである $\lambda$ の選択に非常に敏感であるため、よりロバストな方法が求められる。

G \rightarrow\left\{\begin{array}{ll} \lambda \frac{G}{\|G\|} & \text { if }\|G\|>\lambda \\ G & \text { otherwise. } \end{array}\right.

AGCでは、あるレイヤー $l$ のパラメータ $W^l \in \mathbb{R}^{N \times M}$ の元々のノルムも考慮し、 $\frac{\left\|G^{\ell}\right\|_{F}}{\left\|W^{\ell}\right\|_{F}}$ に対する比率で勾配のクリッピングを行う。ここで、 $\left\|\cdot\right\|_{F}$ はフロベニウムノルムを表し、全要素の自乗和の平方根である。なお、 $\left\|W^{\ell}\right\|_{F}$ が極端に小さいまたはゼロである時に、安定性のために $\left\|W_{i}\right\|_{F}^{\star}=\max \left(\left\|W_{i}\right\|_{F}, \epsilon\right)$ で置き換える。

さらに、経験的に判明したこととして、レイヤーごとの正規化ではなく、ユニットごとの正規化が適切であるということがわかっている。ユニットごとというのは、出力次元・チャネルごとに、という意味である。細かいことは公式の実装（JAX）を参照。最終的に、レイヤー・ユニットごとのパラメータと勾配 $W_i^{\ell}, G_i^{\ell}$ を使って、以下のように書ける。

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.

AGCは、近年のnormalized optimizerの文脈からも理解できる。LARSは、勾配の正規化を、現状のパラメータのノルムに対する比率のみ（勾配のノルムは使用しない）で行う手法であり、考え方は類似している。

takoroy

既存の BNフリーモデルであるNF-ResNetに、AGCを組み合わせた実験を行っている。AGCによって、大きなバッチサイズでも訓練が行えていることがわかった。また、 $\lambda$ が十分小さい時に訓練がうまくいくこともわかった。

また、AGCは、全結合層のLinear層に適用すると劣化することがわかった。最初のConv層については、適用する時としない時で良くなったり悪くなったり、どちらとも言い難い結果となっている。以降の実験では、全結合層のLinear層をのぞき、AGCを適用している。

takoroy

ネットワーク構造

NFNetはResNeXtをベースとして作られている。大きく、以下のような構造になっている。

stem
stage1
stage2
stage3
stage4
classifier

stemは、いくつかのConv層からなる局所特徴を抽出する部分で、3x3サイズのカーネルのConv層が4つ使われており、それぞれのストライドと出力チャネル数は、16,32,64,128と2,1,1,2である。これにより、入力画像の1/4のサイズの128チャネルの特徴マップが得られる。最後のConv層以外では、GELU活性化が適用される。

stage1~4は、 $[N, 2N, 6N, 3N]$ 個のブロックからなるメインモジュールである（ $N$ はモデルの容量をコントロールするパラメータ）。ステージ2~3の各ステージの最初はtransitionブロックと呼ばれ、特徴マップのサイズを1/2にする操作が行われる（下図左側）。それ以外のブロックは、下図右側のような形になっている。

ブロックの最初の1x1Convではチャネル数は出力チャネル数の1/2にされ、その後の2つの3x3Convでは128チャネルを1グループとするGroouped Convが適用される。その後、1x1Convによって出力のチャネル数に合わせられる。各ブロックの最後にはSqueeze & Excite層が続く。SE層は、数式で書くと $2 \sigma(F C(G E L U(F C(\operatorname{pool}(h))))) \times h$ となる。各ブロックの出力チャネル数は、 $[256, 512, 1536, 1536]$ である。

多岐にわたる最新の成果を取り入れているので、論文とそこから取り入れている工夫を列挙する。

Characterizing signal propagation to close the performance gap in unnormalized ResNets. Brock et al. (2021)
- 活性化関数としてGeLUを使っているが、出力の分散を1に保持するためのリスケールが行われる。リスケールの決め方をこの論文に依拠している。上図のScaled Activationに対応する。
- Residual構造のメインパスに対してスカラー倍をかけることで、分散を1に保持する工夫を行っている。上図の $1/\beta$ と $\alpha$ に該当する。
- Scaled Weight Standardizationと呼ばれるConv層のパラメータに対するアフィンパラメータとバイアスパラメータ（いずれも訓練対象）を導入する。上図のWS-Convの"WS"に該当する。
Batch normalization biases residual blocks towards the identity function in deep networks. De & Smith, 2020
- $\alpha$ を0に初期化し、学習対象とすることで、transitブロック以外の通常ブロックを恒等変換のように初期化する。これによって深いネットワークの安定性を改善することができる。

takoroy

モデルのバリエーションと訓練上の工夫

モデルのstage1からstage4までのブロック数をコントロールする $N$ によって、以下のようなバリエーションを作ることができる。

また、正則化のために、モデルのサイズを大きくするに従い、最後のLinear層の直前に挿入するDropoutの割合を徐々に増やしている。正則化は、WeightDecayでも実現できるが、実験の結果効果がなかったとしている。また、Dropoutなしでは、overfitが認められたという。

訓練時とテスト時の解像度は以下のように、訓練時で小さくしている。テスト時の解像度でファインチューニングするということは行っていない。

このスクラップは2021/02/27にクローズされました