[논문읽기] Self-training with Noisy Student improves ImageNet classification

Original Paper ) https://arxiv.org/abs/1911.04252

Self-training with Noisy Student improves ImageNet classification

---

Situation and Problem

• SOTA 비전 모델들은 지도 학습으로 trained되는데, 라벨링된 데이터셋이 클스록 성능이 잘 나온다.
  • 즉, 데이터셋이 클수록 성능이 오른다.
• 모델 성능을 높이는데 있어서 우리는 사용 가능한 엄청난 양의 라벨링되지 않은 데이터셋을 이용하지 않았다.

What to Accomplish

• 우리는 라벨링되지 않은 데이터셋을 사용하여 모델 성능을 높이고자 한다. 이렇게 고안된 방법이 바로 Noisy Student Training이다.
• Noisy Student 학습은 라벨링된 데이터가 충분할 때도 좋은 성능을 내는 semi-supervised learning(준지도) 학으로, 높은 정확도와 견고함을 보여준다.
• Noisy Student는 높은 confidence가 나오는 데이터만 가져오는 self-training과 큰 신경망(teacher network)과 비슷한 성능을 내는 작은 신경망(student network)을 학습시키는 distillation의 아이디어를 확장하였다.
• Noisy Student는 teacher network와 같거나 더 큰 student network를 가지며, 학습하는 동안 student network에 noise를 주입한다.

Key Elements of the Approach

• Iterative Learning

1. noise가 있는 labeled dataset에 대한 CE를 낮추는 방향으로 teacher model을 학습시킨다.
2. 학습된 teacher model을 가지고 unlabeled dataset의 라벨을 예측한다. 이 때 unlabeled dataset은 어떤 noise도 존재하지 않는다.
3. teacher model보다 크거나 같은 student model을 labeled dataset과 unlabeled dataset의 CE를 모두 낮추는 방향으로 학습시킨다. 이 때 dataset에 noise가 존재한다.
4. 학습된 student model을 teacher model로 대체하고 2번으로 돌아간다.

• Adding Noise to Student Model
  • 여기서 사용되는 noise는 크게 input noise와 model noise 두 종류로 분류할 수 있다.
    • input noise는 RandAugment를 사용하여 data augmentation을 수행한 것을 말한다.
    • model noise로는 일정 확률로 뉴런을 드롭시키는 dropout과 layer를 드롭시키는 stochastic depth가 있다.
    • RandAugment를 줌으로써 student model은 다양한 버전의 이미지에도 견고함을 얻을 수 있고 dropout과 stochastic depth를 사용함으로써 추론 시 앙상블과 같은 효과를 낼 수 있다.
• Data Filtering and Balancing
  • unlabeled dataset에 대해 data filtering과 balancing을 적용하였다.
  • unlabeled dataset(JFT dataset)를 labeled dataset(ImageNet dataset)으로 학습된 EfficientNet-B0로 추론했을 때 confidence score가 0.3보다 높은 것만 가져간다. ⇒ data filtering
  • data filtering을 적용하면 각 클래스당 최대 130K개의 이미지를 갖는데 130K보다 적은 클래스의 경우 랜덤으로 이미지를 복사하여 130K를 채워준다. 그렇게 되면 모든 각 클래스는 130K의 이미지를 갖게 된다.
  • 단순 복사를 하기 때문에 130M 데이터셋 중 오직 81M개의 고유 이미지를 갖는다.

Experiments

Details

• labeled dataset으로는 ImageNet 2012 ILSVRC challenge 데이터셋을 사용하였다.
• unlabeled dataset으로는 300M개의 이미지를 가진 JFT dataset을 사용하였다. 원래 라벨이 있는 데이터셋이지만 라벨을 무시하고 unlabeled dataset으로 취급하였다.
  • JFT dataset에 포함된 ImageNet의 validation dataset을 제외시켰다.
  • data filtering과 balancing을 적용하여 각 클래스당 130K 이미지를 갖도록 하였다.
• baseline model로 EfficientNet-B0 ~ EfficientNet-B7을 사용하였으며, EfficientNet-B7보다 더 깊고 넓은 EfficientNet-L2도 사용하였다.
  • EfficientNet-L2의 큰 모델 크기 때문에 Efficientnet-B7과 학습 시간을 비교했을 때 거의 5배 이상 긴 학습 시간을 보여주었다.
• labeled dataset의 경우 batch size를 2048로 설정하였다. 다만 Out-Of-Memory 에러가 나는 경우 batch size를 줄여주었다.
  • 512, 1024, 2048을 batch size로 사용하였는데 같은 성능을 보여주었다.
  • EfficientNet-B4 이상의 큰 모델은 350 epoch 동안, 그 이하의 작은 모델은 700 epoch 동안 student model을 학습시켰다.
  • batch size 2048인 labeled data에 대해서 learning rate는 0.128부터 시작해 0.97의 decay를 주었다.
    • 350 epoch로 학습하는 경우 매 2.4 epoch마다, 700 epoch로 학습하는 경우 매 4.8 epoch마다 decay가 일어났다.
• unlabeled dataset은 labeled dataset보다 큰 batch size를 갖는데, labeled dataset의 batch size보다 14배 큰 batch size를 적용하였다.
• stochastic depth의 경우 생존 확률을 0.8, dropout은 드롭 확률을 0.5, RandAugment의 magnitude는 27을 적용하였다.
• Iterative Training 시 EfficientNet-B7을 초기 teacher model로 두고 EfficientNet-L2를 student model로 사용하여 총 3번의 iteration을 수행하였다.
  • 이 때 마지막 iteration에서 unlabeled dataset의 batch size는 labeled dataset의 batch size보다 28배 더 크다.

Results

ImageNet의 validation set의 정확도 비교

• SOTA 모델보다 정말 좋은 성능을 내고 있으며, EfficientNet-L2와 비교했을 때 정확도가 꽤 향상된 것을 볼 수 있다.
• 3.5B의 labeled dataset을 추가로 사용하여 학습한 모델보다 훨씬 더 적은 양의 unlabeld dataset을 사용함에도 더 좋은 성능을 보여준다.

Iterative Training이 없을 때 EfficientNetB0-B7의 성능 비교

1번의 iteration만 수행하고 teacher model과 student model을 동일 모델로 사용했을 때의 결과를 나타낸 그래프.

• 이 때 EfficientNetB1-B7은 unlabeled dataset의 batch size는 labeled dataset의 3배 더 큰 batch size를 사용했다.
• EfficientNetB0는 unlabeled datatset과 labeled dataset의 batch size를 같에 설정하였다.

• Noisy Student Training을 적용한 모델이 더 좋은 성능을 얻은 것을 볼 수 있다.

ImageNet-A, ImageNet-C, ImageNet-P에 대한 견고성(Robustness) 실험

 

• ImageNet-A, ImageNet-C, ImageNet-P는 각각 블러, 회전, 스케일링과 같이 이미지에 심한 변형을 가한 테스트 데이터셋이다.
• Noisy Student Training을 적용한 EfficientNet-L2 모델이 뛰어난 성능을 보여주는 것을 볼 수 있다. 즉 견고함이 정말 좋다고 할 수 있음! 
•  

ImageNet-A, ImageNet-C, ImageNet-P의 일부 샘플 데이터로 추론한 결과를 나타낸 것이다. 검정색 볼드체는 Noisy Student Training을 적용한 EfficientNet이며, 빨간색 볼드체는 적용하지 않은 모델이다.

Adversial Perturbation에 잘 대응하는지에 대한 실험

: FGSM attack을 가했을 때 Noisy Student Training을 적용했을 때 모델이 더 견고함을 볼 수 있다.

Unlabeled Dataset의 크기와 Noise 유무에 따른 성능 비교

: iteration이 증가할수록 정확도가 증가하는 것을 볼 수 있다.

Iterative Learning의 효과

: iteration이 증가할수록 정확도가 증가하는 것을 볼 수 있다.

Additional Ablation Study

다음과 같은 요인은 더 좋은 성능으로 이어졌다.

• 더 큰 teacher model, student model
• 더 많은 양의 unlabeled dataset
• soft pseudo label
• data balancing
• labeled dataset과 unlabeled dataset을 합쳐 학습하는 것
• labeled dataset보다 더 큰 unlabeled dataset의 batch size
• (somtimes) student model을 랜덤으로 초기화시킨 것

Conclusion

 

• unlabeled dataset을 사용해 더 좋은 정확도와 견고함을 가질 수 있음을 보여주었다.
• unlabeled dataset을 충분히 사용할 수 있는 간단하고 효과적인 self-training 기법을 찾음
  • student model에 noise를 추가함으로써 teacher model 보다 더 좋은 성능을 낼 수 있었다.