[논문 리뷰] Towards Efficient and Scalable Sharpness-Aware Minimization

Yong Liu, Siqi Mai외 3명|arXiv (Cornell University)|2022. 03. 05.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 날카러움에 민감한 최소화(SAM)의 계산 효율적인 변종인 LookSAM을 제안한다. LookSAM은 주기적으로 내부 기울기 상승 단계를 계산하며, 반복 간에 날카움 증진 방향을 재사용한다. 이 방법은 첫 번째 순서 최적화기 수준의 효율성으로 SAM 수준의 일반화 성능를 달성하여, 64,000 배치 크기로 비전 트랜스포머(Vision Transformers)를 한 시간 이내에 훈련시킬 수 있게 하여, ViT 훈련 분야에서 새로운 속도 기록을 수립한다.

ABSTRACT

Recently, Sharpness-Aware Minimization (SAM), which connects the geometry of the loss landscape and generalization, has demonstrated significant performance boosts on training large-scale models such as vision transformers. However, the update rule of SAM requires two sequential (non-parallelizable) gradient computations at each step, which can double the computational overhead. In this paper, we propose a novel algorithm LookSAM - that only periodically calculates the inner gradient ascent, to significantly reduce the additional training cost of SAM. The empirical results illustrate that LookSAM achieves similar accuracy gains to SAM while being tremendously faster - it enjoys comparable computational complexity with first-order optimizers such as SGD or Adam. To further evaluate the performance and scalability of LookSAM, we incorporate a layer-wise modification and perform experiments in the large-batch training scenario, which is more prone to converge to sharp local minima. We are the first to successfully scale up the batch size when training Vision Transformers (ViTs). With a 64k batch size, we are able to train ViTs from scratch in minutes while maintaining competitive performance.

연구 동기 및 목표

• SAM의 계산 오버헤드를 줄이기 위해, 각 단계에서 두 번의 순차적 기울기 계산으로 인해 훈련 시간이 두 배로 증가하는 문제를 해결하기 위해.
• 대규모 모델인 비전 트랜스포머를 위한 훈련 효율성을 크게 향상시키면서도 SAM의 일반화 이점을 유지하기 위해.
• 날카로운 최소값과 계층 간 비균일한 불안정성에 취약한 대배치 훈련 환경으로 SAM을 확장하기 위해.
• 매우 높은 배치 크기에서 안정적이고 효율적인 훈련을 가능하게 하는 LookSAM의 계층별 적응 방식인 Look-LayerSAM을 개발하기 위해.
• 기존 최고 수준의 설정보다 8배 빠른 속도로 비전 트랜스포머 훈련을 수행하여 새로운 훈련 속도 기록을 설정하기 위해.

제안 방법

• LookSAM은 매 단계에서 내부 기울기 상승을 다시 계산하는 대신, 이전 반복에서 얻은 날카움 증진 방향을 재사용한다.
• 이 방법은 SAM의 업데이트를 표준 SGD 방향과 평행하고 수직인 성분으로 분해하며, 수직 성분은 안정할 경우 재사용된다.
• 기울기 재사용 빈도를 제어하는 하이퍼파rameter $k$가 있으며, 이는 속도와 정확도 사이의 균형을 조절한다.
• Look-LayerSAM은 LookSAM의 가중치 편향에 계층별 스케일링을 적용하여 대배치 훈련에서의 안정성을 향상시킨다.
• 일반화와 효율성 사이의 트레이드오프를 미세 조정하기 위해 기울기 재사용 가중치 $\alpha$와 편향 강도 $\rho$를 사용한다.
• ImageNet-1k를 사용한 ViT-B-16에서의 실증적 검증을 통해 $k$, $\alpha$, $\rho$에 대한 분석 연구를 수행하였다.

Figure 1 : Accuracy of SAM-5, SAM and vanilla ViT on ImageNet-1k. SAM-5 indicates the method that calculating SAM gradients every 5 steps.

실험 결과

연구 질문

• RQ1 SAM의 계산 비용을 일반화 성능을 훼손하지 않고 줄일 수 있는가?
• RQ2 날카움 증진 방향을 주기적으로 재사용할 경우 모델 정확도와 훈련 속도에 어떤 영향을 미치는가?
• RQ3 LookSAM은 매우 큰 배치 크기(예: 64,000)로 확장되어도 경쟁력 있는 성능를 유지할 수 있는가?
• RQ4 계층별 편향 스케일링은 대배치 환경에서 훈련 안정성을 어떻게 향상시키는가?
• RQ5 최적의 속도-정확도 트레이드오프를 달성하기 위한 하이퍼파rameter 설정($k$, $\alpha$, $\rho$)은 무엇인가?

주요 결과

• LookSAM은 재사용 빈도 $k=5$일 때 ViT-B-16에서 검증 정확도 79.8%를 달성하여, 원본 SAM과 동일한 성능를 보였다.
• $k=5$일 때 LookSAM은 처리량을 초당 12,800개에서 19,051개로 증가시켜 훈련 속도를 크게 향상시켰다.
• $k=10$일 때 처리량은 초당 20,480개로 증가했으며, 정확도는 78.7%로 약간 감소하였다.
• Look-LayerSAM은 $\alpha=0.7$ 및 $\rho=1.0$일 때 16,384 배치 크기에서 78.4%의 정확도, 32,768 배치 크기에서 77.1%의 정확도를 달성하여 뛰어난 안정성을 입증하였다.
• 이 방법을 통해 ViT-B-16를 64,000 배치 크기로 1시간 이내에 처음부터 훈련시킬 수 있었으며, 정확도 75% 이상을 달성하였다.
• Look-LayerSAM은 [10]의 훈련 설정보다 약 8배 빠른 속도를 기록하여 ViT-B-16 훈련을 0.7시간 내에 완료하였으며, 새로운 속도 기록을 수립하였다.

Figure 2 : Difference of gradients between every 5 steps for $\boldsymbol{g_{s}}$ , $\boldsymbol{g_{h}}$ , and $\boldsymbol{g_{v}}$ (i.e., $||g_{s}^{t}-g_{s}^{t+k}||$ ) . $\boldsymbol{g_{v}}$ that leads to a smoother region changes much slower than $\boldsymbol{g_{s}}$ and $\boldsymbol{g_{h}}$ .