Microsoft Research의 BitNet 아키텍처를 파인튜닝으로 적용해 기존 Llama3 8B 모델을 1.58비트 극단 양자화 모델로 변환하고 MMLU 벤치마크에서 Llama 1 7B를 초과하는 성능 달성

Context

대규모 언어 모델의 계산 비용과 에너지 소비가 증가하면서 모델 크기를 줄여야 하는 과제가 발생했다. 기존 8비트 또는 4비트 양자화는 메모리와 계산 속도는 개선하지만 정확도 손실을 초래한다. BitNet 아키텍처는 파라미터당 1.58비트 극단 양자화를 제공하지만 처음부터 학습해야 하므로 자원이 제한된 조직에는 실용적이지 못했다.

Technical Solution

• 파라미터를 -1, 0, 1의 세 값으로 표현하는 BitLinear 레이어로 기존 Linear 레이어 교체: Multi-Head Attention과 Feed-Forward Networks의 모든 선형층 변경
• 가중치 양자화는 대칭 Per-Tensor 양자화 사용: 가중치 행렬의 절댓값 평균을 스케일로 설정 후 반올림으로 -1, 0, 1로 양자화
• Straight Through Estimator(STE) 기법 적용: 비미분 가능한 반올림 연산을 통해 그래디언트 흐름을 가능하게 하여 표준 역전파로 가중치 업데이트
• 전체 정밀도로 학습하되 진행 중 가중치를 테르너리 값으로 양자화: 활성화는 8비트 정밀도로 유지
• Hugging Face Transformers에 'bitnet' 양자화 메서드 추가: BitLinear 레이어로 자동 변환 및 동적 활성화 양자화 처리, API 변경 없음

Impact

행렬 곱셈의 산술 연산 에너지를 Llama 대비 71.4배 감소. Llama3 8B 모델을 10B 토큰 및 100B 토큰으로 파인튜닝한 결과 MMLU 벤치마크에서 Llama 1 7B 모델을 초과.

Key Takeaway

극단 양자화는 처음부터 학습할 필요 없이 기존 사전학습 모델의 파인튜닝으로도 달성 가능하며, STE와 같은 미분 근사 기법을 활용하면 이산 가중치 표현에서도 그래디언트 기반 최적화가 가능하다. 이는 제한된 자원으로도 에너지 효율적인 LLM을 구축할 수 있는 경로를 제시한다.