Coordinate Ascent Variational Inference (CAVI)

Okay, let's explore Coordinate Ascent Variational Inference (CAVI). We've established that Variational Inference (VI) aims to find the best approximation $q(\mathbf{z})$ to the true posterior $p(\mathbf{z}|\mathbf{x})$ within a chosen family $\mathcal{Q}$, typically by maximizing the Evidence Lower Bound (ELBO). The common choice, the mean-field variational family, assumes independence between latent variables (or groups of variables) in the approximation: $q(\mathbf{z}) = \prod_{j=1}^M q_j(z_j)$ where $\mathbf{z} = \{z_1, \dots, z_M\}$ are the latent variables. The question now is: how do we find the optimal factors $q_j(z_j)$ that maximize the ELBO, $\mathcal{L}(q)$?

CAVI provides an iterative approach. It maximizes the ELBO with respect to each factor $q_j(z_j)$ one at a time, holding all other factors $q_{i \neq j}(z_i)$ fixed. This is analogous to coordinate ascent optimization applied in the functional space of distributions.

Deriving the CAVI Update Rule

Let's focus on optimizing the ELBO with respect to a single factor, say $q_j(z_j)$. Recall the ELBO: $\mathcal{L}(q) = \mathbb{E}_{q} [\log p(\mathbf{x}, \mathbf{z})] - \mathbb{E}_{q} [\log q(\mathbf{z})]$ Substituting the mean-field assumption $q(\mathbf{z}) = \prod_{i=1}^M q_i(z_i)$: $\mathcal{L}(q) = \int \dots \int \left( \prod_{i=1}^M q_i(z_i) \right) \left( \log p(\mathbf{x}, \mathbf{z}) - \sum_{i=1}^M \log q_i(z_i) \right) dz_1 \dots dz_M$ We want to find the $q_j(z_j)$ that maximizes this expression, holding $q_{i \neq j}$ fixed. Let's isolate the terms involving $q_j(z_j)$: $\mathcal{L}(q) = \int q_j(z_j) \left( \int \dots \int \left( \prod_{i \neq j} q_i(z_i) \right) \log p(\mathbf{x}, \mathbf{z}) \left( \prod_{i \neq j} dz_i \right) \right) dz_j - \int q_j(z_j) \log q_j(z_j) dz_j + \text{const}$ Here, "const" represents terms that do not depend on $q_j(z_j)$. The term inside the large parentheses in the first integral is the expectation of the log joint distribution with respect to all factors except $q_j$. Let's denote this as $\mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})]$. So we have: $\mathcal{L}(q) = \int q_j(z_j) \mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})] dz_j - \int q_j(z_j) \log q_j(z_j) dz_j + \text{const}$ Let $\log \tilde{p}_j(z_j) = \mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})]$. The expression becomes: $\mathcal{L}(q) = \int q_j(z_j) \log \tilde{p}_j(z_j) dz_j - \int q_j(z_j) \log q_j(z_j) dz_j + \text{const}$ $\mathcal{L}(q) = - \int q_j(z_j) \log \frac{q_j(z_j)}{\tilde{p}_j(z_j)} dz_j + \text{const}$ $\mathcal{L}(q) = -KL(q_j(z_j) || \tilde{p}_j(z_j)) + \text{const}$ Maximizing $\mathcal{L}(q)$ with respect to $q_j(z_j)$ is equivalent to minimizing the Kullback-Leibler (KL) divergence between $q_j(z_j)$ and $\tilde{p}_j(z_j)$. The KL divergence is minimized (becomes zero) when $q_j(z_j) = \tilde{p}_j(z_j)$. Therefore, the optimal solution $q_j^*(z_j)$ satisfies: $\log q_j^*(z_j) = \mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})] + \text{const}$ Exponentiating both sides gives: $q_j^*(z_j) \propto \exp \left( \mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})] \right)$ The constant term is determined by the normalization requirement $\int q_j^*(z_j) dz_j = 1$.

This is the core update rule for CAVI. It states that the optimal distribution for the $j$-th factor $q_j(z_j)$ is proportional to the exponentiated expected log of the joint probability $p(\mathbf{x}, \mathbf{z})$, where the expectation is taken with respect to the current distributions of all other factors $q_{i \neq j}(z_i)$.

The CAVI Algorithm

The CAVI algorithm proceeds as follows:

1. Initialization: Initialize the parameters of each variational factor $q_j(z_j)$ for $j=1, \dots, M$. This might involve initializing means, variances, or other relevant parameters depending on the distributional form chosen for each $q_j$.
2. Iteration: Repeat until convergence:
   • For each factor $j = 1, \dots, M$:
     • Compute the expectation $\mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})]$. This requires using the current parameters of the distributions $q_{i \neq j}(z_i)$.
     • Update the factor $q_j(z_j)$ using the rule: $q_j(z_j) \leftarrow \frac{\exp \left( \mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})] \right)}{\int \exp \left( \mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})] \right) dz_j}$ In practice, this often means updating the parameters of the distribution $q_j(z_j)$ based on the computed expectation.
3. Convergence Check: Monitor the ELBO or the parameters of the variational factors. Stop when the change between iterations falls below a predefined threshold.

The CAVI algorithm iteratively updates each factor $q_j(z_j)$ based on the expectation of the log joint probability, conditioned on the current state of the other factors, until convergence.

The Role of Conjugacy

The practicality of CAVI often hinges on the structure of the model $p(\mathbf{x}, \mathbf{z})$. If the complete conditional $p(z_j | \mathbf{x}, \mathbf{z}_{\neg j})$ (where $\mathbf{z}_{\neg j}$ denotes all variables except $z_j$) is in the same family as the prior for $z_j$, the model exhibits conditional conjugacy. In such cases, the CAVI update for $q_j^*(z_j)$ often results in a distribution of the same parametric form as the initial choice for $q_j(z_j)$.

For instance, if $q_j(z_j)$ is chosen to be Gaussian, and the model structure allows, the update $q_j^*(z_j)$ derived from the expectation $\mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})]$ might also be Gaussian. The update step then simplifies to calculating the new mean and variance parameters based on the expected values (moments) derived from the other factors $q_{i \neq j}$. This makes the implementation significantly easier, as we only need to track and update the parameters (e.g., means, variances) of the factors.

Advantages and Disadvantages of CAVI

Advantages:

• Guaranteed Convergence: The ELBO is guaranteed to increase (or stay the same) at each step of CAVI, ensuring convergence to a local optimum.
• Deterministic: Unlike MCMC, CAVI is a deterministic optimization algorithm, yielding the same result given the same initialization.
• Simplicity (with Conjugacy): When conditional conjugacy applies, the updates can often be derived analytically and involve updating parameters of standard distributions.

Disadvantages:

• Local Optima: CAVI optimizes the ELBO, which is generally non-convex. The algorithm may converge to a local optimum, potentially providing a poor approximation of the true posterior. Different initializations might lead to different results.
• Mean-Field Limitation: The accuracy is fundamentally limited by the expressiveness of the chosen variational family. The mean-field assumption (independence between factors) might be too restrictive for models with strong posterior dependencies.
• Analytical Calculations: Deriving the update rule requires calculating $\mathbb{E}_{i \neq j} [\log p(\mathbf{x}, \mathbf{z})]$. This can be complex or analytically intractable for some models.
• Scalability: Traditional CAVI requires iterating through the entire dataset to compute expectations for global variable updates (parameters governing all data points), which doesn't scale well to very large datasets.

CAVI provides a foundational understanding of variational optimization. While effective for moderately sized problems with suitable model structures, its limitations, particularly regarding scalability and the need for analytical derivations, motivate the development of more advanced techniques like Stochastic Variational Inference (SVI) and Black Box Variational Inference (BBVI), which we will explore next.