[논문 리뷰] Mesa: A Memory-saving Training Framework for Transformers

Zizheng Pan, Peng Chen외 4명|arXiv (Cornell University)|2021. 11. 22.

Advanced Neural Network Applications참고 문헌 25인용 수 4

한 줄 요약

Mesa는 역전파 중에 활성화를 저정밀도로 양자화함으로써 GPU 메모리 사용량을 최대 50% 감소시키는 메모리 효율적인 훈련 프레임워크이다. 이는 정확한 전방전파를 유지하면서도, 헤드별 양자화 전략을 도입하여 이질적인 어텐션 분포를 처리하고, 런닝 추정치를 통해 양자화 파rameter를 학습함으로써 성능 손실 없이 더 큰 배치 크기나 모델 크기를 허용한다.

ABSTRACT

There has been an explosion of interest in designing high-performance Transformers. While Transformers have delivered significant performance improvements, training such networks is extremely memory intensive owing to storing all intermediate activations that are needed for gradient computation during backpropagation, especially for long sequences. To this end, we present Mesa, a memory-saving training framework for Transformers. Specifically, Mesa uses exact activations during forward pass while storing a low-precision version of activations to reduce memory consumption during training. The low-precision activations are then dequantized during back-propagation to compute gradients. Besides, to address the heterogeneous activation distributions in the multi-head self-attention layers, we propose a head-wise activation quantization strategy, which quantizes activations based on the statistics of each head to minimize the approximation error. To further boost training efficiency, we learn quantization parameters by running estimates. More importantly, by re-investing the saved memory in employing a larger batch size or scaling up model size, we may further improve the performance under constrained computational resources. Extensive experiments on ImageNet, CIFAR-100 and ADE20K demonstrate that Mesa can achieve flexible memory-savings (up to 50%) during training while achieving comparable or even better performance. Code is available at https://github.com/ziplab/Mesa.

연구 동기 및 목표

• 대규모 비전 트랜스포머 훈련 시 높은 메모리 소비 문제, 특히 역전파 중에 전체 정밀도 활성화를 저장함으로써 발생하는 문제를 해결하기 위해.
• 제한된 GPU 자원을 가진 연구자들이 성능 손실 없이 활성화 메모리 프로파일을 줄임으로써 메모리 효율적인 훈련을 가능하게 하기 위해.
• 행렬 곱셈, 소프트맥스, GELU, 레이어노멀라이제이션을 포함한 트랜스포머의 핵심 연산들과 호환 가능한 일반적인 활성화 양자화 프레임워크를 개발하기 위해.
• 다중 헤드 자기주의 어텐션 헤드 간 이질적인 활성화 분포로 인해 표준 양자화 전략이 성능에 영향을 미치는 문제를 해결하기 위해.
• 기존 기법들인 AMP, 체크포인팅, 기울기 양자화와의 상호작용이 가능하도록, 독립적인 메모리 절감 기법을 설계하기 위해.

제안 방법

• Mesa는 정확한 전방전파 활성화를 사용하지만, 훈련 중 메모리 사용량을 줄이기 위해 저정밀도(8비트) 양자화된 복제본을 저장한다.
• 역전파 동안 저정밀도 활성화는 기울기를 계산하기 위해 복원(디퀀타이징)되며, 이로 인해 훈련 정확도가 유지된다.
• 헤드별 활성화 양자화 전략을 제안하여 각 자기주의 어텐션 헤드가 고유의 양자화 파rameter(클리핑 범위 α 및 오프셋 β)를 가지도록 하여 근사 오차를 최소화한다.
• 양자화 파rameter(α 및 β)는 훈련 중 런닝 추정치를 통해 학습되며, 샘플별 통계나 기울기 기반 최적화를 피함으로써 오버헤드를 줄인다.
• 이 프레임워크는 트랜스포머의 주요 연산들인 행렬 곱셈, 소프트맥스, GELU, 레이어노멀라이제이션을 모두 지원하여 종단 간 메모리 절감을 가능하게 한다.
• 절감된 메모리는 더 큰 배치 크기나 모델 크기로 재투자되어 제한된 하드웨어 환경에서 성능 향상에 기여한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머를 위한 메모리 절약형 훈련 프레임워크를 설계할 수 있는가? 이는 성능 저하 없이 활성화 메모리 소비를 줄이는 데 목적이 있다.
• RQ2 헤드별 양자화는 다중 헤드 어텐션 레이어에서 이질적인 활성화 분포를 가진 경우에 비해 균일 양자화나 레이어별 양자화보다 정확도를 어떻게 향상시키는가?
• RQ3 샘플별 통계나 기울기 기반 학습보다 낮은 계산 비용으로도 유사하거나 더 나은 성능을 내는 양자화 파ram터에 대한 런닝 추정치는 가능한가?
• RQ4 실제로 Mesa에서 절감된 메모리는 얼마나 많은 모델 크기나 배치 크기 확장을 가능하게 하는가?
• RQ5 체크포인팅, 기울기 누적, AMP와 같은 기존 메모리 절약 기법과 비교할 때 Mesa는 메모리 감소와 훈련 효율성 측면에서 어떤가?

주요 결과

• Mesa는 Swin-T, Swin-B, DeiT와 같은 비전 트랜스포머에서 훈련 중 메모리 소비를 최대 50% 감소시키며, ImageNet, CIFAR-100, ADE20K에서 성능 저하 없이 유지한다.
• 헤드별 양자화는 이질적인 활성화 분포를 가진 다중 헤드 어텐션 레이어에서 균일 양자화보다 훈련 안정성과 성능 향상에 뚜렷한 기여를 한다.
• 양자화 파ram터에 대한 런닝 추정치는 샘플별 통계보다 유사하거나 더 나은 성능을 내며, 동시에 훈련 오버헤드와 메모리 비용을 감소시킨다.
• 이 프레임워크는 절감된 메모리를 더 큰 배치 크기나 모델 스케일로 재투자할 수 있게 하여 일반화 성능 향상과 더 빠른 수렴을 이끌어낸다.
• CIFAR-100에서 Mesa는 특정 연산을 압축했을 때 기준 훈련보다 略 높은 정확도를 달성하여 유리한 속도-메모리-정확도 트레이드오프를 보여준다.
• 시각화 결과, 양자화 파arameter(α 및 β)는 헤드와 레이어 간에 다르게 진화하며, MSA 레이어에서는 β 값이 음수로 기울어져 음수 활성화의 흔함을 시사한다.