Amine Echraibi

Joachim Flocon-Cholet
email: cholet<joachim.floconcholet@orange.com>

Stéphane Gosselin

Sandrine Vaton
email: <sandrine.vaton@imt-atlantique.fr>

On the Variational Posterior of Dirichlet Process Deep Latent Gaussian Mixture Models

Thanks to the reparameterization trick, deep latent Gaussian models have shown tremendous success recently in learning latent representations. The ability to couple them however with nonparametric priors such as the Dirichlet Process (DP) hasn't seen similar success due to its non parameterizable nature. In this paper, we present an alternative treatment of the variational posterior of the Dirichlet Process Deep Latent Gaussian Mixture Model (DP-DLGMM), where we show that the prior cluster parameters and the variational posteriors of the beta distributions and cluster hidden variables can be updated in closed-form. This leads to a standard reparameterization trick on the Gaussian latent variables knowing the cluster assignments. We demonstrate our approach on standard benchmark datasets, we show that our model is capable of generating realistic samples for each cluster obtained, and manifests competitive performance in a semi-supervised setting.

Introduction

Nonparametric Bayesian priors, such as the Dirichlet Process (DP), have been widely adopted in the probabilistic graphical community. Their ability to generate an infinite amount of probability distributions using a discrete latent variable makes them ideally suited for automatic model selection. The most famous applications of the DP have been however limited to classical probabilistic graphical models such as Dirichlet Process Mixture Models and Hierarchical Dirichlet Process Hidden Markov Models [START_REF] Blei | Variational inference for dirichlet process mixtures[END_REF][START_REF] Fox | An hdp-hmm for systems with state persistence[END_REF][START_REF] Zhang | Stochastic variational inference for the hdp-hmm[END_REF].

Second workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models (ICML 2020), Virtual Conference Recently, deep generative models such as Deep Latent Gaussian Models (DLGMs) and Variational AutoEncoders (VAEs) [START_REF] Kingma | Auto-encoding variational bayes[END_REF][START_REF] Rezende | Stochastic backpropagation and approximate inference in deep generative models[END_REF] have shown huge success in modeling and generating complex data structures such as images. Various proposals to generalize these models to the mixture and nonparametric mixture cases have been made (Nalisnick et al., 2016;Nalisnick & Smyth, 2016;[START_REF] Dilokthanakul | Deep unsupervised clustering with gaussian mixture variational autoencoders[END_REF][START_REF] Jiang | Variational deep embedding: An unsupervised and generative approach to clustering[END_REF]. Introducing such priors on top of the deep generative model can improve its generative capabilities, preserve class structure in the latent representation space, and offer a nonparametric way of performing model selection with respect to the size of the generative model.

The main challenge posed by such models lies in the inference process. Deep generative models with continuous latent variables owe their success mainly to the reparameterization trick [START_REF] Kingma | Auto-encoding variational bayes[END_REF][START_REF] Rezende | Stochastic backpropagation and approximate inference in deep generative models[END_REF]. This approach provides an efficient and scalable method for obtaining low variance estimates of the gradient of the variational lower bound with respect to variational posterior parameters. Applying this approach directly to the variational posterior of the DP is not straightforward, due to the fact that a reparameterization trick for the beta distributions is hard to obtain [START_REF] Ruiz | The generalized reparameterization gradient[END_REF]. One approach to bypass this issue have been proposed by (Nalisnick & Smyth, 2016), where the authors used the Kumaraswamy distribution [START_REF] Kumaraswamy | A generalized probability density function for double-bounded random processes[END_REF] as a higher entropy alternative for the beta distribution in the variational posterior. However, by deriving the nature of the variational posterior directly from the variational lower bound, we can show that the appropriate distribution is in fact the beta distribution.

In this paper we provide an alternative treatment of the variational posterior of the DP-DLGMM, where we combine classical variational inference to derive the variational posteriors of the beta distributions and cluster hidden variables, and neural variational inference for the hidden variables of the latent Gaussian model. This leads to gradient ascent updates over the parameters present in nonlinear transformations where the reparameterization trick can be applied knowing the cluster assignment. As for the remaining parameters, closed-form solutions can be obtained by maximization of the evidence lower bound.

Dirichlet Process Deep Latent Gaussian Mixture Models

Generalizing deep latent Gaussian models to the Dirichlet process mixture case can be obtained by adding a Dirichlet process prior on the hidden cluster assignments. We denote these cluster assignments by z. Following the assignment of a cluster hidden variable, a deep latent Gaussian model is defined for the assigned cluster similar to [START_REF] Rezende | Stochastic backpropagation and approximate inference in deep generative models[END_REF]. We adopt the stick-breaking construction of the Dirichlet Process [START_REF] Sethuraman | A constructive definition of dirichlet priors[END_REF]. The generative process of the model (figure 1) is given by:

β k ∼ Beta(•; 1, η) π k = β k k-1 l=1 (1 -β l ) z n |π ∼ Cat(•|π) (l) n ∼ N (•; 0, I) ∀l h (L) n = m (L) zn + s (L) zn (L) n h (l) n = f W (l) zn h (l+1) n + s (l) zn (l) n x n |h (1) n , z n ∼ p X • |f W (0) zn h (1) n where h (l)
n ∈ R p l is the l th layer hidden representation constructed using a nonlinear transformation f W (l) zn represented by a neural network for the cluster assignment z n . For simplicity, we consider diagonal covariance matrices for each layer where the diagonal elements are (s (l) zn,j ) 2 1≤j≤p l , hence represents the element-wise product. The generalization to full covariance matrices is straightforward using the Cholesky decomposition.

We denote by η the concentration parameter of the Dirichlet process which is a hyperparameter to be tuned manually. The term p X represents the emission distribution of the observable x n , usually chosen to be a normal distribution for continuous variables or the Bernoulli distribution for binary variables. We denote the parameters of the generative model by:

Θ = {m (L) 1:∞ , s (L) 1:∞ , W (0:L-1) 1:∞ , s (1:L-1) 1:∞ }
The model thus has an infinite number of parameters due to the Dirichlet process prior. Furthermore, the posterior distribution of the hidden variables cannot be computed in closed-form. In order to perform inference on the model we need to use approximate methods such as Markov Chain Monte Carlo (MCMC) or Variational Inference. MCMC methods are not suitable for high dimensional models such as the DP-DLGMM, where convergence of the Markov chain to the true posterior can prove to be slow and hard to diagnose [START_REF] Blei | Variational inference: A review for statisticians[END_REF].

β z n h (l+1) n1 h (l+1) np l+1 h (l) n1 h (l) np l l ∈ {0, ..., L -1} n ∈ {1, ..., N } Figure 1.
The graphical representation of the generative process of the model, with the convention x = h (0) .

In the next section, we develop a structured variational inference algorithm for DP-DLGMM. We show that by choosing a suitable structure for the variational posterior, closed-form solutions can be obtained for the updates of the truncated variational posteriors of the beta distributions, the variational posteriors of the cluster hidden variables, and the optimal prior parameters {m (L) , s (L) } maximizing the evidence lower bound.

Structured Variational Inference

For a brief review of variational methods, we denote by x 1:N the N samples present in the dataset supposed to be independent and identically distributed. The log-likelihood of the model is intractable due to the required marginalization of all the hidden variables. In order to bypass this marginalization, we introduce an approximate distribution q Φ and use Jensen's inequality to obtain a lower bound [START_REF] Jordan | An introduction to variational methods for graphical models[END_REF]:

l(Θ) = ln p Θ (x 1:N ) = ln z 1:N p Θ (x 1:N , z 1:N , h (1:L) 1:N , β)dh (1:L) 1:N dβ ≥ E z 1:N ,h (1:L) 1:N ,β∼qΦ ln p Θ (x 1:N , z 1:N , h (1:L) 1:N , β) q Φ (z 1:N , h (1:L) 1:N , β|x 1:N ) L(Θ, Φ).
(1)

We can show that if the distribution q Φ is a good approximation of the true posterior, maximizing the evidence lower bound (ELBO) with respect to the model parameters Θ is equivalent to maximizing the log-likelihood. For deep generative models, most state-of-the-art methods use inference networks to construct the posterior distribution [START_REF] Rezende | Stochastic backpropagation and approximate inference in deep generative models[END_REF]Nalisnick & Smyth, 2016). For deep mixture models with discrete latent variables, this approach leads to a mixture density variational posterior where the reparameterization trick requires additional investigation [START_REF] Graves | Stochastic backpropagation through mixture density distributions[END_REF]. Our approach combines standard variational Bayes and neural variational inference. We approximate the true posterior using the following structured variational posterior:

q Φ (z 1:N , h (1:L) 1:N , β|x 1:N ) = N n=1 L l=1 q ψ (l) zn (h (l) n |x n , z n ) × q φn (z n |x n ) T t=1 q γt (β t |x 1:N ), ( 2
)
where T is a truncation level for the variational posterior of the beta distributions obtained by supposing that q(β T = 1) = 1 [START_REF] Blei | Variational inference for dirichlet process mixtures[END_REF]. We assume a factorized posterior over the hidden layers h (1:L) n

, where the intra-layer dependencies are conserved.

3.1. Deriving the variational posteriors q φn and q γt Deriving the nature of the posterior distributions of the hidden layers h (1:L) n using the variational approach is intractable due to the nonlinearities present in the model. Thus, we take a similar approach to [START_REF] Rezende | Stochastic backpropagation and approximate inference in deep generative models[END_REF], and we assume that the variational posterior is specified by an inference network, where the parameters of the distribution are the outputs of deep neural networks µ ψ (l)

t and Σ ψ (l) t of parameters ψ (l)
t for the l th layer and the t th cluster:

q ψ (l) t (h (l) n |x n , z n = t) = N h (l) n ; µ ψ (l) t (x n ), Σ ψ (l) t (x n ) .
In contrast to the hidden layers, we can use the proposed variational posterior of equation ( 2) to derive closed-form solutions for q φn and q γt . Let us consider the Kullback-Leibler definition of the ELBO L:

L(Θ, Φ) = -D KL [q Φ (•|x 1:N )||p Θ (•, x 1:N )] .
By plugging the variational posterior and isolating β t terms and z n terms, we can analytically derive the optimal distributions q γt and q φn maximizing L:

q γt (β t |x 1:N ) = Beta(β t ; γ 1,t , γ 2,t ) q φn (z n |x n ) = Cat(z n ; φ n ),
where the fixed point equations for the variational parameters φ n and γ t are:

γ 1,t = 1 + N n=1 φ n,t (3)
γ 2,t = η + N n=1 T r=t+1 φ n,r (4)
ln φ n,t = const + E β∼q [ln π t ] + E h (1:L) n ∼q ψ (1:L) t ln p X (x n , h (1:L) n |z n = t) + l H q ψ (l) t (•|z n = t, x n ) s.t. T t=1 φ n,t = 1, (5)
The fixed point equation of φ n,t , requires the evaluation of the expectation over the hidden layers, this can be performed by sampling from the variational posterior of each hidden layer and then forwarding the sample using the generative model:

E h (1:L) n ∼q ψ (1:L) t ln p X (x n , h (1:L) n |z n = t) ≈ 1 S S s=1 ln p X x n , h (1:L)(s) n,t |z n = t where: h (l)(s) n,t ∼ q ψ (l) t (h (l) n |x n , z n = t). (6)
A key insight here is the following: if a cluster t is incapable of reconstructing a sample x n from the variational posterior, this will reinforce the belief that x n should not be assigned to that cluster. Furthermore, the estimation of the expectation can be performed using the same reparameterization trick that we will develop in section 3.3.

Closed-Form updates for m (L)

1:T and s (L) 1:T

In addition to the variational posteriors of the beta distributions and the cluster assignments, closed-form solutions can be obtained for the updates of m (L)

1:T and s (L) 1:T . Let us reconsider the evidence lower bound of equation ( 1), where we isolate only terms dependent on the prior parameters. We have:

L(m (L) 1:T , s (L) 1:T ) = const - n,t φ n,t × D KL N (µ ψ (L) t (x n ), Σ ψ (L) t (x n ))||N (m (L) t , V (L) t ) .
where

V (L) t = diag (s (L) t,j ) 2 1≤j≤p L
represents the covariance matrix of the L th layer. By setting the derivative of L with respect to the parameters to zero, we obtain:

m (L) t = 1 N t N n=1 φ n,t µ ψ (L) t (x n ) N t = N n=1 φ n,t (7)
V (L) t = 1 N t N n=1 φ n,t I Σ ψ (L) t (x n ) + (µ ψ (L) t (x n ) -m (L) t )(µ ψ (L) t (x n ) -m (L) t ) T , (8)
where to extract the diagonal elements we perform an elementwise multiplication by the identity matrix I. The update rules obtained are similar to the M-Step of a classical Gaussian Mixture Model, except in this case the updates are performed on the last hidden layer of the generative model, and the E-step of equation ( 5) takes into account all the hidden layers. Detailed derivation of the previous equations are presented in the supplementary material.

Stochastic Backpropagation

We next show how to perform stochastic backpropagation in order to maximize L with respect to the parameters ψ and Λ = {W (0:L-1) 1:T , s

(1:L) 1:T }. Similarly to the previous section, we isolate the terms in the evidence lower bound dependent on ψ and Λ. We have:

L(ψ, Λ) = const+ n,t φ n,t l H q ψ (l) t (•|z n = t, x n ) + E h (1:L) n ∼q ψ (1:L) t ln p X (x n , h (1:L) n |z n = t) . (9)
By taking the expectation over the hidden cluster variables z n , we obtain conditional expectations over the hidden layers h (1:L) n knowing the cluster assignment. In order to backpropagate gradients of Λ and ψ, it suffices to perform a reparameterization trick for each cluster assignment at each hidden layer (proof in Appendix A). We can achieve this by sampling:

(l) n,t ∼ N (0, I), a sample from the posterior of the l th hidden layer can then be obtained by the following transformation:

h (l) n,t = µ ψ (l) t (x n ) + (l) n,t Σ ψ (l) t (x n ),
where Σ ψ (l) t (x n ) is supposed to be a diagonal matrix for simplicity. Following the previous analysis, we can derive an algorithm to perform inference on the proposed model, where between iterations of the fixed point update steps, E epochs of gradient ascent are performed to obtain a local maximum of the ELBO with respect to Λ and ψ. Algorithm 1 summarizes the process.

Semi-Supervised Learning (SSL)

4.1. SSL using the DP-DGLMM In this section, similarly to [START_REF] Kingma | Semi-supervised learning with deep generative models[END_REF] we consider a partially labeled dataset x 1:N = D l ∪ D u , where Algorithm 1 Variational Inference for the DP-DLGMM Input: x 1:N , T, η, α Initialize φ, Λ, ψ while not converged do update: γ t ∀t {(3), (4)} update:

m (L) t ∀t {(7)} update: V (L) t ∀t {(8)} for each epoch do Λ ← Λ + α∂ Λ L ψ ← ψ + α∂ ψ L end for update: φ n,t ∀n, ∀t {(5)} end while D l = {x n , y n } n
is the labeled part, y n represents the label of the sample x n , and D u represents the unlabeled part. The log likelihood can be divided for the labeled and unlabeled parts as:

l(Θ) = ln p Θ (x 1:N ) = xn∈D l ln p Θ (x n ) + xn∈Du ln p Θ (x n ) = xn∈D l ln p Θ (x n , z n = y n ) + xn∈Du ln p Θ (x n ).
The last equation follows from the fact that p Θ (x n |z n = y n ) = 0. By dividing the labeled and unlabeled parts of the dataset, we can follow the same approach presented in section 3 in order to derive a variational inference algorithm. In this case, the fixed point updates and the gradient ascent steps remain unchanged if we set φ n,y n = 1 for a labeled x n sample.

The predictive distribution

In order to make predictions using the model, we need to evaluate the predictive distribution. Given a new sample x N +1 , the objective is to evaluate the following quantity p(z N +1 = k|x 1:N +1 ). This task requires an intractable marginalization over all the other hidden variables. However, similarly to [START_REF] Blei | Variational inference for dirichlet process mixtures[END_REF], we can use the variational posterior to approximate the true posterior, which in turn leads to simpler expectation terms:

p(z N +1 = k|x 1:N +1 ) ∝ p(z N +1 = k, x N +1 |x 1:N ) ∝ ∼ E β∼q [π k (β)] ×E h (1:L) N +1 ∼q ψ (1:L) k p X x N +1 |f Λ (h (1:L) N +1 ), z n = k (10)
where f Λ (•) represents the forward pass over the generative model. The expectation with respect to the beta terms can be computed in closed-form as a product of expectations over the beta posteriors. The second expectation can be evaluated using the Monte-Carlo estimator of equation ( 6).

Experiments

Evaluation of the semi-supervised classification

We evaluate the semi-supervised classification capabilities of the model. We train our DP-DLGMM model on the MNIST dataset (LeCun & Cortes, 2010) with train-valid-test splits equal to {45000, 5000, 10000} similarly to (Nalisnick & Smyth, 2016), with 10 % labelisation randomly drawn. We run the process for 5 iterations, and we evaluate our model on the test set. We report the mean and standard deviation of the classification error in percentages in Table 1. Our method produces a competitive score with existing state-of-the art methods: Deep Generative Models (DGM) [START_REF] Kingma | Semi-supervised learning with deep generative models[END_REF] and Stick-Breaking Deep Generative Models (SB-DGM) (Nalisnick & Smyth, 2016). Unlike the previous approaches, the loss was not up-weighted for the labeled samples. Figure 2 shows the t-SNE projections [START_REF] Maaten | Visualizing data using t-sne[END_REF] obtained with 10 % of the labels provided. We notice that by introducing a small fraction of labels the class structure was highly preserved in the latent space. 6.13 ± .13 4.86 ± .14 3.95 ± .15 2.90 ± .17

Table 1. Semi-supervised classification error (%) on the MNIST test set with 10 % labelisation. Comparison with (Nalisnick & Smyth, 2016) .

Data generation and visualization

To further test our model, we generate samples for each cluster from the models trained on both the MNIST and SVHN [START_REF] Netzer | Reading digits in natural images with unsupervised feature learning[END_REF] datasets. The MNIST model is trained in an unsupervised manner, and the SVHN model is trained with semi-supervision where we provide 1000 randomly generated labels. The samples obtained are represented in figure 4. For the unsupervised model, we notice that the clusters are representative of the shape of each digit. We plot the t-SNE projections of the MNIST test set of the unsupervised model in Figure 3. We notice that the digits belonging to the same true class tend to group with each other. However, two groups of the same class can be very separated in the embedding space. The interpretation we can draw from this effect is that the DP-DLGMM tends to separate the latent space in order to distinguish between the variations of hidden representations of the same class. The clusters obtained are not always representative of the true classes which is a common effect with infinite mixture models. In a full unsupervised setting, data can be explained by multiple correct clusterings. This effect can simply be countered by adding a small supervision (figure 2).

Conclusion

In this paper, we have presented a variational inference method for Dirichlet Process Deep Latent Gaussian Mixture Models. Our approach combines classical variational inference and neural variational inference. The algorithm derived is thus a standard variational inference algorithm, with fixed point updates over a subset of the parameters presenting linear dependencies. The parameters present in nonlinear transformations are updated using standard gradient ascent where the reparameterization trick can be applied for the variational posterior of the stochastic hidden layers knowing the cluster assignments. Our approach shows promising results both for the unsupervised and semi-supervised cases.

In future work, stochastic variational inference can be explored to speed-up the training procedure. Our approach can also be generalized to other types of deep probabilistic graphical models.

A. Proof of the reparameterization trick knowing the cluster assignment

The evidence lower bound of our model can be written in its general form as: where ˆ is a sample drawn from N (0, I), thus we can backpropagate stochastic gradients for each class assignment.

L(θ, ψ) = T t=1 φ t × E h∼N (µ ψ t (x),σ ψ t (x) 2 I) [f θ (h)]

B. Stochastic Variational Inference

Updating the Λ and ψ parameters using E epochs of gradient ascent significantly adds to the complexity of Algorithm 1. One possible approach is to perform stochastic variational inference [START_REF] Hoffman | Stochastic variational inference[END_REF] for fixed point update equations. This allows for the use of the same batch of data for the gradient ascent steps of Λ and ψ and the stochastic updates of the fixed point equations. Let us consider a batch x 1:B , the updates in this case are:

γ (t+1) 1,k = (1 -ρ t )γ (t) 1,k + ρ t N B γ1,k γ (t+1) 2,k = (1 -ρ t )γ (t) 2,k + ρ t N B γ2,k m (t+1) k = (1 -ρ t )m (t) k + ρ t N B mk V (t+1) k = (1 -ρ t )V (t) k + ρ t N B Vk ,
where γ1,k , γ2,k , mk , Vk , are computed for the minibatch x 1:B using equations (3),( 4),(7), and (8) respectively. In order to guarantee convergence ρ t must satisfy:

t ρ t = ∞ and t ρ 2 t < ∞.

Figure 2

 2 Figure 2. t-SNE plot of the second stochastic hidden layer on the MNIST test set for the semi-supervised (10% labels) version of the DP-DLGMM.

Figure 3

 3 Figure 3. t-SNE plot of the second stochastic hidden layer on the MNIST test set for the unsupervised version of the DP-DLGMM.

Figure 4 .

 4 Figure 4. Generated samples from the DP-DLGMM model for the unsupervised version on the MNIST dataset (left) and the semisupervised version on the SVHN dataset (right).

NN

  (h, µ ψt (x), σ ψt (x) 2 I)f θ (h(h, µ ψt (x), σ ψt (x) 2 I) × f θ (h)|∂ h t |d .By introducing the following transformation:h t ( ) = µ ψt (x) + σ ψt (x)∼ N (0, I), and using the density transformation lemma:N (h, µ ψt (x), σ ψt (x) 2 I)|∂ h t | = N ( ; 0, I),we have:L(θ, ψ) = T t=1 φ t E [f θ (µ ψt (x) + σ ψt (xθ (µ ψt (x) + σ ψt (x) ˆ ).

Supplementary material

A. Detailed derivation of the variational inference algorithm:

The evidence lower bound of equation ( 1) can be written as :

in the following we develop this equation in order to derive the nature of the variational posteriors and their fixed point updates.

A.1. Computing q γt (β t |x 1:N ):

The distribution q * (β t ) maximizing L, minimizes the kullback-leibler term. Hence:

By isolating the terms dependent on z n in L, we obtain:

Supplementary material

Hence, the optimal distribution q * φn (z n |x n ) maximizing L satisfies:

where, 1:T present in the evidence lower bound, we have:

where, V

. By setting the derivative of L to zero:

Hence, the closed-form update equations:

is a diagonal matrix we extract the diagonal updates by multiplying elementwise by the identity matrix.

B. The predictive distribution:

In order to perform inference for a new sample x N +1 , we need to compute the predictive distribution p(z N + 1|x 1:N +1 ) : By approximating the true posteriors with the variational posteriors:

p(β|x 1:N ) ≈ q(β) and p(h|x 1:N , z N +1 = k) ≈ q(h|z N +1 = k)

We conclude: