💻 computer science

Entropy-Constrained Adaptive Stochastic Quantization

This paper introduces Entropy-Constrained Adaptive Stochastic Quantization (ECASQ), a novel framework that jointly optimizes adaptive quantization values to minimize Mean Squared Error under entropy and unbiasedness constraints, offering both an optimal dynamic programming solution and a highly efficient, GPU-friendly approximation with strong theoretical guarantees.

Original authors: Ran Ben Basat, Yaniv Ben-Itzhak, Michael Mitzenmacher, Shay Vargaftik

Published 2026-08-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ran Ben Basat, Yaniv Ben-Itzhak, Michael Mitzenmacher, Shay Vargaftik

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of modern computing, machines are constantly trying to learn from vast oceans of data. To do this, they perform billions of calculations, moving numbers back and forth between memory chips and processors. These numbers are usually stored with extreme precision, like a photograph with every possible shade of gray. However, this high precision comes at a steep cost: it requires massive amounts of memory and creates a bottleneck in the network, slowing down the entire learning process. To solve this, engineers use a technique called quantization. Think of it as simplifying a complex image into a limited palette of colors. Instead of keeping every tiny shade, the system rounds each number to the nearest value in a small, pre-defined list. This makes the data much smaller and faster to move. But there is a catch. If you simply round numbers to the nearest value, you introduce errors that can accumulate and ruin the machine's ability to learn. To prevent this, researchers use a method called stochastic quantization, which adds a layer of randomness. Instead of always rounding a number up or down, the system flips a coin based on how close the number is to the two available options. This ensures that, on average, the rounded numbers are perfectly accurate, even if individual numbers are not.

The challenge that researchers at VMware Research, University College London, and Harvard University recently tackled is how to make this rounding process even smarter when the data is about to be compressed further. In many practical systems, after the numbers are rounded, they are compressed using a lossless encoder, similar to how a ZIP file works. This encoder assigns shorter codes to values that appear frequently and longer codes to rare values. The goal is to minimize the total size of the data. Previous methods for rounding numbers were excellent at minimizing error, but they ignored the fact that some rounded values would be more common than others, leading to inefficient compression. Other methods that tried to optimize for compression size often sacrificed the crucial property of being unbiased, meaning the average of the rounded numbers would drift away from the true average, causing the machine learning model to fail. The researchers set out to find a way to do both at once: choose the best rounding values to minimize error while ensuring the resulting data compresses as efficiently as possible, all without losing that vital statistical accuracy.

The team developed a new approach called Entropy-Constrained Adaptive Stochastic Quantization. They treated the problem as a complex puzzle where they had to select a specific set of values to represent the data. The rules were strict: the set of values had to be small enough to keep the system fast, the average of the rounded numbers had to match the original numbers exactly, and the resulting pattern of values had to be compressible within a specific size limit. To solve this, they created a sophisticated mathematical strategy that works like a careful explorer mapping a terrain. They built a system that looks at the data and decides exactly which values to use, balancing the need for accuracy against the need for a small file size. They proved that their method could find the absolute best solution for this problem, but doing so required a massive amount of computer memory and time, making it impractical for very large datasets.

To make the solution usable in the real world, the researchers also designed a faster, approximate version. This version sacrifices a tiny bit of theoretical perfection for a massive gain in speed and efficiency. It works by making a slightly simpler assumption about how the data behaves, which allows it to run on standard graphics processors used in modern computers. They showed that this faster method produces results that are almost as good as the perfect solution, but it runs dozens of times faster. In their tests, they found that this new method significantly outperformed existing techniques. When they applied it to real-world data from large language models, the new approach reduced the error in the compressed data by a wide margin compared to older methods, while keeping the file sizes small. They also discovered that by running their fast approximation and then making a few small, targeted adjustments to the values, they could get results that were nearly indistinguishable from the perfect, slow solution, but in a fraction of the time.

The researchers were careful to note that their method does not work by magic or by guessing. It is a rigorous mathematical process that guarantees the data remains accurate on average. They also explored whether combining two different rounding strategies could yield even better results, a technique known as time-sharing. Their analysis showed that while this could theoretically help in some very specific edge cases, the single, optimized strategy they developed was sufficient for almost all practical situations. The work provides a new, highly efficient tool for anyone building large-scale machine learning systems. By solving the problem of how to round numbers for both accuracy and compression simultaneously, the team has removed a significant barrier to training and deploying powerful artificial intelligence models on limited hardware. The result is a system that can handle more data, move it faster, and learn more effectively, all without requiring a fundamental change in the underlying hardware.

Technical Summary: Entropy-Constrained Adaptive Stochastic Quantization (ECASQ)

1. Problem Formulation

The paper addresses the limitations of existing quantization methods in modern machine learning workloads, specifically regarding the trade-off between unbiasedness, adaptivity, and entropy efficiency.

Context: Adaptive Stochastic Quantization (ASQ) optimizes Mean Squared Error (MSE) for a given input while preserving unbiasedness ( $E[\hat{X}] = X$ ). However, standard ASQ selects quantization values without considering the subsequent lossless entropy encoding stage, which assigns shorter bit strings to frequent values.
The Gap: Existing entropy-aware methods (e.g., Entropy-Constrained Vector Quantization) do not enforce unbiasedness, while unbiased methods (e.g., QSGD, ASQ) are often non-adaptive or ignore entropy constraints.
The ECASQ Problem: The authors formulate the Entropy Constrained Adaptive Stochastic Quantization (ECASQ) problem. Given an input vector $X \in \mathbb{R}^d$ $X \in R^{d}$ , a permissible alphabet $P$ $P$ , a maximum codebook size $s$ $s$ , and an entropy budget $b$ $b$ , the goal is to select a codebook $Q \subseteq P$ $Q \subseteq P$ ( $|Q| \le s$ $∣ Q ∣ \leq s$ ) to minimize MSE subject to:
1. Unbiasedness: $E[\hat{X}] = X$ .
2. Entropy Constraint: The average-entry entropy $H(\hat{X}) \le b$ .

2. Methodology

2.1 Lagrangian Relaxation

The authors approach the hard-constrained problem by introducing a Lagrange multiplier $\lambda \ge 0$ , transforming the objective into minimizing:
$\text{MSE}(Q, X) + \lambda \cdot H(\hat{X})$
subject to $|Q| \le s$ and unbiasedness. By varying $\lambda$ , one can trace the Pareto-optimal trade-off curve between MSE and entropy. The paper notes that while a single $\lambda$ may not hit an exact entropy budget $b$ (due to the discrete nature of the problem), two-way time-sharing (running two quantizers on a random partition of the input) allows for the recovery of the optimal solution for any $b$ .

2.2 Optimal Dynamic Programming (DP)

The standard ASQ DP cannot be extended to ECASQ because the entropy of a quantization value depends on both its left and right neighbors, making the cost function non-separable.

State Definition: The authors define a new DP state $dp_{OPT}(i, \ell, j)$ representing the minimal cost for quantizing the first $\ell$ entries using $i-1$ quantization values (including $x_\ell$ ), plus the MSE for the interval $[x_\ell, x_j]$ assuming $x_\ell$ and $x_j$ are consecutive quantization values.
Complexity: A direct implementation requires $O(d^3 s)$ time and $O(d^2 s)$ space.
Accelerations:
- SMAWK Algorithm: Used to find row minima in $O(d)$ time by exploiting the quadrangle inequality (Monge property) of the cost matrix, reducing time to $O(d^2 s)$ .
- Hirschberg's Algorithm: Used to reduce space complexity to $O(d^2)$ by employing a divide-and-conquer strategy that only stores two layers of the DP table at a time.
Generalization: The algorithm extends to arbitrary permissible sets $P$ (size $p$ ) with complexity $O(p^2 s + d)$ time and $O(p^2 + d)$ space.

2.3 Space-Efficient Approximation Algorithm

To handle large inputs and permissible sets where $O(d^2)$ space is prohibitive, the authors propose an approximation that optimizes a surrogate entropy metric $H(I, B)$ (joint entropy of interval index and rounding bit) rather than the exact output entropy $H(\hat{X})$ .

Guarantee: The solution satisfies $H(\hat{X}) \le H(I, B) \le H(\hat{X}) + 1$ . Consequently, the approximation with entropy budget $b$ yields an MSE no larger than the optimal solution with budget $b-1$ .
Complexity: This approach reduces space to $O(d)$ (or $O(p+d)$ for general $P$ ) while maintaining $O(d^2 s)$ time. It is designed to be GPU-friendly.

2.4 Iterative Refinement

To bridge the accuracy gap between the approximation and the optimal solution, the authors introduce a refinement procedure. Starting from the approximate solution, it iteratively optimizes subsets of quantization values ( $Q_0, Q_1, Q_2$ ) in parallel using binary search within their local intervals. This process maintains unbiasedness and typically converges to near-optimal results within a few iterations.

3. Key Contributions

Problem Formulation: Defined ECASQ, the first framework to jointly optimize adaptive stochastic quantization for MSE, unbiasedness, and entropy constraints.
Optimal Algorithm: Developed an $O(d^2 s)$ time, $O(d^2)$ space dynamic programming algorithm for the exact solution, utilizing SMAWK and Hirschberg techniques.
Approximation with Guarantees: Proposed a space-efficient ( $O(d)$ ) approximation algorithm with a provable bound: the MSE is bounded by the optimal solution with one fewer bit of entropy.
Refinement Strategy: Introduced a fast, parallelizable iterative refinement that achieves near-optimal accuracy with significantly lower computational cost than the optimal solver.
Theoretical Completeness: Demonstrated that two-way time-sharing allows the Lagrangian relaxation to recover the optimal solution for any entropy budget, even when a single quantizer cannot achieve the exact constraint.

4. Experimental Results

The authors evaluated their methods on synthetic vectors (Normal, Lognormal, Student-t, Sparse outlier) and real-model tensors (weights, activations, KV-caches from models like Qwen2.5, Gemma, Llama, Phi).

Accuracy:
- The Refined Approx ECASQ reduces the mean excess vNMSE (relative to Optimal ECASQ) from ~4.97% (unrefined) to 1.37% on synthetic data.
- On real model tensors, the refined approximation consistently outperforms baselines (QSGD, QUIVER, Uniform + Entropy) across various bit-rates (1.5 to 4 bits/value).
Runtime:
- At vector size $d = 262,144$ , the Optimal ECASQ takes ~6.39 seconds.
- The Approx ECASQ takes ~0.26 seconds (24x speedup).
- The Refined Approx takes ~0.33 seconds (19x speedup).
Time-Sharing: Experiments indicate that the potential gain from time-sharing over a single optimal quantizer is small (median < 2%, max ~21% in narrow low-rate intervals), suggesting that a single well-tuned quantizer is often sufficient in practice.

5. Significance and Claims

The paper claims that ECASQ provides a principled solution to the "accuracy on the table" problem where existing unbiased quantizers ignore entropy encoding costs.

Practicality: The combination of the space-efficient approximation and the fast refinement procedure makes the approach viable for large-scale ML workloads (e.g., gradient compression, KV-cache reduction) where optimal DP is too slow or memory-intensive.
Generality: The framework supports arbitrary permissible alphabets (e.g., BF16 or FP32 representable values), allowing for hardware-aware constraints.
Modesty: The authors acknowledge that while time-sharing theoretically guarantees optimality for any entropy budget, the empirical gains are often marginal, and the single-quantizer approach (via Lagrangian optimization) is generally sufficient. They also note that the approximation's "one-bit" guarantee is a worst-case bound, with empirical results showing much tighter performance.