Sliced Wasserstein Sketches with Error Certificates

Elaheh Akbari; James Liu; Keaton Hamm; Soheil Kolouri

cs.LGOct 14, 2026

Sliced Wasserstein Sketches with Error Certificates

Authors: Elaheh Akbari, James Liu, Keaton Hamm, Soheil Kolouri

Organizations: College of Connected Computing, Vanderbilt University · Department of Mathematics, Division of Data Science, University of Texas at Arlington

Abstract

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+δ)$ moment for some $δ>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.

Figures & tables

Figure 1: Sketch construction. Top: two bags in

\mathbb{R}^{2}

(one weighted) and two shared directions. Middle: projected atoms,

M=4

optimal contiguous cells (dotted), and weighted centroids per direction; marker area is proportional to cell mass, and cells differ between bags. Bottom: sketch step quantiles over mass

u

, with original quantiles in the background and shading marking coupled pairs at each mass level. This monotone coupling has quadratic cost

W_{2}^{2}(\hat{\mu}_{\ell},\hat{\nu}_{\ell})

; the squared sketch distance

D_{\Theta,M}^{2}

averages these costs over directions.

Panel of Figure 2 — Figure 2: Equal-mass and optimal cells ( $M=2$ ) for one AML cytometry subject’s side-scatter channel ( $10^{4}$ cells; Section 5 ). Atoms are cell means of $Q_{\mu}$ ; the shaded residual has energy $\varepsilon^{2}$ . At the equal-mass boundary, $Q_{\mu}$ differs from the atom midpoint, so the cells are not nearest-neighbour cells. The optimal boundary $s_{*}=0.32$ satisfies the midpoint condition and lowers energy by $63\%$ .

Panel of Figure 3 — Figure 3: Partition-mismatch identity ( 5 ) for forward-scatter projections of a normal and a leukaemia subject in the AML cytometry data ( $10^{4}$ cells each; Section 5 ), with optimal $M=3$ sketches. Top and middle: quantiles, step quantiles, and shaded residuals; sketch cells differ. Bottom: residual primitives ( 4 ), each zero at its own cell boundaries. Dots pair each sketch’s jumps with the other’s primitive at the same boundary. Here $W_{2}^{2}-D^{2}=-0.0127$ , comprising $\lVert r_{\mu}-r_{\nu}\rVert^{2}=0.0085$ and boundary terms $-0.0124$ and $-0.0088$ .

benchmark	$(N,M)$	BW	Exact OT	unc. $\mathrm{SW}_{2,\Theta}$	equal-mass	ours	error ratio
		test error ( $\downarrow$ )
BBCSport	$(117,4)$	$6.2\pm 1.9$	$5.1\pm 1.2$	$1.1\pm 0.8$	$1.9\pm 1.1$	$\mathbf{1.1\pm 0.8}$	$1.8$
Twitter	$(10,4)$	$28.7\pm 1.4$	$29.7\pm 1.5$	$27.0\pm 1.4$	$27.0\pm 0.9$	$\mathbf{26.7\pm 0.2}$	$4.0$
Classic	$(39,4)$	$3.2\pm 0.5$	$\mathbf{2.9\pm 0.4}$	$3.1\pm 0.3$	$3.8\pm 0.4$	$3.3\pm 0.3^{*}$	$1.8$
Recipe	$(48,8)$	$43.2\pm 1.0$	$42.9\pm 0.8$	$37.4\pm 0.9$	$38.7\pm 0.7$	$\mathbf{37.4\pm 1.0}^{*}$	$4.1$
Reuters	$(36,2)$	$4.7$	$3.9$	$2.9$	$3.4$	$\mathbf{2.8}$	$1.2$

Table 1: Classification with

L=16

shared learned slices.

N

: points per distribution (mean distinct words for documents);

M

: smallest optimal-sketch size matching uncompressed validation accuracy within

0.5

percentage points, also used for equal mass sketches. Scores are percentages. Bold marks the best method, and ∗ a significant paired difference between ours and equal-mass. Error ratio: median equal mass/optimal quantization energy. (error ratio > 1 shows that equal-mass results in larger error compared to ours.)

		AML subjects in training
summary	stored values	2	4	8	16	all
uncompressed	320,000	$92.7\pm 0.6$	$95.4\pm 0.4$	$97.1\pm 0.4$	$98.1\pm 0.4$	$98.6\pm 0.4$
mean, covariance	35	$92.4\pm 0.8$	$96.2\pm 0.5$	$97.1\pm 0.6$	$97.9\pm 0.5$	$98.1\pm 0.5$
equal-mass	96	$92.4\pm 0.6$	$95.1\pm 0.5$	$96.7\pm 0.5$	$97.6\pm 0.4$	$97.9\pm 0.5$
ours	96	$\mathbf{96.8\pm 0.4}$	$\mathbf{98.1\pm 0.2}$	$\mathbf{98.8\pm 0.2}$	$\mathbf{99.3\pm 0.2}$	$\mathbf{99.6\pm 0.1}$
equal-mass	480	$92.9\pm 0.6$	$95.6\pm 0.5$	$97.3\pm 0.4$	$98.2\pm 0.4$	$98.5\pm 0.4$

Table 2: AML, logistic regression from few labelled subjects. Test AUC (%), mean

\pm

standard error over the

25

folds, with ten draws of the training subjects per fold.

L=32

random slices (seed

0

N=10^{4}

cells. Normal subjects are added in the ratio of the dataset; “all” is the whole training fold, about

34

AML subjects. Stored values: number of values stored per subject. Linear SVM in Table 5 .

Panel of Figure 4 — Figure 4: ModelNet40, following ( Naderializadeh et al., 2021 ) : no backbone (top), and MLP backbone (bottom). Test accuracy (%), mean and standard deviation (error bars) over $3$ seeds. Equal-mass and ours: sketches with $M=4$ cells per slice; the number above each pair is ours minus equal-mass. Lines: GAP, PMA and FSPool, which use no slices, with $\pm$ one standard deviation shaded.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Result	Numerical check	Recorded CSV files
Theorem 2	Independent residual/boundary decomposition; both signs of the distance gap	certificates_atomic
Corollary 3	$H_{\rm cell}\leq H_{0}$ , signed bounds and complete distance interval	certificates_atomic
Theorem 4	Averaged signed and unsquared bounds; shared linear and nonlinear maps	certificates_sliced
Theorem 5	Pierce scaling, Gaussian sphere supremum, Zador constants	continuous_rates , sliced_rate
Proposition 11	Gaussian and all three ball-projection rate regimes	continuous_rates
Corollary 7	Known quantile Lipschitz constants and nonmatching cells	lipschitz

Appendix

Table 3: Theory-to-experiment correspondence. CSV files are in results/theory_validation/ ; all are regenerated by one command.

Figure 5: Exact atomic identity and certificate checks, and direction sampling. Left: independently computed sides of the identity coincide. Middle: the signed gap lies in

[-H_{\rm cell},S^{2}]

, rescaled to

[0,1]

. Right: Monte Carlo mean absolute errors with bars of two standard errors over

8,192

independent direction sets; the dashed line is the analytic squared-distance bound.

Figure 6: Continuous quantization rates. Top left: comparison with analytic Zador constants. Top right: Gaussian squared errors. Bottom left: equal-mass errors under the law-specific normalizations in Table 4 ; colors match the top-left legend. Bottom right: the cost of prescribing equal masses at the same

M

, on the common range

M\leq 512

Law	$M^{2}\varepsilon_{*}^{2}/C_{\rm Zador}$	Equal-mass normalization	Rescaled equal-mass error
Gaussian	0.99597	$M\log M$	1.07909
Ball d=1	1.00000	$M^{2}$	0.33333
Ball d=2	0.99939	$M^{2}$	0.40858
Ball d=3	0.99892	$M^{2}/\log M$	0.07224
Ball d=5	0.99827	$M^{10/6}$	0.11563
Ball d=9	0.99755	$M^{14/10}$	0.03175

Appendix

Table 4: Endpoints of the continuous rate experiment:

M=512

for optimal quantizers and

M=2^{18}

for equal-mass cells. The last column is the equal-mass squared error multiplied by the stated normalization; it is not claimed to equal its limiting constant.

Figure 7: Checks under analytic regularity assumptions: polynomial quantiles and their Lipschitz certificate; optimal widths and masses for

f_{+}

; and the squared-distance gap between

f_{+}

and

f_{-}

with the Lipschitz boundary budget. No empirical quantile is treated as Lipschitz.

Figure 8: Exact integration of the sharpness family. The expansion decays more slowly than the residual-only budget, its

M^{3/2}

rescaling approaches

-1

, and its ratio to

S^{2}

grows. The target

\nu_{M}

varies with

M

Figure 9: Moment checks. Left and middle: fourth-order scalar and tensor errors with the bounds of Proposition 6 ; corrected errors decrease faster than their

M^{-3}

bounds in this smooth example. Right: covariance errors before and after adding residual energies for an unbounded Gaussian.

Figure 10: Documents, learned slices: optimal sketch with

M

atoms against equal-mass cells with

M

cells, fixed

L

(rows), sweep over

M

; test error.

Figure 11: AML and ModelNet40, learned slices: optimal sketch with

M

atoms against equal-mass cells with

M

cells, fixed

L

, sweep over

M

; test accuracy.

Figure 12: Documents, learned slices: optimal sketch with

M

atoms against equal-mass cells with

M

cells, fixed

M

(rows), sweep over

L

; test error.

Figure 13: AML and ModelNet40, learned slices: optimal sketch with

M

atoms against equal-mass cells with

M

cells, fixed

M

, sweep over

L

; test accuracy.

		AML subjects in training
summary	stored values	2	4	8	16	all
uncompressed	320,000	$92.2\pm 0.6$	$94.9\pm 0.4$	$96.8\pm 0.4$	$97.9\pm 0.4$	$98.5\pm 0.4$
mean, covariance	35	$92.6\pm 0.7$	$95.4\pm 0.6$	$96.1\pm 0.7$	$97.0\pm 0.6$	$97.2\pm 0.5$
equal-mass	96	$91.2\pm 0.7$	$93.7\pm 0.4$	$95.7\pm 0.5$	$96.7\pm 0.5$	$97.8\pm 0.5$
ours	96	$\mathbf{96.1\pm 0.4}$	$\mathbf{97.4\pm 0.3}$	$\mathbf{98.4\pm 0.2}$	$\mathbf{99.0\pm 0.2}$	$\mathbf{99.4\pm 0.2}$
equal-mass	480	$92.0\pm 0.6$	$94.7\pm 0.4$	$96.7\pm 0.4$	$97.8\pm 0.4$	$98.4\pm 0.4$

Appendix

Table 5: AML, linear SVM from few labelled subjects. As Table 2 , with a class-balanced linear SVM in place of logistic regression. Stored values: number of values stored per subject.

Sliced Wasserstein Distance

Sliced Wasserstein distance compares probability distributions by averaging one-dimensional optimal transport costs over projections.

See all 17 papers

MaxSketch: Robust Distinct Counting in Streams via Random Projections
May 15, 2026Nikos Tsikouras, Constantine Caramanis, Christos TzamosRandomized Sketching Streaming Algorithms
National and Kapodistrian University of Athens, Greece · Archimedes, Athena Research Center, Greece · The University of Texas at Austin, USA
Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation
May 19, 2026Peter Matthew Jacobs, Jeff M. PhillipsRandomized Sketching Approximation Algorithms
Department of Statistics University of Wisconsin-Madison Madison, WI 53706 · Kahlert School of Computing University of Utah Salt Lake City, UT 84101
Highly Data Parallelizable Estimation of the Sliced-Wasserstein Distance Using Cumulative Distribution Functions
Jun 29, 2026Christophe Vauthier, Quentin Mérigot, Anna KorbaSliced Wasserstein Distance Optimal Transport
Laboratoire de Mathématiques d’Orsay, Université Paris-Saclay Gif-sur-Yvette, France · Centre de recherche en économie et statistique, ENSAE Palaiseau, France

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

MaxSketch: Robust Distinct Counting in Streams via Random Projections

Nikos Tsikouras, Constantine Caramanis, Christos Tzamos

National and Kapodistrian University of Athens, Greece · Archimedes, Athena Research Center, Greece · The University of Texas at Austin, USA

Randomized Sketching Streaming Algorithms

Open paper AskPDF arXiv

Estimating the number of distinct elements in a data stream is well understood when repeated elements are identical. In modern settings, however, observations are high-dimensional and noisy, so repeated instances of the same object are only approximately similar -- for example, different images of the same individual may vary significantly at the pixel level. Classical sketches such as HyperLogLog rely on consistent hash values for identical elements and break down in this regime. Recent work on robust distinct counting in general metric spaces achieves

\widetildeΘ(\sqrt{n})

memory, which is tight in the worst case. We show that substantially improved memory guarantees are possible under geometric structure common in learned representations. We introduce MaxSketch, a simple max-linear sketch built from random Gaussian projections, and prove that it succeeds in estimating the number of distinct latent objects. Concretely, we show that under this assumption

m = \widetilde{O} (\log n / \varepsilon^2)

random projections (and hence

\widetilde{O} (\log n/\varepsilon^2)

memory) suffice to recover the true distinct count within a

(1+\varepsilon)

factor. Experiments on image streams confirm that MaxSketch accurately estimates distinct counts and generalizes beyond the training regime. Our results bridge classical streaming algorithms and modern representation learning, showing how geometric structure can fundamentally reduce the complexity of distinct counting.

Abstract

Sliced Wasserstein Distance

Explore similar work

MaxSketch: Robust Distinct Counting in Streams via Random Projections

Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

Highly Data Parallelizable Estimation of the Sliced-Wasserstein Distance Using Cumulative Distribution Functions