Sharpness-Aware Minimization(SAM)은 신경망 훈련을 위한 최적화 방법으로, 손실 경관에서 평평한 최소점을 명시적으로 탐색한다. 2021년 Pierre Foret와 동료들에 의해 소개된 SAM은 표준 경사 하강법 업데이트를 수정하여 손실이 급격히 변하는 영역을 페널티화함으로써, 보이지 않는 데이터에 대한 모델 일반화를 개선한다. 확률적 경사 하강법(SGD)이나 Adam과 같은 기존 최적화 도구가 훈련 손실만 최소화하는 반면, SAM은 각 파라미터 점의 이웃을 고려하여 섭동에 강건한 해를 찾는다.
SAM의 핵심 아이디어는 신경망이 훈련 데이터를 동일하게 잘 맞추는 많은 지역 최소점을 가지지만, 그중 일부는 다른 것보다 더 잘 일반화된다는 관찰에서 비롯된다. 평평한 최소점 - 손실 표면이 파라미터 변화에 상대적으로 둔감한 영역 - 은 경험적으로 더 나은 일반화와 연관된다. SAM은 이를 min-max 최적화 문제를 풀어 공식화한다: 각 단계에서 먼저 현재 파라미터 주변의 작은 공 내에서 손실을 최대화하는 섭동을 계산한 다음, 그 섭동된 점에서의 기울기를 사용하여 파라미터를 업데이트한다. 이는 최적화 도구가 단일 점이 아닌 이웃에서 균일하게 낮은 인력 분지(basin of attraction)로 이동하도록 장려한다.
SAM은 이미지 분류, 언어 모델링, 강화 학습을 포함한 다양한 아키텍처와 작업에서 정확도를 향상시키는 것으로 나타났다. 특히 데이터 증강 및 가중치 감쇠와 같은 다른 정규화 기법과 결합할 때 효과적이다. 이 방법은 최소한의 계산 오버헤드 - 대략 순전파-역전파 비용의 두 배 - 를 추가하지만 종종 상당한 정확도 향상을 가져온다. 소개 이후 적응형 SAM과 lookahead SAM을 포함한 여러 변형이 제안되어 계산 부담을 더 줄였다.
동기: 날카로운 최소점 대 평평한 최소점
딥러닝의 경험적 성공은 종종 대규모 모델이 고차원 손실 경관에서 좋은 해를 찾는 능력에 기인한다. 그러나 모든 최소점이 동일하지는 않다. 연구에 따르면 SGD는 평평한 최소점으로 수렴하는 경향이 있으며, 이는 작은 파라미터 섭동에 더 강건하고 종종 더 나은 일반화에 해당한다. 반대로, 날카로운 최소점은 최소점에서 멀어질수록 손실이 급격히 증가하는 좁은 분지로, 공격적인 최적화에서 종종 발생하며 과적합으로 이어질 수 있다.
평평한 최소점에 대한 초기 이론적 연구는 1990년대 신경망 일반화 연구로 거슬러 올라가지만, 엄격한 프레임워크는 부족했다. SAM은 명시적 정규화 항 없이 날카로움을 직접 페널티화하는 명확한 최적화 목표를 제공한다. 엡실론 반경 내 최악의 섭동을 계산함으로써, SAM은 손실 경관이 해당 이웃에서 균일하게 낮도록 효과적으로 장려하여 평평함을 촉진한다.
공식 정의 및 알고리즘
\(w\)를 모델 파라미터, \(L(w)\)를 훈련 손실이라고 하자. SAM은 다음 문제를 푼다:
\[ \min_w \max_{\|\epsilon\| \le \rho} L(w + \epsilon) \]
여기서 \(\rho\)는 섭동 반경을 제어하는 하이퍼파라미터이다. 실제로 내부 최대화는 경사 상승의 한 단계로 근사된다: \(\hat{\epsilon} = \rho \frac{\nabla L(w)}{\|\nabla L(w)\|}\). 그러면 업데이트 규칙은 다음과 같다:
\[ w_{t+1} = w_t - \eta \nabla L(w_t + \hat{\epsilon}_t) \]
이는 단계당 두 번의 기울기 계산이 필요하다: 하나는 섭동을 찾기 위한 것이고 다른 하나는 파라미터를 업데이트하기 위한 것이다. 이 방법은 기본 최적화 도구에 구애받지 않으며, SGD, Adam 또는 모든 기울기 기반 알고리즘과 함께 사용할 수 있다. 실제로 SAM은 종종 가중치 감쇠 및 모멘텀과 함께 사용되어 추가 개선을 얻는다.
하이퍼파라미터 및 튜닝
주요 하이퍼파라미터는 섭동 반경 \(\rho\)로, 고려되는 이웃의 크기를 제어한다. 더 큰 \(\rho\)는 더 평평한 최소점을 장려하지만 수렴을 늦출 수 있다. 정규화된 데이터셋의 경우 일반적인 값은 0.01에서 0.1 사이이다. 학습률과 가중치 감쇠는 \(\rho\)와 함께 튜닝해야 한다. 일부 연구는 SAM이 더 큰 배치 크기와 학습률 스케줄의 이점을 얻는다고 제안한다. 적응형 변형인 Adaptive SAM(ASAM)은 파라미터의 규모에 따라 반경을 조정하여 안정성을 개선한다.
변형 및 확장
원래 논문 이후 여러 개선이 제안되었다. Adaptive SAM(ASAM)은 파라미터 크기로 섭동을 정규화하여 스케일 불변성을 달성하고, CIFAR-10 및 ImageNet과 같은 작업에서 성능을 향상시킨다. Lookahead SAM은 Lookahead 최적화 도구를 SAM과 결합하여 분산을 줄인다. 다른 변형으로는 확률적 가중치 평균화(SWA)를 사용한 Sharpness-Aware Minimization과 라벨 스무딩을 사용한 SAM이 있다. 이러한 방법은 계산 비용을 줄이거나 일반화를 더욱 향상시키는 것을 목표로 한다.
다양한 분야에서의 응용
SAM은 광범위한 Machine learning 작업에 성공적으로 적용되었다. 컴퓨터 비전에서는 CIFAR-10 및 ImageNet과 같은 데이터셋에서 Residual Network (ResNet) 아키텍처의 이미지 분류 정확도를 향상시킨다. 자연어 처리에서는 데이터가 부족할 때 특히 Large language model과 같은 트랜스포머를 미세 조정하는 데 도움이 된다. 또한 의료 영상, 약물 발견, 강화 학습에도 사용되었다. Generative AI에서는 안정적인 이미지 생성을 위한 모델 훈련에 SAM이 도움을 준다. 그 다양성은 학술 연구와 산업 모두에서 귀중한 도구로 만들며, PyTorch 및 TensorFlow와 같은 인기 있는 라이브러리에 구현되어 있다.
다른 최적화 기법과의 관계
SAM은 Adam (Optimizer)나 SGD의 대안이 아니라 손실 경관을 수정하는 추가 기능이다. 훈련을 안정화하는 기울기 클리핑과 같은 방법과 자주 비교되지만, SAM은 수렴뿐만 아니라 일반화를 다룬다. Batch Normalization과도 유사점을 공유하는데, 둘 다 손실 표면의 기하학에 암시적으로 영향을 미치지만 메커니즘은 다르다. SAM은 Data Augmentation 및 Dropout과 결합하여 모델을 더 정규화할 수 있다. 평평한 최소점의 개념은 BAIR (Berkeley AI Research) 커뮤니티의 손실 경관 시각화 및 엔트로피 기반 정규화 연구와도 관련이 있다.
이론적 통찰
연구는 평평한 최소점이 더 잘 일반화되는 이유에 대한 이론적 근거를 제공했다. 과도하게 파라미터화된 모델의 경우, 함수 클래스의 복잡성은 종종 최소점의 날카로움으로 측정된다. 평평한 최소점은 일반적으로 과적합하지 않는 저복잡성 해에 해당한다. SAM의 min-max 공식은 모델을 최악의 섭동에 강건하게 만드는 적대적 훈련의 한 형태로 해석될 수 있다. 일부 연구는 SAM을 베이지안 추론과 연결하여 파라미터에 대한 사후 분포를 근사한다고 제안한다. 그러나 완전한 이론적 이해는 여전히 열려 있으며, 활발한 연구가 계속되고 있다.
계산 비용 및 트레이드오프
SAM의 주요 단점은 업데이트당 두 번의 순전파-역전파 패스가 필요하여 계산 비용이 두 배가 된다는 점이다. 이는 매우 큰 모델이나 실시간 응용 프로그램에서는 금지적일 수 있다. 그러나 정확도 향상은 종종 오버헤드를 정당화하며, 특히 훈련 시간보다 최종 모델 품질이 더 중요할 때 그렇다. 일부 변형은 섭동을 덜 자주 업데이트하거나(예: 몇 단계마다) 2차 근사를 사용하여 비용을 분할하려고 시도한다. 표준 변형과 적응형 변형의 선택은 특정 사용 사례에 따라 달라진다.
산업에서의 채택
주요 AI 연구 기관은 다양한 응용 프로그램에 SAM을 채택했다. Google DeepMind는 강화 학습 및 대규모 훈련에서 SAM을 탐구했다. OpenAI는 모델 강건성 개선 맥락에서 평평한 최소점을 언급했다. Apple과 Samsung Electronics는 온디바이스 머신러닝에서 효율적인 훈련 및 미세 조정을 위해 SAM을 적용했다. Amazon Web Services와 Google Cloud는 클라우드 기반 훈련을 위한 SAM 구현을 포함하는 라이브러리를 제공한다. 이 방법은 인기 있는 딥러닝 프레임워크에도 통합되어 있으며 Kaggle 대회에서 널리 사용된다.
다른 정규화 방법과의 비교
SAM은 가중치 감쇠, 드롭아웃, Batch Normalization과 같은 전통적인 정규화 도구와 상호 보완적이다. 이러한 방법은 개별 파라미터나 활성화에 작용하는 반면, SAM은 전체 손실 경관에 작용한다. 이는 손실을 증가시키도록 설계된 섭동을 사용하는 적대적 훈련의 한 형태로 볼 수 있다. 이는 기울기의 크기만 제한하는 기울기 클리핑과 구별된다. SAM은 훈련 샘플의 순서에 초점을 맞춘 Curriculum Learning과도 다르다. 계산 비용 측면에서 SAM은 단순 정규화보다 비싸지만 종종 정확도에서 더 나은 수익을 제공한다.
한계 및 고려 사항
이점에도 불구하고 SAM에는 한계가 있다. 단계당 추가 순전파-역전파 패스는 훈련 시간을 대략 두 배로 늘리며, 이는 대규모 모델이나 데이터셋에서 금지적일 수 있다. \(\rho\)의 선택은 항상 간단하지 않으며 작업별로 신중한 튜닝이 필요할 수 있다. 매우 노이즈가 많은 라벨과 같은 일부 설정에서는 SAM이 개선을 제공하지 않을 수 있으며 오히려 해로울 수 있다. 또한 모델이 이미 잘 정규화된 경우 SAM의 효과는 덜 두드러진다. 연구자들은 섭동 계산에 데이터 하위 집합을 사용하거나 두 기울기 단계를 분할하는 등 비용을 줄이는 방법을 탐구했다.
딥러닝 실무에 미치는 영향
SAM은 딥러닝 실무자의 도구 상자에서 표준 도구가 되었다. 도메인 이동이 우려될 때 사전 훈련된 모델을 미세 조정하는 기본 선택인 경우가 많다. 비전 벤치마크의 많은 최첨단 결과는 훈련 파이프라인의 구성 요소로 SAM에 의존한다. 산업에서는 Apple, Intel, Nokia Bell Labs와 같은 회사가 다양한 응용 프로그램에 SAM을 채택했으며, AWS Trainium 및 Google Cloud와 같은 클라우드 AI 플랫폼에 통합되어 있다. 이 방법의 단순성과 효과성은 연구 및 생산 환경에서 인기 있는 선택이 되게 했다.
한계 및 과제
성공에도 불구하고 SAM에는 한계가 있다. 두 배의 훈련 시간은 Natural language processing에 사용되는 Transformer (architecture)와 같은 매우 큰 모델에서 금지적일 수 있다. 섭동 반경은 신중하게 튜닝해야 하며, 최적 값은 데이터셋과 아키텍처에 따라 다를 수 있다. 일부 설정에서는 특히 기본 모델이 이미 잘 정규화된 경우 SAM의 개선이 미미하다. 또한 SAM의 이론적 보장은 여전히 불완전하며, 시계열이나 그래프 데이터와 같은 다른 양식에서의 성능은 덜 탐구되었다. 진행 중인 연구는 더 효율적인 근사 및 적응형 체계를 제안하여 이러한 문제를 해결한다.
미래 방향
SAM의 성공은 기하학적 관점에서 일반화를 이해하는 데 대한 관심을 불러일으켰다. 연구자들은 정보 이론, PAC-Bayes, 강건성과의 연결을 탐구하고 있다. Amazon Web Services 또는 Microsoft Azure 클라우드 환경과 같이 계산 리소스가 제한된 온디바이스 훈련을 위한 효율적인 SAM 변형이 개발되고 있다. SAM과 Transformer (architecture) 아키텍처 및 Deep learning 프레임워크의 통합은 활발한 연구 영역이다. 모델이 커질수록 추가 데이터 없이 일반화를 개선하는 SAM의 능력은 점점 더 가치 있게 된다.
결론
Sharpness-Aware Minimization은 좋은 일반화의 핵심인 평평한 최소점을 명시적으로 목표로 하여 신경망 최적화에서 중요한 진전을 나타낸다. 최악의 섭동에 대한 손실 최소화라는 단순하면서도 강력한 아이디어는 여러 분야에서 채택되었으며 정확도와 강건성 향상에 효과적임이 입증되었다. 계산 비용에도 불구하고 SAM과 그 변형은 이제 실무자의 도구 상자에서 표준 도구가 되었으며, 기존 최적화 도구로 달성할 수 있는 것 이상으로 모델 성능을 끌어올리는 신뢰할 수 있는 방법을 제공한다. 딥러닝이 계속 진화함에 따라 SAM의 추가 개선은 더 큰 이득을 가져올 수 있다.