Entropy-Constrained Adaptive Stochastic Quantization

Ran Ben Basat; Yaniv Ben-Itzhak; Michael Mitzenmacher; Shay Vargaftik

cs.LGAug 10, 2026

Entropy-Constrained Adaptive Stochastic Quantization

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

Organizations: University College London · VMware Research by Broadcom · Harvard University

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

Figures & tables

Figure 1: The structure of $dp_{OPT}(i,\ell,j)$ : we look for the $k$ that minimizes the cost of the prefix up to $x_{\ell}$ and only the MSE for $[x_{\ell},x_{j}]$ .

Figure 2: Three-color refinement of a quantizer. (a) The quantization values in

Q

are partitioned by index into

Q_{0}

(blue),

Q_{1}

(red), and

Q_{2}

(green). (b) During the

Q_{1}

step, each red value is optimized over its range while neighboring values remain fixed. Consecutive red values have disjoint ranges, so all values in

Q_{1}

are updated in parallel; the other two substeps are analogous.

Figure 3: Matched-rate excess vNMSE (five seeds) relative to Optimal ECASQ on synthetic BF16 vectors (

d=16{,}384

, at most

s=64

quantization values). Lower is better.

Figure 4: Runtime for BF16 inputs,

s=64

, and

\lambda=0.1

. Marks average the four distribution-specific means (five seeds each), and error bars span their min-max range.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Matched-rate vNMSE over 25 BF16 real-model weight tensors. Subfigures show

1.5

2

2.5

3

3.5

, and

4

bits/value in reading order. Filled circles are individual tensors and the larger white-centered circle is their arithmetic mean. Every full subfigure uses a linear horizontal axis sized to its observed range; its inset magnifies Optimal ECASQ, Approx ECASQ with refinement, Approx ECASQ, and uniform quantization with entropy coding on a common linear scale within this figure. Values are relative to the best observed method for the same tensor and total rate, and no method is extrapolated.

Figure 6: Matched-rate vNMSE over 45 BF16 real-model activation tensors. Subfigure order and axis conventions match Figure 5 . Filled squares are individual tensors and the larger white-centered square is their arithmetic mean. Each inset uses a common linear scale within this figure and magnifies the three ECASQ variants together with uniform quantization with entropy coding.

Figure 7: Matched-rate vNMSE over 30 BF16 real-model KV-cache tensors. Subfigure order and axis conventions match Figure 5 . Filled triangles are individual tensors and the larger white-centered triangle is their arithmetic mean. Each inset uses a common linear scale within this figure and magnifies the three ECASQ variants together with uniform quantization with entropy coding.

Figure 8: Matched-rate comparison on paired uniform permissible grids. The subfigures show weights, activations, and KV-cache tensors. Curves report the arithmetic mean of each tensor’s excess vNMSE relative to Optimal ECASQ at the same total-rate budget. Approx ECASQ with refinement stays close to Optimal ECASQ even when all methods use the stored grids selected by the uniform baseline; lower is better.

Figure 9: Synthetic matched-rate comparison at six total rates, using the same method colors and distribution markers as Figure 3 . Circles, squares, triangles, and diamonds denote Normal, Lognormal, Student-

t

, and sparse-outlier distributions, respectively. Insets magnify the first three rows on a shared

0

–

25\%

scale. Total rate includes output entropy and the 32-bit per-level metadata. Smaller values nearer zero are better.

Panel of Figure 10 — Figure 10: Approximation diagnostics over the 20 distribution–seed cases at each budget. Gray points show individual cases. Orange circles show the arithmetic mean in subfigure (a) and the geometric mean in subfigure (b). The dashed horizontal line marks the proved upper bound of one. Subfigure (a) reports the bits saved by exact output-entropy accounting for the same Approx ECASQ codebook selected under surrogate budget $b$ , at unchanged vNMSE: $H(I,B)-H(\widehat{X})=H(B\mid\widehat{X})$ . Subfigure (b) reports the vNMSE ratio between the surrogate-feasible Approx solution at budget $b$ and Optimal ECASQ at budget $b-1$ . Both subfigures linearly interpolate the lower convex envelopes of the measured Lagrangian sweeps, corresponding to whole-vector time sharing between neighboring retained frontier points.

Panel of Figure 11 — Figure 11: Refinement from four initializations. Subfigure (a) shows the median objective gap to Optimal ECASQ. Subfigure (b) shows median end-to-end runtime, including initialization and all executed refinement rounds; requested rounds are an upper bound because refinement stops after a round without objective improvement. To reflect nondecreasing work despite run-to-run timing noise, each trajectory uses its running-maximum time before aggregation.

Figure 12: Runtime and objective scaling for BF16 inputs with at most

s=16

quantization values and fixed

\lambda=0.1

. Points show the mean and two-sided 95% Student-

t

confidence interval over five seeds. The lower subfigures report

\mathrm{vNMSE}(Q,X)+\lambda H(\widehat{X})

. All runtime subfigures share one y-axis range, and all objective subfigures share another.

Figure 13: Runtime and objective scaling for BF16 inputs with at most

s=64

quantization values and fixed

\lambda=0.1

. Points show the mean and two-sided 95% Student-

t

confidence interval over five seeds. The lower subfigures report

\mathrm{vNMSE}(Q,X)+\lambda H(\widehat{X})

. All runtime subfigures share one y-axis range, and all objective subfigures share another.

Figure 14: Runtime and objective scaling for BF16 inputs with at most

s=256

quantization values and fixed

\lambda=0.1

. Points show the mean and two-sided 95% Student-

t

confidence interval over five seeds. The lower subfigures report

\mathrm{vNMSE}(Q,X)+\lambda H(\widehat{X})

. All runtime subfigures share one y-axis range, and all objective subfigures share another.

Figure 15: Runtime–quality tradeoff induced by the permissible alphabet

P

. Each input is a BF16 vector

X\in\mathbb{R}^{d}

with

d=16{,}384

; columns show Normal, Lognormal, Student-

t

, and sparse-outlier inputs. We use at most

s=16

quantization values,

\lambda=0.1

, requested candidate counts

p\in\{16,32,64,128,256\}

. The upper row reports runtime, including histogram construction, candidate-set construction, and optimization, but excluding vector generation and metric serialization. The lower row reports

J_{P}/J_{P_{X}}

, where

J_{P}=\mathrm{vNMSE}(Q_{P},X)+\lambda H(\widehat{X})

. Points and error bars show the mean and two-sided 95% Student-

t

confidence interval. The horizontal axis uses the actual candidate count after duplicate removal.

Figure 16: Entropy-constrained vNMSE and an upper bound on the benefit of two-way time sharing. Columns show BF16 Normal, Lognormal, Student-

t

, and sparse-outlier vectors with

d=16{,}384

P=P_{X}

, at most

s=16

quantization values, and entropy budgets

b\in[1,4]

bits/value, sampled every

0.001

bits and augmented with retained feasible-frontier breakpoints. For each input, Optimal ECASQ solutions supported by

\mathrm{vNMSE}(Q,X)+\lambda H(\widehat{X})

are adaptively refined from

\lambda=0

through an endpoint attaining at most one bit/value. Curves in the upper row are arithmetic means over the five seeds: solid lines show the dual lower bound

L(b)

, while dashed steps and open diamonds show the best retained feasible distortion

D_{\rm single,feasible}(b)

. For each seed, these quantities give the relative-gain upper bound

U(b)=100[1-L(b)/D_{\rm single,feasible}(b)]

shown in the lower row. Thin lines are the five seeds and the thick line is their pointwise median. All lower subfigures here and in Figures 17 and 18 use the same

0

–

22\%

linear scale. Entropy is

H(\widehat{X})

and excludes reconstruction-value and entropy-model metadata, global framing, and finite-stream redundancy.

Figure 17: Entropy-constrained vNMSE and the time-sharing gain upper bound with at most

s=64

quantization values. The setup and curve encodings match Figure 16 : columns show the four BF16 synthetic distributions with

d=16{,}384

P=P_{X}

; budgets

b\in[1,6]

bits/value are sampled every

0.001

bits and augmented with retained feasible-frontier breakpoints. The upper row shows the mean dual lower bound and best retained feasible distortion, and the lower row shows the five seed-wise gain upper bounds and their pointwise median.

Figure 18: Entropy-constrained vNMSE and the time-sharing gain upper bound with at most

s=256

quantization values. The setup and curve encodings match Figure 16 : columns show the four BF16 synthetic distributions with

d=16{,}384

P=P_{X}

; budgets

b\in[1,8]

bits/value are sampled every

0.001

Adaptive Quantization

See all 19 papers

Inner Product Aware Quantization: Provably Fast, Accurate, and Adaptive Algorithms
May 29, 2026Nathan White, Krish SingalVector Quantization Neural Network Quantization
University of Pennsylvania
Benchmarking Neural Speech Compression from a Rate-Distortion Perspective
Jun 10, 2026Jun Xu, Zhengxue Cheng, Fengxi Zhang +3Neural Audio Codecs Speech Processing
School of Information Science and Electronic Engineering, Shanghai Jiao Tong University, Shanghai 200240, China
Provable Quantization with Randomized Hadamard Transform
May 13, 2026Ying Feng, Piotr Indyk, Michael Kapralov +2Vector Quantization Randomized Sketching

May 29, 2026J/K move · Enter open · S save

Inner Product Aware Quantization: Provably Fast, Accurate, and Adaptive Algorithms

Nathan White, Krish Singal

University of Pennsylvania

Vector Quantization Neural Network Quantization Adaptive Quantization

Open paper AskPDF arXiv

Quantization is a fundamental tool used to compress datasets, neural network weights, and memory usage in a range of computational tasks. Many downstream applications of vector quantization perform inner products with arbitrary inputs. This motivates the study of inner product aware quantization schemes that approximately preserve inner products with unseen vectors -- in contrast to simply minimizing the mean-squared error. In this work, we formulate objectives that capture natural desiderata and develop adaptive and unbiased quantization methods that approximately preserve inner products with worst-case and average-case inputs. An analysis of these objectives shows a tight connection with the well-studied notion of Adaptive Stochastic Quantization (ASQ). We develop provably fast exact and approximate algorithms for our objectives. Our theoretical results inspire efficient practical algorithms that perform well across a variety of workload distributions. They also lead to practical algorithms for standard ASQ which are 2-10

\times

faster than prior state-of-the-art methods while maintaining quality. These theoretical and empirical results contribute towards making adaptive quantization techniques more efficient and tractable in practical settings.

Abstract

Adaptive Quantization

Explore similar work

Inner Product Aware Quantization: Provably Fast, Accurate, and Adaptive Algorithms

Benchmarking Neural Speech Compression from a Rate-Distortion Perspective

Provable Quantization with Randomized Hadamard Transform