[논문 리뷰] Surrogate Gap Minimization Improves Sharpness-Aware Training

Juntang Zhuang, Boqing Gong외 7명|arXiv (Cornell University)|2022. 03. 15.

Image Enhancement Techniques인용 수 13

한 줄 요약

이 논문은 예측 가능한 날카기 정도를 측정하기 위한 보조 간격(surrogate gap)을 도입하여 샤프nees-aware minimization(SAM)의 성능을 햖थन하는 새로운 훈련 방법인 GSAM을 제안한다. 이 보조 간격은 헤시안 기반 날카기 정도를 계산적으로 효율적인 대체 측정법으로 사용하며, 손실함수와 함께 날카기 정도를 직접 최소화할 수 있게 한다. GSAM은 각 업데이트 단계에서 두 단계를 수행한다: SAM과 동일하게 변형된 손실함수를 최소화하는 것과, 수직 방향으로 상승하여 보조 간격을 줄이는 것. 이로 인해 더 평탄한 최소값을 얻고 일반화 성능이 향상되며, ImageNet에서 ViT-B/32에 대해 AdamW 대비 5.4%의 top-1 정확도 향상을 달성한다.

ABSTRACT

The recently proposed Sharpness-Aware Minimization (SAM) improves generalization by minimizing a extit{perturbed loss} defined as the maximum loss within a neighborhood in the parameter space. However, we show that both sharp and flat minima can have a low perturbed loss, implying that SAM does not always prefer flat minima. Instead, we define a extit{surrogate gap}, a measure equivalent to the dominant eigenvalue of Hessian at a local minimum when the radius of the neighborhood (to derive the perturbed loss) is small. The surrogate gap is easy to compute and feasible for direct minimization during training. Based on the above observations, we propose Surrogate extbf{G}ap Guided extbf{S}harpness- extbf{A}ware extbf{M}inimization (GSAM), a novel improvement over SAM with negligible computation overhead. Conceptually, GSAM consists of two steps: 1) a gradient descent like SAM to minimize the perturbed loss, and 2) an extit{ascent} step in the extit{orthogonal} direction (after gradient decomposition) to minimize the surrogate gap and yet not affect the perturbed loss. GSAM seeks a region with both small loss (by step 1) and low sharpness (by step 2), giving rise to a model with high generalization capabilities. Theoretically, we show the convergence of GSAM and provably better generalization than SAM. Empirically, GSAM consistently improves generalization (e.g., +3.2\% over SAM and +5.4\% over AdamW on ImageNet top-1 accuracy for ViT-B/32). Code is released at \url{ https://sites.google.com/view/gsam-iclr22/home}.

연구 동기 및 목표

• SAM의 한계를 해결하기 위해, 날카기 정도를 더 신뢰할 수 있는 척도로 측정함으로써 날카롭고 평탄한 최소값 간의 구분을 향상시키는 것.
• 특히 작은 이웃 반경에서 국소 최소값에서 주요 헤시안 고유값을 근사할 수 있는 계산적으로 다룰 수 있는 보조 간격을 정의하는 것.
• 변형된 손실함수와 보조 간격을 동시에 최소화하는 훈련 방법을 개발하여 더 평탄한 최소값과 더 나은 일반화 성능를 달성하는 것.
• 제안된 GSAM 알고리즘의 수렴성을 증명하고, 일반화 성능에서 SAM에 비해 이론적으로 우월함을 입증하는 것.
• 다양한 아키텍처(ResNets, Vision Transformers, MLP-Mixer)에서 GSAM의 성능을 실증적으로 검증하여 SAM 및 기준 최적화기 대비 일관된 성능 향상을 보이는 것.

제안 방법

• 보조 간격 $ h(w) = f_p(w) - f(w) $ 를 정의한다. 여기서 $ f_p(w) $ 는 변형된 손실함수이고 $ f(w) $ 는 표준 손실함수이며, 이는 국소 최소값에서 주요 헤시안 고유값을 근사한다.
• 기울기 $ \nabla f(w) $ 를 $ \nabla f_p(w) $ 와 평행하고 수직인 성분으로 분해함으로써, 수직 방향으로의 최적화를 정밀하게 제어할 수 있도록 한다.
• 변형된 손실함수 $ f_p(w) $ 에 대해 경사 하강법 단계를 수행하여 변형된 손실함수를 최소화함으로써 SAM의 업데이트 규칙을 유지한다.
• 보조 간격 $ h(w) $ 를 최소화하기 위해 수직 성분 $ \nabla_\perp f(w) $ 에 따라 상승 단계를 적용함으로써, $ f_p(w) $ 는 수직성 덕분에 변화하지 않도록 한다.
• 상승 단계의 크기를 제어하기 위해 하이퍼파rameter $ \alpha $ 를 도입하여 손실과 날카기 정도 최소화 간의 유연한 트레이드오프를 가능하게 한다.
• 감소하는 이웃 반경 $ \rho_t $ 를 가정함으로써 이론적 수렴성을 확보하고, GSAM의 수렴 분석을 뒷받침한다.

실험 결과

연구 질문

• RQ1 SAM에서 변형된 손실함수는 날카로운 최소값과 평탄한 최소값을 신뢰성 있게 구분할 수 있는가, 아니면 둘 다 낮은 값을 가질 수 있는가?
• RQ2 훈련 중에 사용할 수 있는, 헤시안 기반 날카기 정도 측정법의 계산적으로 효율적인 대체 방법이 존재하는가?
• RQ3 변형된 손실함수와 표준 손실함수의 차이로 정의된 보조 간격을 최소화하는 것이 평탄한 최소값으로의 최적화를 효과적으로 이끄는가?
• RQ4 변형된 손실함수 감소와 보조 간격 감소를 위한 수직 방향 상승을 조합한 이중 단계 업데이트 전략이 SAM보다 더 나은 일반화 성능을 낳는가?
• RQ5 표준 가정 하에 GSAM이 수렴함을 이론적으로 증명하고, SAM에 비해 일반화 성능에서 명백한 우월성을 입증할 수 있는가?

주요 결과

• 작은 이웃 반경 $ \rho $ 를 가진 경우, 보조 간격 $ h(w) = f_p(w) - f(w) $ 는 국소 최소값에서 주요 헤시안 고유값과 이론적으로 동치이며, 이는 유효하고 효율적인 날카기 정도의 대체 측정법임을 의미한다.
• GSAM은 ImageNet에서 ViT-B/32를 사용할 때 SAM 대비 3.2%의 top-1 정확도 향상을, AdamW 대비 5.4%의 향상을 기록하여 일관된 일반화 성능 향상을 입증한다.
• ablation 실험을 통해 GSAM의 상승 단계가 성능 향상의 주요 원동력임을 입증하였으며, 이는 일정한 $ \rho_t $ 와 감소하는 $ \rho_t $ 스케줄 모두에서 동일하게 관찰된다.
• 실증적 검증 결과, $ \cos\theta_t $ (즉, $ \nabla f(w) $ 와 $ \nabla f_p(w) $ 간의 각도余현수) 는 훈련 전반에 걸쳐 0.9 이상을 유지하여 고차원 매개변수 공간에서 기울기가 거의 일치한다는 이론적 가정을 지지한다.
• 보조 간격은 $ \alpha $ 가 증가할수록 감소하고, 훈련 단계가 진행될수록 증가함을 관찰하여 GSAM이 시간이 지남에 따라 모델이 점점 더 평탄한 최소값으로 이동하도록 성공적으로 이끌고 있음을 시사한다.
• GSAM는 광범위하게 적용 가능하며, SAM 대비 거의 무시할 수 없는 계산 오버헤드를 유발하여 대규모 훈련에 실용적이다.