Revisiting Sharpness-Aware Minimization: A More Faithful and Effective Implementation

이 논문은 Sharpness-Aware Minimization (SAM) 의 기존 구현 방식에 대한 직관적 해석을 제시하고, 근사 오차와 다단계 상승 시의 품질 저하 문제를 해결하기 위해 명시적 방향 추정과 최적화된 탐색 공간을 도입한 새로운 알고리즘인 eXplicit SAM (XSAM) 을 제안하여 기존 방법보다 우수한 일반화 성능을 입증합니다.

핵심

1. 배경: 왜 '날카로운 최소점'을 피해야 할까요?

머신러닝 모델을 훈련시킬 때, 우리는 손실 (오차) 이 가장 작은 지점을 찾습니다. 그런데 이 지점이 두 가지 종류가 있을 수 있습니다.

• 날카로운 최소점 (Sharp Minimum): 마치 뾰족한 바위 꼭대기에 서 있는 상태입니다. 여기서는 손실은 아주 작지만, 조금만 발을 움직여도 (데이터가 조금만 달라져도) 미끄러져서 큰 오차로 떨어집니다. 즉, 새로운 상황에 약한 (일반화 능력이 낮은) 모델입니다.
• 넓고 평평한 최소점 (Flat Minimum): 마치 넓은 평야에 서 있는 상태입니다. 여기서는 손실도 작지만, 발을 조금 움직여도 여전히 낮은 곳에 있습니다. 즉, 새로운 상황에도 잘 적응하는 (일반화 능력이 높은) 모델입니다.

우리는 당연히 **넓은 평야 (Flat Minimum)**를 원합니다.

2. 기존 방법 (SAM) 의 문제점: "눈을 감고 방향을 잡다"

기존의 SAM이라는 방법은 "주변을 훑어보면서 가장 높은 곳 (최대 손실) 을 찾아내고, 그 반대 방향으로 이동하자"는 아이디어입니다.

하지만 SAM 은 실제로는 이렇게 작동합니다:

1. 현재 위치에서 조금 올라가서 (기울기를 따라) 가상의 고지대를 찾습니다.
2. 그 가상의 고지대에서 바라본 방향을 보고, 현재 위치로 돌아와서 그 방향으로 이동합니다.

여기서 문제가 생깁니다.
논문은 SAM 의 이 방식을 **"눈을 감고, 멀리 떨어진 곳에서 들은 소리를 듣고 방향을 잡는 것"**에 비유합니다.

• 비유: 당신이 산에 서 있다고 칩시다. SAM 은 "조금만 올라가서 (가상 고지대) 주변을 둘러보고, '저기 저쪽이 가장 위험해!'라고 외친 뒤, 그 소리를 듣고 현재 위치에서 그 방향으로 뛰어내리는 것"입니다.
• 문제점:
  1. 정확하지 않음: 멀리서 들은 소리가 정확한 방향을 알려주지 못할 때가 많습니다. (가상 고지대의 기울기가 실제 최대 손실 방향과 다를 수 있음)
  2. 멀어질수록 더 나쁨: 한 번만 올라가도 문제가 있는데, 여러 번 (Multi-step) 올라가서 방향을 잡으면, 그 소리는 더 이상 현재 위치와 무관해져서 엉뚱한 곳으로 가게 됩니다.

3. 새로운 해결책: XSAM (Explicit Sharpness-Aware Minimization)

저자들은 "그냥 소리를 듣고 뛰는 게 아니라, 직접 가서 가장 높은 곳을 찾아보자"고 제안합니다. 이것이 XSAM입니다.

XSAM 의 작동 원리 (창의적인 비유):

1. 직접 탐사 (Explicit Estimation):

   • 기존 SAM 이 "가상 고지대에서 방향을 짐작"했다면, XSAM 은 "현재 위치에서 사방팔방을 직접 훑어보며 가장 높은 곳을 찾습니다."
   • 하지만 모든 방향을 다 볼 수는 없으니, **가장 유력한 두 가지 방향 (현재 기울기 방향과 가상 고지대 방향)**이 만드는 2 차원 평면만 집중해서 봅니다.

2. 스마트한 나침반 (Dynamic Interpolation):

   • 이 두 방향 사이에서 **"어느 각도로 가야 가장 높은 곳에 닿을까?"**를 계산합니다. (수학적으로는 구면 보간법을 사용)
   • 이 계산은 매번 할 필요 없이, 학습이 진행되는 동안 천천히 변하므로 가끔만 업데이트하면 됩니다. (계산 비용이 거의 들지 않음)

3. 정확한 탈출:

   • 이렇게 찾아낸 정확한 '위험 지점'의 반대 방향으로 이동합니다.
   • 결과적으로 모델은 날카로운 바위 꼭대기에서 벗어나, 넓은 평야로 안전하게 이동하게 됩니다.

4. 왜 XSAM 이 더 좋은가요?

• 정확성: "멀리서 들은 소리"가 아니라 "직접 찾아낸 지도"를 따르므로, 가장 위험한 곳 (최대 손실) 을 훨씬 정확히 피할 수 있습니다.
• 안정성: 기존 SAM 은 여러 단계로 올라갈수록 엉뚱한 방향으로 갔지만, XSAM 은 단계 수와 상관없이 항상 정확한 방향을 찾습니다.
• 효율성: 직접 찾아보는 과정이 있지만, 아주 가끔만 업데이트하므로 학습 속도는 기존 SAM 과 거의 같습니다. (계산 비용 증가가 미미함)

5. 결론: "눈을 뜨고 걷는 것"

이 논문의 핵심 메시지는 이렇습니다.

"기존의 SAM 은 눈을 감고 멀리서 방향을 짐작하는 방식이라서 가끔 엉뚱한 곳으로 가거나, 여러 걸음 걸을수록 더 헷갈려졌습니다.
하지만 우리가 만든 XSAM은 눈을 뜨고 주변을 직접 훑어보며 가장 위험한 곳을 정확히 찾아내고, 그 반대 방향으로 똑바로 이동합니다.
그 결과, 모델은 더 넓은 평야 (더 좋은 성능) 에 정착하게 되었고, 계산 속도도 거의 느려지지 않았습니다."

이 방법은 다양한 이미지 인식, 번역, 언어 모델 등 여러 분야에서 기존 방법들보다 더 좋은 성능을 보여주었습니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

Sharpness-Aware Minimization (SAM) 은 모델의 일반화 성능을 향상시키기 위해 매개변수 주변의 사전 정의된 이웃 (neighborhood) 내에서 최대 훈련 손실 (worst-case loss) 을 최소화하는 것을 목표로 합니다. 그러나 SAM 의 실제 구현은 다음과 같은 근사 (approximation) 를 사용합니다:

1. 현재 매개변수에서 기울기 상승 (gradient ascent) 을 1 단계 또는 몇 단계 수행하여 '상승점 (ascent point)'을 찾습니다.
2. 이 상승점에서 계산된 기울기를 사용하여 현재 매개변수를 업데이트합니다.

기존 연구들은 이 근사가 상승점과 현재 매개변수 사이의 야코비안 (Jacobian) 을 무시함으로써 목적 함수를 근사적으로 최적화한다고 설명해 왔습니다. 하지만, 왜 상승점에서 계산된 '비국소적 (nonlocal)' 기울기가 현재 매개변수를 업데이트하는 데 더 효과적으로 작용하는지에 대한 직관적이고 직접적인 설명은 부족했습니다.

또한, 저자들은 기존 SAM 구현의 두 가지 근본적인 한계를 발견했습니다:

• 부정확한 근사: 단일 단계 상승점의 기울기로 최대값 방향을 근사하는 것이 종종 부정확하고 불안정합니다.
• 다중 단계 (Multi-step) 의 성능 저하: 상승 단계 수 ($k$) 가 증가할수록 근사 품질이 오히려 저하되어, 기존 다중 단계 SAM 변형들이 예상보다 성능이 떨어지는 현상이 발생합니다.

2. 제안 방법: XSAM (Methodology)

이러한 한계를 극복하기 위해 저자들은 eXplicit Sharpness-Aware Minimization (XSAM) 을 제안했습니다. XSAM 은 SAM 의 근사 부정확성을 해결하고, 학습 중 변화하는 손실 지형 (loss landscape) 에 적응하기 위해 다음과 같은 새로운 접근법을 취합니다.

핵심 아이디어

기존 SAM 이 상승점의 기울기를 단순히 적용하는 대신, 현재 매개변수에서 이웃 내 최대값 (maximum) 으로 향하는 방향을 명시적으로 (explicitly) 추정합니다.

알고리즘의 주요 단계

1. 2 차원 탐색 공간 정의: 전체 고차원 이웃을 탐색하는 것은 계산 비용이 너무 큽니다. 따라서 XSAM 은 최종 상승점 ($\vartheta_k$) 에서의 기울기 ($g_k$) 와 현재 매개변수 ($\vartheta_0$) 에서 $\vartheta_k$ 로 향하는 벡터 ($v_0$) 가 span 하는 2 차원 초평면 (hyperplane) 으로 탐색 공간을 제한합니다.
   • 이 정의는 $g_k$가 가리키는 최대 손실 지점이 반드시 이 평면 안에 있음을 보장하며, 다중 단계 기울기의 정보 가치를 활용하면서도 직접 적용 시 발생하는 부정확성을 피합니다.
2. 구면 보간 (Spherical Interpolation) 을 통한 방향 탐색:
   • $v_0$와 $v_1$ ($g_k$의 정규화 벡터) 사이를 구면 선형 보간 (spherical linear interpolation) 하여 새로운 방향 $v(\alpha)$를 생성합니다.
   • $v(\alpha) = \frac{\sin((1-\alpha)\psi)}{\sin(\psi)}v_0 + \frac{\sin(\alpha\psi)}{\sin(\psi)}v_1$ (여기서 $\psi$는 두 벡터 간의 각도).
3. 명시적 최대값 방향 추정:
   • 정의된 반경 $\rho_m$ 내에서 손실 $L(\vartheta_0 + \rho_m \cdot v(\alpha))$을 최대화하는 최적의 보간 계수 $\alpha^*$를 찾습니다.
   • 이 $\alpha^*$는 현재 매개변수에서 이웃 내 최대값으로 향하는 가장 정확한 방향을 나타냅니다.
4. 동적 업데이트 및 효율성:
   • $\alpha^*$는 학습 과정에서 매우 천천히 변하므로 (Figure 2 참조), 모든 반복마다 계산하지 않고 에포크 단위로만 업데이트합니다.
   • 이로 인해 SAM 대비 계산 오버헤드는 무시할 수준 (약 2.5% 미만) 입니다.

3. 주요 기여 (Key Contributions)

1. SAM 의 새로운 직관적 해석:
   • 단일 단계 상승점의 기울기가 국소 기울기보다 이웃 내 최대값 방향을 더 잘 근사한다는 것을 시각화 및 이론적으로 증명했습니다.
   • 이는 SAM 이 국소 최대값에서 더 직접적으로 탈출할 수 있게 해준다는 것을 의미합니다.
2. SAM 의 한계 규명:
   • 단일 단계 기울기 근사가 부정확하고 불안정하며, 상승 단계 수가 증가할수록 근사 품질이 저하된다는 것을 분석하여 다중 단계 SAM 의 성능 저하 원인을 설명했습니다.
3. XSAM 제안 및 검증:
   • 명시적인 방향 추정을 통해 SAM 의 한계를 해결하는 새로운 알고리즘을 제안했습니다.
   • 단일 단계와 다중 단계 설정 모두에 적용 가능한 통합된 수식을 제공하며, negligible 한 계산 비용으로 구현됩니다.

4. 실험 결과 (Results)

XSAM 은 다양한 모델, 데이터셋, 설정에서 기존 SAM 및 그 변형들 (ASAM, LSAM, MSAM 등) 보다 일관되게 우수한 성능을 보였습니다.

• 단일 단계 설정 (Single-step):
  • CIFAR-10, CIFAR-100, Tiny-ImageNet 데이터셋에서 VGG-11, ResNet-18, DenseNet-121 등 다양한 아키텍처에서 SAM 보다 높은 테스트 정확도를 달성했습니다 (예: CIFAR-100/ResNet-18 에서 SAM 80.93% vs XSAM 81.24%).
  • 다양한 $\rho$ (perturbation radius) 값에 대해 WSAM(가중 SAM) 보다 우월한 성능을 보였습니다.
• 다중 단계 설정 (Multi-step):
  • 기존 다중 단계 SAM 변형들은 $k$가 증가함에 따라 성능이 저하되는 경향이 있었으나, XSAM 은 $k$가 증가할수록 오히려 성능이 향상되거나 유지되었습니다.
  • 이는 XSAM 이 다중 단계 상승 과정에서 얻은 정보를 효과적으로 활용하여 최대값 방향을 정확히 추정하기 때문입니다.
• 일반화 및 견고성:
  • ImageNet (ResNet-50), NMT 작업 (Transformer), ViT-Ti 등 대규모 및 다양한 작업에서도 SAM 보다 우수한 성능을 보였습니다.
  • 노이즈가 있는 데이터셋 (CIFAR-C) 에 대한 견고성 (robustness) 또한 향상되었습니다.
• 손실 지형 분석:
  • Hessian 고유값 분석과 손실 지형 시각화를 통해 XSAM 이 SAM 과 SGD 보다 더 평탄한 최소점 (flatter minima) 에 수렴함을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 SAM 의 작동 원리에 대한 깊은 통찰을 제공하며, 단순한 근사적 구현을 넘어 명시적이고 적응적인 최대값 방향 추정을 통해 Sharpness-Aware Minimization 을 더 정직하고 효과적으로 구현하는 방법을 제시했습니다.

• 이론적 기여: SAM 의 비국소적 기울기 업데이트가 왜 작동하는지에 대한 명확한 기하학적 설명을 제공했습니다.
• 실용적 기여: 계산 비용은 거의 증가시키지 않으면서 일반화 성능을 지속적으로 향상시키는 XSAM 을 제안하여, 실제 딥러닝 모델 훈련에 즉시 적용 가능한 강력한 최적화 도구가 되었습니다.
• 미래 방향: 다중 단계 SAM 의 성능 저하 문제를 해결하고, 다양한 SAM 변형 (ASAM 등) 과의 결합 가능성을 보여주어 향후 연구의 새로운 방향을 제시했습니다.

결론적으로, XSAM 은 Sharpness-Aware Minimization 의 본질적인 목표를 더 정확하게 달성하도록 설계된 차세대 최적화 알고리즘으로 평가됩니다.