Source-linked AI summary
Stochastic Stein Discrepancies
Jackson Gorham, Anant Raj, Lester Mackey
TL;DR
Approximate MCMC and Stein discrepancies can both be costly to evaluate: the former uses subsampling but introduces persistent bias, while the latter may require evaluating an expensive operator. The paper subsamples decomposable Stein operators to form stochastic discrepancies and stochastic SVGD, reporting comparable or more accurate inferences with orders of magnitude fewer datapoint accesses.
Problem
Approximate MCMC uses subsampling to reduce computation but introduces persistent bias, while computing a Stein discrepancy can be prohibitive when its operator is expensive to evaluate.
Source: §1 · pp. 1–2Method
The paper subsamples decomposable Stein operators to define SSDs and uses minibatch score estimates for stochastic SVGD updates.
Source: §3 · p. 3, §3.2 · p. 4Results
Across biased MCMC hyperparameter tuning, approximate sampler selection, and particle-based inference, SSDs deliver equivalent or more accurate inferences than standard SDs with orders of magnitude fewer datapoint accesses.
Source: §1 · p. 2Takeaways & Limitations
The paper identifies SSDs as a practical lower-cost sample-quality measure and names other tasks based on decomposable Stein operators as potential applications.
Source: §6 · p. 9The stated non-convergence guarantee is conditional on standard dissipativity and growth conditions; the cited kernel examples satisfy growth assumptions under specified score conditions.
Source: §4.2 · p. 6, §4.3 · p. 7
Abstract
from arXiv · showhide
Stein discrepancies (SDs) monitor convergence and non-convergence in approximate inference when exact integration and sampling are intractable. However, the computation of a Stein discrepancy can be prohibitive if the Stein operator - often a sum over likelihood terms or potentials - is expensive to evaluate. To address this deficiency, we show that stochastic Stein discrepancies (SSDs) based on subsampled approximations of the Stein operator inherit the convergence control properties of standard SDs with probability 1. Along the way, we establish the convergence of Stein variational gradient descent (SVGD) on unbounded domains, resolving an open question of Liu (2017). In our experiments with biased Markov chain Monte Carlo (MCMC) hyperparameter tuning, approximate MCMC sampler selection, and stochastic SVGD, SSDs deliver comparable inferences to standard SDs with orders of magnitude fewer likelihood evaluations.
1 Introduction
Approximate MCMC scales by using data subsamples, but its persistent bias undermines standard diagnostics that rely on asymptotic exactness. The paper develops stochastic Stein discrepancies to retain convergence tracking while reducing evaluation cost.
1 Introduction
- Biased approximate MCMC uses data subsamples to generate more samples in less time, but its persistent bias limits diagnostics that rely on asymptotic exactness.Standard MCMC methods provide asymptotically correct sample estimates, whereas approximate methods trade that property for speed.Source: §1 · p. 1
- Stein discrepancies track sample convergence even when exact integration and direct sampling from the target are intractable.They have also been applied to compare approximate MCMC, test goodness of fit, and train generative models.Source: §1 · p. 1
- Stochastic Stein discrepancies inherit standard SDs’ convergence-tracking properties with probability 1, while experiments report equivalent or more accurate inferences using orders of magnitude fewer datapoint accesses.The paper also proves convergence of SVGD on unbounded domains and that stochastic SVGD converges to the same limit as SVGD with probability 1.Source: §1 · p. 2
2 Measuring Sample Quality
Sample-quality measures should detect both convergence and non-convergence while remaining computationally tractable. Stein discrepancies address the intractability of many conventional probability metrics by using operators that generate functions with known zero expectation under the target.
2 Measuring Sample Quality
- A sample-quality measure should determine whether empirical measures converge to a target, detect non-convergence, and be computationally tractable.The paper considers discrete measures formed from sample points without assuming whether they come from sampling, a Markov chain, or deterministic quadrature.Source: §2 · p. 2
- Wasserstein and Dudley distances can determine convergence in distribution, but typically require exact integration under the target.Stein discrepancies avoid that requirement by using test functions known to have mean zero under the target.Source: §2 · p. 2
- A decomposable Stein operator is a sum of cheaper base operators; in the Langevin posterior example, each base operator accesses one datapoint.This decomposition is the structure later used to subsample operator terms.Source: §2 · p. 3
3 Stochastic Stein Discrepancies
The paper subsamples the terms of decomposable Stein operators to reduce discrepancy computation while retaining convergence control. It also gives a closed form for stochastic kernel Stein discrepancies and applies stochastic updates to SVGD.
3 Stochastic Stein Discrepancies
- Independent uniformly sampled subsets of base operators define SSDs, reducing operator evaluations while preserving convergence detection with probability 1.Reusing one minibatch for all sample points would instead track convergence to a posterior conditioned on that minibatch.Source: §3 · p. 3
3.1 Stochastic kernel Stein discrepancies
- When the Stein set is a unit ball in an RKHS, the SSD has a closed-form solution for the stated regularity condition on the kernel.The paper names these stochastic kernel Stein discrepancies (SKSDs).Source: §3.1 · pp. 3–4
3.2 Stochastic Stein variational gradient descent
- SSVGD replaces each full score evaluation with an independent minibatch estimate, reducing per-round gradient evaluations from Ln to mn.The update moves particles to minimize a stochastic KSD; the paper analyzes this approach alongside the distinct burden of evaluating an expensive Stein operator.Source: §3.2 · p. 4, §3.1 · p. 4
4 Convergence Guarantees
The section establishes probability-1 convergence guarantees for stochastic Stein discrepancies and extends convergence results for SVGD and SSVGD to unbounded domains. It also shows how SSDs detect non-convergence under stated operator and Stein-set conditions.
4.1 Detecting convergence with SSDs
- If Qn converges to P in Wasserstein distance, the SSD converges to zero almost surely under the stated regularity conditions on its base operators.The result covers Stein operator-set pairings with standard SD convergence-detection results.Source: §4.1 · p. 5
4.2 Detecting non-convergence with SSDs
- With coercive test functions, SSDs remain bounded away from zero for non-tight sequences; under Theorem 6 conditions, they detect any failure of Qn to converge to P with probability 1.Theorem 6 applies to specified kernel Stein sets and positive definite matrices under its stated assumptions.Source: §4.2 · p. 6
- SSDs detect non-convergence with probability 1 when bounded SDs detect it, and they also detect sequences that are not tight under the stated conditions.For the Langevin operator and convergence-determining Stein sets, the result holds under standard dissipativity and growth conditions on the subsampled operator.Source: §4.2 · pp. 5–6
4.3 Convergence of SVGD and SSVGD
- On Rd, SVGD and SSVGD outputs converge to their limiting distributions with probability 1 at each round under assumptions commonly satisfied by p and k.The result addresses an earlier analysis gap: its bounded-and-Lipschitz requirement typically fails because of the growth of the target score times the kernel.Source: §4.3 · p. 6
5 Experiments
The experiments test stochastic Stein discrepancies as lower-cost diagnostics for approximate MCMC and stochastic SVGD. Across sampler tuning, biased-sampler selection, and particle inference, subsampling preserves useful comparisons and stochastic SVGD improves approximation accuracy.
5.1 Hyperparameter selection for approximate MCMC
- SGLD step-size tuning must balance slow mixing from overly small values against severe distributional bias from overly large values.The experiment recreates an earlier setup for assessing whether stochastic KSD can tune the step size.Source: §5.1 · p. 7
- With L = 100 likelihood terms, ϵ = 5 × 10^-3 minimizes every KSD variant’s average across 50 trials, and the fastest SKSD uses one hundredth the likelihood evaluations.Figure 1 shows the same candidate ranking and selected step size under subsampling.Source: §5.1 · p. 7, Figure 1 · p. 8
5.2 Selecting biased MCMC samplers
- SGFS-f uses a full matrix inverse, while SGFS-d zeroes the off-diagonal entries before inversion; both are uncorrected and asymptotically biased from the target.The comparison asks whether SGFS-d’s greater sample count under a fixed budget offsets its additional asymptotic bias.Source: §5.2 · p. 7
- The stochastic KSD prefers SGFS-f over SGFS-d, matching the exact KSD while using 100 times fewer likelihood evaluations at the most aggressive subsampling rate.The stochastic variants subsample 100 or 1,000 of the 10,000 likelihood terms for each sample iterate.Source: Figure 2 · p. 8, §5.2 · p. 8
5.3 Particle-based variational inference with SSVGD
- For every likelihood budget and all three datasets, SSVGD gives more accurate approximations than full-batch SVGD, with the effect magnified on naval.Figure 3 reports test RMSE and log likelihood across 20 train-test splits.Source: §5.3 · p. 9, Figure 3 · p. 9
6 Discussion and Future Work
The discussion presents SSDs as a lower-cost way to assess and improve sample quality, and identifies extensions to other inference tasks and Stein operators.
6 Discussion and Future Work
- SSDs retain the convergence-determining properties of standard discrepancies with probability 1 while requiring orders of magnitude fewer likelihood evaluations.The paper proposes applying these developments to tasks built on decomposable Stein operators.Source: §6 · p. 9
- Potential applications include goodness-of-fit testing, KSD-based sampling, Monte Carlo control variates, sample reweighting or thinning, and parameter estimation in intractable models.The discussion lists these as candidates for future application.Source: §6 · p. 9
- The analysis also extends beyond the Langevin operator to diffusion and discrete Stein operators; variance reduction is proposed as a way to make SSDs more like standard SDs while retaining savings.Source: §6 · p. 9
Broader Impact
The paper frames its contribution as scalable diagnostics for approximate inference and cautions that downstream conclusions depend on appropriate data and models.
Broader Impact
- The work offers producers and consumers of approximate inference techniques a diagnostic for assessing approximations at scale.It also analyzes SSVGD as a scalable algorithm for improving approximate inference.Source: Broader Impact · p. 9
- The authors expect existing Stein-discrepancy users to adopt SSDs to reduce computational costs, and suggest scalable diagnostics may encourage wider use of approximate MCMC.Source: Broader Impact · p. 9
- Incorrect data or an inappropriate model can still lead an inferential tool to harmful conclusions, so downstream results require careful interpretation.Source: Broader Impact · p. 9
B Proof of Theorem 2: SSDs detect convergence
These proof sections establish convergence results by relating random-measure convergence to uniform control of expectations, then applying approximation arguments to Stein sets and bounded discrepancies.
B Proof of Theorem 2: SSDs detect convergence
- The proof uses a lemma to turn almost-sure convergence in distribution into almost-sure uniform convergence of expectations for functions meeting integrability and derivative-growth conditions.The functions are dominated by a uniformly integrable, locally bounded |f0|, with derivatives dominated by a locally bounded |f1|.Source: B Proof of Theorem 2: SSDs detect… · p. 13
C.1 Kernel Stein set
- For the kernel Stein set, if the SSD tends to zero, the corresponding IPM tends to zero and therefore the sample measures converge weakly to the target.The proof passage identifies the Stein equation, smoothing, and truncation steps used in the argument.Source: §C.1 · pp. 15–16
C.2 Classical Stein set
- For the classical Stein set, the result follows by embedding the kernel Stein set into a scaled copy of the bounded Stein set.Source: §C.2 · p. 16
C.3 Graph Stein set
- For the graph Stein set, the result follows because it contains the classical Stein set.Source: §C.3 · p. 16
D Proof of Theorem 4: SSDs detect bounded SD non-convergence
- The bounded-function convergence lemma gives an almost-sure control result for triangular arrays when each function in the sequence is bounded by one.The proof applies the lemma to compare subsampled and full operator evaluations.Source: D Proof of Theorem 4: SSDs detect… · pp. 16–17
D.1 Proof of Lemma 10: Bounded function convergence
Lemma 11 establishes high-probability convergence rates for subsampled measure expectations to their non-subsampled counterparts. Choosing a summable failure-probability sequence yields the advertised almost-sure result.
D.1 Proof of Lemma 10: Bounded function convergence
- Lemma 11 bounds the difference between subsampled and non-subsampled measure expectations with probability at least 1 − δ.The stated result applies for a ∈ [1, 2] and functions h : R^d → R.Source: §D.1 · p. 17
- The expectation bound depends on δ, n, a, and ν_n(|h|^a), as specified in Lemma 11.The supplied display is incomplete, so its exact bound is not recoverable here.Source: §D.1 · p. 17
- Taking δ_n = 1/(n log^2(n)) and applying the Borel–Cantelli lemma gives the almost-sure conclusion.Source: §D.1 · p. 17
E Proof of Prop. 5: Coercive SSDs enforce tightness
The proof uses coercivity to turn failure of tightness into a positive lower bound on the SSD. Thus, a coercive SSD cannot converge to zero when the sequence of measures is not tight.
E Proof of Prop. 5: Coercive SSDs enforce tightness
- A coercive, bounded-below function f yields a nonnegative, non-decreasing coercive function γ measuring f’s growth outside radius r.The construction is γ(r) = inf{f(x) − C : ∥x∥₂ ≥ r}, where C = inf_x f(x).Source: E Proof of Prop. 5: Coercive SSDs… · p. 17
- If (Q_n) is not tight, some radius R has lim sup_n Q_n(∥X∥₂ > R) ≥ ϵ for ϵ > 0.Monotonicity of γ and Markov’s inequality bound this tail probability using (Q_n(f) − C)/γ(R).Source: E Proof of Prop. 5: Coercive SSDs… · p. 17
- The SSD is bounded below by the coercive Stein expectation, so lim sup_n SS(Q_n, T, G_n) ≥ γ(R)ϵ + C > 0.Source: E Proof of Prop. 5: Coercive SSDs… · p. 17
F Proof of Theorem 6: Coercive SSDs detect non-convergence
For the considered Stein sets, the proof shows that non-convergence of Q_n to P forces the SSD not to converge to zero. It handles kernel, classical, and graph Stein sets, and extends the argument to general positive-definite Γ by a change of coordinates.
F Proof of Theorem 6: Coercive SSDs detect non-convergence
- For the specified kernel Stein sets, non-convergence either leaves the bounded Stein discrepancy nonzero or makes (Q_n) not tight; both routes imply the SSD does not vanish.The argument constructs a shared Stein-set function whose Stein expectation dominates a coercive, bounded-below function.Source: F Proof of Theorem 6: Coercive SSDs… · p. 18
- The result extends to any Γ ≻ 0 by changing coordinates and applying Theorem 6 to the transformed distributions with a non-preconditioned kernel.The proof uses the transformed score and dissipativity properties, citing [10, Lem. 4].Source: F Proof of Theorem 6: Coercive SSDs… · p. 18
- For a classical Stein set, containment of the kernel Stein set bounds its SSD by a constant times the classical-set SSD.For a graph Stein set, the result follows because it contains the classical Stein set.Source: F Proof of Theorem 6: Coercive SSDs… · p. 18
G.2 Proof of Theorem 7
The proof of Theorem 7 establishes convergence of SSVGD to SVGD by induction over iterations. Its key ingredients are pseudo-Lipschitz control of the SVGD mapping and almost-sure convergence of each stochastic step to its deterministic counterpart.
G.2 Proof of Theorem 7
- The convergence claims are proved by induction on r ≥ 0; specifically, for W1 convergence, the base case W1(Qn,0, Q∞,0) → 0 holds by assumption.Source: §G.2 · pp. 18–19
G.3 Proof of Lemma 12: Wasserstein pseudo-Lipschitzness of SVGD
- Under the stated growth assumptions, the SVGD mapping is pseudo-Lipschitz in Wasserstein distance, with its bound depending on the measures’ first moments.The bound uses constants c1 and c2 from derivative growth conditions.Source: §G.2 · p. 19, §G.3 · p. 20
G.4 Proof of Lemma 13: One-step convergence of SSVGD to SVGD
- For uniformly integrable starting measures and the stated continuity condition, one SSVGD step converges in Wasserstein distance almost surely to one SVGD step.This one-step result supports the induction establishing convergence across iterations.Source: §G.2 · p. 19, §G.4 · p. 20