While parametric models are characterized by fast convergence rates and interpretability, they have the limitation that real data are generally not well-approximated by a parametric model. The mixture model is a popular approach that generalizes conventional parametric models to a flexible class of models that can fit the data better.
Again, we consider a random sample \(X_1,\cdots, X_n\sim p\), where \(p\) is the underlying PDF. Consider a parametric model \(p(x;\theta)\) with \(\theta\in\Theta\). Instead of using a single parametric model, the mixture model consists of \(K\) models and considers the following PDF: \[\begin{equation} p(x;\eta) = \sum_{\ell=1}^K \pi_\ell p(x;\theta_\ell), \label{eq::MM1} \end{equation}\] where \(\theta_1,\cdots, \theta_K\in\Theta\) are \(K\) parameters and \(\pi_1,\cdots, \pi_K\geq0\) with \(\sum_{\ell=1}^K \pi_\ell = 1\). The parameter \(\pi_\ell\) is the weight/proportion of the \(\ell\)-th component. The quantity \[\eta = (\theta_1,\pi_1,\cdots, \theta_K, \pi_K)\] is the set of all parameters in this model. The PDF in equation [eq::MM1] is called a \(K\)-mixture model. As can be seen easily, the PDF \(p(x;\eta)\) is indeed a mixture of \(K\) individual parametric models.
The mixture model can be generated by the following two-stage procedure. First, we generate a discrete random variable \(Z \in\{1,2,3,\cdots, K\}\) such that \(P(Z=\ell) = \pi_\ell\). Then we generate \(X|Z\sim p(x;\theta_Z)\). The marginal PDF of \(X\) is \(p(x;\eta)\).
Thus, a random variable following \(p(x;\eta)\) can always be viewed as a random variable sampling from the above two-stage procedure and the variable \(Z\) is unobserved. Since \(Z\) is unobserved, it is often called a latent/hidden variable.
Example (Gaussian mixture). One of the most famous examples is the Gaussian mixture model (GMM). An example of \(3\)-GMM (\(K=3\)) is the case where \[p(x;\eta) = \pi_1 \phi(x; \mu_1,\sigma^2_1)+\pi_2 \phi(x; \mu_2,\sigma^2_2)+\pi_3 \phi(x; \mu_3,\sigma^2_3),\] where \(\phi(x;\mu,\sigma^2) = \frac{1}{\sqrt{2\pi \sigma^2}}\exp(-\frac{1}{2\sigma^2}(x-\mu)^2)\) is the PDF of \(N(\mu,\sigma^2)\). In this case, the parameter \[\eta = (\mu_1,\sigma_1^2,\pi_1, \mu_2,\sigma_2^2,\pi_2, \mu_3,\sigma_3^2,\pi_3).\] Sampling from this Gaussian mixture can be done by first generating \(Z\) from \[P(Z=\ell) = \pi_\ell, \qquad\ell=1,2,3,\] and then sampling \[X|Z \sim N(\mu_Z, \sigma^2_Z).\]
Example (kernel density estimator). The kernel density estimator (KDE) is one of the most popular nonparametric density estimators. It can be viewed as a special type of mixture model under some choices of kernel functions. To see this we consider Gaussian KDE. Given observations \(X_1,\cdots, X_n\), the KDE is \[\widehat p_h(x) = \frac{1}{nh}\sum_{i=1}^n K\left(\frac{X_i-x}{h}\right),\] where \(K(x)\) is a smooth function known as the kernel function and \(h>0\) is the smoothing bandwidth. Suppose we choose \(K(x)\) to be the Gaussian kernel, i.e., \(K(x) = \frac{1}{\sqrt{2\pi}}e^{-\frac{1}{2}x^2}\). Then the KDE is \[\begin{align*} \widehat p_h(x) & = \frac{1}{nh}\sum_{i=1}^n \frac{1}{\sqrt{2\pi}}e^{-\frac{1}{2}\left(\frac{X_i-x}{h}\right)^2}\\ & =\sum_{i=1}^n \frac{1}{n}\frac{1}{\sqrt{2\pi h^2}} e^{-\frac{1}{2h^2}\left(X_i-x\right)^2}\\ &= \sum_{i=1}^n \frac{1}{n} \phi(x; X_i, h). \end{align*}\] So it is an \(n\)-Gaussian mixture with each component’s mean being \(\mu_1=X_1,\cdots, \mu_n=X_n\) and variance being \(h^2\) and the weight being \(\frac{1}{n}\). The fact that we can generate from a mixture model implies a simple approach of generating from the KDE \(\widehat p_h(x)\). Using samples from the KDE is called the smooth bootstrap method.
Remark (continuous mixture). Here is an interesting way of re-expressing the mixture model. Let \(H(z)\) be a probability measure that places discrete probability mass \(\pi_\ell\) at value \(z=\ell\). Then equation [eq::MM1] can be written as \[p(x;\eta) = \int p(x;\theta_z) H(dz).\] As a result, if we rewrite the model \(p(x;\theta_z)\) as \(p(x|z;\theta)\), then we obtain \[p(x;\theta) = \int p(x|z;\theta) H(dz).\] This allows us to assume that \(H(z)\) has a PDF \(p(z),\) which leads to \[p(x;\theta) = \int p(x|z;\theta) p(z)dz,\] which is known as a continuous mixture model and is related to Bayesian inference and measurement error models.
There are three key issues that we need to address to use the MLE of a mixture model.
Identifiability. Without additional constraints, the model may not be identifiable. For instance, when using the Gaussian mixture, swapping the sets of parameters between two components leads to the same distribution. A common approach to deal with this is to add an additional ordering constraint so that parameters are identifiable (e.g., \(\mu_1<\mu_2<\mu_3<\cdots\)).
No closed form of MLE. In general, a mixture model will not have a closed-form MLE so we need a numerical procedure to optimize the likelihood function. A common approach is the EM algorithm (see next section)1 but you can use a gradient ascent algorithm to find the MLE as well. However, a challenge of both the EM algorithm and the gradient ascent method is that the likelihood function often has several local modes–both EM and the gradient method will get stuck at the local modes. So we often have to re-initialize the algorithm several times to increase the chance of getting the MLE, although there is no guarantee of that. See,
Chen, Y. C. (2023). Statistical inference with local optima. Journal of the American Statistical Association, 118(543), 1940-1952.
Selection of the number of mixtures. We need to choose the number of mixtures when using a mixture model. In general, how to choose this is a very difficult problem. A principled approach is to choose it via the BIC although other information criteria are also applicable. See the following paper for more details:
Fraley, C., & Raftery, A. E. (1998). How many clusters? Which clustering method? Answers via model-based cluster analysis. The computer journal, 41(8), 578-588.
When we specify a mixture model \(p(x;\eta)\), the parameter \(\eta\) is often unknown to us and has to be estimated from the data \(X_1,\cdots, X_n\). A classical approach to estimating \(\eta\) is via the MLE. Namely, \[\widehat\eta = {\sf argmax}_\eta \ell(\eta|X_1,\cdots,X_n) ={\sf argmax}_\eta \sum_{i=1}^n\ell(\eta|X_i) ={\sf argmax}_\eta \sum_{i=1}^n\log\left( \sum_{\ell=1}^K \pi_\ell p(X_i;\theta_\ell)\right).\] While the theoretical analysis of this MLE follows from the usual MLE theory, the MLE generally does not have a closed form. So we need a numerical method to compute the MLE.
The EM algorithm (expectation-maximization algorithm) is one of the most popular algorithms for computing the MLE. It is an iterative algorithm that starts with an initial guess of the MLE \(\widehat\eta^{(0)}\) and then gradually improves until convergence.
To derive the EM algorithm, we first recall the latent variable representation of this problem. Under the mixture model, each \(X_i\) can be viewed as an outcome from the two-stage sampling procedure. So \(X_i\) is associated with a latent variable \(Z_i\in\{1,2,\cdots, K\}\). When both \(X_i\) and \(Z_i\) can be observed, the joint PDF of \((x,z)\) is (recall \(\eta = (\theta,\pi)\)) \[\begin{align*} p(x,z;\eta) & = \sum_{\ell=1}^K I(z=\ell)\cdot \pi_\ell\cdot p(x;\theta_\ell)\\ &= \prod_{\ell=1}^K [ \pi_\ell\cdot p(x;\theta_\ell) ]^{I(z=\ell)}. \end{align*}\] Thus, the complete-data log-likelihood is \[\begin{equation} \begin{aligned} \ell_c(\eta|X_i=x_i,Z_i=z_i) &= \log p(x_i,z_i;\eta)\\ & = \sum_{\ell=1}^K I(z_i=\ell) \log [p(x_i;\eta|z_i=\ell)p(z_i;\eta)] \\ %&= \log p(x_i|z_i;\eta)p(z_i;\eta) \\ &= \sum_{\ell=1}^K I(z_i=\ell)\log \left( p(x;\theta_{\ell})\right) + \sum_{\ell=1}^KI(z_i=\ell) \log \pi_{\ell}. \end{aligned} \label{eq::cc} \end{equation}\]
Let \(\widehat\eta^{(t)}\) be a prior guess of the parameter. Given the observation \(X_i\), the expected complete-likelihood using this prior guess of parameter will be \[\begin{align*} Q(\eta|X_i; \widehat\eta^{(t)}) = \mathbb{E}_{Z_i\sim p(z|X_i; \widehat\eta^{(t)})}(\ell_c(\eta|X_i,Z_i)|X_i). \end{align*}\] So with all observations \(X_1,\cdots, X_n\), we define the \(Q\)-function \[Q_n(\eta;\widehat\eta^{(t)}) = \sum_{i=1}^n Q(\eta|X_i; \widehat\eta^{(t)}).\] We then improve the current estimate \(\widehat\eta^{(t)}\) to be \[\widehat\eta^{(t+1)} = {\sf argmax}_\eta Q_n(\eta;\widehat\eta^{(t)}).\]
While the above procedure may seem complicated, here is an elegant way to write it as a two-step updating procedure.
E-step. Given \(X_i\) and a previous step parameter \(\widehat\eta^{(t)} = (\widehat\pi_\ell^{(t)}, \widehat\theta_\ell^{(t)}: \ell=1,\cdots, K)\), we first compute \[\begin{equation} \tau_{i,\ell}^{(t)} = P(Z_i = \ell|X_i; \widehat\eta^{(t)}) = \frac{\widehat\pi^{(t)}_{\ell} p(X_i;\widehat\theta^{(t)}_{\ell})}{\sum_{j=1}^K\widehat\pi^{(t)}_{j} p(X_i;\widehat\theta^{(t)}_{j})}. \label{eq::tau} \end{equation}\] This is because the complete data log-likelihood in equation [eq::cc] only involves indicator functions of \(Z_i\). So the conditional expectation step in each \(Q(\eta|X_i; \widehat\eta^{(t)})\) can be done easily with the above weight. With \(\tau_{i,\ell}^{(t)}\), we compute \[Q(\eta|X_i; \widehat\eta^{(t)}) = \sum_{\ell=1}^K\tau_{i,\ell}^{(t)} [ \log p(X_i;\theta_{\ell}) + \log \pi_{\ell}]\] and \[\begin{equation} Q_n(\eta;\widehat\eta^{(t)}) = \sum_{i=1}^n \sum_{\ell=1}^K\tau_{i,\ell}^{(t)} [ \log p(X_i;\theta_{\ell}) + \log \pi_{\ell}]. \label{eq::Q1} \end{equation}\]
M-step. With the above \(Q_n\) function, we find the maximizer \[\widehat\eta^{(t+1)} = {\sf argmax}_\eta Q_n(\eta;\widehat\eta^{(t)}).\]
While the M-step may seem hard, maximizing \(Q\) can be done easily if we know how to maximize the individual model \(\ell(\theta_\ell|X_i) = \log p(x;\theta_\ell)\). To see this, we can rewrite equation [eq::Q1] as \[\begin{equation} \begin{aligned} Q_n(\eta;\widehat\eta^{(t)}) &= \sum_{i=1}^n \sum_{\ell=1}^K\tau_{i,\ell}^{(t)} [ \log p(X_i;\theta_{\ell}) + \log \pi_{\ell}]\\ & = \sum_{\ell=1}^K \underbrace{\left(\sum_{i=1}^n \tau_{i,\ell}^{(t)} \log p(X_i;\theta_{\ell}) \right)}_{A_\ell(\theta_\ell)} + \underbrace{\sum_{\ell=1}^K \left( \sum_{i=1}^n \tau_{i,\ell}^{(t)} \log \pi_{\ell}\right)}_{B}. \end{aligned} \label{eq::Q2} \end{equation}\] Each quantity in \(A_\ell(\theta_\ell)\) only involves parameter \(\theta_\ell\) so the maximization can be done easily without affecting other \(\theta_j\) with \(j\neq \ell\)! Namely, \[\widehat\theta_\ell^{(t+1)} = {\sf argmax}_{\theta_\ell} \sum_{i=1}^n \tau_{i,\ell}^{(t)} \log p(X_i;\theta_{\ell}).\] This is essentially the MLE when we are using the single model \(p(x;\theta_\ell)\) and each observation is assigned with a weight \(\tau_{i,\ell}^{(t)}\).
For the component \(B\), there is a closed-form solution to it using the Lagrange multiplier method: \[\widehat\pi_{\ell}^{(t+1)}= \frac{\sum_{i=1}^n \tau_{i,\ell}^{(t)} }{n}.\]
Thus, we actually only need to compute the weight \(\tau_{i,\ell}^{(t)}\) in equation [eq::tau] and the remaining part is just the simple MLE problem with no mixture.
Example: two-Gaussian mixture. Now we consider a simple case of two-Gaussian mixture that \[p(x;\eta) = \rho \phi(x;\mu_1,\sigma_1^2) + (1-\rho)\phi(x; \mu_2,\sigma_2^2),\] where the parameter \(\eta = (\rho,\mu_1,\mu_2,\sigma_1^2,\sigma_2^2)\). Let \(z \in\{1,2\}\) be the class label. Then the complete-case likelihood can be expressed as \[p(x,z;\eta) = [\rho \cdot \phi(x;\mu_1,\sigma_1^2)]^{I(z=1)} [(1-\rho) \cdot \phi(x;\mu_2,\sigma_2^2)]^{I(z=2)}.\] Note that the complete-case likelihood is represented in a different form as the mixture model. With this, the log-complete likelihood will be \[\log p(x,z;\eta) = I(z=1) \log [\rho \cdot \phi(x;\mu_1,\sigma_1^2)] + I(z=2)\log [(1-\rho) \cdot \phi(x;\mu_2,\sigma_2^2)].\]
E-step. Given a parameter \(\eta^{(t)} = (\rho^{(t)},\mu^{(t)}_1,\mu^{(t)}_2,\sigma_1^{(t)2},\sigma_2^{(t)2})\), the E-step will need \[\tau(x; \eta^{(t)}) = P(Z=1|X=x; \eta^{(t)}) = \frac{\rho^{(t)} \cdot \phi(x;\mu^{(t)}_1,\sigma_1^{(t)2})}{\rho^{(t)} \cdot \phi(x;\mu^{(t)}_1,\sigma_1^{(t)2})+(1-\rho^{(t)}) \cdot \phi(x;\mu^{(t)}_2,\sigma_2^{(t)2})}.\] Given observations \(X_1,\cdots, X_n\), we will compute \[\tau^{(t)}_{i} = \tau(X_i;\eta^{(t)})\] in the E-step for each \(i=1,\cdots, n\).
M-step. Given the output from E-step, we will then update our parameters via \[\begin{align*} \rho^{(t+1)} &= \frac{1}{n}\sum_{i=1}^n \tau^{(t)}_{i}\\ \mu_1^{(t+1)}& = \frac{\sum_{i=1}^n \tau^{(t)}_{i} X_i}{\sum_{j=1}^n \tau^{(t)}_j}\\ \mu_2^{(t+1)}& = \frac{\sum_{i=1}^n (1-\tau^{(t)}_{i}) X_i}{\sum_{j=1}^n (1-\tau^{(t)}_j)}\\ \sigma_1^{(t+1)2}& = \frac{\sum_{i=1}^n \tau^{(t)}_{i} (X_i- \mu_1^{(t+1)})^2}{\sum_{j=1}^n \tau^{(t)}_j}\\ \sigma_2^{(t+1)2}& = \frac{\sum_{i=1}^n (1-\tau^{(t)}_{i} )(X_i- \mu_2^{(t+1)})^2}{\sum_{j=1}^n (1-\tau^{(t)}_j)}. \end{align*}\] Note that in this nice scenario, we have a closed form in both E-step and M-step. Such an elegant property may not always hold when an EM algorithm is applied.
Here is a simple derivation showing that when we update \(\widehat\eta^{(t)}\) to \(\widehat\eta^{(t+1)}\), the observed-data likelihood function will not decrease. Formally, we have the following result:
Proposition 1. The observed-data log-likelihood value is non-decreasing during the EM updates, i.e., \[\begin{equation} \ell(\widehat\eta^{(t+1)}|X_1,\cdots,X_n) \geq \ell(\widehat\eta^{(t)}|X_1,\cdots, X_n), \label{eq::obs2} \end{equation}\] where \[\ell(\eta|X_1,\cdots, X_n) = \sum_{i=1}^n \ell_o(\eta|X_i) = \sum_{i=1}^n \log p(X_i;\eta).\]
For simplicity, here we consider the \(n=1\) case. The proof can be easily generalized to any \(n\).
The joint PDF \[p(X,Z;\eta) = p(Z|X;\eta) p(X;\eta)\] so the log-likelihood function \[\ell(\eta|X) = \log p(X,Z;\eta) - \log p(Z|X;\eta).\] This implies \[\begin{align*} \ell(\widehat\eta^{(t+1)}|X) - \ell(\widehat\eta^{(t)}|X) = [\log p(X,Z;\widehat\eta^{(t+1)}) - \log p(X,Z;\widehat\eta^{(t)})] - [\log p(Z|X;\widehat\eta^{(t+1)}) - \log p(Z|X;\widehat\eta^{(t)})]. \end{align*}\] Note that \[Q_1(\eta;\widehat\eta^{(t)}) = \mathbb{E}(\log p(X,Z; \eta)|X; Z\sim p(\cdot|X; \widehat\eta^{(t)}))\] so taking expectation on both sides of the equality with \(Z\sim p(\cdot|X; \widehat\eta^{(t)})\) leads to \[\begin{align*} \ell(\widehat\eta^{(t+1)}|X) - \ell(\widehat\eta^{(t)}|X) &= Q_1(\widehat\eta^{(t+1)};\widehat\eta^{(t)}) - Q_1(\widehat\eta^{(t)};\widehat\eta^{(t)}) - \mathbb{E}\left(\log \frac{p(Z|X;\widehat\eta^{(t+1)})}{p(Z|X;\widehat\eta^{(t)})}|X; Z\sim p(\cdot|X; \widehat\eta^{(t)})\right)\\ &\geq - \mathbb{E}\left(\log \frac{p(Z|X;\widehat\eta^{(t+1)})}{p(Z|X;\widehat\eta^{(t)})}|X; Z\sim p(\cdot|X; \widehat\eta^{(t)})\right). \end{align*}\]
Finally, by Jensen’s inequality, \[\begin{align*} \mathbb{E}\left(\log \frac{p(Z|X;\widehat\eta^{(t+1)})}{p(Z|X;\widehat\eta^{(t)})}|X; Z\sim p(\cdot|X; \widehat\eta^{(t)})\right) & \leq \log\mathbb{E}\left( \frac{p(Z|X;\widehat\eta^{(t+1)})}{p(Z|X;\widehat\eta^{(t)})}|X; Z\sim p(\cdot|X; \widehat\eta^{(t)})\right)\\ & = \log \int \frac{p(z|X;\widehat\eta^{(t+1)})}{p(z|X;\widehat\eta^{(t)})} p(z|X;\widehat\eta^{(t)})dz\\ & = \log \int p(z|X;\widehat\eta^{(t+1)})dz = \log (1)=0. \end{align*}\] Thus, we conclude that \[\ell(\widehat\eta^{(t+1)}|X) - \ell(\widehat\eta^{(t)}|X)\geq 0,\] which proves the desired result.
While Proposition 1 shows that the EM algorithm will not decrease the observed-likelihood value, it does not imply that the EM algorithm will converge to the MLE. In fact, the EM algorithm may not converge to an MLE but could get stuck at a local optimum if we do not initialize \(\widehat\eta^{(0)}\) nicely. This spurious local modes problem is a notorious issue for EM and many other similar algorithms. Thus, in practice, we often re-initialize the algorithm multiple times and choose the convergent point with the highest observed-likelihood value.
A classical convergence result of the EM algorithm is in the following paper:
Wu, C. J. (1983). On the convergence properties of the EM algorithm. The Annals of statistics, 95-103.
It proves that the EM algorithm converges to a local mode of the observed-data log-likelihood function, as long as we do not initialize the algorithm at a saddle point or a local minimum.
The following paper provides a stronger statement on the behavior of the EM algorithm when it is initialized around the MLE:
Balakrishnan, S., Wainwright, M. J., & Yu, B. (2017). Statistical guarantees for the EM algorithm: From population to sample-based analysis. Annals of Statistics, 45(1), 77-120.
It proves the so-called linear convergence property of the EM algorithm; namely, there exists a constant \(c\in(0,1)\) such that \[\|\widehat\eta^{(t)}-\widehat\eta\| \leq c^{t}\|\widehat\eta^{(0)}-\widehat\eta\|\] when the initial point \(\widehat\eta^{(0)}\in B(\widehat\eta, r)\) for some constant \(r\) and \(\widehat\eta\) is the theoretical MLE (not computable).
The mixture model offers an elegant solution to handle dependency among variables of different types (some can be continuous, some are categorical, and some are discrete).
Example. Suppose each observation has three variables \((X,Y,Z)\) such that \(X\in \mathbb{R}\), \(Y\in\{0,1,2,\cdots,\}\) is a counting number, and \(Z \in\{0,1,2,3,\cdots, N\}\) represents a score, where \(N\) is a known upper bound. It is known that these three variables are associated with each other. The question is: how can we model the joint distribution? The mixture model offers a simple solution. We assume that the joint PDF/PMF of \((X,Y,Z)\) is \[p(x,y,z) = \sum_{k=1}^K \pi_k \cdot \frac{1}{\sqrt{2\pi\sigma_k^2} }e^{-\frac{1}{2\sigma_k^2}(x-\mu_k)^2}\cdot \frac{\lambda_k^y}{y!}e^{-\lambda_k}\cdot {N \choose z} \theta_k^z (1-\theta_k)^{N-z}.\] In this case, the parameters are \[(\pi_k, \mu_k, \sigma^2_k, \lambda_k, \theta_k): k=1,\cdots,K.\] Similar to the usual mixture model, we can estimate the parameters by the MLE and compute them via the EM algorithm. The elegance of this model is that it is applicable to scenarios where variables can be of different types and the mixture structure handles the dependency among variables.
Note: this model implies conditional independence – let the latent variable \(G\) represent the component indicator. Namely, \(G=k\) means that \((X,Y,Z)\) is from the \(k\)-th component. Then we have \(X\perp Y\perp Z|G\), which is known as a local independence assumption. While \((X,Y,Z)\) are conditionally independent, they are unconditionally (marginally) dependent!
Formally, if we have a random vector \(X_1,\cdots, X_p\), the mixture of products model is \[p(x_1,\cdots, x_p) = \sum_{k=1}^K \pi_k \prod_{j=1}^p p(x_j;\theta_{j, k}),\] where each \(p(x_j;\theta_{j,k})\) is a model of variable \(X_j\) with parameter \(\theta_{j,k}\). Similar to the previous example, we can compute the MLE via the EM algorithm.
This model can be used with missing data as well and it will offer a natural clustering structure. See the following paper for more information:
Suen, D., & Chen, Y. C. (2023). Modeling Missing at Random Neuropsychological Test Scores Using a Mixture of Binomial Product Experts. arXiv preprint arXiv:2310.09384.
Here are some other variants of mixture models. For readers who are interested in more details, I recommend the following book chapter:
Gormley, I. C., & Frühwirth-Schnatter, S. (2019). Mixture of experts models. Handbook of Mixture Analysis, 271-307.
Since here we may involve some concepts of regression, we will change our notation a bit so that \(Y\) (rather than \(X\)) plays the central role in the mixture models. Let \(Y\in\mathbb{R}\) be a continuous random variable that is our primary response variable and \(Z\in\{1,2,\cdots, K\}\) be a discrete/categorical variable and \(X\in\mathbb{R}^d\) be a multivariate covariate. We only observe \((X,Y)\) and \(Z\) is unobserved; here \(Z\) often refers to the latent class label or the label of an expert. In mixture models or mixtures of experts, we often use a parametric form of the conditional densities. Depending on the relation among \(X,Y,Z\), there are 4 popular mixture-type models:
Mixture model. In the usual mixture model, there is no covariate \(X\) so we only observe \(Y\). The mixture model can be written as a graphical model with a direct arrow \(Z\rightarrow Y\). Suppose we observe both \((Y,Z)\), then \[p(y,z) = p(y|z) p(z) = p_z(y) \pi_z\Rightarrow p(y) = \sum_k p_k(y) \pi_k,\] where \(p_k(y)\) is the conditional distribution of \(Y\) given \(Z=k\) and \(\pi_k = P(Z=k)\) is the proportion of the \(k\)-th component. Let \(\theta_k\) be the parameter of \(p_k(y)\), then the marginal distribution is \[p(y;\theta) = \sum_k p(y;\theta_k) \pi_k,\] which is the usual mixture model. The Gaussian mixture model is that each \(p(y;\theta_k)\) is a Gaussian, i.e., \(p(y;\theta_k) = p(y;\mu_k,\sigma^2_k)\), where \(\mu_k\) and \(\sigma^2_k\) are the mean and variance of the \(k\)-th component.
Mixture of experts. In the mixture of experts, the model can be expressed as a graphical model with two arrows \(X\rightarrow Z\) and \(Z\rightarrow Y\). Note that \(Z\) is unobserved–we only observe \(X,Y\). In this case, \[\begin{align*} p(x,y,z) = p(y|z) p(z|x) p(x) = p_z(y) \pi_z(x) p(x)&\Rightarrow p(y,z|x) = p_z(y) \pi_z(x)\\ &\Rightarrow p(y|x) = \sum_k p_k(y) \pi_k(x). \end{align*}\] Namely, in the mixture of experts, the density of \(Y\) at each component remains the same across different \(X\). What changes with respect to \(X\) is the proportion \(\pi_k(x)\).
In this case, we need parameters for both \(p_k(y)\) and \(\pi_k(x)\), which leads to \[p(y|x;\theta,\eta) = \sum_k p(y;\theta_k) \pi_k(x;\eta).\] A popular model is to place a Gaussian model over \(p(y;\theta_k)\) and a logistic model over \(\pi_k(x;\eta)\), i.e., \[\pi_k(x;\eta) = \frac{\exp(\eta_{0,k}+ \eta_{1,k}^T x)}{\sum_{m} \exp(\eta_{0,m}+ \eta_{1,m}^T x)}.\]
Mixture of regression. The mixture of regression (a.k.a. regression mixture) looks very similar to the mixture of experts from a graphical perspective. The mixture of regression has two arrows: \(X\rightarrow Y\) and \(Z\rightarrow Y\). Thus, the difference compared to the mixture of experts is that the arrow \(X\rightarrow Z\) becomes \(X\rightarrow Y\). In this case, \[\begin{align*} p(x,y,z) = p(y|x,z) p(z) p(x) = p_z(y|x) \pi_z p(x)&\Rightarrow p(y,z|x) = p_z(y|x) \pi_z\\ &\Rightarrow p(y|x) = \sum_k p_k(y|x) \pi_k. \end{align*}\] In particular, the conditional mean (regression function) becomes \[m(x) = \mathbb{E}(Y|X=x) = \int \sum_k p_k(y|x) \pi_k dy = \sum_{k} \pi_k\cdot m_k(x),\] where \(m_k(x) = \mathbb{E}(Y|Z=k,X=x)\) is the regression function of the \(k\)-th component. So the regression function is written as a mixture of several regression functions. Note that the proportion \(\pi_k\) is independent of \(X\).
Mixture of expert regression. The mixture of expert and the mixture of regression can be combined into the mixture of expert regression. It corresponds to the graph with three arrows: \(X\rightarrow Y\), \(X\rightarrow Z\), and \(Z\rightarrow Y\). In this case, \[\begin{align*} p(x,y,z) = p(y|x,z) p(z|x) p(x) = p_z(y|x) \pi_z(x) p(x)&\Rightarrow p(y,z|x) = p_z(y|x) \pi_z(x)\\ &\Rightarrow p(y|x) = \sum_k p_k(y|x) \pi_k(x). \end{align*}\] The conditional mean (regression function) is \[m(x) = \mathbb{E}(Y|X=x) = \int \sum_k p_k(y|x) \pi_k(x) dy = \sum_{k} \pi_k(x)\cdot m_k(x).\] So it is the mixture of regression with the proportion \(\pi_k(x)\) being allowed to change with respect to \(x\).
Variational inference (VI; also known as variational approximation) is a popular tool in machine learning. It has become more and more popular in statistics communities as well. In short, VI is a method to approximate an intractable quantity using a tractable quantity. It can be used in both Frequentist estimation as well as Bayesian inference.
However, the fact that VI can be used in both Frequentist and Bayesian inference has made VI sometimes confusing. Here we will discuss how VI can be used in both cases and how the two problems are associated.
Consider a regular latent variables problem where we observe IID \[X_1,\cdots, X_n\sim p\] and each observation has a latent variable \(Z\) that is unobserved. Namely, the complete data should be \[(X_1,Z_1),\cdots, (X_n,Z_n)\] but we only observe \(X_1,\cdots, X_n\).
In the latent variable problem, we often place a parametric model on the complete-data distribution: \[p(x,z;\lambda),\quad \lambda\in\Lambda.\] This parametric model implies the observed model \[p(x;\lambda) = \int p(x,z;\lambda)dz\] that generates our observations.
In this case, the MLE is defined as \[\widehat\lambda = {\sf argmax}_\lambda \frac{1}{n}\sum_{i=1}^n\log p(X_i;\lambda).\] Namely, we maximize the observed log-likelihood \(\ell(\lambda|x) = p(x;\lambda)\). A population version of this problem is \[\lambda^* = {\sf argmax}_\lambda \mathbb{E}(\log p(X_1;\lambda)).\]
In most cases, the MLE does not have a closed-form so we need to use a numerical procedure such as the EM algorithm to compute it. However, the EM algorithm may not have a simple form and can still be intractable (this occurs a lot when the \(Q\) function in the EM algorithm is complicated).
The VI offers a remedy to this problem. Note that we can write the observed model as \[\begin{align*} p(x;\lambda) &= \int p(x,z;\lambda)dz \\ &= \int \frac{p(x,z;\lambda)}{q(z;\omega)}q(z;\omega)dz \\ &= \mathbb{E}_{Z\sim q(\cdot;\omega)}\left(\frac{p(x,Z;\lambda)}{q(Z;\omega)}\right), \end{align*}\] where \(\omega\in \Omega\) is another class of parameters and \(Q = \{q(\cdot;\omega):\omega\in\Omega\}\) is called the variational family. The variational family often consists of parametric densities that are easy to compute and sample.
Using Jensen’s inequality, the observed log-likelihood function becomes \[\begin{align*} \ell(\lambda|x)& = \log p(x;\lambda)\\ &= \log \int p(x,z;\lambda)dz\\ &= \log \int \frac{p(x,z;\lambda)}{q(z ;\omega)} q(z;\omega)dz\\ & = \log \mathbb{E}_{Z\sim q(\cdot;\omega)}\left(\frac{p(x,Z;\lambda)}{q(Z;\omega)}\right)\\ &\geq \mathbb{E}_{Z\sim q(\cdot;\omega)}\left(\log\frac{p(x,Z;\lambda)}{q(Z;\omega)}\right)\\ & = \mathbb{E}_{Z\sim q(\cdot;\omega)} \log p(x,Z;\lambda) - \mathbb{E}_{Z\sim q(\cdot;\omega)} (\log q(Z;\omega))\\ & = {\sf ELBO}(\omega,\lambda|x). \end{align*}\] The quantity \({\sf ELBO}(\omega,\lambda|x)\) is called the evidence lower bound. Note that sometimes people write it as \({\sf ELBO}(q)\) and abbreviate \(\lambda, x\) but here for completeness I will keep both of them.
The quantity \(\mathbb{E}_{Z\sim q(\cdot;\omega)} \log p(x,Z;\lambda)\) in the ELBO can be interpreted as an approximation to the E-step in the EM algorithm. This quantity behaves like we are ‘imputing’ the missing/latent variable \(Z\) by sampling from the distribution \(q(\cdot;\omega)\). When \(q(\cdot; \omega) = p(z|x;\lambda)\), this is essentially the E-step in the EM algorithm.
When we observe \(n\) observations, we write \[{\sf ELBO}_n(\omega_1,\cdots, \omega_n,\lambda) = \frac{1}{n}\sum_{i=1}^n{\sf ELBO}(\omega_i,\lambda|X_i),\] where we allow a different variational parameter \(\omega_i\) for different observed \(X_i\). This is a particularly important feature in variational inference–we are approximating each single conditional distribution \(p(z|x;\lambda)\) by another distribution \(q(z;\omega)\). Thus, one can see that the optimal value of \(\omega\) depends on \(x\) and \(\lambda\), so we may informally write the optimal value as \[\omega^*(x,\lambda) = {\sf argmax}_\omega {\sf ELBO}(\omega, \lambda|x).\]
With the \({\sf ELBO}\), the idea of variational inference is very simple. Instead of maximizing the intractable log-likelihood, we maximize \({\sf ELBO}\) to find our estimator. Namely, our estimator is \[(\widehat\omega_{\sf VI},\widehat\lambda_{\sf VI}) = {\sf argmax}_{\omega_1,\cdots, \omega_n,\lambda} {\sf ELBO}_n(\omega_1,\cdots, \omega_n,\lambda),\] where \[\widehat\omega_{\sf VI} = (\widehat\omega_{ \sf VI, 1}, \cdots, \widehat\omega_{\sf VI, n})\] such that \(\widehat\omega_{\sf VI, j} = \omega^*(X_i, \widehat\lambda_{\sf VI})\).
Again, in general there is no closed-form to the above estimators. But we can solve this maximization problem numerically. The power of VI is that the variational family is designed so that evaluation and computation are easy. So the expectation \(\mathbb{E}_{Z\sim q(\cdot;\omega)} \cdots)\) that appears in the \({\sf ELBO}\) is a tractable quantity.
When \(z\in\mathbb{R}^d\) is multivariate, choosing a tractable family \(q(\cdot; \omega)\) may be challenging. A famous method called mean field variational family is the collection of densities \[q(z;\omega) = \prod_{j=1}^d q(z_j;\omega_j),\] where \(z = (z_1,\cdots, z_d)\). In this case, sampling of \(z\) can be decomposed into sampling each coordinate independently and the optimization of \(\omega_j\) can be obtained by the coordinate ascent variational inference algorithm in the following book:
Bishop, C. (2006). Pattern Recognition and Machine Learning. Springer New York.
One thing to keep in mind is that when using VI, we are no longer solving the original problem. In general, the VI estimator will not converge to the MLE (in the Frequentist setting) but instead, it will still converge to a population quantity; see the following paper:
Chen, Y. C., Wang, Y. S., & Erosheva, E. A. (2018). On the use of bootstrap with variational inference: Theory, interpretation, and a two-sample test example. The Annals of Applied Statistics, 12(2), 846-876.
A good review from the Bayesian perspective is
Blei, D. M., Kucukelbir, A., & McAuliffe, J. D. (2017). Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518), 859-877.
Consider a simple Bayesian setting where \[X|\theta\sim p(x|\theta),\quad \theta\sim \pi(\theta).\] Here, \(X\) is the random variable for the observation and \(\theta\) is the underlying parameter of the distribution and \(\pi\) is the prior. The Bayesian inference relies heavily on the posterior distribution (density): \[\pi(\theta|x) = \frac{p(x|\theta) \pi(\theta)}{p(x)} \propto p(x|\theta) \pi(\theta) = p(x,\theta).\] Here we write \(p(x,\theta) = p(x|\theta) \pi(\theta)\) for simplicity. In addition to the posterior distribution, the evidence \(p(x) = \int p(x,\theta)d\theta\) is also sometimes of interest.
Since \(p(x,\theta) = p(x|\theta) \pi(\theta)\), it is often very easy to evaluate the value of \(p(x,\theta)\) for any given \(x,\theta\). Although the joint distribution is easy to compute, both the posterior and the evidence are often intractable. Take the evidence as an example, it can be written as the integral \(p(x) = \int p(x,\theta)d\theta\). When the dimension of \(\theta\) is large (say a mixture model), this integration is very difficult to evaluate even if we can easily compute \(p(x,\theta)\).
Here is an interesting note. The problem of evaluating the evidence and the problem of evaluating the posterior distribution are the same via the following relation: \[\pi(\theta|x) = \frac{p(x,\theta)}{p(x)} .\] \(p(x,\theta)\) is tractable so we can easily convert the evidence and the posterior to each other.
To obtain a tractable approximation of \(\pi(\theta|x)\), we use the idea of VI. The quantity \(\theta\) now plays the role of latent variable \(z\) in the previous section. Let \(q(\theta; \omega)\) be a computable density and \[Q = \{q(\cdot;\omega): \omega\in\Omega\}\] be the variational family. We attempt to find the best density \(q(;\omega^*)\in Q\) such that \[q(\cdot;\omega^*) = {\sf argmin}_{q\in Q} {\sf KL}(q(\cdot)||\pi(\cdot|x)),\] where \({\sf KL}\) is the Kullback-Leibler divergence. Namely, \[\omega^* = {\sf argmin}_{\omega} {\sf KL}(q(\cdot;\omega)||\pi(\cdot|x)),\]
However, minimizing the KL divergence may again run into the same computation problem (need to evaluate the integral). The idea of VI uses the following insight: \[\begin{align*} \log p(x) &= {\sf KL}(q(\cdot;\omega)||\pi(\cdot|x)) +\mathbb{E}_{\theta\sim q(\cdot;\omega)}(\log p(x,\theta))-\mathbb{E}_{\theta\sim q(\cdot;\omega)}(\log q(\theta;\omega)) \\ &= {\sf KL}(q(\cdot;\omega)||\pi(\cdot|x)) +{\sf ELBO}(\omega|x). \end{align*}\] Therefore, minimizing the KL divergence is equivalent to maximizing \({\sf ELBO}\).
The variational approximation chooses \[\omega^*_{\sf VI} = {\sf argmax}_{\omega}{\sf ELBO}(\omega|x)\] and uses \(q(\theta;\omega^*_{VI} )\) as an approximation to \(\pi(\theta|x)\).
In the case of observing \(X_1,\cdots, X_n\), the \({\sf ELBO}\) will be \[\begin{align*} {\sf ELBO}_n(\omega) &= \mathbb{E}_{\theta\sim q(\cdot;\omega)}(\log p(X_1,\cdots, X_n,\theta))-\mathbb{E}_{\theta\sim q(\cdot;\omega)}(\log q(\theta;\omega)) \\ &= \mathbb{E}_{\theta\sim q(\cdot;\omega)}\left(\log\pi(\theta) + \sum_{i=1}^n\log p(X_i|\theta)\right)-\mathbb{E}_{\theta\sim q(\cdot;\omega)}(\log q(\theta;\omega)) \end{align*}\] and we estimate \(\omega\) using \[\widehat\omega_{\sf VI} = {\sf argmax}_{\omega}{\sf ELBO}_n(\omega).\] Again, this maximization is often tractable since the expectation is with respect to the variational distribution \(q\), which is by design easy to compute.
The posterior is then approximated by \[\pi(\theta|X_1,\cdots, X_n) \approx p(\theta; \widehat\omega_{\sf VI})\] and the evidence is approximated by \[p(X_1,\cdots, X_n)\approx \frac{\pi(\theta)\prod_{i=1}^np(X_1|\theta)}{p(\theta; \widehat\omega_{\sf VI})}.\] Note that the approximated evidence may depend on \(\theta\) because it is an approximation rather than an exact value.
There are many good tutorials on this topic. I will recommend the following tutorial:
An Introduction to Variational Autoencoders by Diederik P. Kingma and Max Welling. https://arxiv.org/pdf/1906.02691
A VAE is a generative model that learns a mapping from a high-dimensional data space \(x \in \mathbb{R}^d\) to a potentially lower-dimensional latent space \(z \in \mathbb{R}^p\) and back. A common example of such \(x\) is image data where \(x\) is the vectorized pixel values of an image.
Before the VAE, the standard (or denoising) autoencoder was a widely used neural network for unsupervised learning. Its primary goal is not generation but dimensionality reduction and feature learning. The architecture is composed of two main parts: an encoder and a decoder.
Encoder: The encoder, \(f_\phi(x)\), is a neural network that takes high-dimensional input data \(x\) and maps it to a compressed, lower-dimensional representation \(z\), often called the bottleneck or latent code. This mapping is deterministic: \(z = f_\phi(x)\).
Decoder: The decoder, \(g_\theta(z)\), is another neural network that takes the latent code \(z\) and attempts to reconstruct the original input data. The output is denoted as \(\widehat{x} = g_\theta(z) = g_\theta(f_\phi(x))\).
The entire network is trained end-to-end to minimize the reconstruction error, which is the difference between the original input \(x\) and its reconstruction \(\widehat{x}\). A common loss function for this is the Mean Squared Error (MSE): \[\begin{equation} L(\phi, \theta) = \frac{1}{n} \sum_{i=1}^{n} ||x_i - \widehat{x}_i||^2 \end{equation}\]
While the idea of autoencoder is powerful, it is not easily trained because there are too many choices of \(f_\phi\) and \(g_\theta\), making it hard to use in practice. Also, the deterministic mapping in the encoder and decoder limits their performance in practice.
While the variational autoencoder sounds similar to the autoencoder, it is constructed in a very different way.
Starting from a distribution over the latent space \(p_{\theta}(z)\) (or \(p(z)\) since we often choose this distribution as a uniform distribution or a standard Gaussian), we first specify a decoder model \(p_\theta(x|z)\) that generates data from a latent \(z\). The model parameter \(\theta\) is often very high-dimensional. A popular example is \[p_\theta(x|z) \sim N(x; \mu_\theta(z), \Sigma_\theta(z)) .\] Since \(x\) is high-dimensional (an image), \(\mu_\theta\) and \(\Sigma_\theta\) will also be very high-dimensional objects so the underlying parameter \(\theta\) will naturally be a very high-dimensional object.
The marginal likelihood of the observation \(x\) is given by integrating over all possible latent variables: \[\begin{equation} p_\theta(x) = \int p_\theta(x|z) p_\theta(z) dz \label{eq:vae_marginal} \end{equation}\] We want to find the parameters \(\theta\) that maximize the log-likelihood for the dataset, \(\sum_i \log p_\theta(X_i)\). A key problem is that this integral is numerically intractable for complex decoders like neural networks.
To solve this, we introduce an approximation to the reversed conditional probability \(p_\theta(z|x)\). This approximation is another neural network, \(q_\phi(z|x)\), which is the encoder in the VAE.
We then optimize the Evidence Lower Bound (ELBO) on the log-likelihood: \[\begin{equation} \begin{aligned} \log p_\theta(x) &= \log \int p_\theta(x,z) \frac{q_\phi(z|x)}{q_\phi(z|x)} dz \\ &\geq \int q_\phi(z|x) \log \frac{p_\theta(x,z)}{q_\phi(z|x)} dz \\ & = \int q_\phi(z|x) \log p_\theta(x,z)dz - \int q_\phi(z|x) \log q_\phi(z|x)dz\\ &= \mathbb{E}_{Z \sim q_\phi(Z|x)}[\log p_\theta(x,Z)] - \mathbb{E}_{Z \sim q_\phi(Z|x)}[\log q_\phi(Z|x)]\\ &= \mathbb{E}_{Z \sim q_\phi(Z|x)}\left[\log \frac{p_\theta(x,Z)}{q_\phi(Z|x)}\right]\\ & = {\sf ELBO}(\theta,\phi|x). %D_{KL}(q_\phi(Z|x) || p_\theta(Z)) \end{aligned} \label{eq::VAE} \end{equation}\] This ELBO is now tractable and can be optimized with gradient descent using Monte Carlo estimation (i.e., sampling from the encoder). When we have a random sample \(X_1,\cdots, X_n\), we then find the estimator \[(\widehat\theta, \widehat\phi) = {\sf argmin}_{\theta, \phi} \frac{1}{n}\sum_{i=1}^n {\sf ELBO}(\theta,\phi|X_i).\] The model \(p_{\widehat\theta}(x|z)\) is our trained decoder and the model \(q_{\widehat\phi}(z|x)\) is our trained encoder from the data.
In the VAE framework, the encoder is constructed naturally during the variational approximation process to make the intractable marginal likelihood computable. Note that the decoder model \(p_{\theta}(x|z)\) and the latent variable model \(p_\theta(z)\) imply the conditional probability \(p_\theta(z|x)\), which can be viewed as an ‘idealized encoder’. However, such an encoder cannot be used unless we have an estimate of \(\theta\). The variational distribution \(q_\phi(z|x)\) offers a numerically tractable approximation to such an idealized encoder.
Now we compare VAE to the problem in Section 5.1. Suppose we want to model the distribution of a random variable/vector \(X\) with a latent variable \(Z\) such that \(p_\theta(x|z)\) is a given family with unknown parameter \(\theta\) and \(p(z)\) is a prior distribution that we pre-specify. Then the log-likelihood function of \(\theta\) given \(x\) is \[\ell(\theta|x) = \log p_\theta(x) = \log \int p_\theta(x|z)p(z)dz.\] The conventional variational inference will use the ELBO from \[\begin{align*} \ell(\theta|x) &= \log \int p_\theta(x,z)dz\\ &\geq \int q_\omega(z) \left[\log \frac{p_\theta(x,z)}{q_\omega(z)}\right]dz\\ & = {\sf ELBO}(\theta, \omega|x) \end{align*}\] such that the optimal \(\omega^*(x,\theta) = {\sf argmax}_\omega {\sf ELBO}(\theta, \omega|x)\) is a function of \(x\) and \(\theta\). This is because our variational family is a distribution of \(z\) only. When we have \(n\) observations, The ELBO will become \[\begin{align*} \sum_{i=1}^n \ell(\theta|x_i) &= \sum_{i=1}^n\log \int p_\theta(x_i,z_i)dz_i\\ &\geq \sum_{i=1}^n\int q_{\omega_i}(z_i) \left[\log \frac{p_\theta(x_i,z_i)}{q_{\omega_i}(z_i)}\right]dz_i\\ & = {\sf ELBO}(\theta, \omega_1,\cdots, \omega_n|x_1,\cdots, x_n). \end{align*}\] This is because we need a different parameter for each \(z_i\) to improve the evidence lower bound. So we need \(n\) distinct values of \(\omega\), which could be computationally slow.
In the VAE, the variational family is a conditional distribution of \(z|x\) and we obtain a slightly different ELBO \[\begin{align*} \ell(\theta|x) &= \log \int p_\theta(x,z)dz\\ &\geq \int q_\phi(z|x) \left[\log \frac{p_\theta(x,z)}{q_\phi(z|x)}\right]dz\\ & = {\sf ELBO}_{VAE}(\theta, \phi|x). \end{align*}\] In the VAE, the parameter \(\phi\) does NOT depend on \(x\). It is a fixed set of parameters that works for all \(x\) because the distribution for our variational model is already a conditional \(q_\phi(z|x)\). This is called amortized inference. The method from VAE is computationally easier than the traditional ELBO when \(n\) is large, but there will be an amortized gap between the best \({\sf ELBO}_{VAE}\) can achieve relative to the conventional \({\sf ELBO}\). To see this, when we have \(n\) observations \(x_1,\cdots,x_n\), the VAE simply optimizes \[\begin{align*} \sum_{i=1}^n\ell(\theta|x_i) &= \sum_{i=1}^n\log \int p_\theta(x_i,z_i)dz_i\\ &\geq \sum_{i=1}^n\int q_\phi(z_i|x_i) \left[\log \frac{p_\theta(x_i,z_i)}{q_\phi(z_i|x_i)}\right]dz_i\\ & = {\sf ELBO}_{VAE}(\theta, \phi|x_i). \end{align*}\] This is because the variational distribution \(q_\phi(z|x)\) already accounts for the fact that the likelihood of a different \(x\) should be approximated by a different distribution
The diffusion model is another popular approach for learning a generative model under the deep learning framework. The core idea is to construct a generative model \(p(x)\) via a sequential process that learns the true data-generating distribution \(p_0(x)\). Due to the sequential generating procedure, this approach can be viewed as creating a diffusion process that moves from one point to another.
The denoising/variational diffusion model is a popular diffusion model that has strong connections to the VAE. In the denoising diffusion model, we create a forward and a reverse process that move an observation (often an image) \(X_0\) to pure noise (forward) so that we can move the noise backward to an observation (reverse).
The diffusion model is a popular topic in modern AI/ML. For interested readers, I highly recommend the following tutorial:
Understanding Diffusion Models: A Unified Perspective by Calvin Luo. https://arxiv.org/pdf/2208.11970
The forward process gradually adds Gaussian noise to an observation/image \(X_0\) over a series of \(T\) steps. This creates a sequence of noisy observations/images \(X_1, \dots, X_T\). \[\begin{equation} X_t = \sqrt{1 - \beta_t} X_{t-1} + \sqrt{\beta_t} Z_t, \quad \text{where } Z_t \sim N(0, I) \label{eq:forward_step} \end{equation}\] Due to the properties of Gaussians, we can sample \(x_t\) at any arbitrary step \(t\) directly from the original observation \(x_0\): \[\begin{equation} X_T | X_0 \sim N(\alpha_T X_0, \sigma_T^2 I) \label{eq:forward_direct} \end{equation}\] where \(\alpha_T\) and \(\sigma_T^2\) are functions of the noise schedule \(\beta_t\). Therefore, once we have learned all the parameters (learning these parameters will be done in the variational approximation), we do not need to generate each \(x_t\) sequentially – we can just directly jump from the original observation to any step. After enough steps, \(X_T\) becomes indistinguishable from pure Gaussian noise.
The interesting part of the diffusion model is the reverse process. The reverse process will reverse the above procedure that we start from a noise and then gradually come back to an observation/image. This allows us to start with a random noise \(X_T \sim N(0, I)\) and then gradually denoise it back to a clean sample \(X_0\).
Here is an interesting part. While the forward model \(X_t|X_{t-1}\) is a Gaussian, the reverse model \(X_{t-1}|X_t\) is generally NOT a Gaussian. However, we still model the reverse part as a Gaussian: \[\begin{equation} p_\theta(x_{t-1} | X_t=x_t) = N(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)), \label{eq:reverse_step} \end{equation}\] where the mean \(\mu_\theta(x_t, t)\) and variance \(\Sigma_\theta(x_t, t)\) are learned by a neural network.
To learn the parameters in equation \(\ref{eq:reverse_step}\), we will consider the maximum likelihood approach. Note that the forward model implies that \[p(x_t|x_{t-1},\cdots, x_0) = p(x_t|x_{t-1});\] namely, we have a Markov property of the process. Thus. the joint distribution of \(x_0,\cdots, x_T\) is \[\begin{equation} p(x_0,x_1,\cdots, x_T) = p(x_0)p(x_1|x_0)\cdots p(x_T|x_{T-1}) \label{eq::MC} \end{equation}\] from the forward process. The Markov property of equation \(\ref{eq::MC}\) also implies the decomposition of the reverse process. To see this, equation \(\ref{eq::MC}\) implies that \[p(x_t|x_{t+1},\cdots, x_T) = \frac{p(x_t)p(x_{t+1}|x_t)p(x_{t+2}|x_{t+1})\cdots p(x_{T}|x_{T-1})}{p(x_{t+1})p(x_{t+2}|x_{t+1})\cdots p(x_{T}|x_{T-1})} = \frac{p(x_t)p(x_{t+1}|x_t)}{p(x_{t+1})} = p(x_t|x_{t+1}).\] Thus, we can rewrite equation \(\ref{eq::MC}\) using the reverse model \[\begin{equation} p(x_0,x_1,\cdots, x_T) = p(x_T)p(x_{T-1}|x_T)\cdots p(x_0|x_1) =p(x_T)p_\theta(x_{T-1}|x_T)\cdots p_\theta(x_0|x_1) \label{eq::MC2} \end{equation}\] and we know that \(p(x_T)\) is a pure Gaussian noise when \(T\) is sufficiently large. Thus, the probability (density) of the observed data is \[p_\theta (x_0 ) = \int p(x_T)p_\theta(x_{T-1}|x_T)\cdots p_\theta(x_0|x_1) dx_1dx_2\cdots dx_T.\] So the log-likelihood is \[\ell(\theta|x_0) = \log \int p(x_T)p_\theta(x_{T-1}|x_T)\cdots p_\theta(x_0|x_1) dx_1dx_2\cdots dx_T.\] Ideally, we want to use the MLE but similar to the VAE problem, this likelihood function involves an intractable integral over \(x_1,\cdots, x_T\). To resolve this issue, we use variational approximation again.
Using a similar derivation to equation \(\ref{eq::VAE}\), we are able to construct an ELBO to \(\ell(\theta|x_0)\) as \[\begin{equation} \begin{aligned} \ell(\theta|x_0) &= \log \int \underbrace{p(x_T)p_\theta(x_{T-1}|x_T)\cdots p_\theta(x_0|x_1)}_{p(x_0,\cdots, x_T)} dx_1dx_2\cdots dx_T\\ &\geq \int q_\phi(x_1,\cdots, x_T|x_0) \log \frac{p_\theta(x_0,\cdots, x_T)}{q_\phi(x_1,\cdots, x_T|x_0)} dx_1\cdots dx_T\\ & = \mathbb{E}_{X_{1:T}|x_0\sim q_\phi}\left[\log \frac{p_\theta(x_0, X_1\cdots, X_T) }{q_\phi(X_1,\cdots, X_T|x_0)}\right]\\ & = {\sf ELBO}(\theta, \phi|x_0). \end{aligned} \label{eq::elbo::diffusion} \end{equation}\] When we have IID observations \(\mathbf{X}_1,\cdots, \mathbf{X}_n\), the we will compute the estimators as \[(\widehat\theta, \widehat\phi) = {\sf argmin}_{\theta,\phi}\frac{1}{n}\sum_{i=1}^n {\sf ELBO}(\theta, \phi|\mathbf{X}_i).\]
The variational distribution \(q_\phi\) is our model in equation [eq:forward_step], so it implies that we can decompose \(q_\phi\) as \[\begin{align*} q_\phi(x_0,\cdots, x_T) = q_\phi(x_0) q_\phi(x_1|x_0)\cdots q_\phi(x_T|x_{T-1}) \end{align*}\] and it is easy to sample from any \(q(x_t|x_0)\) as in equation [eq:forward_direct].
Note that the variational parameters under equation [eq:forward_step] are \(\beta_1,\cdots, \beta_T\). Thus, \(\phi = (\phi_1,\cdots, \phi_T)\) such that \(\phi_t = \beta_t\). This further implies that each \(q_\phi(x_t|x_{t-1}) = q_{\phi_t}(x_t|x_{t-1})\) so optimizing the ELBO in equation [eq::elbo::diffusion] can be further decomposed into several subproblems.
Suppose \(\theta = (\theta_1,\cdots, \theta_T)\) such that each submodel \(p_\theta(x_{t-1}|x_t) = p_{\theta_t}(x_{t-1}|x_t)\) only depends on local parameter \(\theta_t\). Due to the Markov property that we can decompose \[p_\theta(x_0,x_1,\cdots, x_T) =p(x_T)p_\theta(x_{T-1}|x_T)\cdots p_\theta(x_0|x_1),\] we further have \[\log p_\theta(x_0,x_1,\cdots, x_T) = \log p(x_T)+ \sum_{t=T}^1 \log p_{\theta_t}(x_{t-1}|x_t).\]
Then the parameter \(\theta\) in ELBO can be further decomposed as \[\begin{align*} \mathbb{E}_{X_{1:T}|x_0\sim q_\phi}\left[\log {p_\theta(x_0, X_1\cdots, X_T) }\right] & = \mathbb{E}_{X_{1:T}|x_0\sim q_\phi}\left[\log p(x_T)+ \sum_{t=T}^1 \log p_\theta(X_{t-1}|X_t)\right] \\ & = \int q_\phi(x_T|x_0) \log p(x_T)dx_T + \sum_{t=T}^1 q_\phi(x_t,x_{t-1}|x_0)\log p_{\theta_t}(x_{t-1}|x_t) dx_tdx_{t-1}. \end{align*}\] Given \(\phi\), it is very easy to perform a Monte Carlo approximation to the above quantity because sampling from \(q_\phi(x_T|x_0)\) and \(q_\phi(x_t,x_{t-1}|x_0) = q_\phi(x_t|x_{t-1})q_\phi(x_{t-1}|x_0)\) can be very easy when we choose the variational family similar to the forward process (see equation [eq:forward_direct]).
The denoising diffusion model is a variational autoencoder. The forward process is the encoder and the reverse process is the decoder. And the latent variable \(z\) in the VAE is the noisy variable \(X_T\).
Another way to see how denoising diffusion is a VAE is that in our likelihood model \(p_\theta(x_0)\), we do not use any information about the forward process. This is like in VAE where our original likelihood model only involves the decoder model. The forward process in the denoising diffusion model comes from variational approximation, similar to how the encoder model in VAE is created for variational approximation.
An alternative diffusion model is the score-based diffusion model. The objective aims at learning the score function of the data distribution, defined as the gradient of the log-probability density with respect to the data. \[\begin{equation} s(x) = \nabla_x \log p(x), \end{equation}\] where \(x\in\mathbb{R}^d\). Note that this score function is NOT the conventional score function in likelihood inference; in likelihood inference, the score is the gradient of the log-likelihood with respect to parameter \(\theta\), not the variable \(x\).
In practice, we model the score function via a neural network \(f_\theta(x) \in\mathbb{R}^d\) where \(\theta\) is learned by the following risk minimization: \[\begin{align*} \min_\theta \int \|f_\theta(x) - s(x)\|^2 p(x)dx & = \min_\theta \int [\|f_\theta(x)\|^2 +2 \nabla_x\cdot f_\theta(x)] p(x)dx + C\\ & = \min_\theta \mathbb{E}\left[\|f_\theta(X)\|^2 +2 \nabla_x\cdot f_\theta(X) \right]+C, \end{align*}\] where the first equality follows from integration by parts. With the above objective function, estimating \(\theta\) can be done via simple empirical risk minimization: \[\widehat\theta = {\sf argmin}_\theta \frac{1}{n}\sum_{i=1}^n \|f_\theta(X_i)\|^2 +2 \nabla_x\cdot f_\theta(X_i).\]
Once we have a good approximation of the score \(s(x)\), we can generate samples using Langevin dynamics. This is an iterative process that follows the gradient of the log density: \[\begin{equation} x_{t+1} = x_t + \frac{\epsilon}{2} s(x_t) + \sqrt{\epsilon} z_t, \quad z_t \sim N(0, I) \label{eq:langevin} \end{equation}\] The stationary distribution of this process (when \(\epsilon\rightarrow0\) and \(t\rightarrow \infty\)) is the target data distribution \(p(x)\). The model objective is to train a neural network \(f_\theta\) to approximate this score function.
An interesting feature of score-based diffusion is that the generation of a conditional model can be done in a cool way. To build a conditional model that can sample from \(p(x|y)\) (e.g., generate an image \(x\) given a text label \(y\)), we need the conditional score \(s(x|y) = \nabla_x \log p(x|y)\). Using Bayes’ rule, this can be decomposed: \[\begin{equation} \underbrace{s(x|y)}_{\text{Conditional Score}} = \underbrace{s(x)}_{\text{Unconditional Score}} + \underbrace{\nabla_x \log p(y|x)}_{\text{Classifier Gradient}} \end{equation}\] The unconditional score \(s(x)\) can be learned by a standard pre-trained diffusion model. The second term can be provided by training a separate classifier that predicts \(y\) from a noisy input \(x\) and then using its gradients to guide the diffusion process.
Normalizing flows are very popular deep generative models. I recommend the following two tutorials for interested readers:
1. Kobyzev, I., Prince, S. J., & Brubaker, M. A. (2020). Normalizing flows: An introduction and review of current methods. IEEE transactions on pattern analysis and machine intelligence, 43(11), 3964-3979.
2. Papamakarios, G., Nalisnick, E., Rezende, D. J., Mohamed, S., & Lakshminarayanan, B. (2021). Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57), 1-64.
Normalizing flows are based on the inverse CDF formula of a continuous random variable. To start with, we first discuss an interesting property of the inverse CDF. Suppose \(Z\in\mathbb{R}\) is a continuous random variable with a CDF \(F_Z\) that has a well-defined inverse \(F^{-1}_Z\). Then it is well-known that you can generate \(Z\) by sampling from a uniform random variable \(U\sim {\sf Uni}[0,1]\) and then using the following property \[Z \overset{d}{=} F^{-1}_Z(U).\] In other words, \(F_Z(Z) \overset{d}{=} U\).
This phenomenon also occurs in the multivariate case, but with a small modification.
Suppose \(Z\in\mathbb{R}^d\) is a multivariate continuous random variable with a CDF \(F_Z\) and PDF \(p_Z\). For any vector \(z\in\mathbb{R}^d\), we denote \[z_{<j} = (z_1,\cdots, z_{j-1}),\qquad z_{\leq j} = (z_1,\cdots, z_j).\] We cannot directly apply the inverse CDF in the multivariate case because the inverse \(F_{Z}^{-1}\) will be a set, not a point (\(d-1\) dimensional manifold under suitable conditions).
Now we consider the partial CDF: \[\begin{equation} G_j(z) = F_j(z_j|z_{<j}) = P(Z_j\leq z_j|Z_{<j} = z_{<j}). \label{eq::NF1} \end{equation}\] Namely, \(G_j(z)\) only depends on the variables \(z_1,\cdots, z_j\) and is the conditional CDF of \(z_j\) given the ‘past’ \(z_1,\cdots, z_{j-1}\). Clearly, the distributions \(F_j\) and \(G_j\) are well defined.
Given a vector \(z\in\mathbb{R}^d\) and the functions \(G_1,\cdots, G_d\), we then define \[\begin{equation} w_j = F_j(z_j|z_{<j}). \label{eq::NF2} \end{equation}\] Since \(F_j\) is a CDF, we have the following iterative relation \[\begin{equation} z_j = F^{-1}_j(w_j|z_{<j}). \label{eq::NF3} \end{equation}\] By iteratively applying equation \(\ref{eq::NF2}\), we are able to construct \(w =(w_1,\cdots, w_d)^T\).
Here is an interesting question:
Suppose we apply equation \(\ref{eq::NF2}\) to the random vector \(Z\) and let the resulting vector be \(W\), what is the distribution of \(W\)?
Using equations \(\ref{eq::NF1}\) and \(\ref{eq::NF2}\), \(W = (G_1(Z),\cdots, G_d(Z))^T\) so we can apply the Jacobian method to investigate the PDF of \(W\). Note that each function \(G_j\) has the following interesting differential: \[\frac{\partial G_j(z)}{\partial z_j} = \frac{\partial P(Z_j\leq z_j|Z_{<j} = z_{<j})}{\partial z_j} = p(z_j|z_{<j}),\qquad \frac{\partial G_j(z)}{\partial z_i} = 0 \mbox{ for $i>j$}.\] As a result, the Jacobian matrix \[J_G(z) = \left\{\frac{\partial G_j(z)}{\partial z_j}\right\}\] is a lower-triangular matrix with diagonal being \(p(z_j|z_{<j}).\) Thus, the determinant \[{\sf det}[J_G(z)] = \prod_{j=1}^d p(z_j|z_{<j}) = p_Z(z).\] Because \(W = (G_1(Z),\cdots, G_d(Z))^T\), the change of variable formula shows that the PDF of \(w\) is \[p_W(w) = p_Z(G^{-1}(w)) ({\sf det}[J_G(G^{-1}(w))])^{-1} = 1\] for \(w\in [0,1]^d\). Therefore, \[W\overset{d}{=} {\sf Uni}[0,1]^d\] is a uniform random variable in the unit cube.
With the above result, given any continuous random variable \(Z\in \mathbb{R}^d\), we can create a mapping \(G:R^d\rightarrow \mathbb{R}^d\) such that \(G(Z)\) follows a uniform distribution. Moreover, we can apply this procedure inversely that maps a uniform distribution over \([0,1]^d\) to any continuous random variable by iteratively applying equation \(\ref{eq::NF3}\): \[\begin{equation} Z_1 = F_1^{-1}(U_1) = Q_1(U),Z_2 = F_2^{-1}(U_2|Z_1)= Q_2(U),\cdots, Z_d = F_d^{-1}(U_d|Z_{<d})= Q_d(U) \label{eq::NF4} \end{equation}\] and \(U_1,\cdots, U_d\) are IID \({\sf Uni}[0,1]\). Namely, any random vector can be generated by \[\begin{equation} Z = Q(U), \label{eq::NF5} \end{equation}\] where \(U\sim {\sf Uni}[0,1]^d\) and \(Q = (Q_1,\cdots, Q_d)^T\).
Moreover, we can morph one continuous random variable to another. Suppose \(Y\) has a CDF \(F_Y\) and we define \(G_Y\) to be the function created by the conditional CDF of \(Y\) in equation \(\ref{eq::NF1}\). Let \(Q_Z\) be the function \(Q\) in equation \(\ref{eq::NF5}\) constructed from a CDF \(F_Z\). Then the function \[\phi_{Y\rightarrow Z} = Q_Z \circ G_Y\] will morph the random variable \(Y\) into the random variable \(Z\), i.e., \[Z \overset{d}{=}\phi_{Y\rightarrow Z}(Y) = Q_Z(G_Y(Y)).\]
The normalizing flow is based on the idea in the previous section–it starts with a well-known distribution such as uniform or multivariate Gaussian and then follow a sequence of parametric transformations in the above form to eventually match the data’s distribution. For simplicity, we start with a uniform random variable \(U\in \mathbb{R}^d\) and construction a sequence of mappings \[\phi_1(u; \theta_1),\cdots, \phi_K(u; \theta_K)\] such that \(\phi_i(\cdot; \theta_i):\mathbb{R}^d \rightarrow \mathbb{R}^d\) and \[\begin{equation} X = \phi_1(\cdot \theta_1)\circ \phi_2(\cdot; \theta_2) \circ\cdots\phi_K(\cdot; \theta_K) (U) = \phi_1(\phi_2(\cdots\phi_K(U; \theta_K); \theta_2) \theta_1) \label{eq::ONF1} \end{equation}\] matches the distribution of the observed data. Note that \(\theta_1,\cdots, \theta_K\) are the parameters for each transformation.
In theory, we only need one transformation that is the inverse CDF. However, it is infeasible to estimate the inverse CDF when the dimension \(d\) is large, so the normalizing flow attempts to approximate the inverse CDF via a sequence of parametric transformations.
Examples of flows. Here are some common flows:
Linear flow. \(\phi(y;\theta) = A y + b\), where \(A\in\mathbb{R}^{d\times d}, b\in\mathbb{R}^d\) are the parameters.
Planar flow. \(\phi(y;\theta) = y + \eta \cdot h(\omega^T y + b)\), where \(\eta, \omega \in\mathbb{R}^d\) and \(b\in \mathbb{R}\) are the parameters and \(h\) is a given function.
We will discuss more examples in Section 9.4.
For simplicity, we first consider the scenario that our flow only has one layer, i.e., \(K=1\).
Let \(X_1,\cdots, X_n\) be the observed data. Under equation \(\ref{eq::ONF1}\), we model the distribution of \(X_i\) as \[X\sim \phi_1(U;\theta_1).\] This is called forward modeling since we are moving from \(U\) toward \(X\) like moving forward.
Suppose \(U\) has a PDF \(p_U\). The resulting PDF of \(X\) is \[\begin{equation} p_X(x;\theta_1) = p_U(\phi^{-1}_1(x; \theta_1))\cdot \left|{\sf det}[J_{\phi^{-1}_1(;\theta_1)}(x)]\right| = L_1(\theta_1|x), \label{eq::ONF2} \end{equation}\] where \[\left[J_{\phi^{-1}_1(;\theta_1)}(z)\right]_{ij} = \frac{\partial \phi^{-1}_{1, j}(z;\theta_1)}{\partial z_i},\qquad i,j=1,\cdots, d\] is the Jacobian matrix. Note that the mapping \(\phi^{-1}_{1, j}(\cdot;\theta_1)\) is the \(j\)-th component of the mapping \(\phi^{-1}_{1}(\cdot;\theta_1): \mathbb{R}^d\rightarrow \mathbb{R}^d\).
Equation \(\ref{eq::ONF2}\) describes the likelihood function \(L(\theta_1|x)\). With this, we then estimate \(\theta_1\) by \[\widehat\theta_1 = {\sf argmax}_{\theta_1} \frac{1}{n}\sum_{i=1}^n \ell_1(\theta_1|X_i),\qquad \ell_1(\theta_1|X_i) = \log L_1(\theta_1|X_i).\]
Note that when we choose \(U\) to be from uniform distribution, \(p_U(u)=1\) so the log-likelihood reduces to \[\ell_1(\theta_1|x) = \log \left|{\sf det}[J_{\phi^{-1}_1(;\theta_1)}(x)]\right|.\]
Forward trick. Using property of inverse Jacobian, we can rewrite \[\begin{equation} {\sf det}[J_{\phi^{-1}_1(;\theta_1)}(x)] = \left\{{\sf det}[J_{\phi_1(;\theta_1)}(u)]\right\}^{-1},\qquad u = \phi^{-1}_1(x;\theta_1). \label{eq::ONF3} \end{equation}\] In the construction of forward model, \(\phi_1(u;\theta_1)\) generally has a simple form so we often have a closed-form for its derivative. Thus, computing the Jacobian \[\left\{{\sf det}[J_{\phi_1(;\theta)}(W_i)]\right\}^{-1},\qquad W_i = \phi^{-1}_1(X_i;\theta_1)\] can be done easily.
Forward trick: linear flow. As an example, consider a linear flow \(\phi_1(u;\theta) = A u + b\), where \(\theta = (A,b)\). The Jacobian \(J_{\phi_1(;\theta)}(u) = A\), so \[\left\{{\sf det}[J_{\phi_1(;\theta_1)}(u)]\right\}^{-1} = {\sf det}(A)^{-1}.\]
Forward trick: planar flow. Suppose we use the planar flow \[\phi_1(u;\theta) = u + \rho \cdot h(\omega^T u + b), \theta = (\rho,\omega, b)\] and \(h\) is a given function. Then the Jacobian matrix \[[J_{\phi_1(\cdot;\theta)}(u)]_{ij} = 1+ \rho_i \cdot h'(\omega^Tu +b) \omega_j\] or equivalently, \[J_{\phi_1(\cdot;\theta)}(u) = \mathbf{I}_d + h'(\omega^Tu +b) \cdot \rho \omega^T.\] Thus, \[{\sf det}[J_{\phi_1(\cdot;\theta)}(u) ] = 1 + h'(\omega^Tu +b) \cdot \rho^T \omega.\]
In equation [eq::ONF2], we see that the key to evaluate the likelihood function is the Jacobian, which involves the inverse of the model \(\phi_1^{-1}\). To simplify this, one may consider putting a parametric model directly on the inverse, i.e., \[\phi_1^{-1}(\cdot;\theta) = \Psi_1(\cdot;\lambda_1).\] In this case, we find the model \[\Psi_1(X;\lambda_1) \overset{d}{=} U\] that maps the data’s distribution back to the generating random variable \(U\). Since this goes backward from observed data to the generating distribution, it is called the backward modeling.
Under the backward modeling, the Jacobian \[J_{\phi^{-1}_1(;\theta_1)}(x) = J_{\Psi_1(;\lambda_1)}(x)\] is generally easy to compute.
The only downside of this idea is that the backward model \(\Psi:\mathbb{R}^d\rightarrow [0,1]^d\), so the choice of models is restricted.
An interesting fact is that if the PDF of \(X\) follows from a parametric model \(p(x;\rho)\) and and we construct our backward modeling via the corresponding conditional PDF \[p_j(x_j|x_{<j};\rho),\qquad F_j(x_j|x_{<j};\rho)= \int_{-\infty}^{x_j} p_j(t|x_{<j};\rho)dt,\qquad \Psi_{1, j}(x;\rho) = F_j(x_j|x_{<j};\rho).\] Then the backward model \(\Psi_1\) leads to a Jacobian \[J_{\Psi_1(;\lambda_1)}(x) = \prod_{j=1}^dp_j(x_j|x_{<j};\rho) = p(x;\rho),\] which is the usual likelihood model. Thus, the usual MLE approach can be viewed as a backward modeling normalizing flow.
When we have \(K\)-layers, the normalizing flow in equation \(\ref{eq::ONF1}\) can be expressed as an iterative sampling procedure:
\(Y_0 = U\sim p_U\).
For \(k=1,\cdots, K\), compute \(Y_{k} = \phi_k(Y_{k-1};\theta_k)\)
\(X = Y_K\).
The above sampling procedure also allows us to generate \(X\) if we have learned/estimated the parameters \(\phi_1,\cdots, \phi_K\).
In this case, we have a set of parameters \(\theta = (\theta_1,\cdots, \theta_K)\). To estimate the parameters, we use the likelihood approach. However, the usual likelihood function in equation [eq::ONF2] will be very complicated and not easy to work with, so people often use the forward trick in equation [eq::ONF3] so that we obtain the following elegant form \[\begin{equation} \begin{aligned} \ell(\theta|X) &= \log p_U(Y_0) - \sum_{k=1}^K \log \left|{\sf det}\left[J_{\phi_k(\cdot;\theta_k)}(Y_{k-1})\right]\right|\\ & = \log \Omega_0(\theta|X) - \sum_{k=1}^K \log \Omega_k (\theta_{k:K}|X),\\ Y_k & = \phi^{-1}_{k}(\cdot; \theta_k)\circ \phi^{-1}_{k+1}(\cdot; \theta_{k+1})\circ \phi^{-1}_{K}(X; \theta_K). \end{aligned} \label{eq::KNF1} \end{equation}\] The first term \(\Omega_0\) can be eliminated if we pick \(p_U\) to be the uniform distribution so it is generally not a big problem. Each \(\Omega_k (\theta_{k:K}|X)\) is the \(k\)-th Jacobian evaluated at \(Y_{k-1}\) so it only depends on the parameters ‘forward’ \(\theta_{k:K} = (\theta_k, \theta_{k+1},\cdots, \theta_K)\).
Maximizing \(\ell(\theta|X)\) can be done by numerical methods. Note that in our forward modeling the Jacobian \[J_{\phi_k(\cdot;\theta_k)}(y_{k-1}) = \frac{\partial \phi_k(y_{k-1};\theta_k)}{\partial y_{k-1}}\] often has a closed-form (see, e.g., linear and planar flows in Section 9.2.1). So the computational cost is mostly in the evaluation of \(Y_{k-1}\) since it involves all the forward models \(\phi_k,\cdots, \phi_K\).
The design of flow \(\phi(\cdot;\theta): \mathbb{R}^d \rightarrow \mathbb{R}\) plays a key role in Normalizing Flows. While we want to choose a family of flows that is rich, we also have to be cautious about the computational costs on its derivative. In particular, the determinant of Jacobian \[{\sf det} [J_{\phi(;\theta)}(u)], \qquad [J_{\phi(;\theta)}(u)]_{ij} = \frac{\partial \phi_j(u;\theta)}{\partial u_i}\] is crucial in the inference.
The autoregressive flows refer to the family where \(\phi_\ell(u;\theta)\), the \(\ell\)-th component of the function \(\phi(u;\theta)\), only depends on \(u_1,\cdots, u_\ell\). Namely, \[\begin{equation} \phi_\ell(u;\theta) = \kappa_\ell(u_{\leq\ell}; \theta) = \kappa_\ell(u_1,\cdots, u_\ell; \theta) \label{eq::AF} \end{equation}\] for some function \(\kappa_\ell\).
The power of autoregressive flow is the fact that for any \(i>j\), \[\frac{\partial \phi_j(u;\theta)}{\partial u_i} = 0.\] Thus, the Jacobian \(J_{\phi(;\theta)}(u)\) is upper-triangular, making the determinant a product form: \[{\sf det} [J_{\phi(;\theta)}(u)] = \prod_{j=1}^d \frac{\partial \phi_j(u;\theta)}{\partial u_j}.\] The log-likelihood function involves the logarithm of the determinant, which further leads to an elegant expression: \[\log {\sf det} [J_{\phi(;\theta)}(u)] = \sum_{j=1}^d \log \left(\frac{\partial \phi_j(u;\theta)}{\partial u_j}\right).\]
The transformer flow is a special class of autoregressive flow. In the autoregressive flow, we have \[\phi_\ell(u;\theta) =\kappa_\ell(u_1,\cdots, u_\ell; \theta).\] The transformer flow further requires \(\kappa_\ell\) to be \[\kappa_\ell(u_1,\cdots, u_\ell; \theta) = \zeta (u_\ell; \xi_\ell),\qquad \xi_\ell = \xi_\ell(u_1,\cdots, u_{\ell-1}; \theta),\] where \(\zeta\) is a given function called the transformer and \(\xi_\ell\) is a given function called the conditioner. Note that in the transformer function \(\zeta\), the output from the conditioner \(\xi_\ell\) behaves like the parameter of \(\zeta\).
A feature of the transformer is that the function \(\zeta\) is a univariate function and we may construct it by an integral \[\zeta(t; \xi) = \int_{-\infty}^t g(s;\xi) ds,\] where \(g\) is a non-negative function. In this choice, \[\frac{\partial \phi_j(u;\theta)}{\partial u_j} = \frac{\partial \zeta (u_j; \xi_j)}{\partial u_j} = \frac{\partial \int_{-\infty}^{u_j} g(s;\xi_j) ds}{\partial u_j} = g(u_j;\xi_j) = g(u_j;\xi_j(u_1,\cdots, u_{j-1}; \theta)) .\] Thus, the log-determinant is \[\log {\sf det} [J_{\phi(;\theta)}(u)] = \sum_{j=1}^d \log \left(\frac{\partial \phi_j(u;\theta_j)}{\partial u_j}\right) =\sum_{j=1}^d g(u_j;\xi_j(u_1,\cdots, u_{j-1}; \theta_j)),\] a very elegant form.