확률적 가중치 평균(SWA)은 딥러닝에서 모델 평균화 기법으로, 신경망의 가중치를 훈련 궤적의 여러 지점에 걸쳐 평균내어 일반화 성능을 향상시킨다. 최종 가중치만 유지하는 전통적인 훈련과 달리, SWA는 훈련 후반부에 정기적인 간격으로 가중치를 수집하고 이들의 산술 평균을 계산한다. 이 간단한 절차는 종종 손실 표면에서 더 평평한 해를 도출하며, 이는 더 나은 테스트 성능과 견고성과 관련이 있다. SWA는 Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, Andrew Gordon Wilson이 2018년 논문 "Averaging Weights Leads to Wider Optima and Better Generalization"에서 소개했다. 이 방법은 기본 모델 아키텍처나 손실 함수를 변경할 필요가 없어 기존 훈련 파이프라인에 쉽게 통합할 수 있다.
SWA의 핵심 아이디어는 확률적 경사 하강법(SGD) 또는 그 변형으로 훈련하는 동안 모델 매개변수가 낮은 손실 영역 주변에서 진동한다는 것이다. 이러한 매개변수를 평균화하면 진동을 완화하고 넓고 평평한 최솟값에 있는 해를 얻을 수 있다. 이러한 평평한 최솟값은 훈련 분포와 테스트 분포 간의 섭동 및 이동에 덜 민감하기 때문에 예리한 최솟값보다 일반화가 더 잘 된다고 가정된다. SWA는 이미지 분류, 객체 감지, 언어 모델링을 포함한 다양한 아키텍처와 작업에서 정확도와 보정(calibration)을 개선하는 것으로 나타났다.
메커니즘 및 알고리즘
SWA는 두 단계로 작동한다. 첫 번째 단계에서는 모델이 표준 학습률 스케줄(예: 코사인 어닐링 또는 단계 감쇠)로 훈련되어 최솟값 근처 영역에 도달한다. 두 번째 단계에서는 일정하거나 주기적인 학습률로 훈련을 계속하고, 각 에폭이 끝나거나 일정한 반복 횟수 후에 가중치를 수집한다. SWA 가중치는 실행 평균으로 업데이트된다: n_swa = n_swa + 1; w_swa = (w_swa * (n_swa - 1) + w) / n_swa. 여기서 w는 현재 가중치이고 w_swa는 평균화된 가중치이다. 최종 모델은 w_swa를 사용한다.
핵심 세부 사항은 평균화 후 배치 정규화 통계를 다시 계산해야 한다는 것이다. 배치 정규화 계층의 실행 평균과 분산은 가중치와 함께 평균화되지 않기 때문이다. 실제로 SWA 가중치를 계산한 후 훈련 데이터에 대한 순전파를 수행하여 배치 정규화 통계를 업데이트한다. 이 단계는 ResNet과 같이 배치 정규화를 사용하는 모델에 중요하다.
이론적 동기
SWA의 성공은 종종 손실 표면 기하학의 관점에서 설명된다. Izmailov 등의 연구와 후속 연구는 SGD가 낮은 손실 영역의 경계 근처 지점으로 수렴하는 경향이 있는 반면, 여러 지점을 평균화하면 해를 해당 영역의 중심으로 이동시킨다는 것을 보여주었다. 이 중심점은 일반적으로 더 평평한 분지에 있으며, 이는 더 나은 일반화와 관련이 있다. 평평한 최솟값과 일반화 사이의 연결은 PAC-Bayes 경계 및 손실 함수의 예리함과의 연결을 포함하여 광범위하게 연구되었다.
또 다른 관점은 SWA를 베이즈 추론과 관련시킨다. 특정 가정 하에서 SGD의 궤적은 사후 분포에서 샘플링하는 마르코프 체인으로 볼 수 있다. 이 궤적을 따라 평균화하는 것은 사후 평균을 근사하며, 이는 제곱 오차 손실에 대한 베이즈 최적 예측이다. 이 해석은 SWA를 확률적 경사 마르코프 체인 몬테카를로 방법과 연결하지만, SWA는 더 간단하고 노이즈 스케일의 세심한 조정이 필요하지 않다.
변형 및 확장
성능을 더욱 향상시키기 위해 여러 SWA 변형이 제안되었다. 주목할 만한 확장 중 하나는 2019년 Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, Andrew Gordon Wilson이 소개한 확률적 가중치 평균 가우시안(SWAG)이다. SWAG는 수집된 가중치의 1차 및 2차 모멘트를 사용하여 가우시안을 피팅함으로써 가중치의 사후 분포를 근사한다. 이를 통해 불확실성 추정과 보정 개선이 가능하며 베이즈 딥러닝에 사용할 수 있다.
또 다른 변형은 Fast Geometric Ensembling(FGE)으로, 주기적인 학습률 스케줄을 사용하여 손실 표면의 서로 다른 모드에서 가중치를 수집한다. SWA는 FGE와 결합하여 더 나은 성능을 얻을 수 있다. 또한 SWA는 학습률 스케줄, 데이터 증강, 가지치기와 같은 기술과 통합되어 견고성과 효율성을 향상시킨다.
응용 분야
SWA는 다양한 분야에서 성공적으로 적용되었다. 컴퓨터 비전에서는 CIFAR-10, CIFAR-100, ImageNet과 같은 데이터셋에서 이미지 분류기의 정확도를 향상시킨다. 예를 들어, SWA는 아키텍처 변경 없이 CIFAR-100에서 PreAct ResNet-164의 테스트 정확도를 81.15%에서 82.90%로 끌어올리는 것으로 나타났다. 객체 감지에서는 훈련을 안정화하고 평균 평균 정밀도(mean average precision)를 개선하는 데 도움이 된다.
자연어 처리에서는 트랜스포머와 대규모 언어 모델을 더 효과적으로 훈련하는 데 SWA가 사용되었다. 언어 모델링 작업에서 과적합을 줄이고 혼란도(perplexity)를 개선하는 것으로 나타났다. SWA는 또한 GAN 및 확산 모델과 같은 생성 모델에 이점을 제공하여 더 부드러운 가중치 업데이트를 통해 더 안정적인 훈련과 더 높은 품질의 샘플을 생성할 수 있다.
표준 지도 학습 외에도 SWA는 준지도 학습, 전이 학습, 도메인 적응에 적용되었다. 그 단순성과 일반성 덕분에 머신러닝 실무자의 도구 상자에서 귀중한 도구가 되었다.
다른 기법과의 비교
SWA는 종종 다른 정규화 및 앙상블 방법과 비교된다. 여러 개의 독립적으로 훈련된 모델의 예측을 평균화하는 전통적인 앙상블과 달리, SWA는 가중치를 평균화하므로 계산 비용이 더 저렴하고 단일 훈련 실행만 필요하다. 이는 모델 저장 또는 추론 비용이 문제가 될 때 특히 매력적이다.
드롭아웃 또는 배치 정규화와 같은 기법과 비교할 때, SWA는 훈련 역학을 수정하지 않고 가중치를 사후 처리한다. 이러한 방법과 함께 사용할 수 있다. SWA는 또한 매개변수별 학습률을 유지하는 Adam 및 기타 적응형 최적화 도구와 다르다. SWA는 일반적으로 SGD 또는 Adam 궤적 위에 적용된다.
SWA의 한계 중 하나는 가중치 수집을 시작할 시점을 결정하는 스케줄이 필요하다는 것이다. 수집이 너무 일찍 시작되면 평균이 최적이 아닌 가중치에 의해 지배될 수 있고, 너무 늦게 시작하면 이점이 줄어든다. 그러나 실제로는 학습률이 낮은 값으로 감쇠된 후 수집을 시작하는 간단한 경험 법칙이 있다.
실용적 고려 사항
SWA 구현은 간단하다. PyTorch 및 TensorFlow와 같은 대부분의 딥러닝 프레임워크는 SWA를 위한 유틸리티 또는 예제를 제공한다. 주요 단계는 다음과 같다: (1) 표준 스케줄로 모델을 훈련하고, (2) 일정하거나 주기적인 학습률로 훈련을 계속하고, (3) 간격으로 가중치를 수집하고, (4) 실행 평균을 계산하고, (5) 배치 정규화 통계를 다시 계산한다.
SWA는 하이퍼파라미터 선택에 견고한 것으로 나타났다. 두 번째 단계의 학습률은 일반적으로 초기 스케줄의 최솟값 또는 작은 상수로 설정된다. 가중치 수집 빈도는 매 에폭 또는 수백 회 반복마다일 수 있다. 실제로 5~20 에폭 동안 가중치를 수집하는 것이 종종 충분하다.
영향 및 평가
SWA의 도입은 딥러닝 커뮤니티에 상당한 영향을 미쳤다. 추론 중 추가 계산 비용 없이 일반화를 개선하는 간단하면서도 효과적인 방법을 제공했다. 이 논문은 널리 인용되었으며 가중치 평균화와 평평한 최솟값에 대한 많은 후속 연구에 영감을 주었다. SWA는 이제 특히 학술 연구와 대회에서 많은 훈련 파이프라인의 표준 기법이 되었다.
SWA는 또한 생성 모델 및 대규모 언어 모델 훈련에서 일반적으로 사용되는 지수 이동 평균(EMA)과 같은 다른 평균화 방법의 개발에도 영향을 미쳤다. EMA는 최근 가중치에 더 많은 가중치를 부여하는 변형으로, 단순성과 효과성 때문에 실제로 자주 사용된다. SWA와 EMA는 상호 보완적이며 일부 프레임워크에서는 둘 사이를 전환할 수 있다.
한계 및 향후 방향
이점에도 불구하고 SWA에는 한계가 있다. 손실 표면이 비교적 매끄럽고 가중치 평균화가 의미 있다고 가정하지만, 이는 모든 아키텍처나 작업에 해당하지 않을 수 있다. 예를 들어, 배치 정규화가 있는 모델은 세심한 처리가 필요하며, 복잡한 매개변수 공간(예: 순환 신경망)을 가진 모델은 이점이 적을 수 있다. 또한 SWAG와 같은 확장이 없으면 SWA는 불확실성 추정을 제공하지 않는다.
향후 연구 방향으로는 적응형 평균화 스케줄 개발, SWA와 베이즈 방법의 결합, 강화 학습 및 연합 학습과 같은 새로운 분야에 SWA 적용 등이 있다. 딥러닝이 계속 진화함에 따라 SWA와 같은 가중치 평균화 기법은 모델 성능과 신뢰성을 향상시키는 기본 도구로 남아 있다.