[논문 리뷰] Why Does Sharpness-Aware Minimization Generalize Better Than SGD?

Zixiang Chen, Junkai Zhang외 4명|arXiv (Cornell University)|2023. 10. 11.

Machine Learning and Data Classification인용 수 4

한 줄 요약

이 논문은 두 층의 ReLU 컨볼루션 네트워크에서 레이블 노이즈 상황에서 SGD보다 Sharpness-Aware Minimization(SAM)이 더 잘 일반화되는 이유에 대한 이론적 설명을 제공한다. SAM은 평탄한 최소값을 선호함으로써 초기 단계에서 노이즈 기억을 방지하고, 신호 강도가 약할 때 실패하는 SGD와는 달리 효과적인 특징 학습을 가능하게 한다. 핵심 결과는 신호 대 노이즈 비율에 기반한 일반화 성능의 형식적 단계 전이이다.

ABSTRACT

The challenge of overfitting, in which the model memorizes the training data and fails to generalize to test data, has become increasingly significant in the training of large neural networks. To tackle this challenge, Sharpness-Aware Minimization (SAM) has emerged as a promising training method, which can improve the generalization of neural networks even in the presence of label noise. However, a deep understanding of how SAM works, especially in the setting of nonlinear neural networks and classification tasks, remains largely missing. This paper fills this gap by demonstrating why SAM generalizes better than Stochastic Gradient Descent (SGD) for a certain data model and two-layer convolutional ReLU networks. The loss landscape of our studied problem is nonsmooth, thus current explanations for the success of SAM based on the Hessian information are insufficient. Our result explains the benefits of SAM, particularly its ability to prevent noise learning in the early stages, thereby facilitating more effective learning of features. Experiments on both synthetic and real data corroborate our theory.

연구 동기 및 목표

• 비선형이고 비가속도인 신경망에서, 특히 레이블 노이즈가 존재할 경우 SAM이 SGD보다 더 잘 일반화되는 이유를 이해하는 것.
• 두 층의 ReLU 컨볼루션 네트워크에서 SGD가 유해 과적합을 유도하는 조건과 유익 과적합을 유도하는 조건을 규명하는 것.
• SAM이 초기 학습 단계에서 노이즈 학습을 방지함으로써, SGD가 실패할 때조차도 유익 과적합을 달성할 수 있음을 형식적으로 증명하는 것.
• 비가속도 손실 곡면에서 일반화 오차 측면에서 SAM과 SGD 간의 이론적 차이를 설정하는 것, 특히 비가속도 손실 곡면에서의 일반화 성능을 중심으로 한다.

제안 방법

• 두 층의 ReLU 컨볼루션 네트워크의 손실 곡면을 분석하여, 비가속도이므로 헤시안 기반 설명이 무효화됨을 보여준다.
• 일반화 성능을 레이블 노이즈 하에서 연구하기 위해 신호 강도 $\|\bm{\mu}\|_2$와 입력 차원 $d$로 구성된 데이터 모델을 도입한다.
• 확률적 농도 부등식을 사용하여 SAM과 SGD의 테스트 오차를 근사하며, 예측의 마진에 초점을 맞춘다.
• 네트워크 가중치의 노름과 활성화 출력의 노름을 분석하여, 예측 마진이 큰 상태를 유지할 조건을 유도한다.
• 신호 강도 $\|\bm{\mu}\|_2$가 $d^{1/4}$에 상대적으로 어떤지에 따라 일반화 성능의 급격한 경계가 존재함을 보여주는 단계 전이 분석을 수행한다.
• SAM은 $\|\bm{\mu}\|_2 \geq \widetilde{\Omega}(1)$일 때 테스트 오차가 $\leq p + \epsilon$로 유지됨을 증명하며, 반면 SGD는 $\|\bm{\mu}\|_2 \leq O(d^{1/4})$일 때 실패함을 보여준다.

실험 결과

연구 질문

• RQ1 두 층의 ReLU 네트워크에서 SGD가 유해 과적합으로 인해 일반화에 실패하는 조건은 무엇인가?
• RQ2 비가속도이고 비선형적인 신경망에서, 특히 레이블 노이즈가 존재할 경우 SAM이 왜 SGD보다 더 잘 일반화되는가?
• RQ3 신호 강도 $\|\bm{\mu}\|_2$는 SAM과 SGD의 일반화 성능을 결정하는 데 어떤 역할을 하는가?
• RQ4 SAM은 초기 학습 단계에서 노이즈 학습을 어떻게 방지함으로써 더 나은 특징 학습을 가능하게 하는가?
• RQ5 신호 대 노이즈 비율에 따라 SGD와 SAM 간의 일반화 오차에 단계 전이가 존재하는가?

주요 결과

• SGD의 경우 $\|\bm{\mu}\|_2 \leq O(d^{1/4})$일 때, 테스트 오차는 최소 $p + 0.1$이 되어 유해 과적합을 나타낸다.
• SGD의 경우 $\|\bm{\mu}\|_2 \geq \Omega(d^{1/4})$일 때, 테스트 오차는 $p + \epsilon$ 이하로 제한되어 유익 과적합을 나타낸다.
• SAM의 경우 $\|\bm{\mu}\|_2 \geq \widetilde{\Omega}(1)$일 때, 테스트 오차는 $p + \epsilon$ 이하로 제한되어, SGD가 실패할 때조차도 유익 과적합을 보장한다.
• SAM은 초기 학습 단계에서 노이즈 학습을 효과적으로 방지하여, SGD보다 더 효율적인 특징 학습을 가능하게 한다.
• 논문은 일반화 성능에 대한 형식적 단계 전이를 확립하였으며, 특히 신호가 약한 영역에서 SGD가 실패할 때 SAM이 뛰어난 성능을 보임을 보여준다.
• 이론적 분석은 SAM의 성공이 헤시안 기반의 가속도 가정에 기반하지 않음을 확인하였으며, 이는 비가속도 ReLU 네트워크에도 적용 가능함을 의미한다.