[논문 리뷰] BitNet: Scaling 1-bit Transformers for Large Language Models

Hongyu Wang, Shuming Ma외 8명|arXiv (Cornell University)|2023. 10. 17.

Topic Modeling인용 수 26

한 줄 요약

BitNet은 BitLinear가 적용된 1-비트 트랜스포머를 도입하여 이진화 가중치를 사용해 처음부터 학습을 가능하게 하면서도 경쟁력 있는 perplexity를 유지하고 FP16 및 사후 학습 8-비트 기준선 대비 상당한 메모리 및 에너지 절감을 달성합니다. 또한 전체 정밀도 트랜스포머와 유사한 스케일링 법칙을 보입니다.

ABSTRACT

The increasing size of large language models has posed challenges for deployment and raised concerns about environmental impact due to high energy consumption. In this work, we introduce BitNet, a scalable and stable 1-bit Transformer architecture designed for large language models. Specifically, we introduce BitLinear as a drop-in replacement of the nn.Linear layer in order to train 1-bit weights from scratch. Experimental results on language modeling show that BitNet achieves competitive performance while substantially reducing memory footprint and energy consumption, compared to state-of-the-art 8-bit quantization methods and FP16 Transformer baselines. Furthermore, BitNet exhibits a scaling law akin to full-precision Transformers, suggesting its potential for effective scaling to even larger language models while maintaining efficiency and performance benefits.

연구 동기 및 목표

• 성능을 저하시키지 않으면서 대형 언어 모델의 메모리 및 에너지 비용을 줄이는 것을 목표로 한다.
• 처음부터 학습 가능하며 교체형으로 사용할 수 있는 1-비트 트랜스포머 구성요소(BitLinear)를 도입한다.
• BitNet의 안정성, 확장성 및 효율성을 FP16 및 사후 퀀타이제이션 기준선과 비교하여 시연한다.
• BitNet이 전체 정밀도 트랜스포머와 유사한 스케일링 법칙을 따르고 더 큰 모델로 효과적으로 확장될 수 있음을 보인다.

제안 방법

• 최적화기 상태와 그래디언트를 높은 정밀도로 유지하면서 이진화 가중치를 사용하도록 nn.Linear를 BitLinear로 교체한다.
• 가중치를 ±1로 이진화하고 sign(W - alpha)로 중심화하여 제로 평균으로 만든 뒤 이진화 후 스케일링 beta를 적용한다.
• 훈련 중 활성화를 absmax로 8비트로 양자화하고 텐서당 양자화를 수행한다; 추론 시에는 토큰당 양자화한다.
• 양자화 경로를 통해 분산을 안정시키기 위해 활성화 양자화 전에 LayerNorm(SubLN)을 적용한다.
• 추가 커뮤니케이션 없이도 효율적인 모델 병렬화를 가능하게 하기 위해 Group Quantization과 Group Normalization을 구현한다.
• 비분화 가능 단계에 대해 Straight-Through Estimator(STE)로 학습하고 혼합 정밀도(가중치/활성화는 저정밀도; 그래디언트/최적화기는 고정밀도)를 사용한다; 수렴을 개선하기 위해 대형 학습률을 사용한다.

실험 결과

연구 질문

• RQ1 양자화 인지 학습으로 학습된 1-비트 가중치가 대형 언어 모델에서 경쟁력 있는 perplexity 및 다운스트림 태스크 성능을 달성할 수 있는가?
• RQ2 모델 크기가 커질 때 BitNet은 FP16 트랜스포머 및 사후 학습 양자화 방법과 에너지 소비 및 메모리 사용 측면에서 어떻게 비교되는가?
• RQ3 BitNet은 전체 정밀도 트랜스포머와 유사한 스케일링 법칙을 보이는가, 그리고 추론 에너지 스케일링 동작은 어떠한가?
• RQ4 1-비트 가중치로 수렴을 달성하기 위해 어떤 학습 안정화 도구(예: SubLN, absmax 활성화 양자화)가 중요한가?

주요 결과

• BitNet은 FP16 트랜스포머 및 8-비트 양자화 기준선과 비교하여 메모리 및 에너지를 크게 절감하면서 경쟁력 있는 perplexity 및 다운스트림 태스크 정확도를 달성한다.
• BitNet의 에너지 절감은 1-비트 가중치와 스케일 팩터로의 곱셈을 한정시켜 행렬 곱셈 연산에서 특히 두드러진다.
• BitNet은 전체 정밀도 트랜스포머와 유사한 스케일링 법칙을 따르며 FP16 대비 추정 성능 차이는 모델 크기가 커질수록 줄어든다.
• BitNet은 동일한 학습률 설정에서 FP16보다 더 큰 학습률과 더 빠른 수렴을 가능하게 하는 더 높은 학습 안정성을 보여준다.
• 제로샷 및 파샷 이하 다운스트림 태스크 성능은 계산 예산과 함께 스케일링되며, 더 큰 규모에서 BitNet의 추론 효율성이 우수함을 보여준다.