Optimal Thinning of MCMC Output
Abstract
The use of heuristics to assess the convergence and compress the output of Markov chain
Monte Carlo can be sub-optimal in terms of the empirical approximations that are produced.
Typically a number of the initial states are attributed to “burn-in” and removed, whilst the remainder of the chain is “thinned” if compression is also required.
In this paper, we consider the problem of retrospectively selecting a subset of states, of ﬁxed cardinality, from the sample path such that the approximation provided by their empirical distribution is close to optimal.
We report a method based on greedy minimisation of a kernel Stein discrepancy, that is suitable for problems where heavy compression is required.
Theoretical results guarantee consistency of the method and its eﬀectiveness is demonstrated in the challenging context of parameter inference for ordinary diﬀerential equations.
1. Introduction
The most popular computational tool for non-conjugate Bayesian inference is Markov chain
Monte Carlo (MCMC; Robert and Casella, 2013).
The output ( X i ) ni =1 from MCMC is usually post-processed to remove burn-in and, if compression is desired, to reduce the total amount of data that are stored (useful, for example, when the MCMC output is used as experimental design in multi-scale models).
This results in just a subset of the states being retained, say indexed by π ∈ { 1 , . . . , n } m , corresponding to an empirical approximation
Q m := 1 m (cid:80) mj
=1 δ ( X π ( j ) ) of the distribution P being targeted.
Here δ ( x ) denotes a point mass centred at x .
From the perspective of statistical eﬃciency, this post-processing can be sub-optimal: for a ﬁxed computational budget, imposing tight control on bias through burn-in removal leads to estimators with high variance, due to the limited number of states that remain once the initial portion have been removed.
In this abstract we ask whether an “optimal” index set π can be selected.
To this end, we summarize Stein Thinning , a method introduced by the authors (Riabiz et al., 2020), that selects an index set π , of speciﬁed cardinality m , such that the associated Q m is close to optimal, and it is a consistent approximation of P .
This is achieved by adopting a kernel Stein discrepancy (KSD) as the optimality criterion.
The minimisation of KSD is performed using a greedy sequential algorithm and we report results that study the interplay of the greedy algorithm with the randomness inherent to the MCMC output.
This work lies in an active area of research that attempts to cast discrete approximation of a posterior distribution as an optimisation problem.
Liu and Lee (2017) considered the use of KSD to optimally weight an arbitrary set ( X i ) ni
=1 ⊂ R d of states in an manner loosely analogous to importance sampling, and Hodgkinson et al.
(2020) studied the eﬀect of applying this to the output of MCMC.
Outside of the MCMC framework, various criteria have been proposed to capture how well a discrete measure approximates P , so that global numerical optimisation methods can be used to arrive at a suitable point set: this is the case for the “minimum energy design” (MED) of Joseph et al.
(2015, 2019), the “Stein points” (SP) of Chen et al.
and “Support Points” of Mak and Joseph (2018).
The reliance on global optimisation renders the theoretical analysis of these methods diﬃcult, and, to obviate this, Chen et al.
(2019) considered using Markov chains, in the context of SP, to approximately perform numerical optimisation, allowing a tractable analytic treatment.
Stein Thinning diﬀers from the contributions cited above, in one or more of the following senses: (1) It is primarily intended for compression, not re-weighting, of the MCMC output.
(2) Continuous optimisation is avoided by working directly on the MCMC output.
Finite sample size error bounds and consistency of the algorithm can be established.
In the remainder we report Stein Thinning , together with a theoretical and empirical assessment.
Notation: Throughout we assume that the target distribution P admits a positive and continuously diﬀerentiable density p on R d .
Let Q n ⇒ P denote weak convergence of a sequence ( Q n ) of measures to P .
H ( k ) denotes a reproducing kernel Hilbert space (RKHS) of functions on R d , namely a Hilbert space, equipped with a function k : R d × R d → R , called a kernel , such that ∀ x ∈ R d k ( · , x ) ∈ H ( k ) and ∀ x ∈ R d , h ∈ H ( k ), h ( x )
= (cid:104) h, k ( · , x ) (cid:105) H ( k ) , where (cid:104)· , ·(cid:105)
H ( k ) denotes the inner product in H ( k ) and the induced norm is denoted (cid:107)·(cid:107) H ( k ) .
On the other hand, (cid:104)· , ·(cid:105) is the standard Euclidean inner product, with induced Euclidean norm (cid:107) x (cid:107) =
(cid:104) x, x (cid:105) 1 / 2 .
We denote ∇· the divergence operator in R d .
2. Methods
In Section 2.1 we recall KSD and in Section 2.2 we summarize our Stein Thinning algorithm for minimisation of KSD.
2.1. Kernel Stein Discrepancy
To select states from the MCMC output we require a notion of optimal approximation for probability distributions.
To this end, recall the notion of a Stein discrepancy , proposed in Gorham and Mackey (2015).
This was based on Stein’s method (Stein, 1972), which consists of ﬁnding a diﬀerential operator A P , depending on P and acting on d -dimensional vector ﬁelds on R d , and a set G of suﬃciently diﬀerentiable d -dimensional vector ﬁelds on
R d such that (cid:82) R d A P g d P = 0 for all g ∈ G .
The seminal paper of Gorham and Mackey (2015) introduced the Stein discrepancy
[Equation]
In this paper we focus on a particular form of (1) due to Liu et al.
Chwialkowski et al.
; Gorham and Mackey (2017), called a kernel Stein discrepancy (KSD).
In this case, A P is the Langevin Stein operator A P g :
=
p − 1 ∇· ( pg ) derived in Gorham and Mackey (2015), where G := { g : R d → R d | (cid:80) di =1 (cid:107) g
i (cid:107) 2 H ( k ) ≤ 1 } is the unit ball in a Cartesian product of RKHS.
It follows from construction that the set A P G is the unit ball of another RKHS, denoted H ( k P ), whose elements h satisfy (cid:82) R d h d P = 0.
The explicit form of the kernel k P was derived in Oates et al.
and depends only on k and the gradient of the log-density of P .
KSD is recognised as a maximum mean discrepancy in H ( k P ) (Gretton et al., 2006) and is therefore fully characterised by the kernel k P ; we can thus unambiguously adopt the shorthand D k P ( Q ) for D A P G ( P, Q ).
Under suitable conditions on the kernel k , elements of H ( k P ) have zero mean with respect to P , and under further conditions on k and P , the KSD can be explicitly computed for an empirical measure Q n = 1 n (cid:80) ni
=1 δ ( x i ), supported on states x i ∈ R d :
[Equation]
Theoretical analysis in Gorham and Mackey (2017); Chen et al.
; Huggins and Mackey (2018)
; Chen et al.
; Hodgkinson et al.
; Gorham et al.
(2020) has established suﬃcient conditions for when D k P ( Q n ) → 0 implies Q n ⇒ P (a property called convergence control ).
Assuming that such conditions for convergence control are met, KSD is a suitable optimality criterion to minimise for the post-processing of MCMC output.
2.2. Greedy Minimisation of KSD
Continuous optimisation algorithms over R d were proposed for greedy minimisation of KSD in Chen et al.
(2018, 2019), wherein at iteration n a new state x n is appended to the current sequence ( x 1 , . . .
, x n − 1 ) by searching over a continuous domain in R d .
Our proposed method does not attempt to solve the fundamentally diﬃcult continuous optimisation problem for selection of the next point x n ∈ R d .
Instead, we exactly solve the discrete optimisation problem of selecting a suitable element x n from supplied MCMC output.
The method, called
Stein Thinning , is straight-forward to implement and is succinctly stated in Algorithm 1.
(The convention (cid:80)
i =1 = 0 is employed.)
The algorithm is illustrated on a simple bivariate Gaussian mixture in Figure 1. Observe that, when m is ﬁxed, the computational complexity of Stein Thinning is equal to O ( n ), identical to the linear complexity of MCMC, while it is higher when m is increasing with n , being at most O ( nm 2 ). Notice moreover that, in general, the indices in π need not be distinct. That is, Algorithm 1 may prefer to include a duplicate state rather than to include a state which is not useful for representing P .
Data: The output { x i } ni =1 from an MCMC method, a kernel k P for which convergence control holds, and a desired cardinality m ∈ N .
Result:
The indices π of a sequence ( x π ( j ) ) mj
=1 ⊂ { x i } ni =1 where π ∈ { 1 , . .
. , n } m .
for j = 1 , . . .
, m do π ( j ) ∈ arg min i =1 ,...,n k P ( x i , x i ) + j − 1 (cid:88) j (cid:48)
=
k P ( x π ( j (cid:48) ) , x i ) end
[Figure]
[Figure]
The suitability of KSD to quantify how well Q m , the empirical measure resulting from
Stein Thinning , approximates P is determined by the choice of the kernel k that is used in the deﬁnition of k P .
Several choices are possible and we follow Chen et al.
(2019), who advocated the pre-conditioned inverse multi-quadric kernel k ( x, y ) :
= (1 + (cid:107) Γ − 1 / 2 ( x − y )
(cid:107) 2 ) − 1 / 2 , where we take a data-driven approach to select Γ, based on the MCMC output.
To explore diﬀerent strategies for the selection of Γ, we focus on the following candidates: (1) Median ( med ) Γ = (cid:96)
I , where (cid:96) = med :=
median {(cid:107) X i − X j (cid:107) : 1 ≤ i < j ≤ n 0 } is the median Euclidean distance between states; (2) Scaled median ( sclmed ) Γ = (cid:96) 2 I , where (cid:96) =
med / (cid:112) log( m ), proposed in the context of Stein variational gradient descent in Liu and
Wang (2016); (3) Sample covariance ( smpcov ), Γ = ˆΣ provided that the sample covariance matrix ˆΣ of the MCMC output is non-singular.
The experiments in Section 4 shed light on which of these settings is the most eﬀective, but we acknowledge that many other settings could also be considered.
In what follows, we set n 0 = min( n, 10 3 ) for the med and sclmed settings, to avoid an O ( n 2 ) cost of computing (cid:96) , and otherwise set n 0 = n , so that the whole of the MCMC output is used to select Γ.
3. Theoretical Assessment
This section clariﬁes the limiting behaviour of Stein Thinning as m, n → ∞ , reporting convergence in mean-square of the KSD, along with a ﬁnite sample size error bound.
Let
V be a function V : R d → [1 , ∞ ) and, for a function f : R d → R and a measure µ on R d , let
[Equation]
Then a ψ -irreducible and aperiodic
Markov chain ( X i ) i ∈ N ⊂ X with n th step transition kernel P n is V -uniformly ergodic if and only if ∃ R ∈ [0 , ∞ ) , ρ ∈ (0 , 1) such that
[Equation]
for all initial states x ∈ R d and all n ∈ N .
The notation E will be used to denote expectation with respect to the law of the Markov chain.
The proof of this result can be found in Riabiz et al.
Theorem
Let P be a probability distribution on R d .
Let k P : R d × R d → R be a reproducing kernel with (cid:82) R d k P ( x, · )d P ( x ) = 0 for all x ∈ R d .
Consider a P -invariant, time-homogeneous, reversible Markov chain ( X i ) i ∈ N ⊂ R d generated using a V -uniformly ergodic transition kernel, such that (3) is satisﬁed with V ( x ) ≥ (cid:112) k P ( x, x ) for all x ∈ R d .
Suppose that, for some γ > 0 ,
[Equation]
Let π be an index sequence of length m produced by Algorithm 1 applied to the Markov chain output ( X i ) ni =1 .
Then, with C =
Rρ 1 − ρ , we have that
[Equation]
Observe that the upper bound in (4) is asymptotically minimised when (up to log factors)
m is proportional to n .
In practice we are interested in the case m (cid:28) n , so we may for example set m = (cid:98) n 1000 (cid:99) if we aim for substantial compression.
It is not claimed that the bound in (4) is tight and indeed empirical results in Section 4 endorse the use of Stein Thinning in the small m context.
4. Empirical Assessment
In this section we report a subset of the experiments presented in Riabiz et al. (2020), which compare the performance of Stein Thinning with traditional approaches to postprocessing MCMC output.
Our focus is a Bayesian inference problem deﬁned by a system of ordinary diﬀerential equations (ODEs), due to Goodwin (1965), that has a d = 4 dimensional parameter to be inferred.
Random Walk (RW) MCMC was implemented ( n = 10 6 ) and trace plots reveal a clear a burn-in period.
As in the rest of this abstract, we assume that compression of the MCMC output to m (cid:28) n states is desired.
We therefore compare the following methods: (1) the traditional approach, which estimates a burn-in period using the convergence diagnostics of Vats and Knudson (2018) based on L independent [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
chains 1 and discards it, followed by equidistant thinning of the remainder of the chain to obtain precisely m states; (2) the Support Points algorithm proposed in Mak and Joseph (2018), implemented in the R package support ; 2 (3) the Stein Thinning algorithm that we have proposed, with each of the kernel choices described in Section 2.2.
For a principled assessment, we computed two quantitative measures for how well the resulting empirical distributions approximate the posterior: (a) the energy distance (ED; Sz´ekely and Rizzo, 2004; Baringhaus and Franz, 2004), given up to an additive constant by
[Equation]
where we use the norm (cid:107) x (cid:107) Σ := (cid:107) Σ − 1 / 2 x (cid:107) induced by the covariance matrix of P , with both Σ and (5) being estimated from MCMC output, and (b) the KSD based on med , the simplest setting for Γ.
The results for ED and KSD are shown in Figure 2.
Stein Thinning based on med and sclmed performed at least as well as the other methods considered with respect to ED, even if ED is closely related to the quantity that the Support Points algorithm attempts to minimise (Mak and Joseph (2018) used the (cid:107) · (cid:107) norm in place of (cid:107) · (cid:107) Σ ).
The same holds for KSD, for all but the largest values of m considered.
The smpcov setting performed well for small m but for large m its performance degraded (this may be because in smpcov there are more degrees of freedom in Γ that must be estimated).
Note that neither ED nor KSD values tend to 0 as m → ∞ in this experiment, since the number n of
MCMC iterations was ﬁxed.
5. Conclusion
In this abstract we summarized Stein Thinning , a method that seeks a subset of the
MCMC output, of ﬁxed cardinality, such that the associated empirical approximation is close to optimal.
Our theoretical analysis is able to handle the eﬀect of the post-processing procedure jointly with the randomness involved in simulating from the Markov chain, such that ﬁnite sample size error bounds and the consistency of the overall estimator can be established