PaperlayerView exact revision on arXiv

Source-linked AI summary

Entropy-Constrained Adaptive Stochastic Quantization

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

arXiv:2608.18147v1cs.LGcs.AIcs.DScs.IT

TL;DR

Entropy encoding is not directly accounted for when existing unbiased quantizers select quantization values. ECASQ addresses this with entropy-constrained dynamic programs and refinement, achieving near-optimal accuracy with a faster, space-efficient approximation.

  • Problem

    Selecting quantization values for entropy-encoded outputs requires accounting for entropy overhead alongside quantization error.

  • Method

    ECASQ jointly imposes an entropy budget and permissible-value constraints, then solves the problem with optimal and GPU-friendly approximate dynamic programs plus iterative refinement.

  • Results

    Refinement reduces Approx ECASQ’s mean excess vNMSE from 4.97% to 1.37% across 12 distribution-rate comparisons, with a 4.06% maximum.

  • Takeaways & Limitations

    The approximation and refinement provide near-optimal quantization while retaining substantial practical speed and space advantages over the optimal dynamic program.

  • Takeaways & Limitations

    The Lagrangian relaxation may not attain the exact entropy constraint for a given b unless time-sharing between two quantizers is allowed.

Abstract

from arXiv · show

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.

1 Introduction

ECASQ jointly selects input-adaptive quantization values to minimize MSE under entropy, codebook-size, and unbiasedness constraints. The paper develops exact and approximate dynamic programs, plus iterative refinement that achieves near-optimal results efficiently.

  • Motivation: Quantization addresses memory, computation, and bandwidth constraints across machine-learning workloads, but desirable methods must combine unbiasedness, input adaptivity, and entropy-efficient encoding.Existing methods do not provide all three properties simultaneously.
  • Related work: ASQ adaptively minimizes MSE over codebooks of size at most s and is optimal for fixed-length identifiers, whereas entropy-aware methods generally lack unbiasedness or adaptivity.Entropy-aware scalar, vector, and neural compression methods do not enforce unbiasedness; several unbiased entropy-encoded schemes use predetermined grids.
  • Problem formulation: ECASQ minimizes MSE subject to |Q| ≤s, H(bX) ≤b, and E[bX] = X, while limiting codebook size to reduce representation and processing overhead.Its inputs include a permissible alphabet, maximum codebook size s, and average-entry entropy budget b.
  • Solution overview: O(p^2 · s + d) time and O(p^2 + d) space solve λ-ECASQ optimally after DP acceleration, versus O(p^3 · s + d) time and O(p^2 · s + d) space naively.The Lagrange multiplier λ traces the convex hull of the Pareto-optimal MSE-entropy trade-off, though an exact λ need not exist for every entropy budget.
  • Approximation: At most 1 bit away from optimality, the approximation uses O(p^2 · s + d) time and O(p + d) space after Hirschberg reduction.For b > 1, its MSE is no larger than an optimal λ-ECASQ solution with entropy bound b −1 and the same λ.
  • Evaluation: Near-optimal results after a handful of iterative-refinement steps retain substantial speed and space advantages over the optimal dynamic-program variant.The refinement follows the approximate optimization and is evaluated across multiple input distributions.

2 Preliminaries

This section defines stochastic quantization, its MSE and scale-free vNMSE objectives, and adaptive stochastic quantization (ASQ) for selecting input-specific quantization values. It also presents the optimal dynamic-programming formulation, which relies on quantization values drawn from the input and runs in O(s · d) time and space.

  • Stochastic Quantization: Stochastic quantization unbiasedly maps each sorted, distinct input entry between its two encompassing quantization values.If entries are unsorted, a sorted copy is used while preserving original indices; repeated entries can be handled with weights.
  • Error Metrics: MSE equals the sum of individual stochastic quantization variances, with each interval contribution given by (b_x − x)(x − a_x).The variance for an entry quantized between a_x and b_x is derived from its two possible quantization outcomes.
  • Adaptive Stochastic Quantization: ASQ selects s ≥2 quantization values for a specific X to minimize MSE, equivalently vNMSE.vNMSE is scale free: vNMSE(Q, X) = vNMSE(c·Q, c·X) for any c ≠ 0.
  • Adaptive Stochastic Quantization: O(s · d) time and space suffice for the optimal ASQ dynamic program, which assumes an optimum whose quantization values are entries of X.The program minimizes the MSE for the j smallest entries using i quantization values and recursively chooses the next quantization value.

3 Entropy Constrained ASQ

ECASQ extends adaptive stochastic quantization by jointly constraining entropy and permissible quantization values, then solves the resulting nonseparable optimization with optimal and approximate dynamic programs.

  • Problem formulation: ECASQ adds an entropy bound b and permissible values P to ASQ, enforcing H(bX) ≤ b and Q ⊆ P.The permissible set supports quantization values outside the input and practical encoding constraints.
  • Problem formulation: 0.948 bits per entry is the entropy for X = ⟨0, 1, 10⟩ with Q = {0, 10} under the stated stochastic assignments.The middle entry maps to 0 with probability 9/10 and to 10 otherwise.
  • Optimization: The Lagrangian cost combines MSE(Q, X) with λ′ · H(bX), or vNMSE(Q, X) + λ · H(bX), and binary search over λ enforces the entropy constraint.λ = 0 recovers ASQ.
  • Optimization: The standard ASQ dynamic program fails because entropy is nonseparable: each quantization value’s expected frequency depends on both adjacent quantization values.The optimal and approximate DPs are designed to address this dependence.
  • Algorithms: O(d^2 · s) time and O(d^2) space characterize the optimal algorithm, while the approximation uses O(d) space and, for b > 1, has MSE bounded by the optimum with b − 1 bits.The generalization to P ≠ X appears in Appendix D.

4 The Optimal Dynamic Program

The optimal dynamic program tracks adjacent quantization values to account for expected entropy contributions while minimizing quantization cost. SMAWK and Hirschberg reduce its time and space requirements, and the formulation extends to general permissible value sets.

  • Dynamic-program formulation: Tracking the last two quantization values captures each quantization value’s expected frequency and entropy contribution.The expected frequency depends only on adjacent values in Q.
  • Dynamic-program formulation: Helper arrays U[a,b] and D[a,b] aggregate round-up and round-down probabilities between adjacent quantization values.For consecutive xk, xℓ, xj ∈ Q, the expected frequency at xℓ is U(k,ℓ) + D(ℓ,j).
  • SMAWK acceleration: SMAWK acceleration reduces the optimal DP runtime to O(d^2 · s).Column reversal makes the relevant matrix satisfy the quadrangle inequality, enabling O(d)-time minimization for each i and ℓ.
  • Space reduction: Hirschberg’s divide-and-conquer implementation achieves O(d^2) space while retaining O(d^2 · s) time complexity.It keeps only pairs of DP layers and recursively reuses memory.
  • General permissible sets: For a sorted permissible set of size p, the optimal algorithm runs in O(p^2 · s + d) time and uses O(p^2 + d) space.The extension indexes DP states by permissible values and adjusts boundary conditions.

5 A Space-Efficient Approximation Algorithm

The space-efficient approximation replaces nonlocal entropy accounting with a locally separable surrogate, enabling a dynamic program that uses less space. Its solution is feasible for the original entropy budget and has MSE no worse than the optimum with one fewer entropy bit per entry.

  • Space-efficient approximation: A local interval-based entropy surrogate avoids storing two adjacent quantization values in the dynamic-programming state.The surrogate optimizes H(I, B) instead of H(bX), making interval contributions separable.
  • Approximation guarantee: The returned quantizer satisfies H(bX) ≤ b and has MSE no larger than OPT(b − 1), the best quantizer using one fewer entropy bit per entry.The guarantee follows because any quantizer feasible at budget b − 1 is feasible for the surrogate budget b.
  • Space-efficient approximation: The actual entropy overhead is H(B | bX), which depends on the input X and quantization set Q and is bounded by one bit per entry.The method optimizes the joint entropy to select Q, then compresses the quantized values bX.
  • Space-efficient approximation: O(sp^2 + d) time and O(p + d) stored state support all s dynamic-programming layers after preprocessing.Each layer costs O(p^2) time, while only previous and current layers plus prefix information are stored.

6 The Refinement Algorithm

The refinement algorithm iteratively improves a feasible ECASQ solution while preserving unbiasedness, using GPU-parallel updates to reduce the accuracy gap efficiently. Although optimality is not guaranteed, refinement always improves or matches its starting solution and is usually near-optimal while remaining very quick.

  • 6 The Refinement Algorithm: Refinement iteratively improves a feasible solution while maintaining unbiasedness, targeting the accuracy gap between the approximation and optimal algorithms.It is motivated by the approximation algorithm’s space efficiency and GPU-friendliness on large inputs.
  • 6 The Refinement Algorithm: Each iteration updates one of three interleaved quantization-value sets, using binary search to optimize a changed value between its fixed neighboring values.The three sets are Q0, Q1, and Q2, partitioned by quantization-value indices modulo 3.
  • 6 The Refinement Algorithm: The three sets can be optimized in parallel on the GPU because moving a value in one set does not change another set’s expected frequency.This independence enables parallel optimization while keeping the other quantization values fixed.
  • 6 The Refinement Algorithm: Consecutive values within each set separate small and large entries, yielding a complete separation of MSE and entropy across that set.Entries near the lower boundary cannot use the higher consecutive value, while entries near the upper boundary cannot use the lower one.
  • 6 The Refinement Algorithm: Refinement is never worse than its starting solution and is usually near-optimal while being very quick compared to the approximate solution itself.The algorithm is not guaranteed to reach the optimal solution.

7 Optimality with Two-Way Time-Sharing

Lagrange relaxation recovers supported entropy-vNMSE frontier points, while two-way time-sharing extends optimality to every entropy budget. The method preserves adjacent stochastic quantization by applying two quantizers to a random partition of the input.

  • From Lagrange Multipliers to Entropy Constraints: Lagrange relaxation is globally optimal when its multiplier search exactly matches the entropy budget, but otherwise recovers only supported frontier points.No time-sharing is needed when H(b̂_X) = b.
  • Optimality with Two-Way Time-Sharing: Two-way time-sharing attains an optimal solution for every entropy constraint by mixing at most two supported quantizers.A sweep over λ recovers the required mixture.
  • Quantizing to non-adjacent values under two-way time-sharing: Time-sharing runs two quantizers on a random partition of the input, recovering optimality for nonsupported entropy budgets while retaining adjacent rounding in Q.The partition is determined by a shared pseudorandom number generator seed.
  • Quantizing to non-adjacent values under two-way time-sharing: Proposition 5 shows that arbitrary unbiased rounding rules can be matched by two ordinary adjacent stochastic quantizers using sets of at most s values.The construction applies for every α ∈ [0, 1].

8 Evaluation

Evaluation shows that refined Approx ECASQ approaches Optimal ECASQ in matched-rate vNMSE while retaining substantial runtime advantages, and improves results across real-model tensor tasks.

  • Setup: Experiments use synthetic BF16 vectors of size d = 16,384 across five seeds and four distributions, plus 100 real-model tensor samples.Comparisons include Optimal ECASQ, Approx ECASQ, five-round refined Approx ECASQ, and entropy-coded non-ECASQ baselines.
  • Matched-rate vNMSE: Refinement reduces Approx ECASQ’s mean excess vNMSE from 4.97% to 1.37%, with a 4.06% maximum across 12 distribution-rate comparisons.Unrefined Approx ECASQ beats every non-ECASQ baseline in mean at every rate and beats QSGD and QUIVER pointwise in all comparisons.
  • Runtime: At d = 262,144, Optimal ECASQ, Approx, and refined Approx take 6.39, 0.26, and 0.33 seconds on average, yielding 24× and 19× speedups.Across tested distributions, speedups range from 13-34× for Approx and 8-28× after refinement.

A Proof of Lemma 1

The proof establishes that G(k, j) = g(v(k) + m(j)) satisfies the quadrangle inequality when g is concave and v, m are monotonically decreasing. It does so by comparing ordered function arguments and applying concavity.

  • Proof: For every k < k′ and j < j′, the proof targets G(k, j) + G(k′, j′) ≤ G(k, j′) + G(k′, j).This is the quadrangle inequality to be established.
  • Proof: Monotonic decrease gives v(k) ≥ v(k′) and m(j) ≥ m(j′), ordering the increments used in the concavity argument.The proof introduces A, B, C, and D with A + D = B + C.
  • Proof: Concavity yields g(B) + g(C) ≥ g(A) + g(D), which is equivalent to the desired inequality after substituting the definitions.The zero-increment case is immediate; otherwise, the proof assumes ∆v + ∆m > 0 before applying concavity.
  • Proof: Therefore, G satisfies the quadrangle inequality.The conclusion follows directly from the preceding substitution and the definition of G.

B Proof of Lemma 2

For fixed i and ℓ, the column-reversed matrix f M satisfies the quadrangle inequality. The proof reduces to the entropy term, using monotonicity of v and D and concavity of g, while row-dependent terms cancel or preserve the inequality.

  • For fixed i and ℓ, the column-reversed matrix f M satisfies the quadrangle inequality.
  • The proof suffices to establish the inequality for the entropy term because adding the row-dependent term A(k) preserves it.
  • Concavity of h(y) = −y log2 y on [0, 1] implies that g(z) = λh(z/d) is concave over all matrix arguments.
  • The proof establishes that v(k) decreases with k and D(ℓ, j) increases with j because the relevant summands and additional terms are nonnegative.
  • With v and em decreasing and g concave, Lemma 1 yields the quadrangle inequality; row terms A(k) then cancel from both sides.

C Proof of Lemma 3 … E.2 Quantizing to Non-Adjacent Values in Q

The paper proves a Hirschberg-based dynamic program with O(d^2) space and O(d^2 · s) time, extends it to arbitrary permissible values, and identifies limits of Lagrangian and adjacent-rounding formulations. Allowing time-sharing, non-adjacent quantization, and differential treatment of identical entries can improve constrained vNMSE.

  • C Proof of Lemma 3: O(d^2) space and O(d^2 · s) time follow from Hirschberg’s algorithm for the dynamic program.Hirschberg reduces memory while preserving the dynamic program’s time complexity.
  • C Proof of Lemma 3: The reconstruction splits at a minimizing crossing pair into independent recursive subproblems with disjoint unfixed candidate sets.Forward and backward layers identify the crossing pair, after which recursion proceeds on the two sides.
  • D Generalizing to Arbitrary P: For arbitrary permissible values P, the optimal algorithm replaces input indices with permissible-value indices and runs in O(p^2 · s + d) time.SMAWK and Hirschberg continue to apply after replacing d candidate indices by p permissible-value indices.
  • D Generalizing to Arbitrary P: Pairwise costs for arbitrary P are computed after prefix-sum construction in O(p^2 + d) time and O(p^2) space.The required prefix quantities count, sum, and sum of squares of sorted input entries.
  • E.1 From Lagrange Multipliers to Entropy Constraints: When b = 0.95, quantizer B is the hard-constrained optimum, yet no value of λ produces B as the Lagrangian optimum.This demonstrates that Lagrangian optimization need not recover the optimum for a fixed entropy constraint.
  • E.1 From Lagrange Multipliers to Entropy Constraints: Time-sharing achieves entropy H( b X) = 0.95 and vNMSE ≈0.26 < 1/2, improving over quantizer B.The construction mixes quantizers C and A using a random subset of entries.
  • E.2 Quantizing to Non-Adjacent Values in Q: Under adjacent rounding, the entropy constraint forces Q = {0, 2}, yielding entropy 1 ≤b and MSE d/3.For the constructed input, this solution has vNMSE 1/5.
  • E.2 Quantizing to Non-Adjacent Values in Q: Allowing non-adjacent quantization and differential treatment of identical entries reduces vNMSE to 1/10 under the same entropy budget.Half of the entries equal to 1 are mapped deterministically to 1, while the remainder are stochastically quantized between 0 and 2.

E.3 Proof of Proposition 5 … F.2 Matched-Rate vNMSE

The proof shows that arbitrary unbiased rounding cannot improve the entropy–vNMSE frontier beyond adjacent stochastic quantization with two-way time-sharing. The evaluation protocol and matched-rate real-model results show that refined Approx ECASQ remains close to optimal, while unrefined Approx ECASQ outperforms the baselines at 3 bits/value.

  • E.3 Proof of Proposition 5: Arbitrary unbiased rounding does not improve the entropy–vNMSE frontier when two-way time-sharing between adjacent quantizers is allowed.This remains true even when identical input entries may be treated differently.
  • E.3 Proof of Proposition 5: Every two-way time-share using arbitrary unbiased rounding is dominated by a two-way time-share using ordinary adjacent stochastic quantization with at most s values per quantizer.The proof represents expected output frequencies as mixtures of adjacent-SQ frequency vectors and reduces the mixture to at most two quantizers.
  • F.1 Protocol and Rate Accounting: The experiments use BF16 inputs and scale-free vNMSE, comparing methods at common storage budgets with ideal entropy-coded payloads plus 32 bits of metadata per active quantization value.The ideal payload uses H(bX), excluding global framing and finite-stream redundancy.
  • F Evaluation Details: The evaluation compares three ECASQ variants with entropy-coded QUIVER, QSGD, and uniform stochastic quantization across synthetic and real-model data.The real-model study covers weights, activations, and KV-cache tensors from five model families.
  • F.2 Matched-Rate vNMSE: Figures 5–7 evaluate matched-rate vNMSE on all 100 real-model tensors at every target rate.The tensors comprise 25 weight matrices, 45 activations, and 30 KV-cache tensors.
  • F.2 Matched-Rate vNMSE: 1.35% average and 4.8% maximum excess vNMSE are achieved by refined Approx ECASQ at 3 bits/value.Refined Approx ECASQ remains close to the best observed method across the matched-rate real-model study.
  • F.2 Matched-Rate vNMSE: 8.12% average excess vNMSE is achieved by unrefined Approx ECASQ at 3 bits/value, versus 44.35% for uniform, 54.9% for QUIVER, and 390.57% for QSGD.Unrefined Approx ECASQ is lower than each listed baseline.

F.3 Paired Uniform-Grid Control … F.7 Entropy–vNMSE Frontiers and Time Sharing

Across controlled comparisons, ECASQ’s refined approximation stays close to Optimal ECASQ while substantially improving over uniform quantization, with exact entropy accounting and refinement providing strong empirical guarantees. Uniform candidate grids offer the best permissible-alphabet tradeoff, while time sharing has limited potential gains at most tested budgets.

  • F.3 Paired Uniform-Grid Control: 3.35% mean excess vNMSE for refined Approx ECASQ compares with 4.41% for unrefined Approx ECASQ and 12.25% for uniform quantization across 18 matched-rate cells.Refined Approx ECASQ is within 6.62% of Optimal ECASQ in every cell, while uniform quantization reaches 33.46% excess.
  • F.3 Paired Uniform-Grid Control: 1.37% mean excess vNMSE after five refinement rounds, down from 4.97% across 12 synthetic distribution–rate points.Unrefined Approx ECASQ beats QSGD and QUIVER at all 12 points and uniform quantization at 11 of 12.
  • F.4 Approximation and Refinement Reality vs. Guarantee: Exactly H(B | b X) bits/value are saved by replacing the optimized entropy surrogate with exact output entropy at unchanged vNMSE, and all 100 checks satisfy the empirical one-bit guarantee.The saving varies with the Approx surrogate budget b and remains within the proved one-bit bound.
  • F.4 Approximation and Refinement Reality vs. Guarantee: 0.12% median objective gap after five refinement rounds, with a 0.58% maximum across 20 distribution/seed cases.Refinement stops when a complete round fails to decrease the objective, so most runs converge before 20 requested rounds.
  • F.5 Runtime and Objective Scaling: 24.4× and 19.2× speedups are achieved by Approx and refined Approx over Optimal ECASQ at d = 262,144 and s = 64.The pooled mean runtimes are 6.39 s for Optimal ECASQ, 0.26 s for Approx, and 0.33 s for refined Approx.
  • F.6 Effect of the Permissible Alphabet P: 0.25% objective deviation from observed-support results is maintained by 256-candidate uniform grids for every distribution while reducing runtime by 20–88×.Empirical quantiles incur objective penalties of 26.9%, 38.5%, and 25.2% on Lognormal, Student-t, and sparse-outlier inputs, respectively.
  • F.7 Entropy–vNMSE Frontiers and Time Sharing: Two-way time sharing is evaluated on output-entropy-constrained frontiers, because output entropy is the rate constraint convexified by the time-sharing result.The upper bound compares optimal time-shared and deterministic frontiers derived from globally minimized Lagrangian objectives.
  • F.7 Entropy–vNMSE Frontiers and Time Sharing: 98.2%, 94.9%, and 91.4% of the 60k, 100k, and 140k checks have time-sharing gain upper bounds at or below 5%.The maximum is 21.7% for each quantization-value budget and occurs on a narrow low-rate interval.
Loading 2608.18147v1…