[논문 리뷰] Optimal Scalar Quantization for Matrix Multiplication: Closed-Form Density and Phase Transition

Calvin Ang, Sungyoon Kim외 1명|arXiv (Cornell University)|2026. 03. 20.

Stochastic Gradient Optimization Techniques인용 수 0

한 줄 요약

이 연구는 행렬 곱셈에 대한 고속의 폐쇄형 스칼라 양자화 전략을 도출하여, 상관 Gaussian 쌍에 대한 밀도 기반의 상전이를 밝히고 합성 데이터 및 LLM 활성화에서 검증한다.

ABSTRACT

We study entrywise scalar quantization of two matrices prior to multiplication. Given $A\in R^{m imes k}$ and $B\in R^{k imes n}$, we quantize entries of $A$ and $B$ independently using scalar quantizers with $K_X$ and $K_Y$ levels per entry, and form $\widehat C=\widehat A\,\widehat B$. The objective is to minimize the matrix multiplication mean-squared error (MSE) $E[\|{AB-\widehat A\widehat B}\|_F^2]$ under a pair-i.i.d.\ inner-product model. In the high-resolution regime $K_X,K_Y o\infty$, we derive a sharp $K^{-2}$ asymptotic expansion for $\mathcal{E}$, identify the exact optimal leading constants, and characterize asymptotically optimal quantization center densities in terms of conditional second moments. We then specialize to correlated Gaussian multiplicative pairs, obtaining a closed-form optimal point density \[ λ^\star(u)\ \propto\ \exp\!\left(-\frac{u^2}{6} ight)\bigl((1-ρ^2)+ρ^2u^2\bigr)^{1/3}, \qquad u=\frac{x}{σ_X}, \] with the same form for $y/σ_Y$, and prove a correlation-driven phase transition: the density is unimodal at the origin for $|ρ|\leq 1/\sqrt{3}$ and becomes bimodal for $|ρ|>1/\sqrt{3}$ with peaks at $u_{\mathrm{peak}}=\pm\sqrt{3-1/ρ^2}$. We show our method's applicability in synthetic experiments such as matrix multiplication quantization and least squares optimization, as well as quantization of large language model key and query activations.

연구 동기 및 목표

• 행 distortion이 입력 재구성이 아닌 곱 AB에서 측정되는 행렬 곱셈에 대한 양자화 설계를 동기를 부여한다.
• 쌍-i.i.d. 내적 모델 하에서 양자화된 A와 B에 대한 고속, 선도항 MSE 전개를 도출한다.
• 일반 분포에 대한 최적 컴팩싱 밀도와 명시적 상수들을 얻고, 상관 Gaussian 쌍에 대해 폐쇄형 해를 제시한다.
• 합성 실험 및 트랜스포머 모델의 활성화 양자화를 통해 방법의 실용성을 입증한다.

제안 방법

• 조건부 제2모멘트 w_X 및 w_Y에 의해 이끈 두 개의 가중 스칼라 MSE 문제로 행렬 곱셈 MSE를 축소한다.
• 고속 양자화기에 대한 E의 급격한 K^{-2} 점근적 스케일링을 보이고 최적의 점 밀도 lambda_X^* 및 lambda_Y^*를 도출한다.
• 상관 Gaussian 쌍으로 특수화하여 |rho|=1/sqrt(3)에서 위상이전이 있는 폐쇄형 lambda^*를 얻는다.
• 비율-형과 최적의 X와 Y 간 비트 분할을 표현하는 보손들(Corollaries)과 Gaussian 고속 상수 J(rho)를 계산한다.
• 합성 행렬 곱셈, 양자화된 최소제곱, 및 GPT-2와 Qwen3 모델의 활성화 양자화를 통해 이론을 검증한다.

Figure 1 : Optimal density phase transition. When $\rho=0$ , there is only a single mode. As $\rho$ increases, an additional mode emerges at the critical value $\left\lvert\rho\right\rvert=1/\sqrt{3}$ .

실험 결과

연구 질문

• RQ1 A와 B가 항목별로 양자화되었을 때, 제품 AB의 Frobenius MSE를 최소화하는 스칼라 양자화 스킴은 무엇인가?
• RQ2 쌍-i.i.d. 내적에서 최적 행렬 곱 MSE의 고속 점근행동과 선도 상수는 무엇인가?
• RQ3 조건부 제2모멘트가 A와 B의 최적 컴팩싱 밀도에 어떻게 영향을 미치는가?
• RQ4 상관 Gaussian의 경우 최적 밀도의 명시적 형태는 무엇이며 상관에 대해 위상이전이가 발생하는가?
• RQ5 제안된 양자화기가 양자화된 행렬 곱셈, 양자화된 최소제곱, 트랜스포머의 활성화 양자화와 같은 실용적 작업을 개선하는가?

주요 결과

• 행렬 곱셈 MSE는 E = mnk(I_X^3/(12 K_X^2) + I_Y^3/(12 K_Y^2)) + o(K^{-2})로 스케일링된다.
• 최적의 엔트리당 양자화기는 밀도 lambda_X^*(x) ∝ (f_X(x) w_X(x))^{1/3} 및 lambda_Y^*(y) ∝ (f_Y(y) w_Y(y))^{1/3}인 컴팩싱 양자화다.
• 상관 Gaussian 쌍의 경우, 폐쇄형 최적 밀도는 lambda^*(x) ∝ exp(-x^2/6sigma_X^2) ((1-ρ^2)+ρ^2 x^2/sigma_X^2)^{1/3}이며, |ρ|=1/√3에서 단모드에서 이분모드로의 상전이가 있다.
• 상전이는 작은 |ρ|에서는 하나의 모드를, 더 큰 |ρ|에서는 두 개의 대칭 모드를 생성하며 피크 위치는 u_peak = ±√(3 - 1/ρ^2)이다.
• 보손은 레이트 형식 표현과 최적의 비트 분할을 제공한다: 총 레이트가 고정되면 K_X와 K_Y는 (alpha_X/alpha_Y)^{1/4}에 따라 균형을 이룬다.

Figure 2 : Performance of our optimal quantizer vs. other commonly used quantizers.