High-Performance Large-Scale Image Recognition Without Normalization (2021.02.11)
https://arxiv.org/abs/2102.06171
High-Performance Large-Scale Image Recognition Without Normalization
Batch normalization is a key component of most image classification models, but it has many undesirable properties stemming from its dependence on the batch size and interactions between examples. Although recent work has succeeded in training deep ResNets
arxiv.org
Research direction & Motivation
Main Purpose
Efficient large batch training & High perfromance without normalization
Image Recognition의 분야에서 모델을 구성할때, Batch Normalization의 사용은 매우 대중적인 것으로 여겨졌다. BN이 가져오는 효과로 인해 기존보다 학습속도가 빨라졌고 대부분의 모델에서는 이를 주저없이 차용해왔다. 하지만 본 논문에서는 BN을 적용했을때의 단점에 주목했다. BN은 다음과 같은 단점을 가진다.
- Expensive computational primitive
- Discrepency between the model during training and at inference time
- Breaks the independence between training examples in the mini batch
- Cannot be used for some tasks (Interaction between train examples in a batch enables the network to cheat loss function)
제시된 단점들을 극복하기 위해 다양한 종류의 normalizer들이 제안되었으나 추론시 계산비용의 증가와 같은 normalization layer가 가지는 고유한 문제점들은 극복되지 못했다. 일례로, Deep ResNets에 대해 normalization layer 없이 학습하는 모델에 대한 연구가 이루어졌고 어느정도는 성공했지만 BN을 적용한 경우와 비교했을 때 성능이 떨어질 뿐만 아니라 높은 learning rate와 강도높은 data augmentation에 대해 안정적으로 학습하지 못하는 모습을 보였다.
마찬가지로, 최근 연구에서 Brock.et.al은 'Normalizer-Free' ResNets를 고안했고 ImageNet에 대해 BN ResNets의 성능에 상응하는 수준을 보여주었지만 큰 batch size에 불안정하며 sota인 EfficientNets에는 미치지 못하는 모습을 보인바가 있다.
논문에서는 위에서 언급한 'Normalizer-Free' ResNets(NFNets) 구조를 기반으로 큰 batch size에도 안정적으로 학습하며 높은 성능과 속도를 가지게하는 방법을 고안했다. 여기서 중심이 되는 기법이 바로 Adaptive Gradient Clipping이다.
Adaptive Gradient Clipping ( AGC )
AGC는 특정 임계값을 설정해 모델 학습과정에서 gradient descent 과정에서 gradient가 임계값을 초과하지 않도록해 모델의 안정적인 학습을 돕는 방법이다. 일종의 clipping 방법 중 하나로 더 큰 learning rate와 batch size, 강한 data augmentation에도 안정적으로 NFNets을 학습할 수 있도록한다. 그러나 일반적인 clipping과는 다른데, 이를 비교했다.
좌측은 일반적인 GC, 우측은 AGC를 표기한 것이다. 일반적인 GC는 학습 중에 gradient norm을 gradient에 대해 나눈 값에 λ라는 clipping threshold를 설정해 unit 단위의 clipping을 진행했다.
하지만 AGC는 gradient norm이 아닌 변화하는 weight에 대해 gradient의 변화도를 지표로 사용해 clipping을 수행한다. 해당 지표는 gradient descent가 기존 weight를 얼마나 변경시키는 지에 대한 지표 역할을 하게되고, gradient가 지나치게 커지면
Experiments
실험은 크게 세가지로 나뉜다. AGC의 효과, NFNets의 구조, 평가로 구성된다.
Ablations for Adaptive Gradient Clipping
첫번째로, AGC의 효과를 검증하기위해 ImageNet에 대해 batch size를 256부터 4096까지 높여가며 NF-ResNet, NF-ResNet+AGC를 중점으로 정확도 지표를 비교했다. 이때 λ의 범위는 [0.01, 0.02, 0.04, 0.08, 0.16] 의 범위를 고려했다.
AGC를 적용하지 않은 경우, batch size가 4096일때 정확도는 급격히 떨어져 학습이 정상적으로 이루어지지 않는 모습을 보였다. 반면 AGC를 적용한 경우는 BN이 없이도 큰 batch size에 성공적으로 학습했다. 부가적으로 λ의 최적값은 optimizer, learning rate, batch size에 따라 달라지는데 일반적으로 batch size가 클수록 λ는 작은 수를 채택하는 것이 효과적인 것으로 나타났다.
두번째로, AGC가 모든 layer에 효과적인지를 실험했다. 첫번째 convolution layer, final linear layer, residual stage의 모든 block에 대해 AGC를 제거하며 다양한 조합을 실험해본 결과, final linear layer는 clipping을 하지 않는 것이 더 좋으며 첫번째 convolution layer는 clipping을 하지 않아도 안정적으로 학습하는 경우가 많다는 것을 알아냈다.
종합적으로 논문에서 AGC에 대해 알아낸 사실은 정리하면 다음과 같다.
1. 최적의 λ는 optimizer, learning rate, batch size에 의존하며, batch size가 클수록 작은 λ를 채택하는 것이 효과적임
2. 작은 batch size에서는 AGC의 효과가 감소함
3. Final linear layer에는 AGC를 적용하지 않는 것이 좋음
4. 첫번째 convolution layer는 AGC가 없어도 안정적으로 학습이 가능한 경우가 많음
NF Architectures with Improves Accuracy and Training Speed
논문에서는 모델 설계에있어서, 기존에 존재하는 accelerator에 대해 training latency를 최적화하는 방향으로 모델을 설계하고자했다. 기본구조는 GELU activation을 사용한 SE-ResNeXt-D를 사용했으며 stage depth를 [1, 2, 6, 3] 의 배수로 늘려가며 작은 모델부터 큰 모델까지를 구상했다. 동시에 stage width pattern도 기존의 [256, 512, 1024, 2048]의 width에서 [256, 512, 1536, 1536]을 기본으로 적용했다.
이는 세 번째 stage에 capacity를 늘려주고 4번째 stage의 capacity를 축소해주는 것이 학습 속도의 보존에 유리한 동시에 세번째 stage가 capacity를 확대하기에 적합하기 때문이다. 또한 bottleneck residual block에도 변형을 적용했는데, 첫 번째 layer 다음에 추가적인 3x3의 spatial convolution layer를 추가해주는 것이 더 좋은 효과를 가져왔다.
해당 구조에 추가적으로 몇가지 augmentation 방법을 적용해 성능향상을 이뤄냈는데, 그 결과는 이어지는 ImageNet에 대한 평가부분에서 같이 다뤄지므로 evaluation에서 소개하도록 하겠다.
Evaluating NFNets on ImageNet
논문에서는 ImageNet을 이용해 NFNets을 평가했는데 결과적으로 sota를 달성했다.
4096의 batch size로 360 epochs 동안 학습했으며 momentum coefficient 0.9로 Nesterov's Momentum을 사용했다. 가장 중요한 기법인 AGC의 clipping threshold는 0.01을 적용했고 learning rate는 0부터 1.6까지 5epochs에 걸쳐 점진적으로 증가시켰다.
동시에 여러가지 augmentation technique을 적용했다. MixUp, RandAugment, CutMix를 적용해 성능향상을 이뤄냈다. 특히 RA를 적용했을때 가장 큰 폭의 향상을 보였다. 이를 토대로 ImageNet에 대해서 NFNet-F0 부터 NFNet-06 까지 성능을 기존의 EffNet, BoTNet, LambdaNet, DeIT등과 비교했다.
NFNet-F5모델이 86.0%의 정확도로 sota를 달성했고, NFNet-F1은 EffNet-B7과 유사한 성능을 보였으나, 학습속도는 약 8.7배 빠르다. 또한 최근 제안된 Sharpness-Aware Minimization (SAM)을 활용해 NFNet-F6를 보면 86.5%로 큰 폭의 개선을 이끌어냈다. 하지만 SAM은 trainin speed를 크게 저하시켜 기존의 약 2배의 시간이 소요되므로 논문에서 제시하는 pipeline에는 포함되지 않는다. 동시에 전이학습에 있어서도 약 3억개의 데이터로 학습해 fine-tuning을 한 뒤, 89.2%로 높은 수준의 성능을 얻었다.
마지막으로, 앞서 적용한 data augmentation 전략을 EffNets에서도 적용해 평가했다. RA는 큰 개선을 가져왔지만 그 외 방법은 효과가 미미했으며 RMSProp 대신 cosine annealing한 SGD를 사용한 것이 오히려 성능 하락의 모습을 보였다.
이는 본 논문의 실험 결과가 단순히 강한 data augmentation에 기인해 초래된 결과는 아니며 마찬가지로 optimizer의 차이에서 기인한 결과가 아니라는 것을 보여준다.
