초기화가 분지를 결정한다: 극한의 LLM 양자화를 위한 효율적인 코드북 최적화

대형 언어 모델을 스마트폰이나 CPU 환경에서 실행하려면 2비트 수준의 극한 압축이 필요하지만, 기존 방식은 성능이 급격히 떨어지는 문제가 있었다. 이 논문은 복잡한 계산을 늘리는 대신 '초기 설정'이라는 근본적인 병목을 해결하여, 추가 연산 비용 없이도 압축 모델의 품질을 획기적으로 높이는 방법을 제시한다.

딥러닝 모델의 가중치를 아주 적은 비트로 압축할 때, 가중치들을 대표하는 값들의 집합인 '코드북'을 어떻게 구성하느냐가 핵심이다. 기존에는 가중치 하나하나를 순차적으로 가장 가까운 대표값에 할당하는 방식을 사용했는데, 이는 마치 미로의 입구에서 당장 눈앞의 길만 보고 들어가는 것과 같아서 나중에 아무리 좋은 길을 찾으려 해도 이미 막다른 길(나쁜 최적화 영역)에 갇히게 되는 문제를 야기한다.

이 논문은 특히 2비트와 같은 극한의 압축 상황에서는 선택할 수 있는 대표값의 가짓수가 급격히 줄어들어 이 '첫 단추'를 잘못 끼우는 효과가 치명적임을 발견했다. 이를 해결하기 위해 단순히 가중치 숫자 자체의 차이를 줄이는 것이 아니라, 그 가중치가 모델의 최종 결과물에 얼마나 큰 영향을 주는지(Hessian)를 고려하여 초기 대표값들을 배치한다.

결과적으로 모델은 학습이나 미세 조정(Fine-tuning)을 시작하기도 전에 이미 성능이 잘 나올 수밖에 없는 유리한 지형(Basin)에 자리를 잡게 된다. 이는 이후에 복잡한 탐색 알고리즘을 수행하는 것보다 훨씬 효율적이며, 한 번 정해진 좋은 지형은 후속 학습 과정에서도 유지되어 최종 모델의 품질을 결정짓는다.

가산 양자화(Additive Quantization) 프레임워크인 AQLM을 기반으로 하며, 각 가중치 그룹 w를 M개의 코드북에서 추출한 코드워드들의 합으로 근사한다. 기존의 잔차 k-means 방식이 코드북 간의 결합 구조를 무시하고 순차적으로 최적화하여 발생하는 서브옵티멀(suboptimal) 문제를 해결하고자 한다.

제안된 OA-EM(Output-Aware EM)은 두 단계의 반복으로 구성된다. M-step(Centroid Optimisation)에서는 할당된 가중치 그룹들을 고정하고, Hessian 가중치가 적용된 재구성 오차를 최소화하도록 코드북의 중심점(Centroid)을 업데이트한다. 이때 [입력 활성화 값의 상관행렬 XᵀX → 가중치 오차와 곱셈 → 출력 오차 계산] 과정을 통해 모델 출력에 민감한 가중치 영역을 더 정밀하게 보존한다.

E-step(Hard Reassignment)에서는 업데이트된 중심점들을 바탕으로 각 가중치 그룹을 가장 가까운 코드워드 조합에 다시 할당한다. 이때도 단순 유클리드 거리가 아닌 Hessian 가중 마할라노비스 거리를 사용하여 [가중치 차이 벡터 → Hessian 행렬과 내적 → 거리 점수 산출] 과정을 거침으로써 모델 성능 관점에서 최적의 할당을 수행한다. 이 과정을 R=3 라운드 반복하여 최적의 초기 상태를 구축한다.

Llama 3.2 3B 모델의 2비트 양자화 실험에서, 기존 탐욕적 초기화 방식은 WikiText-2 Perplexity 352.39라는 처참한 결과를 보였으나, OA-EM은 동일한 연산 자원에서 16.82를 기록하며 정상적인 모델 작동을 가능케 했다. 빔 서치(Beam Search) 폭을 4에서 16으로 늘려도 기존 방식은 46.01에 그쳤으나, OA-EM은 좁은 빔 폭(b=4)만으로도 이를 압도했다.

Llama 3.1 8B와 Qwen 2.5 3B 모델에서도 일관된 개선이 확인되었다. 특히 PV-tuning이라는 강력한 후속 최적화를 거친 후에도 초기화 단계에서 발생한 성능 격차가 사라지지 않고 유지되는 '분지 지속성(Basin Persistence)' 현상을 발견했다. 이는 초기화가 단순히 시작점이 아니라 최종 도달 가능한 성능의 한계를 결정함을 시사한다.

효율성 측면에서 OA-EM은 기존 방식보다 38% 적은 양자화 시간으로도 더 낮은 Perplexity를 달성했다. 제로샷(Zero-shot) 작업 정확도에서도 Llama 3.2 3B 기준 평균 1.7%p 이상의 향상을 보이며 극한 압축 환경에서의 실용성을 입증했다.

본 연구는 가산 양자화의 성능 저하가 단순한 탐색 부족이 아닌, 초기화 단계에서 발생하는 '조기 확정(Premature Commitment)' 오류 때문임을 수학적으로 분석했다. Proposition 1을 통해 탐욕적 순차 할당 시 발생하는 서브옵티멀 갭(Suboptimality Gap)을 직접 비용, 결합 항, 잔차 불일치 세 가지 요소로 분해하여 설명했다.

핵심 지표인 표현 비율 ρ = N/KM은 가중치 그룹 수(N) 대비 코드북이 표현 가능한 조합 수(KM)의 비율을 의미한다. ρ < 1인 과잉(Overcomplete) 상태에서는 초기화 오류가 흡수될 여지가 있으나, 2비트와 같이 ρ ≫ 1인 부족(Undercomplete) 상태에서는 초기 코드워드 배치가 후속 최적화의 성패를 좌우하는 결정적 요인이 됨을 규명했다.

OA-EM은 AQLM의 잔차 프레임워크 내에서 작동하며, Adam 옵티마이저와 코사인 스케줄링을 사용하여 중심점을 업데이트한다. 특히 블록 대각 Hessian 근사치를 사용하여 연산 효율성을 유지하면서도 출력 민감도를 반영했다. 실험 결과, 이러한 '출력 인지형' 배치가 캘리브레이션 데이터에 대한 과적합을 방지하고 도메인 변화(Domain Shift)에 대한 강건성을 높이는 효과도 확인되었다.

본 방법론은 자유 형식(Free-form) 가산 양자화에 특화되어 있으며, 격자 기반(Lattice-based)이나 트렐리스 코드(Trellis-coded)와 같이 수학적으로 고정된 구조를 가진 양자화 방식에는 직접적으로 적용되지 않는다. 또한 영어 모델과 벤치마크를 중심으로 평가되었으므로 다국어 환경에서의 성능은 추가 검증이 필요할 수 있다.