[논문 리뷰] Self-distillation with Batch Knowledge Ensembling Improves ImageNet Classification

Yixiao Ge, Xiao Zhang외 7명|arXiv (Cornell University)|2021. 04. 27.

Advanced Neural Network Applications참고 문헌 58인용 수 16

한 줄 요약

이 논문은 이미지 분류 성능을 햖을 수 있도록 동일한 미니배치 내 다른 샘플들 간의 상호 샘플 유사도를 활용하여 지식을 온라인으로 앙상블하는 자기-정규화(self-distillation) 방법인 배치 지식 앙상블(BAKE)을 제안한다. 이 방법은 추가 파라미터 없이 단지 1.5%의 추가 FLOPs만으로도 최신 기준 성능을 달성하며, Swin-T의 정확도를 +0.7% 향상시킨다.

ABSTRACT

The recent studies of knowledge distillation have discovered that ensembling the "dark knowledge" from multiple teachers or students contributes to creating better soft targets for training, but at the cost of significantly more computations and/or parameters. In this work, we present BAtch Knowledge Ensembling (BAKE) to produce refined soft targets for anchor images by propagating and ensembling the knowledge of the other samples in the same mini-batch. Specifically, for each sample of interest, the propagation of knowledge is weighted in accordance with the inter-sample affinities, which are estimated on-the-fly with the current network. The propagated knowledge can then be ensembled to form a better soft target for distillation. In this way, our BAKE framework achieves online knowledge ensembling across multiple samples with only a single network. It requires minimal computational and memory overhead compared to existing knowledge ensembling methods. Extensive experiments demonstrate that the lightweight yet effective BAKE consistently boosts the classification performance of various architectures on multiple datasets, e.g., a significant +0.7% gain of Swin-T on ImageNet with only +1.5% computational overhead and zero additional parameters. BAKE does not only improve the vanilla baselines, but also surpasses the single-network state-of-the-arts on all the benchmarks.

연구 동기 및 목표

• 1-핫 레이블의 한계로 인해 감독 이미지 분류에서 성능 향상이 제한되는 문제를 해결하기 위해.
• 지식 정규화의 계산 비용과 메모리 비용을 줄이면서도 성능을 유지하거나 향상시키기 위해.
• 단일 네트워크를 사용하여 미니배치 내 샘플 간에 온라인으로 경량 지식 앙상블을 가능하게 하기 위해.
• 다른 샘플 간의 지식 전파를 통해 정교화된 소프트 타겟을 도출하여 이미지 분류의 일반화 및 강인성 향상시키기 위해.

제안 방법

• BAKE는 동일한 미니배치 내 다른 샘플들의 예측을 상호 샘플 유사도로 가중하여 실시간으로 지식을 앙상블한다.
• 상호 샘플 유사도는 현재 네트워크의 특징을 사용하여 추정되며, 유사한 샘플이 지식 앙상블에 더 큰 기여를 하도록 보장한다.
• 지식 전파 및 앙상블는 수렴할 때까지 반복적으로 적용되며, 높은 계산 비용을 피하기 위해 효율적인 추론을 통해 근사된다.
• 이 방법은 복수의 교사 네트워크나 추가 브랜치가 필요 없이 단일 네트워크를 사용해 정교화된 소프트 타겟을 생성함으로써 자기-정규화를 수행한다.
• 이 프레임워크는 ResNet 및 Swin Transformer와 같은 다양한 아키텍처와 호환되며, 동기화 배치 정규화를 사용한 분산 학습을 지원한다.

실험 결과

연구 질문

• RQ1 미니배치 내 샘플 간의 지식 앙상블이 파라미터나 계산 부담을 추가하지 않고도 정규화 성능을 향상시킬 수 있는가?
• RQ2 실시간으로 수행되는 상호 샘플 간 지식 전파가 자기-정규화에서 소프트 타겟의 품질에 어떤 영향을 미치는가?
• RQ3 BAKE는 최신 단일 네트워크 기반 방법들과 비교해 ImageNet 및 후속 작업에서 얼마나 정확도를 향상시킬 수 있는가?
• RQ4 BAKE는 레이블 노이즈와 분포 이탈이 발생하는 테스트 시점 강인성 벤치마크에서 얼마나 강인한가?

주요 결과

• BAKE는 훈련 시 3.7%의 추가 FLOPs만으로 ResNet-50의 ImageNet 정확도를 +1.2% 향상시켜 76.8% → 78.0%로 개선한다.
• BAKE는 단지 +1.5%의 계산 부담과 추가 파라미터 없이 Swin-T의 ImageNet 정확도를 +0.7% 향상시켜 81.3% → 82.0%로 개선한다.
• BAKE는 CIFAR-100, TinyImageNet, CUB-200-2011, MIT67, Stanford Dogs 등 평가된 모든 벤치마크에서 단일 네트워크 최신 기준 방법들을 뛰어넘는 성능을 보인다.
• 이 방법은 COCO 검출 및 세그멘테이션 작업에서의 전이 학습 성능을 향상시켜 학습된 표현의 일반화 능력 향상을 보여준다.
• BAKE는 ImageNet-R, ImageNet-A, ImageNet-P와 같은 변형된 데이터셋에서의 테스트 시점 강인성을 향상시켜 모델 강인성 향상의 잠재력을 보여준다.
• 절단 분석 결과, 근사 추론을 통한 무한 반복 지식 앙상블이 단일 반복 앙상블보다 더 강인한 성능을 내며, 특히 per-class 샘플링을 사용한 분산 학습에서 sync BatchNorm이 핵심임을 확인한다.