[논문 리뷰] The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits

Shuming Ma, Hongyu Wang외 8명|arXiv (Cornell University)|2024. 02. 27.

Natural Language Processing Techniques인용 수 37

한 줄 요약

논문은 모든 매개변수가 {-1, 0, 1}인 1.58비트 LLM인 BitNet b1.58을 소개한다. 비슷한 크기에서 FP16/LLaMA 성능과 일치하면서도 속도, 메모리 및 에너지 측면에서 큰 이점을 제공하여 새로운 확장 법칙과 1비트 LLM에 대한 잠재적 하드웨어 최적화를 가능하게 한다.

ABSTRACT

Recent research, such as BitNet, is paving the way for a new era of 1-bit Large Language Models (LLMs). In this work, we introduce a 1-bit LLM variant, namely BitNet b1.58, in which every single parameter (or weight) of the LLM is ternary {-1, 0, 1}. It matches the full-precision (i.e., FP16 or BF16) Transformer LLM with the same model size and training tokens in terms of both perplexity and end-task performance, while being significantly more cost-effective in terms of latency, memory, throughput, and energy consumption. More profoundly, the 1.58-bit LLM defines a new scaling law and recipe for training new generations of LLMs that are both high-performance and cost-effective. Furthermore, it enables a new computation paradigm and opens the door for designing specific hardware optimized for 1-bit LLMs.

연구 동기 및 목표

• 성능을 희생하지 않으면서 LLM의 추론 비용과 에너지를 감소시키는 것을 동기화한다.
• 제로 상태를 포함하는 1.58비트 가중치 양자화 체계를 제안하여 모델링 능력을 향상시킨다.
• BitNet b1.58가 3B 크기부터 perplexity 및 엔드태스크 지표 전반에서 풀-정밀도 기준선과 일치할 수 있음을 보여준다.
• 하드웨어 및 에너지 영향력을 보여주고 더 큰 모델 및 엣지 디바이스로의 확장 가능성을 논의한다.

제안 방법

• BitNet b1.58을 도입한다. 선형 계층이 BitLinear로 대체되고 가중치가 {-1, 0, 1}로 양자화되는 Transformer이다.
• 가중치를 제약하기 위해 absmean 양자화를 사용하고, 감마를 평균 절대 가중치 값으로 정의한다.
• 편향 없이 RMSNorm, SwiGLU, 로터리 임베딩 등 LLaMA와 유사한 아키텍처 구성요소를 채택한다.
• RedPajama 데이터로 1.58비트 가중치와 8비트 활성화를 사용해 처음부터 학습한다.
• 모델 크기(700M–3.9B) 전반에 걸쳐 FP16 LLaMA 기준선과 비교하여 perplexity 및 제로샷 엔드 태스크 성능을 평가한다.
• FasterTransformer 및 2비트/Ladder 커널을 사용해 메모리, 대기 시간, 처리량 및 추정 에너지를 측정하고, 확장성 및 에너지 절감을 분석한다.

실험 결과

연구 질문

• RQ1 동일한 모델 크기와 학습 설정에서 1.58비트 LLM이 perplexity와 엔드 태스크에서 풀-정밀도 LLM의 성능과 일치할 수 있는가?
• RQ2 모델 크기가 커질 때 1.58비트 LLM의 메모리, 대기 시간 및 에너지 영향은 무엇인가?
• RQ3 크기에 따라 비용과 정확도 측면에서 FP16 기준선에 비해 1.58비트 LLM이 파레토 개선을 보이는가?
• RQ4 2T 토큰으로의 학습이 BitNet b1.58를 표준 벤치마크에서 일반화 능력을 유지하거나 개선하도록 확장할 수 있는가?

주요 결과

• BitNet b1.58은 동일 구성에서 3B 사이즈부터 perplexity에서 FP16 LLaMA 기준선과 일치한다.
• BitNet b1.58은 3B에서 FP16 LLaMA에 비해 2.71배 더 빠르고 GPU 메모리 사용은 3.55배 더 적다.
• 3.9B에서 BitNet b1.58은 2.4배 더 빠르고 메모리 효율은 3.32배 더 높으며 엔드 태스크에서 LLaMA 3B를 능가한다.
• 크기에 걸쳐 BitNet b1.58은 FP16 기준선과의 성능 차이를 좁히며 3B 이후부터 엔드 태스크에서 유사한 결과를 달성한다.
• BitNet b1.58 70B는 LLaMA 70B에 비해 최대 4.1배의 속도 향상을 달성하고, 두 대의 A100 GPU에서 최대 11배의 처리량을 보이며 메모리도 감소한다.
• 에너지 분석 결과 BitNet b1.58은 7nm 칩에서 행렬 곱셈의 산술 에너지를 FP16 기준선 대비 약 71.4배 절감하며, 더 큰 모델에서 이점이 커진다.