MoBiQuant: Mixture-of-Bits Quantization for Token-Adaptive Elastic LLMs

1. 문제 제기 (Problem Statement)

대규모 언어 모델 (LLM) 의 클라우드 및 엣지 디바이스 배포 시, 가용한 컴퓨팅 자원에 따라 다양한 정밀도 (Quantization Precision) 로 추론을 수행할 수 있는 탄성적 (Elastic) 배포가 필수적입니다. 그러나 기존 Post-Training Quantization (PTQ) 방법론에는 다음과 같은 근본적인 한계가 존재합니다.

• 정밀도 의존적 보정 파라미터: 기존 PTQ 는 특정 비트 폭 (예: 4-bit) 에 맞춰 보정 파라미터 (스케일링 팩터, 클리핑 범위 등) 를 학습합니다. 이로 인해 보정된 모델이 다른 정밀도 (예: 3-bit 또는 5-bit) 로 추론될 때 성능이 급격히 저하되는 일반화 (Generalization) 실패가 발생합니다.
• 아웃라이어 이동 현상 (Outlier Migration): 저정밀도 변환 시 활성화 (Activation) 의 아웃라이어 (Outlier) 가 발생하는 토큰의 집합이 정밀도에 따라 달라집니다. 4-bit 에서는 정상적인 토큰이 3-bit 에서는 주요 오차 원인이 되거나 그 반대의 경우가 발생할 수 있습니다.
• 고정된 정밀도 할당: 기존 "Any-Precision" 방법론들은 모델 전체 또는 레이어 단위로 정밀도를 고정하거나, 비균일 양자화를 사용하여 메모리 오버헤드와 전환 비용을 발생시킵니다. 토큰 단위의 민감도 (Sensitivity) 를 고려하지 않아 탄성적 추론에 비효율적입니다.

2. 방법론 (Methodology)

저자들은 MoBiQuant를 제안하여, 토큰의 민감도에 따라 가중치 정밀도를 동적으로 조절하는 새로운 프레임워크를 구축했습니다. 이 프레임워크는 두 가지 핵심 구성 요소로 이루어집니다.

A. MoBiSlice: 재귀적 잔차 양자화 (Recursive Residual Quantization)

• Many-in-One 구조: 하나의 가중치 체크포인트에서 여러 정밀도 (2~8-bit 등) 를 효율적으로 지원하기 위해 가중치를 비트 슬라이스 (Bit Slices) 로 분해합니다.
• 재귀적 분해: 가중치 $W$를 가장 중요한 비트 (MSB) 슬라이스 $W_1$과 잔차 (Residual) $R_1$로 나눈 후, 잔차를 다시 다음 슬라이스 $W_2$와 $R_2$로 재귀적으로 분해합니다.
  • $R_1 = W, W_e = Q(R_e), R_{e+1} = R_e - W_e$
• 유연한 구성: $k$개의 슬라이스를 합산하여 원하는 정밀도 $b$를 구성할 수 있습니다 ($W^{(b)} = \sum_{e=1}^k W_e$). 이는 메모리 중복 없이 정밀도를 전환할 수 있게 합니다.

B. MoBiRoute: 토큰 인지 라우팅 (Token-Aware Routing)

• 이진 결정 라우터: 각 토큰이 몇 개의 슬라이스를 활성화해야 하는지 결정하는 경량 학습 가능 라우터 (MLP 기반) 입니다.
• 동적 할당: 입력 토큰의 양자화 난이도 (Sensitivity) 를 예측하여, 민감한 토큰은 더 많은 슬라이스 (고정밀도) 를, 단순한 토큰은 적은 슬라이스 (저정밀도) 를 할당합니다.
• 탄성 제어: 학습 중에는 온도를 점진적으로 낮추어 이진 마스크를 학습하고, 추론 시에는 임계값 ($\delta$) 을 조정하여 전체 평균 비트 폭을 목표 부하에 맞춰 실시간으로 조절합니다.

C. 공동 최적화 (Joint Optimization)

• PTQ 과정에서 MoBiSlice 의 양자화 파라미터 ($\Theta_q$) 와 MoBiRoute 의 라우팅 파라미터 ($\Theta_r$) 를 함께 최적화합니다.
• 정규화 목적 함수: 평균 비트 폭이 목표값을 초과하거나 미달할 때 슬라이스를 가지치기하거나 활성화하도록 유도하는 예산 인식 (Budget-aware) 정규화 항을 추가하여 라우터가 다양한 정밀도 조합을 탐색하도록 합니다.

3. 주요 기여 (Key Contributions)

1. 아웃라이어 이동 현상의 규명: 기존 PTQ 의 일반화 실패 원인이 정밀도 의존적인 토큰 민감도 변화 (Outlier Migration) 에 있음을 발견하고, 이를 해결하기 위해 토큰 단위 적응형 정밀도 조절이 필요함을 증명했습니다.
2. MoBiSlice (Many-in-One): 재귀적 잔차 양자화를 통해 단일 모델에서 계층적 재구성을 가능하게 하여, 메모리 오버헤드 없이 2~6 비트 등 다양한 정밀도를 지원하도록 설계했습니다.
3. MoBiRoute (Token-Adaptive): 토큰별 민감도를 기반으로 동적으로 비트 슬라이스를 선택하는 라우터를 도입하여, 정밀도 전환 시 발생하는 일반화 오류를 방지하고 지연 시간 (Latency) 목표에 맞춰 평균 정밀도를 조절합니다.
4. 효율적인 커널 구현: 비트 단위 패킹 (Bit-major packing) 과 이진 행렬 곱셈 (BMMA), 스트림 기반 병렬 실행을 통해 NVIDIA A100 GPU 에서 FP16 대비 최대 2.7 배의 속도 향상을 달성했습니다.

4. 실험 결과 (Results)

• 탄성적 추론 성능: LLaMA-3-8B 및 LLaMA-2 계열 모델에서 MoBiQuant 는 보정 (Calibration) 시 3-bit 로 학습하더라도 4-bit 추론 시 기존 PTQ(OmniQuant 등) 보다 월등히 우수한 일반화 성능을 보였습니다. 특히 2~3 비트의 극저정밀도 구간에서도 안정적인 성능을 유지합니다.
• 정적 PTQ 대비 경쟁력: MoBiQuant 는 3.9~4.0 비트 범위로 제한했을 때, 4-bit 로 전용 보정된 최첨단 정적 PTQ 방법론 (AWQ, GPTQ, SpinQuant 등) 과 동등하거나 더 나은 퍼플렉시티 (Perplexity) 및 제로샷 (Zero-shot) 태스크 정확도를 달성했습니다.
• 아웃라이어 이동 감소: 시각화 결과, MoBiQuant 는 다양한 정밀도 간 토큰별 양자화 오차 분포가 일관되어 있어, 기존 방법론에서 발생하는 아웃라이어 이동 현상이 현저히 감소했음을 확인했습니다.
• 하드웨어 가속: 커널 최적화를 통해 긴 시퀀스 추론 시 메모리 병목 현상을 완화하고 FP16 대비 2.7 배의 속도 향상을 달성했습니다.

5. 의의 및 결론 (Significance)

MoBiQuant 는 LLM 의 탄성적 배포를 위한 새로운 패러다임을 제시합니다.

• 실용성: 별도의 모델 체크포인트를 로드하거나 재학습 없이, 런타임에 시스템 자원 (배터리, 요청 부하 등) 에 따라 정밀도를 유연하게 조절할 수 있습니다.
• 효율성: 토큰 단위의 미세한 정밀도 조절을 통해 불필요한 연산을 줄이고, 하드웨어 자원을 최적화하여 추론 속도를 높입니다.
• 일반화: 정밀도 간 이동 시 발생하는 성능 저하를 해결함으로써, 단일 모델이 다양한 환경에서 고품질 서비스를 제공할 수 있는 기반을 마련했습니다.

결론적으로, MoBiQuant 는 정적 양자화의 경직성과 기존 Any-Precision 방법론의 비효율성을 극복하고, 토큰 민감도에 기반한 적응형 양자화를 통해 LLM 의 실용적 배포를 혁신하는 중요한 기술적 진보입니다.