Fitting Generative Models

How to Fit a Generative Model?

• The goal of a generative model \(\mathcal{M}\) is to generate new data points that resemble the training data.
• We achieve this by minimizing some measure of distance/divergence between the true data distribution \(p_{\text{data}}(\mathbf{y})\) and the model distribution (the marginal distribution of the response) \(p(\mathbf{y} \mid \mathcal{M})\).
• We only have access to samples from the true data distribution, i.e., the training data \(\mathcal{D} = \{\mathbf{y}^{(i)}\}_{i=1}^N\) but not the distribution itself.

How to align the two distributions or their samples?

Kullback-Leibler Divergence

• The Kullback-Leibler (KL) divergence is a measure of how one probability distribution diverges from a second, it is not symmetric.
• For two distributions \(P\) and \(Q\) defined on the same probability space, the KL divergence from \(P\) to \(Q\) is defined as: \[ D_{KL}(P \| Q) = \int p(x) \log\left(\frac{p(x)}{q(x)}\right) dx \]
• One valid way to align the model and data distributions is to minimize the KL divergence from the data distribution to the model distribution. \[ D_{KL}(p_{\text{data}}(\mathbf{y}) \| p(\mathbf{y} \mid \mathcal{M})) \]

Kullback-Leibler Divergence

Denote the best model that minimizes the KL divergence by \(\mathcal{M}^*\).

\[ \begin{aligned} \mathcal{M}^* &= \arg\min_{\mathcal{M}} D_{KL}(p_{\text{data}}(\mathbf{y}) \| p(\mathbf{y} \mid \mathcal{M})) \\ &= \arg\min_{\mathcal{M}} \int p_{\text{data}}(\mathbf{y}) \log\left(\frac{p_{\text{data}}(\mathbf{y})}{p(\mathbf{y} \mid \mathcal{M})}\right) d\mathbf{y} \\ &= \arg\min_{\mathcal{M}} \left[ \int p_{\text{data}}(\mathbf{y}) \log(p_{\text{data}}(\mathbf{y})) d\mathbf{y} - \int p_{\text{data}}(\mathbf{y}) \log(p(\mathbf{y} \mid \mathcal{M})) d\mathbf{y} \right] \\ &= \arg\max_{\mathcal{M}} \int p_{\text{data}}(\mathbf{y}) \log(p(\mathbf{y} \mid \mathcal{M})) d\mathbf{y} \\ & \approx \arg\max_{\mathcal{M}} \frac{1}{N} \sum_{i=1}^N \log(p(\mathbf{y}^{(i)} \mid \mathcal{M})) \end{aligned} \]

• The first term is independent of the model \(\mathcal{M}\) and can be ignored during optimization.
• \(N\) is the number of training samples.

Kullback-Leibler Divergence

\[ \begin{aligned} \mathcal{M}^* &= \arg\min_{\mathcal{M}} D_{KL}(p_{\text{data}}(\mathbf{y}) \| p(\mathbf{y} \mid \mathcal{M})) \\ & \approx \arg\max_{\mathcal{M}} \frac{1}{N} \sum_{i=1}^N \log(p(\mathbf{y}^{(i)} \mid \mathcal{M})) \end{aligned} \]

• \(\log(p(\mathbf{y}^{(i)} \mid \mathcal{M}))\) is the contribution of the \(i\)-th data point to the log marginal likelihood of the model \(\mathcal{M}\).

Other Distances and Divergences

Other common distances/divergences used to align the two distributions are:

1. Jensen–Shannon divergence, used in generative adversarial networks (GANs): \[ \begin{aligned} D_{JS}(p_{\text{data}}(\mathbf{y}) \| p(\mathbf{y} \mid \mathcal{M})) &= \tfrac{1}{2} D_{KL}\!\left(p_{\text{data}}(\mathbf{y}) \| m\right) + \tfrac{1}{2} D_{KL}\!\left(p(\mathbf{y} \mid \mathcal{M}) \| m\right) \\ m &= \tfrac{1}{2}\big(p_{\text{data}}(\mathbf{y}) + p(\mathbf{y} \mid \mathcal{M})\big) \end{aligned} \]

2. Wasserstein (Earth Mover’s) distance, used in Wasserstein GANs: \[ W\!\left(p_{\text{data}}(\mathbf{y}),\, p(\mathbf{y} \mid \mathcal{M})\right) = \inf_{\gamma \in \Pi \left(p_{\text{data}},\, p(\cdot \mid \mathcal{M})\right)} \mathbb{E}_{(\mathbf{y}, \tilde{\mathbf{y}})\sim \gamma} \big[ \| \mathbf{y} - \tilde{\mathbf{y}} \| \big] \] where \(\Pi\!\left(p_{\text{data}}, p(\cdot \mid \mathcal{M})\right)\) is the set of all couplings (joint distributions) with the specified marginals.

Other Distances and Divergences

Other common distances/divergences used to align the two distributions are:

3. Fisher divergence, used in energy-based and (Stein) score-based generative models: \[ D_F\!\left(p_{\text{data}}(\mathbf{y}) \,\|\, p(\mathbf{y} \mid \mathcal{M})\right) = \int p_{\text{data}}(\mathbf{y})\, \left\| \nabla_{\mathbf{y}} \log p_{\text{data}}(\mathbf{y}) + \nabla_{\mathbf{y}} \log p(\mathbf{y} \mid \mathcal{M}) \right\|^2 d\mathbf{y} \]

The gradient \(\nabla_{\mathbf{y}} \log p(\mathbf{y} \mid \mathcal{M})\) is the model’s Stein score. \(\nabla_{\mathbf{y}} \log p_{\text{data}}(\mathbf{y})\) is the Stein score of the unknown data distribution.

Which Divergence Goes With Which Model?

The choice of divergence to minimize during training is closely tied to which operations the model supports tractably:

• KL divergence (forward) \(\to\) maximum likelihood. Requires evaluating \(\log p(\mathbf{y} \mid \mathcal{M})\). Natural for normalizing flows, PPCA, classical distributions; via ELBO for VAEs and DDPMs.
• Jensen–Shannon divergence \(\to\) GANs. Avoids requiring \(\log p(\mathbf{y})\) entirely — uses an auxiliary discriminator instead.
• Wasserstein distance \(\to\) Wasserstein GANs. More stable than JS for distributions with disjoint supports.
• Fisher divergence \(\to\) score matching, energy-based models, score-based diffusion. Only needs the score of the model, not the normalized density.

Marginal Likelihood

The marginal likelihood, also known as the model evidence, is a key quantity in latent variable models. It is the probability of the observed data under a given model, integrating over all possible values of the latent variables:

\[ p(\mathbf{y} \mid \mathcal{M}) = \int p(\mathbf{y} \mid \boldsymbol{z}, \mathcal{M}) \space p(\boldsymbol{z} \mid \mathcal{M}) \space d\boldsymbol{z} \]

• \(p(\mathbf{y} \mid \mathcal{M})\): Marginal likelihood of the model \(\mathcal{M}\) (model evidence).
• \(p(\mathbf{y} \mid \boldsymbol{z}, \mathcal{M})\): Conditional likelihood of the data given specific values of the latent variable \(\boldsymbol{z}\).
• \(p(\boldsymbol{z} \mid \mathcal{M})\): Prior distribution over the latent variables.

Pharmacometrics Notation

• In pharmacometrics, the latent variable \(\boldsymbol{z}\) is denoted by \(\boldsymbol{\eta}\), representing individual-specific random effects.
• The observed data \(\mathbf{y}\) corresponds to the measurements taken from individuals.
• The model \(\mathcal{M}\) encompasses the structural model and population parameters: \(\theta\), \(\Omega\), and \(\sigma\).

Balancing Individual Fit and Population Simulation

• Say we have panel data from \(N\) individuals, \(\mathcal{D} = \{\mathbf{y}_i\}_{i=1}^N\).
• We now consider another view of the marginal likelihood that shows how it balances fitting the individual time-series data while ensuring simulation accuracy at the population level.
• Denote the latent variable (individual random effects) by \(\boldsymbol{\eta}_i\) for individual \(i\) and the population parameters by \(\boldsymbol{\theta}\).
• We start with Bayes’ rule:

\[ \begin{aligned} p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta}) &= \frac{p(\mathbf{y} \mid \boldsymbol{\eta}, \boldsymbol{\theta}) \, p(\boldsymbol{\eta} \mid \boldsymbol{\theta})}{p(\mathbf{y} \mid \boldsymbol{\theta})} \\ p(\mathbf{y} \mid \boldsymbol{\theta}) &= \frac{p(\mathbf{y} \mid \boldsymbol{\eta}, \boldsymbol{\theta}) \, p(\boldsymbol{\eta} \mid \boldsymbol{\theta})}{p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta})} \end{aligned} \]

Balancing Individual Fit and Population Simulation

We then take the log of both sides:

\[ \log p(\mathbf{y} \mid \boldsymbol{\theta}) = \log p(\mathbf{y} \mid \boldsymbol{\eta}, \boldsymbol{\theta}) + \log p(\boldsymbol{\eta} \mid \boldsymbol{\theta}) - \log p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta}) \]

Then we take the expectation with respect to the posterior \(p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta})\):

\[ \log p(\mathbf{y} \mid \boldsymbol{\theta}) = \underbrace{\mathbb{E}_{p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta})} \left[\log p(\mathbf{y} \mid \boldsymbol{\eta}, \boldsymbol{\theta})\right]}_{\text{fit quality under posterior}} + \underbrace{\mathbb{E}_{p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta})} \left[\log p(\boldsymbol{\eta} \mid \boldsymbol{\theta}) - \log p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta})\right]}_{-\mathrm{KL}\left(p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta}) \| p(\boldsymbol{\eta} \mid \boldsymbol{\theta})\right)} \]

The first term is a measure of how well the model fits the data, averaged over the posterior distribution of the random effects. It encourages good individual fit. We denote it by \(Q\).

The second term is the negative KL divergence between the posterior distribution of the random effects \(p(\boldsymbol{\eta} \mid \mathbf{y}, \boldsymbol{\theta})\) and the prior distribution \(p(\boldsymbol{\eta} \mid \boldsymbol{\theta})\).

Balancing Individual Fit and Population Simulation

• For a population of \(N\) individuals, we can write the same quantity as: \[ \log p(\mathbf{y} \mid \boldsymbol{\theta}) = \sum_{i=1}^N \left( Q_i - \mathrm{KL}\!\left(p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta}) \| p(\boldsymbol{\eta}_i \mid \boldsymbol{\theta})\right) \right) \] where

  • \(Q_i\) is the fit quality for individual \(i\), and
  • \(p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta})\) is the posterior distribution of the random effects for individual \(i\).

• In NLME, the prior distribution \(p(\boldsymbol{\eta}_i \mid \boldsymbol{\theta})\) is typically the same for all individuals, e.g. \(\mathcal{N}(\mathbf{0}, \boldsymbol{\Omega})\), where \(\boldsymbol{\Omega}\) is a component of the population parameters \(\boldsymbol{\theta}\). \[ \log p(\mathbf{y} \mid \boldsymbol{\theta}) = \sum_{i=1}^N Q_i \;-\; \sum_{i=1}^N \mathrm{KL}\!\left( p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta}) \,\|\, \mathcal{N}(\mathbf{0}, \boldsymbol{\Omega}) \right) \]

Balancing Individual Fit and Population Simulation

The KL divergence is convex with respect to its first argument, the posterior \(p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta})\). Therefore, the following inequality holds for the second term above:

\[ \frac{1}{N} \sum_{i=1}^N \mathrm{KL}\!\left(p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta}) \,\|\, \mathcal{N}(\mathbf{0}, \boldsymbol{\Omega})\right) \;\ge\; \mathrm{KL}\!\left( \frac{1}{N} \sum_{i=1}^N p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta}) \,\Big\|\, \mathcal{N}(\mathbf{0}, \boldsymbol{\Omega}) \right) \]

where \(\frac{1}{N} \sum_{i=1}^N p(\boldsymbol{\eta}_i \mid \mathbf{y}_i, \boldsymbol{\theta})\) is the mixture distribution of the individual posteriors.

Maximizing the marginal likelihood therefore balances between:

1. Maximizing the ability of the individual posteriors to fit the individual data by maximizing the sum of the individual fit quality (\(Q_i\)) terms, and
2. Minimizing the sum of the individual KL divergences which indirectly aligns the mixture of the individual posteriors with the prior distribution, \(\mathcal{N}(\mathbf{0}, \boldsymbol{\Omega})\), ensuring a calibrated model and realistic simulations.

Balancing Individual Fit and Population Simulation

When simulating from an NLME model with population parameters \(\boldsymbol{\theta}\), we:

1. Sample \(\boldsymbol{\eta}\) from the prior distribution \(p(\boldsymbol{\eta} \mid \boldsymbol{\theta})\), e.g. \(\mathcal{N}(\mathbf{0}, \boldsymbol{\Omega})\),
2. Then sample \(\mathbf{y}\) from the conditional distribution \(p(\mathbf{y} \mid \boldsymbol{\eta}, \boldsymbol{\theta})\).

If the prior is close to the mixture of the individual posterior distributions, and the individual posteriors themselves fit the individual data well, this ensures that the simulations are realistic and consistent with the observed data.

This is especially useful for pharmacometrics, where it is common to simulate counter-factual scenarios for decision-making, such as predicting the effect of a new dosing regimen.

Predicting the Next Observation

Autoregressive Factorization

Let an individual have \(o\) observations at times \(\mathbf{t} = (t_1,\ldots,t_o)\) with responses \(\mathbf{y} = (y_1,\ldots,y_o)\).

The marginal likelihood factors autoregressively as

\[ \begin{aligned} p(\mathbf{y} \mid \mathbf{t}, \theta) &= \prod_{j=1}^o p(y_j \mid y_{1:j-1}, \mathbf{t}, \theta) \\ &= \prod_{j=1}^o p(y_j \mid y_{1:j-1}, t_{1:j}, \theta) \\ \log p(\mathbf{y} \mid \mathbf{t}, \theta) &= \sum_{j=1}^o \log p(y_j \mid y_{1:j-1}, t_{1:j}, \theta) \end{aligned} \]

Here \(y_{1:j}\) denotes the first j observations; \(t_{1:j}\) their time points.

Predicting the Next Observation

First Observation (j = 1)

\[ \begin{aligned} p(y_1 \mid t_1, \theta) &= \int p(y_1 \mid t_1, \boldsymbol{\eta}, \theta)\, p(\boldsymbol{\eta} \mid \theta)\, d\boldsymbol{\eta} \\ &= \mathbb{E}_{p(\boldsymbol{\eta} \mid \theta)} \big[ p(y_1 \mid t_1, \boldsymbol{\eta}, \theta) \big] \end{aligned} \]

Average predictive probability of the first observation under the prior distribution of random effects.

Subsequent Observations (j > 1)

\[ \begin{aligned} p(y_j \mid y_{1:j-1}, t_{1:j}, \theta) &= \int p(y_j \mid t_j, \boldsymbol{\eta}, \theta)\, p(\boldsymbol{\eta} \mid y_{1:j-1}, t_{1:j-1}, \theta)\, d\boldsymbol{\eta} \\ &= \mathbb{E}_{p(\boldsymbol{\eta} \mid y_{1:j-1}, t_{1:j-1}, \theta)} \big[ p(y_j \mid t_j, \boldsymbol{\eta}, \theta) \big] \end{aligned} \]

Uses the updated (posterior) distribution of random effects given past data.

Predicting the Next Observation

Predictive / Generalization View

Maximizing \(\log p(\mathbf{y} \mid \mathbf{t}, \theta)\) maximizes, on average over j:

\[ \log p(y_j \mid y_{1:j-1}, t_{1:j}, \theta) \]

This is the (one–step–ahead) predictive log probability.

Thus the marginal likelihood measures sequential generalization: how well the model predicts each next observation given the past while integrating over uncertainty in \(\boldsymbol{\eta}\) (prior for j=1, posterior thereafter).

Implication

• Encourages good individual fit (accurate one–step predictions).
• Discourages over-fitting by averaging over \(\boldsymbol{\eta}\) instead of conditioning on point estimates.

Summary

• Fitting a generative model means aligning the model and data distributions under some divergence; the choice of divergence depends on which operations the model supports.
• The marginal likelihood is the central quantity in latent-variable models, and maximizing it is asymptotically equivalent to minimizing the forward KL divergence.
• For panel data, the marginal likelihood balances individual-level fit and population-level calibration, and admits an autoregressive view as one-step-ahead prediction performance.

Fitting PPCA

In PPCA, the marginal likelihood is available in closed form, so fitting reduces to direct gradient-based or analytical MLE.

Recall the PPCA generative process with parameters \(\theta = \{\mathbf{W}, \boldsymbol{\mu}, \sigma^2\}\): \[ \begin{aligned} \mathbf{z} &\sim \mathcal{N}(\mathbf{0}, \mathbf{I}), \quad \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \sigma^2 \mathbf{I}) \\ \mathbf{y} &= \mathbf{Wz} + \boldsymbol{\mu} + \boldsymbol{\epsilon} \end{aligned} \]

• The marginal likelihood is available in closed form as a Gaussian: \[ p(\mathbf{y} \mid \theta) = \int p(\mathbf{y} \mid \mathbf{z}, \theta) \, p(\mathbf{z}) \, d\mathbf{z} = \mathcal{N}(\mathbf{y} \mid \boldsymbol{\mu}, \mathbf{C}), \quad \mathbf{C} = \mathbf{W}\mathbf{W}^T + \sigma^2 \mathbf{I} \]
• The marginal likelihood has a closed form so we can directly maximize it with respect to \(\theta\).
• PPCA admits a closed-form solution via eigen-decomposition of the data covariance matrix, similar to classical PCA, but we can also use gradient-based optimization to find the MLE.

Fitting NLME Models

In NLME models, the marginal likelihood is not available in closed form due to the nonlinearity of the structural model and the presence of individual random effects.

\[ p(\mathbf{y} \mid \boldsymbol{\theta}) = \int p(\mathbf{y} \mid \boldsymbol{\eta}, \boldsymbol{\theta}) \, p(\boldsymbol{\eta} \mid \boldsymbol{\theta}) \, d\boldsymbol{\eta} \]

We typically use one of the following 2 approaches to fit NLME models:

1. Laplace approximation: Approximate the integral over the random effects with a Gaussian centered at the mode of the integrand.
2. Expectation-Maximization (EM): Treat the random effects as latent variables and iteratively optimize the expected complete-data log-likelihood.

In pharmacometrics, it is common to further approximate the Laplace approximation by using the first-order (FO) or first-order conditional estimation (FOCE) methods.

Variational Autoencoders (VAEs)

Variational Autoencoders (VAEs)

The decoder and latent space of a VAE can be viewed as an NLME model with:

• A nonlinear structural model defined by a neural network (decoder).
• A Gaussian prior distribution over the latent variables.

\[ \begin{aligned} \boldsymbol{z} &\sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ \mathbf{y} &\sim p(\mathbf{y} \mid \boldsymbol{z}) = \mathcal{N}(\mathbf{\mu}(\boldsymbol{z}), \boldsymbol{\Sigma}(\boldsymbol{z})) \end{aligned} \]

• The encoder of the VAE can be viewed as an inference network that approximates the posterior distribution of the latent variables given the observed data.
• VAEs are fitted by maximizing the log marginal likelihood (aka evidence) indirectly by maximizing a lower bound called the evidence lower bound (ELBO).
• The ELBO is derived using variational inference and the estimation algorithm is known as variational expectation-maximization (VEM).

Variational Inference

• Variational inference is a method for approximating complex posterior distributions in probabilistic models (\(p(\boldsymbol{z} \mid \mathbf{y})\)) particularly in models with latent variables.
• The key idea is to introduce a family of simpler distributions \(q(\boldsymbol{z})\), called the variational family, to approximate the true posterior distribution.
• The variational family is often chosen to be a parametric distribution, such as a Gaussian distribution with mean and covariance that are neural network functions of the observed data \(\mathbf{y}\).
• The goal is to find the member of the variational family (by tuning its parameters) that is closest to the true posterior distribution, typically by minimizing the KL divergence between the two distributions. \[ D_{KL}(q(\boldsymbol{z}) \| p(\boldsymbol{z} \mid \mathbf{y})) = \int q(\boldsymbol{z}) \log\left(\frac{q(\boldsymbol{z})}{p(\boldsymbol{z} \mid \mathbf{y})}\right) d\boldsymbol{z} \]

Variational Inference

• In VAEs, we want to maximize the marginal likelihood of the observed data \(\mathbf{y}\): \[ p(\mathbf{y}) = \int p(\mathbf{y} \mid \boldsymbol{z}) \, p(\boldsymbol{z}) \, d\boldsymbol{z} \]

• However, this integral is often intractable.

• Instead of directly maximizing the marginal likelihood, we use variational inference to:

  1. Approximate the posterior distribution of the latent variable \(\boldsymbol{z}\) given the observed data \(\mathbf{y}\), and
  2. Construct a lower bound on the marginal likelihood that can be optimized.

Variational Inference

Assume the variational family/distribution for each subject \(i\) is defined as: \[ \begin{aligned} \boldsymbol{\xi}_i &\sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ \boldsymbol{z}_i &= T(\boldsymbol{\xi}_i; \boldsymbol{\phi}_i) \end{aligned} \] where \(T(.; \boldsymbol{\phi}_i)\) is an invertible parameterized transformation/function (e.g., an invertible neural network) that:

1. Has the same structure for all subjects, but
2. Has different parameters \(\boldsymbol{\phi}_i\) for each subject \(i\).

We denote the variational distribution and its probability density function by \(q(\boldsymbol{z}; \boldsymbol{\phi}_i)\).

Variational Inference

The probability density of the variational distribution can be computed using the change of variables formula: \[ \begin{aligned} q(\boldsymbol{z}; \boldsymbol{\phi}_i) &= p(\boldsymbol{\xi}) \left| \det\left( \frac{\partial T(\boldsymbol{\xi}; \boldsymbol{\phi}_i)}{\partial \boldsymbol{\xi}} \right) \right|^{-1} \\ \boldsymbol{\xi} &= T^{-1}(\boldsymbol{z}; \boldsymbol{\phi}_i) \end{aligned} \]

The Jacobian is analytical for simple transformations (e.g., affine transformations) and can be computed using automatic differentiation tools for more complex transformations (e.g., neural networks).

Amortized Inference

• In VAEs, we use amortized inference to share information across different subjects and reduce the number of parameters to be learned.
• Instead of learning separate variational parameters \(\boldsymbol{\phi}_i\) for each subject \(i\), we use a shared inference network (encoder) to map the observed data \(\mathbf{y}_i\) to the variational parameters \(\boldsymbol{\phi}_i\).
• The inference network is typically a neural network that takes the observed data \(\mathbf{y}_i\) as input and outputs the parameters of the variational distribution \(q(\boldsymbol{z}; \boldsymbol{\phi}_i)\).
• The inference network is trained jointly with the generative model (decoder) by maximizing the evidence lower bound (ELBO) on the marginal likelihood of the observed data.
• In NLME models, we typically learn separate variational parameters for each subject without using a shared inference network. This is more accurate but can require more parameters to be learned.

Evidence Lower Bound (ELBO)

• The evidence lower bound (ELBO) is a lower bound on the marginal likelihood of the observed data that can be obtained using variational inference.
• We start from the log marginal likelihood of the observed data \(\mathbf{y}\): \[ \log p(\mathbf{y}) = \log \int p(\mathbf{y} \mid \boldsymbol{z}) \, p(\boldsymbol{z}) \, d\boldsymbol{z} \]
• We introduce the variational distribution \(q(\boldsymbol{z}; \boldsymbol{\phi})\) (per subject but subscript \(i\) is omitted) by multiplying and dividing the integrand by \(q(\boldsymbol{z}; \boldsymbol{\phi})\): \[ \log p(\mathbf{y}) = \log \int p(\mathbf{y} \mid \boldsymbol{z}) \, p(\boldsymbol{z}) \, \frac{q(\boldsymbol{z}; \boldsymbol{\phi})}{q(\boldsymbol{z}; \boldsymbol{\phi})} \, d\boldsymbol{z} = \log \mathbb{E}_{q(\boldsymbol{z}; \boldsymbol{\phi})} \left[ \frac{p(\mathbf{y} \mid \boldsymbol{z}) \, p(\boldsymbol{z})}{q(\boldsymbol{z}; \boldsymbol{\phi})} \right] \]

Evidence Lower Bound (ELBO)

• We can then apply Jensen’s inequality to obtain a lower bound on the log marginal likelihood: \[ \log p(\mathbf{y}) \geq \mathbb{E}_{q(\boldsymbol{z}; \boldsymbol{\phi})} \left[ \log \left(\frac{p(\mathbf{y} \mid \boldsymbol{z}) \, p(\boldsymbol{z})}{q(\boldsymbol{z}; \boldsymbol{\phi})} \right) \right] \]
• Recall Bayes’ rule \[ p(\mathbf{y} \mid \boldsymbol{z})\, p(\boldsymbol{z}) = p(\boldsymbol{z} \mid \mathbf{y})\, p(\mathbf{y}) \]
• One can alternatively write the ELBO as: \[ \begin{aligned} \text{ELBO} = \int \log\!\left( \frac{ p(\boldsymbol{z} \mid \mathbf{y})\, p(\mathbf{y}) }{ q(\boldsymbol{z}; \boldsymbol{\phi}) } \right) q(\boldsymbol{z}; \boldsymbol{\phi})\, d\boldsymbol{z} &= \log p(\mathbf{y}) + \int \log\!\left( \frac{ p(\boldsymbol{z} \mid \mathbf{y}) }{ q(\boldsymbol{z}; \boldsymbol{\phi}) } \right) q(\boldsymbol{z}; \boldsymbol{\phi})\, d\boldsymbol{z} \\ &= \log p(\mathbf{y}) - \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}) \,\|\, p(\boldsymbol{z} \mid \mathbf{y}) \right) \end{aligned} \]

Evidence Lower Bound (ELBO)

\[ \begin{aligned} \text{ELBO} &= \log p(\mathbf{y}) - \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}) \,\|\, p(\boldsymbol{z} \mid \mathbf{y}) \right) \end{aligned} \]

Minorization Maximization (MM)

• The minorization maximization (MM) principle is an iterative optimization technique that can be used to maximize the marginal likelihood.

  • First, we construct a surrogate function (ELBO) that is a lower bound on (minorizes) the objective function (log marginal likelihood). The surrogate function is easier to compute and maximize than the original objective function.
  • We then iteratively maximize the lower bound to find its local maximizer and then update the lower bound at the new maximizer.

• Maximizing the lower bound indirectly maximizes the original objective function.

• When minimizing (instead of maximizing) an objective function, the MM principle is also known as majorization minimization. An upper bound (majorizer) is constructed instead of a lower bound (minorizer).

• When the surrogate function is an expectation, the MM principle is also known as expectation maximization (EM).

Minorization Step

• Let’s include the model parameters \(\boldsymbol{\theta}\) explicitly and rewrite the ELBO as: \[ \begin{aligned} \text{ELBO}(\boldsymbol{\phi}, \boldsymbol{\theta}) &= \log p(\mathbf{y} \mid \boldsymbol{\theta}) - \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}) \,\|\, p(\boldsymbol{z} \mid \mathbf{y}, \boldsymbol{\theta}) \right) \end{aligned} \]
• The minorization step finds the best variational parameters \(\boldsymbol{\phi}\) for fixed model parameters \(\boldsymbol{\theta}\). \[ \begin{aligned} \boldsymbol{\phi}^* &= \arg \max_{\boldsymbol{\phi}} \text{ELBO}(\boldsymbol{\phi}, \boldsymbol{\theta}) \\ &= \arg \max_{\boldsymbol{\phi}} \log p(\mathbf{y} \mid \boldsymbol{\theta}) - \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}) \,\|\, p(\boldsymbol{z} \mid \mathbf{y}, \boldsymbol{\theta}) \right) \\ &= \arg \min_{\boldsymbol{\phi}} \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}) \,\|\, p(\boldsymbol{z} \mid \mathbf{y}, \boldsymbol{\theta}) \right) \end{aligned} \]
• In practice, we don’t need the real posterior \(p(\boldsymbol{z} \mid \mathbf{y}, \boldsymbol{\theta})\) to perform the minorization step. We use another equivalent form of the ELBO which does not require the true posterior.
• However, conceptually, maximizing the ELBO with respect to \(\boldsymbol{\phi}\) is equivalent to minimizing the KL divergence to the true posterior.

Maximization Step

\[ \begin{aligned} \text{ELBO}(\boldsymbol{\phi}, \boldsymbol{\theta}) &= \log p(\mathbf{y} \mid \boldsymbol{\theta}) - \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}) \,\|\, p(\boldsymbol{z} \mid \mathbf{y}, \boldsymbol{\theta}) \right) \end{aligned} \]

• The maximization step finds the best model parameters \(\boldsymbol{\theta}\) for fixed variational parameters \(\boldsymbol{\phi}\). \[ \begin{aligned} \boldsymbol{\theta}^* &= \arg \max_{\boldsymbol{\theta}} \text{ELBO}(\boldsymbol{\phi}, \boldsymbol{\theta}) \end{aligned} \]

Joint Optimization

• Since both the minorization and maximization attempt to maximize the ELBO with respect to different parameters, we can combine them into a single optimization problem: \[ \begin{aligned} \boldsymbol{\phi}^*, \boldsymbol{\theta}^* &= \arg \max_{\boldsymbol{\phi}, \boldsymbol{\theta}} \text{ELBO}(\boldsymbol{\phi}, \boldsymbol{\theta}) \end{aligned} \]

Multiple Subjects

For multiple subjects \(i = 1, \ldots, N\), the ELBO decomposes as a sum of individual ELBOs: \[ \text{ELBO}(\boldsymbol{\phi}, \boldsymbol{\theta}) = \sum_{i=1}^N \text{ELBO}_i(\boldsymbol{\phi}_i, \boldsymbol{\theta}) \] where \[\text{ELBO}_i(\boldsymbol{\phi}_i, \boldsymbol{\theta}) = \mathbb{E}_{q(\boldsymbol{z}; \boldsymbol{\phi}_i)} \left[ \log p(\mathbf{y}_i \mid \boldsymbol{z}, \boldsymbol{\theta}) \right] - \mathrm{KL}\!\left( q(\boldsymbol{z}; \boldsymbol{\phi}_i) \,\|\, p(\boldsymbol{z} \mid \boldsymbol{\theta}) \right)\]

Amortized Inference Revisited

• When using amortized inference, the variational parameters \(\boldsymbol{\phi}_i\) for each subject \(i\) are outputs of a shared inference network (encoder) with parameters \(\boldsymbol{\psi}\).
• Denote this function by \(\boldsymbol{\phi}(\mathbf{y}_i; \boldsymbol{\psi})\), which takes the observed data \(\mathbf{y}_i\) as input and outputs the variational parameters \(\boldsymbol{\phi}_i\).
• The ELBO for a dataset with \(N\) subjects becomes: \[ \text{ELBO}(\boldsymbol{\psi}, \boldsymbol{\theta}) = \sum_{i=1}^N \text{ELBO}_i(\boldsymbol{\phi}(\mathbf{y}_i; \boldsymbol{\psi}), \boldsymbol{\theta}) \]
• In this case, we optimize the ELBO with respect to the shared inference network parameters \(\boldsymbol{\psi}\) and the model parameters \(\boldsymbol{\theta}\): \[ \boldsymbol{\psi}^*, \boldsymbol{\theta}^* = \arg \max_{\boldsymbol{\psi}, \boldsymbol{\theta}} \text{ELBO}(\boldsymbol{\psi}, \boldsymbol{\theta}) \]

Summary

Denoising Diffusion Probabilistic Models (DDPMs)

Denoising Diffusion Probabilistic Models (DDPMs)

Denoising Diffusion Probabilistic Models (DDPMs)

• DDPMs are analogical to a \(T\)-step VAE, for \(T > 1\).
• Instead of defining the prior and conditional likelihood, we define the posterior distribution as a noise adding step.
• The prior naturally tends to a standard normal distribution as \(T\) increases.
• The conditional likelihood is learned to undo the noise added in the posterior.
• DDPMs are trained by maximizing the ELBO, just like VAEs.
• As \(T \to \infty\), the optimal conditional likelihood can be shown to converge to an isotropic Gaussian form with a closed form expression for the variance, which simplifies training. \[ p(\mathbf{y}_{t-1} \mid \mathbf{y}_t) = \mathcal{N}(\mathbf{y}_{t-1} \mid \boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{y}_t, t), \sigma_t^2 \mathbf{I}) \]

Denoising Diffusion Probabilistic Models (DDPMs)

• \(\sigma_t^2\) has a closed form optimal solution as \(T \to \infty\).
• The only parameters to be learned are those of the mean function \(\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{y}_t, t)\), which is typically parameterized using a neural network.
• We can re-parameterize the mean function in terms of a noise function \(\boldsymbol{\epsilon}\): \[ \boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{y}_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{y}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{y}_t, t) \right) \] where \(\alpha_t\) and \(\bar{\alpha}_t\) are known functions of \(t\) that control the noise schedule. \[ \alpha_t = 1 - \beta_t, \quad \bar{\alpha}_t = \prod_{s=1}^t \alpha_s \]

Denoising Diffusion Probabilistic Models (DDPMs)

• Note that the noisy data \(\mathbf{y}_t\) at time step \(t\) can be expressed as a linear combination of the original data \(\mathbf{y}_0\) and the noise \(\boldsymbol{\epsilon}\): \[ \mathbf{y}_t = \sqrt{\bar{\alpha}_t} \mathbf{y}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon} \]
• With the re-parameterization of \(\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{y}_t, t)\) and the above expression for \(\mathbf{y}_t\), the ELBO objective reduces to a weighted sum: \[ \text{ELBO} = \sum_{t=1}^T w_t \, \mathbb{E}_{\mathbf{y}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{y}_t, t) \right\|^2 \right] + \text{const} \] where \[ w_t = \frac{(1 - \alpha_t)^2}{\sigma_t^2 (1 - \bar{\alpha}_t)} \]

Denoising Diffusion Probabilistic Models (DDPMs)

• More generally, one can define a distribution over the time step \(t\) and re-write the ELBO up to a constant as:

\[ \mathbb{E}_{t, \mathbf{y}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{y}_t, t) \right\|^2 \right] = \mathbb{E}_{t, \mathbf{y}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\sqrt{\bar{\alpha}_t} \mathbf{y}_0 + \sqrt{1 - \bar{\alpha}_t} \boldsymbol{\epsilon}, t) \right\|^2 \right] \]

• \(\epsilon_{\boldsymbol{\theta}}(\mathbf{y}_t, t)\) is a neural network that takes the noisy data \(\mathbf{y}_t\) and the time step \(t\) as input and outputs an estimate of the noise \(\boldsymbol{\epsilon}\).
• The ELBO objective can be optimized using stochastic gradient descent by sampling \(t\), \(\mathbf{y}_0\), and \(\boldsymbol{\epsilon}\) from their respective distributions.

Continuous-Time Diffusion Models