Raise your ELBO
Published on 2019-02-14

Gaussian mixture model (GMM) is an example of mixture models.

The space of the data is \(\mathbb R^n\). We use the hypothesis that the data is Gaussian conditioned on the latent variable:

\[(x_i; \eta_k) \sim N(\mu_k, \Sigma_k),\]

so we write \(\eta_k = (\mu_k, \Sigma_k)\),

During E-step, the \(q(z_i)\) can be directly computed using Bayes' theorem:

\[r_{ik} = q(z_i = k) = \mathbb P(z_i = k | x_{i}; \theta_t) = {g_{\mu_{t, k}, \Sigma_{t, k}} (x_{i}) \pi_{t, k} \over \sum_{j = 1 : n_z} g_{\mu_{t, j}, \Sigma_{t, j}} (x_{i}) \pi_{t, j}},\]

where \(g_{\mu, \Sigma} (x) = (2 \pi)^{- n / 2} \det \Sigma^{-1 / 2} \exp(- {1 \over 2} (x - \mu)^T \Sigma^{-1} (x - \mu))\) is the pdf of the Gaussian distribution \(N(\mu, \Sigma)\).

During M-step, we need to compute

\[\text{argmax}_{\Sigma, \mu, \pi} \sum_{i = 1 : m} \sum_{k = 1 : n_z} r_{ik} \log (g_{\mu_k, \Sigma_k}(x_{i}) \pi_k).\]

This is similar to the quadratic discriminant analysis, and the solution is

\[\begin{aligned} \pi_{k} &= {1 \over m} \sum_{i = 1 : m} r_{ik}, \\ \mu_{k} &= {\sum_i r_{ik} x_{i} \over \sum_i r_{ik}}, \\ \Sigma_{k} &= {\sum_i r_{ik} (x_{i} - \mu_{t, k}) (x_{i} - \mu_{t, k})^T \over \sum_i r_{ik}}. \end{aligned}\]

Remark. The k-means algorithm is the \(\epsilon \to 0\) limit of the GMM with constraints \(\Sigma_k = \epsilon I\). See Section 9.3.2 of Bishop 2006 for derivation. It is also briefly mentioned there that a variant in this setting where the covariance matrix is not restricted to \(\epsilon I\) is called elliptical k-means algorithm.

As a transition to the next models to study, let us consider a simpler mixture model obtained by making one modification to GMM: change \((x; \eta_k) \sim N(\mu_k, \Sigma_k)\) to \(\mathbb P(x = w; \eta_k) = \eta_{kw}\) where \(\eta\) is a stochastic matrix and \(w\) is an arbitrary element of the space for \(x\). So now the space for both \(x\) and \(z\) are finite. We call this model the simple mixture model (SMM).

As in GMM, at E-step \(r_{ik}\) can be explicitly computed using Bayes' theorem.

It is not hard to write down the solution to the M-step in this case:

\[\begin{aligned} \pi_{k} &= {1 \over m} \sum_i r_{ik}, \qquad (2.7)\\ \eta_{k, w} &= {\sum_i r_{ik} 1_{x_i = w} \over \sum_i r_{ik}}. \qquad (2.8) \end{aligned}\]

where \(1_{x_i = w}\) is the indicator function, and evaluates to \(1\) if \(x_i = w\) and \(0\) otherwise.

Two trivial variants of the SMM are the two versions of probabilistic latent semantic analysis (pLSA), which we call pLSA1 and pLSA2.

The model pLSA1 is a probabilistic version of latent semantic analysis, which is basically a simple matrix factorisation model in collaborative filtering, whereas pLSA2 has a fully Bayesian version called latent Dirichlet allocation (LDA), not to be confused with the other LDA (linear discriminant analysis).

The pLSA model (Hoffman 2000) is a mixture model, where the dataset is now pairs \((d_i, x_i)_{i = 1 : m}\). In natural language processing, \(x\) are words and \(d\) are documents, and a pair \((d, x)\) represent an ocurrance of word \(x\) in document \(d\).

For each datapoint \((d_{i}, x_{i})\),

\[\begin{aligned} p(d_i, x_i; \theta) &= \sum_{z_i} p(z_i; \theta) p(d_i | z_i; \theta) p(x_i | z_i; \theta) \qquad (2.91)\\ &= p(d_i; \theta) \sum_z p(x_i | z_i; \theta) p (z_i | d_i; \theta) \qquad (2.92). \end{aligned}\]

Of the two formulations, (2.91) corresponds to pLSA type 1, and (2.92) corresponds to type 2.

The hidden markov model (HMM) is a sequential version of SMM, in the same sense that recurrent neural networks are sequential versions of feed-forward neural networks.

HMM is an example where the posterior \(p(z_i | x_i; \theta)\) is not easy to compute, and one has to utilise properties of the underlying Bayesian network to go around it.

Now each sample is a sequence \(x_i = (x_{ij})_{j = 1 : T}\), and so are the latent variables \(z_i = (z_{ij})_{j = 1 : T}\).

The latent variables are assumed to form a Markov chain with transition matrix \((\xi_{k \ell})_{k \ell}\), and \(x_{ij}\) is completely dependent on \(z_{ij}\):

\[\begin{aligned} p(z_{ij} | z_{i, j - 1}) &= \xi_{z_{i, j - 1}, z_{ij}},\\ p(x_{ij} | z_{ij}) &= \eta_{z_{ij}, x_{ij}}. \end{aligned}\]

Also, the distribution of \(z_{i1}\) is again categorical with parameter \(\pi\):

\[p(z_{i1}) = \pi_{z_{i1}}\]

So the parameters are \(\theta = (\pi, \xi, \eta)\). And HMM can be shown in plate notations as:

Now we apply EM to HMM, which is called the Baum-Welch algorithm. Unlike the previous examples, it is too messy to compute \(p(z_i | x_{i}; \theta)\), so during the E-step we instead write down formula (2.5) directly in hope of simplifying it:

\[\begin{aligned} \mathbb E_{p(z_i | x_i; \theta_t)} \log p(x_i, z_i; \theta_t) &=\mathbb E_{p(z_i | x_i; \theta_t)} \left(\log \pi_{z_{i1}} + \sum_{j = 2 : T} \log \xi_{z_{i, j - 1}, z_{ij}} + \sum_{j = 1 : T} \log \eta_{z_{ij}, x_{ij}}\right). \qquad (3) \end{aligned}\]

Let us compute the summand in second term:

\[\begin{aligned} \mathbb E_{p(z_i | x_{i}; \theta_t)} \log \xi_{z_{i, j - 1}, z_{ij}} &= \sum_{k, \ell} (\log \xi_{k, \ell}) \mathbb E_{p(z_{i} | x_{i}; \theta_t)} 1_{z_{i, j - 1} = k, z_{i, j} = \ell} \\ &= \sum_{k, \ell} p(z_{i, j - 1} = k, z_{ij} = \ell | x_{i}; \theta_t) \log \xi_{k, \ell}. \qquad (4) \end{aligned}\]

Similarly, one can write down the first term and the summand in the third term to obtain

\[\begin{aligned} \mathbb E_{p(z_i | x_{i}; \theta_t)} \log \pi_{z_{i1}} &= \sum_k p(z_{i1} = k | x_{i}; \theta_t), \qquad (5) \\ \mathbb E_{p(z_i | x_{i}; \theta_t)} \log \eta_{z_{i, j}, x_{i, j}} &= \sum_{k, w} 1_{x_{ij} = w} p(z_{i, j} = k | x_i; \theta_t) \log \eta_{k, w}. \qquad (6) \end{aligned}\]

plugging (4)(5)(6) back into (3) and summing over \(j\), we obtain the formula to maximise over \(\theta\) on:

\[\sum_k \sum_i r_{i1k} \log \pi_k + \sum_{k, \ell} \sum_{j = 2 : T, i} s_{ijk\ell} \log \xi_{k, \ell} + \sum_{k, w} \sum_{j = 1 : T, i} r_{ijk} 1_{x_{ij} = w} \log \eta_{k, w},\]

where

\[\begin{aligned} r_{ijk} &:= p(z_{ij} = k | x_{i}; \theta_t), \\ s_{ijk\ell} &:= p(z_{i, j - 1} = k, z_{ij} = \ell | x_{i}; \theta_t). \end{aligned}\]

Now we proceed to the M-step. Since each of the \(\pi_k, \xi_{k, \ell}, \eta_{k, w}\) is nicely confined in the inner sum of each term, together with the constraint \(\sum_k \pi_k = \sum_\ell \xi_{k, \ell} = \sum_w \eta_{k, w} = 1\) it is not hard to find the argmax at time \(t + 1\) (the same way one finds the MLE for any categorical distribution):

\[\begin{aligned} \pi_{k} &= {1 \over m} \sum_i r_{i1k}, \qquad (6.1) \\ \xi_{k, \ell} &= {\sum_{j = 2 : T, i} s_{ijk\ell} \over \sum_{j = 1 : T - 1, i} r_{ijk}}, \qquad(6.2) \\ \eta_{k, w} &= {\sum_{ij} 1_{x_{ij} = w} r_{ijk} \over \sum_{ij} r_{ijk}}. \qquad(6.3) \end{aligned}\]

Note that (6.1)(6.3) are almost identical to (2.7)(2.8). This makes sense as the only modification HMM makes over SMM is the added dependencies between the latent variables.

What remains is to compute \(r\) and \(s\).

This is done by using the forward and backward procedures which takes advantage of the conditioned independence / topology of the underlying Bayesian network. It is out of scope of this post, but for the sake of completeness I include it here.

Let

\[\begin{aligned} \alpha_k(i, j) &:= p(x_{i, 1 : j}, z_{ij} = k; \theta_t), \\ \beta_k(i, j) &:= p(x_{i, j + 1 : T} | z_{ij} = k; \theta_t). \end{aligned}\]

They can be computed recursively as

\[\begin{aligned} \alpha_k(i, j) &= \begin{cases} \eta_{k, x_{1j}} \pi_k, & j = 1; \\ \eta_{k, x_{ij}} \sum_\ell \alpha_\ell(j - 1, i) \xi_{k\ell}, & j \ge 2. \end{cases}\\ \beta_k(i, j) &= \begin{cases} 1, & j = T;\\ \sum_\ell \xi_{k\ell} \beta_\ell(j + 1, i) \eta_{\ell, x_{i, j + 1}}, & j < T. \end{cases} \end{aligned}\]

Then

\[\begin{aligned} p(z_{ij} = k, x_{i}; \theta_t) &= \alpha_k(i, j) \beta_k(i, j), \qquad (7)\\ p(x_{i}; \theta_t) &= \sum_k \alpha_k(i, j) \beta_k(i, j),\forall j = 1 : T \qquad (8)\\ p(z_{i, j - 1} = k, z_{i, j} = \ell, x_{i}; \theta_t) &= \alpha_k(i, j) \xi_{k\ell} \beta_\ell(i, j + 1) \eta_{\ell, x_{j + 1, i}}. \qquad (9) \end{aligned}\]

And this yields \(r_{ijk}\) and \(s_{ijk\ell}\) since they can be computed as \({(7) \over (8)}\) and \({(9) \over (8)}\) respectively.

Definition (Fully Bayesian mixture model). The relations between \(\pi\), \(\eta\), \(x\), \(z\) are the same as in the definition of mixture models. Furthermore, we assume the distribution of \((x | \eta_k)\) belongs to the exponential family (the definition of the exponential family is briefly touched at the end of this section). But now both \(\pi\) and \(\eta\) are random variables. Let the prior distribution \(p(\pi)\) is Dirichlet with parameter \((\alpha, \alpha, ..., \alpha)\). Let the prior \(p(\eta_k)\) be the conjugate prior of \((x | \eta_k)\), with parameter \(\beta\), we will see later in this section that the posterior \(q(\eta_k)\) belongs to the same family as \(p(\eta_k)\). Represented in a plate notations, a fully Bayesian mixture model looks like:

Given this structure we can write down the mean-field approximation:

\[\log q(z) = \mathbb E_{q(\eta)q(\pi)} (\log(x | z, \eta) + \log(z | \pi)) + C.\]

Both sides can be factored into per-sample expressions, giving us

\[\log q(z_i) = \mathbb E_{q(\eta)} \log p(x_i | z_i, \eta) + \mathbb E_{q(\pi)} \log p(z_i | \pi) + C\]

Therefore

\[\log r_{ik} = \log q(z_i = k) = \mathbb E_{q(\eta_k)} \log p(x_i | \eta_k) + \mathbb E_{q(\pi)} \log \pi_k + C. \qquad (9.1)\]

So the posterior of each \(z_i\) is categorical regardless of the $p$s and $q$s.

Computing the posterior of \(\pi\) and \(\eta\):

\[\log q(\pi) + \log q(\eta) = \log p(\pi) + \log p(\eta) + \sum_i \mathbb E_{q(z_i)} p(x_i | z_i, \eta) + \sum_i \mathbb E_{q(z_i)} p(z_i | \pi) + C.\]

So we can separate the terms involving \(\pi\) and those involving \(\eta\). First compute the posterior of \(\pi\):

\[\log q(\pi) = \log p(\pi) + \sum_i \mathbb E_{q(z_i)} \log p(z_i | \pi) = \log p(\pi) + \sum_i \sum_k r_{ik} \log \pi_k + C.\]

The right hand side is the log of a Dirichlet modulus the constant \(C\), from which we can update the posterior parameter \(\phi^\pi\):

\[\phi^\pi_k = \alpha + \sum_i r_{ik}. \qquad (9.3)\]

Similarly we can obtain the posterior of \(\eta\):

\[\log q(\eta) = \log p(\eta) + \sum_i \sum_k r_{ik} \log p(x_i | \eta_k) + C.\]

Again we can factor the terms with respect to \(k\) and get:

\[\log q(\eta_k) = \log p(\eta_k) + \sum_i r_{ik} \log p(x_i | \eta_k) + C. \qquad (9.5)\]

Here we can see why conjugate prior works. Mathematically, given a probability distribution \(p(x | \theta)\), the distribution \(p(\theta)\) is called conjugate prior of \(p(x | \theta)\) if \(\log p(\theta) + \log p(x | \theta)\) has the same form as \(\log p(\theta)\).

For example, the conjugate prior for the exponential family \(p(x | \theta) = h(x) \exp(\theta \cdot T(x) - A(\theta))\) where \(T\), \(A\) and \(h\) are some functions is \(p(\theta; \chi, \nu) \propto \exp(\chi \cdot \theta - \nu A(\theta))\).

Here what we want is a bit different from conjugate priors because of the coefficients \(r_{ik}\). But the computation carries over to the conjugate priors of the exponential family (try it yourself and you'll see). That is, if \(p(x_i | \eta_k)\) belongs to the exponential family

\[p(x_i | \eta_k) = h(x) \exp(\eta_k \cdot T(x) - A(\eta_k))\]

and if \(p(\eta_k)\) is the conjugate prior of \(p(x_i | \eta_k)\)

\[p(\eta_k) \propto \exp(\chi \cdot \eta_k - \nu A(\eta_k))\]

then \(q(\eta_k)\) has the same form as \(p(\eta_k)\), and from (9.5) we can compute the updates of \(\phi^{\eta_k}\):

\[\begin{aligned} \phi^{\eta_k}_1 &= \chi + \sum_i r_{ik} T(x_i), \qquad (9.7) \\ \phi^{\eta_k}_2 &= \nu + \sum_i r_{ik}. \qquad (9.9) \end{aligned}\]

So the mean field approximation for the fully Bayesian mixture model is the alternate iteration of (9.1) (E-step) and (9.3)(9.7)(9.9) (M-step) until convergence.

A typical example of fully Bayesian mixture models is the fully Bayesian Gaussian mixture model (Attias 2000, also called variational GMM in the literature). It is defined by applying the same modification to GMM as the difference between Fully Bayesian mixture model and the vanilla mixture model.

More specifically:

• \(p(z_{i}) = \text{Cat}(\pi)\) as in vanilla GMM
• \(p(\pi) = \text{Dir}(\alpha, \alpha, ..., \alpha)\) has Dirichlet distribution, the conjugate prior to the parameters of the categorical distribution.
• \(p(x_i | z_i = k) = p(x_i | \eta_k) = N(\mu_{k}, \Sigma_{k})\) as in vanilla GMM
• \(p(\mu_k, \Sigma_k) = \text{NIW} (\mu_0, \lambda, \Psi, \nu)\) is the normal-inverse-Wishart distribution, the conjugate prior to the mean and covariance matrix of the Gaussian distribution.

The E-step and M-step can be computed using (9.1) and (9.3)(9.7)(9.9) in the previous section. The details of the computation can be found in Chapter 10.2 of Bishop 2006 or Attias 2000.

As the second example of fully Bayesian mixture models, Latent Dirichlet allocation (LDA) (Blei-Ng-Jordan 2003) is the fully Bayesian version of pLSA2, with the following plate notations:

It is the smoothed version in the paper.

More specifically, on the basis of pLSA2, we add prior distributions to \(\eta\) and \(\pi\):

\[\begin{aligned} p(\eta_k) &= \text{Dir} (\beta, ..., \beta), \qquad k = 1 : n_z \\ p(\pi_\ell) &= \text{Dir} (\alpha, ..., \alpha), \qquad \ell = 1 : n_d \\ \end{aligned}\]

And as before, the prior of \(z\) is

\[p(z_{\ell, i}) = \text{Cat} (\pi_\ell), \qquad \ell = 1 : n_d, i = 1 : m\]

We also denote posterior distribution

\[\begin{aligned} q(\eta_k) &= \text{Dir} (\phi^{\eta_k}), \qquad k = 1 : n_z \\ q(\pi_\ell) &= \text{Dir} (\phi^{\pi_\ell}), \qquad \ell = 1 : n_d \\ p(z_{\ell, i}) &= \text{Cat} (r_{\ell, i}), \qquad \ell = 1 : n_d, i = 1 : m \end{aligned}\]

As before, in E-step we update \(r\), and in M-step we update \(\lambda\) and \(\gamma\).

But in the LDA paper, one treats optimisation over \(r\), \(\lambda\) and \(\gamma\) as a E-step, and treats \(\alpha\) and \(\beta\) as parameters, which is optmised over at M-step. This makes it more akin to the classical EM where the E-step is Bayesian and M-step MLE. This is more complicated, and we do not consider it this way here.

Plugging in (9.1) we obtain the updates at E-step

\[r_{\ell i k} \propto \exp(\psi(\phi^{\pi_\ell}_k) + \psi(\phi^{\eta_k}_{x_{\ell i}}) - \psi(\sum_w \phi^{\eta_k}_w)), \qquad (10)\]

where \(\psi\) is the digamma function. Similarly, plugging in (9.3)(9.7)(9.9), at M-step, we update the posterior of \(\pi\) and \(\eta\):

\[\begin{aligned} \phi^{\pi_\ell}_k &= \alpha + \sum_i r_{\ell i k}. \qquad (11)\\ %%}}$ %%As for $\eta$, we have %%{{$%align% %%\log q(\eta) &= \sum_k \log p(\eta_k) + \sum_{\ell, i} \mathbb E_{q(z_{\ell i})} \log p(x_{\ell i} | z_{\ell i}, \eta) \\ %%&= \sum_{k, j} (\sum_{\ell, i} r_{\ell i k} 1_{x_{\ell i} = j} + \beta - 1) \log \eta_{k j} %%}}$ %%which gives us %%{{$ \phi^{\eta_k}_w &= \beta + \sum_{\ell, i} r_{\ell i k} 1_{x_{\ell i} = w}. \qquad (12) \end{aligned}\]

So the algorithm iterates over (10) and (11)(12) until convergence.

The Dirichlet process mixture model (DPMM) is like the fully Bayesian mixture model except \(n_z = \infty\), i.e. \(z\) can take any positive integer value.

The probability of \(z_i = k\) is defined using the so called stick-breaking process: let \(v_i \sim \text{Beta} (\alpha, \beta)\) be i.i.d. random variables with Beta distributions, then

\[\mathbb P(z_i = k | v_{1:\infty}) = (1 - v_1) (1 - v_2) ... (1 - v_{k - 1}) v_k.\]

So \(v\) plays a similar role to \(\pi\) in the previous models.

As before, we have that the distribution of \(x\) belongs to the exponential family:

\[p(x | z = k, \eta) = p(x | \eta_k) = h(x) \exp(\eta_k \cdot T(x) - A(\eta_k))\]

so the prior of \(\eta_k\) is

\[p(\eta_k) \propto \exp(\chi \cdot \eta_k - \nu A(\eta_k)).\]

Because of the infinities we can't directly apply the formulas in the general fully Bayesian mixture models. So let us carefully derive the whole thing again.

As before, we can write down the ELBO:

\[L(p(x, z, \theta), q(z, \theta)) = \mathbb E_{q(\theta)} \log {p(\theta) \over q(\theta)} + \mathbb E_{q(\theta) q(z)} \log {p(x, z | \theta) \over q(z)}.\]

Both terms are infinite series:

\[L(p, q) = \sum_{k = 1 : \infty} \mathbb E_{q(\theta_k)} \log {p(\theta_k) \over q(\theta_k)} + \sum_{i = 1 : m} \sum_{k = 1 : \infty} q(z_i = k) \mathbb E_{q(\theta)} \log {p(x_i, z_i = k | \theta) \over q(z_i = k)}.\]

There are several ways to deal with the infinities. One is to fix some level \(T > 0\) and set \(v_T = 1\) almost surely (Blei-Jordan 2006). This effectively turns the model into a finite one, and both terms become finite sums over \(k = 1 : T\).

Another walkaround (Kurihara-Welling-Vlassis 2007) is also a kind of truncation, but less heavy-handed: setting the posterior \(q(\theta) = q(\eta) q(v)\) to be:

\[q(\theta) = q(\theta_{1 : T}) p(\theta_{T + 1 : \infty}) =: q(\theta_{\le T}) p(\theta_{> T}).\]

That is, tie the posterior after \(T\) to the prior. This effectively turns the first term in the ELBO to a finite sum over \(k = 1 : T\), while keeping the second sum an infinite series:

\[L(p, q) = \sum_{k = 1 : T} \mathbb E_{q(\theta_k)} \log {p(\theta_k) \over q(\theta_k)} + \sum_i \sum_{k = 1 : \infty} q(z_i = k) \mathbb E_{q(\theta)} \log {p(x_i, z_i = k | \theta) \over q(z_i = k)}. \qquad (13)\]

The plate notation of this model is:

As it turns out, the infinities can be tamed in this case.

As before, the optimal \(q(z_i)\) is computed as

\[r_{ik} = q(z_i = k) = s_{ik} / S_i\]

where

\[\begin{aligned} s_{ik} &= \exp(\mathbb E_{q(\theta)} \log p(x_i, z_i = k | \theta)) \\ S_i &= \sum_{k = 1 : \infty} s_{ik}. \end{aligned}\]

Plugging this back to (13) we have

\[\begin{aligned} \sum_{k = 1 : \infty} r_{ik} &\mathbb E_{q(\theta)} \log {p(x_i, z_i = k | \theta) \over r_{ik}} \\ &= \sum_{k = 1 : \infty} r_{ik} \mathbb E_{q(\theta)} (\log p(x_i, z_i = k | \theta) - \mathbb E_{q(\theta)} \log p(x_i, z_i = k | \theta) + \log S_i) = \log S_i. \end{aligned}\]

So it all rests upon \(S_i\) being finite.

For \(k \le T + 1\), we compute the quantity \(s_{ik}\) directly. For \(k > T\), it is not hard to show that

\[s_{ik} = s_{i, T + 1} \exp((k - T - 1) \mathbb E_{p(w)} \log (1 - w)),\]

where \(w\) is a random variable with same distribution as \(p(v_k)\), i.e. \(\text{Beta}(\alpha, \beta)\).

Hence

\[S_i = \sum_{k = 1 : T} s_{ik} + {s_{i, T + 1} \over 1 - \exp(\psi(\beta) - \psi(\alpha + \beta))}\]

is indeed finite. Similarly we also obtain

\[q(z_i > k) = S^{-1} \left(\sum_{\ell = k + 1 : T} s_\ell + {s_{i, T + 1} \over 1 - \exp(\psi(\beta) - \psi(\alpha + \beta))}\right), k \le T \qquad (14)\]

Now let us compute the posterior of \(\theta_{\le T}\). In the following we exchange the integrals without justifying them (c.f. Fubini's Theorem). Equation (13) can be rewritten as

\[L(p, q) = \mathbb E_{q(\theta_{\le T})} \left(\log p(\theta_{\le T}) + \sum_i \mathbb E_{q(z_i) p(\theta_{> T})} \log {p(x_i, z_i | \theta) \over q(z_i)} - \log q(\theta_{\le T})\right).\]

Note that unlike the derivation of the mean-field approximation, we keep the \(- \mathbb E_{q(z)} \log q(z)\) term even though we are only interested in \(\theta\) at this stage. This is again due to the problem of infinities: as in the computation of \(S\), we would like to cancel out some undesirable unbounded terms using \(q(z)\). We now have

\[\log q(\theta_{\le T}) = \log p(\theta_{\le T}) + \sum_i \mathbb E_{q(z_i) p(\theta_{> T})} \log {p(x_i, z_i | \theta) \over q(z_i)} + C.\]

By plugging in \(q(z = k)\) we obtain

\[\log q(\theta_{\le T}) = \log p(\theta_{\le T}) + \sum_{k = 1 : \infty} q(z_i = k) \left(\mathbb E_{p(\theta_{> T})} \log {p(x_i, z_i = k | \theta) \over q(z_i = k)} - \mathbb E_{q(\theta)} \log {p(x_i, z_i = k | \theta) \over q(z_i = k)}\right) + C.\]

Again, we separate the \(v_k\)'s and the \(\eta_k\)'s to obtain

\[q(v_{\le T}) = \log p(v_{\le T}) + \sum_i \sum_k q(z_i = k) \left(\mathbb E_{p(v_{> T})} \log p(z_i = k | v) - \mathbb E_{q(v)} \log p (z_i = k | v)\right).\]

Denote by \(D_k\) the difference between the two expetations on the right hand side. It is easy to show that

\[D_k = \begin{cases} \log(1 - v_1) ... (1 - v_{k - 1}) v_k - \mathbb E_{q(v)} \log (1 - v_1) ... (1 - v_{k - 1}) v_k & k \le T\\ \log(1 - v_1) ... (1 - v_T) - \mathbb E_{q(v)} \log (1 - v_1) ... (1 - v_T) & k > T \end{cases}\]

so \(D_k\) is bounded. With this we can derive the update for \(\phi^{v, 1}\) and \(\phi^{v, 2}\):

\[\begin{aligned} \phi^{v, 1}_k &= \alpha + \sum_i q(z_i = k) \\ \phi^{v, 2}_k &= \beta + \sum_i q(z_i > k), \end{aligned}\]

where \(q(z_i > k)\) can be computed as in (14).

When it comes to \(\eta\), we have

\[\log q(\eta_{\le T}) = \log p(\eta_{\le T}) + \sum_i \sum_{k = 1 : \infty} q(z_i = k) (\mathbb E_{p(\eta_k)} \log p(x_i | \eta_k) - \mathbb E_{q(\eta_k)} \log p(x_i | \eta_k)).\]

Since \(q(\eta_k) = p(\eta_k)\) for \(k > T\), the inner sum on the right hand side is a finite sum over \(k = 1 : T\). By factorising \(q(\eta_{\le T})\) and \(p(\eta_{\le T})\), we have

\[\log q(\eta_k) = \log p(\eta_k) + \sum_i q(z_i = k) \log (x_i | \eta_k) + C,\]

which gives us

\[\begin{aligned} \phi^{\eta, 1}_k &= \xi + \sum_i q(z_i = k) T(x_i) \\ \phi^{\eta, 2}_k &= \nu + \sum_i q(z_i = k). \end{aligned}\]

As an example of AEVB, the paper introduces variational autoencoder (VAE), with the following instantiations:

• The prior \(p(z_i) = N(0, I)\) is standard normal, thus independent of \(\theta\).
• The distribution \(p(x_i; \eta)\) is either Gaussian or categorical.
• The distribution \(q(z_i; \mu, \Sigma)\) is Gaussian with diagonal covariance matrix. So \(g_\phi(z_i) = (\mu_\phi(x_i), \text{diag}(\sigma^2_\phi(x_i)_{1 : d}))\). Thus in the reparameterisation trick \(\epsilon \sim N(0, I)\) and \(T_\phi(x_i, \epsilon) = \epsilon \odot \sigma_\phi(x_i) + \mu_\phi(x_i)\), where \(\odot\) is elementwise multiplication.
• The KL divergence can be easily computed analytically as \(- D(q(z_i; g_\phi(x_i)) || p(z_i)) = {d \over 2} + \sum_{j = 1 : d} \log\sigma_\phi(x_i)_j - {1 \over 2} \sum_{j = 1 : d} (\mu_\phi(x_i)_j^2 + \sigma_\phi(x_i)_j^2)\).

With this, one can use backprop to maximise the ELBO.

Let us turn to fully Bayesian version of AEVB. Again, we first recall the ELBO of the fully Bayesian mixture models:

\[L(p(x, z, \pi, \eta; \alpha, \beta), q(z, \pi, \eta; r, \phi)) = L(p(x | z, \eta) p(z | \pi) p(\pi; \alpha) p(\eta; \beta), q(z; r) q(\eta; \phi^\eta) q(\pi; \phi^\pi)).\]

We write \(\theta = (\pi, \eta)\), rewrite \(\alpha := (\alpha, \beta)\), \(\phi := r\), and \(\gamma := (\phi^\eta, \phi^\pi)\). Furthermore, as in the half-Bayesian version we assume \(p(z | \theta) = p(z)\), i.e. \(z\) does not depend on \(\theta\). Similarly we also assume \(p(\theta; \alpha) = p(\theta)\). Now we have

\[L(p(x, z, \theta; \alpha), q(z, \theta; \phi, \gamma)) = L(p(x | z, \theta) p(z) p(\theta), q(z; \phi) q(\theta; \gamma)).\]

And the objective is to maximise it over \(\phi\) and \(\gamma\). We no longer maximise over \(\theta\), because it is now a random variable, like \(z\). Now let us transform it to a neural network model, as in the half-Bayesian case:

\[L\left(\left(\prod_{i = 1 : m} p(x_i; f_\theta(z_i))\right) \left(\prod_{i = 1 : m} p(z_i) \right) p(\theta), \left(\prod_{i = 1 : m} q(z_i; g_\phi(x_i))\right) q(\theta; h_\gamma(x))\right).\]

where \(f_\theta\), \(g_\phi\) and \(h_\gamma\) are neural networks. Again, by separating out KL-divergence terms, the above formula becomes

\[\sum_i \mathbb E_{q(\theta; h_\gamma(x))q(z_i; g_\phi(x_i))} \log p(x_i; f_\theta(z_i)) - \sum_i D(q(z_i; g_\phi(x_i)) || p(z_i)) - D(q(\theta; h_\gamma(x)) || p(\theta)).\]

Again, we assume the latter two terms can be computed analytically. Using reparameterisation trick, we write

\[\begin{aligned} \theta &= R_\gamma(\zeta, x) \\ z_i &= T_\phi(\epsilon, x_i) \end{aligned}\]

for some transformations \(R_\gamma\) and \(T_\phi\) and random variables \(\zeta\) and \(\epsilon\) so that the output has the desired distributions.

Then the first term can be written as

\[\mathbb E_{\zeta, \epsilon} \log p(x_i; f_{R_\gamma(\zeta, x)} (T_\phi(\epsilon, x_i))),\]

so that the gradients can be computed accordingly.

Again, one may use Monte-Carlo to approximate this expetation.