PaperlayerView exact revision on arXiv

Source-linked AI summary

Sliced Wasserstein Distance for Learning Gaussian Mixture Models

Soheil Kolouri, Gustavo K. Rohde, Heiko Hoffmann

arXiv:1711.05376v2cs.CVcs.LGstat.ML

TL;DR

EM can converge to poor stationary points, motivating alternatives for estimating GMM parameters. This paper fits GMMs by minimizing sliced Wasserstein distance, and reports greater empirical robustness than KL-divergence-based methods including EM.

  • Problem

    GMM parameter estimation commonly uses EM, which guarantees only convergence to a stationary point and can be sensitive to random initialization.

    Source: §1 · p. 1

  • Method

    The paper estimates GMM parameters by minimizing sliced p-Wasserstein distance between the mixture model and the data distribution.

    Source: §1 · pp. 1–2

  • Results

    Empirical results suggest the method is more robust than KL-divergence-based methods, including EM, for maximizing log-likelihood.

    Source: §5 · p. 8

  • Takeaways & Limitations

    The paper presents sliced p-Wasserstein distance as a framework for fitting GMMs, rather than only comparing GMM models.

    Source: §1 · p. 2

    The method is not shown to be convex or to have fewer local minima, and comparable smooth-distribution approximation can require O(Ld−1) projections in dimension d.

    Source: §2.3 · p. 3, §5 · p. 8

Abstract

from arXiv · show

Gaussian mixture models (GMM) are powerful parametric tools with many applications in machine learning and computer vision. Expectation maximization (EM) is the most popular algorithm for estimating the GMM parameters. However, EM guarantees only convergence to a stationary point of the log-likelihood function, which could be arbitrarily worse than the optimal solution. Inspired by the relationship between the negative log-likelihood function and the Kullback-Leibler (KL) divergence, we propose an alternative formulation for estimating the GMM parameters using the sliced Wasserstein distance, which gives rise to a new algorithm. Specifically, we propose minimizing the sliced-Wasserstein distance between the mixture model and the data distribution with respect to the GMM parameters. In contrast to the KL-divergence, the energy landscape for the sliced-Wasserstein distance is more well-behaved and therefore more suitable for a stochastic gradient descent scheme to obtain the optimal GMM parameters. We show that our formulation results in parameter estimates that are more robust to random initializations and demonstrate that it can estimate high-dimensional data distributions more faithfully than the EM algorithm.

1. Introduction

This paper focuses on estimating GMM parameters, an important task because EM can be sensitive to initialization. It proposes fitting GMMs by minimizing sliced Wasserstein distance instead of relying on likelihood-based methods.

1. Introduction

  • GMMs are flexible parametric density models used across machine learning, computer vision, and signal and image analysis.Examples include image representation, point-set registration, and time-series classification.Source: §1 · p. 1
  • The paper focuses on estimating mixture-component parameters, distinct from the separate problem of choosing the number of components.Source: §1 · p. 1
  • EM guarantees convergence only to a stationary point and can reach bad critical points from random initializations, making its estimates sensitive to initialization.The cited work reports that EM can converge to a bad critical point with high probability under random initialization.Source: §1 · p. 1
  • The paper proposes estimating GMM parameters by minimizing sliced p-Wasserstein distance between the data distribution and the mixture model.In the infinite-i.i.d.-sample limit, minimizing negative log-likelihood is equivalent to minimizing KL divergence between the data distribution and the GMM.Source: §1 · pp. 1–2
  • Random projections reduce high-dimensional distributions to one-dimensional slices, whose Wasserstein distances are minimized with respect to GMM parameters.The paper presents this as a way to address the computational burden of Wasserstein minimization for high-dimensional distributions.Source: §1 · p. 2
  • Prior uses of p-Wasserstein distances for GMMs compare models, whereas this work uses sliced distances as a framework for fitting a GMM to data.Source: §1 · p. 2

2. Preliminary

The preliminary section defines Wasserstein distance through optimal transport, introduces the Radon transform for forming projections, and defines sliced Wasserstein distance using one-dimensional marginals. It also notes the increased projection requirements for approximating distributions as dimensionality grows.

2. Preliminary

  • The p-Wasserstein distance is an optimal transport cost between probability measures, with plans constrained to reproduce each measure’s marginals.For absolutely continuous measures, the text also gives an equivalent transport-map formulation.Source: §2.1 · p. 2

2.3. Sliced p-Wasserstein Distance

  • In one dimension, Wasserstein distance has a closed-form solution, which supports computing sliced distances through one-dimensional transport problems.Source: §2.1 · p. 2, §2.3 · p. 3

2.2. Radon transform

  • The Radon transform maps a function to its integrals over hyperplanes, and a projected multivariate Gaussian remains a one-dimensional Gaussian.Consequently, a projected GMM is a one-dimensional GMM with projected component means and variances.Source: §2.2 · pp. 2–3

2.3. Sliced p-Wasserstein Distance

  • Sliced p-Wasserstein distance compares distributions by aggregating Wasserstein distances between their one-dimensional linear-projection marginals.The text characterizes it as a true metric; approximating comparable smooth distributions requires O(Ld−1) projections in dimension d when L suffice in two dimensions.Source: §2.3 · p. 3

3. Sliced Wasserstein Means and Gaussian Mixture Models

The section extends Wasserstein means from K-means to fitting general GMMs, then motivates sliced Wasserstein optimization as a more tractable alternative to likelihood fitting. In simple scenarios, the Wasserstein means landscape shares the global optimum with likelihood while being smoother or less sensitive to initialization.

3. Sliced Wasserstein Means and Gaussian Mixture Models

  • The Wasserstein means formulation extends the connection from K-means to GMMs by fitting a parametric mixture distribution to the empirical data distribution.The resulting optimization estimates the GMM component means, covariances, and mixture weights.Source: §3 · pp. 3–4

3.1. Wasserstein Means vs. Maximum Log-Likelihood

  • Unlike KL and related divergences, Wasserstein distance is continuous and differentiable in GMM parameters and can compare distributions across spatial bins.The cited motivation notes that KL-type divergences can have limited capture range when distributions lie on low-dimensional manifolds.Source: §3.1 · p. 4
  • Both scenarios share the same global optimum for likelihood and Wasserstein means; the Wasserstein landscape is less initialization-sensitive, and in scenario 2 it is smoother with no local minima.The landscapes in Figure 1 are scaled and shifted for visualization.Source: Figure 1 · p. 4, §3.1 · p. 5

3.2. Sliced Wasserstein Means

  • Sliced Wasserstein means approximate the GMM-to-data Wasserstein objective by comparing random one-dimensional projections, avoiding high-dimensional transport optimization.For each projection, the transport map is uniquely defined and has a closed-form solution.Source: §3.2 · p. 5
  • Randomly sampling projection directions at each step yields a stochastic-gradient, EM-like scheme that alternates transport-map calculation and GMM parameter updates.The method uses RMSProp; component derivatives and projection transport maps are independent, allowing parallelization.Source: §3.2 · pp. 5–6

4. Numerical Experiments

Experiments compare SW-GMM with EM-GMM on a two-dimensional distribution and on MNIST and CelebA. SW-GMM is more robust across initializations and performs better on the reported high-dimensional evaluations.

4.1. Robustness to initialization

  • The robustness experiment uses a ring, square, and connecting line, with K = 10 modes and 100 random initializations for both algorithms.The same three datasets used across the experiments are the two-dimensional distribution, MNIST, and CelebA.Source: §4.1 · p. 6, §4 · p. 6
  • SW-GMM reaches the optimal negative log-likelihood in 100% of 100 random initializations, compared with 29% for EM-GMM.Both negative log-likelihood and sliced-Wasserstein distance were lower for SW-GMM than for EM-GMM in this experiment.Source: §4.1 · p. 7

4.2. High-dimensional datasets

  • The MNIST and CelebA comparisons use GMMs fitted in learned autoencoder embedding spaces, with generated samples decoded into images.The full pipeline is unsupervised.Source: §4.2 · p. 7
  • On MNIST, SW-GMM components average 86.98% purity, versus 80.48% for EM-GMM components.Purity was measured by generating 10,000 samples from each component and classifying them; each component ideally represents a single digit.Source: §4.2 · p. 7
  • Figure 5 provides a qualitative comparison of generated CelebA component samples, and the authors explicitly distinguish these visual comparisons from their quantitative fitness measures.The samples are contrast enhanced for visualization.Source: §4.2 · p. 7, Figure 5 · p. 8
  • SW-GMM consistently achieves higher discriminator fooling rates than EM-GMM, with a significant advantage on MNIST and CelebA but reasonable fits from both methods on the low-dimensional dataset.The fooling rate is used as a measure of GMM fitness.Source: §4.2 · p. 8, Figure 6 · p. 8

5. Discussion

The discussion presents sliced p-Wasserstein minimization as a GMM-fitting method based on random one-dimensional projections. Empirical results indicate greater robustness than EM, while the method's convexity and local-minimum properties remain unproven.

5. Discussion

  • Across a two-dimensional distribution and the high-dimensional MNIST and CelebA datasets, experiments indicate greater robustness than KL-divergence-based methods, including EM.The authors also report more well-behaved p-Wasserstein energy landscapes.Source: §5 · p. 8
  • Each iteration uses random projections to create one-dimensional distributions and estimates GMM parameters from those projections.This is the discussion's description of the sliced p-Wasserstein procedure.Source: §5 · p. 8
  • The authors provide no theoretical guarantee that the method is convex or has fewer local minima, and report using deep convolutional encoders for two datasets while noting applicability to other embeddings, including original data space.The method's empirical robustness does not establish those theoretical properties.Source: §5 · p. 8

7. Supplementary material

The supplement covers the KL-divergence framing, RMSProp updates, experimental implementation details, and generated GMM component samples.

7.1. Maximum log-likelihood and KL-divergence

  • The supplement defines KL divergence between Ix and Iy and considers maximum log-likelihood in the limit as N→∞.Source: §7.1 · p. 11

7.2. RMSProp update equations

  • RMSProp updates GMM parameters using a learning rate, decay parameter, and momentum; its first and second moments are based on the parameter gradient.Similar update equations apply to Σk and αk.Source: §7.2 · p. 11

7.3. Experimental Details

  • The autoencoder uses equally weighted reconstruction and adversarial losses; adversarial networks and classifiers use cross-entropy, with RMSProp training for 100 epochs at batch size 250.Source: §7.3 · p. 11
  • The supplement documents convolutional autoencoders for MNIST and CelebA and deep binary classifiers used to score GMM fitness.Source: §7.3 · p. 11, Figure 7 · p. 12, Figure 8 · p. 13

7.4. CelebA Generated Images

  • The figures show GMM-generated samples from EM-GMM and SW-GMM, with each column representing samples from one component.Source: §7.4 · p. 11, Figure 9 · p. 13
Loading 1711.05376v2…