Fitting large mixture models using stochastic component selection

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 ) can not 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

Summary: 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.