跳转至

文章背景与核心概要

在现代机器学习工作负载中，模型、梯度以及KV缓存的压缩对于缓解通信和内存瓶颈至关重要。自适应随机量化（ASQ）是一种旨在优化均方误差（MSE）同时保持无偏性的主流技术。然而，传统的ASQ方法在设计时并未考虑下游的无损熵编码阶段，导致压缩准确率未能达到最优。

为了弥补这一研究空白，本文正式提出了熵约束自适应随机量化（ECASQ）问题。该问题在同时满足熵预算和无偏性约束的前提下，联合优化量化值以最小化MSE。作者提出了一种运行时间为 \(O(sd^2)\)、空间复杂度为 \(O(d^2)\) 的最优动态规划算法（适用于长度为 \(d\) 的向量以及最多 \(s\) 个量化值），并进一步设计了对GPU友好的近似动态规划算法（时间复杂度 \(O(sd^2)\)，空间复杂度降至 \(O(d)\)），同时证明了其理论误差界。实验表明，文章提出的迭代精炼程序能够以显著的计算速度优势实现接近最优的MSE结果。

Entropy-Constrained Adaptive Stochastic Quantization

Summary

Adaptive Stochastic Quantization (ASQ) is a technique used to optimize Mean Squared Error (MSE) while maintaining unbiasedness, helping alleviate communication and memory bottlenecks in machine learning workloads (such as model, gradient, and KV-cache compression). However, traditional ASQ methods do not factor in downstream lossless entropy encoding, resulting in suboptimal compression accuracy.

To bridge this gap, this paper introduces the Entropy-Constrained Adaptive Stochastic Quantization (ECASQ) problem, which jointly optimizes quantization values to minimize MSE under both an entropy budget and an unbiasedness constraint. The authors propose: * An optimal dynamic programming algorithm running in \(O(sd^2)\) time and \(O(d^2)\) space for a length-\(d\) vector and up to \(s\) quantization values. * A GPU-friendly approximate dynamic program with \(O(sd^2)\) time and \(O(d)\) space, providing theoretical error bounds. * An iterative refinement procedure that achieves near-optimal MSE results with a significant speed advantage over the exact solver.

Metadata & Reference Information

  • arXiv Identifier: arXiv:2608.18147 [cs.LG]
  • Submission Date: August 10, 2026
  • Authors: Ran Ben Basat, Yaniv Ben-Itzhak, Michael Mitzenmacher, Shay Vargaftik
  • Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Data Structures and Algorithms (cs.DS); Information Theory (cs.IT)
  • DOI: 10.48550/arXiv.2608.18147

Abstract

自适应随机量化（ASQ）是一种近期引入的量化方法，它在保持无偏性的同时优化给定输入的均方误差（MSE）。该方法旨在缓解现代数据和机器学习工作负载中的通信与内存瓶颈，包括模型、梯度、KV缓存压缩以及最近邻搜索。此外，实际系统随后可以使用无损熵编码器压缩量化后的数据。然而，现有的无偏方法（包括ASQ）在选择量化值时并未考虑后续的编码阶段，从而损失了部分准确率。

我们正式建立了熵约束自适应随机量化（ECASQ）问题，该问题在熵预算和无偏性约束下联合选择自适应量化值以最小化MSE。针对长度为 \(d\) 的向量和至多 \(s\) 个量化值，我们提出了一种最优动态规划算法，其时间复杂度为 \(O(sd^2)\)，空间复杂度为 \(O(d^2)\)，以及一种对GPU友好的近似动态规划算法，其时间复杂度为 \(O(sd^2)\)，空间复杂度为 \(O(d)\)。该近似算法保证了解的MSE不会超过每个条目少使用一个比特熵的最优解。我们还为近似解提供了一种迭代精炼程序，在我们的实验中，该程序在保持对最优解求解器的大幅速度优势的同时，产出了近乎最优的结果。

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. 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.

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. 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.

Full-Text & External Links

全文与外部链接：