Source-linked AI summary
Multilevel Clustering via Wasserstein Means
Nhat Ho, XuanLong Nguyen, Mikhail Yurochkin, Hung Hai Bui, Viet Huynh, Dinh Phung
TL;DR
The paper addresses simultaneous local and global clustering in hierarchically grouped data. It proposes Wasserstein-based joint optimization with barycenter-linked algorithms, establishes consistency, and demonstrates flexibility and scalability on synthetic and real data.
Problem
Multilevel clustering must partition data within groups while also discovering grouping patterns among the groups.
Method
The method jointly optimizes over discrete probability measures at multiple hierarchical levels using Wasserstein metrics, with variants that exploit Wasserstein barycenters.
Results
The paper establishes consistency for local and global cluster estimates and demonstrates efficiency and flexibility on simulated and real datasets.
Takeaways & Limitations
The approach directly solves multilevel clustering while also offering a fast approximation to nested Dirichlet process mixing-measure inference.
Takeaways & Limitations
The method assumes continuous data, requires the numbers of local and global clusters, and does not directly model noise distributions.
Abstract
from arXiv · showhide
We propose a novel approach to the problem of multilevel clustering, which aims to simultaneously partition data in each group and discover grouping patterns among groups in a potentially large hierarchically structured corpus of data. Our method involves a joint optimization formulation over several spaces of discrete probability measures, which are endowed with Wasserstein distance metrics. We propose a number of variants of this problem, which admit fast optimization algorithms, by exploiting the connection to the problem of finding Wasserstein barycenters. Consistency properties are established for the estimates of both local and global clusters. Finally, experiment results with both synthetic and real data are presented to demonstrate the flexibility and scalability of the proposed approach.
1. Introduction
The paper formulates multilevel clustering as a Wasserstein-based optimization problem that jointly discovers local clusters within groups and global clusters among groups. It develops efficient algorithms, consistency results, flexible variants, and empirical demonstrations.
- Motivation: Multilevel clustering jointly partitions data within each group and partitions collections of groups into global clusters.The motivating settings include words within documents and documents within corpora, as well as annotated regions within images.
- Approach: The paper takes a purely optimization-based approach to discovering hidden multilevel clustering structures in grouped data.The formulation inserts Wasserstein distances across different levels of the hierarchical data structure.
- Contributions: The proposed algorithms exploit a connection to Wasserstein barycenters and may provide fast approximations to inference under the nested Dirichlet process.The paper also establishes consistency theory for the multilevel clustering estimates.
- Contributions: The contributions include a new Wasserstein formulation, fast algorithms, consistency theorems, strength-sharing alternatives, and experiments on simulated and real data.The alternatives introduce constraints encouraging borrowing of strength among local and global clusters.
- Organization: The paper presents background, optimization formulations and algorithms, consistency results, simulation studies, and a concluding discussion.The stated experiments include both synthetic and real data.
2. Background
The background develops Wasserstein distances on probability measures and measures of measures, then connects Wasserstein barycenters and K-means quantization to the paper’s optimization framework.
- Probability-measure spaces: P(Θ) denotes the space of Borel probability measures on a subset Θ of R^d.The background introduces Wasserstein spaces as probability-measure spaces equipped with Wasserstein metrics.
- Wasserstein distances: The Wasserstein distance is the optimal mass-transport cost, with unit-mass movement charged according to the r-power of Euclidean distance.For finite discrete measures, the passage notes that fast computation is available.
- Measures of measures: Wasserstein metrics extend recursively to probability measures on probability-measure spaces, with transport cost determined by distances between supported measures.The same notation is used for both levels, with the level inferred from context.
- Wasserstein barycenters: A Wasserstein barycenter minimizes a weighted sum of Wasserstein distances to input probability measures.The barycenter problem is presented as a key computational connection for the paper’s method.
- Computational connection: Efficient local algorithms for Wasserstein barycenters over measures with bounded support provide building blocks for the proposed method.The background also notes applications of Wasserstein barycenters to approximate Bayesian inference.
- K-means and quantization: K-means can be viewed as a quantization problem that finds a finite discrete probability measure minimizing Wasserstein distance to an empirical distribution.This connection motivates using Wasserstein barycenter algorithms as an alternative to Lloyd’s algorithm for local minima.
3. Clustering with multilevel structure data
The paper formulates multilevel clustering as a joint problem that simultaneously learns local clusters within groups and global clusters among groups. Multilevel Wasserstein Means uses hierarchical Wasserstein objectives, efficient barycenter-based optimization, and variants that can encourage sharing across local clusters.
- The goal is to learn local clusters for each data group and global clusters among all groups simultaneously.
- MWM jointly optimizes local measures Gj and global measures H using distances at different levels of the hierarchical data structure.The formulation combines distances between empirical group measures and local measures with distances between local measures and global Wasserstein means.
- The original formulation can be simplified to an optimization involving only measures on Θ by exploiting the structure of H.
- Properties in special cases: When empirical measures from two groups are close, their optimal local solutions can also be close, whereas sufficiently separated groups need not share atoms.
- Algorithm: Algorithm 1 alternates updates of local measures and global means, using Wasserstein barycenter procedures for efficient local optimization.The algorithm is initialized using K-means-based procedures and uses barycenter algorithms to update local measures.
- Algorithm: Algorithm 1 monotonically decreases the MWM objective function.
- Sharing variant: MWMS adds constraints on local-measure supports to encourage component sharing and borrowing of strength among groups.
4. Consistency results
The paper establishes consistency results for the multilevel Wasserstein means estimators, showing convergence of estimated local and global structures under stated assumptions. When the population optimum is unique, the estimated local measures and global measure converge to that optimum.
- Consistency framework: The consistency analysis focuses on the MWM method, with convergence considered as every group sample size n_j tends to infinity.Consistency for MWMS is stated to be obtainable similarly.
- Consistency results: Theorem 4.1 establishes the first consistency property of the WMW formulation under P_j ∈ P_2(Θ).
- Consistency results: Theorem 4.2 proves almost-sure convergence of the estimated local and global quantities to the set of optimal solutions as n tends to infinity.The convergence is expressed through d(Ĝ_n, Ĥ_n, F) → 0 almost surely.
- Assumptions: The consistency theorem assumes that Θ is bounded and each P_j belongs to P_2(Θ).
- Unique optimum: When the optimal solution is unique, estimated local measures converge in W_2 to their population counterparts and the estimated global measure converges in W_2 to the unique optimum.
5. Empirical studies
Synthetic experiments compare MWM, MWMS, and baseline clustering across group structure and variance conditions, while real-data studies demonstrate multilevel clustering on LabelMe and StudentLife. The proposed algorithms outperform three-stage K-means when shared structure or non-constant variance is present and scale to large datasets.
- Synthetic data: Synthetic experiments vary group counts, group sizes, dimensionality, global clusters, local atoms, and variance or sharing structure.The comparison metric is Wasserstein distance between estimated and data-generating distributions.
- Synthetic data: Three-stage K-means performs best only for no-constraint data with constant variance, whereas MWM and MWMS outperform it when variance is non-constant or atoms are shared.The advantage is strongest for local-constraint data, including constant-variance settings where three-stage K-means fails to model shared structure.
- Real data: LabelMe contains 2,688 annotated images across 8 scene categories, and preprocessing retains 1,800 images with at least 4 annotated regions.Regions are represented by 30-dimensional PCA-projected GIST descriptors and clustered within images before image-level clustering.
- Real data: StudentLife comprises 49 students and approximately 4.6 million 500-dimensional WiFi data points, making it a large-scale multilevel clustering test.The algorithms require approximately 1 hour and discover global student groups together with local shared activity-location clusters.
- Real data: MWMS discovers five representative LabelMe image clusters whose accumulated tag-clouds are consistent with the grouped images.StudentLife results similarly expose student clusters, student-to-cluster assignments, and activity-location relationships.
6. Discussion
The discussion identifies scope boundaries and extensions for the Wasserstein multilevel clustering approach. These include discrete data, unknown cluster counts, robustness to noise, and richer hierarchical context.
- Limitations and extensions: The current formulation considers only continuous data and requires the numbers of local and global clusters to be known.The authors suggest extending it to discrete data and penalizing model complexity when cluster counts are unknown.
- Limitations and extensions: The formulation does not directly account for noise distributions away from Wasserstein means, motivating possible use of the first-order Wasserstein metric.The paper also proposes extending the method to hierarchical settings with group-level context.
Appendix A
Appendix A reviews Wasserstein distances and barycenters for discrete probability measures. It explains computational representations and notes analytical and computational simplifications for finite supports.
- Wasserstein distance: For finite discrete measures, Wasserstein distance is represented through an optimal transport coupling and a pairwise-distance cost matrix.The optimal coupling encodes transport between the two measures.
- Wasserstein distance: When the measures have equal finite support size k, best-known optimal-transport algorithms have complexity O(k^3 log k), while Sinkhorn approximations have complexity O(k^2).The regularized Sinkhorn formulation provides an approximation rather than the exact transport solution.
- Wasserstein barycenter: A Wasserstein barycenter combines probability measures using weights λ in the simplex and can be obtained through a multimarginal optimal-transport problem.The appendix introduces measure-preserving maps between the input measures in this construction.
- Wasserstein barycenter: The maps defining Wasserstein barycenters can be analytically intractable when the input measures lack special constraints.This limits direct closed-form treatment of the barycenter construction.
- Wasserstein barycenter: For finite discrete input measures with uniform weights, finding a barycenter over the large space P2(Θ) reduces to searching a smaller space.The appendix attributes this reduction to a sharp bound on the number of barycenter atoms.
Appendix B
Appendix B proves equivalences and convergence properties for the multilevel Wasserstein objectives and algorithms. The proofs establish monotone objective descent and consistency of estimated local and global mixing measures.
- Objective equivalence: Projection onto a closed subset of probability-measure space connects expected squared distance to Wasserstein distance from a measure to its projected image.The appendix constructs a coupling through the projection map to establish the relation.
- Objective equivalence: The multilevel Wasserstein means objectives are equivalent through lemmas relating projection, discrete support, and infima over local and global measures.The equivalence is summarized by C = D for the corresponding optimization formulations.
- Algorithmic convergence: Each algorithmic iteration does not increase the objective: f(G(t+1), H(t+1)) ≤ f(G(t), H(t)).The decrease follows from alternating Wasserstein-barycenter updates for local measures and optimization steps for global measures.
- Statistical consistency: As sample sizes grow, empirical objective values converge to population values, and estimated mixing measures approach the population solution set almost surely.The consistency argument combines convergence of empirical Wasserstein terms with the theorem for estimated objectives.
Appendix C
Appendix C details Algorithm 2 for Multilevel Wasserstein Means with Sharing (MWMS), including its initialization, alternating updates, and convergence guarantee. The algorithm monotonically decreases the MWMS objective function.
- Algorithm 2: The update of a(t+1)j is obtained as the closed-form minimum after differentiating the objective with respect to a(t)j.The updated measures are formed by replacing atoms of G(t)j with elements of S(t+1).
- Algorithm 2: The local weights of G(t+1)j are updated after its atoms are updated.The appendix separately identifies this as step (iii) of the MWMS update procedure.
- Algorithm 2: Algorithm 2 initializes a global set, local measures, and elements of H from the data and input parameters.Its inputs are Xj,i, K, and M; its outputs are the global set SK, local measures Gj, and elements Hi of H.
- Theoretical guarantee: Algorithm 2 monotonically decreases the objective function of the MWMS formulation.The proof establishes f(G(t+1), H(t+1)) ≤ f(G(t), H(t)) for every t ≥ 0.
Appendix D
Appendix D defines the notation used in the simulation-data generation process, including atom counts, an all-ones vector, and an identity matrix.
- Notation: The simulation notation includes Ki and kj for the numbers of supporting atoms of Hi and atoms of Gj, respectively.It also defines 1d as the d-dimensional all-ones vector and Id as the d-dimensional identity matrix.
Comparison metric (Wasserstein distance to truth)
The appendix uses a minimum-matching Wasserstein distance to compare estimated global measures with the truth and describes a multilevel Wasserstein means simulation setting. Global clusters are generated from Gaussian atoms centered at separated means, while each group generates local measures and observations from a selected global cluster.
- Comparison metric: The comparison metric is dM(Ĥ, H), a minimum-matching distance defined using pairwise Wasserstein distances between global measures.It is expressed as min over matchings of W2(Hi, Ĥj).
- Local measures and data: Each group selects a cluster label zj and samples local atom means τji from the corresponding global measure Hzj.Local atoms θji are then sampled from N(τji, Id), and observations Xj,i are sampled from N(µi, Id).
- Local measures and data: For non-constrained variances, the variance used to generate local atoms is set proportional to the global cluster label assigned to the local measure.This provides a variance condition for the simulation setting.
Multilevel Wasserstein means with sharing setting
The sharing setting generates global clusters and assigns shared atoms and local groups to global labels. It also describes a three-stage K-means procedure that estimates local measures, clusters their atoms globally, and estimates global elements.
- Global clusters: The sharing construction uses global atoms φij ∼ N(µi1d, Id), with µi := 5(i − 1) and weights πi ∼ Dir(1Ki).These define the global clusters used in the simulation setting.
- Sharing construction: Each shared atom receives a uniformly sampled global cluster label zk, while each group receives a uniformly sampled label ˜zj.The group selects shared atoms whose labels match its assigned label.
- Local measures and data: For non-constrained variances, the variance of local atoms θi is set proportional to the assigned global cluster label ˜zj.The local data-generation notation also samples observations as Xj,i ∼ N(µi, Id).
- Three-stage K-means: The three-stage K-means baseline first estimates each local measure, then clusters all local atoms into M groups, and finally estimates each Hi from its assigned atom group.The final stage uses exactly L clusters for each group of local atoms.