[Paper Review] Grassmann Stein Variational Gradient Descent

Xing Liu, Harrison Zhu+3 more|arXiv (Cornell University)|Feb 7, 2022

Markov Chains and Monte Carlo Methods4 citations

TL;DR

This paper proposes Grassmann Stein Variational Gradient Descent (GSVGD), a novel particle-based inference method that improves uncertainty quantification in high-dimensional Bayesian inference by adaptively projecting both data and score functions onto m-dimensional subspaces via a coupled ODE-SDE system on the Grassmann manifold. GSVGD outperforms SVGD and S-SVGD in estimating epistemic uncertainty, particularly in problems with intrinsic low-dimensional structure, while maintaining accurate mean estimation.

ABSTRACT

Stein variational gradient descent (SVGD) is a deterministic particle inference algorithm that provides an efficient alternative to Markov chain Monte Carlo. However, SVGD has been found to suffer from variance underestimation when the dimensionality of the target distribution is high. Recent developments have advocated projecting both the score function and the data onto real lines to sidestep this issue, although this can severely overestimate the epistemic (model) uncertainty. In this work, we propose Grassmann Stein variational gradient descent (GSVGD) as an alternative approach, which permits projections onto arbitrary dimensional subspaces. Compared with other variants of SVGD that rely on dimensionality reduction, GSVGD updates the projectors simultaneously for the score function and the data, and the optimal projectors are determined through a coupled Grassmann-valued diffusion process which explores favourable subspaces. Both our theoretical and experimental results suggest that GSVGD enjoys efficient state-space exploration in high-dimensional problems that have an intrinsic low-dimensional structure.

Motivation & Objective

• Address the variance underestimation issue in Stein Variational Gradient Descent (SVGD) under high-dimensional target distributions.
• Overcome the limitations of 1D projections in S-SVGD, which can overestimate epistemic uncertainty due to fixed score function projections.
• Develop a unified framework that jointly optimizes projectors for both data and score functions to enhance state-space exploration.
• Enable accurate uncertainty estimation in high-dimensional problems with intrinsic low-dimensional structure by using higher-dimensional subspaces.
• Introduce a Grassmann-valued stochastic diffusion process to balance exploration and exploitation in subspace selection.

Proposed method

• Propose a Grassmann Kernel Stein Discrepancy (GKSD) to measure discrepancy between proposal and target distributions in m-dimensional subspaces.
• Formulate a coupled ODE-SDE system where SVGD particles and Grassmann projectors evolve simultaneously, with the projectors guided by a diffusion process on the Grassmann manifold.
• Use a stochastic process on the Grassmann manifold to adaptively learn optimal m-dimensional subspaces that maximize discrepancy, enabling dynamic exploration of favorable directions.
• Simultaneously update projectors for both the data and the score function, avoiding the deterministic bias of prior methods that fix one set of projections.
• Allow for arbitrary m-dimensional projections (1 ≤ m ≤ d), capturing correlations between components and improving uncertainty quantification over 1D slices.
• Introduce a batched extension of GSVGD that uses multiple projectors to improve diversity and convergence in high-dimensional inference.

Experimental results

Research questions

• RQ1 Can higher-dimensional subspaces (m > 1) improve uncertainty quantification in high-dimensional Bayesian inference compared to 1D projections?
• RQ2 Does jointly optimizing projectors for both data and score functions reduce bias and improve estimation accuracy in SVGD?
• RQ3 How does the Grassmann-valued diffusion process balance exploration and exploitation in subspace selection?
• RQ4 Can GSVGD achieve better agreement with HMC reference posteriors in terms of mean and covariance estimation in high-dimensional problems?
• RQ5 What is the sensitivity of GSVGD performance to the choice of projection dimension m and hyperparameters of the diffusion process?

Key findings

• GSVGD achieves more accurate uncertainty estimates than SVGD and S-SVGD in a high-dimensional stochastic process inference problem, with credible intervals closely matching HMC reference.
• In the Covertype dataset experiment, GSVGD with m ≈ 10 achieves the most accurate posterior covariance estimation, outperforming S-SVGD, which severely overestimates uncertainty.
• GSVGD reduces variance underestimation in high-dimensional problems by leveraging m-dimensional subspaces that capture inter-component correlations.
• Energy distance between GSVGD results and HMC reference is consistently lower than SVGD and S-SVGD across different projection dimensions, indicating better approximation quality.
• The optimal projection dimension for GSVGD is non-trivial but lies in a moderate range (e.g., m ≈ 20 for the stochastic process problem), with performance stable across a broad range of m values.
• GSVGD demonstrates robustness to hyperparameter choices when tuned via heuristics, though sensitivity analysis suggests case-specific tuning may be needed for optimal performance.