PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Discrete Distribution Clustering Using Wasserstein Barycenter with Sparse Support

Jianbo Ye, Panruo Wu, James Z. Wang, Jia Li

arXiv:1510.00012v4stat.COcs.LGstat.ML

TL;DR

D2-clustering is hard to scale because computing its Wasserstein barycenters is costly. The paper adapts B-ADMM for approximate barycenter computation and reports that accelerated D2-clustering often clearly outperforms other widely used methods in their respective domains.

  • Problem

    D2-clustering's centroid computation is a major scalability challenge, while Wasserstein distance computation is costly.

    Source: §I · p. 1, §II · p. 4

  • Method

    The paper uses modified B-ADMM for centroid optimization and separates the update into per-distribution constrained quadratic subproblems.

    Source: A. Alternating Direction Method of… · p. 5, §VIII · p. 14

  • Results

    Accelerated D2-clustering often clearly outperforms other widely used methods in the respective data domains.

    Source: §VIII · p. 14

  • Takeaways & Limitations

    The resulting clustering tool is easy to use and requires no tuning of optimization parameters.

    Source: §VIII · p. 14

    The modified B-ADMM algorithm's theoretical convergence properties are not well understood, and the work emphasizes data-size scalability over support-point scalability.

    Source: §VIII · p. 14

Abstract

from arXiv · show

In a variety of research areas, the weighted bag of vectors and the histogram are widely used descriptors for complex objects. Both can be expressed as discrete distributions. D2-clustering pursues the minimum total within-cluster variation for a set of discrete distributions subject to the Kantorovich-Wasserstein metric. D2-clustering has a severe scalability issue, the bottleneck being the computation of a centroid distribution, called Wasserstein barycenter, that minimizes its sum of squared distances to the cluster members. In this paper, we develop a modified Bregman ADMM approach for computing the approximate discrete Wasserstein barycenter of large clusters. In the case when the support points of the barycenters are unknown and have low cardinality, our method achieves high accuracy empirically at a much reduced computational cost. The strengths and weaknesses of our method and its alternatives are examined through experiments, and we recommend scenarios for their respective usage. Moreover, we develop both serial and parallelized versions of the algorithm. By experimenting with large-scale data, we demonstrate the computational efficiency of the new methods and investigate their convergence properties and numerical stability. The clustering results obtained on several datasets in different domains are highly competitive in comparison with some widely used methods in the corresponding areas.

I. INTRODUCTION

Discrete distributions represent complex data, but Wasserstein-based D2-clustering is difficult to scale because computing cluster centroids is costly. The paper develops and evaluates scalable barycenter methods, including a modified B-ADMM approach, for this setting.

I. INTRODUCTION

  • Wasserstein-based D2-clustering extends K-means to discrete distributions, but costly centroid computation has limited its use in large-scale learning.Discrete distributions also represent sparse histograms and bags of words without requiring a shared fixed codebook across instances.Source: §I · p. 1

B. Related Work

  • Prior approaches include divide-and-conquer and standard ADMM, while the paper modifies B-ADMM because a direct adaptation does not properly solve the barycenter problem.The modified approach is not guaranteed to reach a local optimum, though it often produces a solution close to one and empirically exceeds IBP accuracy for finite supports.Source: B. Related Work · pp. 2–3

A. Our Contributions

  • Modified B-ADMM jointly optimizes barycenter support locations and weights, and AD2-clustering handles many instances when each support has only several hundred points.The paper also develops a parallel version with adequate scaling efficiency on hundreds of CPUs.Source: A. Our Contributions · pp. 1–2, B. Related Work · p. 2

B. Related Work

  • Optimizing support locations can avoid a high-dimensional fixed support set, while using fewer barycenter points can reduce computation; experiments suggest large support sizes are usually unnecessary.The authors also note that tuning robustness across changing cluster partitions matters alongside the speed of computing a single centroid.Source: B. Related Work · pp. 2–3

II. DISCRETE DISTRIBUTION CLUSTERING

D2-clustering alternates between assigning distributions to their nearest centroids and updating those centroids. The update is the main computational challenge because it requires solving a Wasserstein barycenter problem.

II. DISCRETE DISTRIBUTION CLUSTERING

  • D2-clustering minimizes total within-cluster squared Wasserstein variation by alternating nearest-centroid assignment with centroid updates.Each centroid is a discrete distribution with weights and support points.Source: §II · pp. 3–4
  • Each centroid update seeks weights and support points minimizing summed squared Wasserstein distances to the cluster's N distributions.The update is the inner loop; unlike K-means, the optimal centroid is not available in simple closed form.Source: §II · p. 4
  • A full-batch LP becomes impractical as cluster size grows, motivating approximate scalable methods because the matching variables dominate the update.The weights are part of the output centroid, whereas matching weights are not, so a sufficiently accurate LP solution can be adequate for an outer-loop update.Source: §II · p. 4

III. SCALABLE CENTROID COMPUTATION

The paper proposes scalable methods for the centroid-update inner loop, including subgradient descent, standard ADMM, and modified B-ADMM. These methods aim to make centroid computation comparable in cost to cluster assignment.

III. SCALABLE CENTROID COMPUTATION

  • The proposed scalable centroid-update methods are subgradient descent with mini-LPs, standard ADMM with mini-QPs, and modified B-ADMM with closed-form inner-loop iterations.The methods are compared in speed and memory.Source: §III · pp. 4–5
  • The methods target the inner loop, identified as D2-clustering's computational bottleneck: updating each centroid.The inner loop is the centroid-update step in Algorithm 2.Source: §III · p. 5
  • The new methods aim to reduce centroid-update computation to a level comparable to or below label assignment, which was usually negligible in the original algorithm.The paper also describes measures to expedite the labeling step.Source: §III · p. 5

A. Alternating Direction Method of Multipliers

The ADMM formulation relaxes equality constraints into augmented Lagrangians and alternates updates over transport plans and barycenter weights. Splitting transport-plan updates by member gives linear complexity in the number of distributions, but its quadratic subproblems motivate exploring B-ADMM.

A. Alternating Direction Method of Multipliers

  • ADMM relaxes equality constraints into augmented Lagrangians and optimizes the split variables over convex sets.The formulation has multiple equality and inequality constraints, unlike the typical two-variable ADMM setting.Source: A. Alternating Direction Method of… · p. 5
  • Splitting the transport-plan update into N independent constrained QPs gives computational complexity linear in N.This is the main motivation for applying ADMM to the clustering problem.Source: A. Alternating Direction Method of… · p. 5
  • The centroid routine repeatedly updates transport plans, weights, and dual variables until convergence, with ρ and iteration count as hyperparameters.The method empirically chooses ρ proportional to averaged transportation costs.Source: A. Alternating Direction Method of… · p. 5
  • ADMM uses costlier QP subproblems than subgradient descent’s LPs, and whether faster convergence offsets that cost is unclear.This computational limitation motivates B-ADMM, which avoids QPs in each iteration.Source: A. Alternating Direction Method of… · p. 5

B. Bregman ADMM

B-ADMM replaces quadratic augmented Lagrangians with Bregman divergence to update split variables in the barycenter problem. Its consensus update avoids a numerically unstable geometric mean, but convergence is unproved for this formulation and the centroid optimization is non-convex when support points are updated.

B. Bregman ADMM

  • B-ADMM replaces quadratic augmented Lagrangians with Bregman divergence when updating split variables.The approach adapts prior designs to solve the optimal-transport problem with many support points.Source: B. Bregman ADMM · pp. 5–6
  • The geometric-mean consensus update is numerically unstable when some component values approach zero, so the method uses an alternative update rule.The alternative rule is described as overcoming this numerical instability.Source: B. Bregman ADMM · p. 6
  • Neither tested consensus update rule ensures convergence to a local minimum.Experiments found little difference in converged objective values among the tested rules, including the Fisher-Rao Riemannian center.Source: B. Bregman ADMM · p. 7
  • Convergence of B-ADMM is unproved for this formulation, and centroid optimization is non-convex when support points are updated.Empirically, the method usually converges quickly to moderate accuracy; the implementation uses 100 B-ADMM iterations by default.Source: B. Bregman ADMM · p. 7

IV. COMPARISON BETWEEN B-ADMM AND ITERATIVE BREGMAN PROJECTION

The section compares IBP and B-ADMM for sparse-support Wasserstein barycenters, examining convergence, objective quality, and practical limits. The comparison is specific to barycenters embedded in D2-clustering, not a general ranking of optimal-transport solvers.

Comparison between B-ADMM and iterative Bregman projection

  • IBP and B-ADMM use different mirror-map updates: IBP projects in the primal simplices, whereas B-ADMM projects in a product space that also includes support variables.The B-ADMM move is not monotonic, unlike standard ADMM; the paper reports that its convergence is therefore not evident from conventional optimization results.Source: §IV · p. 7
  • IBP has linear convergence and more accurately satisfies coupling constraints than B-ADMM at the same computational time, but those advantages concern the stand-alone OT problem.The study does not establish which algorithm is better for OT in general; in barycenter optimization, both methods are embedded in an outer loop.Source: §IV · p. 7
  • At 5,000 iterations, B-ADMM’s transport output is nearly indistinguishable from the unregularized LP solution, while IBP at ε = 0.1/N remains noticeably different.Smaller IBP ε can cause double-precision overflow; the tested thresholding workaround produced incorrect coupling artifacts.Source: §IV · pp. 7–8, Figure 1 · p. 8
  • The experiments compare fixed-support and jointly optimized support-and-weight tracks, using support sizes m = 6 and m = 60 for the sparse-support setting.The barycenters are evaluated against the objective obtained by solving the full-batch LP.Source: §IV · p. 8, Table I · p. 9
  • In both experimental tracks, B-ADMM achieves lower distortion than IBP and stays closer to the full-batch LP, while IBP can be faster when ε₀ is relatively large.The comparison uses barycenters with fixed shared support locations in one track and optimized locations and weights in the other.Source: §IV · pp. 8–9
  • The modified B-ADMM is approximate and does not reach the exact local minimum obtained by full LP, so it cannot fully replace subgradient descent for that goal.The passage identifies subgradient descent as minimizing the objective to a true local minimum.Source: §IV · p. 9

V. ALGORITHM INITIALIZATION AND IMPLEMENTATION

The implementation initializes sparse centroid supports and B-ADMM couplings, then uses distance pruning and data partitioning to support efficient serial and parallel execution. It also describes fixed inner-loop iteration counts and extensions to multiple distributions per instance.

Algorithm initialization and implementation

  • Centroid initialization sets support size m to the cluster’s average and selects a member with at least m supports, merging pairs until the target size is reached.The merge replaces two support vectors with their weighted mean and combined weight, choosing the pair by an optimal criterion.Source: §V · p. 9
  • For unchanged cluster assignments, B-ADMM warm-starts each coupling from its cached value; otherwise, it initializes the coupling from the corresponding weights.The scheme also applies in the first outer iteration, when no previous solution is available.Source: §V · p. 9
  • Triangle-inequality pruning skips unnecessary Wasserstein-distance calculations during relabeling, where the naive step requires N̄K distances.The paper identifies this step as a sizable cost in its algorithms.Source: §V · p. 9
  • A fixed inner-loop count is used for simplicity; B-ADMM typically uses tens to a hundred iterations, while subgradient descent and ADMM can use around or below ten.The outer-loop count is determined adaptively by a termination criterion.Source: §V · p. 9
  • The algorithms parallelize by partitioning data across processors; centroid updates are mostly separate per-instance problems, and synchronization communicates centroid updates.The implementation uses MPI Allreduce, with evenly partitioned data and negligible synchronization time reported.Source: §V · p. 10
  • Instances may contain arrays of distributions, whose centroid updates are processed separately and whose label-assignment distance combines their squared Wasserstein distances.For symbolic data, support-vector updates can be skipped because the symbol set is fixed; grid-based supports can be stored as indices.Source: §V · p. 10

VI. COMPLEXITY AND PERFORMANCE COMPARISONS

The section analyzes clustering and centroid-update complexity, memory use, and algorithm selection. It recommends modified B-ADMM as the default pipeline while noting its accuracy and convergence trade-offs.

Complexity and performance comparisons

  • For cluster assignment, triangle-inequality pruning reduces empirical complexity from O(εₒN̄K l(N̄m,d)) to roughly O(εₒ(N̄ + K²)l(N̄m,d)).Here l(N̄m,d) is the average time to solve a Wasserstein distance between distributions in a d-dimensional metric space.Source: §VI · p. 10
  • For serial centroid updates, the stated B-ADMM complexity is O(εₒεᵢN̄md/τ) + O(εₒεᵢN̄m̄m), and the implementation’s memory allocation is O(N̄m̄m) + O(dN̄m̄ + dKm).Serial centroid-update complexity is described as depending on data size N̄, not cluster count K.Source: §VI · p. 10
  • ADMM is significantly slower than the other two methods on large datasets because it requires many constrained QP subproblems; B-ADMM performs consistently in convergence and clustering quality.The paper notes ADMM’s convergence properties are better understood, while B-ADMM’s convergence theory is weaker.Source: §VI · p. 10
  • Modified B-ADMM is recommended as the default pipeline for scalability, stability, and speed, assuming large memory is available.The paper does not treat it as a universal replacement because strong theoretical convergence results are lacking.Source: §VI · p. 10
  • When highly accurate centroids are needed, the paper recommends subgradient descent initialized with centroids first obtained by B-ADMM.B-ADMM resets its Lagrangian multipliers each inner round and performs a fixed number of iterations, rather than aggressively pursuing high accuracy at every round.Source: §VI · p. 10

VII. EXPERIMENTS

The experiments examine convergence, stability, computational and memory efficiency, scalability, and clustering quality on large datasets from several domains. The dataset collection includes synthetic, image, digit, and document data.

Experiments

  • The experiments assess convergence, stability, computational and memory efficiency, scalability, and clustering quality across large datasets from several domains.Table II summarizes dataset size, support-vector dimension, centroid support count, and maximum tested cluster count.Source: §VII · p. 10, Table II · p. 11
  • Synthetic distributions use multivariate-normal support vectors with added heavy-tailed Student’s t noise and perturbed, normalized Dirichlet probabilities; they are used only for scalability studies.The excerpt says details of the synthetic-data generation are omitted for lack of space.Source: §VII · pp. 10–11
  • The real-data experiments include image color or texture distributions, USPS pixel-location histograms, and document collections including 20 Newsgroups, BBC news abstracts, and Wiki events abstracts.The BBC and Wiki sets retain only titles and first sentences to test a more challenging document or sentence analysis setting.Source: §VII · p. 11

A. Convergence Analysis

Modified B-ADMM is numerically stable across tested hyperparameters and converges at a rate comparable to subgradient descent, while ADMM is much slower on large datasets. AD2-clustering also scales well across hundreds of processors in strong and weak scaling tests.

A. Convergence Analysis

  • Modified B-ADMM has the best numerical stability across tested hyperparameters while retaining a CPU-time convergence rate comparable to subgradient descent; ADMM is much slower on large datasets.The stability comparison covers modified B-ADMM, ADMM, and subgradient descent on four datasets.Source: A. Convergence Analysis · p. 11
  • B-ADMM converges comparably across ρ0 values from 0.5 to 16.0, with similar residual-ratio patterns across datasets; Figure 2 plots objective versus CPU time and dual versus primal residuals.The plotted objective is the B-ADMM objective in Eq. (15), not the original clustering objective in Eq. (2).Source: A. Convergence Analysis · p. 11, Figure 2 · p. 12

B. Efficiency and Scalability

  • The scalability experiments address both strong and weak scaling in a multi-core setting, where AD2-clustering can be CPU-bound or memory-bound.The experimental configuration keeps each iteration within one hour and fits within typical cluster memory.Source: B. Efficiency and Scalability · p. 11
  • AD2-clustering scales well in strong scaling tests using 32–512 processors and shows good weak scalability as processor counts increase.The strong-scaling problem uses N̄ = 250000, d = 16, m = 64, and k = 16; weak scaling is tested with respect to N̄ and m.Source: Table III · p. 11, B. Efficiency and Scalability · p. 12

C. Quality of Clustering Results

Across the tested domains, accelerated AD2-clustering often outperforms widely used methods, including on USPS digits and short documents. Its performance depends on the data representation, and theoretical convergence of modified B-ADMM remains insufficiently understood.

C. Quality of Clustering Results

  • For noisy USPS quantization, AD2-clustering consistently lowers test classification error versus Kmeans++, especially with many clusters and high blankout rates.Figure 3b plots test error against cluster count and blankout rate.Source: C. Quality of Clustering Results · p. 12, Figure 3 · p. 13
  • On USPS, AD2-clustering outperforms Kmeans++ in homogeneity and completeness across tested cluster counts.The comparison is shown in Figure 3a; homogeneity and completeness trade off by construction.Source: C. Quality of Clustering Results · p. 12, Figure 3 · p. 13
  • AD2-clustering achieves the best AMI and ARI on both 20Newsgroups datasets, with no further observed gain for m ≥64.The reported comparisons use the GV and WV vocabularies.Source: Table IV · p. 13, C. Quality of Clustering Results · p. 13
  • On short documents, AD2-clustering beats bag-of-words-only and embedding-only methods; random embeddings yield AMI 0.187 on BBC news abstract and 0.369 on Wiki events abstract, versus 0.759 and 0.545 with carefully trained embeddings.AD2-clustering takes several minutes on an 8-core machine for these sparse-support datasets.Source: C. Quality of Clustering Results · p. 13, Table V · p. 14

VIII. CONCLUSIONS AND FUTURE WORK

  • Across several real-world datasets, accelerated D2-clustering often clearly outperforms widely used methods in the respective domains.The paper’s conclusion presents this as an empirical pattern, not a universal guarantee.Source: §VIII · p. 14
  • The theoretical convergence property of modified B-ADMM is not well understood, and the work prioritizes scalability with respect to data size over support-point count.The authors identify support-count efficiency as an area for further improvement.Source: §VIII · p. 14

APPENDIX

The subgradient centroid method treats the barycenter objective as a bilevel optimization and uses LP-derived subgradients to update weights. It alternates weight and support updates, while a time-allocation protocol limits inner-loop work during clustering.

A. Subgradient Descent Method

  • The centroid objective is recast as a bilevel optimization over weights and transport variables, with LP dual solutions providing weight subgradients.With fixed support points, updating the weights is a linear program.Source: A. Subgradient Descent Method · p. 15
  • A predetermined step size replaces per-iteration line search, which would require repeated LP solves across cluster members.The method reparameterizes weights as normalized exponentials to avoid negative components and enforce their sum-to-one constraint.Source: A. Subgradient Descent Method · p. 15
  • With fixed support, weight optimization is convex; updating support points makes the full centroid problem non-convex.The method alternates updates to weights and support points, with τ specifying their relative update frequency.Source: A. Subgradient Descent Method · pp. 15–16
  • The protocol caps each inner centroid update at ηTa/K because centroid accuracy and hyperparameter choices do not guarantee faster outer-loop objective decrease.Ta is assignment-step time and K is the number of clusters; the protocol accounts for the assignment step's role in overall progress.Source: A. Subgradient Descent Method · p. 16
  • Under η = 2.0, subgradient descent completes at most two inner iterations, while B-ADMM averages more than 60 on color and texture data and more than 20 on protein data.Standard ADMM cannot complete a single iteration within this time allocation.Source: A. Subgradient Descent Method · p. 16

B. Stability analysis of AD2-clustering

Across four datasets, B-ADMM at a fixed parameter yields lower clustering objective values over time than subgradient descent. The latter requires step-size tuning and shows numerical instability on image color data.

B. Stability analysis of AD2-clustering

  • B-ADMM achieves consistently lower clustering objective values over time and converges substantially faster on protein sequence 1-gram data than fine-tuned subgradient descent.Figure 4 plots clustering objective versus CPU time for four datasets, with K = 10 and η = 2.0.Source: Figure 4 · p. 16, B. Stability analysis of AD2-clustering · p. 16
  • With fixed ρ0 = 2.0, B-ADMM converges well on all four datasets, while subgradient descent needs manual step-size tuning for comparable convergence speed.The subgradient step size was selected from eight values based on the lowest clustering objective in the first round.Source: B. Stability analysis of AD2-clustering · p. 16
  • On image color data, the fine-tuned subgradient step size still produces noticeable later-round objective fluctuations, making stability and speed difficult to balance.The step size was chosen using early-round performance.Source: B. Stability analysis of AD2-clustering · p. 16

Related papers

Loading 1510.00012v4…