[Paper Review] Gradient-based training of Gaussian Mixture Models in High-Dimensional Spaces

Alexander Gepperth, Benedikt Pfülb|arXiv (Cornell University)|Sep 25, 2019

Gaussian Processes and Bayesian Inference21 references4 citations

TL;DR

This paper proposes a gradient-based training method for Gaussian Mixture Models (GMMs) in high-dimensional spaces using Stochastic Gradient Descent (SGD), leveraging a numerically stable lower bound on log-likelihood to enable efficient, memory-efficient training on large, non-stationary datasets. The approach achieves log-likelihood performance comparable to EM while supporting small batch sizes and reduced training time.

ABSTRACT

We present an approach for efficiently training GMMs solely with Stochastic Gradient Descent (SGD) on huge amounts of non-stationary, high-dimensional data. In such scenarios, SGD is superior to the traditionally Expectation-Maximization (EM) algorithm w.r.t. execution time and memory usage, and additional admits the use of small batch sizes. To use SGD in high-dimensional spaces, we propose to maximize a lower bound of a GMMs log-likelihood, which we prove to be feasible, justifiable by experiments and numerically stable. Since SGD seems more prone to get stuck in local optima than EM during early training phases, we introduce an annealing procedure that initially penalizes a large class of degenerate solutions before transitioning into a normal training regime. Experiments on several image datasets show that our approach is realizable, efficient and achieves comparable log-likelihood values as EM in a variety of scenarios. A TensorFlow implementation is provided to allow for reproduction.

Motivation & Objective

• To address the inefficiency of the Expectation-Maximization (EM) algorithm in high-dimensional, large-scale, non-stationary data settings.
• To enable end-to-end training of GMMs using Stochastic Gradient Descent (SGD) in high-dimensional spaces where EM is computationally prohibitive.
• To design a numerically stable objective function that serves as a lower bound on the true GMM log-likelihood for use in SGD.
• To mitigate SGD's tendency to converge to poor local optima during early training by introducing a structured annealing procedure.
• To demonstrate that SGD-based GMM training achieves log-likelihood performance comparable to EM on real-world image datasets.

Proposed method

• Maximize a lower bound of the GMM log-likelihood instead of the exact likelihood to ensure numerical stability during SGD optimization.
• Formulate the objective function such that it remains differentiable and suitable for mini-batch training in high-dimensional settings.
• Introduce an annealing schedule that initially penalizes degenerate solutions (e.g., collapsed components) to improve optimization stability in early training.
• Gradually reduce the penalty term over time, transitioning from a regularized regime to standard SGD training.
• Use the reparameterization trick and Monte Carlo sampling to estimate gradients of the lower bound efficiently.
• Implement the method in TensorFlow to support reproducibility and integration into deep learning pipelines.

Experimental results

Research questions

• RQ1 Can SGD be effectively used to train GMMs in high-dimensional spaces where EM is too slow and memory-intensive?
• RQ2 Is a lower bound on the GMM log-likelihood sufficient to achieve competitive performance compared to EM?
• RQ3 Does the proposed annealing procedure significantly improve convergence and avoid poor local optima in early training phases?
• RQ4 How does the proposed method scale in terms of training time and memory usage compared to EM on large, non-stationary datasets?
• RQ5 Can the method achieve log-likelihood values comparable to EM while supporting small batch sizes and online learning?

Key findings

• The proposed lower bound on the GMM log-likelihood is numerically stable and enables effective training via SGD in high-dimensional settings.
• The annealing procedure successfully reduces the risk of converging to degenerate solutions during early training, improving optimization reliability.
• The method achieves log-likelihood values on image datasets that are comparable to those obtained with the EM algorithm.
• Training with SGD is significantly faster and more memory-efficient than EM, especially on large-scale, non-stationary data.
• The approach supports small batch sizes and is suitable for online learning, making it practical for real-world deployment.
• A public TensorFlow implementation is provided, enabling full reproducibility and integration into existing machine learning workflows.