기술 요약: JacQuant

문제 정의

양자화 인식 학습 (QAT) 은 특히 초저비트 영역 (예: $\le 2$비트 가중치) 에서 엄격한 메모리 및 지연 시간 제약 하에 대규모 언어 모델 (LLM) 을 배포하기 위한 표준 접근법입니다. 그러나 현대의 QAT 파이프라인은 거의 보편적으로 **Straight-Through Estimator (STE)**에 의존합니다. STE 는 양자화기가 항등 함수인 것처럼 그래디언트를 통과시킴으로써 미분 불가능한 양자화기를 근사화합니다.

본 논문은 이 접근법에서 치명적인 결함을 지적합니다. 하드 양자화기는 구간 경계와 포화/클리핑 영역을 가진 조각별 상수 함수입니다. 초저비트 설정에서는 많은 매개변수가 이러한 경계 근처에 위치하거나 클리핑 한계에 도달합니다. STE 의 항등 야코비안 (identity Jacobian) 에 대한 가정은 지속적인 불일치를 초래하여, 양자화기가 국소적으로 무감각한 포화 영역에서 그래디언트가 과도하게 전파되게 합니다. 이는 취약한 학습, 구간 경계 근처의 불안정성, 그리고 실제 저정밀도 모델의 동작과 약하게 정렬된 해를 초래합니다.

방법론: JacQuant

JacQuant 은 STE 의 고정된 항등 야코비안을 **학습 가능한 경량 대리 야코비안 (learned, lightweight surrogate Jacobian)**인 $B(W)$로 대체하는 프레임워크를 제안합니다. 핵심 철학은 QAT 를 "야코비안 누락" 문제로 간주하는 것입니다. 즉, 순전파 (forward pass) 에는 배포된 하드 양자화기 $Q_\Delta(\cdot)$를 유지하는 반면, 역전파 (backward pass) 에서는 정밀도 가중치의 섭동이 양자화된 대리 변수에 미치는 영향을 학습합니다.

핵심 메커니즘

표준 STE 업데이트 $\nabla_W \ell \approx \nabla_{\hat{W}} \ell$ 대신, JacQuant 는 다음을 사용합니다:
$$\nabla_W \ell \approx B(W) \nabla_{\hat{W}} \ell$$
여기서 $\hat{W} = Q_\Delta(W)$이며 $B(W)$는 학습된 민감도 맵입니다.

주요 설계 선택

1. 매개변수화: $B(W)$는 대각 또는 블록 대각으로 제한되며, 양자화 그룹당 (예: 128 개 가중치) 하나의 스칼라 이득을 가집니다. 이로 인해 메모리 오버헤드는 무시할 수준으로 유지됩니다 (Adam 옵티마이저 상태에 약 1.5% 만 추가).
2. 추정 전략: 논문은 $B(W)$를 업데이트하기 위한 두 가지 실용적이고 데이터 기반의 추정기를 소개합니다:
   • JacQuant-Probe: 국소 기울기 $\Delta q_g = Q_\Delta(W_g + \delta_g) - Q_\Delta(W_g)$를 측정하기 위해 가우시안 랜덤 프로브 $\delta_g$를 사용합니다. 최소제곱법 (least-squares fit) 으로 국소 민감도를 추정합니다.
   • JacQuant-Dither: 양자화기의 미분 가능한 대리 변수를 만들기 위해 감산적 디더링 (subtractive dithering, 균일 노이즈 추가) 을 사용하며, 기대 응답을 기반으로 민감도를 추정합니다.
3. 상각 업데이트 (Amortized Updates): 학습 처리량을 유지하기 위해 $B(W)$는 모든 단계에서 업데이트되지 않습니다. 지수 이동 평균 (EMA) 을 사용하여 가끔씩 (예: 50~200 단계마다) 새로고침됩니다. 새로고침 사이에는 현재 $B(W)$가 재사용됩니다.
4. 분산 감소: 프레임워크는 프로빙/디더링 과정에서 도입된 확률성을 안정화시키기 위해 분산 감소 옵티마이저 (예: loopless-SVRG, SAGA) 와 통합됩니다.

이론적 기여

논문은 비볼록 목적 함수 하에서 JacQuant 에 대한 수렴 분석을 제공하며, 네 가지 주요 논증으로 구성됩니다:

1. 목표 그래디언트 정의: 저자들은 "평균장 (mean-field)" 목표 그래디언트 $g^\dagger(W) = J(W)\bar{v}$를 정의합니다. 여기서 $J(W)$는 디더 평균화된 양자화기의 실제 민감도입니다. 그들은 JacQuant 의 편향이 $\|B(W) - J(W)\|$에 비례하는 반면, STE 의 편향은 $\|I - J(W)\|$에 비례함을 증명합니다.
2. 추정기의 일관성: Probe 및 Dither 추정기 모두 목표 민감도 $J(W)$와 일관성이 있음을 보여줍니다. 구체적으로, Probe 업데이트는 평균장 민감도로 수렴하는 국소 회귀이며, Dither 업데이트는 $J(W)$를 고유한 인구 고정점으로 가집니다.
3. 안정성: 학습된 $B(W)$는 적응형 이득 제어기로 작용합니다. 반응이 좋은 구간 내부에서는 항등에 가깝게 유지되지만, 양자화기가 무감각한 포화 영역에서는 그래디언트를 감쇠시킵니다 (이득 $< 1$). 이는 STE 기반 학습에서 흔히 발생하는 진동을 방지합니다.
4. 수렴 속도: Polyak-Łojasiewicz (PŁ) 조건과 "코드 보존" 속성 (국소 윈도우 내에서 양자화 코드가 자주 변경되지 않음) 을 가정할 때, 논문은 JacQuant 가 선형 수렴 속도를 달성함을 증명합니다. 수렴 오차는 야코비안 추적 오차와 확률적 분산 바닥에 의해 제한됩니다.

실험 결과

저자들은 LLaMA3-1B/3B/8B, Qwen3-1.7B 등 여러 LLM 계열을 초저비트 영역 ($W \le 2$, $A8/A16$) 에서 평가했습니다.

• 성능: JacQuant 는 STE 기반 베이스라인 (ParetoQ 및 WinQ 포함) 과 "Rotation Trick"(고정 기하학적 변환 방법) 보다 일관되게 우수한 성능을 발휘했습니다.
  • **LLaMA3-1B (W2A16)**에서 JacQuant 는 가장 좋은 STE 기반 베이스라인 대비 퍼플렉시티를 13.5 에서 12.3 으로, 정확도를 55.7% 에서 56.6% 로 개선했습니다.
  • **LLaMA3-8B (W2A16)**에서 JacQuant 는 30K, 60K, 90K 학습 단계에 걸쳐 ParetoQ 보다 단조로운 개선을 보였습니다.
• 효율성: 런타임 오버헤드는 무시할 수준입니다. LLaMA3-1B 에서 오버헤드는 그룹 크기에 따라 1~5% 범위이며, LLaMA3-8B 에서 실용적인 그룹 크기의 경우 1.6% 미만으로 유지됩니다.
• 메커니즘 검증: 유한 차분 진단은 JacQuant 가 클리핑/포화 영역에서 STE 가 과도하게 전파할 그래디언트를 효과적으로 감쇠시키는 그룹 이득 $< 1$을 학습함을 확인시켜 주었습니다.

중요성 및 주장

논문은 JacQuant 를 기존 QAT 파이프라인을 위한 **즉시 적용 가능한 업그레이드 (drop-in upgrade)**로 위치시킵니다. 그 중요성은 다음과 같습니다:

1. STE 없는 학습: 순전파 양자화기나 모델 아키텍처 변경 없이 휴리스틱 항등 야코비안에 대한 의존성을 제거합니다.
2. 기하학적 인식: 국소 민감도를 학습함으로써 양자화 지형의 기하학에 자연스럽게 적응하며, 특히 초저비트 학습을 괴롭히는 포화 및 클리핑 문제를 해결합니다.
3. 실용성: 최소한의 계산 비용으로 이러한 이점을 달성하여 대규모 LLM 학습에 실용적입니다.

저자들은 경험적 주장을 명시적으로 초저비트 LLM QAT($\le 2$비트) 로 한정합니다. 역방향 규칙에 대한 즉각적인 대체재이지만 순전파 양자화기는 여전히 배포된 하드 연산자로 남는다는 점을 인정합니다. 논문은 모든 양자화 과제를 해결한다고 주장하지는 않지만, STE 기반 최적화에 내재된 그래디언트 불일치를 구체적으로 대상으로 합니다. 향후 연구 방향으로는 혼합 정밀도 정책으로의 확장, 비균일 양자화기 적용, 그리고 추정 스케줄링 정교화가 언급됩니다.