[논문 리뷰] Sharpness-Aware Training for Free

Jiawei Du, Daquan Zhou외 3명|arXiv (Cornell University)|2022. 05. 27.

Advanced Neural Network Applications인용 수 15

한 줄 요약

이 논문은 현재 및 과거 가중치에서의 모델 출력 간의 KL 발산을 기반으로 한 경로 손실을 최소화하는 방식으로 근본적인 계산 오버헤드 없이 날카움 인식 학습을 구현하는 새로운 방법인 Sharpness-Aware Training for Free (SAF)를 제안한다. SAF는 표준 SGD와 거의 동일한 속도로 학습하면서 ImageNet에서 SAM과 동일한 일반화 성능을 달성하여 대규모 모델에 매우 효율적이다.

ABSTRACT

Modern deep neural networks (DNNs) have achieved state-of-the-art performances but are typically over-parameterized. The over-parameterization may result in undesirably large generalization error in the absence of other customized training strategies. Recently, a line of research under the name of Sharpness-Aware Minimization (SAM) has shown that minimizing a sharpness measure, which reflects the geometry of the loss landscape, can significantly reduce the generalization error. However, SAM-like methods incur a two-fold computational overhead of the given base optimizer (e.g. SGD) for approximating the sharpness measure. In this paper, we propose Sharpness-Aware Training for Free, or SAF, which mitigates the sharp landscape at almost zero additional computational cost over the base optimizer. Intuitively, SAF achieves this by avoiding sudden drops in the loss in the sharp local minima throughout the trajectory of the updates of the weights. Specifically, we suggest a novel trajectory loss, based on the KL-divergence between the outputs of DNNs with the current weights and past weights, as a replacement of the SAM's sharpness measure. This loss captures the rate of change of the training loss along the model's update trajectory. By minimizing it, SAF ensures the convergence to a flat minimum with improved generalization capabilities. Extensive empirical results show that SAF minimizes the sharpness in the same way that SAM does, yielding better results on the ImageNet dataset with essentially the same computational cost as the base optimizer.

연구 동기 및 목표

• 두 번의 역전파를 수행하는 바람에 학습 시간이 두 배로 증가하는 날카움 인식 최소화(SAM)의 높은 계산 비용을 해결하기 위해.
• 표준 최적화보다 거의 추가 계산이 없는 날카움 인식 학습 방법을 개발하기 위해.
• SGD와 동일한 속도로 SAM과 동등하거나 이를 초월하는 일반화 성능를 유지하거나 향상시키기 위해.
• 메모리와 계산 자원이 중요한 대규모 데이터셋인 ImageNet-1k 및 이를 초월하는 데이터셋에서 효율적인 날카움 인식 학습을 가능하게 하기 위해.
• ImageNet-21K나 JFT-300M와 같은 매우 대규모 데이터셋에 적용 가능한 메모리 효율적인 변형인 MESA를 도입하기 위해.

제안 방법

• 현재 가중치와 과거 가중치를 사용한 신경망의 출력 분포 간의 KL 발산을 기반으로 한 새로운 경로 손실을 제안한다.
• 이 경로 손실을 SAM의 날카움 측정 기준의 대체 수단으로 사용하여 학습 경로를 따라 손실의 변화 속도를 캡처한다.
• 기본 학습 손실과 함께 이 경로 손실을 최소화함으로써 표준 최적화에 통합하여 손실의 급격한 하강을 효과적으로 억제한다.
• 학습 중 추가 전방향 프로파일 없이도 이전 단계의 모델 출력을 기억하여 경로 손실을 계산한다.
• 모델 가중치의 지수이동평균(EMA)을 사용해 과거 모델 출력을 생성하는 메모리 효율적인 변형인 MESA를 도입한다.
• 두 단계 학습 과정을 사용한다: 첫 번째 단계에서는 표준 SGD로 모델을 학습하고, 두 번째 단계에서는 이전 단계의 출력을 저장한 상태에서 경로 손실을 최소화한다.

실험 결과

연구 질문

• RQ1 표준 SGD와 비교해 근본적인 추가 계산 오버헤드로 날카움 인식 학습을 달성할 수 있는가?
• RQ2 출력 분포 이동을 기반으로 한 경로 기반 손실이 SAM에서 사용하는 날카움 측정 기준을 효과적으로 근사할 수 있는가?
• RQ3 제안된 방법인 SAF가 SAM의 일반화 성능를 유지하거나 초월하면서도 학습 속도를 그대로 유지할 수 있는가?
• RQ4 정확도와 메모리 효율성 측면에서 ImageNet-1k 및 ImageNet-21k와 같은 대규모 데이터셋으로의 확장성은 어떻게 되는가?
• RQ5 ImageNet-21K나 JFT-300M와 같은 매우 대규모 학습 작업에까지 적용 가능한 메모리 효율적인 변형인 MESA를 설계할 수 있는가?

주요 결과

• ResNet-50를 사용해 ImageNet-1k에서 SAF는 96.37%의 top-1 정확도를 달성했으며, SAM의 96.50%와 유사하지만 표준 SGD의 97.6% 수준의 학습 속도로 학습된다.
• ResNet-101에서 SAF는 96.93%의 top-1 정확도를 기록해 SAM의 성능을 그대로 유지하면서도 SGD의 99.2% 수준의 속도로 학습된다.
• Wide-ResNet-28-10에서 SAF는 97.08%의 top-1 정확도를 기록해 SAM(97.07%)를 略히 초월하며 거의 동일한 학습 속도를 유지한다.
• MESA는 PyramidNet-110에서 97.46%의 top-1 정확도를 기록해 SAM과 SAF를 모두 초월하지만, SGD의 84.3% 수준의 학습 속도만을 사용한다.
• SAF는 SAM의 100% 대비 0.7%의 추가 계산 오버헤드만을 유발하여 실질적으로 근본적인 오버헤드를 갖는다.
• KL 발산 기반의 경로 손실은 날카움 역학을 효과적으로 캡처하고 추가 역전파 없이 더 평탄한 최소값으로 수렴하는 데 기여한다.