[논문 리뷰] Training Transformers with 4-bit Integers

Haocheng Xi, Changhao Li외 2명|arXiv (Cornell University)|2023. 06. 21.

Advanced Neural Network Applications인용 수 8

한 줄 요약

논문은 forward passes에 Hadamard 양자화를 사용하고 backpropagation에 비트 분할과 활용 점수 샘플링으로 INT4 기반 트랜스포머 학습 방법을 제시하여 다양한 작업에서 경쟁력 있는 정확도와 하드웨어 친화적 속도 향상을 달성한다.

ABSTRACT

Quantizing the activation, weight, and gradient to 4-bit is promising to accelerate neural network training. However, existing 4-bit training methods require custom numerical formats which are not supported by contemporary hardware. In this work, we propose a training method for transformers with all matrix multiplications implemented with the INT4 arithmetic. Training with an ultra-low INT4 precision is challenging. To achieve this, we carefully analyze the specific structures of activation and gradients in transformers to propose dedicated quantizers for them. For forward propagation, we identify the challenge of outliers and propose a Hadamard quantizer to suppress the outliers. For backpropagation, we leverage the structural sparsity of gradients by proposing bit splitting and leverage score sampling techniques to quantize gradients accurately. Our algorithm achieves competitive accuracy on a wide range of tasks including natural language understanding, machine translation, and image classification. Unlike previous 4-bit training methods, our algorithm can be implemented on the current generation of GPUs. Our prototypical linear operator implementation is up to 2.2 times faster than the FP16 counterparts and speeds up the training by up to 35.1%.

연구 동기 및 목표

• 트랜스포머를 초저정밀(INT4)로 학습시켜 계산 속도를 높이고 메모리 사용량을 줄이려는 동기를 제시한다.
• Hadamard 기반 변환을 사용해 활성화의 이상치(outliers)를 완화하는 순방향 양자화를 개발한다.
• 그래디언트의 구조적 희소성을 활용하여 비트 분할과 활용 점수 샘플링으로 백프로파게이션을 효과적으로 양자화한다.
• 현재 GPU와 호환되는 실용적인 INT4 MM 기반 학습 파이프라인을 제공한다.
• NLP, 번역, 비전 벤치마크에서 경쟁력 있는 정확도와 속도 향상을 입증한다.

제안 방법

• 트랜스포머의 모든 선형 연산을 행렬 곱(MM)으로 구성하고 LSQ 기반 양자화를 통해 INT4 산술을 적용한다.
• 활성화를 양자화 전에 블록 대각 Hadamard 행렬로 변환하여 순방향 활성화의 이상치를 억제하는 Hadamard 양자화(HQ)를 도입한다.
• 역전파에서 그래디언트의 구조적 희소성을 활용하기 위해 비트 분할(BS)을 적용해 그래디언트를 상위/하위 4비트 파트로 분할하고 정보성 행을 MM 계산에 선택하기 위해 활용 점수 샘플링(LSS)을 사용한다.
• 필요에 따라 FP16 누적을 위한 최종 결과를 양자화된 상태에서 역전파를 위한 각 MM에 대해 그래디언트를 동적으로 양자화하고 활성화/가중치는 INT4를 유지한다.
• 순방향용 HQ-MM과 역전파용 LSS-MM 두 가지 MM 절차를 제공하여 엔드 투 엔드 INT4 MM 기반 학습을 가능하게 한다.
• NLP 벤치마크(GLUE, SQuAD, CoNLL, MT)와 이미지 분류(ViT, ImageNet)에서 FP, INT8, FP4, LSQ+LUQ 기초선과의 비교를 평가한다.

실험 결과

연구 질문

• RQ1 표준 하드웨어만으로 맞춤형 수치 형식을 사용하지 않고 4비트 정수 산술로 트랜스포머를 엔드투엔드 학습할 수 있는가?
• RQ2 초저정밀 학습에서 정확도를 유지하기 위해 순방향 활성화 이상치와 역전파 그래디언트 희소성을 어떻게 해결할 수 있는가?
• RQ3 Hadamard 기반 순방향 양자화와 활용 점수 샘플링 기반 백프로파게이션의 정확도 및 효율성 트레이드오프는 무엇인가?

주요 결과

• HQ+LSS는 NLP, MT, 비전 작업 전반에서 FP 및 기존 4비트 학습 방법과 비교해 경쟁력 있는 정확도를 달성한다.
• GLUE 및 SQuAD 스타일 평가에서 HQ+LSS는 LSQ+LUQ를 개선하고 특히 Bert-base 및 Bert-large 변형에서 이를 능가한다.
• 기계 번역(WMT)에서 HQ+LSS는 약 1.0% BLEU 저하를 보이며 일부 초저정밀 설정보다 우수하고 다른 설정과 비슷하며 하드웨어 친화적이다.
• ViT 모델을 이용한 이미지 분류는 전체 정밀도 벤치마크에 비해 정확도 저하가 작고(대개 1-2% 미만) HQ+LSS가 여러 경우 LSQ+LUQ를 능가한다.
• 이 방법은 현재 GPU에서 INT4 MM 구현을 가능하게 하며 FP16 MM보다 최대 2.2배의 속도 향상 및 학습 속도 증가 최대 35.1%를 제공한다.
• CoLA에 대한 소거 연구에서 순방향 Hadamard 양자화의 정확도 유지의 중요성과 백프로파게이션은 BS+LSS의 이점을 확인한다.