Gradient-Based Training of Gaussian Mixture Models for High-Dimensional Streaming Data

This contribution focuses on Gaussian Mixture Model (GMM), which represent a probabilistic unsupervised model for clustering and density estimation, allowing sampling and outlier detection. GMMs have been used in a wide range of scenarios, see [18]. Commonly, free parameters of a GMM are estimated using the Expectation-Maximization (EM) algorithm [7], which does not require learning rates and automatically enforces all GMM constraints. A popular online variant is stochastic EM [2], which can be trained mini-batch wise and is thus more suited for large datasets or streaming data.

GMMs are probabilistic models that try to explain the observed data X \(=\) \(\{{{\varvec{x}}}_n\}\) by expressing their density as a weighted mixture of K Gaussian component densities \({\mathcal {N}}({{\varvec{x}}};{{\varvec{\mu }}}_k,{{\varvec{P}}}_k)\) \(\equiv \) \({\mathcal {N}}_k({{\varvec{x}}})\):Here, we parameterize Gaussian densities by precision matrices \({{\varvec{P}}}_k\) \(=\) \({{\varvec{\Sigma }}}_k^{-1}\) instead of covariances \({{\varvec{\Sigma }}}_k\). The component weights \(\pi _k\) represent another set of GMM parameters, which modulate the overall influence of each Gaussian density. For a derivation of Eq. 1, we must introduce the probabilistic foundations of GMMs. GMMs assume that each observed data sample \(\{{{\varvec{x}}}_n\}\) is drawn from one of the Gaussian component densities \({\mathcal {N}}_k\). The selection of this component density is assumed to depend on an unobserved (and unobservable) latent variable \(z_n\in \{1,\dots , K\}\) which follows an unknown distribution. This is formalized for a GMM with K components by formulating the complete-data likelihood for a single data sample as:Since the latent variables are, by construction, unobservable, we must marginalize them out of Eq. 2 in order to obtain an expression suitable for optimization. This gives us the incomplete-data likelihood for a single data sample \({{\varvec{x}}}_n\):which depends on observable quantities only. Please compare this to Eq. 1. The incomplete-data likelihood for all samples is thus given by:where we have inserted Eq. 2 in the last step. Passing to the log-domain (as it is common for probabilistic models), we obtain the total incomplete-data log-likelihood for all observed data samples:The function \({\mathcal {L}}\) contains only observable quantities and is a suitable loss function for optimization. For convenience and numerical stability, the sum is usually replaced by an expectation, and we follow this convention:Please note that \({\mathcal {L}}\) represents the likelihood of the observed data under the GMM with current parameters, and must therefore be maximized to obtain a better explanation of the data.

Unless stated otherwise, the experiments in this section will be conducted with the following parameter values for sEM and SGD (where applicable): mini-batch size B \(=\) 1, K \(=\) 8 \(\times \) 8, \(\mu ^i\) \(=\) 0.1, \(\sigma _0\) \(=\) 2, \(\sigma _\infty \) \(=\) 0.01, \(\epsilon \) \(=\) 0.001, \(D_{\text {max}}\) \(=\) 20. Each experiment is repeated 10 times with identical parameters but different random seeds for parameter initialization. See Sect. 5 for a justification of these choices. Due to high input dimensionalities, all precision matrices are assumed to be diagonal. The training/test data comes from the datasets shown below (see Sect. 3.1).

The EM algorithm assumes that the observed data samples \(\{{{\varvec{x}}}_n\}\) depend on unobserved latent variables \({{\varvec{z}}}_n\) in a non-trivial manner, see Sect. 2. The derivation of the EM algorithm starts out with the total incomplete-data log-likelihoodDue to the assumption that \({\mathcal {L}}\) is obtained by marginalizing out the latent variables, an explicit dependency on \(z_n\) can be re-introduced. For the last expression, Jensen’ inequality can be used to construct a lower bound:Since the realizations of the latent variables are unknown, we can assume any form for their distribution. In particular, for the choice \(p(z_n)\sim p({{\varvec{x}}}_n,z_n)\), the lower bound becomes tight. Simple algebra and the fact that the distribution \(p(z_n)\) must be normalized gives us:where we have used Eq. 2 in the last step. \(p(z_n\!=\!k|{{\varvec{x}}}_n)\) is a quantity that can be computed from data with no reference to the latent variables. For GMM it is usually termed responsibility and we write it as \(p(z_n\!=\!k|{{\varvec{x}}}_n)\!\equiv \!\gamma _{nk}\).

However, the construction of a tight lower bound, which is actually different from \({\mathcal {L}}\), only works when \(p({{\varvec{x}}}_n,z_n)\) depends non-trivially on the latent variable \(z_n\). If this is not the case, we have \(p({{\varvec{x}}}_n,z_n)\) \(=\) \(K^{-1}p({{\varvec{x}}}_n)\) and the derivation of Eq. 16 goes down very differently:where \({\mathcal {H}}\) represents the Shannon entropy of \(p({{\varvec{z}}})\). The highest value this can have is \(\log K\) for an uniform distribution of the \(z_n\), finally leading to a lower bound for \({\mathcal {L}}\) ofwhich is trivial by Jensen’s inequality, but not tight. In particular, no closed-form solutions to the associated extremal value problem can be computed.

This shows that optimizing GMM by EM assumes that each sample has been drawn from a single component in a set of K uni-modal Gaussian distributions. Which distribution is selected for sampling depends on a latent random variable. On the other hand, optimization by SGD uses the incomplete-data log-likelihood \({\mathcal {L}}\) as basis for optimization, without assuming the existence of hidden variables at all. This may be advantageous for problems where the assumption of Gaussianity is badly violated, although empirical studies indicate that optimization by EM works very well in a very wide range of scenarios.

The Relevance of this Article is outlined by the fact that training GMMs by SGD was recently investigated in the community by [13, 14]. We go beyond, since our approach does not rely on off-line data-driven model initialization, and works for high-dimensional streaming data. The presented SGD scheme is simple and very robust to initial conditions due to the proposed annealing procedure, see Sects. 3.2 and 3.3. In addition, our SGD approach compares favorably to the reference model for online EM [2] in terms of achieved log-likelihoods, which was verified on multiple real-world datasets. Superior SGD performance is observed for the high-dimensional datasets.

The Analysis of the Results suggests that SGD performs better than sEM on average, see Sect. 3.5, although the differences are modest. It is neither expected, nor is it the goal to show here, that sEM is outperformed by SGD. Instead, we aim at achieving a similar performance. However, if sEM is used without an initialization, e.g., k-means, components may not converge (see Fig. 4) for high-dimensional data like Fruits 360 and SVHN datasets. In such cases, SGD does outperform sEM. Another important advantage of SGD over sEM is the fact that the only parameter that needs to be tuned is the learning rate \(\epsilon \). In contrast to that, sEM has a complex and non-intuitive dependency on \(\rho _0\), \(\rho _\infty \) and \(\alpha _0\).

Low-Dimensional Data can be processed as well using our SGD-based approach, which makes numerical issues and undesirable local minima less relevant. Thus, we recommend optimizing the full incomplete-data log-likelihood in this case. Although there is no real need to use the max-component approximation and annealing for low-dimensional data, it is nevertheless possible. To this effect, we performed experiments with synthetic data drawn from various 2D Gaussian mixture distributions, using the same parameters as in the experiments of Sect. 3. We always observed rapid convergence if the initialization range of the centroids, \(\mu ^i\), fits the ranges of the individual data components. This ensures that initial centroids cover the data space sufficiently so that each cluster in the data has at least one centroid close to it.

No Assumptions About Data Generation are made by SGD in contrast to the EM and sEM algorithms. The latter guarantees that the loss will not decrease due to an M-step. This, however, assumes a non-trivial dependency of the data on an unobservable latent variable (shown in Sect. 4). In contrast, SGD makes no hard-to-verify assumptions, which is a rather philosophical point, but may be an advantage in certain situations where data are strongly non-Gaussian.

Numerical Stability is assured by our SGD training approach. It does not optimize the log-likelihood but its max-component approximation. This approximation contains no exponentials at all, and is well justified by the results of Table 3 which shows that component probabilities are strongly peaked. In fact, it is the gradient computations where numerical problems occurred, e.g., NaN values. The “logsumexp” trick mitigates the problem, but does not eliminate it (see Sect. 2.2). It cannot be used when gradients are computed automatically, which is what most machine learning frameworks do.

Hyper-Parameter Selection Guidelines are as follows: the learning rate \(\epsilon \) must be set by cross-validation (a good value is 0.001). We empirically found that initializing precisions to the cut-off value \(D_\text {max}\) and an uniform initialization of the \(\pi _i\) are beneficial, and that centroids are best initialized to small random values. A value of \(D_\text {max}=20\) always worked in our experiments. Generally, the cut-off must be much larger than the inverse of the data variance. In many cases, it should be possible to estimate this roughly, even in streaming settings, especially when samples are normalized. For density estimation, choosing higher values for K leads to higher final log-likelihoods. For clustering, K should be selected using standard techniques for GMMs. The parameter \(\delta \) controls loss stationarity detection for the annealing procedure and was shown to perform well for \(\delta \) \(=\) 0.05. Larger values will lead to a faster decrease of \(\sigma (t)\), which may impair convergence. Smaller values are always admissible but lead to longer convergence times. The annealing time constant \(\alpha \) should be set to the GMM learning rate \(\epsilon \) or lower. Smaller values of \(\alpha \) lead to longer convergence times since \(\sigma (t)\) will be updated less often. The initial value \(\sigma _0\) needs to be large in order to enforce convergence for all components. A typical value is \(\sqrt{K}\). The lower bound on \(\sigma _\infty \) should be as small as possible in order to achieve high log-likelihoods (e.g., 0.01, see Sect. 2.6 for a proof).

The presented work can be extended in several ways: First of all, annealing control could be simplified further by inferring good \(\delta \) values from \(\alpha \). Likewise, increases of \(\sigma \) might be performed automatically when the loss rises sharply, indicating a task boundary. As we found that GMM convergence times grow linear with the number of components, we will investigate hierarchical GMM models that operate like a Convolutional Neural Network (CNN), in which individual GMM only see a local patch of the input and can therefore have low K.