[논문 리뷰] Switch EMA: A Free Lunch for Better Flatness and Sharpness

Siyuan Li, Zicheng Liu외 10명|arXiv (Cornell University)|2024. 02. 14.

Soil Mechanics and Vehicle DynamicsEngineering인용 수 3

한 줄 요약

이 논문은 각 학습 에포크의 끝에서 메인 모델에 동적으로 지수이동평균(EMA) 파arameter를 전환하는 간단하면서도 효과적인 방법인 Switch EMA(SEMA)를 제안한다. 빠르고 날카로운 수렴을 보이는 업데이트와 평탄하고 일반화에 최적화된 EMA 상태를 번갈아가며 활용함으로써, SEMA는 추가적인 계산 비용 없이도 시각 및 언어 작업 전반에서 뛰어난 모델 일반화 성능과 더 빠른 수렴 속도를 달성한다. 이는 딥러닝 학습을 위한 즉시 사용 가능한 '무료 점심'을 제공한다.

ABSTRACT

Exponential Moving Average (EMA) is a widely used weight averaging (WA) regularization to learn flat optima for better generalizations without extra cost in deep neural network (DNN) optimization. Despite achieving better flatness, existing WA methods might fall into worse final performances or require extra test-time computations. This work unveils the full potential of EMA with a single line of modification, i.e., switching the EMA parameters to the original model after each epoch, dubbed as Switch EMA (SEMA). From both theoretical and empirical aspects, we demonstrate that SEMA can help DNNs to reach generalization optima that better trade-off between flatness and sharpness. To verify the effectiveness of SEMA, we conduct comparison experiments with discriminative, generative, and regression tasks on vision and language datasets, including image classification, self-supervised learning, object detection and segmentation, image generation, video prediction, attribute regression, and language modeling. Comprehensive results with popular optimizers and networks show that SEMA is a free lunch for DNN training by improving performances and boosting convergence speeds.

연구 동기 및 목표

• 기존의 가중치 평균화(WA) 방법이 평탄함을 향상시키지만 종종 최종 성능 향상에 실패하거나 추가적인 추론 비용을 수반한다는 한계를 해결하기 위해.
• 빠르고 날카로운 수렴과 평탄하고 일반화에 최적화된 성능의 이점을 하나의 효율적인 학습 전략으로 통합하기 위해.
• 각 에포크 후 메인 모델에 EMA 파arameter를 전환하는 최소한의 수정을 제안함으로써 최적화 동역학과 최종 모델 성능를 동시에 향상시키기 위해.
• SEMA가 추가적인 계산 비용 없이 다양한 작업과 아키텍처에서 정확도와 수렴 속도를 향상시켜 '무료 점심'처럼 작용함을 입증하기 위해.

제안 방법

• SEMA는 동적 전환 메커니즘을 도입한다: 각 학습 에포크의 끝에서 EMA 평균화된 가중치가 메인 모델의 파arameter로 복사된다.
• 이로 인해 두 단계의 학습 루프가 형성된다. 메인 모델은 날카로운 손실 영역을 신속하게 탐색하고, EMA 모델은 더 평탄하고 일반화에 유리한 경로를 추적한다.
• 전환 과정은 각 새로운 에포크가 평탄하고 일반화에 최적화된 상태에서 시작되도록 보장하여 더 나은 최소값으로 더 빠르게 수렴할 수 있도록 한다.
• 이 방법은 완전히 즉시 사용 가능한 방식이며, 기존의 EMA 구현에 단 한 줄의 코드 변경만으로도 적용 가능하다.
• SEMA는 학습 기간 동안만 작동하므로, SWA나 앙상블 기반 방법에서 흔히 발생하는 테스트 시 추론 오버헤드를 피한다.
• 빠르고 날카로운 업데이트와 안정적인 평탄한 EMA 상태의 상호보완적 강점을 활용하여 최적화 동역학을 균형 잡는다.

실험 결과

연구 질문

• RQ1 EMA와 메인 모델 가중치 간의 단순한 전환 메커니즘이 딥 네트워크 최적화에서 평탄함과 날카로움을 동시에 향상시킬 수 있는가?
• RQ2 각 에포크 후 EMA 파arameter를 메인 모델로 전환하는 것이 표준 EMA나 기본 학습보다 더 나은 일반화 성능과 더 빠른 수렴 속도를 제공하는가?
• RQ3 SEMA는 추가적인 계산 비용이나 추론 오버헤드 없이 성능 향상을 달성할 수 있는가?
• RQ4 다양한 작업에서 기존의 SWA, SAM, 또는 EMA와 비교할 때 SEMA의 최적화 동역학과 최종 모델 정확도는 어떻게 다른가?

주요 결과

• SEMA는 이미지 분류, 객체 검출, 자기지도 학습, 비디오 예측, 언어 모델링 등 12개의 다양한 벤치마크에서 일관되게 테스트 정확도를 향상시켰다.
• DeiT-S를 사용한 ImageNet-1K에서 SEMA는 84.2%의 Top-1 정확도를 기록하여 베이스라인(83.8%)과 EMA(84.0%)를 모두 초월했으며, 더 빠른 수렴 속도를 보였다.
• ResNet-50 Cascade Mask R-CNN를 사용한 COCO 객체 검출에서 SEMA는 베이스라인 대비 0.7% 향상된 AP를 기록했고, EMA 대비 0.5% 향상되었다.
• MoCo.V3를 사용한 CIFAR-100에서 SEMA는 선형 프로브 정확도를 88.9%로 끌어올려 EMA(88.3%)와 베이스라인(87.6%)을 모두 뛰어넘었다.
• AgeDB에서의 연령 회귀 작업에서 SEMA는 EMA 대비 0.4년, 베이스라인 대비 0.6년의 MAE 감소를 기록했다.
• 모든 평가된 작업에서 SEMA는 수렴 속도를 가속화했으며, EMA나 베이스라인 모델보다 더 이르고 급격한 성능 향상을 보이는 학습 곡선을 보였다.