A Stein variational Newton method

1 Introduction

• SVGD occupies a middle ground between cheap but potentially inaccurate mean-field variational inference and potentially slow-to-converge MCMC.The paper focuses on nonparametric variational inference, with close links to density estimation.Source: §1 · p. 1
• SVGD builds transport maps by perturbing the identity along an RKHS-projected steepest-descent direction for the KL objective.The maps are composed sequentially to transform samples from a tractable reference distribution toward the target.Source: §1 · p. 2
• The paper’s Stein variational Newton method replaces steepest descent with a Newton-like direction based on a local quadratic approximation of the objective.Second-order information can accelerate convergence, at the cost of additional work per iteration; several numerical examples favor this tradeoff.Source: §1 · p. 2
• Geometry-aware Gaussian kernels use second-order information and yield substantially faster convergence than SVGD or SVN with an isotropic kernel.Kernel choice matters to the efficiency of both SVGD and SVN.Source: §1 · p. 2

2 Background

• A transport map sends samples from a tractable reference density to approximate an intractable target, and the pushforward samples form an empirical approximation.The map is not unique: infinitely many can yield the target.Source: §2 · p. 2
• SVGD minimizes KL divergence by iteratively composing identity perturbations chosen from a vector-valued RKHS.At each step, the next map is selected to further decrease the divergence between the current pushforward and the target.Source: §2 · pp. 2–3
• The SVGD iteration uses a functional gradient at the zero map to choose the next transport perturbation.The resulting sequence of pushforward densities becomes increasingly closer to the target in KL divergence.Source: §2 · p. 3
• The particle update combines a weighted descent direction for the log-target density with a repulsion force that spreads particles across the target’s support.The descent term moves particles toward high-probability regions, while repulsion helps prevent collapse around a mode.Source: §2 · p. 3

3 Stein variational Newton method

• The SVN update replaces SVGD’s steepest-descent direction with an approximation of the Newton direction.The direction minimizes a local quadratic approximation of the objective, and the next transport map perturbs the identity along it.Source: §3 · pp. 3–4
• Galerkin approximation restricts the Newton-direction search to a finite-dimensional space spanned by kernel functions evaluated at the particles.The coefficients are characterized by a linear system involving the second variation and the first variation evaluated at the particles.Source: §3 · p. 4
• Inexact Newton–conjugate gradient avoids explicitly constructing the full Newton matrix and can guarantee a descent direction through its stopping rule.The full system has nd unknowns and may be impractical for many particles or high-dimensional parameters; direct Newton solutions may also fail to descend when the target is not log-concave.Source: §3 · p. 4
• A block-diagonal second-variation approximation breaks the system into n decoupled d × d systems that can be solved in parallel.Both block-diagonal approximation and inexact NCG are more efficient than solving for the full Newton direction, and provide similar progress per SVN iteration.Source: §3 · p. 5
• Second-order information gives a natural step-size scale of ε = O(1), and ε = 1 performs reasonably well in the reported experiments.The authors identify line search or a trust-region step as future refinements.Source: §3 · p. 5

4 Scaled Hessian kernel

• The kernel’s choice matters because it weights local descent contributions and spreads particles, so it should capture the target’s geometry.The paper uses curvature information from the log-target Hessian to design anisotropic kernels.Source: §4 · p. 5
• The curvature metric is built from an approximation to the local Hessian of the negative log-target, averaged under the target distribution.Because the expectation under the target cannot be computed directly, the particle ensemble is used for an approximation; the Newton update already calculates approximations to the required Hessian.Source: §4 · p. 5
• The scaled Hessian kernel uses an average-curvature metric to define anisotropic distances, stretching or compressing parameter-space directions according to curvature.The metric shapes both locally averaged gradients and particle repulsion, and tends to spread particles more effectively over high-probability regions.Source: §4 · pp. 5–6
• The dimension-dependent scaling g(d) can make the rescaled norm converge to a functional L2 norm as discretization dimension increases, under stated technical conditions.For targets without a well-defined infinite-dimensional limit, an appropriate scaling can still improve discrimination among inter-particle distances.Source: §4 · p. 6

5 Test cases

• The test cases compare SVN and SVGD with scaled-Hessian or heuristically scaled isotropic kernels on multimodal and high-dimensional tasks.The two main cases are a two-dimensional bimodal banana-shaped posterior and a high-dimensional model arising from a Langevin SDE.Source: §5 · p. 6, §5.1 · p. 6, §5.2 · p. 7

5.1 Two-dimensional double banana

• On the double-banana target, SVN-H converges by 10 iterations with structured particles; SVN-I converges by 47 iterations but remains rougher.Figure 1 shows particle configurations over the target-density contours.Source: Figure 1 · p. 7, §5.1 · p. 7
• The Hessian kernel takes less computational time than the isotropic kernel, which recalculates particle distances at each iteration for heuristic rescaling.The comparison rescales iteration counts to make total CPU time approximately equal.Source: §5.1 · pp. 6–7

5.2 100-dimensional conditioned diffusion

• With comparable computational effort, SVN-H reconstructs the posterior mean and credible intervals accurately, while SVGD-H and SVGD-I depart significantly from the true values.Figure 2 encodes the true path, reconstructed posterior mean, and marginal 90% credible intervals.Source: Figure 2 · p. 8, §5.2 · pp. 8–9
• SVN-H's posterior means and marginal 90% credible intervals agree well with reference DILI MCMC using an effective sample size of 105.Both methods' posterior means are also comparable to the data-generating path.Source: Figure 3 · p. 9, §5.2 · p. 9

• Gaussian reproducing kernels may be problematic in high dimensions because of their locality.Source: §6 · p. 9
• Properly rescaling a Gaussian kernel can improve SVN performance in high dimensions.Source: §6 · p. 9
• The authors suggest designing algorithms to detect and exploit target structure, such as decaying correlation, conditional independence, low rank, or multiple scales.Source: §6 · p. 9

The first variation takes the form

• A Taylor expansion of terms in the directional-derivative expression supplies intermediate terms for the variation calculation.Source: The first variation takes the form · p. 10

A Proof of Theorem 1

• The proof uses a proposition about the second variation of Jp at zero along directions V and W to establish Theorem 1.Source: The first variation takes the form · pp. 10–11

The first variation takes the form

• The RKHS reproducing property is applied to V and W and their gradients when expressing the second variation.Source: The first variation takes the form · p. 11

B Proof of Corollary 1

• The Galerkin approach seeks the Newton direction in the span of kernel functions evaluated at the particles.Source: B Proof of Corollary 1 · pp. 11–12
• Projecting the Newton direction onto particle-indexed test functions yields equations involving the Hessian kernel and the gradient of Jpl at zero.Source: B Proof of Corollary 1 · p. 12
• The resulting block system is Hs,k αk = ∇Js for each particle index s.Source: B Proof of Corollary 1 · p. 12

C.1 Comparison between the full and inexact Newton methods

• SVNfull solves the fully coupled Newton system, SVNCG applies inexact Newton–CG to it, and SVNbd uses a block-diagonal approximation.Source: §C.1 · p. 12
• On the two-dimensional double banana and a second test, both the block-diagonal approximation and inexact Newton–CG produced results of similar quality to full Newton; each used N=100 particles and 20 iterations.The authors propose these approximate SVN strategies for large-scale problems where computing the full Newton direction is not feasible.Source: §C.1 · p. 12

C.2 Bayesian neural network

• The experiment uses 308 input–prediction pairs, dividing 247 pairs for training and 61 additional pairs for validation; its neural network has a 2,951-dimensional weight vector with independent Gaussian priors.Source: §C.2 · p. 13
• At the posterior mode, the neural network’s predictions almost perfectly match the validation data.Source: Figure 6 · p. 14, §C.2 · p. 14
• For the high-dimensional model, SVN uses inexact Newton–CG, compared with SVGD-I at 10, 20, and 30 iterations matched for computational cost.Source: §C.2 · p. 14
• SVN-H yields validation prediction-error distributions with smaller means and considerably reduced variances than SVGD-I, while both have small overall error ranges.The comparison matches computational costs by rescaling SVGD-I iterations.Source: §C.2 · pp. 14–15

Scalability experiments

• The tests use 40, 60, 80, and 100 dimensions, with 1,000 particles and 50 SVN iterations; the non-invariant case uses identity prior covariance and uniformly drawn forward-operator coefficients.Source: §C.3 · p. 16

Discretization-invariant posterior distribution

• For the discretization-invariant posterior, SVN-H estimates covariance trace consistently across dimensions; SVN-I underestimates it, indicating under-dispersed particles.Source: §C.3 · p. 16, Table 2 · p. 16
• All algorithms recover the average posterior mean to the first three significant figures in the discretization-invariant case.Source: §C.3 · p. 16, Table 1 · p. 16

A posterior distribution that is not discretization invariant

• In the non-discretization-invariant problem, the scaled Hessian kernel remains reasonably effective as dimension rises, while the isotropic kernel underestimates variances at every tested dimension.Source: §C.3 · p. 16, Table 4 · p. 16