Source-linked AI summary
Projected Stein Variational Newton: A Fast and Scalable Bayesian Inference Method in High Dimensions
Peng Chen, Keyi Wu, Joshua Chen, Thomas O'Leary-Roseberry, Omar Ghattas
TL;DR
High-dimensional Bayesian inference can make posterior sampling prohibitive. pSVN uses an adaptively constructed Hessian-informed subspace for parallel SVN updates, and experiments report fast convergence, improved accuracy over SVGD and SVN, and scalability across parameter and sample counts and processor cores.
Problem
Posterior sampling becomes prohibitive in high dimensions, where conventional methods can suffer exponential computational growth or deteriorating convergence.
Source: §1 · p. 1Method
pSVN adaptively constructs a Hessian-informed, data-informed subspace and performs parallel SVN updates on parameter samples projected into it.
Source: §1 · p. 2, §3.3 · p. 5Results
Across linear and nonlinear inference problems, pSVN showed faster convergence and higher accuracy than SVGD and SVN, with complexity independent of parameter and sample dimensions and scalability across processor cores.
Source: §5 · p. 8Takeaways & Limitations
The experiments support using a low-rank approximation of averaged Hessian information to exploit posterior geometry in high-dimensional inference.
Source: §5 · p. 8The convergence statement assumes Gaussian observation noise and an existence bound with constant Cf for the specified norm.
Source: §3.1 · p. 4
Abstract
from arXiv · showhide
We propose a fast and scalable variational method for Bayesian inference in high-dimensional parameter space, which we call projected Stein variational Newton (pSVN) method. We exploit the intrinsic low-dimensional geometric structure of the posterior distribution in the high-dimensional parameter space via its Hessian (of the log posterior) operator and perform a parallel update of the parameter samples projected into a low-dimensional subspace by an SVN method. The subspace is adaptively constructed using the eigenvectors of the averaged Hessian at the current samples. We demonstrate fast convergence of the proposed method and its scalability with respect to the number of parameters, samples, and processor cores.
1 Introduction
High-dimensional Bayesian inference faces costly posterior sampling, while existing methods can lose efficiency or accuracy in challenging settings. The paper introduces pSVN, which projects parameter samples into a posterior-informed low-dimensional subspace and reports fast, scalable convergence.
1 Introduction
- High-dimensional posterior sampling is often prohibitive because conventional methods can scale exponentially or converge more slowly as parameter dimension grows.The introduction frames this as the curse of dimensionality in Bayesian inference.Source: §1 · p. 1
- Geometry-aware MCMC can improve sampling and sparse quadrature can converge rapidly on smooth, sparse problems, but expensive likelihoods and local posterior variation remain challenges.The passage notes that sparse-quadrature convergence can be lost when the posterior has significant local variations.Source: §1 · p. 1
- Variational inference recasts posterior sampling as KL-minimizing optimization, but Stein gradient and Newton methods can lose convergence and accuracy in high dimensions.These methods are reported to converge quickly in relatively low dimensions.Source: §1 · pp. 1–2
- pSVN projects parameters into a data-informed low-dimensional subspace and applies SVN there, with experiments reporting convergence independent of parameter and sample counts.The reported intrinsic dimensions are 6 and 40 for nominal dimensions over 1K and 16K, respectively.Source: §1 · p. 2
2 Background
The background frames Bayesian inference as updating prior beliefs with observations and describes Stein variational methods as sequences of transport-map updates. Their updates use gradient information and preconditioning, with kernels and metrics encoding posterior geometry.
2.1 Bayesian inference
- Bayesian inference conditions prior parameter beliefs on observations through Bayes' rule, but its normalization constant is typically computationally intractable, especially at large d.The setup considers a parameter, a parameter-to-observable map, and noisy observations.Source: §2.1 · p. 2
2.2 Stein variational methods
- Stein variational methods approximate a transport from prior to posterior as a sequence of simpler maps, each a perturbation of the identity.The transport is framed as minimizing a Kullback–Leibler divergence.Source: §2.2 · p. 2
- The update uses a preconditioned cost-functional gradient; SVGD uses the identity preconditioner, whereas SVN approximates the Hessian.This distinguishes the gradient and Newton variants within the shared update framework.Source: §2.2 · p. 3
- The transport perturbation is represented with basis functions and unknown coefficient vectors, whose combined vector is obtained from a linear system.The coefficient-vector formulation is given for the chosen basis functions.Source: §2.2 · p. 3
- Kernel metrics can incorporate averaged Hessian geometry, and randomized low-rank decomposition is proposed when direct Hessian computation is intractable.The cited passage reports that Hessian-informed metrics accelerated SVGD and SVN relative to a rescaled identity metric.Source: §2.2 · p. 3
3 Projected Stein variational Newton
pSVN addresses high-dimensional inference by updating samples in a low-dimensional, data-informed subspace. It combines Hessian-based projection with Stein variational Newton updates, and adapts the subspace using current samples.
3.1 Dimension reduction by projection
- The projected posterior converges to the full posterior as the projected parameter converges in X-norm, with convergence bounded by that of the parameter.In practice, projected convergence may be faster because it depends only on the data-informed subspace, unlike the full parameter’s data-independent X-norm.Source: §3.1 · p. 4
3.2 Projected Stein variational Newton
- When r ≪ d, mass-lumping reduces the projected system from size Nr × Nr to N systems of size r × r.The projection targets the low-dimensional parameter subspace; the complementary component is kept unchanged during the update.Source: §3.2 · pp. 4–5, §3.1 · p. 4
3.3 Hessian-based subspace
- The subspace uses leading eigenvectors of an averaged negative-log-likelihood Hessian, whose spectrum identifies parameter directions most informed by data.A randomized SVD solves the generalized eigenproblem in O(NrCh + dr^2) flops, where Ch is the cost of a Hessian action in one direction.Source: §3.3 · p. 5
3.4 Parallel and adaptive pSVN algorithm
- The algorithm adaptively rebuilds the basis from samples pushed toward the posterior and uses MPI for parallel updates.Its stopping criteria include update or gradient norms below tolerances, or reaching a preset iteration limit.Source: §3.4 · p. 6
4 Numerical experiments
The numerical experiments compare pSVN with SVGD and SVN on a linear inference problem, then test pSVN on a nonlinear problem for accuracy, dimension and sample independence, and processor-core scalability. They report faster convergence and higher accuracy for pSVN in the comparisons described.
4.1 A linear inference problem
- In the linear problem, pSVN converges faster and estimates mean and especially variance more accurately than SVGD and SVN; SVN variance deteriorates with dimension.With d = 256, pSVN converges rapidly in a subspace of dimension 6; increasing the sample count improves its accuracy.Source: §4.1 · p. 7, Figure 1 · p. 7
4.2 A nonlinear inference problem
- At d = 1089, pSVN converges faster and achieves higher accuracy than SVN for both mean and variance estimates.The comparison also reports that SVN’s high-dimensional kernel is more expensive per iteration than pSVN.Source: §4.2 · p. 7, Figure 2 · p. 8
- As d increases from 289 to over 16K, the Hessian-based subspace dimension r does not change; convergence is similar as sample count increases, and pSVN shows nearly perfect strong scaling with processor cores.Variation dominates component time; kernel and Newton-system solve costs depend on r, not d.Source: §4.2 · p. 8, Figure 3 · p. 8
5 Conclusion
The conclusion presents pSVN as a fast, scalable high-dimensional Bayesian inference method and reports demonstrations on linear and nonlinear problems. It identifies intrinsically high-dimensional problems and broader models as directions for further work.
5 Conclusion
- pSVN exploits posterior geometry through a low-rank approximation of the averaged Hessian of the log likelihood, computed with randomized matrix-free SVD.The method targets the intrinsic low-dimensionality of the change from prior to posterior.Source: §5 · p. 8
- Linear and nonlinear inference examples demonstrate pSVN’s fast convergence, higher accuracy relative to SVGD and SVN, dimension- and sample-independent complexity, and processor-core scalability.Source: §5 · p. 8
- Further work includes intrinsically high-dimensional problems, more general probability distributions and forward models, and additional convergence and scalability analysis.The conclusion gives local-dimensional message passing and complement-subspace MCMC updates as possible approaches for intrinsically high-dimensional problems.Source: §5 · p. 8
Appendix A: Proof of Theorem 1
The proof bounds the discrepancy between the posterior and projected posterior by treating an integral term and a normalizing-constant term separately.
Appendix A: Proof of Theorem 1
- The proof begins from the definitions of the posterior and projected posterior densities and separates the bound into two terms.Source: Appendix A · p. 11
- Under Assumption 1, the first integral term is bounded by a constant times the norm of the difference between x and x^r.The stated bound is (I) ≤ C_I ||x − x^r||_X.Source: Appendix A · p. 11
- The proof bounds the normalizing-constant term using the first integral bound and the inequality log(1 + τ) ≤ τ for τ ≥ 0.Source: Appendix A · p. 11
Appendix B: Globalization by line search
Because the Kullback–Leibler objective is generally nonconvex, the method uses an Armijo line search to choose step sizes. The line search evaluates a sample-average approximation and can use sample-dependent step sizes.
Appendix B: Globalization by line search
- Except for linear inference, the Kullback–Leibler objective is nonconvex, making constant step sizes inefficient for its generally complex, non-quadratic local geometry.The passage identifies step-size choice as important for convergence speed and stability.Source: Appendix B · p. 11
- An Armijo line search chooses ε by seeking to minimize the Kullback–Leibler divergence, evaluating the relevant term with a sample-average approximation.Source: Appendix B · p. 12
- Sample-dependent step sizes can be chosen so the second Kullback–Leibler term is reduced for each sample.The passage notes that this term is close to zero when kernel values between distinct samples are close to zero.Source: Appendix B · p. 12
Appendix C: Complexity analysis for parallel pSVN
Parallel pSVN distributes sample-level calculations and communicates the data needed to assemble its update. Its communication costs are bounded for the main algorithm, while adaptive basis construction offers a separate parallel strategy.
Appendix C: Complexity analysis for parallel pSVN
- Algorithm 1 communication volumes are bounded by max(Mr^2, MN) floats.The communicated data include low-dimensional samples, gradients and Hessians, kernel values, and sums used to assemble the lumped Hessian.Source: Appendix C · p. 12
- For adaptive pSVN, parallel basis construction uses averaged Hessian actions and randomized SVD, with dr floats communicated.The basis step costs O(M(rCh)) flops per core for Hessian actions and O(dr^2) flops for randomized SVD; when d is large, dr may dominate communication.Source: Appendix C · p. 12
- Alternatively, each core can construct its basis Ψ from its local samples without communication for Ψ.Source: Appendix C · p. 12
Appendix D: Bayesian Autoencoder Example
The Bayesian autoencoder example applies pSVN to infer neural-network parameters under noisy MNIST data. In this 133-parameter case, the effective Hessian rank is 40, and pSVN minimizes training error faster than SVN.
Appendix D: Bayesian Autoencoder Example
- The autoencoder maps input data through m encoding layers to a low-dimensional latent representation and through m decoding layers back to data space.Its 2m-layer mapping uses convolution kernels and nonlinear activation functions, with operations representing convolution and downsampling.Source: Appendix D · pp. 12–13
- The inference uses 1,000 randomly selected MNIST images, with 5% i.i.d. noise added to the min-max-normalized targets.The objective is a least-squares misfit between reconstructed inputs and noisy target images.Source: Appendix D · p. 13
- The convolution-kernel support is fixed at 4 × 4 while the number of filters varies across 2, 4, and 8, with m = 2 layers.Source: Appendix D · p. 13
- The subspace rank depends on the decay of the absolute Hessian eigenvalues, which are used to choose the pSVN subspace in Figure 4.The example reports low-rank Hessian structure for neural-network training problems.Source: Appendix D · p. 13, Figure 4 · p. 13
- pSVN minimizes training error faster than SVN in the 133-parameter example, using 128 particles and rank r = 40.The chosen rank is reported as the effective rank of the prior-preconditioned Hessian, whose absolute eigenvalues appear in Figure 4.Source: Appendix D · p. 13, Figure 4 · p. 13
Appendix E: Code
The authors implemented the Stein variational methods and DILI MCMC in hIPPYlib, which relies on FEniCS.
Appendix E: Code
- The implementations and test code are available through hIPPYlib and the pSVN GitHub repository.hIPPYlib relies on FEniCS for solving partial differential equations.Source: Appendix E · p. 13