EigenVI: score-based variational inference with orthogonal function expansions (2410.24054v1)

Abstract: We develop EigenVI, an eigenvalue-based approach for black-box variational inference (BBVI). EigenVI constructs its variational approximations from orthogonal function expansions. For distributions over $\mathbb{R}^D$, the lowest order term in these expansions provides a Gaussian variational approximation, while higher-order terms provide a systematic way to model non-Gaussianity. These approximations are flexible enough to model complex distributions (multimodal, asymmetric), but they are simple enough that one can calculate their low-order moments and draw samples from them. EigenVI can also model other types of random variables (e.g., nonnegative, bounded) by constructing variational approximations from different families of orthogonal functions. Within these families, EigenVI computes the variational approximation that best matches the score function of the target distribution by minimizing a stochastic estimate of the Fisher divergence. Notably, this optimization reduces to solving a minimum eigenvalue problem, so that EigenVI effectively sidesteps the iterative gradient-based optimizations that are required for many other BBVI algorithms. (Gradient-based methods can be sensitive to learning rates, termination criteria, and other tunable hyperparameters.) We use EigenVI to approximate a variety of target distributions, including a benchmark suite of Bayesian models from posteriordb. On these distributions, we find that EigenVI is more accurate than existing methods for Gaussian BBVI.

• The paper introduces EigenVI, a novel BBVI method that replaces traditional gradient descent with eigenvalue minimization using score matching.
• It systematically approximates non-Gaussian, multimodal distributions by employing orthogonal function expansions, including weighted Hermite polynomials.
• Experimental results demonstrate that EigenVI outperforms standard Gaussian methods, enhancing robustness and reducing hyperparameter tuning in Bayesian models.

EigenVI: Score-based Variational Inference with Orthogonal Function Expansions

The paper presents EigenVI, a new approach to black-box variational inference (BBVI) that utilizes eigenvalue-based score matching in conjunction with orthogonal function expansions. The approach is tailored for approximating target distributions that can exhibit substantial non-Gaussian behavior. EigenVI offers a systematic method to develop variational approximations through orthogonal function bases, providing a framework for modeling complex and diverse distributions beyond simple Gaussian behavior.

Background and Motivation

Variational inference (VI) has been pivotal for scalable Bayesian inference, offering a structured optimization approach to approximate posterior distributions with certain tractable families. Despite its efficacy, traditional gradient-based VI approaches face challenges with tuning hyperparameters such as learning rates and iteration thresholds, which can complicate the optimization process. EigenVI addresses these by using orthogonal expansions and Fisher divergence minimization, avoiding the typical iterative gradient descent.

Methodology

EigenVI leverages orthogonal basis functions to create a family of polynomial expansions that captures diverse distribution properties. The expansions allow for systematic modeling of non-Gaussian features by adding higher-order terms. The choice of orthogonal functions, including weighted Hermite polynomials, facilitates the handling of different types of support and distribution characteristics.

Rather than employing traditional stochastic gradient descent, EigenVI minimizes a score-based Fisher divergence through a minimum eigenvalue problem. This is a notable shift, as it mitigates tuning sensitivity and convergence issues typical in gradient-based methods. The stochastic estimate of the Fisher divergence is calculated via importance sampling, and the solution reduces to an eigenvalue problem, making it computationally efficient.

Results and Computational Experiments

EigenVI demonstrates superior performance over standard Gaussian BBVI methods across several examples, highlighting its capability to model multimodal distributions, asymmetries, and heavy tails. The experiments encompass synthetic distributions and benchmarks from the posteriordb suite, including complex Bayesian models. These highlight EigenVI’s ability to deliver more accurate posterior approximations compared to leading BBVI techniques, confirming its utility for diverse and challenging inference tasks.

The numerical results revealed EigenVI's robustness in approximating both skewed and heavy-tailed targets effectively, significantly outperforming traditional Gaussian methods. The handling of complex distributions, verified through careful comparisons with ground truth and other inference methods, underscores its practical impact.

Implications and Future Directions

EigenVI positions itself as a valuable alternative in the BBVI toolkit. Its independence from hyperparameter tuning for optimization advances practical usability, providing a fertile area for further exploration in variational inference. Future research may investigate its applicability to even higher-dimensional spaces and the development of adaptive schemes for importance sampling within the EigenVI framework.

Furthermore, there is potential to explore different orthogonal expansions and extend this approach to various types of models and data structures. Examining the theoretical bounds of the eigenvalue minimize under different conditions and decomposing larger matrices for efficient parallel computation opens future development paths.

In summary, EigenVI's novel use of orthogonal function expansions and score-based optimization via eigenvalue problems yields a compelling approach for approximating intricate distributions in variational inference, marking it as a promising advancement in the field.