Entropy-Constrained Adaptive Stochastic Quantization

Abstract: Adaptive stochastic quantization (ASQ) is a recently introduced quantization approach that optimizes the Mean Squared Error (MSE) for a given input while preserving unbiasedness. It is designed to alleviate the communication and memory bottlenecks of modern data and machine learning workloads, including model, gradient, and KV-cache compression and nearest-neighbor search.

摘要： 自适应随机量化（ASQ）是一种近期提出的量化方法，它在保持无偏性的同时，优化了给定输入的均方误差（MSE）。该方法旨在缓解现代数据和机器学习工作负载中的通信与内存瓶颈，包括模型、梯度、KV 缓存压缩以及最近邻搜索等应用场景。

Further, practical systems can then compress quantized data with a lossless entropy encoder. However, existing unbiased methods, including ASQ, choose their quantization values without considering this later encoding stage, leaving accuracy on the table.

此外，实际系统随后可以使用无损熵编码器对量化后的数据进行压缩。然而，现有的无偏方法（包括 ASQ）在选择量化值时并未考虑这一后续编码阶段，从而导致了精度上的损失。

We formulate the Entropy Constrained Adaptive Stochastic Quantization (ECASQ) problem, which jointly selects adaptive quantization values to minimize MSE under an entropy budget and an unbiasedness constraint. We give an optimal dynamic program with $O(sd^2)$ time and $O(d^2)$ space for a length-d vector and at most s quantization values, and a GPU-friendly approximate dynamic program with $O(sd^2)$ time and $O(d)$ space.

我们提出了熵约束自适应随机量化（ECASQ）问题，该问题旨在熵预算和无偏性约束下，联合选择自适应量化值以最小化 MSE。针对长度为 $d$ 的向量和最多 $s$ 个量化值的情况，我们给出了一个时间复杂度为 $O(sd^2)$、空间复杂度为 $O(d^2)$ 的最优动态规划算法，以及一个对 GPU 友好的、时间复杂度为 $O(sd^2)$、空间复杂度为 $O(d)$ 的近似动态规划算法。

The approximation guarantees that the solution has an MSE no larger than the optimal solution that uses one fewer bit of entropy per entry. We also provide an iterative refinement procedure for the approximation solution that, in our experiments, yields near-optimal results while retaining a substantial speed advantage over our solver for the optimal solution.

该近似算法保证了其解的 MSE 不会超过每个条目少使用 1 比特熵的最优解。我们还为该近似解提供了一种迭代细化过程；实验表明，该过程在保持对最优解求解器显著速度优势的同时，能够获得接近最优的结果。