Technical Summary: Wasserstein Contraction of Coordinate Ascent Variational Inference

Problem Statement

The paper addresses the theoretical convergence guarantees for the Coordinate Ascent Variational Inference (CAVI) algorithm. CAVI is a fundamental method for approximating intractable target distributions $\pi$ (typically posterior distributions in Bayesian statistics) by minimizing the Kullback–Leibler (KL) divergence over a family of tractable, factorized distributions (e.g., $\mu \otimes \nu$). While CAVI is widely used due to its computational efficiency, rigorous convergence guarantees have historically been limited. Existing literature often relies on strong assumptions, such as global strong log-concavity of the target distribution, which excludes many practical models (e.g., mixture models, logistic regression) where the target is not strongly log-concave or where the space involves discrete or non-smooth components.

The authors aim to establish local exponential convergence of CAVI iterates $(\mu_k, \nu_k)$ to a fixed point in the Wasserstein distance, under conditions that are weaker and more verifiable than global strong log-concavity.

Methodology

The authors analyze the CAVI updates as a sequence of operators acting on probability measures. Let the target be $\pi = e^{-f} \in \mathcal{P}(Z \times B)$. The CAVI updates are defined via optimality maps:
$$\mu_k = \mu^\star[\nu_{k-1}], \quad \nu_k = \nu^\star[\mu_k]$$
where $\mu^\star[\nu] \propto \exp(-\mathbb{E}_\nu f(\cdot, \beta))$ and $\nu^\star[\mu] \propto \exp(-\mathbb{E}_\mu f(Z, \cdot))$.

The convergence analysis relies on two primary conditions defined on the space of probability measures equipped with the Wasserstein-$p$ distance ($W_p$):

1. Fisher-Smoothness (Cross-Smoothness): The optimality maps $\mu^\star[\cdot]$ and $\nu^\star[\cdot]$ are required to be "Fisher-smooth" around the fixed points. An operator $\gamma^\star[\cdot]$ is Fisher-smooth with constant $L$ if:
   $$I(\gamma^\star[\varrho] \| \gamma^\star[\varrho']) \leq L^2 W_p^2(\varrho, \varrho')$$
   where $I$ is the relative Fisher Information. This condition essentially bounds how much the gradient of the potential changes with respect to the other coordinate (cross-smoothness of the interaction term $\omega$ in $f = f_1 + f_2 + \omega$).

2. Transport-Information (TI) Inequality: The fixed point distributions $\mu^\star$ and $\nu^\star$ must satisfy a TI inequality with constants $\lambda_1, \lambda_2$:
   $$\lambda^2 W_p^2(\varrho, \gamma) \leq I(\varrho \| \gamma)$$
   This is a curvature condition on the fixed points. It is weaker than log-Sobolev inequalities and strong log-concavity, and it tensorizes, making it suitable for high-dimensional settings.

The authors also extend this framework to Data Augmentation settings (where $Z$ represents auxiliary variables). In this case, they show that convergence of the parameter marginal $\nu_k$ can be established solely based on the smoothness of the composed map $\nu^\star[\mu^\star[\cdot]]$ and the TI inequality of $\nu^\star$, without requiring curvature conditions on the auxiliary fixed point $\mu^\star$. This is crucial when $Z$ is discrete or non-smooth.

Key Contributions

1. General Convergence Theorem (Theorem 1)

The paper proves that if the cross-smoothness constants ($L_{1|2}, L_{2|1}$) and the TI constants ($\lambda_{1|2}, \lambda_{2|1}$) satisfy the condition $L_{1|2}L_{2|1} < \lambda_{1|2}\lambda_{2|1}$, then the CAVI iterates converge exponentially fast to the fixed point in Wasserstein distance with a rate of $\frac{L_{1|2}L_{2|1}}{\lambda_{1|2}\lambda_{2|1}}$.

• Generality: The result holds for general smooth Riemannian manifolds and allows for non-smooth spaces (via the data augmentation extension).
• Local vs. Global: The theorem provides local convergence guarantees. Global convergence is possible if the conditions hold globally, but the authors emphasize that local guarantees are sufficient for many practical scenarios where multiple fixed points exist.
• Sharpness: For Gaussian targets, the derived rate coincides with the spectral norm of the correlation matrix, demonstrating the sharpness of the bound.

2. Data Augmentation Framework (Theorem 4)

A significant contribution is the relaxation of assumptions for data augmentation problems. The authors show that if the interest lies only in the parameter $\beta$ (marginal $\nu$), convergence can be proven without requiring the auxiliary fixed point $\mu^\star$ to satisfy a TI inequality. This allows the application of CAVI convergence theory to models with discrete latent variables (e.g., Gaussian Mixture Models) where the auxiliary space is non-smooth.

3. Application to Specific Bayesian Models

The authors apply their theoretical framework to three specific models, deriving explicit convergence rates and conditions:

• Unbalanced Gaussian Mixture Models (GMM):

  • The convergence is governed by a posterior phase transition.
  • The rate depends on the separation of posterior modes. If the modes are well-separated, the algorithm converges locally and exponentially fast.
  • In the balanced case ($p=0.5$), global convergence requires a specific signal-to-noise ratio condition ($r < 1$). If this fails, the posterior becomes multimodal, and CAVI converges to different modes depending on initialization.
  • The results highlight that unbalanced mixtures (asymmetric $p$) facilitate faster convergence by reducing label ambiguity.

• Bayesian Probit Regression:

  • The authors analyze the CAVI updates for Probit regression with data augmentation (Albert and Chib, 1993).
  • They prove global exponential convergence with a rate determined by the Bayesian fraction of missing information: $\lambda_{\max}((Q_0 + X^\top X)^{-1}X^\top X)$.
  • They show that in high-dimensional settings ($n, p \to \infty$), the rate remains bounded away from 1 under appropriate prior scaling (e.g., $g$-priors or direct scaling), ensuring stability even as dimensionality grows.

• Bayesian Logistic Regression (Jaakkola-Jordan Algorithm):

  • The paper analyzes the algorithm using Pólya-Gamma augmentations.
  • They derive a local convergence rate $r_\varepsilon$ that depends on the derivative of the logistic link function and the curvature of the posterior.
  • Crucially, they establish that the convergence rate for Logistic regression is strictly better (smaller) than for Probit regression under comparable conditions, due to the properties of the Pólya-Gamma augmentation and the specific form of the interaction terms. This is the first available convergence rate result for the Jaakkola-Jordan algorithm.

Significance and Claims

The paper claims to provide a sharp, direct, and short proof of CAVI convergence that relies on functional inequalities rather than global strong log-concavity.

• Relaxation of Assumptions: By utilizing Transport-Information inequalities and Fisher-smoothness, the authors bypass the need for the target distribution $\pi$ to be strongly log-concave. This is a significant relaxation, as standard models like mixture models and logistic regression do not satisfy strong log-concavity.
• Handling Non-Smooth Spaces: The data augmentation extension allows for convergence analysis in settings with discrete latent variables, a common feature in mixture models that was previously difficult to analyze under smooth manifold assumptions.
• Accuracy-Convergence Duality: The paper notes that the conditions ensuring fast convergence (small interaction terms relative to curvature) also imply that the variational approximation is of high quality (low KL divergence). This establishes a duality between the speed of the algorithm and the accuracy of the approximation.
• Practical Relevance: The derived rates for high-dimensional regression (Probit and Logistic) demonstrate that CAVI remains stable and efficient even as the dimension $p$ and sample size $n$ grow, provided the prior is scaled correctly.

The authors explicitly state that their current analysis is restricted to the two-block ($m=2$) setting and that extensions to multi-block dynamics are left for future work. They also note that while their results are local, they align with and explain the empirical behavior of CAVI in complex, non-convex landscapes.