기술 요약: 폴리악 유형(Polyak-type) 스텝 사이즈를 이용한 적응형 샤프니스 인식 최소화(Adaptive Sharpness-Aware Minimization)

문제 정의

샤프니스 인식 최소화(SAM)와 그 변형인 비정규화 SAM(USAM)은 손실 곡면(loss landscape)의 샤프니스를 최소화함으로써 심층 신경망(DNN)의 일반화 성능을 향상시키는 유망한 최적화 알고리즘으로 주목받고 있습니다. 그러나 이러한 방법들의 중대한 실무적 한계는 학습률(learning rate) 선택에 매우 민감하다는 점입니다. 대부분의 학습 알고리즘과 마찬가지로, SAM 변형 모델들은 최적의 성능을 달기 위해 광범위한 하이퍼파라미터 튜닝이나 사전 정의된 스케줄러(예: 코사인 어닐링)를 필요로 합니다. 이러한 수동 튜닝에 대한 의존은 비용이 많이 들고 시간이 오래 걸리는 작업입니다.

저자들은 수동적인 학습률 튜닝을 요구하지 않는 SAM 변형을 설계하는 과제를 해결하고자 합니다. 특히, 최근 확률적 경사 하강법(SGD)에서 매우 효과적임이 입증된 확률적 폴리악 스텝 사이즈(stochastic Polyak step sizes)의 적용 가능성을 조사합니다.

방법론

본 논문은 고전적인 폴리악 스텝 사이즈로부터 유도된 USAM(및 확장된 SAM)을 위한 새로운 클래스의 적응형 스케줄러를 제안합니다. 핵심 방법론은 다음과 같습니다:

1. 적응형 스텝 사이즈 유도:

   • 결정론적 설정(Deterministic Setting): USAM 업데이트 규칙에서 시작하여, 최적점까지의 제곱 거리($\|x_{t+1} - x^*\|^2$)의 상한을 최소화함으로써 폐쇄형(closed-form) 스텝 사이즈를 유도합니다. 이를 통해 섭동된 지점에서의 손실 $f(e_t)$, 최적값 $f^*$, 섭동된 지점에서의 그래디언트, 그리고 샤프니스 반경 $\rho_t$에 의존하는 스텝 사이즈 $\gamma_t$를 얻습니다.
   • 확률적 설정(Stochastic Setting): 결정론적 유도 과정을 유한 합(finite-sum) 설정(미니 배치)으로 확장합니다. 제안된 **확률적 폴리악 스케줄러(Stochastic Polyak Scheduler)**는 전체 배치 양을 미니 배치 추정치($f_{S_t}, \nabla f_{S_t}$)로 대체하며, 안정성을 보장하기 위해 하한 $\ell^*_{S_t}$(비음수 손실의 경우 흔히 0)와 캡(cap) $\gamma_b$를 포함합니다.
   • 정규화(Normalization): 주요 이론적 전개는 USAM(비정규화)에 초점을 맞추고 있으나, 저자들은 동일한 로직을 정규화된 SAM 업데이트에도 적용하며, USAM의 이론적 비음수 보장이 정규화된 경우로 직접 전달되지 않기 때문에 비음수성을 보장하기 위한 안전장치(ReLU)를 유지합니다.

2. 이론적 프레임워크:

   • 분석은 표준적인 매끄러움($L$-smooth) 및 볼록성($\mu$-강볼록성) 조건을 가정합니다.
   • 저자들은 샤프니스 반경 $\rho_t$가 충분히 작을 때(구체적으로 $\rho_t \leq 1/L$), 유도된 스텝 사이즈가 자연스럽게 비음수가 되어 ReLU 안전장치가 불필요해짐을 입증합니다.
   • 이 프레임워크는 표준 그래디언트 방법과 샤프니스 인식 방법을 통합합니다: $\rho_t = 0$으로 설정하면 GD를 위한 고전적 폴리악 스텝 사이즈와 그 확률적 유사체(SPS$_{max}$)를 회복합니다.

주요 기여

1. 새로운 적응형 스케줄러: 본 논문은 SAM 스타일의 업데이트에 특화된 최초의 폴리악 유형 스텝 사이즈를 도입합니다. 이 스케줄러들은 완전히 적응형이며, 수동 학습률 튜닝을 필요로 하지 않고, 오직 현재 반복 단계에서 사용 가능한 양들(섭동된 지점의 손실 및 그래디언트)에만 의존합니다.
2. 결정론적 수렴 보장:
   • $\mu$-강볼록(strongly convex) 목적 함수에 대해, 저자들은 제곱 거리($\|x_t - x^*\|^2$)에 대한 **선형 수렴(linear convergence)**을 증명합니다. 수렴 속도는 GD의 폴리악 스텝 사이즈와 일치하며, 수렴 인자는 샤프니스 반경의 영향을 명시적으로 정량화합니다.
   • 볼록(convex) 목적 함수에 대해, 함수 값에 대한 서브리니어(sublinear) $O(1/T)$ 수렴 속도를 증명합니다.
   • 반경 스케줄이 비증가하는 형태($\rho_t \to 0$)일 때, 이 방법이 그래디언트 노름을 0으로 유도함을 보여줍니다.
3. 확률적 수렴 보장:
   • 확률적 설정에서, 이 방법은 최적점의 근방(neighborhood)으로 수렴합니다. 이 근방의 크기는 미니 배치 하한과 관련된 분산 측정값($\sigma^2$)에 의해 정량화됩니다.
   • 결정적으로, 이러한 보장은 유계 그래디언트(bounded gradients), 유계 분산(bounded variance) 또는 성장 조건(growth conditions)과 같이 기존의 적응형 SAM 분석에서 흔히 요구되는 추가 가정 없이도 성립합니다.
   • 보간 영역(interpolated regimes)(과잉 매개변수화된 모델에서 흔한 $\sigma^2 = 0$인 경우)에서는 근방이 사라지며, 최적점으로 선형 수렴합니다.
   • 또한, 이 분석은 확률적 환경에서의 상수 스텝 사이즈 USAM에 대한 수렴 보장을 제공하며, 이전 연구들보다 더 큰 스텝 사이즈를 허용함을 보여줍니다.
4. 경험적 검증:
   • 합성 실험: 리지 회귀(ridge regression) 문제에서, 제안된 스케줄러는 예측된 선형 수렴 경향을 보이며 튜닝된 상수 스텝 베이스라인 및 최근의 다른 적응형 SAM 변형들(예: AdaSAM, LightSAM)과 대등하거나 더 우수한 성능을 보입니다.
   • 딥러닝 벤치마크: ResNet 아키텍처를 사용한 CIFAR-10 및 CIFAR-100 실험 결과, 제안된 폴리악 스케줄러가 정교하게 튜닝된 상수 학습률 및 코사인 어닐링과 대등하거나 더 나은 테스트 정확도를 달리는 것을 입증했습니다. 특히, 제안된 스케줄러는 샤프니스 반경 $\rho$가 증가함에 따라 견고한 성능을 유지하는 반면, 코사인 어닐링은 상당한 정확도 저하를 겪습니다.

결과 및 의의

본 논문은 제안된 폴리악 스케줄러가 SAM 유형의 방법에서 학습률 튜닝 부담을 줄이기 위한 원칙적인 접근 방식을 제공한다고 주장합니다.

• 성능: 수치 실험은 적응형 스케줄러가 튜닝된 베이스라인과 대등하거나 더 나은 성능을 달성하면서도 하이퍼파라미터 탐색의 필요성을 실질적으로 줄여준다는 것을 보여줍니다.
• 견고성: 이 방법은 표준 SAM 구현에서 주의 깊은 튜닝이 필요한 파라미터인 샤프니스 반경 $\rho$의 선택에 대해 견고함을 보입니다.
• 이론적 진보: 이 연구는 폴리악 유형의 스텝 사이즈에 기반한 적응형 USAM에 대한 최초의 수렴 이론을 제공합니다. 이는 폴리악 스텝 사이즈의 SGD에 대한 이론적 이해와 샤프니스 인식 최적화에 대한 응용 사이의 간극을 메우며, 기존의 적응형 SAM 문헌보다 더 적은 가정을 가진 강력한 보장을 제공합니다.

저자들은 자신들의 이론이 볼록한 설정에 초점을 맞추고 있지만, 딥러닝 벤치마크에서의 실질적인 성공은 광범위한 적용 가능성을 시사한다고 결론짓습니다. 또한 비볼록(non-convex) 목적 함수(예: PL 조건)로 분석을 확장하고, 샤프니스 반경 $\rho_t$를 스케줄링하는 원칙적인 전략을 개발하는 것을 중요한 향후 연구 방향으로 제시합니다.