2025-11-13T04:07:09.837900

Optimal Quantization for Matrix Multiplication

Ordentlich, Polyanskiy
Recent work in machine learning community proposed multiple methods for performing lossy compression (quantization) of large matrices. This quantization is important for accelerating matrix multiplication (main component of large language models), which is often bottlenecked by the speed of loading these matrices from memory. Unlike classical vector quantization and rate-distortion theory, the goal of these new compression algorithms is to be able to approximate not the matrices themselves, but their matrix product. Specifically, given a pair of real matrices $A,B$ an encoder (compressor) is applied to each of them independently producing descriptions with $R$ bits per entry. These representations subsequently are used by the decoder to estimate matrix product $A^\top B$. In this work, we provide a non-asymptotic lower bound on the mean squared error of this approximation (as a function of rate $R$) for the case of matrices $A,B$ with iid Gaussian entries. Algorithmically, we construct a universal quantizer based on nested lattices with an explicit guarantee of approximation error for any (non-random) pair of matrices $A$, $B$ in terms of only Frobenius norms $\|\bar{A}\|_F, \|\bar{B}\|_F$ and $\|\bar{A}^\top \bar{B}\|_F$, where $\bar{A},\bar{B}$ are versions of $A,B$ with zero-centered columns, respectively. For iid Gaussian matrices our quantizer achieves the lower bound and is, thus, asymptotically optimal. A practical low-complexity version of our quantizer achieves performance quite close to optimal. In addition, we derive rate-distortion function for matrix multiplication of iid Gaussian matrices, which exhibits an interesting phase-transition at $R\approx 0.906$ bit/entry, showing necessity of Johnson-Lindestrauss dimensionality reduction (sketching) in the low-rate regime.
academic

Optimal Quantization for Matrix Multiplication

基本信息

  • 论文ID: 2410.13780
  • 标题: Optimal Quantization for Matrix Multiplication
  • 作者: Or Ordentlich (Hebrew University of Jerusalem), Yury Polyanskiy (MIT)
  • 分类: cs.IT cs.AI cs.CL cs.LG math.IT
  • 发表时间: 2024年10月 (arXiv预印本)
  • 论文链接: https://arxiv.org/abs/2410.13780

摘要

本文针对大规模矩阵乘法的量化问题进行了深入研究。与传统的向量量化不同,该研究的目标不是近似矩阵本身,而是近似它们的矩阵乘积。给定两个实矩阵A、B,编码器对每个矩阵独立进行压缩,每个条目使用R比特描述,然后解码器利用这些压缩表示来估计矩阵乘积A⊤B。论文为具有独立同分布高斯条目的矩阵情况提供了近似均方误差的非渐近下界,构建了基于嵌套格的通用量化器,并发现了在R≈0.906比特/条目处的有趣相变现象,表明在低码率情况下需要Johnson-Lindenstrauss降维技术。

研究背景与动机

问题定义

随着深度神经网络和大语言模型的兴起,矩阵乘法成为计算的主要瓶颈。现代计算硬件往往受到内存带宽限制,而非计算能力限制。因此,对矩阵进行有损压缩以减少内存传输成为一个重要问题。

实际需求

对于大语言模型,作者估算了所需的量化率:

  • 在生成阶段,CPU需要1-3比特/条目的量化率才能充分利用计算资源
  • 在预填充阶段,对于在快速GPU上运行的小型LLM,需要约11.7比特/条目的量化率

现有方法局限性

  1. 经典向量量化:直接对矩阵A和B进行独立量化,然后计算量化矩阵的乘积,会导致O(n²)的误差
  2. 素描方法:虽然能提供无偏估计,但方差仍为O(n²)
  3. 确定性量化器:在球面上的向量存在Ω(n²)的下界

核心贡献

  1. 理论下界:为具有iid高斯条目的矩阵提供了矩阵乘法量化的非渐近下界
  2. 通用量化器:构建了基于嵌套格的通用量化器,对任意矩阵都有明确的误差保证
  3. 渐近最优性:证明了所提出的量化器对iid高斯矩阵达到下界,因此是渐近最优的
  4. 相变现象:发现了在R≈0.906比特/条目处的相变,揭示了低码率下降维的必要性
  5. 实用算法:提供了接近最优性能的低复杂度实现

方法详解

任务定义

给定矩阵A ∈ R^(n×a)和B ∈ R^(n×b),目标是设计编码器f₁: R^(n×a) → 2^(naR)和f₂: R^(n×b) → 2^(nbR)以及解码器g,使得:

EABg(f1(A),f2(B))F2E\|A^⊤B - g(f_1(A), f_2(B))\|_F^2

最小化,其中每个矩阵条目使用R比特描述。

核心函数Γ(R)

论文定义了关键的率失真函数:

1 - \left(1 - (2 \cdot 2^{-2R^*} - 2^{-4R^*})\right) \frac{R}{R^*} & R < R^* \\ 2 \cdot 2^{-2R} - 2^{-4R} & R \geq R^* \end{cases}$$ 其中R* ≈ 0.906是固定点方程R = ½log₂(1 + 4R ln 2)的解。 ### 嵌套格量化方案 #### 1. 内积量化(基础构建块) 对于球面上的ρ-相关向量U、V,使用嵌套格Λc ⊂ Λf进行量化: **编码过程**: - 对U和V分别添加独立的抖动向量Z₁、Z₂ - 量化到精细格Λf - 输出在粗糙格Λc中的陪集表示 **解码过程**: - 从陪集恢复量化点 - 移除抖动 - 计算内积估计 #### 2. 矩阵乘法量化 **预处理步骤**: 1. **零中心化**:计算Ā = A - (1/n)1·1^⊤A,B̄ = B - (1/n)1·1^⊤B 2. **范数量化**:高精度描述每列的均值和范数 3. **随机旋转**:应用相同的正交矩阵S到Ā和B̄ **量化步骤**: - 对旋转后的每列应用内积量化器 - 使用时间共享参数κ和MMSE缩放参数α ### 技术创新点 1. **抖动技术**:使量化误差独立于输入,避免了确定性量化器的O(n²)误差 2. **嵌套格结构**:提供有限码率的同时保持良好的量化性能 3. **时间共享**:在低码率下通过降维实现最优性能 4. **随机旋转**:将任意向量转换为球面均匀分布,便于分析 ## 实验设置 ### 理论验证实验 **数据生成**: - 矩阵A, B ∈ R^(n×n),条目为iid N(0,1) - n = 3 × 2¹¹ **实现细节**: - 基础格:D₃格(3维) - 嵌套比:q = 6 - 查找表大小:< 64KB(可放入L1缓存) - 有效码率:≈ 3.015 比特/符号 ### 对比方法 - 3比特标量量化器(ℓ∞归一化) - 理论下界Γ(R) ## 实验结果 ### 主要结果 **性能对比**: - 提出方法:1/n³ ∥Â⊤B - A⊤B∥²F ≈ 0.0593 - 3比特标量量化:≈ 0.1668(约3倍差距) - 理论下界:Γ(3.015) = 0.0304 **关键发现**: 1. 基于D₃格的方案显著优于标量量化 2. 性能接近理论最优(约2倍差距) 3. 与n的增长相比,性能差距会进一步缩小 ### 复杂度分析 **编码复杂度**:O(n log n)(使用快速Hadamard变换) **解码复杂度**:O(1)(使用查找表) **存储开销**:每个量化器需要O(log n)额外比特描述缩放因子 ## 相关工作 ### 随机线性代数 - **Monte Carlo矩阵乘法(MCMM)**:随机采样行进行近似 - **局部敏感哈希(LSH)**:用于余弦相似度的低维素描 - **局限性**:相对误差随∥A∥²F∥B∥²F/∥A⊤B∥²F增长 ### 神经网络量化 - **训练后量化**:OPTQ, GPTQ等方法 - **旋转技术**:QuIP, QuaRot使用Hadamard变换 - **格量化**:QUIP#使用E₈格进行权重量化 ### 信息理论 - **分布式压缩**:针对线性函数计算的压缩 - **码本设计**:Voronoi码和嵌套格码 ## 结论与讨论 ### 主要结论 1. **最优性**:对于iid高斯矩阵,提出的方案达到信息论下界 2. **通用性**:对任意矩阵都有明确的性能保证 3. **相变现象**:R* ≈ 0.906比特/条目是关键阈值 4. **实用性**:低复杂度实现接近理论性能 ### 局限性 1. **共享随机性**:需要编码器和解码器共享随机种子 2. **矩阵条件**:要求矩阵条目有界(M = n^(10^22000)) 3. **高维格**:理论最优需要高维"好"格,实际中使用低维格的乘积 4. **确定性方案**:不清楚是否存在不需要随机性的最优确定性方案 ### 未来方向 1. **多矩阵乘积**:扩展到k>2个矩阵的乘积 2. **其他距离度量**:KL散度等非MSE度量 3. **确定性量化器**:探索无需共享随机性的方案 4. **深度网络应用**:在实际LLM中的部署和优化 ## 深度评价 ### 优点 1. **理论严谨性**:提供了完整的信息论分析,包括上界和下界 2. **实用价值**:解决了LLM推理中的实际瓶颈问题 3. **技术创新**:巧妙结合了格量化、随机旋转和时间共享 4. **通用性**:不依赖于特定的矩阵分布假设 ### 不足 1. **复杂性**:理论分析较为复杂,实际实现需要多个技术组件 2. **常数因子**:虽然渐近最优,但有限样本下的常数可能较大 3. **硬件适配**:需要针对不同硬件平台优化实现 4. **扩展性**:从2个矩阵扩展到多个矩阵的乘积非平凡 ### 影响力 **理论贡献**: - 建立了矩阵乘法量化的信息论基础 - 揭示了相变现象和降维的必要性 - 为该领域提供了重要的理论基准 **实际应用**: - 为LLM量化提供了新的理论指导 - 后续工作NestQuant已在实际LLM上取得SOTA性能 - 为硬件加速器设计提供了理论依据 ### 适用场景 1. **大语言模型推理**:权重和激活的联合量化 2. **边缘计算**:内存受限环境下的矩阵运算 3. **分布式计算**:通信受限的矩阵乘法 4. **科学计算**:大规模数值线性代数问题 ## 参考文献 论文引用了44篇相关文献,涵盖了信息论、格理论、随机线性代数和神经网络量化等多个领域的重要工作。特别值得关注的包括: - Zamir的格编码专著提供了理论基础 - Erez和Zamir关于嵌套格的开创性工作 - 近期的LLM量化方法如OPTQ、QuIP等 - 随机矩阵理论和球面几何的相关结果 --- **总体评价**:这是一篇在理论和实践上都有重要贡献的优秀论文,为矩阵乘法量化问题提供了坚实的信息论基础,并展示了接近最优的实用算法。该工作对于理解和改进大规模机器学习系统中的量化技术具有重要意义。