跳转至

Block Sphere Vector Quantization

会议: NeurIPS 2026
arXiv: 2605.19972
OpenReview: https://openreview.net/forum?id=U13XD8SxhN
代码: https://github.com/sang3798/BlockQuant
领域: 模型压缩
关键词: 矢量量化, 随机旋转, 球面几何, 块量化, KV缓存压缩

一句话总结

针对基于随机旋转的向量量化方法多依赖逐坐标标量码本而割裂球面几何结构的问题,本文统一分析了 EDEN、RabitQ 与 TurboQuant 的理论边界,并提出基于真实超球边缘分布构造质心的块球面量化方法(BlockQuant),在重构均方误差与内积失真上均严格超越现有基准。

研究背景与动机

在大规模机器学习系统中,高维向量的高效存储、传输与相似度检索构成了核心性能瓶颈。从数十亿规模的向量检索数据库、分布式与联邦学习中的高频梯度同步,到大语言模型(LLM)长文本推理中随序列长度与批次大小急剧膨胀的 KV 缓存,低比特向量量化(VQ)技术均是削减显存占用与内存带宽开销的基石。近期以 EDEN、RabitQ 和 TurboQuant 为代表的旋转量化方法受到广泛关注,它们利用 Haar 随机旋转将任意输入单位向量各向同性化,进而在低比特极端压缩场景下实现了极强的实证表现与理论保障。

然而,现有旋转量化方法之间的横向对比长期处于碎片化状态:不同方法依赖相互割裂的优化目标与理论假设,例如 TurboQuant 和 EDEN 主要依托均方误差(MSE)与内积期望失真进行刻画,而 RabitQ 则建立在高概率比特复杂度界限之上,使得学术界对各方法的相对优势认知模糊。更为关键的是,EDEN 和 TurboQuant 在旋转后均采用逐坐标独立的标量量化器,由一维标量质心构成的笛卡尔乘积码本中,有相当比例的码字远离单位超球面,导致在 \(2^{bd}\) 的庞大编码空间中大量比特被浪费在远离真实输入流形的无效区域。

本文首先在统一的理论框架下厘清了现有旋转量化方案的性能边界,揭示了各自在不同准则下的权衡取舍;随后跳出逐坐标标量量化的传统思路,从旋转单位向量的真实超球几何先验出发重构多维块质心。核心 idea:利用单位球面随机旋转诱导的精确低维欧氏球边缘概率分布构造离线块码本(Block-Sphere Quantization),从几何流形层面消除逐坐标量化的结构性失配,在保持极低在线计算复杂度的同时使重构误差与内积失真逼近理论香农下界。

方法详解

整体框架

BlockQuant 的核心流程由两大部分组成:基于超球边缘密度的离线码本预构建,以及基于笛卡尔查找表(LUT)的在线旋转块量化与目标感知反量化。给定待量化的 \(d\) 维单位输入向量 \(\mathbf{x} \in \mathbb{S}^{d-1}\),算法首先通过 Haar 分布生成的正交随机旋转矩阵 \(R \in \mathbb{R}^{d \times d}\) 将其映射为各向同性旋转向量 \(\mathbf{z} = R\mathbf{x}\)。随后,算法不进行逐坐标标量切分,而是将 \(\mathbf{z}\) 划分为 \(m = d/p\) 个 \(p\) 维局部向量块 \(\mathbf{z}_j \in \mathbb{B}^p\)。每个数据块通过预计算的 \(p\) 维球面边缘码本独立检索最近邻质心并生成 \(bp\) 比特的量化索引。

在反量化阶段,解码端根据索引拼接出旋转空间的重构块向量 \(\bar{\mathbf{z}}\),施加逆旋转矩阵 \(R^\top\) 映射回原始坐标系得到未缩放重构向量 \(\bar{\mathbf{x}} = R^\top \bar{\mathbf{z}}\)。最后,算法根据具体的下游任务目标(重构误差最小化、内积无偏估计或极速原始重构)匹配闭式标量重缩放系数 \(S\),得到最终输出 \(\hat{\mathbf{x}} = S \cdot \bar{\mathbf{x}}\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入单位向量<br/>x ∈ S^(d-1)"] --> B["超球先验与离线块码本构建<br/>基于球边缘分布预计算 p 维质心"]
    B --> C["Haar 随机正交旋转<br/>z = R x 映射为均匀超球分布"]
    C --> D["低维块空间划分<br/>拆分为 m = d/p 个 p 维块 z_j"]
    D --> E["查找表辅助近似最近邻量化<br/>Cartesian LUT 加速索引检索"]
    E --> F["逆旋转与目标感知标量重缩放<br/>根据 MSE 或 IP 目标自适应缩放"]
    F --> G["量化紧凑表征输出<br/>用于近似检索与 KV 缓存推理"]

关键设计

1. 超球先验与离线块码本构建:依据精确球边缘分布求解块最优质心

针对传统逐坐标量化忽略坐标间耦合导致码字脱离超球面流形的缺陷,BlockQuant 将维度划分为大小为 \(p\) 的子块,直接建模子块在 \(p\) 维欧氏球 \(\mathbb{B}^p\) 内的边缘分布。当 \(d\) 维向量 \(\mathbf{x} \sim \text{Unif}(\mathbb{S}^{d-1})\) 时,任意 \(p\) 维子块 \(\mathbf{z}_j\) 的边缘概率密度函数具备如下解析闭式解: $$ f_{p,d}(\mathbf{z}_j) = \frac{\Gamma(d/2)}{\pi^{p/2}\Gamma((d-p)/2)}\left(1 - |\mathbf{z}j|2^2\right)^{\frac{d-p-2}{2}}, \quad \mathbf{z}j \in \mathbb{B}^p $$ 该分布具有完美的极坐标分解性质:\(\mathbf{z}_j = r_j \boldsymbol{\theta}_j\),其中方向向量 \(\boldsymbol{\theta}_j \sim \text{Unif}(\mathbb{S}^{p-1})\),径向模长平方服从 Beta 分布 \(r_j^2 \sim \text{Beta}(p/2, (d-p)/2)\),且两者相互独立。基于此解析密度,BlockQuant 将码本设计转化为连续概率空间下的 \(K\)-means 优化问题: $$ \mini} \int}^{2^{bp}} \min}^p_i|}]} |\mathbf{u} - \mathbf{o2^2 \, f $$ 由于该分布完全由几何先验确定而独立于具体数据分布,优化出的 }(\mathbf{u}) \, d\mathbf{u\(2^{bp}\) 个 \(p\) 维质心只需离线求解一次,即可在各类异构数据集与推理任务中通用复用。

2. 查找表辅助近似最近邻量化:用笛卡尔网格剪枝化解高维码本检索开销

当块大小 \(p\) 或每维度比特数 \(b\) 增加时,每个块的候选质心总数 \(K = 2^{bp}\) 呈指数级增长(例如 \(p=3, b=4\) 时 \(K=4096\)),导致暴力全码本距离搜索复杂度 \(O(m K p)\) 成为在线量化的吞吐瓶颈。为此,BlockQuant 引入了基于空间网格的轻量级查找表(Cartesian LUT)机制。在离线阶段,将 \(p\) 维块空间划分为各轴分辨率为 \(L\) 的超立方体网格,为每个网格中心预先索引距离最近的 \(k\) 个候选质心集 \(\mathcal{C}_{\text{top}k}(u)\)。在实际量化时,输入块 \(\mathbf{z}_j\) 首先根据坐标快速定位所在网格单元,随后仅在该网格的 \(k\) 个候选子集中执行局部最近邻搜索。该机制将在线距离计算复杂度从 \(O(m K p)\) 骤降至 \(O(m k p)\),消除了延迟瓶颈,且实验表明该近似对最终重构质量的影响几乎可忽略。

3. 逆旋转与目标感知标量重缩放:基于双重校准因子解耦不同任务失真

在完成逆正交变换 \(\bar{\mathbf{x}} = R^\top \bar{\mathbf{z}}\) 后,未缩放重构向量 \(\bar{\mathbf{x}}\) 的模长往往收缩。为了在不同下游任务中达成理论最优,BlockQuant 根据内积对齐系数 \(\rho = \langle \mathbf{x}, \bar{\mathbf{x}} \rangle = \langle \mathbf{z}, \bar{\mathbf{z}} \rangle\) 派生了三种自适应重缩放规则: $$ S = \begin{cases} \dfrac{\rho}{|\bar{\mathbf{x}}|2^2}, & \text{针对重构均方误差最小化(BlockQuant} \[10pt] \dfrac{1}{\rho}, & \text{针对内积无偏估计(BlockQuant}}}\text{){\text{UB}}\text{)} \[10pt] 1, & \text{针对追求极低延迟的原始重构(BlockQuant} \end{cases} $$ 理论证明表明,对于内积估计,利用比例缩放因子 }}\text{)\(1/\rho\) 能精确消除残差在平行于原向量方向的径向偏移,使期望内积估计误差严格受控于垂直方向的正交扰动 \(\frac{1-\langle\mathbf{x},\mathbf{y}\rangle^2}{d-1} \mathbb{E}\left[\frac{\|\bar{\mathbf{x}}\|_2^2 - \rho^2}{\rho^2}\right]\),从而在不耗费额外校准比特的情况下逼近理论极限。

实验关键数据

主实验

论文首先在统一理论框架下给出了高维渐近下各旋转量化方法的理论失真常数对比,展现了 BlockQuant 在各类位宽下的理论领先优势。

量化算法 重构 MSE (\(b=2\)) 重构 MSE (\(b=3\)) 重构 MSE (\(b=4\)) 内积期望失真 (\(b=2\)) 内积期望失真 (\(b=4\))
EDEN 0.117 0.0345 0.0095 \(0.133 / (d-1)\) \(0.0096 / (d-1)\)
RabitQ 0.119 0.0374 0.0115 \(0.135 / (d-1)\) \(0.0117 / (d-1)\)
TurboQuant 0.117 0.0300 0.0090 \(0.560 / d\) \(0.0470 / d\)
BlockQuant (\(p=2\)) 0.108 0.0297 0.0078 \(0.120 / (d-1)\) \(0.0078 / (d-1)\)
BlockQuant (\(p=3\)) 0.101 0.0271 0.0071 \(0.113 / (d-1)\) \(0.0071 / (d-1)\)

在真实落地场景中,作者将 BlockQuant 应用于 Llama-3.1-8B-Instruct 的长文本 KV 缓存量化(Key 采用 3.5-bit 混合精度,Value 采用 2-bit 量化),在 LongBench-E 长文本基准上评估全任务表现(5 次随机种子均值与标准差):

方法 SingleQA MultiQA Summarization Few-shot Synthetic Code 综合平均
RabitQ 19.12 (0.54) 15.88 (0.26) 29.40 (0.11) 68.23 (0.10) 55.75 (0.35) 60.35 (0.13) 43.52 (0.15)
TurboQuant 18.60 (0.39) 15.79 (0.15) 29.11 (0.11) 68.19 (0.10) 55.31 (0.81) 59.72 (0.49) 43.20 (0.20)
EDEN 19.02 (0.27) 16.30 (0.26) 29.42 (0.17) 68.46 (0.12) 56.29 (0.22) 61.41 (0.21) 43.87 (0.07)
BlockQuant (\(p=3\)) 19.55 (0.46) 16.31 (0.15) 29.72 (0.14) 68.49 (0.16) 56.48 (0.57) 61.42 (0.34) 44.03 (0.10)
Full Cache (16-bit) 19.53 16.54 30.28 68.41 56.41 61.60 44.15

消融实验

论文分析了渐近高比特区间下块大小 \(p\) 对重构 MSE 前导常数的影响,并与纠正后的香农失真下界进行对照:

块大小配置 高比特 MSE 渐近界 高比特内积失真渐近界 说明
\(p=1\) (TurboQuant/EDEN) \(2.721 \cdot 4^{-b}\) \(2.721 \cdot 4^{-b} / (d-1)\) 传统逐坐标量化
\(p=2\) (BlockQuant) \(2.015 \cdot 4^{-b}\) \(2.015 \cdot 4^{-b} / (d-1)\) 2维平面超球投影
\(p=3\) (BlockQuant) \(1.770 \cdot 4^{-b}\) \(1.770 \cdot 4^{-b} / (d-1)\) 3维空间超球投影
\(p=d\) (理想极限) \(C_d \cdot 4^{-\frac{bd}{d-1}}\) (\(C_d \to 1.0\)) – 逼近香农球面下界 \(c_d \cdot 4^{-\frac{bd}{d-1}}\) (\(c_d \to 1.0\))

关键发现

  • 块维度增大单调逼近香农极限:随着块大小从 \(p=1\) 扩展到 \(p=3\),重构失真界的前导常数从 2.721 显著下降至 1.770;在理想情形 \(p=d\) 下,上界前导系数趋近于 1.0,与基于球体内在维度 \(d-1\) 建立的严格香农下界完全吻合,证明了块球面建模的理论完备性。
  • 内积失真相比 TurboQuant 取得数倍优势:TurboQuantPROD 因专门划分 1 比特用于 QJL 残差纠偏,导致主量化比特预算受损;而 BlockQuantUB 借由旋转对称性与比例重缩放因子,实现了真正无偏估计,内积均方误差比 TurboQuant 降低了近 4 倍。
  • 超长文本大海捞针性能几乎无损:在 8B 大模型的 Needle-In-A-Haystack 测试中,全精度 KV 缓存得分为 0.965,BlockQuant 达到 0.962,显著高于 EDEN(0.953)、RabitQ(0.944)和 TurboQuant(0.933),表明其能有效防止长上下文注意力累积漂移。

亮点与洞察

  • 将数据无关的球面几何先验融入乘积量化:传统乘积量化(PQ)高度依赖基于训练集的子空间聚类,在动态生成的 KV 缓存场景中不可行;BlockQuant 证明了只要引入随机旋转,低维块向量在欧氏球内的边缘分布即具备解析解,实现了“零训练样本”的免调优最优码本生成。
  • 笛卡尔查找表消解多维码本的计算诅咒:巧妙利用欧氏球空间的连续局部性,将多维高位宽码本检索转换为网格候选子集查询,以极小的离线内存表代价换取了实时量化的吞吐可行性。
  • 澄清了关于旋转量化评估标准的认知分歧:指出了先前关于 TurboQuant 与 RabitQ 优劣论争的根源在于评估准则的错位(期望准则 vs 高概率准则),为领域后续的研究建立了严谨公平的参照系。

局限与展望

  • 块尺寸增大受限于维度可除性与离线表膨胀:算法要求块大小 \(p\) 必须整除向量维度 \(d\);此外当 \(p > 3\) 且位宽 \(b \ge 4\) 时,笛卡尔网格的单元数呈指数增长,使得查找表显存占用和检索代价激增,限制了向更大块大小拓展。
  • 对局部异构异常通道的处理仍需外挂启发式设计:在 LLM KV 缓存量化中,BlockQuant 依然需要沿用将高范数异常通道单独隔离保留高比特(如 4-bit)的混合位宽配置,尚未将通道离群值完全内化到统一度量流形中。
  • 硬件友好的定点化核函数待完善:当前反量化依赖基于浮点内积的缩放因子计算,后续需进一步设计高度适配 GPU Tensor Core 的专用整数点积与反量化融合算子。

相关工作与启发

  • vs EDEN: EDEN 基于标准正态分布近似构建标量 Lloyd-Max 码本并采用内积重缩放,而 BlockQuant 将坐标独立假设拓展为多维块球面边缘分布,在重构与内积两项指标上均取得了严格更优的前导常数。
  • vs TurboQuant: TurboQuant 使用精确一维超球坐标边缘分布,但在内积估计时消耗 1 比特给 QJL 纠偏导致常数偏大;BlockQuant 统一在块级别优化几何质心,以比例估计量实现了无损比特的无偏内积估计。
  • vs RabitQ: RabitQ 依赖超球投影的均匀网格量化,具备更强的高概率界限控制,但在期望 MSE 和低位宽内积精度上略逊于考虑精确边缘密度的 BlockQuant。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从随机旋转超球诱导的精确块多维边缘分布出发重构量化码本,打破了长期的逐坐标标量思维定势。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖完备的数学定理推导、渐近香农下界分析、合成与真实向量检索评测,以及完整的 8B 模型长上下文 LLM 推理验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密自洽,理论推导详实清晰,附录完备且实验复现设置明确。
  • 价值: ⭐⭐⭐⭐⭐ 为旋转向量量化领域确立了统一的理论桥梁,并在实际 LLM KV 缓存无损压缩中展现出显著的实用落地潜力。