Stochastic Variational Inference (SVI) for Large Data

Coordinate Ascent Variational Inference (CAVI) is a technique for optimizing the Evidence Lower Bound (ELBO) in probabilistic models. A limitation of CAVI, however, is its requirement to iterate through the entire dataset to update each model parameter. For modern datasets, which can contain millions or even billions of data points, this full-dataset pass becomes a significant computational bottleneck. Imagine needing to process terabytes of data just to make a single adjustment to your model's parameters; the process would be prohibitively slow.

Stochastic Variational Inference (SVI) addresses this scalability challenge by borrowing ideas from stochastic optimization, particularly Stochastic Gradient Descent (SGD). Instead of calculating the true gradient of the ELBO using the full dataset, SVI uses noisy but computationally cheap gradients estimated from small, randomly selected subsets of the data, often called minibatches.

From Batch Gradients to Stochastic Gradients

Recall the ELBO: $\mathcal{L}(q) = \mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] - \mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$ Assuming the data points $\mathbf{x} = \{\mathbf{x}_1, ..., \mathbf{x}_N\}$ are conditionally independent given the latent variables $\mathbf{z}$, the joint probability often decomposes: $\log p(\mathbf{x}, \mathbf{z}) = \log p(\mathbf{z}) + \sum_{n=1}^N \log p(\mathbf{x}_n | \mathbf{z}_n)$. Let's simplify and consider models where $\mathbf{z}$ includes global parameters $\beta$ shared across all data points, and local latent variables $\mathbf{z}_n$ specific to each data point $\mathbf{x}_n$. The variational distribution $q$ also factors accordingly, often $q(\beta, \mathbf{z}_{1:N}) = q(\beta; \lambda) \prod_{n=1}^N q(\mathbf{z}_n; \phi_n)$, where $\lambda$ and $\phi_n$ are the variational parameters for global and local variables, respectively.

The ELBO can then be written involving a sum over data points: $\mathcal{L}(\lambda, \phi_{1:N}) = \mathbb{E}_q [\log p(\beta)] + \sum_{n=1}^N \mathbb{E}_q [\log p(\mathbf{x}_n | \mathbf{z}_n, \beta)] - \mathbb{E}_q [\log q(\beta; \lambda)] - \sum_{n=1}^N \mathbb{E}_q [\log q(\mathbf{z}_n; \phi_n)]$ CAVI requires optimizing this full objective. SVI, however, focuses on the global parameters $\lambda$. To update $\lambda$, we need the gradient $\nabla_{\lambda} \mathcal{L}$. Notice that the summation term $\sum_{n=1}^N \mathbb{E}_q [\log p(\mathbf{x}_n | \mathbf{z}_n, \beta)]$ prevents easy calculation without touching all data.

SVI's core idea is to approximate the full gradient using just a minibatch. If we sample a data point $\mathbf{x}_n$ uniformly at random, we can construct an unbiased stochastic estimate of the gradient related to the summation term. By sampling a minibatch $M$ of indices from $\{1, ..., N\}$, we can form a noisy gradient estimate: $\nabla_{\lambda} \hat{\mathcal{L}}_M \approx \nabla_{\lambda} \mathbb{E}_q [\log p(\beta)] + \frac{N}{|M|} \sum_{n \in M} \nabla_{\lambda} \mathbb{E}_q [\log p(\mathbf{x}_n | \mathbf{z}_n, \beta)] - \nabla_{\lambda} \mathbb{E}_q [\log q(\beta; \lambda)]$ The scaling factor $N/|M|$ corrects for the fact that we are summing over a smaller set $M$ instead of the full dataset $N$, ensuring the stochastic gradient remains an unbiased estimator of the true gradient (up to dependencies involving expectations over $q$).

The SVI Update Algorithm

SVI proceeds iteratively, updating global and local variational parameters using these stochastic gradients. At each step $t$:

1. Sample Minibatch: Randomly select a small subset (minibatch) $M_t$ of data points from the full dataset.
2. Update Local Parameters: For each data point $\mathbf{x}_n$ in the minibatch $M_t$, update its corresponding local variational parameters $\phi_n$ to optimize its contribution to the ELBO, holding the current global parameters $\lambda^{(t)}$ fixed. This might involve one or more CAVI-like updates specific to that data point. $\phi_n^{(t+1)} \leftarrow \arg \max_{\phi_n} \mathbb{E}_q [\log p(\mathbf{x}_n | \mathbf{z}_n, \beta)] - \mathbb{E}_q [\log q(\mathbf{z}_n; \phi_n)]$ (where expectations are taken holding $\lambda^{(t)}$ and other $\phi_{m \neq n}$ fixed).
3. Compute Stochastic Gradient: Calculate the noisy gradient $\nabla_{\lambda} \hat{\mathcal{L}}_{M_t}$ of the ELBO with respect to the global parameters $\lambda$, using the updated local parameters $\phi_n^{(t+1)}$ for $n \in M_t$.
4. Update Global Parameters: Update the global variational parameters $\lambda$ using a gradient ascent step with a learning rate $\rho_t$: $\lambda^{(t+1)} \leftarrow \lambda^{(t)} + \rho_t \nabla_{\lambda} \hat{\mathcal{L}}_{M_t}$

This process repeats, cycling through minibatches of data. The learning rate $\rho_t$ typically decreases over iterations to ensure convergence.

A view of an SVI update step. A small minibatch is sampled from the large dataset. This minibatch, along with the current global parameters $\lambda$, is used to calculate a noisy gradient, which then updates $\lambda$.

Learning Rates and Convergence

The choice of the learning rate schedule $\rho_t$ is important for SVI's performance. For the stochastic updates to converge properly, the learning rates must satisfy the Robbins-Monro conditions: $\sum_{t=1}^{\infty} \rho_t = \infty \quad \text{and} \quad \sum_{t=1}^{\infty} \rho_t^2 < \infty$ A common choice is a polynomial decay schedule: $\rho_t = (\tau_0 + t)^{-\kappa}$ where $\kappa \in (0.5, 1]$ controls the decay rate, and $\tau_0 \geq 0$ down-weights early iterations. Tuning $\kappa$ and $\tau_0$, along with the minibatch size $|M|$, often requires experimentation. Too large a learning rate can lead to instability, while too small a rate can result in slow convergence.

Natural Gradients for Faster Convergence

Standard gradient ascent updates the parameters $\lambda$ in the direction of steepest ascent in the Euclidean geometry of the parameter space. However, variational parameters often define probability distributions, and the Euclidean geometry might not be the most appropriate. The space of distributions has its own geometry, captured by the Fisher information matrix $F(\lambda)$.

Natural gradients modify the update direction by pre-multiplying the standard gradient with the inverse of the Fisher information matrix: $\tilde{\nabla}_{\lambda} \mathcal{L} = F(\lambda)^{-1} \nabla_{\lambda} \mathcal{L}$ The SVI update rule becomes: $\lambda^{(t+1)} \leftarrow \lambda^{(t)} + \rho_t \tilde{\nabla}_{\lambda} \hat{\mathcal{L}}_{M_t}$ For variational distributions belonging to the exponential family (a common choice), calculating the natural gradient can sometimes be simpler than the standard gradient and often leads to significantly faster convergence because it accounts for the information geometry of the variational parameter space.

Advantages of SVI

Advantages:

• Scalability: SVI's primary advantage is its ability to handle massive datasets that do not fit into memory, as it only processes small minibatches at each step.
• Speed: It often converges much faster in terms of wall-clock time compared to batch VI (CAVI) or MCMC methods, especially early in the optimization process.
• Online Learning: SVI can naturally incorporate new data points as they arrive without needing to retrain on the entire dataset.

Considerations:

• Tuning: Requires careful tuning of the learning rate schedule and minibatch size.
• Noisy Gradients: The stochastic nature of the gradients introduces noise, which can slow down convergence in later stages compared to batch methods or lead to oscillations around the optimum.
• Approximation Quality: Like all VI methods, SVI finds an approximation within the chosen variational family, which might not perfectly capture the true posterior. The quality depends on the flexibility of the family $q$.

SVI provides a powerful tool for applying Bayesian inference to large-scale problems where traditional methods become computationally infeasible. It forms the backbone of many modern probabilistic modeling applications, particularly in areas like topic modeling (e.g., Latent Dirichlet Allocation on large text corpora) and is a precursor to techniques used in Bayesian deep learning. While MCMC methods might offer asymptotically exact samples, and CAVI provides deterministic updates, SVI strikes a practical balance, enabling approximate Bayesian inference at scale by leveraging the power of stochastic optimization.