기술 요약: LLM의 다중 비트 너비 양자화를 위한 Drop-by-Drop

문제 정의

대규모 언어 모델(LLM)은 수십억 개의 파라미터로 인해 발생하는 막대한 메모리 및 계산 자원 요구량으로 인해 배포 시 상당한 어려움에 직면해 있습니다. 양자화는 이러한 모델을 유한한 정밀도 표현으로 압축하지만, 기존 방식은 일반적으로 고정된 압축 수준(예: 정적 4비트 또는 8비트 모델)의 모델을 생성합니다. 이러한 경직성은 클라우드 서버부터 자원이 제한된 엣지 디바이스에 이르는 이기종 하드웨어 환경에 최적화되지 못합니다. 다양한 비트 너비를 위해 별도로 양자화된 여러 모델 변체를 유지하는 것은 과도한 저장 공간과 관리 비용을 초래합니다. 또한, 기존의 적응형 방법들은 재학습에 의존하거나, 구성 요소를 제거할 때 성능이 급격히 저하되는 현상 없이 정밀도와 효율성 사이의 매끄러운 트레이드오프를 제공하지 못하는 경우가 많습니다.

방법론: Drop-by-Drop (DbyD)

저자들은 재학습 없이 단일 학습 모델이 여러 정밀도 수준에서 작동할 수 있게 하는 새로운 사후 학습 양자화(PTQ) 프레임워크인 Drop-by-Drop을 제안합니다. 이 방법은 정보 이론에 근거하며 **가산 양자화(Additive Quantization, AQLM)**를 기반으로 합니다.

이론적 토대

이 접근 방식은 연속적 소스 정제(successive source refinement) 개념에서 동기를 얻었습니다. 본 논문은 가우시안 소스(LLM의 가중치가 경험적으로 근사하는 형태)가 가중 평균 제곱 오차(WMSE) 왜곡 척도 하에서 연속적으로 정제 가능하다는 이론적 정리를 확립합니다. 이 왜곡 척도는 입력 데이터 행렬 $X$에 의해 오차가 가중되는 LLM 손실 함수로부터 유도됩니다. 이 정리는 소스를 인코딩할 때, 비트를 점진적으로 추가함으로써 타겟 왜곡 수준에서 직접 인코딩했을 때의 레이트 페널티(rate penalty) 없이 점진적으로 정확한 재구성을 얻을 수 있음을 입증합니다.

알고리즘 구현

Drop-by-Drop은 가중치를 여러 개의 학습된 코드북으로부터의 벡터 합으로 표현하는 AQLM 프레임워크를 변형합니다.

1. 가산 코드북(Additive Codebooks): 단일 양자화 단계 대신, 가중치는 $M$개의 서로 다른 코드북에서 선택된 벡터들의 합으로 재구성됩니다.
2. 마트료시카 스타일 감독(Matryoshka-style Supervision): 이 코드북들의 부분 집합이 정확한 부분 재구성을 보장하도록 하기 위해, 저자들은 학습 과정에서 특화된 손실 함수를 도입합니다. 이 손실 함수는 처음 $k$개의 코드북을 사용한 부분 재구성의 WMSE 왜곡을 합산합니다.
   • $L_{DbyD} = \sum_{k=1}^{K} \lambda_k D^{(k)}$, 여기서 $D^{(k)}$는 처음 $k$개의 코드북을 사용한 왜곡입니다.
   • 이 "마트료시카" 감독은 처음 몇 개의 코드북이 거칠고 근본적인 특징을 포착하고, 후속 코드북이 미세한 정제를 제공하도록 명시적으로 유도합니다.
3. 추론 시 유연성: 추론 시 모델은 코드북을 드롭(drop)함으로써 정밀도를 동적으로 조절할 수 있습니다. 제약이 엄격한 디바이스는 처음 몇 개의 코드북(예: $k=3$)만 사용할 수 있고, 자원이 풍부한 디바이스는 모든 $M$개의 코드북을 활용할 수 있습니다. 이를 통해 단일 체크포인트가 여러 비트 너비를 지원할 수 있습니다.

주요 기여

1. 이론적 증명: 본 논문은 가우시안 소스가 LLM과 관련된 WMSE 왜곡 척도 하에서 연속적으로 정제 가능함을 증명하여, 레이트 손실 없는 점진적 양자화에 대한 이론적 근거를 제공합니다.
2. Drop-by-Drop 프레임워크: 마트료시카 감독을 가산 코드북 학습에 통합하여, 코드북 드롭을 통해 단일 모델이 여러 정밀도 수준을 지원할 수 있게 하는 실용적인 PTQ 프레임워크를 제시합니다.
3. 경험적 검증: Qwen, LLaMA, Gemma, Mistral 등 다양한 아키텍처에 걸친 광범에한 실험을 통해, Drop-by-Drop이 코드북이 드롭될 때 표준 AQLM과 비교하여 우수한 성능 저하(graceful degradation)를 보이며, 전 범위의 비트 너비에서 경쟁력 있는 퍼플렉서티(perplexity)와 태스크 정확도를 유지함을 입증했습니다.

실험 결과

저자들은 0.5B에서 32B 파라미터 규모의 모델을 대상으로 WikiText-2, C4 및 5가지 제로샷 추론 벤치마크(ARC, HellaSwag, Winogrande, PIQA)를 통해 Drop-by-Drop을 평가했습니다.

• 퍼플렉서티(Perplexity): Drop-by-Drop 변체들은 표준 "AQLM (Drop)" 방식(표준 AQLM 모델을 학습시킨 후 추론 시 단순히 코드북을 드롭하는 방식)보다 일관되게 우수한 성능을 보였습니다. 예를 들어, Mistral-7B의 3비트 정밀도에서 표준 AQLM (Drop)은 퍼플렉서티가 112.75로 급격히 증가한 반면, Drop-by-Drop (Weighted)은 개별적으로 학습된 모델의 성능을 밀접하게 추종하며 6.19의 퍼플렉서티를 유지했습니다.
• 정확도: 제로샷 태스크에서 Drop-by-Drop (Weighted)은 모든 비트 너비에 대해 개별적으로 학습된 모델(AQLM Ind.)과 대등한 정확도를 달축했습니다.
• 효율성: 이 방법은 서로 다른 비트 너비를 위해 별도의 모델을 학습하고 저장할 필요를 없애줍니다. 그림 8은 세 가지 별도의 모델(3비트, 4비트, 5비트)을 위해 누적된 학습 시간과 디스크 공간이 필요한 기존 방식과 달리, Drop-by-Drop은 단 한 번의 학습 실행과 저장 인스턴스만으로 유사한 성능을 유지하면서도 이를 해결함을 보여줍니다.

의의 및 주장

본 논문은 Drop-by-Drop이 정보 이론적 연속 정제 이론과 실질적인 LLM 양자화 사이의 간극을 메운다고 주장합니다. 이 연구의 주요 의의는 재학습이나 재보정 없이 추론 시 적응형 정밀도 제어를 가능하게 한다는 점에 있습니다. 단일 모델이 배포 디바이스의 역량에 맞춰 효과적인 자원 사용량을 조절할 수 있게 함으로써, 이 방법은 이기종 하드웨어에 LLM을 배포할 때 발생하는 실질적인 문제를 해결합니다. 저자들은 이를 다중 비트 너비 PTQ와 연속 정제 이론 간의 새로운 연결 고리로 제시하며, 실제 시나리오에서의 런타임 하드웨어 인식 동적 로딩 및 혼합 정밀도 양자화를 위한 솔루션을 제공합니다.