PaperlayerView exact revision on arXiv

Source-linked AI summary

Sliced Wasserstein Sketches with Error Certificates

Elaheh Akbari, James Liu, Keaton Hamm, Soheil Kolouri

arXiv:2610.20788v1cs.LG

TL;DR

Repeated transport comparisons on original samples can dominate retrieval, clustering, and related tasks. The paper builds adaptive quantization sketches over shared projections and derives exact distance-error certificates. Across benchmarks, the sketches preserve uncompressed sliced-distance accuracy with far fewer stored numbers.

  • Problem

    Repeatedly computing transport distances on original samples can dominate the cost of tasks comparing many distributions, motivating compact summaries that also quantify distance deviation.

    Source: §1 · p. 1

  • Method

    The paper quantizes each shared projection with an adaptive free-weight M-atom summary and derives an exact decomposition of squared distance error for centroid sketches.

    Source: §1 · p. 2

  • Results

    1.5 to more than 3000 times fewer stored numbers reproduce uncompressed sliced-distance accuracy across document, point-cloud, and flow-cytometry benchmarks.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Optimal sketches can preserve uncompressed sliced-distance classification performance under substantial compression.

    Source: §6 · p. 12

    The deterministic decomposition leaves sampling rates to assumptions on the data, because its last two terms are unobservable and depend on sample sizes and data.

    Source: §4.2 · p. 8

Abstract

from arXiv · show

Compact representations that support geometrically meaningful comparisons are valuable for learning from probability distributions. We study optimal, free-weight $M$-atom quantization of one-dimensional probability measures in quadratic Wasserstein distance and extend it to higher dimensions via slicing. Using $L$ shared projections yields a sketch with $O(LM)$ comparison time. Our main result is an exact decomposition of squared distance error into residual and partition-mismatch terms, yielding computable certificates for arbitrary centroid sketches. For optimal sketches, a finite $(2+\delta)$ moment for some $\delta>0$ guarantees $O(M^{-1})$ error in the finite-direction sliced distance, while equal-mass sketches can converge more slowly. Each slice preserves its mean, while storing residual energies enables covariance recovery from sufficiently many generic directions. Experiments on document, point-cloud, and flow-cytometry benchmarks show that optimal sketches substantially reduce storage while preserving the classification performance of uncompressed sliced distances. They also improve generalization from few labelled subjects and support end-to-end learning as compact pooling layers.

1 Introduction

The paper develops adaptive, free-weight quantization sketches for cheap distribution comparisons and analyzes how independently chosen partitions affect distance error. It gives computable certificates and approximation guarantees, alongside benchmark evidence of strong compression with preserved classification accuracy.

1 Introduction

  • Each projection is summarized by an optimal measure with at most M atoms, with both atom locations and weights adapting to the projected distribution.For weighted point measures, this is one-dimensional weighted k-means; L shared directions yield O(LM) storage and comparison time, and each slice’s quantization error is returned.Source: §1 · pp. 1–2
  • An exact identity for arbitrary centroid sketches separates residual and partition-mismatch effects, yielding computable certificates for distance error.Adaptive partitions can differ, so the shared-partition contraction argument does not apply; the triangle inequality bounds error but does not describe mismatched-cell effects.Source: §1 · p. 2
  • O(M^-2) squared-distance error holds when quantiles are Lipschitz and cell masses are O(M^-1), while general centroid sketches can reach the M^-3/2 error scale.The partition-mismatch identity yields two-sided bounds and shared-partition contraction.Source: §1 · p. 2
  • A finite (2+δ) moment for some δ > 0 guarantees O(M^-1) error in the finite-direction sliced distance for optimal sketches.Gaussian projections and projections of uniform balls in dimensions d ≥3 have slower equal-mass rates.Source: §1 · p. 2
  • Centroid identities preserve directional means and relate second-moment deficit to residual energy, supporting recovery of means and covariances from finitely many directions.The paper also bounds higher-moment tensor errors, with a per-cell correction that improves the rate.Source: §1 · p. 2
  • 1.5 to more than 3000 times fewer stored numbers preserve uncompressed sliced-distance accuracy across document, point-cloud, and flow-cytometry benchmarks.Learned slices also make the sketched distance competitive with full-dimensional transport.Source: §1 · p. 3

2 Related Work

The related work spans exact scalar quantization, quantization theory, sliced transport, and distribution embeddings. The paper distinguishes its analysis of mismatched centroid partitions from prior approaches that bound pairwise error through individual approximation errors.

Exact one-dimensional quantization

  • Dynamic programming gives globally optimal scalar quantization, while Lloyd–Max iterations can stabilize at suboptimal partitions for general distributions.The cited exact methods achieve O(MN) time after sorting; one divide-and-conquer alternative takes O(MN log N).Source: §2 · p. 3

Quantization theory

  • Classical quantization results provide centroid-stationarity bounds, asymptotic formulas, and extended Pierce bounds used in the paper’s analysis.Separate theories cover equal-mass rates and empirical-quantizer consistency and stability.Source: §2 · p. 3

Sliced transport and compressed comparisons

  • Prior Gaussian projection summaries, k-means transport, common-codebook quantization, and structured approximations control individual errors and derive pairwise bounds through the triangle inequality.This paper instead resolves residual and boundary contributions for centroid sketches with different mass partitions.Source: §2 · p. 3

Distribution sketches and embeddings

  • Stream-SW combines shared directions, KLL rank sketches, and exact one-dimensional transport for streaming updates with probabilistic support-diameter bounds.This paper instead performs offline quantization and records its distortion.Source: §2 · p. 3
  • Equal-mass cell means are optimal when weights are prescribed, while directional moment reconstruction recovers polynomial information rather than identifying an entire law through projections.Related work also includes quantile-grid embeddings, midpoint grids, and geometric or Wasserstein coresets.Source: §2 · p. 4

3 W2-Optimal Sketches

The section builds one-dimensional centroid quantizers and extends their comparisons to shared projection slices. Adaptive cell masses define optimal quantizers, while centroid conditions also cover prescribed partitions.

3.2 Centroid sketches

  • Optimal cells are nearest-neighbour intervals separated by adjacent-atom midpoints, unlike equal-mass cells that need not satisfy this condition.For the illustrated AML side-scatter channel, the optimal boundary s*=0.32 satisfies the midpoint condition and lowers energy by 63%.Source: §3.1 · p. 4, Figure 2 · p. 5

3.1 One-dimensional building block

  • For weighted point measures, free-weight M-atom quadratic Wasserstein quantization is weighted k-means on the line, with contiguous clusters and an exact dynamic-programming global optimum.The implementation uses an O(MN log N) divide-and-conquer variant and retains O(MN) traceback storage per slice.Source: §3.1 · p. 4

3.2 Centroid sketches

  • A centroid sketch partitions quantile mass intervals and represents each interval by its cell mean; its residual norm equals its Wasserstein error.The residual integrates to zero on each cell, and every M-cell sketch has error at least that of the optimal quantizer.Source: §3.2 · p. 5

3.3 Sliced sketch and comparison

  • With M atoms per slice, the sketch stores L(2M −1) scalars; one residual energy per slice raises this to 2LM, and all per-cell energies require L(3M −1).The sliced comparison uses the same directions as its reference, separating sketching error from finite-direction approximation.Source: §3.3 · pp. 5–6
  • Comparisons couple sorted weighted atoms using at most 2M −1 nonzero flows per slice, giving O(LM) time per query.Each slice is quantized independently; the analysis also applies to shared learned directions or measurable maps, while moment recovery requires linear slices in general position.Source: §3.3 · pp. 5–6

4 Theory

The theory decomposes sketch-distance error into residual and partition-mismatch terms, yielding computable certificates for adaptive centroid sketches. It also establishes sliced error rates and shows how slice residual energies support moment recovery.

4.1 Exact sketch distance identity

  • The squared distance error separates into a nonnegative residual term and boundary terms that vanish on shared partitions and can otherwise expand sketch distances.Unlike a triangle-inequality bound, the identity distinguishes underestimation from overestimation.Source: §4.1 · p. 6
  • Theorem 3 and Theorem 7 provide computable distance certificates; Lipschitz quantiles with maximum cell mass O(M^-1) yield O(M^-2) squared-distance error.Per-cell residual energies refine the boundary budget and also support higher-moment correction.Source: §4.1 · p. 7

4.2 Sliced certificates

  • Sliced certificates separate sketch quantization error from sampling error, but only the quantization terms are observable from the sketches.For independent uniformly sampled directions, the Monte Carlo estimate has an O(L^-1/2) unsquared error bound under finite second moments.Source: §4.2 · p. 8

4.3 Rates: free weights vs. equal-mass cells

  • A finite (2+δ) moment guarantees an O(M^-1) bound on optimal-quantizer error for every chosen direction.The bound scales with the (2+δ)-moment of X and holds for some δ > 0.Source: §4.3 · p. 8

4.4 Moments

  • Each slice preserves its mean, its second-moment deficit equals residual energy, and stored per-cell energies enable exact second-moment correction and sharper higher-moment bounds.Least-squares combination across directions bounds recovered moment-tensor error; exactness applies to the measure being sketched, including empirical moments for empirical inputs.Source: §4.4 · p. 9

5 Experiments

The experiments test whether compact optimal sketches preserve classification performance, improve learning from scarce labels, and serve as end-to-end pooling layers. Across seven benchmarks they remain close to uncompressed test performance with substantially fewer stored values.

Classification

  • The classification study uses five document datasets, AML samples, and ModelNet40, with L = 16 learned and frozen projections and nearest-neighbour classification.For each benchmark, M is selected as the smallest value within 0.5 percentage points of uncompressed validation accuracy; optimal sketches store L(2M − 1) values versus LN for uncompressed projections.Source: §5 · p. 10
  • 48–240 stored values per distribution keep optimal sketches within 0.3 percentage points of uncompressed test scores across all seven benchmarks.Storage falls by 1.5–18× for documents, 68× for point clouds, and over 3000× for AML samples; optimal sketches outperform equal-mass sketches on all seven benchmarks.Source: §5 · p. 10

Learning from few subjects

  • The optimal sketch has higher test AUC than the uncompressed representation at every AML training size.Equal-mass cells at matched storage remain within 1.2 points of the uncompressed representation, while mean and covariance features do not consistently improve on it.Source: §5 · p. 11

Sketches as a pooling layer

  • The pooling experiment compares equal-mass and optimal sketches with four cells per slice, changing only the partition in the shared pooling setup.The sketch is computed exactly in each forward pass; gradients flow through cell-mean atoms while the partition remains fixed.Source: §5 · pp. 11–12
  • With no backbone, the optimal sketch is more accurate than equal-mass cells at every slice count, by 1.7 to 4.6 points.With the MLP backbone it leads by 0.4 to 1.5 points for L ≥ 16, but is less accurate at L = 4; at L = 256 it comes within 2.5 points of PMA and 1.8 points of FSPool.Source: Figure 4 · p. 11, §5 · p. 12

6 Conclusion

The conclusion presents partition mismatch as the key obstacle to understanding distance error for adaptive centroid sketches and reports that optimal sketches preserve classification performance under substantial compression. The proof passage shows how residual and boundary terms yield bounds and a computable certificate.

6 Conclusion

  • Optimal sketches can preserve uncompressed sliced-distance classification performance under substantial compression, but minimizing quantization error does not guarantee the best distance approximation or downstream performance.Adaptive sketches also require more storage than equal-mass sketches at the same atom count.Source: §6 · p. 12
  • The partition-mismatch identity characterizes how independently chosen cells affect pairwise distances and yields error certificates for centroid sketches.The conclusion identifies faster construction and sharper distance-error guarantees for optimal quantizers as future work.Source: §6 · p. 12

A.1 Proof of Theorem 3

  • In the proof of Theorem 3, squared distance error separates into a residual difference and a boundary term, which is nonpositive and bounded using cell contributions.The boundary budget can be evaluated in O(Mµ + Mν) time by merging the two sorted boundary lists.Source: §A.1 · p. 18

A.2 Regularity refinements of Theorem 3

The appendix refines distance-error bounds under regularity assumptions, gives an example showing the general boundary budget can be sharp, and supplies proof steps for sliced-distance and moment-based rates.

A.2 Regularity refinements of Theorem 3

  • For Lipschitz quantiles, the boundary budget yields O(M^-2) squared-distance error when both largest cell masses are O(M^-1).For bounded densities on finite intervals, optimal free-weight quantizers also have cell widths bounded by a constant times interval length divided by M.Source: §A.2 · pp. 18–19

A.3 Sharpness of the boundary budget

  • For a numerically checked sequence of pairs with general centroid sketches, νM = qh2/4 + O(M−2); whether an O(M−2) gap holds for all optimal pairs with εµ, εν = O(M−1) remains open.Source: §A.3 · p. 20

A.4 Proof of Theorem 4

  • The sliced-distance proof averages the one-dimensional bounds over slices and uses Minkowski’s inequality, then a reverse triangle inequality for the unsquared bound.For empirical inputs, the product-space triangle inequality applies generally; sampling assumptions enter only when bounding the sampling term.Source: §A.4 · p. 20

A.5 Proof of Theorem 5

  • The optimal-quantization rate proof applies the extended Pierce bound to each projection; its stated rate requires a finite moment of order 2 + δ for some δ > 0.A finite second moment alone does not guarantee that rate.Source: §A.5 · p. 21

A.6 Equal-mass approximation rates

This appendix compares equal-mass and optimal one-dimensional quantization rates. Low-density regions can force equal-mass cells to be wide, producing slower decay than optimal quantization in the Gaussian and uniform-ball projection examples.

A.6 Equal-mass approximation rates

  • Under a finite moment of order 2 + δ, optimal quantization has a finite positive asymptotic M^2-scaled squared error constant.The equal-mass comparison depends on the density and can have slower rates when low-density regions are present.Source: §A.6 · p. 21
  • Equal-mass error decays more slowly than M^-2 when the integral of f^-1 is infinite, as in Gaussian tails and near uniform-ball projection endpoints.Low-density cells must be wide to carry mass 1/M; the Gaussian rate is sharp along M = 2^k.Source: §A.6 · p. 21
  • The limiting equal-mass and optimal error constants coincide exactly for distributions uniform on an interval; this comparison concerns limiting constants, not finite-M errors.For other distributions in the stated comparison, the limiting equal-mass constant is larger.Source: §A.6 · pp. 21–22
  • For uniform-ball projections, equal-mass squared-error rates depend on dimension, while optimal squared error remains of order M^-2 for every fixed dimension.The equal-mass rates are M^-2 for d ≤ 2, M^-2 log M for d = 3, and M^-(d+5)/(d+1) for d ≥ 4.Source: §A.6 · pp. 21, 23

A.7 Proof of Theorem 6

Centroid sketches preserve directional means, and their residual energies quantify the second-moment deficit. The proof uses these identities and a per-cell correction to bound higher-moment errors and recover moments from sufficiently many directions.

Centroid and scalar bounds

  • The centroid condition alone gives |Ef(Xθ)−Ef(ˆXθ)| ≤ 1/2∥f′′∥∞ε²µ for twice-differentiable f with bounded second derivative.For bounded support, applying this to f(x)=x^r yields the stated scalar moment bounds.Source: §A.7 · p. 23

Per-cell correction

  • The per-cell correction retains only second-order residual energies, and its higher-order remainder bound requires no residual moments above order two.This follows by binomial expansion and the zero conditional mean of each cell residual.Source: §A.7 · pp. 23–24

Directional reconstruction

  • With sufficiently many generic directions, directional moments reconstruct tensors; d spanning directions recover the mean, and injective second-order reconstruction recovers covariance exactly.The reconstruction matrix has full column rank outside a measure-zero exceptional set of direction tuples.Source: §A.7 · pp. 24–25

Moment error bounds

  • For every fixed moment order, εµ,θ = O(M−1) gives O(M−2) moment error, while hθ = O(M−1) improves the corrected bound to O(M−3).The latter condition holds for optimal quantizers under the bounded-density assumptions, where cell widths are O(M−1).Source: §A.7 · p. 25

B Implementation details

The implementation computes optimal slice sketches and compares them efficiently, then evaluates classification, few-subject learning, and pooling experiments under specified protocols. Separate numerical checks test the paper’s theoretical statements, but finite computations do not establish universal bounds or asymptotic rates.

B Implementation details

  • Sketch distances use a monotone transport sweep in O(M) time per slice, while exact dynamic programming computes each optimal sketch.Weights are normalized, zero-weight points removed, and duplicate locations merged before sketching.Source: B Implementation details · p. 25
  • Classification compares optimal sketches with equal-mass and uncompressed distances across document, AML, and ModelNet40 benchmarks using shared splits and validation-based kNN selection.The experiments also evaluate AML few-subject learning and replace only the pooling layer in the ModelNet40 protocol.Source: B Implementation details · pp. 25–26

C Numerical validation of the theoretical results

  • Numerical experiments check finite identities and inequalities and illustrate rates under analytically specified assumptions, but do not prove universal bounds or asymptotic theorems.The section maps each theoretical result to its experiment and recorded data.Source: C Numerical validation of the… · p. 26

C.1 Partition mismatch and computable certificates

The numerical studies check partition-mismatch identities and certificates, sampling bounds, quantization rates, regularity-based bounds, and a sharpness construction. Their findings illustrate the stated results in specific settings rather than establishing universal rates.

C.1 Partition mismatch and computable certificates

  • The quantity 2 + D2 + S2 + H0 is 1.29 × 10−15; certificate inequalities had no violation beyond scaled tolerance 2 · 10−10, with largest positive scaled discrepancy 4.89 × 10−16.Shared partitions have zero boundary term and contract the distance. An atomic input is recovered exactly once the budget reaches its support size.Source: §C.1 · p. 28

C.2 Sliced, sampling and direction certificates

  • For direction sampling, log–log slopes were −0.504 for squared-distance mean absolute error and −0.506 for distance error, consistent with L−1/2.The experiment used 8,192 independent direction sets for each L.Source: Figure 5 · p. 28, §C.2 · p. 29
  • For the tested Gaussian pair, sketch-distance error stayed below Eµ + Eν, while M(Eµ + Eν) stabilized and the observed error decayed faster.All 64 finite-population sampling errors were also below the deterministic sampling bound; these checks make no sampling-rate claim.Source: §C.3 · p. 29, §C.2 · p. 29

C.3 Optimal and equal-mass quantization rates

  • Equal-mass errors remained of constant order under the tested law-specific normalizations, and the equal/free error ratio grew for the Gaussian and uniform balls in dimensions d ≥ 3.These are finite-range illustrations; logarithmic corrections converge slowly.Source: §C.3 · p. 29, Figure 6 · p. 30

C.4 Regularity and optimal-cell geometry

  • For mismatched partitions of Lipschitz quantiles, both signed distance bounds held and Figure 7 showed the predicted M−2 squared-distance scale.For the bounded-density example at M = 512, f+ had Mwmax = 1.1864 and Mqmax = 1.3446.Source: §C.4 · p. 31, Figure 7 · p. 31
  • For the bounded-density pair, every squared gap was below Kµν/M²; at M = 512, M²|W²−D²| = 0.17901.The explicit constant was Kµν = 1247.6, while the computable bound was much smaller.Source: §C.4 · p. 31

C.5 Sharpness of the boundary budget

  • At M = 65,536, the sharpness construction’s rescaled signed gap was −1.005526 and (D²−W²)/S² was 143.12, for general centroid sketches.This construction does not establish a counterexample for optimal quantizers.Source: §C.5 · p. 31, Figure 8 · p. 32

C.6 Moment fidelity, correction and recovery

The moment experiments test scalar and tensor error bounds, a per-cell correction, and directional recovery. In the tested smooth example the corrected fourth-moment errors decayed faster than the stated bounds, and residual energies enabled accurate covariance recovery.

Finite scalar and tensor bounds

  • Across 17,280 scalar checks, all inequalities passed; the largest corrected-error/bound ratio was 0.192, and tensor ratios were 0.558 uncorrected and 0.0363 corrected.The checks included moment orders 2–6 and the C2 bound for several test functions.Source: §C.6 · p. 32

Moment rates

  • For fourth directional and tensor moments, corrected errors had fitted slope −4.000 versus −1.999 without correction, faster than the M−3 corrected bound in this smooth example.The theoretical budgets scale as M−2 without correction and M−3 with correction.Source: §C.6 · p. 32, Figure 9 · p. 33

Exact recovery, rank and conditioning

  • Adding one residual energy per slice recovered the unbounded Gaussian covariance with maximum Frobenius error 3.78 × 10−14; three spanning directions recovered the mean with error 2.29 × 10−15.The checks also found that nearly aligned directions, although full rank, amplify perturbations.Source: §C.6 · p. 32, Figure 9 · p. 33

D Additional experiments

The additional experiments compare optimal sketches with equal-mass cells across sweeps of L and M, then test optimal sketches with a class-balanced linear SVM on few labelled subjects.

D.1 Sweeps over L and M

  • Optimal sketches store L(2M −1) numbers versus LM for equal-mass cells; each configuration uses its own validation-selected k.Bands show standard errors for documents and AML and bootstrap 95% intervals for ModelNet40.Source: §D.1 · p. 33
  • At M=4, optimal sketches are more accurate in all 23 benchmark–L pairs by 0.8 points on average; at M=8, they lead in 21 pairs by 0.6 points.Figures 10–13 show the comparisons across the L and M sweeps.Source: §D.1 · p. 33, Figure 10 · p. 34, Figure 11 · p. 34
  • At M=2, optimal sketches lead in 15 pairs but trail on Recipe at every L; from M=16, mean differences are 0.1 points or less.The average lead at M=2 is 2.0 points, while the Recipe deficit is 2.3 points.Source: §D.1 · p. 33, Figure 10 · p. 34, Figure 11 · p. 34

D.2 Few labelled subjects: linear SVM

  • With class-balanced linear SVM, optimal sketches have higher test AUC than uncompressed at every training size; the paired difference is 3.9 ± 0.5 points for two AML subjects and 0.9 ± 0.3 for the whole fold.Table 5 repeats the comparison on the same slices, folds, and training-subject draws as the logistic-regression analysis.Source: §D.2 · p. 33, Table 5 · p. 36
Loading 2610.20788v1…