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