기술 요약: 모델 보존형 적응형 반올림 (YAQA)

문제 정의

대규모 언어 모델(LLM) 양자화의 주요 목표는 원래 모델의 출력 분포를 최대한 가깝게 보존하면서 모델 파라미터를 저정밀도 표현으로 압축하는 것입니다. 이는 원래 모델 $M(\theta^*, X)$와 양자화된 모델 $M(\theta, X)$ 사이의 쿨백-라이블러(KL) 발산을 최소화하는 것으로 공식 정의됩니다.

GPTQ 및 LDLQ와 같은 현재의 최첨단 사후 학습 양자화(PTQ) 방법들은 **로컬 레이어별 적응형 반올림(local layerwise adaptive rounding)**에 의존합니다. 이러한 알고리즘들은 전체 엔드 투 엔드(end-to-end) 오차의 대리 지표로서 단일 레이어의 즉각적인 활성화 오차를 최소화합니다. 그러나 이 접근 방식은 후속 레이어를 통해 전파되는 양자화 오차를 무시하므로, 진정한 목적에 비추어 볼 때 하위 최적(suboptimal)인 대리 지표가 됩니다. 다른 방법들이 글로벌 정보(예: 경험적 피셔 행렬 또는 블록 대각 근사)를 통합하려고 시도하기도 하지만, 이들은 종-엔드 오차 경계에 대한 이론적 보장을 결여하고 있으며, 성능을 일관되게 개선하지 못하는 임의적인 헤시안 근사에 의존하는 경우가 많습니다.

방법론: Yet Another Quantization Algorithm (YAQA)

YAQA는 로컬 레이어별 오차가 아닌 직접적으로 엔드 투 엔드 KL 발산을 최소화하도록 설계된 새로운 적응형 반올림 알고리즘입니다. 이는 이론적으로 근거가 확실한 반올림 알고리즘과 확장 가능한 헤시안 스케치(Hessian sketches)라는 두 가지 핵심 구성 요소로 이루어져 있습니다.

1. 엔드 투 엔드 레이어별 적응형 반올림

YAQA는 전체 모델 에러를 고려하도록 LDLQ의 고정점 반복(fixed-point iteration)을 일반화합니다. 즉각적인 활성화 오차의 헤시안($H_1$)을 사용하는 대신, YAQA는 크로네커 인수 분해된 헤시안 근사 $\tilde{H} = H_O \otimes H_I$를 활용하며, 여기서 $H_O$와 $H_I$는 각각 출력 측 및 입력 측 피드백을 나타냅니다.

반올림 업데이트는 다음과 같이 정의됩니다:
$$W = \text{vec}^{-1} \left( Q \left( \text{vec}(W^*) + \text{vec}(W^* - W)(L - I) \right) \right)$$
여기서 $L$은 $\tilde{H}$의 LDL 분해에서 얻은 삼각 성분입니다.

구조적 닐포턴 차수 (Structural Nilpotence Degree, SND):
이 반복 과정을 실행 가능하게 만들기 위해, 저자들은 구조적 닐포턴 차수(SND)라는 개념을 도입합니다. 일반적인 조밀한(dense) 헤시안의 경우 수렴 단계는 $O(mn)$이지만, 크로네커 인수 분해 구조($\tilde{H} = H_O \otimes H_I$)를 사용함으로써 SND는 $m + n - 1$로 제한됩니다. 이를 통해 알고리즘은 $O(m^2n^2)$의 복잡도를 갖는 조밀한 방식 대신, 고도로 병렬화 가능한 행렬 곱셈을 포함하는 $m+n$ 번의 반복 단계 내에 수렴할 수 있습니다. 이 구조는 이전 방법들이 입력 측 피드백만을 활용했던 것과 달리, 입력 및 출력 채널 모두로부터 대칭적인 피드백을 도입합니다.

2. 헤시안 스케치 (Hessian Sketches)

대규모 LLM에 대해 $H_O$와 $H_I$를 실행 가능하게 계산하기 위해, YAQA는 실제 KL 발산의 헤시안($\nabla^2 L(W^*)$)에 대한 파워 반복(power iteration)으로부터 유도된 두 가지 헤시안 스케치를 제안합니다.

• 스케치 A (Sketch A): 시퀀스 내의 토큰들이 독립적이라고 가정합니다. 이는 토큰을 독립적인 샘 샘플로 취급하여 분산을 줄이는 헤시안 추정치에 대해 파워 반복을 수행하지만, 약간의 편향(bias)을 유발합니다. 이 방식은 데이터에 대한 여러 번의 패스(pass)가 필요합니다.
• 스케치 B (Sketch B): 항등 초기화(identity initialization) 상태에서 실제 헤시안에 대해 단 한 번의 파워 반복을 수행합니다. 이 방법은 시퀀스 믹싱(sequence mixing)을 고려하여 전체 시퀀스에 대한 그래디언트를 계산합니다. 이는 Shampoo의 프리컨디셔닝 기저(preconditioning basis)와 개념적으로 유사하지만, 실제 피셔 정보 행렬을 사용합니다.

이론적 분석에 따르면, YAQA의 엔드 투 엔드 오차는 크로네커 근사($H_O \otimes H_I$)와 실제 헤시안($\nabla^2 L(W^*)$) 사이의 코사인 유사도에 의해 제한됩니다. 두 스케치 모두 낮은 무결성(incoherence)과 낮은 랭크(rank)를 유지하면서 이 유사도를 최대화하도록 설계되었습니다.

주요 기여

1. YAQA 알고리즘: 로컬 방법을 일반화하여 전체 모델의 엔드 투 엔드 KL 발산을 최적화하는 새로운 적응형 반올림 알고리즘을 도입했습니다. 이는 엔드 투 엔드 오차에 대한 이론적 경계를 제공하는 최초의 양자화 알고리즘입니다.
2. 이론적 프레임워크: 적응형 반올림의 수렴 시간을 구조적 닐포턴 차수(SND)를 통해 규명하고, 양자화 알고리즘을 위한 최초의 엔드 투 엔드 오차 경계를 도출했습니다. 이 경계는 근사 오차를 헤시안 스케치와 실제 헤시안 사이의 코사인 유사도와 연결합니다.
3. 확장 가능한 헤시안 스케치: 과도한 계산 비용 없이 글로벌 오차 정보를 계산할 수 있게 해주는 두 가지 실행 가능한 근사적 크로네커 인수 분해 헤시안 스케치(Sketch A 및 Sketch B)를 개발했습니다.
4. 경험적 우월성: YAQA가 GPTQ, LDLQ, GuidedQuant 및 심지어 양자화 인식 훈련(QAT)을 포함한 기존의 SOTA 방법들과 비교하여 현저히 낮은 KL 발산과 더 나은 다운스트림 태스크 성능을 달성함을 입증했습니다.

결과

저자들은 Llama 3.1 (1B, 3B, 8B, 70B), Gemma 3 (12B), Qwen 3 (8B) 모델에 대해 YAQA를 평가했습니다.

• KL 발산: YAQA는 LDLQ 및 GPTQ에 비해 KL 발산을 약 30% 감소시킵니다. 주목할 점은, YAQA가 훨씬 더 많은 계산량과 데이터를 필요로 하는 Google의 Gemma 3 QAT 레시피보다 더 낮은 KL 발산을 달로 달성했다는 것입니다.
• 다운스트림 성능: YAQA는 다양한 비트 너비(2-bit ~ 4-bit)에서 제로샷 다운스트림 태스크(예: MMLU, GSM8K, BoolQ)에 대해 새로운 SOTA 결과를 세웠습니다.
• 효율성: 이 방법은 양자화된 모델에 추론 오버헤드를 추가하지 않습니다. 복잡성은 사후 학습 양자화 단계에 국한됩니다.
• 강건성: 실험 결과, Sketch B는 단 2,000개의 시퀀스만으로도 SOTA 결과를 달성하였으며, 이는 전형적인 QAT 방법보다 훨씬 적은 데이터를 요구함을 보여줍니다.

의의 및 주장

본 논문은 YAQA가 로컬 레이어별 오차 최소화에서 글로벌 엔드 투 엔드 오차 최소화로 양자화 전략의 근본적인 전환을 나타낸다고 주장합니다. 그 의의는 다음과 같습니다:

• 이론적 보장: 양자화 알고리즘의 성능을 헤시안 근사의 기하학적 특성(코사인 유사도)과 연결하는 최초의 증명 가능한 엔드 투 엔드 오차 경계를 제공합니다.
• 성능: 로컬 반올림 방법과 계산 비용이 많이 드는 QAT 방법 모두보다 높은 충실도(fidelity)의 양자화를 달성합니다.
• 실용성: 모델의 재학습이나 미세 조정이 필요 없는 확장 가능한 솔루션을 제공하여, 추론 페널티 없이 대규모 모델을 효율적으로 배포하는 데 적합합니다.

저자들은 GuidedQuant와 같은 이전 방법들이 글로벌 정보를 탐색했음에도 불구하고 이론적 근거와 일관된 성능 향상이 부족했다고 강조합니다. YAQA는 수학적으로 원칙에 기반한 접근 방식을 제공함으로써 이 간극을 메우며, 이론적으로 견고할 뿐만 아니라 경험적으로도 우수합니다.