Sharpness-Aware Minimization for Efficiently Improving Generalization
일반화를 효율적으로 개선하기 위한 샤프니스 인식 최소화
Pierre Foret, Ariel Kleiner, Hossein Mobahi, Behnam Neyshabur (2020)
손실 값뿐만 아니라 손실 지형의 '평탄도(flatness)'까지 동시에 최적화하는 Sharpness-Aware Minimization(SAM) 알고리즘을 제안하여, 날카로운 최솟값 대신 넓고 평탄한 최솟값으로 수렴하게 함으로써 일반화 성능을 크게 향상시켰다.
배경
심층 학습의 일반화를 이해하려는 연구에서, 손실 지형(loss landscape)의 기하학적 특성이 중요한 역할을 한다는 것이 알려져 있었다. 특히 날카로운(sharp) 최솟값보다 평탄한(flat) 최솟값에 수렴한 모델이 일반화를 더 잘한다는 이론적, 경험적 증거가 축적되어 있었다. 그러나 기존 옵티마이저(SGD, Adam)는 손실 값 자체만을 최소화하며, 손실 지형의 평탄도를 명시적으로 고려하지 않았다. PAC-Bayes 이론에 기반한 일반화 경계(generalization bound)가 손실의 샤프니스와 직접 연관되어 있음이 이론적으로 알려져 있었다.
핵심 아이디어
SAM은 현재 파라미터 w 근처의 최악의 섭동(worst-case perturbation)에서의 손실을 최소화한다. 목적함수는 min_w max_{||ε||≤ρ} L(w + ε)로, ε-근방에서의 최대 손실을 최소화하는 미니맥스 문제이다. 이를 해석하면, 단순히 현재 지점의 손실을 줄이는 것이 아니라, 근처 어디에서든 손실이 낮게 유지되는 '평탄한' 영역을 찾도록 유도한다. 내부 최대화 문제의 근사해로 ε̂ = ρ · ∇L(w) / ||∇L(w)||를 사용하며, 최종 그래디언트는 ∇L(w + ε̂)로 한 번의 추가 순전파-역전파만 필요하다.
방법론
각 학습 스텝에서 두 번의 순전파-역전파를 수행한다. 첫째, 현재 파라미터 w에서 그래디언트 ∇L(w)를 계산하고, 이 방향으로 ρ 크기의 섭동 ε̂을 가한다. 둘째, 섭동된 파라미터 w + ε̂에서 다시 그래디언트 ∇L(w + ε̂)를 계산한다. 이 두 번째 그래디언트로 파라미터를 업데이트한다. 기존 옵티마이저(SGD, Adam)와 결합하여 사용할 수 있으며, ρ는 섭동 반경을 제어하는 유일한 추가 하이퍼파라미터이다.
주요 결과
CIFAR-10에서 ResNet, WideResNet, PyramidNet 등 다양한 아키텍처에서 기존 최고 성능을 경신했으며, WRN-28-10에서 기존 96.1%를 96.9%로 향상시켰다. CIFAR-100에서도 1% 이상의 일반화 성능 향상을 보였다. ImageNet에서 ResNet-50의 top-1 정확도를 76.3%에서 77.3%로 향상시켰으며, EfficientNet-B7의 경우 84.7%에서 85.0%로 개선했다. 레이블 노이즈가 있는 환경에서 특히 큰 성능 차이를 보여, 노이즈에 대한 강건성도 향상되었다.
임팩트
SAM은 딥러닝 최적화에서 '손실 지형의 기하학'을 명시적으로 활용하는 실용적 방법론을 확립했다. 두 배의 계산 비용이라는 단점에도 불구하고, ViT, BERT 미세조정 등 다양한 설정에서 채택되었다. ASAM(적응적 SAM), LookSAM(효율적 SAM), GSAM 등 효율성과 성능을 개선한 후속 변형이 다수 제안되었다. SAM의 성공은 최적화와 일반화의 관계에 대한 이론적 연구도 촉진하여, 평탄 최솟값 이론의 실증적 근거를 강화했다.