[论文解读] Optimal Scalar Quantization for Matrix Multiplication: Closed-Form Density and Phase Transition

Calvin Ang, Sungyoon Kim+1 more|arXiv (Cornell University)|Mar 20, 2026

Stochastic Gradient Optimization Techniques被引用 0

一句话总结

这项工作推导出高速、闭式标量量化策略用于矩阵乘法，揭示相关高斯对的密度驱动相变并在合成数据和大语言模型激活上进行验证。

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 上衡量，而非输入重建上。
• 在成对独立同分布的内积模型下，推导量化的 A 与 B 的高速率、前导项 MSE 展开。
• 获得一般分布的最优压缩密度与显式常数，对相关高斯对给出闭式结果。
• 通过合成实验和变换器模型中的激活量化来展示方法的实用性。

提出的方法

• 将矩阵乘法的 MSE 降至由条件二阶矩 w_X 与 w_Y 驱动的两个加权纯量 MSE 问题。
• 在高码率压缩量化器下显示 E 的尖锐 K^{-2} 渐近标度并推导最优点密度 lambda_X^* 与 lambda_Y^*。
• 对相关高斯对进行特化，得到闭式的 lambda^*，在 |rho|=1/√3 出现相变。
• 给出推论，表达速率形式与 X、Y 的最优比特分配，并计算高斯高码率常数 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 在成对独立同分布内积下，最优矩阵乘法 MSE 的高码率渐近行为与前导常数是什么？
• RQ3 条件二阶矩如何塑造 A 与 B 的最优压缩密度？
• RQ4 在相关高斯情形下，最优密度的显式形式为何，是否存在关于相关性的相变？
• 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}。
• 对于相关高斯对，闭式最优密度为 lambda^*(x) ∝ exp(-x^2/(6σ_X^2)) ((1-ρ^2)+ρ^2 x^2/σ_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.