Fitting large mixture models using stochastic component selection
Year: 2021

Abstract:
Traditional methods for unsupervised learning of finite mixture models require to 1 evaluate the likelihood of all components of the mixture. This becomes computa2 tionally prohibitive when the number of components is large, as it is, for example, 3 in the sum-product (transform) networks. As a remedy, we propose an approach 4 combining the expectation maximization and the Metropolis-Hastings algorithm 5 to evaluate only a small number of, stochastically sampled, components, thus 6 substantially reducing the computational cost. We put emphasis on generality of 7 our method, equipping it with the ability to train both shallow and deep mixture 8 models which involve complex, and possibly nonlinear, transformations. The 9 performance of our method is illustrated in a variety of synthetic and real-data 10 contexts, considering deep models, such as mixtures of normalizing flows and 11 sum-product (transform) networks. 12

text: Traditional methods for unsupervised learning of finite mixture models require to1 evaluate the likelihood of all components of the mixture. This becomes computa-2 tionally prohibitive when the number of components is large, as it is, for example,3 in the sum-product (transform) networks. As a remedy, we propose an approach4 combining the expectation maximization and the Metropolis-Hastings algorithm5 to evaluate only a small number of, stochastically sampled, components, thus6 substantially reducing the computational cost. We put emphasis on generality of7 our method, equipping it with the ability to train both shallow and deep mixture8 models which involve complex, and possibly nonlinear, transformations. The9 performance of our method is illustrated in a variety of synthetic and real-data10 contexts, considering deep models, such as mixtures of normalizing flows and11 sum-product (transform) networks.12
1 Introduction13
Finite mixture models [40] constitute a fundamental class of density estimation models. They14 have been successfully applied in diverse fields, including bioinformatics [49], econometrics [10],15 engineering [33], etc. A mixture model relies on a weighted sum of probability distributions—here16 referred to as components—to cluster N unlabelled datapoints into K categories. The traditional17 maximum likelihood techniques train the model by optimizing either (i) the marginal likelihood via18 gradient-descent [50] or (ii) the evidence lower bound via variational methods [4], including the19 expectation-maximization (EM) [13]. The dependence structure among approximate, variational,20 distributions then ranges from the fully independent (mean-field) [25] to fully dependent [30]. The21 sampling-based techniques target the posterior distribution using sequential Monte Carlo [9] or22 Markov chain Monte Carlo [52], e.g. via the Gibbs [34] or Metropolis-Hastings sampling [38]. The23 computational cost of these methods typically scales with O(TKND) operations, where N and K24 are defined above, T is the number of iterations and D is the dimension of data.25
Various methods to decrease the computational cost via any factor inO(TKND) have been proposed.26 T can be lowered by proper initialization, e.g. the optimal seeding [5]; an efficient step-size schedule,27 e.g. the line-search [58]; or increased estimation precision, e.g. the variance reduction [8]. N is often28 reduced using the coreset methods, which approximate the original dataset by a weighted dataset29 such that the exact and approximate marginal likelihoods are close. The weighted variants of the30 variational [17, 59, 6] and sampling-based [39] methods then process the coresets. Reducing D relies31 on the compression of data into smaller representations via random projections [53, 2], which is32 achieved in two ways: (i) each data item is projected into an individual representation [11]; (ii) all33 data items are projected into an overall representation, commonly referred to as sketch [28, 22].34
Nevertheless, all the aforementioned techniques—including those with reduced computational cost—35 evaluate all K components. This is very demanding for large models, and the problem is even more36
Submitted to 35th Conference on Neural Information Processing Systems (NeurIPS 2021). Do not distribute.
severe for mixtures involving intricate models, such as neural networks [21, 42], Gaussian processes37 [57], normalizing flows [48]; or deep mixtures, including sum-product (transform) networks [45, 47],38 deep Gaussian mixture models [55], etc. In spite of this, a little attention has been paid to the design39 of algorithms which does not evaluate all K components. The notable exceptions are the sparse EM40 algorithm [24] and the truncated variational EM algorithm [18], see Table 1 and Section 5 for details.41 Moreover, the methods are mostly tailored for a specific class of mixture models, e.g. the Gaussian42 mixture models.43
In this paper, we make the following contributions:44
• We propose an EM-based algorithm which relies on the MH sampler to stochastically evaluate less45 components in mixture models, substantially reducing the computational cost.46
• We design our method to enable optimization of fairly generic EM objective functions, making it47 suitable for training of both shallow and deep mixture models.48
• We apply our approach to Gaussian mixture mdoels (GMMs) and their generalizations: sum-49 product-transform networks (SPTNs) and mixtures of real-valued non-volume preserving (real50 NVP) flows [15], reaching approximately 100× speed-up compared to state-of-the-art methods.51
2 Problem formulation52
A finite mixture model characterizes the relation between an observed (known) variable, x ∈ X ⊆ RD,53 and a latent (unknown) variable, z ∈ Z := {1, . . . ,K}, via the marginal (incomplete-data) likelihood54 in the following form:55
pθ(x) = K∑ k=1 pηk(x|z = k)pπk(z = k), (1)
where θ := (π1, η1, . . . , πK , ηK) ∈ Θ are unknown parameters. Here, ηz are the parameters of the56 conditional likelihood, pηz (x|z), and πz is the weight which parameterizes the prior, pπz (z) = πz ,57 and satisfies 0 ≤ πk ≤ 1 for each k ∈ Z and ∑K k=1 πk = 1.58
Given a set of independent and identically distributed data, x := (xi)Ni=1, our goal is to learn the59 unknown parameters of the marginal log-likelihood,60
L(θ) := log pθ(x) = N∑ i=1 log K∑ k=1 pηk(xi|zi = k)pπk(zi = k). (2)
The marginalization in (2) is tractable for almost all forms of pηz (x|z). Indeed, we consider pηz (x|z)61 to belong to an arbitrary family of ηz-differentiable probability distributions. However, we assume that62 K is high, making the marginalization in (2) computationally costly, thus rendering the optimization63 objective presumably intractable. Therefore, we want to design a computationally efficient algorithm,64 requiring only M < K evaluations of pηz (x|z) at each iteration.65
3 The EM algorithm66
The maximum likelihood estimation seeks the parameters maximizing the marginal log-likelihood,67 θML := argmaxθ∈Θ L(θ). The traditional EM algorithm [13] addresses this task indirectly, i.e. by68 optimizing the evidence lower bound (ELBO),69
L(θ) ≥ Q(θ) +H(θ̂) := ELBO(θ̂), (3)
whereH(θ̂) := −Epθ̂(z|x)[log pθ̂(z|x)] is the differential entropy at an estimate, θ̂ ∈ Θ, and70
Q(θ) := Epθ̂(z|x)[log pθ(z,x)] = N∑ i=1 K∑ k=1 pθ(zi = k|xi) log pθ(zi = k, xi) (4)
is the EM objective function. Here, pθ(z,x) is the joint (complete-data) likelihood, and pθ(z|x) is71 the posterior distribution over the latent variables z := (zi)Ni=1. Given an initial value, θ0, the EM72 algorithm produces a sequence of estimates, (θt)Tt=1, by alternating between the expectation (E) and73 maximization (M) steps,74
E-step: Qt−1(θ), (5) M-step: θt := argmax
θ∈Θ Qt−1(θ). (6)
This sequence is guaranteed to monotonically tighten the ELBO, arriving at a local optimum of (2)75 under mild regularity assumptions [56].76
The EM algorithm is computationally expensive, since (4) evaluates pθ(zi, xi) for each zi ∈ Z and77 i ∈ (1, . . . , N). This has to be performed for all t ∈ (1, . . . , T ) in (5). Albeit the marginal factor,78 pπz (z), is just the cheap categorical distribution, the conditional factor, pηz (x|z), typically involves79 high-dimensional operations (e.g., the inversion of the full D×D-dimensional covariance matrices80 in the GMMs). Moreover, the M-step (6) is also expensive for large K. This holds despite that (6)81 can be reduced to closed-form updates of expected sufficient statistics for pηz (x|z) belonging to the82 exponential family [44] (again, due to high D). All in all, the computational complexity of the EM83 algorithm scales with O(TDNK).84 If (6) cannot be computed under a closed-form solution, one can resort to direct gradient-descent85 optimization of Q(θ), where argmax is replaced by one (or more) step(s) of a gradient descent86 technique. The EM algorithm is then referred to as the generalized EM algorithm [56].87
4 The generalized MHSAEM algorithm88
We design a version of the generalized EM algorithm suitable for scenarios where (4) can represent89 deep, discrete, latent variable models, thus being parameterized by possibly complex nonlinear90 transformations. We particularly focus on decreasing the the number of operations in the generalized91 EM algorithm from O(TDNK) to O(TDBM), where B N and M K.92
4.1 E-step93
We reduce the cost of evaluating the EM objective function (4) by combining the minibatching (as94 used many times before) and the Monte Carlo sampling. Namely, the specific application of the latter95 to generic mixture models is the key contribution of this paper.96
Minibatching. At each iteration, t, we compute the conditional expectation in (4) only for a subset—97 here referred to as a minibatch—of the original full dataset, i.e. (xi)i∈I . Here, I is a set of B N98 indices, i, sampled uniformly without replacement from (1, . . . , N). This substantially decreases the99 necessary computations compared to the full sweep over all N datapoints [23].100
Monte Carlo sampling. For each i ∈ I , we want to draw M K random samples from pθ(zi|xi) in101 order to obtain a Monte Carlo estimate of (4). The straightforward way to do this would be to draw102 the samples directly from pθ(zi|xi). However, direct sampling from pθ(zi|xi) does not lead to any103 substantial decrease in the number of operations. This is caused by the fact that even for a single104 sample of zi, we have to first compute the normalizing factor, pθ(xi), to obtaining the posterior,105
pθ(zi|xi). This requires K expensive evaluations of pθ(zi, xi), which is precisely what we want to106 avoid. Our approach is to resort to the Markov chain Monte Carlo (MCMC), which allows us to107 sample from pθ(zi|xi), with the computational complexity decreasing to only a single evaluation of108 pθ(zi, xi) per a single sample of zi.109
MCMC methods obviate the computation of the normalizing factor in pθ(zi|xi) by simulating a110 Markov chain, (zi,t)Tt=1, from a transition kernel, zi,t ∼ P (zi,t−1, ·), which leaves pθ(zi|xi) as its111 unique stationary (invariant) distribution, starting from an initial value zi,0. The specific form of P112 determines the structure of an MCMC method. We chose the Metropolis-Hastings (MH) sampler,113 which represents P (zi,t−1, zi,t) as follows: given z̄i := zi,t−1, draw a sample from the proposal114 distribution zi ∼ q(·|z̄i), compute the acceptance ratio,115
α(z̄i, zi) := min { 1, pηzi,t−1(xi|zi)πzi,t−1q(z̄i|zi) pηz̄i,t−1(xi|z̄i)πz̄i,t−1q(zi|z̄i) } , (7)
and, if u < α(z̄i, zi)—where u is drawn from a uniform distribution, Uniform(0, 1)—accept the116 sample and set zi,t = zi; otherwise, set zi,t = z̄i. For each i ∈ I and t ∈ (1, . . . , T ), we repeat117 this process M times, construing a set zi,t = (z1i,t, . . . , z M i,t). Therefore, at every current iteration,118 t, we continue to extend the chain from the point where we left at the previous iteration, t− 1, by119 taking z̄i = zMi,t−1. Under mild regularity assumptions [52], the chain passes the transition period120 (the burn-in phase), and the samples can then be used to approximate the conditional expectation in121 (4) as follows:122
Q̂t−1(θ) = 1
M ∑ i∈I ∑ z∈zi,t log pηz (xi|z)πz. (8)
Note that, to ensure this approach is truly efficient, we have to draw only M K samples at each123 iteration, t; otherwise, for M ≈ K, we may rather compute the exact marginalization in (4), since it124 is tractable (but computationally costly).125
4.2 M-step126
Assume for a moment that (6) with Qt−1(θ) given by (8) has a closed-form solution, yielding an127 estimate of θ. Such an estimate would have a high variance, converging only for M → ∞ and128 T → ∞ [19]. The main reason is that the samples would not be reused over the iterations, t,129 thus wasting computational resources. We consider that there is no closed-form solution of (6),130 and—to ensure that the samples (and thus computations) are recycled over the iterations—we use131 the stochastic approximation (SA) [51] to optimize (8). This is analogous to applying a stochastic132 gradient-descent method, θt = θt−1+γt∇θQ̂t−1(θ), where γt is the step-size, satisfying the Robbins-133 Monro constraints, γt ∈ [0, 1], ∑ t≥1 γt =∞, ∑ t≥1 γ 2 t <∞, and∇θ is the gradient w.r.t. θ. In this134 way, the computations made in ∇θQ̂ are accumulated via θt and reused over the iterations.135 The parameters ηz have a different form based on a specific case of pηz (x|z), whereas πz is a136 permanent structure in (1). Therefore, without loss of generality, we split (6) into a generic part and a137 fixed part as follows:138
ηk,t = ηk,t−1 + γt∇ηkQ̂t−1(θ), (9a) νk,t = νk,t−1 + γt∇νkQ̂t−1(θ), (9b)
where—to ensure that the probabilities, (πk,t)Kk=1, satisfy the constraints (Section 2)—we transform139 ∇πkQ̂ via νk = log πk and optimize w.r.t. νk. Then, to obtain (πk,t)Kk=1 from νt := (νk,t)Kk=1, we140 use the softmax function, i.e. πk,t := softmax(νt)k := exp(νk,t)/ ∑K l=1 exp(νl,t).141
Computing the gradients for all pairs of (νk, ηk)Kk=1 would be inefficient, especially since zi,t contains142 only a small number of unique values of Z for M K. Consequently, we compute ∇ηkQ̂ and143 ∇νkQ̂ only for k ∈ unique(zi,t). We summarize the proposed approach in Algorithm 1.144
4.3 Proposal distribution145
The choice of the proposal distribution has a significant impact on the speed of convergence and the146 computational cost of the proposed algorithm. Here, we discuss various possible choices of q(zi|z̄i).147
Algorithm 1 The generalized MHSAEM algorithm Input: θ0, (zi,0)Ni=1, (xi)Ni=1 Output: (θt)Tt=1
for t ∈ (1, . . . , T ) or until convergence do form the set I = (ij)Bj=1 by sampling (without replacement) B indices i ∼ (1, . . . , N) for i ∈ I do
set z̄i as the last element of zi,t−1 for j ∈ (1, . . . ,M) do
sample zi ∼ q(zi|z̄i) sample u ∼ Uniform(0, 1) compute α(z̄i, zi) in (7) if u < α(z̄i, zi) then
set zji,t = zi and z̄i = zi else
set zji,t = z̄i end if
end for set zi,t = (z1i,t, . . . , z M i,t)
end for compute (8) compute (9) for k ∈ unique(zi,t) compute πk,t := softmax(νt)k for k ∈ Z
end for
Optimal proposal (O). The optimal proposal distribution is q(zi|z̄i) := q(zi) := pθ(zi|xi). This148 ensures that the acceptance rate (7) is always α(z̄i, zi) = 1. However, the need to perform K149 expensive evaluations of pθ(zi, xi) before sampling from pθ(zi|xi) is the reason we resorted to150 the MH sampler in the first place. We consider this case only to set the upper limit on admissible151 computational cost and to study the impact of sub-optimal proposal distribtions.152
Uniform proposal (U). The uniform distribution on the discrete interval from 1 to K, i.e. q(zi|z̄i) :=153 q(zi) := Uniform(1,K), is the simplest and computationally cheapest variant of the proposal154 distribution. However, due to poor mixing properties, the algorithm may converge slowly for high K.155
Tabular proposal with forgetting (TF). The key requirement to design a proposal distribution is to156 restrict its computational complexity somewhere between that of the U and O proposals. One way to157 satisfy this constraint is to use the Markov chain, (zi,t)Tt=1, to learn a transition kernel, p(zi|z̄i), see,158 e.g. [3]. Unfortunately, this would require us to store a table with K2 entries for each i ∈ (1, . . . , N),159 which is very demanding even for moderate K and N . Therefore, we break the dependence in160 the Markov chain and define: q(zi|z̄i) := qαi(zi) := C(αi), where C(αi) ∝ ΠKk=1α 1(zi=k) k,i is the161 categorical distribution with the weights αi := (α1,i, . . . , αK,i). For L(αi) := Σtτ=1 log qαi(zi,τ ),162 we obtain an estimate of αi at iteration t as follows: αi,t := argmaxαi L(αi) = ni,t t , with163
ni,t = Σ t τ=1ezi,t , where ek is the standard basis vector (a one-hot vector) with one at kth position164 and zeros otherwise. This can be further rewritten into a recursive form: ni,t = ni,t−1 + ezi,t or,165 using the Robbins-Monro step-size, ni,t = (1− ezi,tγt) ni,t−1 + γtezi,t , where 1 is the vector of166 ones, and is the Hadamard product. We refer to this case simply as “table with forgetting” (TF)167 due to that it represents N ×K table in the memory and γt is a forgetting factor.168
5 Related work169
Stochastic approximation expectation-maximization. The application of SA to prevent the evaluation170 of all K components in mixture models has been overlooked for a long time. The reason is that the171 original motivation to combine the EM algorithm with SA is to address the analytical intractability172 of the expected value under pθ(z|x) in (4), which is, however, almost always tractable for mixture173 models. The intractability issue is addressed by either the Monte Carlo SAEM (MCSAEM) [12]174 or the Markov chain Monte Carlo SAEM (MCMCSAEM) [31]. Applying the former approach to175 mixture models would be inefficient, since it evaluates K joint distributions, pθ(z, x), before drawing176
M samples from pθ(z|x). Therefore, this method reduces only the computational cost of updating the177 sufficient statistics. This is addressed by the latter approach, where M < K samples from a proposal178 distribution, q(z|x), is used to calculate pθ(z, x) and also the sufficient statistics. However, all these179 methods process all data at every iteration, providing only a limited advantage over the conventional180 EM algorithm. Minibatch versions of these techniques have recently been proposed [27, 32, 1].181
All the above methods commonly assume pθ(z, x) belonging to the exponential family. This provides182 a convenient, but limiting, property which allows (6) to be computed under a closed-form solution.183 The main contribution of our work is to release this restrictive assumption by admitting that pθ(z, x)184 (and thus Q) is given by possibly complex and intractable transformations.185 Sparse and truncated variational techniques. There is only a small body of methods explicitly186 reducing the number of evaluated components. Their common aspect is that they follow from the187 variational framework, where the exact posterior, pθ(z|x), is approximated by a variational posterior,188 q(z|x). This sparse, approximate, posterior is defined over a lower number of components, M K,189 such that only the important components are selected, relying on relaxation of the hard EM algorithm190 from taking a single M = 1 assignment [26] to taking multiple M K assignments. The sparse191 SAEM (SSAEM) algorithm [24] selects the components by a quick partial sorting of the posterior192 probabilities, pθ(z|x). Again, this requires K evaluations of pθ(z, x) before the sorting, thus only193 reducing the amount of updated statistics. Similarly, the truncated SAEM (TSAEM) algorithm [18]194 selects M < K cluster-to-cluster and M̄ < K cluster-to-datapoint minimal Euclidean distances,195 preventing the problem in the SSAEM algorithm. However, all these distances are evaluated for all196 components in a pairwise manner, leading to K2-computational complexity, which makes the saving197 dubious. Similarly as before, these methods assume pθ(z, x) to belong to the exponential family.198
We summarize the distinguishing features of the above discussed methods in Table 1.199
6 Experiments200
To demonstrate the key features of our algorithm—its low computational complexity, competitive201 learning performance, and generality—we use it below to train: (i) GMMs on synthetic datasets, and202 (ii) SPTNs [47] and (iii) mixtures of real NVP flows [48] on real datasets. All experiments have been203 performed on a Slurm cluster equipped with Intel Xeon Scalable Gold 6146 with 384GB of RAM.204
6.1 Gaussian mixture models205
Consider the special case of a data-generating distribution given by (1), with the components taking206 the form of the multivariate Gaussian distribution, pηz (x|z) = N (x;µz,Σz), where µz is the mean207 value and Σz is the covariance matrix. The difficulty of learning GMMs heavily depends on the208 degree of interaction among all mixture components, hence having the ability to generate synthetic209 datasets with arbitrary overlap characteristics between all pairs of components is crucial for systematic210
evaluation of performance of learning algorithms [43]. Traditional techniques usually define overlap211 (or separation) of components only in terms of their mean vectors and maximum eigenvalues of the212 covariance matrices, not accounting for their rotation and mixing weights (see [36] for a detailed213 treatment of the problem). We therefore use a more objective measure of the clustering complexity214 defined by the total probability of misclassification [41], which allows to generate data with a215 user-defined degree of maximum pairwise overlap, ω.216
Experiment settings: We generate the parameters of (1), and the corresponding dataset, uniquely for a217 given quadruple (D,K,N, ω). Therefore, the parameters of the generative model are known and we218 can measure and display the convergence of the training log-likelihood, L(θt), compared to the exact219 log-likelihood, L(θ), for t = (1, . . . , T ). We are further interested in the absolute error between the220 training log-likelihood at the iteration of reaching 95% of its maximum value, t95, and the exact221 log-likelihood, i.e. AE = |L(θt95)− L(θ)|.222 We also measure the computational time until reaching t95. We have used 95% of the maximum223 value instead of the maximum value to prevent cases, where the model oscillate around target value,224 making the estimate of convergence time very noisy (for example MCSAEM in Figure 1).225
Algorithms: The GMMs belong to the exponential family of probability distributions. This allows us226 to find a closed-form, recursive, solution of (6), relying on a Robbins-Monro type of the step-size227 sequence, (γt)Tt=1, [7, 44]. In this setting, we compare our MHSAEM algorithm with a number of228 related methods in Table 1. Note we use the acronyms U and TF to specify the proposal distribution of229 the MHSAEM algorithm (Section 4.3). However, we do not use the O-proposal, since the MHSAEM-230 O algorithm is equivalent to the MCSAEM algorithm. All the SA-variants in Table 1 use a minibatch231 of size B. The key quantity to reduce the number of evaluated components and/or sufficient statistics232 in the SSAEM, TSAEM, MCSAEM and MHSAEM algorithms is collectively denoted by M (Section233 5). Note that we always keep M = M̄ in the TSAEM algorithm (see Figure 1 and 2 for concrete234 numbers). We use the step-size given by γt = 1 for t = 1, . . . , 50 and γt = 0.05 otherwise. In235 this section, to counteract the issue of attaining poor local optima, we equip all algorithms with the236 anti-annealing schedule (βt)Tt=1, starting with β1 = 0.1, reaching β2/3T = 1.2, and decreasing back237 to βT = 1.0, see [43] for details. The initial estimates of: (i) µk are uniformly drawn from the unit238 hyper-cube, (ii) Σk are fixed to unit diagonal matrix, and (iii) πk are uniformly drawn from the unit239 interval (followed by normalization).240
Results: Figure 1 shows that the EM [13] and SAEM [44] algorithms take the longest time to241 converge, attaining a poor local optima. On the other hand, the MCSAEM [1] and MHSAEM (U242 and TF) algorithms achieve L(θt95) closest to the likelihood L(θ) of the true model. Moreover, both243 MHSAEM algorithms reach this value in the shortest time compared to all the other methods. The244 SSAEM [24] and TSAEM [18] algorithms are comparable in terms of the computational time, but245 they both provide the lowest L(θt95). In Figure 2, we investigate sensitivity of fitting the model to246
increasing values of K, B and M by measuring the time and the likelihood again. In all the cases,247 the proposed MHSAEM algorithms achieve the lowest AE in the shortest time.248
SSAEM and TSAEM algorithms failed to converge for M > 2 and for K > 50 respectively. We249 believe this is caused by selecting only M maximal probabilities in the SSAEM (or distances in the250 TSAEM) algorithm (Section 5), which prevents certain, but not a negligible number of, components251 from being updated, thus providing only a crude approximation of pθ(z|x). The results then suffer252 from substantial variational gap to the exact log-likelihood (Figure 1). On the contrary, MH sampler253 provides samples which consistently approximate pθ(z|x) despite evaluating much lower number of254 components in each step.255
6.2 Sum-product transform networks256
The sum product networks (SPNs) are a deep learning extension of finite mixture models. They can257 be interpreted as a mixture of trees [60], where each tree corresponds to a component. Therefore,258 they can be cast into the form of (1), but the number of components grows exponentially with their259 depth. In this section, we use recently proposed SPTNs which introduce additional transformation260 nodes to provide better expressiveness than the SPNs (SPTNs effectively generalize SPNs and flow261 models into one large family of models).262
Experimental settings: We use 19 real datasets from the UCI database [16, 37, 35, 54], preprocessed263 in the same way as in [46]. For each experiment, we randomly split the data into 64%, 16% and 20%264 for training, validation and testing, respectively. We calculate the average log-likelihood on the test265 set and measure again the time to reach 95% of the maximal training log-likelihood, L(θt95).266 To evaluate various (possibly shallow and/or deep) architectures of SPTNs, we fit each dataset with267 all the following combinations of hyper-parameters1: s ∈ (8, 32, 128), b ∈ (2, 4, 6, 8), l ∈ (2, 3, 4),268 where s is the number of children of each sum node, b is the number of partitions of each product269 node, and l is the number of layers (one layer contains sum and product nodes). The number of270 components of the SPTN, after its conversion into (1), is given as follows: K = sl. Note that the271 maximum number of components for the investigated parameters of the SPTN is 268,435,456. To272 reduce the space of possible architectures, we restrict ourselves only to (i) the leaf nodes given by273 N (0, I); (ii) affine transformations fixed to the singular value decomposition, choosing the the Givens274 parameterization for the unitary matrices [47]; and (iii) no sharing of any type of nodes [47].275
Algorithms: We evaluate only on the MHSAEM-U algorithm—due to its favourable computational276 complexity and simplicity—and compare it with the stochastic gradient-descent (SGD) algorithm,277 which is routinely used to train SP(T)Ns [45, 47]. In this case, SGD in each iteration performs278 computations over all subtrees of the network, whereas the MHSAEM-U algorithm computes with279 only M = 1 subtrees, thus we should observe speed-up of the computations. In our implementation,280 both these methods perform optimization of their respective objective functions—the log-likelihood281 (2) for SGD and the EM objective (8) for MHSAEM-U—via the use of the automatic differentiation282 and the ADAM optimizer [29], using B = 100 and T = 20000.283
Results: Since each dataset might benefit from a different architecture, Table 6.2 shows the test284 log-likelihood of the architectures selected according to the best likelihood measured on the validation285 set and the corresponding speed-up. The test log-likelihoods reveal that the MHSAEM-U algorithm286 outperforms the SGD algorithm on 10 out of 19 datasets, which was not originally the goal, but287 the added stochasticity helps to escape poor local minima. The speed-up demonstrates lower288 computational complexity of the MHSAEM-U algorithm on 17 out of 19 datasets, which was the289 main goal. The magic-telescope and wine datasets show approximately 102× and 75× speed-up,290 respectively, while on very small datasets (pima-indians and iris), the SGD is faster due to291 effective implementation. In the supplementary material, we present Table 3, exhibiting the same292 trends on a fixed architecture.293

heading: 6.3 Mixtures of real NVP flows294
text: We consider another class of mixture models (1), where each component pηz (x|z) is transformed by295 the flow model—real NVP [15]. These transformations are parameterized via deep neural networks,296 allowing for flexible adjustment of the learning capacity of each component.297
1We have set a hard limit to train a single model to 24h, which is default on our Slurm cluster.
Experimental settings: We use the same experimental settings and evaluation metrics as in Section298 6.2. We apply the mixture model on all datasets, changing the number of components as follows:299 K ∈ (8, 32, 128). Each real NVP-based component in the mixture model has (i) the translation300 function parameterized via multi-layer perceptron with a single hidden layer of dimension 10, using301 the rectified linear activation function; and (ii) the scale function parameterized via the same network302 except with the hyperbolic tangent activation function. We do not use the batch normalization [15] and303 we stack two layers of the translation-scale transformation (we have used implementation from [20]).304
Algorithms: The algorithms and their settings are the same as those in Section 6.2.305
Results: The experimental results are presented in right part of Table 6.2. They are similar to those306 obtained in the previous section. In terms of the test log-likelihood, the MHSAEM-U algorithm307 outperforms the SGD algorithm on all but three datasets, and it provides a substantial speed-up on all308 datasets except one. The test likelihood of models with the real NVP flows is most of the time worse309 than that of SPTNs with the affine transformations. As explained in the supplementary, this is due to310 the overfitting, which has been observed in [47].311
7 Conclusion312
This paper has presented a method to decrease computational complexity of fitting mixture models,313 including their generalizations, such as sum-product-(transform) networks and mixtures of flow314 models. The speed-up is achieved by evaluating and updating only a single component (per iteration),315 where the Metropolis-Hasting algorithm ensures sampling of components from a proper posterior. An316 experimental comparison on all three classes of models mentioned above confirmed the theoretical317 expectations. The method significantly speeds-up the fitting time and, importantly, without sacrificing318 the quality of the fit. In fact, the likelihood was better than that of the models fitted by the EM319 algorithm or the SGD algorithm in more than 50% of cases. We attribute this to higher stochasticity,320 which helps to escape from poor local minima.321
In the experiments, the proposed method has used a uniform proposal distribution in the MH sampler.322 Despite outperforming the alternative methods, we conjecture that this limits the speed of convergence.323 Therefore, we believe that there is still a room for improvement in the implementation. We plan to324 address these issues in future work.325
8 Broader impact statement326
The presented method decreases the computational complexity of fitting large (and deep) mixture327 models, which leads to five to hundred time speed-up depending on a size of the problem (although328 negative exceptions occurs). We believe this line of research, which we want to continue, to have329 important benefits. First, it is directly related to decrease in energy consumption and in production of330 CO2 (we expect similar rates as the speedup). Second, it has a positive effect on financial aspects of331 deploying (and experimenting with) mixture models. Third, it decreases the hardware requirements,332 as in all experiments presented above the model was fitted on a single-core.333
References334 [1] S. Allassonnière and J. Chevallier. A new class of stochastic EM algorithms: Escaping335 local maxima and handling intractable sampling. Computational Statistics & Data Analysis,336 159:107159, 2021.337
[2] S. Ayesha, M. K. Hanif, and R. Talib. Overview and comparative study of dimensionality338 reduction techniques for high dimensional data. Information Fusion, 59:44–58, 2020.339
[3] D. S. Bai. Efficient estimation of transition probabilities in a Markov chain. The Annals of340 Statistics, pages 1305–1317, 1975.341
[4] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians.342 Journal of the American statistical Association, 112(518):859–877, 2017.343
[5] J. Blömer and K. Bujna. Adaptive seeding for Gaussian mixture models. In Pacific-asia344 conference on knowledge discovery and data mining, pages 296–308. Springer, 2016.345
[6] T. Campbell and B. Beronov. Sparse variational inference: Bayesian coresets from scratch.346 arXiv preprint arXiv:1906.03329, 2019.347
[7] O. Cappé and E. Moulines. On-line expectation–maximization algorithm for latent data models.348 Journal of the Royal Statistical Society: Series B (Statistical Methodology), 71(3):593–613,349 2009.350
[8] J. Chen, J. Zhu, Y. Teh, and T. Zhang. Stochastic expectation maximization with variance351 reduction. Advances in Neural Information Processing Systems, page 7967, 2018.352
[9] N. Chopin. A sequential particle filter method for static models. Biometrika, 89(3):539–552,353 2002.354
[10] G. Compiani and Y. Kitamura. Using mixtures in econometric models: A brief review and some355 new results. The Econometrics Journal, 19(3):C95–C127, 2016.356
[11] S. Dasgupta. Learning mixtures of Gaussians. In 40th Annual Symposium on Foundations of357 Computer Science (Cat. No. 99CB37039), pages 634–644. IEEE, 1999.358
[12] B. Delyon, M. Lavielle, E. Moulines, et al. Convergence of a stochastic approximation version359 of the EM algorithm. The Annals of Statistics, 27(1):94–128, 1999.360
[13] A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data361 via the EM algorithm. Journal of the Royal Statistical Society: Series B (Methodological),362 39(1):1–22, 1977.363
[14] J. Demšar. Statistical comparisons of classifiers over multiple data sets. The Journal of Machine364 Learning Research, 7:1–30, 2006.365
[15] L. Dinh, J. Sohl-Dickstein, and S. Bengio. Density estimation using real NVP. In 5th Interna-366 tional Conference on Learning Representations, ICLR 2017, 2017.367
[16] D. Dua and C. Graff. UCI machine learning repository, 2017.368
[17] D. Feldman, M. Faulkner, and A. Krause. Scalable training of mixture models via coresets. In369 Proceedings of the 24th International Conference on Neural Information Processing Systems,370 pages 2142–2150, 2011.371
[18] D. Forster and J. Lücke. Can clustering scale sublinearly with its clusters? A variational EM372 acceleration of GMMs and k-means. In International Conference on Artificial Intelligence and373 Statistics, pages 124–132. PMLR, 2018.374
[19] G. Fort, E. Moulines, et al. Convergence of the Monte Carlo expectation maximization for375 curved exponential families. Annals of Statistics, 31(4):1220–1259, 2003.376
[20] J. Franců. Continuousflows.jl. https://github.com/janfrancu/ContinuousFlows.jl,377 2020.378
[21] K. Greff, S. van Steenkiste, and J. Schmidhuber. Neural expectation maximization. In Proceed-379 ings of the 31st International Conference on Neural Information Processing Systems, pages380 6694–6704, 2017.381
[22] R. Gribonval, A. Chatalic, N. Keriven, V. Schellekens, L. Jacques, and P. Schniter. Sketching382 datasets for large-scale learning (long version). arXiv preprint arXiv:2008.01839, 2020.383
[23] M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley. Stochastic variational inference. Journal384 of Machine Learning Research, 14(5), 2013.385
[24] M. C. Hughes and E. B. Sudderth. Fast learning of clusters and topics via sparse posteriors.386 arXiv preprint arXiv:1609.07521, 2016.387
[25] K. Humphreys and D. Titterington. Approximate Bayesian inference for simple mixtures. In388 COMPSTAT, pages 331–336. Springer, 2000.389
[26] B.-H. Juang and L. R. Rabiner. The segmental K-means algorithm for estimating parameters390 of hidden Markov models. IEEE Transactions on acoustics, speech, and signal Processing,391 38(9):1639–1641, 1990.392
[27] B. Karimi, M. Lavielle, and É. Moulines. On the convergence properties of the mini-batch EM393 and MCEM algorithms, 2019.394
[28] N. Keriven, A. Bourrier, R. Gribonval, and P. Pérez. Sketching for large-scale learning of395 mixture models. Information and Inference: A Journal of the IMA, 7(3):447–508, 2018.396
[29] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint397 arXiv:1412.6980, 2014.398
[30] A. Kucukelbir, D. Tran, R. Ranganath, A. Gelman, and D. M. Blei. Automatic differentiation399 variational inference. The Journal of Machine Learning Research, 18(1):430–474, 2017.400
[31] E. Kuhn and M. Lavielle. Coupling a stochastic approximation version of EM with an MCMC401 procedure. ESAIM: Probability and Statistics, 8:115–131, 2004.402
[32] E. Kuhn, C. Matias, and T. Rebafka. Properties of the stochastic approximation EM algorithm403 with mini-batch sampling. Statistics and Computing, 30(6):1725–1739, 2020.404
[33] A. Lagrange, M. Fauvel, and M. Grizonnet. Large-scale feature selection with Gaussian mixture405 models for the classification of high dimensional remote sensing images. IEEE Transactions on406 Computational Imaging, 3(2):230–242, 2017.407
[34] M. Lavine and M. West. A Bayesian method for classification and discrimination. Canadian408 Journal of Statistics, 20(4):451–461, 1992.409
[35] M. Little, P. McSharry, S. Roberts, D. Costello, and I. Moroz. Exploiting nonlinear recurrence410 and fractal scaling properties for voice disorder detection. Nature Precedings, pages 1–1, 2007.411
[36] R. Maitra and V. Melnykov. Simulating data to study performance of finite mixture modeling412 and clustering algorithms. Journal of Computational and Graphical Statistics, 19(2):354–376,413 2010.414
[37] O. L. Mangasarian and W. H. Wolberg. Cancer diagnosis via linear programming. Technical415 report, University of Wisconsin-Madison Department of Computer Sciences, 1990.416
[38] J.-M. Marin, K. Mengersen, and C. P. Robert. Bayesian modelling and inference on mixtures of417 distributions. Handbook of statistics, 25:459–507, 2005.418
[39] C. A. McGrory, D. C. Ahfock, J. A. Horsley, and C. L. Alston. Weighted Gibbs sampling for419 mixture modelling of massive datasets via coresets. Stat, 3(1):291–299, 2014.420
[40] G. J. McLachlan, S. X. Lee, and S. I. Rathnayake. Finite mixture models. Annual review of421 statistics and its application, 6:355–378, 2019.422
[41] V. Melnykov, W.-C. Chen, and R. Maitra. MixSim: An R package for simulating data to study423 performance of clustering algorithms. Journal of Statistical Software, 51(12):1, 2012.424
[42] T. Monnier, T. Groueix, and M. Aubry. Deep transformation-invariant clustering. In Conference425 on Neural Information Processing Systems (NeurIPS 2020), 2020.426
[43] I. Naim and D. Gildea. Convergence of the EM algorithm for Gaussian mixtures with unbalanced427 mixing coefficients. In Proceedings of the 29th International Coference on International428 Conference on Machine Learning, pages 1427–1431, 2012.429
[44] H. D. Nguyen, F. Forbes, and G. J. McLachlan. Mini-batch learning of exponential family finite430 mixture models. Statistics and Computing, pages 1–18, 2020.431
[45] R. Peharz, A. Vergari, K. Stelzner, A. Molina, X. Shao, M. Trapp, K. Kersting, and Z. Ghahra-432 mani. Random sum-product networks: A simple and effective approach to probabilistic deep433 learning. In Uncertainty in Artificial Intelligence, pages 334–344. PMLR, 2020.434
[46] T. Pevný. Loda: Lightweight on-line detector of anomalies. Machine Learning, 102(2):275–304,435 2016.436
[47] T. Pevný, V. Šmídl, M. Trapp, O. Poláček, and T. Oberhuber. Sum-product-transform networks:437 Exploiting symmetries using invertible transformations. arXiv preprint arXiv:2005.01297,438 2020.439
[48] G. G. Pires and M. A. Figueiredo. Variational mixture of normalizing flows. arXiv preprint440 arXiv:2009.00585, 2020.441
[49] A. Rau, C. Maugis-Rabusseau, M.-L. Martin-Magniette, and G. Celeux. Co-expression analysis442 of high-throughput transcriptome sequencing data with Poisson mixture models. Bioinformatics,443 31(9):1420–1427, 2015.444
[50] R. A. Redner and H. F. Walker. Mixture densities, maximum likelihood and the EM algorithm.445 SIAM review, 26(2):195–239, 1984.446
[51] H. Robbins and S. Monro. A stochastic approximation method. The annals of mathematical447 statistics, pages 400–407, 1951.448
[52] C. Robert and G. Casella. Monte Carlo statistical methods. Springer Science & Business Media,449 2013.450
[53] W. Siblini, P. Kuntz, and F. Meyer. A review on dimensionality reduction for multi-label451 classification. IEEE Transactions on Knowledge and Data Engineering, 2019.452
[54] J. P. Siebert. Vehicle recognition using rule based methods. 1987.453
[55] C. Viroli and G. J. McLachlan. Deep gaussian mixture models. Statistics and Computing,454 29(1):43–51, 2019.455
[56] C. F. J. Wu. On the convergence properties of the EM algorithm. The Annals of statistics, pages456 95–103, 1983.457
[57] D. Wu and J. Ma. An effective EM algorithm for mixtures of Gaussian processes via the MCMC458 sampling and approximation. Neurocomputing, 331:366–374, 2019.459
[58] W. Xiang, A. Karfoul, C. Yang, H. Shu, and R. L. B. Jeannès. An exact line search scheme to460 accelerate the EM algorithm: Application to Gaussian mixture models identification. Journal of461 computational science, 41:101073, 2020.462
[59] M. Zhang, Y. Fu, K. M. Bennett, and T. Wu. Computational efficient variational Bayesian Gaus-463 sian mixture models via coreset. In 2016 International Conference on Computer, Information464 and Telecommunication Systems (CITS), pages 1–5. IEEE, 2016.465
[60] H. Zhao, P. Poupart, and G. Gordon. A unified approach for learning the parameters of sum-466 product networks. In Proceedings of the 30th International Conference on Neural Information467 Processing Systems, pages 433–441, 2016.468
Checklist469
1. For all authors...470 (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s471 contributions and scope? [Yes]472 (b) Did you describe the limitations of your work? [Yes] Our main contribution is compu-473 tational speedup. Cases where it was not achieved are highlighted in the experimental474 section.475
(c) Did you discuss any potential negative societal impacts of your work? [No] We do not476 foresee any potential negative impact.477
(d) Have you read the ethics review guidelines and ensured that your paper conforms to478 them? [Yes]479
2. If you are including theoretical results...480 (a) Did you state the full set of assumptions of all theoretical results? [N/A]481 (b) Did you include complete proofs of all theoretical results? [N/A]482
3. If you ran experiments...483 (a) Did you include the code, data, and instructions needed to reproduce the main exper-484 imental results (either in the supplemental material or as a URL)? [Yes] The code is485 available in a github repository. All dataset are public from the UCI database.486
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they487 were chosen)? [Yes] See Section 6.488
(c) Did you report error bars (e.g., with respect to the random seed after running experi-489 ments multiple times)? [No] We report only average of Monte Carlo repetitions, the490 error bars were too small to have any visual impact in the reported logarithmic scale.491
(d) Did you include the total amount of compute and the type of resources used (e.g., type492 of GPUs, internal cluster, or cloud provider)? [Yes]493
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...494 (a) If your work uses existing assets, did you cite the creators? [Yes] We use 20 datasets495 from UCI, we cite the required papers for each dataset, mostly the UCI database and496 few additional publications.497
(b) Did you mention the license of the assets? [No] The data are publically available, we498 comply with the requirement on citing appropriate publications.499
(c) Did you include any new assets either in the supplemental material or as a URL? [No]500 (d) Did you discuss whether and how consent was obtained from people whose data you’re501 using/curating? [N/A]502 (e) Did you discuss whether the data you are using/curating contains personally identifiable503 information or offensive content? [N/A]504 5. If you used crowdsourcing or conducted research with human subjects...505
(a) Did you include the full text of instructions given to participants and screenshots, if506 applicable? [N/A]507
(b) Did you describe any potential participant risks, with links to Institutional Review508 Board (IRB) approvals, if applicable? [N/A]509
(c) Did you include the estimated hourly wage paid to participants and the total amount510 spent on participant compensation? [N/A]511

Review summaries:
The paper presents a modified version of the EM algorithm that allows for faster fitting of mixture models. The approach is based on the idea of evaluating and updating a single mixture component per iteration. Furthermore, they present a generalized version that extends beyond the standard case of the exponential family.   The authors then follow with an empirical demonstration of the benefits: up to two orders of magnitude faster training while in some cases increasing the likelihood compared to the standard SGD approach.  The evaluation is carried on two architectures, SPTNs and mixture of real NVP flows. 

The paper suggests a method for efficiently evaluating mixture models. The main idea is to use Expectation-Maximization (EM) and Metropolis-Hastings (MH) to evaluate only a small number of mixture components, reducing the computational cost by consequence. The theoretical contributions of the manuscript are two-folded. First, the authors proposed an EM-based algorithm for training, which uses MH for sampling from a smaller set of mixture components. Second, the manuscript generalizes the EM-based algorithm by making it compatible with gradient-based optimization methods. The practical contributions include experiments with Gaussian mixture models (GMMs), sum-product networks (SPNs), and mixtures of real-valued flows. Results show a benchmark speed-up of 100x.

This paper proposes a new method for scalably fitting mixture models. The authors tackle the problem of fitting with a large number of mixture components. While there exist some works handling this problem, these methods are all limited to mixtures of exponential family distributions, whereas the authors’ method applies to generic mixture models (as long as the model provides a likelihood p(data point i | cluster k)). They do so by proposing a modified version of the EM algorithm: instead of exactly computing the expectations in the E step, the authors propose to use a Metropolis Hastings algorithm to approximate the required expectations. In experiments, the authors show that their method (1) outperforms pre-existing large-K methods (where such previous applications are applicable), and (2) show that their method outperforms standard stochastic gradient descent in settings where other large-K methods are not applicable. The authors’ method seems to work well, but I have a few clarification questions about the technical ideas behind the proposed algorithm and its empirical evaluation. Currently, I put a borderline score for the paper, but I’m definitely open to increasing it after discussion.

This paper proposes a stochastic EM algorithm. Specifically, the authors use a minibatch to approximate the data and use samples obtained from Metropolis-Hasting (MH) to approximate the likelihood component per datum. The proposed method is tested on a synthetic dataset and several UCI datasets.

This paper proposes a stochastic approximation to the E-step of the EM algorithm to reduce the computational cost of fitting mixture models when the number of components $K$ is large. Specifically, the paper proposes using $M << K$ samples of the latent variables $z$ from a Metropolis-Hasting (MH) sampler to approximate the E-step. This lets us only use $M$ likelihood evaluations instead of $K$ for each observation in the E-step. The paper compares there proposed method against other EM-based methods on Gaussian mixture models on synthetic data, and sum product  transform networks and non-volume preserving flows on 19 UCI datasets.