[논문 리뷰] Stable and low-precision training for large-scale vision-language models

Mitchell Wortsman, Tim Dettmers외 4명|arXiv (Cornell University)|2023. 04. 25.

Multimodal Machine Learning Applications인용 수 5

한 줄 요약

이 논문은 대규모 비전-언어 모델의 안정적이고 효율적인 8비트(int8 및 float8) 훈련을 위한 새로운 선형 레이어인 SwitchBack을 소개한다. CLIP ViT-Huge에서 bfloat16 훈련 대비 13-25%의 속도 향상을 달성하면서도 정확도는 거의 동일하다. 또한 손실 급등 현상을 제거하기 위해 업데이트 클리핑을 적용하는 StableAdamW라는 AdamW-AdaFactor 하이브리드를 제안하여 대규모 훈련의 안정성을 크게 향상시켰다.

ABSTRACT

We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models. 1) For acceleration, we introduce SwitchBack, a linear layer for int8 quantized training which provides a speed-up of 13-25% while matching the performance of bfloat16 training within 0.1 percentage points for the 1B parameter CLIP ViT-Huge -- the largest int8 training to date. Our main focus is int8 as GPU support for float8 is rare, though we also analyze float8 training through simulation. While SwitchBack proves effective for float8, we show that standard techniques are also successful if the network is trained and initialized so that large feature magnitudes are discouraged, which we accomplish via layer-scale initialized with zeros. 2) For stability, we analyze loss spikes and find they consistently occur 1-8 iterations after the squared gradients become under-estimated by their AdamW second moment estimator. As a result, we recommend an AdamW-Adafactor hybrid which avoids loss spikes when training a CLIP ViT-Huge model and outperforms gradient clipping at the scales we test.

연구 동기 및 목표

• 정확도를 훼손하지 않으면서 저정밀도(int8 및 float8) 계산을 통해 대규모 비전-언어 모델 훈련을 가속화하는 것.
• 모델 성능 저하를 유발하는 손실 급등 현상을 식별하고 완화하여 훈련 안정성을 향상시키는 것.
• 1B 파라미터를 가진 CLIP ViT-Huge 모델에 대해 처음으로 int8 훈련을 성공적으로 수행하여, bfloat16 성능과 0.1% 이내로 일치시키는 것.
• 미래의 저정밀도 훈련 연구를 위한 실용적이고 오픈소스인 도구(예: Triton 커널)를 개발하는 것.
• 적절한 초기화와 적응형 최적화 기법이 존재할 경우, 하드웨어 네이티브 지원 없이도 안정적인 float8 훈련을 가능하게 할 수 있는지 입증하는 것.

제안 방법

• SwitchBack은 고차원의 어텐션 레이어에서의 양자화 노이즈를 줄이기 위해, 가중치 기울기 계산에는 16비트 행렬 곱셈을 사용하고, 순방향 전파 및 입력 기울기 계산은 int8에서 수행한다.
• 이 방법은 하이브리드 정밀도 전략을 적용한다: 추론 및 입력 기울기 계산은 8비트로, 기울기 누적은 16비트로 수행하여 표준 양자화 대비 정확도를 향상시킨다.
• float8 훈련을 위해 SwitchBack은 정확한 fp8 에뮬레이션을 사용하여 시뮬레이션되며, 이는 스케일링된 범위에서 발산하는 텐서 수준 양자화 기반 베이스라인보다 우수한 성능을 보임을 입증한다.
• StableAdamW는 AdamW와 AdaFactor의 업데이트 클리핑 기법을 융합한 것으로, 제곱 기울기의 비율과 그 이동 평균 추정치를 모니터링하여 동적으로 학습률을 조정한다.
• 손실 스케일러는 텐서 수준에서 Inf/NaN 여부를 확인하고 고정함으로써, 고립된 레이어의 불안정성으로 인한 전역 훈련 실패를 방지한다.
• 영역 스케일 초기화를 0으로 설정하여 큰 특징 크기를 억제함으로써, 표준 양자화 기법을 사용한 안정적인 float8 훈련을 가능하게 한다.

실험 결과

연구 질문

• RQ1 SwitchBack과 같은 새로운 선형 레이어 설계가 대규모 비전-언어 모델의 int8 훈련에서 높은 정확도를 달성하면서도 상당한 속도 향상을 이룰 수 있는가?
• RQ2 대규모 CLIP 훈련 중 손실 급등 현상이 발생하는 이유는 무엇이며, 이를 예측하거나 방지할 수 있는가?
• RQ3 네이티브 하드웨어 지원 없이도 float8 훈련을 안정화시킬 수 있는가? 어떤 초기화 또는 최적화 기법이 이를 가능하게 하는가?
• RQ4 AdamW와 AdaFactor의 업데이트 클리핑을 융합한 StableAdamW는 기울기 클리핑 대비 안정성과 최종 성능에서 어떻게 비교되는가?
• RQ5 큰 활성화 및 기울기 값은 저정밀도 훈련 중 불안정성을 유발하는 데 어떤 역할을 하는가?

주요 결과

• SwitchBack은 bfloat16 훈련 대비 CLIP ViT-Huge 훈련에서 종단 간 속도 향상을 13-25% 달성하였으며, 제로샷 ImageNet 정확도는 0.1%포인트 이내로 일치한다.
• SwitchBack은 1B 파라미터를 가진 CLIP ViT-Huge 모델에서 현재까지 보고된 바 가장 큰 int8 훈련을 가능하게 하였으며, LLM.int8()보다 정확도에서 뛰어나다.
• 표준 텐서 수준 양자화를 사용한 float8 훈련은 420M 파라미터를 초과하면 발산하지만, SwitchBack에 레이어 스케일 초기화(0으로 설정)를 조합하면 안정적인 훈련이 가능하다.
• 업데이트 클리핑 기능을 갖춘 AdamW-AdaFactor 하이브리드인 StableAdamW는 손실 급등 현상을 완전히 제거하고, 대규모 훈련에서 기울기 클리핑 대비 제로샷 ImageNet 정확도에서 뛰어난 성능을 보였다.
• 손실 급등 현상은 항상 패치 임bedding 레이어에서 제곱 기울기와 그 이동 평균의 비율 감소로 이어지며, 이는 제2모멘트 추정기의 과소 추정을 시사한다.
• Inf/NaN 여부를 텐서 수준에서 확인하고 고정하는 고정된 텐서 스케일 손실 스케일러를 사용하면 전역 훈련 실패를 방지하고, ViT-Huge 규모에서 안정적인 fp16 혼합 정밀도 훈련을 가능하게 한다.