「从零开始学大模型」TurboQuant
TurboQuant 是最近的热门,Google 一将它发出来,各种自媒体都在疯传说这个技术将内存占用率压缩了 6 倍,内存价格闪崩。出于兴趣以及正好在学大模型,就看了一下,发现其实也没有想象的那么复杂。
简单来说,TurboQuant 做了一个 3bit 量化,将 FP16 降到 3bit (约5.3倍压缩),就是网传的“6 倍”。不过它与众不同的点有这两个:
- 是真·3bits,其他的压缩方法往往还要存额外的信息,说是4bit压缩实际上均摊下来要5~6bit
- 在仅用 3bits 的情况下,很好的保证了量化精度
要看懂 TurboQuant,主要是需要 PolarQuant 和 QJL 这两个压缩方法。
PolarQuant
传统量化痛点
在对 KV Cache 进行量化的时候,传统的量化方法(比如把 FP16 映射到 INT4)都是在笛卡尔坐标下进行的,但这种方法面临两个致命问题:
- 异常值(Outliers)灾难:尤其是 Key 向量在经过旋转位置编码 (RoPE) 后,某些维度会出现极大的异常值。按通道量化时,这些异常值会把整个量化范围撑爆,导致正常值的精度严重丢失。
- 归一化开销(Normalization Overhead):传统量化必须为每个数据块存储量化参数(Scale 和 Zero-point)。在极低比特(比如 3-bit 甚至更低)量化时,这些额外参数占用的内存比例高得吓人,会吃掉量化带来的一大半显存收益。
PolarQuant 的计算逻辑
PolarQuant 很聪明的换了个角度来解决这个问题,抛弃了传统的笛卡尔坐标,转而用极坐标来对向量进行压缩。它的工作流程非常简单却巧妙:
- 随机预处理 (Random Preconditioning): 首先对高维向量施加一个随机旋转矩阵(正交矩阵)。这一步不会改变向量之间的内积(即保留了 Attention 计算所需的相似度信息),但打散了原始特征,将异常值的能量分摊到了各个维度。
- 极坐标变换: 将向量从笛卡尔坐标系转换为极坐标系,即表示为一个半径 (Radius) $r$ 和一组极角 (Polar Angles) $\theta_1, \theta_2, \dots$
- 极度浓缩的角度分布: 最神奇的数学现象出现了。经过随机预处理后,这些极角 $\theta$ 在极坐标空间中会呈现出极其集中、且可以用解析式精确计算的完美先验分布(高维空间中的高斯分布是极其集中的)
随后就可以用非均匀 codebook,进行量化,实现 zero-overhead 的同时很好地处理了异常值。
PolarQuant 的牛逼之处
私以为最厉害的地方是利用低开销的随机旋转矩阵乘法与坐标系变换,将数据分布强行变成一个已知的、非常良好的分布。因此量化的时候可以直接根据这个分布,离线计算好量化表(一个非均匀 codebook),而不需要在实际的量化后数据里额外存任何东西。用数学的魔法实现了极其高效的“时间换空间”。
QJL
QJL 是另一个非常暴力的量化方式,它直接生硬的将数据降维到 1bit,但借助 JL 引理,保留了足够的数据信息。
Johnson-Lindenstrauss (JL) 引理是高维几何和降维领域的一个核心数学定理。简单来说,它告诉我们:我们可以把高维空间中的数据点投影到一个低维空间中,同时几乎完美地保留这些点之间的相对距离。
JL 引理的核心定义
假设你有一组包含 $n$ 个点的数据集 $V$,存在于一个极高维度的空间 $\mathbb{R}^d$ 中。JL 引理证明,对于任意给定的误差容忍度 $0 < \epsilon < 1$,只要目标维度 $k$ 满足:
$$k \ge \frac{8 \ln n}{\epsilon^2}$$即 $k = \mathcal{O}(\epsilon^{-2} \log n)$,就必然存在一个线性映射 $f: \mathbb{R}^d \rightarrow \mathbb{R}^k$,使得对于数据集 $V$ 中的任意两个点 $u, v$,它们在低维空间中的距离与在高维空间中的距离满足以下关系:
$$(1-\epsilon)\|u-v\|^2 \le \|f(u)-f(v)\|^2 \le (1+\epsilon)\|u-v\|^2$$最反直觉的结论是: 目标维度 $k$ 只和点数 $n$ 以及容忍误差 $\epsilon$ 有关,与原始维度 $d$ 完全无关! 这意味着你可以把十万维、百万维的向量压缩到几千维,而距离信息的损失依然被严格控制在 $\epsilon$ 范围内。
如何实现 JL 映射?
数学证明中给出的构造方法非常简单暴力:随机投影 (Random Projection)。
你只需要构造一个随机矩阵 $P \in \mathbb{R}^{k \times d}$,其矩阵元素独立采样自均值为 0、方差为 $1/k$ 的高斯分布 $\mathcal{N}(0, 1/k)$(在工程中,为了计算更快,常简化为从 $\{-1, +1\}$ 中等概率采样的 Rademacher 分布)。
对于高维空间中的任意向量 $x \in \mathbb{R}^d$,它在低维空间的表示就是一次简单的矩阵乘法:
$$y = Px$$从 JL 进阶到 QJL (Quantized JL)
在现代 AI 基础设施(如 LLM 推理引擎、向量数据库)中,遇到海量高维向量时,仅仅“降维”是不够的,我们还需要极大地减少内存占用 (Memory footprint) 和内存带宽消耗。QJL 就是将“JL 随机投影”与“极低比特量化”结合起来的产物。
QJL 的工作流程通常分为两步:
- 降维 (Random Projection): 使用上述的随机矩阵 $P$,将原始向量 $x$ 降维到 $y = Px$。
- 量化 (Quantization): 对降维后的向量 $y$ 进行极其激进的量化,最典型的是 1-bit 量化(仅保留符号):
$$q = \text{sgn}(Px)$$ 此时,原本由 32-bit 浮点数构成的向量,变成了只有 $+1$ 和 $-1$(在计算机中存储为 $1$ 和 $0$)的布尔向量 $q \in \{0, 1\}^k$。
为什么 QJL 有效且在工程中极具价值?
尽管 1-bit 截断看似丢失了大量信息,但由于 JL 引理保证了随机投影后距离的同构性,以及高维空间向量的正交性质,两个 QJL 压缩后的 1-bit 向量的汉明距离 (Hamming Distance),依然能够高度近似它们在原始高维空间中的余弦相似度 (Cosine Similarity)。
这在底层系统开发中带来了巨大的红利:
- 极致的内存压缩: 假设原始向量是 4096 维的 Float16(占用 8192 Bytes)。经过 QJL 映射到 1024 维并做 1-bit 量化后,只需要 1024 bits,即 128 Bytes。压缩比高达 64 倍,这对于缓解 Memory Wall 极其有效。
- 计算指令级加速: 在计算相似度时,原本需要用 GPU/CPU 计算浮点向量内积(FMA指令),现在变成了计算两个位串的异或 (XOR) 和统计 1 的个数 (POPCNT)。在底层架构中,
POPCNT是一条极快的硬件指令,这使得距离计算的速度能有数量级的提升。
TurboQuant
如果理解了 PolarQuant 与 QJL,TurboQuant 就简单了。TurboQuant 的计算逻辑就两步:
- 用 PolarQuant 进行主压缩(2bit)
- 用 QJL 进行残差纠正(1bit),让 attention score 更准
感觉这两步其实某种意义上是类似的,都包含了一个核心思想:“用随机化消除结构依赖”。PolarQuant 里面用随机旋转将向量从原有的分布直接扭转成了哈希分布,JL变换里面也是使用随机矩阵在无需预知数据分布的情况下就进行有损压缩,但保留距离信息。用"随机性"把复杂的、不规则的问题变成简单的、可预测的问题。
下课!
加餐时间 - RHT变换
为了同时优化“存储开销”和“计算延迟”,工程师们通常会使用随机哈达玛变换(Randomized Hadamard Transform, RHT)来代替普通的随机高斯矩阵:
- 计算无需存矩阵: 经典的沃尔什-哈达玛矩阵(Walsh-Hadamard Matrix)是一个高度对称的矩阵,元素只有 $+1$ 和 $-1$。因为其特殊的递归结构,乘以这个矩阵不需要做标准的矩阵乘法($O(d^2)$),而是可以像快速傅里叶变换(FFT)那样,用一种被称为 快速沃尔什变换(FWT) 的算法来计算。FWT 的时间复杂度只有 $O(d \log d)$,而且不需要在显存中实体化这个矩阵。
- 唯一需要存的只是一个“对角线符号向量”: RHT 的完整形式是 $H \times D$,其中 $H$ 是隐式的哈达玛矩阵,$D$ 是一个对角矩阵(对角线上随机分布 $+1$ 或 $-1$)。这意味着,为了实现“随机预处理”,系统只需要在显存里存一个长度为 $d$ 的随机符号向量(甚至可以压缩到几个 Byte)就可以了
参考资料
TurboQuant: Redefining AI efficiency with extreme compression