PaperlayerView exact revision on arXiv

Source-linked AI summary

Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm

Qiang Liu, Dilin Wang

arXiv:1608.04471v3stat.MLcs.LG

TL;DR

Bayesian inference requires tractable approximations to posterior distributions, while existing variational approaches can be restrictive or model-specific. The paper proposes particle-based functional gradient descent that minimizes KL divergence, supported by a KL–Stein discrepancy connection. Across real-world models and datasets, the method is competitive with existing state-of-the-art methods.

  • Problem

    Existing Bayesian inference methods face computational or scalability challenges, and variational methods can require restrictive approximation families or model-specific derivations.

  • Method

    The method iteratively transports particles toward the posterior using functional gradient descent on KL divergence, with a closed-form direction derived from smooth-transform analysis and kernelized Stein discrepancy.

  • Results

    The method is competitive with existing state-of-the-art methods across various real-world models and datasets.

  • Takeaways & Limitations

    The paper presents a simple, general-purpose variational inference algorithm for fast and scalable Bayesian inference.

  • Takeaways & Limitations

    The authors leave concrete theoretical analysis of propagation of chaos for future work, and parametric transform choices must balance accuracy, tractability, and solvability.

Abstract

from arXiv · show

We propose a general purpose variational inference algorithm that forms a natural counterpart of gradient descent for optimization. Our method iteratively transports a set of particles to match the target distribution, by applying a form of functional gradient descent that minimizes the KL divergence. Empirical studies are performed on various real world models and datasets, on which our method is competitive with existing state-of-the-art methods. The derivation of our method is based on a new theoretical result that connects the derivative of KL divergence under smooth transforms with Stein's identity and a recently proposed kernelized Stein discrepancy, which is of independent interest.

1 Introduction

The paper motivates a general-purpose Bayesian inference algorithm by contrasting scalable but restrictive variational methods with the flexibility needed for full posterior inference. It proposes particle-based functional gradient descent grounded in a connection between KL divergence and kernelized Stein discrepancy.

  • Motivation: MCMC can be slow and difficult to assess for convergence, while variational inference offers deterministic KL-divergence optimization that scales through stochastic gradient descent.The paper frames these trade-offs as central computational motivations for a general-purpose alternative.
  • Motivation: Variational inference accuracy and computational cost depend critically on the approximation family, with simple families too restrictive and advanced families harder to optimize.This dependence makes efficient variational methods difficult to design broadly.
  • Motivation: Model-specific derivations remain a major barrier to efficient, user-friendly variational tools applicable across models and accessible to non-ML experts.The paper contrasts this gap with the generic effectiveness of gradient descent for MAP optimization.
  • Contribution: The proposed algorithm uses particles and functional gradient descent to minimize KL divergence and transport particles toward the true posterior.It is presented as a natural counterpart of gradient descent for full Bayesian inference.
  • Contribution: The method is supported by a theoretical connection between KL-divergence derivatives under smooth transforms and kernelized Stein discrepancy.This connection yields a closed-form steepest-descent perturbation within an RKHS unit ball.

2 Background

The background introduces Stein’s identity and Stein discrepancy as tools for comparing distributions through the target score function. Kernelization makes the relevant functional optimization tractable and can distinguish distributions under suitable kernels.

  • Stein’s identity: Stein’s identity states that the Stein operator has zero expectation under the target density for sufficiently regular functions.The result follows under mild zero-boundary conditions and defines the Stein class.
  • Stein discrepancy: Under a different density, the Stein-operator expectation measures violation of the target identity and motivates Stein discrepancy.The discrepancy is formed by maximizing this violation over a chosen function class.
  • Kernelized Stein discrepancy: Traditional bounded-Lipschitz function classes make Stein discrepancy optimization computationally intractable or require special considerations.The choice of function class controls both discriminative power and computational tractability.
  • Kernelized Stein discrepancy: Kernelized Stein discrepancy uses an RKHS unit ball, where the optimization has a closed-form solution.The RKHS is constructed from a positive-definite kernel and its associated function space.
  • Kernelized Stein discrepancy: For strictly positive-definite kernels in a proper sense, Stein discrepancy equals zero if and only if the two distributions are equal.The paper notes that commonly used kernels such as the RBF kernel satisfy the relevant condition.
  • Unnormalized targets: Stein operators and kernelized Stein discrepancy require the target only through its score, which is available from an unnormalized density without its normalization constant.This property supports applications to unnormalized distributions common in machine learning and statistics.

3 Variational Inference Using Smooth Transforms

The paper represents variational approximations as smooth transforms of a tractable reference distribution, then iteratively constructs incremental transforms that minimize KL divergence. A Stein-operator result identifies the steepest RKHS descent direction, yielding a particle algorithm that avoids explicit parametric transforms and Jacobians.

  • Variational approximation: Variational inference selects q*(x) from a distribution set Q by minimizing KL divergence to the target p(x).The approximation family must balance accuracy, tractability, and solvability.
  • Smooth-transform family: The method uses smooth one-to-one transforms z = T(x) of samples from a tractable reference distribution q0.Expectations under transformed distributions can be evaluated by averaging transformed reference samples.
  • Smooth-transform family: Restricting transforms to parametric families creates a selection problem because T must be one-to-one with an efficiently computable Jacobian.The family must simultaneously preserve approximation accuracy, tractability, and optimization solvability.
  • Incremental transforms: The algorithm instead constructs incremental transforms that perform steepest descent on T in an RKHS without explicit parametric forms or Jacobian calculations.Its simple gradient-descent-like form is intended to make implementation accessible to non-experts.
  • Iterative procedure: Repeated small perturbations transport an initial distribution toward p, while functional-gradient evaluation at the identity avoids the inverse Jacobian required at nonzero transforms.The procedure repeatedly updates the transformed distribution and can be implemented with empirical particle expectations.
  • Stein direction: For T(x) = x + ϵφ(x), the derivative of KL divergence connects the Stein operator with kernelized Stein discrepancy, identifying φ* as the steepest-descent perturbation.The optimal direction lies in a zero-centered RKHS ball and the corresponding negative gradient equals the KSD.
  • Particle dynamics: At the particle level, one update term attracts particles toward high-probability regions and another repels neighboring particles to prevent collapse.As the RBF bandwidth h approaches zero, the repulsive term vanishes and the updates reduce to independent MAP-like gradient ascent.
  • Particle dynamics: With one particle, the algorithm reduces to gradient ascent for MAP, while multiple particles provide a full Bayesian particle approximation.The method can start from arbitrary points and deterministically transport them toward the target distribution.

4 Related Works

The paper relates its particle method to transformed-variable variational inference, particle-based methods, and Monte Carlo approaches. It emphasizes simpler optimization without matrix determinants or inversions.

  • Prior transformed-variable methods use specialized parametric transforms whose invertibility and Jacobian computation require additional conditions or case-by-case treatment.
  • The proposed algorithm instead uses a simpler particle-based form that does not require calculating matrix determinants or inversions.
  • Particle-based variational methods represent distributions with particles, while related Gaussian-mixture approaches often approximate entropy and expectation terms in their KL objectives.

5 Experiments

Experiments evaluate the method on toy and real-world Bayesian inference problems using particle-size variations and comparisons with sampling, variational, and particle-based baselines. The reported results include competitive estimation accuracy, strong logistic-regression performance, and improved neural-network accuracy and speed over PBP except on Yacht.

  • Toy Example on 1D Gaussian Mixture: The method moves particles toward a distant target distribution from an almost non-overlapping initialization in the one-dimensional Gaussian-mixture experiment.The experiment initializes particles with q0(x) = N(x; −10, 1) while the target has modes near −2 and 2.
  • Toy Example on 1D Gaussian Mixture: The particle estimates have mean square error similar to or better than exact Monte Carlo across the tested expectation functions.The tested functions are x, x2, and cos(ωx + b), with the latter averaged over 20 random draws of ω and b.
  • Toy Example on 1D Gaussian Mixture: The authors leave formal error-rate analysis for the particle method as an open question.
  • Bayesian Logistic Regression: The Covertype logistic-regression experiments use 80% training data, 20% testing data, 50 random trials, and mini-batches of 50.
  • Bayesian Logistic Regression: Each iteration of parallel SGLD is only 3 times more expensive than sequential SGLD despite the particle-based methods requiring 100 times as many likelihood evaluations in principle.
  • Bayesian Neural Network: The method consistently improves over PBP in accuracy and speed across Bayesian neural-network datasets except Yacht.The reported results are also comparable with more recent results using advanced techniques.

6 Conclusion

The paper proposes a simple, general-purpose variational inference algorithm for fast and scalable Bayesian inference.

  • The proposed algorithm targets fast and scalable Bayesian inference.

A Proof of Theorem 3.1

The proof derives the KL-divergence derivative under a general smooth, invertible transform and then specializes it to an additive perturbation.

  • For a smooth one-to-one transform T, the transformed density q[T] is defined from z = T_ϵ(x) with x sampled from q.
  • The KL derivative is expressed through the derivative of the inverse-transformed density p[T^-1].
  • The derivative of log p[T^-1](x) combines the target score evaluated at T(x) with a Jacobian-variation trace term.
  • For T(x) = x + ϵφ(x) at ϵ = 0, the transform and its derivatives reduce to φ(x), the identity matrix, and ∇xφ(x).

B Proof of Theorem 3.3

The proof computes the functional gradient of KL divergence in a vector-valued RKHS by differentiating the transformed target density and its Jacobian determinant.

  • The functional gradient ∇fF[f] is defined as an element of the vector-valued RKHS Hd.
  • The functional F[f] is represented as KL(q[x+f(x)] || p), equivalently comparing q with the inverse-transformed target density.
  • Differentiating F[f] separates changes in the target log density from changes in the log Jacobian determinant.
  • The first-order expansion converts the perturbation into an RKHS inner product involving the expected target score and kernel function.
  • The resulting gradient contains a target-score term and a trace involving the inverse Jacobian and the kernel derivative; setting f = 0 gives the stated result.

C Connection with de Bruijn’s identity and Fisher Divergence

The paper connects its KL-based transform analysis to Fisher divergence and contrasts its deterministic formulation with the randomized transform used in de Bruijn’s identity.

  • Choosing φq,p(x) = ∇x log p(x) − ∇x log q(x) reduces the result to an expression involving Fisher divergence.
  • Fisher divergence is introduced as the divergence quantity governing this reduced expression.
  • The connection is described as a deterministic version of de Bruijn’s identity, which instead uses T(x) = x + √ϵ · ξ with standard Gaussian noise ξ.

D Additional Experiments

Additional experiments report comparisons on Bayesian logistic regression and posterior prediction uncertainty using toy data. The reported results provide visual and comparative checks across inference approaches.

  • Figure 4 evaluates Bayesian logistic regression on eight datasets studied by Gershman et al.
  • The experiments compare the proposed method with NPV and NUTS on the logistic-regression datasets.
  • Figure 5 visualizes posterior prediction uncertainty inferred by different approaches on toy data.

D.1 Bayesian Logistic Regression on Small Datasets

The appendix evaluates the method on Bayesian logistic regression across eight datasets and a two-feature toy dataset. It reports similar performance to NPV and NUTS while noting practical differences in their computational requirements.

  • Bayesian Logistic Regression on Small Datasets: The Bayesian logistic-regression study compares the method with NUTS and NPV on eight datasets with more than 500 observations.The experiment uses 100 particles, 100 NPV mixture components, and 1,000 NUTS draws after 1,000 burn-in samples.
  • Bayesian Logistic Regression on Small Datasets: All three algorithms tend to perform similarly on the toy dataset's prediction probabilities.The toy dataset has two features, and the prediction probabilities are visualized for the three approaches.
  • Bayesian Logistic Regression on Small Datasets: NPV requires a Hessian matrix, while NUTS has a stated drawback when applied to massive datasets.
Loading 1608.04471v3…