PaperlayerView exact revision on arXiv

Source-linked AI summary

EMS Coreset: An Efficient Expectation-Maximization Algorithm for Sinkhorn Coreset

Haoyun Yin, Chuanhui Liu, Xiao Wang

arXiv:2608.16101v1stat.MLcs.LG

TL;DR

OT-based coreset selection is computationally expensive because it repeatedly solves transport problems, limiting scalability. EMS replaces explicit OT computation with variational EM updates and achieves comparable or improved approximation quality while being over 100× more efficient than gradient-based OT coresets.

  • Problem

    Existing OT-based coreset methods repeatedly solve expensive transport problems, limiting their practicality for large-scale datasets.

  • Method

    EMS uses variational EM updates, including a closed-form softmax E-step and decoupled M-step updates, to construct weighted Sinkhorn coresets efficiently.

  • Results

    Across diverse data structures and downstream tasks, EMS improves distributional accuracy and stability while achieving comparable performance to explicit gradient-based OT coresets at over 100× greater efficiency.

  • Takeaways & Limitations

    EMS provides a fast and stable weighted coreset construction method with theoretical consistency and fidelity to the original data.

  • Takeaways & Limitations

    The approximate OT plan introduces bias relative to the true OT plan, creating a robustness–fidelity trade-off.

Abstract

from arXiv · show

Coresets distill large datasets into small, representative subsets for efficient downstream learning. Yet Optimal Transport (OT)-based selection typically requires intensive computation of transport plans, limiting scalability. We introduce a scalable Sinkhorn coreset method that permits closed-form updates of the entropically regularized OT coupling by allowing non-uniform coreset weights. This produces centroids that generalize k-means via soft assignments. We establish asymptotic consistency of the selected measure and Lipschitz stability to data perturbations, providing accuracy and robustness guarantees. Across synthetic and real-world benchmarks, the proposed method achieves competitive or improved approximation quality while substantially reducing runtime compared to Wasserstein- and standard Sinkhorn-based coreset selection, especially at large scale.

1 Introduction

The introduction motivates EMS by the high computational cost of OT-based coreset construction and presents an EM-style alternative based on efficient variational updates. It also states theoretical guarantees for weighted Sinkhorn coresets and empirical gains in efficiency, accuracy, stability, and scalability.

  • Motivation: Coresets reduce large datasets to small representative subsets while preserving data geometry for downstream clustering, regression, and inference tasks.They aim to address the computational and memory demands of large-scale learning.
  • Motivation: OT-based coresets preserve geometric structure by minimizing transport cost to the full dataset, connecting them to barycenters, representative points, and moment-matching herding.Unlike query-space sub-sampling, this approach directly uses the geometry of the observed data.
  • Computational bottleneck: Existing OT-based coreset methods are impractical at scale because Sinkhorn-distance evaluation requires repeated OT computations, with Wasserstein plans scaling cubically and near-quadratically for Sinkhorn methods.The introduction contrasts this cost with classical coreset construction times that are linear or near-linear for (1 + ε)-approximations.
  • EMS approach: EMS replaces explicit OT computation with EM-style iterative updates: the E-step computes a variational OT plan under a Gibbs distribution, while the M-step updates the coreset representation.The method minimizes a variational upper bound on the Sinkhorn distance.
  • Contributions: EMS combines weighted Sinkhorn-coreset theory, a variational soft-assignment formulation, and evaluations showing comparable OT-based performance while being over 100× more efficient.The stated guarantees include geometry preservation, approximation error, convergence, stability, and asymptotic consistency; a mini-batch variant further improves practical scalability.

2 Related Work

Prior work develops coresets through sensitivity-based sampling, Bayesian likelihood approximation, and optimal-transport formulations. Related perspectives also include barycenters, representative points, and variational or Schrödinger interpretations of entropic OT.

  • Sensitivity-based Coreset Selection: Classical coresets sample existing data points by sensitivity, while Bayesian coresets use sparse weighted subsets to approximate full-data log-likelihoods.These subsets provide efficient surrogate posteriors for inference.
  • Optimal-Transport Coresets: OT-based coresets can extend beyond finite point sets and pointwise target objectives, admitting (1 + ϵ)-approximation guarantees over a function class.This follows from Kantorovich–Rubinstein duality.
  • Barycenters and Representative Points (RPs): Wasserstein barycenters and representative points both compress probability measures, but offer conceptually complementary views centered on transport-cost optimization and representative support.The proposed EMS algorithm has smooth, closed-form gradient updates analogous to free-support Wasserstein barycenters, with a variational and optimization-theoretic derivation.
  • Variational, Schrödinger, and EM views of entropic OT: Entropic OT admits KL-penalized projection, iterative Bregman projection, and Schrödinger bridge interpretations, supporting differentiability, stability, and statistical interpretations of Sinkhorn couplings and costs.These views connect transport plans with Gibbs measures and motivate the analysis of entropic OT.

3 Sinkhorn Coreset

This section formulates Sinkhorn coresets as discrete, potentially non-uniform weighted approximations that preserve the target measure’s geometry. It establishes finite-sample and asymptotic guarantees and explains why adaptive weighting corrects uniform-weight k-means bias.

  • Problem formulation: Sinkhorn coresets approximate a probability measure µ with a discrete measure having at most k weighted atoms.The construction targets discrete empirical measures while extending naturally to absolutely continuous measures.
  • Sinkhorn objective: For λ > 0, Sinkhorn regularization makes the transport objective convex in the transport plan and enables efficient computation through Sinkhorn iterations.The regularized optimum is biased and can be interpreted as a discrete-time Schrödinger bridge.
  • Sinkhorn objective: The coreset objective shifts selection from obtaining an optimal coupling to preserving distributional geometry while summarizing µ with k << n atoms.This formulation treats the coreset as a geometry-preserving approximation under the Sinkhorn metric.
  • Consistency guarantees: Theorem 3.1 provides a finite-sample Sinkhorn-coreset upper bound under bounded-density, compact-support, and finite-(p + δ)-moment assumptions.Its proof combines a Wasserstein quantization bound from deterministic Voronoi assignments with a Taylor-expansion bound on KL divergence.
  • Consistency guarantees: As k →∞, global minimizers converge to µ in Wasserstein distance, establishing asymptotic consistency and recovery of the target’s underlying geometry.The finite-sample guarantee matches the classical quantization rate.
  • Weighted correction: Uniform-weight k-means generally fails to converge to µ, whereas non-uniform weights remove this bias and motivate adaptive weighting in Sinkhorn coresets.Uniform-weight k-means is consistent only when the target distribution has uniform density on its support; the weighted correction recovers µ.

4 An Efficient Algorithm for Sinkhorn Coresets

EMS reformulates weighted Sinkhorn coreset optimization as an EM-style procedure using variational transport plans. Its closed-form soft assignments and monotone centroid-weight updates provide efficient optimization, while softmax smoothing yields stability with a regularization-dependent bias–optimality trade-off.

  • Variational formulation: EMS minimizes a variational objective over transport-plan distributions, using the expected transport plan as a surrogate instead of solving the exact optimal plan.This expectation is the averaged transport plan ¯T = Eq∗[T].
  • E-step: Under a factorized reference measure, the E-step produces closed-form vector-valued softmax assignments that can be computed efficiently.The resulting averaged plan may differ from the exact Sinkhorn-optimal plan but remains tractable.
  • M-step: The M-step updates coreset locations and weights as minimizers of the free energy, with locations given by weighted centers of mass and weights fitted to soft assignments.These updates correspond to maximum-likelihood estimates under the current transport-plan assignments.
  • M-step: Fλ(Y (t+1), W (t+1), ¯T (t+1)) ≤Fλ(Y (t), W (t), ¯T (t+1)), so the free energy is non-increasing after each M-step.Equality holds only when both parameter blocks satisfy first-order optimality conditions.
  • Stability: Larger λ tightens the soft-assignment gradient bound and improves robustness to perturbations, but introduces bias relative to the true OT plan.Uniform-weighted k-means does not share this stability property.

5 Numerical Experiments

Numerical experiments show that EMS matches or improves competing coresets in distributional accuracy and downstream SHAP fidelity while substantially reducing runtime. Its performance remains stable across regularization choices and random seeds, with weighted assignments improving robustness over uniform-weight baselines.

  • Distributional accuracy: Experiments compare EMS with WCSL and uniform-weight K-means using entropic Wasserstein distances and wall-clock time on Gaussian mixtures and MNIST.The benchmarks vary coreset size and report the quality-cost trade-off.
  • Distributional accuracy: EMS achieves distributional accuracy similar to WCSL, while K-means lags on the image task because uniform weights cannot represent mass flexibly.WCSL directly minimizes Sinkhorn loss with gradient descent, making it much slower.
  • Runtime scaling: The EMS family scales near-linearly with data size, coreset size, and dimension, with OT-plan complexity O(nkd) and substantially lower runtime than WCSL.At n=10^6, unregularized EMS finishes in under ten seconds, while WCSL requires thousands of seconds; entropic variants remain in the single-digit-second band.
  • Regularization sensitivity: EMS maintains similar distributional accuracy across λ ∈ {10^-3, 10^-2, 10^-1} and λ=0, while runtime increases smoothly as λ grows.The softmax E-step changes responsibility sharpness, whereas the M-step recomputes weighted centroids and mixture weights.
  • Robustness and interpretability: Across seeds, EMS produces concentrated coresets, and EMS with λ=0.01 achieves the lowest SHAP MeanAE on both evaluated datasets.K-means degrades on separated mixtures, while WCSL remains seed-sensitive; on Census, K-means has slightly lower MaxAE but larger mean distortion and base-value drift.

6 Conclusion … A.5 Proof of Lemma 4.1

EMS provides a fast, robust weighted-coreset method with closed-form softmax E-steps and decoupled M-step updates, while the theoretical results establish convergence and density properties for Sinkhorn coresets. The proofs further characterize quantization, manifold, and Gibbs-measure behavior underlying these guarantees.

  • 6 Conclusion: EMS combines a closed-form softmax E-step with decoupled M-step updates for fast and stable Sinkhorn-loss optimization.The approach maintains strong fidelity to the original data and supports downstream interpretability such as Shapley value computation while running orders of magnitude faster.
  • A Proofs: The proofs establish the theoretical framework supporting EMS’s OT-based coreset guarantees.The subsequent subsections derive quantization, convergence, density, and Gibbs-measure properties.
  • A.1 Proof of Theorem 3.1: For an optimal k-point quantizer, the typical Voronoi diameter scales as hk ∼ k−1/d, and the proof bounds Wasserstein error using Voronoi cells and a Gibbsized local coupling.The resulting bound contains the classical quantization rate and a second-order entropic-bias term.
  • A.2 Proof of Corollary 3.2: As k →∞, both classical quantization error and entropic KL correction vanish, so Sinkhorn coresets converge to the underlying µ in Wasserstein distance.This is the asymptotic consistency conclusion derived in the corollary proof.
  • A.3 High-resolution asymptotics for uniform k-means: Under high-resolution regularity, asymptotically optimal k-means centroids admit a limiting codepoint density g⋆.The assumptions include bounded continuous density, shrinking cell diameters, negligible boundary effects, and asymptotically regular local cell shapes.
  • A.4 Proof of Proposition 3.4: The weighted Voronoi measures νk^w weakly converge to µ as k →∞.The proof uses weak convergence of centroid locations to g⋆ and continuity of y 7→p(y)/g⋆(y).
  • A.5 Proof of Lemma 4.1: The proof of Lemma 4.1 shows Ff,P0 is convex in f and identifies the unique Gibbs maximizer q∗(T) ∝ ef(T)P0(T).The argument represents expectations and KL divergence through the Radon–Nikodym derivative and applies Jensen’s inequality.

A.6 Proof of Corollary 4.2

Corollary 4.2 follows by applying Jensen’s inequality to the convex transport-cost and KL-divergence terms, yielding the regularized optimal transport objective evaluated at the averaged plan. The bound is tight under deterministic collapse or when the averaged plan reaches the global Sinkhorn minimizer, increasingly so as λ decreases.

  • Proof: Convexity of the transport cost and KL divergence enables Jensen’s inequality for the expected transport plan ¯T := E_q∗[T].This produces the corollary’s key inequality.
  • Proof: The resulting right-hand side is F_λ(Y, W, ¯T), the regularized optimal transport objective whose minimum over T_k defines the Sinkhorn loss.F_λ(Y, W, ¯T) = ⟨C, ¯T⟩_µ + λ D_KL(π_¯T ∥ µ ⊗ ν_k).
  • Equality conditions: Equality in the Jensen step holds if and only if T is almost surely constant under q∗, including P_0 = δ(T − T_0) or λ → 0.As λ approaches zero, the Gibbs measure collapses to a single deterministic transport plan.
  • Equality conditions: In the zero-temperature limit, D_KL(q∗∥P_0) vanishes and the free energy becomes a pure energy functional.The limiting collapse eliminates stochasticity in the transport plan.
  • Equality conditions: The second inequality is tight when ¯T coincides with the global minimizer T∗, and approximately tight when q∗ concentrates near T∗.As λ decreases, q∗ becomes increasingly peaked around T∗, tightening the free-energy bound.

A.7 Proof of Theorem 4.3

Theorem 4.3 links entropy-regularized transport’s free-energy formulation to mean-field variational inference. Under a factorized prior, the optimal Gibbs posterior factorizes across locations, yielding a closed-form softmax E-step for transport assignments.

  • Mean-field factorization and conditional independence: A factorized prior P0 over spatial locations implies that the optimal Gibbs posterior q∗ inherits the same factorization.The optimization decouples across x and reduces the free-energy functional to independent local contributions.
  • Mean-field factorization and conditional independence: The KL divergence decomposes additively under this factorization, establishing it as a sufficient condition for the variational optimization.This additive decomposition verifies consistency of the factorized solution.
  • Softmax as the variational expectation: The expected transport plan T̄(x) is a softmax over clusters, weighted by prior masses ωj and penalized by costs Cj(x) with temperature λ.As λ → 0, the soft assignment approaches a hard assignment and recovers the deterministic optimal transport map.
  • Relation to variational inference: The factorized free-energy minimization is a mean-field variational EM E-step, with q∗ as the posterior and T̄ as its expected sufficient statistic.Theorem 4.3 identifies the softmax update in (17) as a closed-form E-step linking Gibbs variational principles, regularized transport, and probabilistic clustering.

A.8 Proof for Proposition 4.5 · A.9 Proof of Proposition 4.6

The proofs derive the closed-form centroid and weight updates for the M-step, establish its objective decrease, and characterize equality. They also prove the soft-assignment gradient identity and a Lipschitz bound valid for non-uniform weights.

  • A.8 Proof for Proposition 4.5: With the E-step transport plan fixed, Proposition 4.5 analyzes the Sinkhorn Loss Fλ(Y, W, ¯T (t+1)) in the M-step.This sets the objective and fixed transport plan used throughout the proof.
  • A.8 Proof for Proposition 4.5: For each cluster j, the M-step optimizes with respect to yj, differentiates, sets the gradient to zero, and obtains the update of Y (t+1).The centroid update follows from the first-order condition for each cluster.
  • A.8 Proof for Proposition 4.5: The proof also optimizes each weight ωj through its linear KL-divergence term, yielding the derivative condition and the next-step mass assigned to cluster j.The weight derivation is performed analogously to the centroid derivation.
  • A.8 Proof for Proposition 4.5: The M-step is a block coordinate optimization over Y and W, so updating either block cannot increase the Sinkhorn objective.The two blockwise inequalities are stated explicitly in equations (83) and (84).
  • A.8 Proof for Proposition 4.5: Combining the blockwise inequalities gives Fλ(Y (t+1), W (t+1), ¯T (t+1)) ≤Fλ(Y (t), W (t), ¯T (t+1)).Equality holds only when the previous Y and W already minimize both blocks and satisfy the first-order conditions.
  • A.9 Proof of Proposition 4.6: For Proposition 4.6, the proof uses the squared–Euclidean cost Cj(x) = ∥x −yj∥2 with p = 2 and recalls the soft assignment from Eq. (17).Applying the product and chain rules establishes the stated gradient identity, equation 22.
  • A.9 Proof of Proposition 4.6: The gradient satisfies ∥∇x ¯Tj(x)∥≤(2/λ) ¯Tj(x) DY ≤2DY /λ, where DY := maxr,s ∥yr −ys∥, and the bound holds for non-uniform W.The weights ω enter only through ¯T and cancel in the derivative formula.

A.10 Discussion for Stability … B.2 Experiment Details for Sec. 5.1

The paper establishes stability guarantees for soft assignments and empirical M-step updates, contrasting them with hard k-means discontinuities. It also specifies reproducible implementation settings and benchmarks on synthetic Gaussian data and MNIST.

  • A.10 Discussion for Stability: Soft assignments are Lipschitz continuous in data, while hard k-means assignments can jump by Θ(1) across Voronoi boundaries.Theorem A.3 gives global data stability for fixed λ > 0; Remark A.7 contrasts this with the absence of a finite global Lipschitz constant for hard assignments.
  • A.10 Discussion for Stability: Soft assignments are locally Lipschitz in centers on compact data domains, with uniform bounds derived from partial derivatives and bounded data and center norms.The result applies when x lies in a compact set K and centers remain bounded.
  • A.10 Discussion for Stability: The empirical M-step admits a per-sample leverage bound when one sample is perturbed, using assignment stability to control changes in sufficient statistics and updated centers.The perturbation is Xk 7→ Xk + h with ∥h∥ ≤ δ.
  • B Extended Experimental Setup Details: The experimental setup reports wall-clock times from GPU-aware routines on an Intel Core i7-12700K workstation with an NVIDIA RTX 5090 GPU.Unless noted otherwise, computations use the cuda backend.
  • B.1 Implementation Details: Experiments use Ubuntu 22.04, Python 3.11.5, PyTorch 2.7.1+cu128, CUDA 12.8, and supporting scientific Python libraries.The listed environment also includes TorchVision, NumPy, SciPy, Pandas, Matplotlib, Seaborn, and scikit-learn versions specified in the setup.
  • B.1 Implementation Details: Unless stated otherwise, methods use max iter=1000, eps=10^-2, batch size=1000, random dataset initialization, and Sinkhorn regularisation parameter 0.01.WCSL additionally uses Adam with learning rate 10^-1 and a StepLR scheduler multiplying the learning rate by 0.99 each epoch.
  • B.2 Experiment Details for Sec. 5.1: The experiments generate n=10^4 Gaussian points in d=100 and evaluate MNIST using its standard train/test split.The Gaussian covariance is assembled from a randomly sampled orthonormal basis.

B.3 Experiment Details for Sec. 5.5

This experiment tests whether a size-64 coreset preserves SHAP interpretability relative to the full training set by comparing attribution reconstruction errors. It uses standardized Census and Diabetes data, a shared gradient-boosted tree model with TreeExplainer, and five randomized runs.

  • Goal and high-level protocol: The evaluation compares SHAP outputs using the full training set versus a k=64 coreset as the background distribution.Attributions are computed on the test split, and reconstruction errors are reported between the reference and approximation outputs.
  • Datasets, splits, and preprocessing: Census classification and Diabetes regression use 80/20 train/test splits, with stratification for classification and training-fitted standardization applied to both splits.The setup follows Jethani et al. [2022] and the shap package datasets.
  • Predictive model and SHAP explainer: Each dataset uses a gradient-boosted tree trained on the 80% split, with TreeExplainer computing attributions under a shared model while only the background distribution changes.Unless otherwise noted, the model uses 400 estimators, learning rate 0.05, and maximum depth 3.
  • Background distributions: full data vs. coreset: EMS learns both coreset support points and weights, whereas K-means and WCSL use learned support points with uniform weights and Random samples uniformly weighted training points.All methods construct discrete weighted measures supported on k atoms.
  • SHAP matrices and reconstruction metrics: The study reports three SHAP reconstruction errors, including maximum absolute attribution deviation and base-value error capturing preservation of the explainer’s reference level.The SHAP base value is the background expected prediction in the decomposition f(x) = b(ν) + Σj ϕj(x; ν).
  • Repetitions and randomness control: Reported results average over five runs with independent random seeds, accounting for data-splitting and/or coreset-initialization randomness.The passage directs readers to table 4 for the results.

B.4 Limitations and Future Directions: Large-Scale Image Datasets

EMS is not evaluated on large-scale image datasets because squared Euclidean and pixel-space entropic OT discrepancies may not reflect semantic similarity. A proposed extension applies EMS in learned feature space, while addressing image realization and semantic evaluation.

  • EMS is not evaluated on large-scale image datasets because its squared Euclidean geometry and entropic OT costs suit tabular and controlled synthetic data better than images.Pixel-wise similarity is generally misaligned with semantic similarity in images.
  • Direct Wasserstein or Sinkhorn objectives in pixel space can produce visually or semantically meaningless summaries despite small transport costs.
  • Applying EMS to a pushforward measure µψ in learned feature space could yield a compact weighted measure whose atoms summarize dataset geometry.ψ may be a fixed or lightly tuned feature extractor, such as a pretrained vision encoder.
  • Future work must realize feature-space atoms as images and evaluate semantic fidelity with downstream performance or representation-level metrics.Candidate realization methods include nearest-neighbor selection and synthetic-image optimization under perceptual losses with appropriate priors.
Loading 2608.16101v1…