arXiv:2410.13780cs.ITcs.AI2024-10被引 20

为矩阵乘法设计最优量化方法,实现低比特下高精度计算。

Optimal Quantization for Matrix Multiplication

  • 基于嵌套格点构造通用量化器,仅依赖矩阵的Frobenius范数。
  • 对独立同分布高斯矩阵,达到理论误差下界,渐近最优。
  • 在低比特率(约0.906位/元素)以下需降维预处理,适合模型压缩场景。

机器学习中针对大矩阵的有损压缩(量化)研究日益增多,其目标是加速矩阵乘法——大型语言模型的核心瓶颈在于从内存加载数据的速度。不同于经典向量量化与率失真理论,新方法不追求逼近原矩阵本身,而是准确估计矩阵乘积 $A^ op B$。给定实矩阵 $A, B$,分别独立进行编码(压缩),每元素用 $R$ 比特表示,解码后重建乘积。本文针对 $A, B$ 具有独立同分布高斯分布的情况,给出了非渐近的均方误差下界。算法上,构造了一种基于嵌套格点的通用量化器,对任意非随机矩阵 $A, B$,其近似误差可由 $ar{A}, ar{B}$ 的列中心化版本的 Frobenius 范数 $ orm{ar{A}}_F, orm{ar{B}}_F$ 及 $ orm{ar{A}^ op ar{B}}_F$ 显式控制。对于高斯矩阵,该量化器达到理论下界,故渐近最优。此外,我们推导了高斯矩阵乘法的率失真函数,在 $R \approx 0.906$ bit/entry 处出现相变,表明在低速率时必须使用 Johnson-Lindestrauss 降维(映射)以维持精度。

原文摘要 · Abstract (English)

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.

量化矩阵乘法低比特格点编码

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。