Title: Kernel Stein Discrepancy thinning: a theoretical perspective of pathologies and a practical fix with regularization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Bayesian inference

Markov chain Monte Carlo

kernelized Stein discrepancy

Stein thinning

kernel methods
Abstract: Stein thinning is a promising algorithm proposed by (Riabiz et al., 2022) for post-processing outputs of Markov chain Monte Carlo (MCMC). The main principle is to greedily minimize the kernelized Stein discrepancy (KSD), which only requires the gradient of the log-target distribution, and is thus well-suited for Bayesian inference. The main advantages of Stein thinning are the automatic remove of the burn-in period, the correction of the bias introduced by recent MCMC algorithms, and the asymptotic properties of convergence towards the target distribution. Nevertheless, Stein thinning suffers from several empirical pathologies, which may result in poor approximations, as observed in the literature. In this article, we conduct a theoretical analysis of these pathologies, to clearly identify the mechanisms at stake, and suggest improved strategies. Then, we introduce the regularized Stein thinning algorithm to alleviate the identified pathologies. Finally, theoretical guarantees and extensive experiments show the high efficiency of the proposed algorithm. An implementation of regularized Stein thinning as the kernax library in python and JAX is available at https://gitlab.com/drti/kernax.

Kernel Stein Discrepancy thinning: a theoretical perspective of pathologies and a practical fix with regularization

Abstract

Stein thinning is a promising algorithm proposed by Riabiz et al. [2022] for postprocessing outputs of Markov chain Monte Carlo (MCMC). The main principle is to greedily minimize the kernelized Stein discrepancy (KSD), which only requires the gradient of the log-target distribution, and is thus well-suited for Bayesian inference. The main advantages of Stein thinning are the automatic remove of the burn-in period, the correction of the bias introduced by recent MCMC algorithms, and the asymptotic properties of convergence towards the target distribution. Nevertheless, Stein thinning suffers from several empirical pathologies, which may result in poor approximations, as observed in the literature. In this article, we conduct a theoretical analysis of these pathologies, to clearly identify the mechanisms at stake, and suggest improved strategies. Then, we introduce the regularized Stein thinning algorithm to alleviate the identified pathologies. Finally, theoretical guarantees and extensive experiments show the high efficiency of the proposed algorithm. An implementation of regularized Stein thinning as the kernax library in python and JAX is available at https://gitlab.com/drti/kernax.

1 Introduction

Bayesian inference is a powerful approach to solve statistical tasks, and is especially efficient to incorporate prior expert knowledge of the studied system, or to provide uncertainties of the estimated quantities. Bayesian methods have thus demonstrated a high empirical performance for a wide range of applications, in particular in the fields of physics and computational biology, to just name a few. However, the Bayesian framework often leads to the evaluation of expectations with respect to a posterior distribution, which is not tractable [Green et al., 2015], except in the specific case of conjugate prior distribution and likelihood, which hardly occurs in practice. To overcome this issue, Markov chain Monte Carlo (MCMC) is one of the most commonly used computational methods to estimate these integrals. Indeed, MCMC algorithms iteratively generate a sample, which follows the targeted posterior distribution, as the Markov chain converges to its stationary state [Robert and Casella, 1999, Brooks et al., 2011]. Consequently, the quality of the resulting estimates strongly depends on the convergence of the MCMC and how its output is post-processed. Standard postprocessing procedures of MCMC outputs consist in removing the first iterations, called the burn-in period, and thinning the Markov chain with a constant frequency. Burn-in removal aims at reducing the bias introduced by the random initialization of the Markov chain. The ˆ R convergence diagnosis of Gelman et al. [1995] is, for instance, a well known method for determining the burn-in period. On the other hand, thinning the Markov chain allows for compressing the MCMC output and may also reduce the correlation between the iteratively selected points. More recently, promising kernel-based

1

procedures were proposed to automatically remove the burn-in period, compress the output, and reduce the asymptotic bias [South et al., 2022]. These approaches consist in minimizing a kernelbased discrepancy measure D ( P , Q m ) between the empirical distribution Q m of a subsample of the MCMCoutput of size m , and the target distribution P . In this respect, minimization of the maximum mean discrepancy (MMD) was investigated by several authors, but these strategies require the full knowledge of the target distribution P , whose density is not tractable in non-conjugate Bayesian inference.

Based on the previous works of Chen et al. [2018] and Chen et al. [2019], Riabiz et al. [2022] propose to minimize the kernelized Stein discrepancy (KSD), to design an efficient kernel-based algorithm to thin MCMC outputs in a non-tractable Bayesian setting. The KSD [Liu et al., 2016] is a score-based discrepancy measure, i.e. , it only requires the knowledge of the score function of the target P , which is readily available in our Bayesian framework. Importantly, Gorham and Mackey [2017] showed that under suitable mild conditions, the KSD enjoys good convergence properties. More precisely, the KSD is a valid distance to detect samples drawn form the target distribution, provided that the sample size is large enough. Therefore, KSD thinning is a highly promising tool for post-processing and measuring the quality of MCMC outputs. This article thus focuses on the Stein thinning algorithm proposed by Riabiz et al. [2022], which consists in selecting m points amongst the n iterations of the MCMC output, by greedily minimizing the KSD distance. Thanks to the convergence properties of the KSD, the empirical measure of the selected points weakly converges towards the posterior law P . However, on the practical side, several articles [Wenliang and Kanagawa, 2020, Korba et al., 2021] have noticed empirical limitations of KSD-based sampling algorithms, especially for multimodal target distributions. In fact, these limitations happen to be quite problematic, even in simple experiments, and have been slightly overlooked in the literature so far, in our opinion. Therefore, this article first focuses on the analysis of KSD pathologies in Section 2, taking both an empirical and theoretical point of view. Then, we propose strategies to mitigate the identified problems, and introduce the regularized Stein thinning in Section 3. We show the efficiency of our algorithm through both a theoretical analysis and extensive experiments in Section 4. Notice that proofs and additional experiments are gathered in Appendices 1 -7 in the Supplementary Material. In the remaining of this initial section, we mathematically formalize the KSD distance and the associated Stein thinning algorithm.

Kernelized Stein discrepancy. Kernelized Stein discrepancy was independently introduced by Chwialkowski et al. [2016], Liu et al. [2016], Gorham and Mackey [2017] as a promising tool for measuring dissimilarities between two distributions P and Q on R d with d ≥ 1 , whenever P admits a continuously differentiable density p , and the normalization constant of p is not tractable. Let k : R d × R d → R be a positive semi-definite kernel and let H ( k ) be the associated reproducing kernel Hilbert space (RKHS) with inner product 〈· , ·〉 H ( k ) and norm ‖ · ‖ H ( k ) . Kernelized Stein discrepancy belongs to the family of maximum mean discrepancies (MMD) [Gretton et al., 2006] defined as

where X ∼ P , Z ∼ Q . If the kernel k is characteristic, then the MMD is a distance between probability distributions. In practice, the MMD may not be computable as it involves mathematical expectations with respect to P , whose density is not tractable. To circumvent this issue, Gorham and Mackey [2015] proposed the Stein discrepancy which relies on Stein's method [Stein, 1972]. It consists in defining an operator T p that maps functions g : R d → R d to real-valued functions such that E [ T p g ( X )] = 0 , with X ∼ P , for all g in G ( k ) = { g : R d → R d : ∑ d i =1 ‖ g i ‖ 2 H ( k ) ≤ 1 } . The probability measure P on R d is assumed to admit a continuously differentiable Lebesgue density p ∈ C 1 ( R d ) , such that E [ ‖∇ log p ( X ) ‖ 2 2 ] < ∞ . The Stein discrepancy is then defined as SD( P , Q ) = sup g ∈G ( k ) | E [( T p g )( Z )] | , where Z ∼ Q . If the Stein operator T p is chosen as the Langevin operator ( T p g )( x ) = 〈 g ( x ) , ∇ log p ( x ) 〉 + 〈∇ , g ( x ) 〉 , then Stein's discrepancy has a closed-form expression known as kernelized Stein discrepancy [Chwialkowski et al., 2016, Liu et al., 2016], KSD 2 ( P , Q ) = E [ k p ( Z , Z ′ )] , where Z ∼ Q , Z ′ ∼ Q , and k p denotes the Langevin Stein kernel defined from the score function s p ( x ) = ∇ log p ( x ) for x , x ′ ∈ R d , as

The main advantage of the KSD is that it only requires the knowledge of the score function, and does not involve any integration with respect to P . Gorham and Mackey [2017] also established

convergence guaranties when the kernel k is chosen as the inverse multi-quadratic (IMQ) kernel function k ( x , x ′ ) = ( c + ‖ x -x ′ ‖ 2 Γ ) -β with c > 0 , β ∈ (0 , 1) , the positive definite matrix Γ is the identity matrix, and the density p is distantly dissipative as defined below. Log-concave distributions outside of a compact set are a typical example of such probability densities.

Stein thinning algorithm. Let P be a target probability measure that admits density p , and let { x i } n i =1 ⊂ R d be a MCMC output. The Stein thinning algorithm [Riabiz et al., 2022] selects m ≤ n particles x π 1 , . . . , x π m by greedily minimizing the kernelized Stein discrepancy. Given t -1 < m particles x π 1 , . . . , x π t -1 , the t -th particle is defined as

where the KSD of an empirical distribution has been used to simplify the objective function. The kernel function k is usually chosen as the IMQ kernel function, defined above, for both its good theoretical properties and empirical efficiency. Indeed, several articles [Chen et al., 2018, Riabiz et al., 2022] have led extensive experiments to show the better practical performance of the IMQ kernel over other choices. Also notice that the bandwidth parameter glyph[lscript] is quite influential on the algorithm performance, but happens to be very difficult to tune, as highlighted by Chopin and Ducrocq [2021]. Indeed, since the normalization constant of the target distribution is unknown, no additional metric is available to assess the precise performance of the thinning procedure when glyph[lscript] varies. Furthermore, the sample quality output by Stein thinning varies in an erratic fashion with respect to glyph[lscript] , making the design of heuristic procedures for the choice of glyph[lscript] notoriously difficult. Following the literature recommendations [Riabiz et al., 2022], we use the median heuristic to set glyph[lscript] in our experiments, and refer to Garreau et al. [2017] for an extensive analysis of this approach for kernel methods.

2 Analysis of KSD Pathologies

Although kernelized Stein discrepancy is a highly promising approach to thin MCMC outputs, several empirical studies have highlighted that KSD-based algorithms may suffer from strong pathologies in simple experiments [Wenliang and Kanagawa, 2020, Korba et al., 2021, Riabiz et al., 2022, Liu et al., 2023]. The most established KSD pathology is that Stein thinning ignores the weights of distant modes of the target distribution, leading to the selection of samples of poor quality by Stein thinning. This problem, called Pathology I throughout the article, is analyzed in Subsection 2.1. Additionally, Korba et al. [2021] also notice that KSD thinning may result in samples concentrated in regions of low probability of p . As opposed to Pathology I, the mechanism leading to this problematic behavior is not well understood in the literature, to our best knowledge. Subsection 2.2 is thus dedicated to the theoretical characterization and illustration of Pathology II. Throughout the article, we illustrate KSD thinning using the running example of a Gaussian mixture, defined in Example 1 below, where initial particles are directly sampled from p to better highlight pathologies. We will come back to the thinning of MCMC outputs in detail in Section 4.

Example 1. Let the density p be a Gaussian mixture model of two components, respectively centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , of weights w and 1 -w , and of variance σ 2 I d . The initial particles { x i } n i =1 are drawn from p . The KSD thinning algorithm selects m<n points to approximate p .

2.1 Pathology I: mode proportion blindness

We first focus on Pathology I, which states that Stein thinning is blind to the relative weights of multiple distant modes of a target distribution. Indeed, Wenliang and Kanagawa [2020] show that the score s p is insensitive to distant mode weights. Consequently, the KSD distance is unable to properly identify samples with different weights than those of the target, in finite sample settings, as long as samples are accurately distributed within each mode. To be more specific, we illustrate this pathology with our Example 1 of a Gaussian mixture in dimension 2 . Weset µ = 3 and σ = 1 to enforce the two modes to be well separated, and take an unbalanced proportion w = 0 . 2 for the left mode, and 1 -w = 0 . 8 for the right mode. We generate n = 3000 observations and select m = 300 particles with Stein thinning. Clearly, the red selected sample displayed in Figure 1 has wrong proportions, with about half of the

particles in each mode, instead of the expected 20 -80 %, reflected by the initial black particles sampled from p . More precisely, over 100 repetitions of the Stein thinning algorithm, we obtain an average proportion of 0 . 53 particles in the left mode, with a standard deviation of 0 . 08 across the 100 runs.

Although Wenliang and Kanagawa [2020] clearly show that the KSD distance is insensitive to the mode weights in the specific case of Gaussian mixtures, the mechanism leading to the selection of about half of the particles in each mode by Stein thinning, as in Example 1, remains unexplained in the literature, to our best knowledge. Therefore, we conduct a theoretical analysis in the general case of any mixture distribution with two distant modes, stated in Assumption 2.1 below. For the sake of clarity, we only study the case of a number of modes of two, without loss of generality. Importantly, notice that a finite sample drawn from a distribution with distant modes, takes the form of clusters of particles around each mode, as illustrated in Figure 1. Then, Stein thinning selects particles among these clusters to approximate p , and these particles define an empirical law of a density q with a compact support around each mode. Wenliang and Kanagawa [2020] explain that the score s p is especially insensitive to the mode weights in these compact areas around modes, which is the root cause of the generation of samples with wrong proportions, as in Figure 1. Therefore, Assumption 2.1 below defines this observed setting, required to have Pathology I to occur, where density q has compact supports around each distant mode. Additionally,

5.0

2.5

0.0

-2.5

-5.0

0

5

c(1)

we also need to formalize Assumption 2.2, which tells that the distributions of the two modes of the mixture q have a close KSD distance with respect to the target p . In particular, this assumption can be easily verified when both p and q have symmetric mode distributions, since the KSD distance is insensitive to the weights of p .

glyph[negationslash]

Assumption 2.1 (Distant bimodal mixture distributions) . Let p and q be two mixture distributions in R d , made of two modes centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , with µ > 0 . The distribution of each mode of p ∈ C 1 ( R d ) has R d as support, whereas each mode distribution of q have a compact support, included in a ball of radius r > 0 , with r < µ . The left mode of p has weight w p = 1 / 2 , and the right mode has weight 1 -w p . Similarly, w and 1 -w are the mode weights of q . Let Q L and Q R be the probability measures that respectively admit the density of the left and right modes of q , and P and Q w be also the probability laws for p and q .

Assumption 2.2. For distant bimodal mixture distributions q and p satisfying Assumption 2.1, and for η ∈ (0 , 1) , we have ∣ ∣ KSD 2 ( P , Q L ) / KSD 2 ( P , Q R ) -1 ∣ ∣ < η .

Theorem 2.3. Let k p be the Stein kernel associated with the radial kernel k ( x , x ′ ) = φ ( ‖ x -x ′ ‖ 2 /glyph[lscript] ) , where x , x ′ ∈ R d , glyph[lscript] > 0 , and φ ∈ C 2 ( R ) , such that φ ( z ) → 0 , φ ′ ( z ) → 0 , and φ ′′ ( z ) → 0 for z →∞ . Let p and q be two bimodal mixture distributions satisfying Assumptions 2.1 and 2.2, for any η ∈ (0 , 1) . We define w glyph[star] as the optimal mixture weight of q with respect to the KSD distance, i.e., w glyph[star] = argmin w ∈ [0 , 1] KSD( P , Q w ) . Then, for µ large enough, we have ∣ ∣ w glyph[star] -1 2 ∣ ∣ < η 2(1 -η ) .

Theorem 2.3, proved in Appendix B, states that the weight w glyph[star] of the optimal mixture q , which minimizes the KSD distance to the target p , is close to 1 / 2 regardless of the true target weight w p , whenever the distributions of the two modes of the mixture q have a close KSD distance to p , and provided that the two modes are distant. In particular, this is the case in the experiment of Example 1 and Figure 1, where the two modes are symmetric and well separated. Additionally, a more specific empirical illustration of Theorem 2.3 can be found in Appendix A.1. In Section 3, we will propose strategies improving Stein thinning to recover samples with accurate mode proportions.

2.2 Pathology II: spurious minimum

The core of this section is dedicated to the theoretical characterization of Pathology II. We first need to introduce additional notations to formalize our analysis. We thus define M s 0 , the region of the input space where the score norm is lower than the threshold s 0 ≥ 0 , formally M s 0 = { x ∈

5.0

2.0

2.5

1.5

2

0.0

1.0

S

-2.5

Io

0.5

-5.0

0.0

0

5

c(1)

-2

0

2

+(1

R d : ‖ s p ( x ) ‖ 2 ≤ s 0 } . We also introduce an independent and identically distributed (iid) sample X 1 , . . . , X m of P , with P m the associated empirical measure for a positive integer m , and X ( j ) the j -th component of X . Then, Theorem 2.4 below shows that samples concentrated in regions of the input space where the norm of the score is low, have smaller KSD than samples drawn from the true target distribution p , for small sample sizes. Additionally, the score norm is low around stationary points of p , including local minimum and saddle points, as shown in Corollary 2.5 below. However, samples concentrated at local minimum of p are bad approximations of the target distribution by definition. Therefore, pathological samples may be generated by Stein thinning, which minimizes the empirical KSD, and thus explains Pathology II observed by Korba et al. [2021], and shown in Figure 2. For the sake of clarity, we formalize our result for the IMQ kernel used in practice, and set c = 1 without loss of generality, since it is equivalent to tune c or glyph[lscript] in the Stein thinning algorithm.

Theorem 2.4 (KSD spurious minimum) . Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . Let { x i } i m =1 ⊂ M s 0 = { x ∈ R d : ‖ s p ( x ) ‖ 2 ≤ s 0 } be a fixed set of points of empirical measure Q m = 1 m ∑ m i =1 δ ( x i ) , with s 0 ≥ 0 and m ≥ 2 . We have KSD 2 ( P , Q m ) < E [KSD 2 ( P , P m ) ] , if the score threshold s 0 and the sample size m are small enough to satisfy m< 1 + ( E [ ‖ s p ( X ) ‖ 2 2 ] -s 2 0 ) / (2 βd/glyph[lscript] 2 +2 βs 0 /glyph[lscript] + s 2 0 ) .

Corollary 2.5 (Low KSD samples at density minimum) . Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . Let p be a density with at least one local minimum or saddle point. For m ≥ 2 , if { x i } i m =1 ⊂ R d is a set of points, all located at local minimum or saddle points of p , then we have KSD 2 ( P , Q m ) < E [KSD 2 ( P , P m ) ] , if m< 1 + glyph[lscript] 2 2 βd E [ ‖ s p ( X ) ‖ 2 2 ] .

The proofs of Theorem 2.4 and Corollary 2.5, reported in Appendix C, are built on the idea that the KSD of the empirical law of { x i } n i =1 , has a bias of the form ∑ m i =1 ‖ s p ( x i ) ‖ 2 2 /m 2 . Consequently, when m is small, the bias has a strong influence on KSD estimates, which favor samples concentrated in regions of low score norm, as stationary points of p . This mechanism is illustrated in Figure 2 and Corollary 2.6 for Gaussian mixtures. In this case, Stein thinning aligns a large number of particles around the line of saddle points defined by x (1) = 0 , an area of low probability of the targeted mixture distribution, because of the variations of the score function, if the sample size m is small enough. From another perspective, for any sample size m , it exists a Gaussian mixture with µ/σ large enough, such that Pathology II occurs. Therefore, Pathology II can appear for arbitrarily large samples m , depending on the target distribution properties.

Corollary 2.6 (KSD spurious minimum for Gaussian mixtures) . Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . Let the density p be a Gaussian mixture model of two components with equal weights, respectively centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , of variance σ 2 I d , and let ν = µ/σ . If ν > 1 and 0 ≤ s 0 < [ ν √ ν 2 -1 -ln( ν + √ ν 2 -1) ] /µ , then for any { x i } i m =1 ⊂ M s 0 of empirical measure Q m , we have

,

(i) KSD 2 ( P , Q m ) < E [KSD 2 ( P , P m ) ] if m and s 0 satisfy m< 1 + E [ ‖ s p ( X ) ‖ 2 2 ] -s 2 0 2 βd/glyph[lscript] 2 +2 βs 0 /glyph[lscript] + s 2 0 , (ii) there exists three disjoint intervals I -µ , I 0 , I µ ⊂ R , respectively centered around -µ , 0 , and µ such that x (1) 1 , . . . , x (1) m ∈ I -µ ∪ I 0 ∪ I µ .

3 Regularized Stein Thinning

Stein thinning suffers from two main pathologies, analyzed in Section 2. In a word, Pathology I comes from the insensitivity of the score to the relative weights of distant modes, whereas Pathology II originates from the variations of the score norm, which do not differentiate local minimum from local maximum of the target distribution. We propose to regularize the KSD distance to fix these two problems, using terms that are highly sensitive to the type of stationary point and the relative weights of modes. The proposed algorithm is first introduced in Subsection 3.1, then theoretical properties are discussed in Subsection 3.2, and finally the good empirical performance will be shown in Section 4.

3.1 Algorithm

Entropic regularization. In order to compensate the blindness of the KSD to mode proportions in multimodal distributions, we introduce the following entropic regularized KSD, denoted by KSD λ , and defined as KSD 2 λ ( P , Q ) = E [ k p ( Z , Z ′ )] -λ E [log( p ( Z ))] , where Z and Z ′ have probability law Q , and P admits the density p . In our Bayesian setting, E [log( p ( Z ))] is known up to an additive constant since the normalization factor of p is not tractable. However, it is possible to use KSD 2 λ ( P , Q ) as the objective function of the Stein thinning algorithm, as the greedy selection of particles to optimize this quantity does not rely on the unknown additive constant. The main idea of this entropic regularization is that -log( p ( x )) takes higher values in modes of smaller probability, and therefore provides the relative mode weight information, which is missing in the KSD distance. More precisely, modes with smaller weights take smaller density values, and are therefore more penalized than modes of higher weights. Therefore, with such entropic penalization, regularized Stein thinning tends to select particles in modes of higher weights more frequently than in modes of smaller weights, and we recover appropriate proportions.

Laplacian correction. Chen et al. [2018] and Riabiz et al. [2022] have noticed that the term k p ( x i , x i ) , which naturally appears in the empirical kernelized Stein discrepancy with the Langevin operator, can be interpreted as a regularization term. For example, Stein thinning does not select particles in the burn-in period of an MCMC output thanks to this regularization. However, this term k p ( x i , x i ) is also responsible for Pathology II, of samples concentrated in stationary points of p , as shown in Theorem 2.4. Therefore, we add a second regularization term to compensate the weaknesses of k p ( x i , x i ) , by penalizing particles located at local minimum and saddle points of the density p . Such points are located in areas of convexity of the target distribution, which can thus be detected with the positive values of the Laplacian of the density. Therefore, using the truncated Laplacian operator ∆ + f ( x ) = ∑ d j =1 ( ∂ 2 f ( x ) /∂x ( j )2 ) + for a function f ∈ C 2 ( R d ) , we propose the L-KSD estimate with a Laplacian correction for densities p ∈ C 2 ( R d ) , defined by glyph[negationslash]

Regularized Stein thinning. Overall, we obtain the following estimate for the entropic regularized KSD with Laplacian correction L-KSD 2 λ ( P , Q m ) = L-KSD 2 ( P , Q m ) -λ m ∑ m i =1 log( p ( x i )) . Then, at each iteration t ∈ { 1 , . . . , m } , the regularized Stein thinning selects the particle index π t ∈ { 1 , . . . , n } to greedily minimize

Finally, Figure 3 illustrates the performance of regularized Stein thinning to fix the two pathologies analyzed in Section 2, in the case of Example 1 with Gaussian mixtures. Indeed, the top panel of Figure 3 shows that the majority of particles are selected in the right mode, as expected from the target distribution with w = 0 . 2 . More precisely, an average proportion of 0 . 11 of the particles are located in the left mode over 100 repetitions of the procedure ( 0 . 89 in the right mode), with a standard deviation of 0 . 03 . For the value choice of λ , we refer to the next subsection and the experimental Section 4. On the bottom panel of Figure 3, we observe that no particle is now selected on the line x (1) = 0 , as expected from the target Gaussian mixture distribution.

Remark 3.1. The truncated Laplacian operator is simply given by the trace of the Hessian matrix, where negative components are set to 0 . It follows that the computational cost the regularized algorithm is similar to the original Stein thinning.

Remark 3.2. The Laplacian correction of k p introduces second-order derivatives of p in the Stein discrepancy, and therefore enables to differentiate local minimum and saddle points of the density p from its local maximum. A natural approach to introduce second-order derivatives of p in KSD estimates, is to define the Stein discrepancy using second-order operators. A Laplacian Stein operator [Oates et al., 2017] is derived in Appendix G, but experiments show that this strategy is not efficient to fix Pathologies I & II.

Regularized Stein Thinning

Regularized Stein Thinning

5.0

5.0

2.5

2.5

0.0

0.0

2.5

2.5

-5.0

-5.0

0

0

c(1)

c(1)

3.2 Theoretical properties

This subsection is dedicated to the theoretical analysis of regularized Stein thinning. First, we show that the proposed algorithm now enjoys good properties regarding Pathologies I and II, and thus mitigates the identified problems of the original Stein thinning. Secondly, we extend the convergence analysis of Riabiz et al. [2022] for the post-treatment of MCMC output, to show the convergence of the empirical law output by regularized Stein thinning towards the target probability measure.

glyph[negationslash]

Entropic regularization. In the previous section, Theorem 2.3 highlights how Pathology I of mode proportion blindness originates from the score insensitivity to mode weights. On the other hand, the entropic regularization is directly built on the target density, and therefore strongly depends on the mode weights. In the same setting of Assumption 2.1, required for Pathology I to occur with the original algorithm, the following Theorem 3.3 shows that the entropic regularized KSD is minimized for the appropriate target weight, with the suitable regularization strength λ . Notice that Theorem 3.3, proved in Appendix D, is valid if E [log( p ( Z L ))] = E [log( p ( Z R ))] with Z L ∼ Q L and Z R ∼ Q R , otherwise the impact of the entropic regularization on w glyph[star] λ vanishes. However, as w p = 1 / 2 is required in Assumption 2.1 for Pathology I to occur, p is asymmetric, and E [log( p ( Z L ))] = E [log( p ( Z R ))] is only possible in very specific cases. Theorem 3.3 clearly shows that the regularized entropic KSD is sensitive to the weights of distant modes. Efficient strategies to choose the regularization strength will be first discussed in the asymptotic analysis below, and then in the experiments of the next section.

glyph[negationslash]

glyph[negationslash]

Theorem 3.3. Let k p be the Stein kernel associated with the radial kernel k ( x , x ′ ) = φ ( ‖ x -x ′ ‖ 2 /glyph[lscript] ) , where x , x ′ ∈ R d , glyph[lscript] > 0 , and φ ∈ C 2 ( R ) . Let p and q be two bimodal mixture distributions satisfying Assumption 2.1. We define w glyph[star] λ as the optimal mixture weight of q with respect to the entropic regularized KSD distance, i.e., w glyph[star] λ = argmin w ∈ [0 , 1] KSD λ ( P , Q w ) . If E [log( p ( Z L ))] = E [log( p ( Z R ))]

Laplacian correction. First, we stress that the L-KSD is a strongly consistent estimate of the KSD distance, where the proof follows from the law of large numbers. Therefore, the Laplacian correction introduced in the L-KSD estimate does not undermine the good asymptotic properties of the KSD distance. Secondly, the following theorem shows that samples concentrated in local minimum or saddle points of the target distribution and of low density values, are well identified by the L-KSD as samples of worse quality than those truly sampled from the target. Consequently, the Laplacian correction fixes Pathology II, previously formalized in Theorem 2.4.

Theorem 3.4. Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . For m ≥ 2 , let { x i } i m =1 ⊂ R d be a set of points located at x 0 , a local minimum or saddle point of p , and of empirical measure Q m . Then, we have L-KSD 2 ( P , Q m ) > E [ L-KSD 2 ( P , P m ) ] , if the density at x 0 satisfies p ( x 0 ) < ∆ + p ( x 0 ) / ( E [ ‖ s p ( X ) ‖ 2 2 ] + E [∆ + log p ( X )] ) .

Convergence of regularized Stein thinning. While regularized Stein thinning fixes finite sample size pathologies, the asymptotic properties of Stein thinning are also preserved. Indeed, if the initial set of particles is drawn from a different distribution than the target using a Markov chain Monte Carlo, Theorem 3.6 states that the empirical measure of the sample obtained with regularized Stein Thinning, converges towards the target measure P , and thus extends the results of Riabiz et al. [2022]. Notice that the weak convergence of a sequence of probability measure is denoted by ⇒ , and that distantly dissipative distributions are defined in Definition 1.1. The required assumption below, essentially states mild integrability conditions, and enforces that the MCMC output is not too far from a sample drawn from p -see Appendix F for additional details.

Assumption 3.5. Let Q be a probability distribution on R d , such that P is absolutely continuous with respect to Q . Let { Z i } i ∈ N ⊂ R d be a Q -invariant, time-homogeneous Markov chain, generated using a V -uniformly ergodic transition kernel, such that V ( x ) ≥ d P d Q √ 2 βd/glyph[lscript] 2 + ‖ s p ( x ) ‖ 2 2 . Suppose +

Theorem 3.6. Let P be a distantly dissipative probability measure, that admits the density p ∈ C 2 ( R d ) , k p be the Stein kernel associated with the IMQ kernel where glyph[lscript],c > 0 , β ∈ (0 , 1) . Let { Z i } i ∈ N ⊂ R d be a Markov chain satisfying Assumption 3.5, π be the index sequence of length m n generated by regularized Stein thinning, and Q m n be the empirical measure of { Z π i } m n i =1 . If log( n ) α < m n < n , with any α > 1 , and λ m n = o (log( m n ) /m n ) , then we have almost surely Q m n = ⇒ n →∞ P .

Theorem 3.6, proved in Appendix F, provides us with interesting insights about the entropic regularization strength λ . We already know that λ should be chosen with a rate at least as fast as O (1 /m ) , to avoid the introduction of a higher a bias in the L-KSD than the original KSD. Indeed, for a sample drawn from the true target distribution p , this bias E [ L-KSD 2 ( P , P m )] takes the form E [ k p ( X , X )] /m + E [∆ + log( p ( X ))] /m -λ E [log( p ( X ))] . Therefore, for slower rates of λ than O (1 /m ) , trivial samples concentrated at a local maximum of the target p , can have smaller L-KSD than samples drawn from p . Then, Theorem 3.6 states that, for our ultimate application of MCMC post-processing, the Stein thinning sample distribution converges towards the target for such λ rates of O (1 /m ) or faster. In practice, in our Bayesian setting, it is not possible to fine tune this parameter λ because no metric is available to assess the Stein thinning quality for various values of λ , as already mentioned in the case of the bandwidth parameter glyph[lscript] . In addition, we cannot theoretically determine which exact range of values of λ leads to good thinned samples in a finite sample regime. However, we will see in the experiments of the following section that both slower and faster λ rates than O (1 /m ) lead to samples of degraded quality. Therefore, we set λ = 1 /m in the regularized Stein thinning, to ensure good empirical performance and the algorithm convergence.

4 Empirical Assessment

This section shows how regularized Stein thinning outperforms the original algorithm through three batches of experiments: mixtures of standard distributions using exact or MCMC sampling, and Bayesian logistic regression on real datasets. For the experiments considered in Sections 4.2 and 4.3, two Metropolis-Hastings samplers are considered with the Metropolis-Adjusted Langevin Algorithm (MALA) and the No-U-Turn sampler (NUTS). We use the IMQ kernel with glyph[lscript] set with the median heuristic, β = 1 / 2 , and c = 1 , as recommended in Chen et al. [2018], Riabiz et al. [2022]. We also set the regularization parameter with the default value of λ = 1 /m . Notice that additional experiments are provided in Appendix A, and that the code is available at https://gitlab.com/drti/kernax.

When the target distribution is known, the efficiency of the Stein thinning algorithms are assessed by computing the MMD distance (see Equation (1)) between a large sample drawn from the target distribution and the thinned samples. More specifically, we use the following closed-form expression of the MMD [Gretton et al., 2006] with X , X ′ ∼ P and Z , Z ′ ∼ Q ,

where the kernel function k is chosen as the distance-induced kernel studied by Sejdinovic et al. [2013] and given by k ( x , x ′ ) = ‖ x ‖ 2 + ‖ x ′ ‖ 2 -‖ x -x ′ ‖ 2 , for x , x ′ ∈ R d . In this setting, the MMD reduces to the well known energy distance, as shown by Sejdinovic et al. [2013].

4.1 Gaussian mixtures with exact sampling

As a first batch of experiments, we build on Example 1 and consider more complicated twodimensional Gaussian mixtures to further illustrate the correction of Pathologies I & II. The first Gaussian mixture is made of four modes located at µ 1 = ( -3 , 3) , µ 2 = ( -3 , 3) , µ 3 = (3 , 3) , and µ 4 = (3 , -3) , and with weights w 1 = w 2 = 0 . 1 , and w 3 = w 4 = 0 . 4 , respectively. The second mixture is taken from [Qiu and Wang, 2023]. It is made of 6 equally weighted Gaussian distributions centered at µ i = (3 cos(2 π ( i -1) / 6) , 3 sin(2 π ( i -1) / 6)) , for i = 1 , . . . , 6 . For both experiments, we rely on exact Monte Carlo sampling to generate n = 3000 observations, and apply Stein thinning and its regularized variant to select m = 300 particles. The observed samples and the selected particles are shown in Figure 4. The first example shows that vanilla Stein thinning does not capture the right proportions, while the regularized variant appropriately penalizes modes with lower weights. The second example illustrates Pathology II, which is corrected by the regularized Stein thinning.

Regularized Stein Thinning

Regularized Stein Thinning

5.0

5.0

5.0

5.0

2.5

2.5

2.5

2.5

2

0.0

0.0

CN

0.0

N

0.0

2.5

2.5

2.5

2.5

5.0

-5.0

5.0

-5.0

0

L

0

5

0

0

5

2(1)

2（1)

a(1)

(1)

4.2 Banana-shaped and Gaussian mixtures with MCMC sampling

We consider a mixture of two distant modes of d -dimensional banana-shaped distributions with t-tails and unbalanced weights [Haario et al., 1999, Pompe et al., 2020], illustrated in Figure 5, and precisely defined in Appendix A.2. We sample this target banana mixture with both MALA and NUTS using three different step sizes ε and 10 5 iterations. The generated samples are post-processed with the Stein thinning and regularized Stein thinning algorithms, and their performances are compared with the MMD between the post-processed samples and large samples drawn from the known target banana mixture. This experiment is run for various thinning sizes m and dimensions d , with 20 repetitions to quantify uncertainties. The results obtained with the MALA sampler are shown in Figure 6: the regularized Stein thinning clearly generates samples of higher quality than the vanilla Stein thinning. Additionally, an example of post-processed MALA output is depicted in Figure 5, together with a heatmap of the Laplacian correction. On the left panel of Figure 5, we see that pathologies are especially strong in this experiment, with a large number of particles lying between the two modes in a region of low probability. On the right panels, we observe that regularized Stein thinning fix pathologies. Similar results were obtained with NUTS and are reported in Appendix A.2 for brevity. Next, we conduct the same experiments for a d -dimensional Gaussian mixture of four modes with different variances, as detailed in Appendix A.2. Again, Figure 6 shows the better

Regularized Stein Thinning

△+log(p)

30

20

10

0

-10

-20

0

20

20

0

20

20

0

20

c(1)

x(1)

1.00

1.00

2

2

Regularized Stein Thinning

MMD(P,Q

0.75

Regularized Stein Thinning

0.75

P

0.50

0.50

0.25

0.25

0

100

300

500

700

900

1100

102030405060708090100

0.00

100

300

500

700

900

1100

0.00

m

p

7m

p

performance of regularized Stein thinning. Besides, we take advantage of this last experiment to explore other regularization rates than our default λ = 1 /m . Figure 12 in Appendix A.2 shows that a slower rate of λ = 1 / log( m ) , which violates the convergence assumptions of Theorem 3.6, has significantly worse performance than the original Stein thinning. On the other hand, with a faster rate than 1 /m such as 1 /m 2 , the effect of the entropic regularization disappears, and we recover similar results than the original Stein thinning. This supports that the default value of λ = 1 /m is an efficient heuristic, since slower and faster rates of λ strongly degrade the algorithm performance.

4.3 Bayesian logistic regression

We now compare the two Stein thinning algorithms in the Bayesian logistic regression setting for binary classification, since such problem usually involves multimodal posterior-see, e.g. , Gershman et al. [2012], Liu and Wang [2016], Fong et al. [2019], Korba et al. [2021]. Given a dataset D N = { ( X i , Y i ) } N i =1 made of N pairs of features X i ∈ R d and labels Y i ∈ { 0 , 1 } , the probability that Y i is of class 1 is given by p ( Y i = 1 | X i , β , β 0 ) = 1 / (1 + exp( -β 0 -β T X i )) , for some parameters θ = ( β 0 , β ) ∈ R d +1 . The prior distributions of the weight vector θ is assumed to be Gaussian, p ( β ( j ) | γ ( j ) ) = N ( β ( j ) | 0 , 1 /γ ( j ) ) , and a Gamma prior with parameters ( a, b ) is chosen for the precision γ ( j ) . Following [Fong et al., 2019], the hyperparameters are chosen as a = b = 1 .

The posterior distribution of the weights θ is sampled with both MALA and NUTS using 48 independent chains, of respectively 10 4 and 10 5 iterations, and four step sizes ε are considered along with three thinning sizes m . Each MCMC sample is post-processed with the two Stein thinning algorithms. For a new input x glyph[star] , the resulting thinned samples are used to approximate the posterior predictive distribution p ( Y = 1 | x glyph[star] , D N ) , defined by ∫ p ( Y = 1 | x glyph[star] , θ ) p ( θ |D N ) d θ . The per-

formance of Stein thinning algorithms are assessed using the standard AUC metric for classification problems, estimated with 10 -fold cross-validation and 10 repetitions for uncertainties. Table 1 gathers the results for five public datasets from the UCI repository [Dua and Graff, 2017], and described in Appendix A.3, where the best AUC obtained for each algorithm over the four MCMC step sizes are reported. Clearly, regularized Stein thinning significantly improves the performance of Bayesian logistic regression.

5 Conclusion

Stein thinning has raised a high interest in recent years, as a powerful tool to post-process MCMC outputs, by the greedy minimization of the kernelized Stein discrepancy. Unfortunately, empirical studies have shown that KSD-based algorithms suffer from strong pathologies. We have conducted an in-depth theoretical analysis to identify the mechanisms at stake. From this understanding, we propose an improved Stein thinning algorithm relying on entropic regularization and Laplacian correction. This approach exhibits relevant theoretical properties regarding pathologies, as well as highly improved empirical performance. Finally, the analysis of these regularization terms for other types of KSD-based algorithms, such a KSD descent, seems a promising route for future work.

References

M Bishop and N.M. Nasrabadi. Pattern Recognition and Machine Learning , volume 4. Springer, 2006.

Brooks, A. Gelman, G. Jones, and X-L. Meng. Handbook of Markov Chain Monte Carlo . CRC press, 2011.

W.Y. Chen, L. Mackey, J. Gorham, F. Briol, and C. Oates. Stein points. In International Conference on Machine Learning , pages 844-853. PMLR, 2018.

W.Y. Chen, A. Barp, F-X. Briol, J. Gorham, M. Girolami, L. Mackey, and C. Oates. Stein point markov chain monte carlo. In International Conference on Machine Learning , pages 1011-1021. PMLR, 2019.

Chopin and G. Ducrocq. Fast compression of MCMC output. Entropy , 23:1017, 2021.

Chwialkowski, H. Strathmann, and A. Gretton. A kernel test of goodness of fit. In International Conference on Machine Learning , pages 2606-2615. PMLR, 2016.

Dua and C. Graff. UCI machine learning repository, 2017. URL http://archive.ics.uci. edu/ml .

Fong, S. Lyddon, and C. Holmes. Scalable nonparametric sampling from multimodal posteriors with the posterior bootstrap. In International Conference on Machine Learning , pages 1952-1962. PMLR, 2019.

Garreau, W. Jitkrittum, and M. Kanagawa. Large sample analysis of the median heuristic. arXiv preprint arXiv:1707.07269 , 2017.

Gelman, J.B. Carlin, H.S. Stern, and D.B. Rubin. Bayesian Data Analysis . Chapman and Hall/CRC, 1995.

Gershman, M. Hoffman, and D. Blei. Nonparametric variational inference. arXiv preprint arXiv:1206.4665 , 2012.

Gorham and L. Mackey. Measuring sample quality with stein's method. Advances in Neural Information Processing Systems , 28, 2015.

Gorham and L. Mackey. Measuring sample quality with kernels. In International Conference on Machine Learning , pages 1292-1301. PMLR, 2017.

P.J. Green, K. Łatuszy´ nski, M. Pereyra, and C.P. Robert. Bayesian computation: a summary of the current state, and samples backwards and forwards. Statistics and Computing , 25:835-862, 2015.

Gretton, K. Borgwardt, M. Rasch, B. Schölkopf, and A. Smola. A kernel method for the twosample-problem. Advances in Neural Information Processing Systems , 19, 2006.

Haario, E. Saksman, and J. Tamminen. Adaptive proposal distribution for random walk metropolis algorithm. Computational Statistics , 14:375-395, 1999.

Korba, P-C. Aubin-Frankowski, S. Majewski, and P. Ablin. Kernel stein discrepancy descent. In International Conference on Machine Learning , pages 5719-5730. PMLR, 2021.

Liu and D. Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. In Advances in neural information processing systems , pages 2378-2386, 2016.

Liu, J. Lee, and M. Jordan. A kernelized stein discrepancy for goodness-of-fit tests. In International Conference on Machine Learning , pages 276-284. PMLR, 2016.

Liu, A.B. Duncan, and A. Gandy. Using perturbation to improve goodness-of-fit tests based on kernelized stein discrepancy. arXiv preprint arXiv:2304.14762 , 2023.

Oates, A. Barp, and M. Girolami. Posterior integration on a riemannian manifold. arXiv preprint arXiv:1712.01793 , 2017.

Pompe, C. Holmes, and K. Łatuszy´ nski. A framework for adaptive mcmc targeting multimodal distributions. The Annals of Statistics , 48:2930-2952, 2020.

Yixuan Qiu and Xiao Wang. Efficient multimodal sampling via tempered distribution flow. Journal of the American Statistical Association , 0(0):1-15, 2023. doi: 10.1080/01621459.2023.2198059.

Riabiz, W.Y. Chen, J. Cockayne, P. Swietach, S.A. Niederer, L. Mackey, and C.J. Oates. Optimal thinning of MCMC output. Journal of the Royal Statistical Society: Series B , in press, 2022.

CP.P Robert and G. Casella. Monte Carlo Statistical Methods , volume 2. Springer, 1999.

Sejdinovic, B. Sriperumbudur, A. Gretton, and K. Fukumizu. Equivalence of distance-based and rkhs-based statistics in hypothesis testing. The Annals of Statistics , pages 2263-2291, 2013.

L.F. South, M. Riabiz, O. Teymur, and C.J. Oates. Postprocessing of mcmc. Annual Review of Statistics and Its Application , 9:529-555, 2022.

Stein. A bound for the error in the normal approximation to the distribution of a sum of dependent random variables. In Proceedings of the sixth Berkeley Symposium on Mathematical Statistics and Probability, Volume 2: Probability Theory , volume 6, pages 583-603. University of California Press, 1972.

L.K. Wenliang and H. Kanagawa. Blindness of score-based methods to isolated components and mixing proportions. arXiv preprint arXiv:2008.10087 , 2020.

Appendix

A Additional Experiments

A.1 Illustration of Theorem 2.3

To better illustrate Theorem 2.3, we run an additional experiment, where p is still defined as in Figure 7 from Example 2 recalled below, with unbalanced mode weights of 0 . 2 and 0 . 8 . The density q is distributed as p , but each mode is truncated outside a circle of two standard deviation radius, and q has weight w . Next, for various values of w ∈ [0 . 1 , 0 . 9] , we draw two samples of size n = 3000 from p and q , and compute KSD( P , Q w ) (with 30 repetitions for each w value). The result is displayed in Figure 8, and shows that the optimal weight is close to 1 / 2 , as predicted by Theorem 2.3, since | KSD 2 ( P , Q L ) / KSD 2 ( P , Q R ) -1 | is estimated as 0 . 01 in this case, implying that | w glyph[star] -1 / 2 | < 0 . 005 .

Example 2. Let the density p be a Gaussian mixture model of two components, respectively centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , of weights w and 1 -w , and of variance σ 2 I d . The initial particles { x i } n i =1 are drawn from p . The KSD thinning algorithm selects m<n points to approximate p .

Stein Thinning

5.0

2.5

2

0.0

-2.5

-5.0

0

5

c(1)

A.2 Gaussian and Banana-shaped Mixtures

This appendix gathers additional results and details for the Gaussian and banana-shaped mixtures, as well as the MMD distance used to evaluate thinning performance, and the regularization parameter λ .

Gaussian mixture. The second batch of experiments in Section 4 considers a d -dimensional Gaussian mixture of four modes of equal weight, with d ≥ 2 , illustrated in Figure 9. The center of modes are chosen as ( -2 , 0) , (2 , 0) , ( -3 , 4) , and (3 , 4) , and null values for the higher dimension coordinates. The first two modes have an identity covariance matrix, while the remaining two modes have a diagonal covariance matrix with variance equal to 2 . The results for regularized Stein thinning and the original Stein thinning are provided in Figure 10 for MALA sampler, and in Figure 11 for NUTS sampler. In both figures, the three tested step size ε are displayed, with a small impact on the resulting performance. Figures 9, 10, and 11 show the high performance improvement of regularized Stein thinning over the original algorithm.

Regularization parameter λ . Figure 12 displays the MMD obtained with regularization parameters λ , set as λ = 1 /m 2 and λ = 1 / log( m ) . These results should be compared with the ones shown

0.12

T

0.11

0.10

3

Q

P'

0.09

D

0.08

0.07

008

0.06

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

m

Regularized Stein Thinning

△+ log(p)

7.5

5.0

2.5

0.0

-2.5

0

5

0

in Figures 10 and 11, which were obtained with a regularization parameter λ = 1 /m . These additional experiments show the importance of choosing the regularization parameter as λ = O (1 /m ) , as suggested by Theorem 3.6. Indeed, slower rates of λ give poor quality samples, and faster rates than λ = O (1 /m ) tend to remove the effect of the entropic regularization, and we then recover similar performance than the original Stein thinning. On the other hand, λ = 1 /m provides a high improvement over the standard thinning, as shown in Figures 10 and 11.

glyph[negationslash]

Banana-shaped mixture with t-tails. The first batch of experiments in Section 4 considers a banana-shaped mixture with t-tails, defined as follows. Let ϕ : R d → R d be the transformation defined by ϕ i ( x ) = x i if i = 2 , and ϕ 2 ( x ) = x 2 + bx 2 1 -100 b . Let Z be a random variable that follows the multivariate t-Student distribution with degrees of freedom 7 . Then, the random variable X = ϕ ( Z ) + µ follows a t-banana-shaped distribution centered at µ . We consider a mixture of two t-banana-shaped distributions centered in 0 d and (0 , 8 , 0 d -2 ) , with weights w 1 = 0 . 25 and w 2 = 0 . 75 , respectively, which is illustrated in Figure 13 for d = 10 . The results for regularized Stein thinning and the original Stein thinning are provided in Figure 14 for MALA sampler, and in Figure 15 for NUTS sampler. In both figures, the three tested step size ε are displayed, which confirms the higher performance of regularized Stein thinning. Figures 13, 14, and 15 show the high performance improvement of regularized Stein thinning over the original algorithm.

=0.05

=0.1

= 0.5

0.4

Q

0.3

MMD(P,

0.2

Regularized Stein Thinning

0.1

100

300

500

700

900

1100

100

300

500

700

006

1100

100

300

500

700

900

1100

m

m

= 0.05

= 0.1

= 0.5

2.0

1.5

MMD(

1.0

Regularized Stein Thinning

MMD(P,

MMD(P,

0.5

0.0

10

20

30

40

50

60

70

80

90100

10

20

30

40

50

60

70

80

90100

10

20

30

40

5060

70

80

90100

p

p

P

=0.05

=0.1

=0.5

0.4

Q

0.3

MMD(P,

0.2

Regularized Stein Thinning

0.1

100

300

500

700

900

1100

100

300

500

700

006

1100

100

300

500

700

900

1100

m

m

m

= 0.05

= 0.1

=0.5

2.0

1.5

MMD(

1.0

Regularized Stein Thinning

MMD(P,

0.5

0.0

10

20

30

40

50

60

70

80

90

100

10

20

30

40

50

60

70

80

90100

10

20

30

40

5060

70

80

90100

d

P

p

A.3 Bayesian Logistic Regression

This appendix gathers additional results for Bayesian logistic regression. In particular, Table 2 provides a description of the tested datasets. Table 3 gives the resulting AUC, for m = 50 , 100 , 300 , using NUTS or MALA sampler. We recall that only the best AUC over the four tested MCMC step size ε is reported.

Recall that the Bayesian logistic regression defines the probability that Y i is of class 1 as p ( Y i = 1 | X i , β , β 0 ) = 1 / (1 + exp( -β 0 -β T X i )) , for some vector of parameters θ = ( β 0 , β ) ∈ R d +1 . The prior distributions of the weight vector θ is assumed to be Gaussian, p ( β ( j ) | γ ( j ) ) = N ( β ( j ) | 0 , 1 /γ ( j ) ) , and a Gamma prior with parameters ( a, b ) is chosen for the precision γ ( j ) . Upon marginalizing, it is found that β ( j ) is distributed as the non-standardized t-distribution Student-t (2 a, 0 , b/a ) [Bishop and Nasrabadi, 2006]. Following [Fong et al., 2019], the hyperparameters are chosen as a = b = 1 .

=0.05

= 0.1

=0.5

0.78

0.76

0.74

0.72

100

300

500

700

900

1100

100

300

500

700

006

1100

100

300

500

700

006

1100

u

m

m

= 0.05

= 0.1

=0.5

1.2

(u

MMD(P,

1.0

Regularized Stein Thinning

0.8

100

300

500

700

006

1100

100

300

500

700

006

1100

100

300

500

700

900

1100

m

Regularized Stein Thinning

△+ log(p)

30

20

2

10

0

-10

-20

0

20

-20

0

20

-20

0

20

c(1)

c()

c(1)

= 0.1

=0.5

=1.0

2.0

2.0

2.0

Regularized Stein Thinning

(

1.5

1.5

1.5

MMD(P,

1.0

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

100

300

500

700

900

1100

100

300

500

700

900

1100

100

300

500

700

900

1100

m

m

m

= 0.1

= 0.5

= 1.0

Regularized Stein Thinning

10

2030

40

50

60

70

80

90100

10

20

30

405060

7080

90

100

1020

08

4050

60

70

8090100

d

P

d

2.0

= 0.1

2.0

= 0.5

2.0

= 1.0

Regularized Stein Thinning

Qm)

1.5

1.5

1.5

MMD(P,

MMD(P,

1.0

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

100

300

500

700

900

1100

100

300

500

700

900

1100

100

300

500

700

900

1100

m

m

m

= 0.1

=0.5

= 1.0

Regularized Stein Thinning

10

203040

50

60

70

80

90100

10

20

30

40

50

60

90

100

10

20

30

50

60

70

80

90100

d

p

d

B Proof of Theorem 2.3

Assumption 2.1 (Distant bimodal mixture distributions) . Let p and q be two mixture distributions in R d , made of two modes centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , with µ > 0 . The distribution of each mode of p ∈ C 1 ( R d ) has R d as support, whereas each mode distribution of q have a compact support, included in a ball of radius r > 0 , with r < µ . The left mode of p has weight w p = 1 / 2 , and the right mode has weight 1 -w p . Similarly, w and 1 -w are the mode weights of q . Let Q L and Q R be the probability measures that respectively admit the density of the left and right modes of q , and P and Q w be also the probability laws for p and q .

glyph[negationslash]

Assumption 2.2. For distant bimodal mixture distributions q and p satisfying Assumption 2.1, and for η ∈ (0 , 1) , we have ∣ ∣ KSD 2 ( P , Q L ) / KSD 2 ( P , Q R ) -1 ∣ ∣ < η .

Theorem 2.3. Let k p be the Stein kernel associated with the radial kernel k ( x , x ′ ) = φ ( ‖ x -x ′ ‖ 2 /glyph[lscript] ) , where x , x ′ ∈ R d , glyph[lscript] > 0 , and φ ∈ C 2 ( R ) , such that φ ( z ) → 0 , φ ′ ( z ) → 0 , and φ ′′ ( z ) → 0 for z →∞ . Let p and q be two bimodal mixture distributions satisfying Assumptions 2.1 and 2.2, for any η ∈ (0 , 1) . We define w glyph[star] as the optimal mixture weight of q with respect to the KSD distance, i.e., w glyph[star] = argmin w ∈ [0 , 1] KSD( P , Q w ) . Then, for µ large enough, we have ∣ ∣ w glyph[star] -1 2 ∣ ∣ < η 2(1 -η ) .

MALA Sampler

m

= 100

ST

RST

0.72 (0.048)

0.53

(0.028)

0.53 (0.024)

0.55 (0.034)

0.80 (0.010)

Dataset

m

= 50

ST

0.68 (0.044)

Diabetes

0.51

(0.012)

Haberman 0.52 (0.034)

Liver

0.54 (0.033)

Sonar

0.80 (0.019)

RST

0.93

(0.010)

0.48 (0.010)

0.60

(0.027)

0.70

(0.008)

0.80 (0.019)

0.93

(0.007)

0.51 (0.014)

0.58

(0.017)

0.69

(0.005)

0.80 (0.010)

m

= 300

ST

RST

0.72 (0.037)

0.53 (0.016)

0.55 (0.024)

0.57 (0.024)

0.81

(0.013)

0.88

(0.026)

0.56

0.61

(0.016)

(0.013)

0.62

(0.032)

0.80 (0.010)

Lemma 1. If k p is the Stein kernel associated with the radial kernel k ( x , y ) = φ ( ‖ x -y ‖ 2 /glyph[lscript] ) , where glyph[lscript] > 0 , φ ∈ C 2 ( R ) , and x , y ∈ R d such that s p ( x ) , s p ( y ) < s 0 , then we have

Proof of Theorem B. We consider the mixture distributions p and q satisfying Assumption 2.1, for µ > 0 and r > 0 , and assume that Assumption 2.2 is satisfied for η ∈ (0 , 1) . More precisely, we denote by q L the distribution of the left mode of the mixture q , and similarly, q R is the distribution of the right mode of q . The probability measures Q L and Q R respectively admits the densities q L and q R .

By definition of the KSD, we can write

Additionally, given the above notations, q takes the form q = wq L + (1 -w ) q R . Then, we can develop the KSD expression to get

where the last term follows from the symmetry of k p . Finally, we have

and we denote by ∆ L,R the last term of this equation, which now writes

We first focus on the last term ∆ L,R of this expression, which can be shown to be arbitrarily small when µ gets large. According to Assumption 2.1, the distance between the centers of the two modes is 2 µ , and both q L and q R have a compact support included in a ball of radius r . Consequently, for x , x ′ ∈ R d such that q L ( x ) > 0 and q R ( x ′ ) > 0 , then ‖ x -x ′ ‖ 2 > 2( µ -r ) . Additionally, since the score s p is continuous, s p is bounded on a compact set, and it exists s 0 > 0 such that s p ( x ) < s 0 and s p ( x ′ ) < s 0 . Then, from Lemma 1, we have

and since q L ( x ) q R ( x ′ ) = 0 for ‖ x -x ′ ‖ 2 < 2( µ -r ) , we get

By assumption, φ ( z ) → 0 , φ ′ ( z ) → 0 , and φ ′′ ( z ) → 0 for z →∞ , and then, we have

Next, we reorder the terms of equation (4) to get a second-order polynomial in w as follows

Notice that the coefficient of w 2 is KSD 2 ( P , Q 1 / 2 ) / 4 , and is therefore positive. Then, KSD 2 ( P , Q w ) admits a unique minimum with respect to w , given by

We rewrite w glyph[star] as follows,

We can deduce the following bound

which gives an upper bound for the numerator of the right hand side of inequality (5). Additionally, for µ large enough, ∆ R,L is arbitrarily small, and in particular, we can have ∆ R,L < KSD 2 ( P , Q R ) / 2 , and then 2(1 -∆ L,R / KSD 2 ( P , Q R )) > 1 . Next, we use the triangle inequality to get

where the last inequality is obtained using 2(1 -∆ L,R / KSD 2 ( P , Q R )) > 1 for µ large enough, and Assumption 2.2 again. Finally, this lower bound on the denominator and the upper bound on the numerator combined with inequality (5) give

glyph[negationslash]

Proof of Lemma 1. We consider x , y ∈ R d , such that x = y , s p ( x ) < s 0 and s p ( y ) < s 0 . From Equation ( 2 ) of the main article, we derive the Stein kernel obtained for a radial kernel φ ( ‖ x -y ‖ 2 /glyph[lscript] ) , where glyph[lscript] > 0 and φ ∈ C 2 ( R ) , and get

Using Cauchy-Schwartz inequality, we have s p ( x ) · s p ( y ) ≤ s 2 0 , and

Overall, we obtain the following bound

C Proofs of Theorem 2.4, Corollary 2.5, and Corollary 2.6

Theorem 2.4 (KSD spurious minimum) . Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . Let { x i } i m =1 ⊂ M s 0 = { x ∈ R d : ‖ s p ( x ) ‖ 2 ≤ s 0 } be a fixed set of points of empirical measure Q m = 1 m ∑ m i =1 δ ( x i ) , with s 0 ≥ 0 and m ≥ 2 . We have KSD 2 ( P , Q m ) < E [KSD 2 ( P , P m ) ] , if the score threshold s 0 and the sample size m are small enough to satisfy m< 1 + ( E [ ‖ s p ( X ) ‖ 2 2 ] -s 2 0 ) / (2 βd/glyph[lscript] 2 +2 βs 0 /glyph[lscript] + s 2 0 ) .

Corollary 2.5 (Low KSD samples at density minimum) . Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . Let p be a density with at least one local minimum or saddle point. For m ≥ 2 , if { x i } i m =1 ⊂ R d is a set of points, all located at local minimum or saddle points of p , then we have KSD 2 ( P , Q m ) < E [KSD 2 ( P , P m ) ] , if m< 1 + glyph[lscript] 2 2 βd E [ ‖ s p ( X ) ‖ 2 2 ] .

Corollary 2.6 (KSD spurious minimum for Gaussian mixtures) . Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . Let the density p be a Gaussian mixture model of two components with equal weights, respectively centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , of variance σ 2 I d , and let ν = µ/σ . If ν > 1 and 0 ≤ s 0 < [ ν √ ν 2 -1 -ln( ν + √ ν 2 -1) ] /µ , then for any { x i } i m =1 ⊂ M s 0 of empirical measure Q m , we have

(i) KSD 2 ( P , Q m ) < E [KSD 2 ( P , P m ) ] if m and s 0 satisfy m< 1 + E [ ‖ s p ( X ) ‖ 2 2 ] -s 2 0 2 βd/glyph[lscript] 2 +2 βs 0 /glyph[lscript] + s 2 0 ,

(ii) there exists three disjoint intervals I -µ , I 0 , I µ ⊂ R , respectively centered around -µ , 0 , and µ , such that x (1) 1 , . . . , x (1) m ∈ I -µ ∪ I 0 ∪ I µ .

Lemma2. Let k p be the Stein kernel associated to the IMQ kernel, with parameters glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . For s 0 ≥ min x ∈ R d ‖ s ( x ) ‖ 2 , and x , y ∈ M s 0 , we have

Proof of Theorem 2.4. By definition of the kernelized Stein discrepancy between the target distribution P and the empirical measure P m = 1 m ∑ m i =1 δ ( X i ) , we have

In what follows, k p denotes the Stein kernel obtained for the inverse multi-quadratric kernel function, i.e. ,

glyph[negationslash]

Given that X i and X j are independent random variables that follow the distribution P , one has E [ k p ( X i , X j )] = 0 for any i = j . Using this property and the closed-form expression of the Stein kernel k p , it is found that

where X ∼ P . On the other hand, the kernelized Stein discrepancy between the target P and the empirical distribution Q m = 1 m ∑ m i =1 δ ( x i ) is given by glyph[negationslash]

glyph[negationslash]

Next, it can be shown that the difference m ( KSD 2 ( P , Q m ) -E [KSD 2 ( P , P m )] ) takes the form glyph[negationslash]

Since x i , x j ∈ M s 0 , we can use Lemma 2 to bound the terms k p ( x i , x j ) , and then obtain

where the right hand side is always negative if

Proof of Lemma 2. The Stein kernel k p obtained for the inverse multi-quadratric kernel function, with parameters glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 , is given by

Since the first term is always negative and (1 + ‖ x -y ‖ 2 2 /glyph[lscript] 2 ) -α ≤ 1 for α = β, β +1 , we obtain

We define the function g for z ≥ 0 as

and a simple function analysis shows that

Since β ∈ (0 , 1) , we have c 1 ≤ 1 . We combine the last two inequalities for k p ( x , y ) and c 1 to get

We can apply Cauchy-Schwartz inequality, and since x , y ∈ M s 0 , we get

and also

Overall, we obtain

Proof of Corollary 2.5. As minimum and saddle points are stationary points of p , we have

and we can apply Theorem 2.4 for s 0 = 0 to get the final result.

Proof of Corollary 2.6. Let the density p be a Gaussian mixture model of two components with equal weights, respectively centered in ( -µ, 0 d -1 ) and ( µ, 0 d -1 ) , and of variance σ 2 I d , and let ν = µ/σ . Weassume that ν > 1 and 0 ≤ s 0 < [ ν √ ν 2 -1 -ln( ν + √ ν 2 -1) ] /µ . Then, according to Theorem 2.4, for any { x i } i m =1 ⊂ M s 0 of empirical measure Q m = 1 m ∑ m i =1 δ ( x i ) , we have

To prove statement (ii), we need to characterize the shape of the set M s 0 ⊂ R d , given by the level lines of the squared score norm ‖ s p ( x ) ‖ 2 2 . The density p is a Gaussian mixture, i.e.,

where x ( -1) is the vector x without the first component. Then, the score is also given by an explicit formula,

where tanh is the standard hyperbolic tangent function. An important property of this score function is that the j -th component of s p only depends on x ( j ) , which makes s ( j ) p ( x ) invariant by any translation orthogonal to the j -th axis. Then, we can compute the squared score norm

7.5 -

5.0-

01

S0-

0.0-

0

min

min

where s (1) p ( z ) 2 = ( z σ 2 -µ σ 2 tanh( µ σ 2 z ) ) 2 . A simple function analysis of this univariate function, illustrated in Figure 16, shows that s (1) p ( z ) 2 has two local maximum in z -max and z + max , and three local minimum in z -min , 0 , and z + min , provided that ν = µ/σ > 1 . We also get that s (1) p ( z ) 2 grows to + ∞ when x (1) → + / -∞ . The extreme values are ordered as follows

The values of z -min , z -max , z + max , and z + min are given by the zeros of the first derivative of s (1) p ( z ) 2 , defined by

This derivative vanishes when one of the two factors is null. Since µ/σ > 1 , the first term is null when

which leads to

The second factor is null when

Obviously, z = 0 is solution. Since µ/σ > 1 , equation (6) has two additional solutions. Although they do not have a closed form, we have

Also notice that, as µ/σ gets larger, z -min is closer to -µ , and z + min to µ . For example in Figure 16, we set µ/σ = 3 , and we hardly see a gap between -µ and z -min , or µ and z + min .

By definition, for x ∈ M s 0 , we have

Therefore, given the variations of s (1) p ( x (1) ) 2 detailed above and illustrated in Figure 16, if s 2 0 < s (1) p ( z -max ) 2 = s (1) p ( z + max ) 2 , it exists three disjoint intervals I -µ , I 0 , I µ ⊂ R , respectively centered around -µ , 0 , and µ , such that

To conclude, we compute the value of s (1) p ( z -max ) , that is

Using the formulas tanh( arcosh ( x )) = √ x 2 -1 x for | x | > 1 , arcosh ( x ) = ln( x + √ x 2 -1) , and with ν = µ/σ , we get

which is always strictly positive since ν > 1 . By assumption, 0 ≤ s 0 < [ ν √ ν 2 -1 -ln( ν + √ ν 2 -1) ] /µ , and therefore, we have s 0 < s (1) p ( z + max ) = s (1) p ( z -max ) , which concludes the proof of statement (ii).

D Proof of Theorem 3.3

Theorem 3.3. Let k p be the Stein kernel associated with the radial kernel k ( x , x ′ ) = φ ( ‖ x -x ′ ‖ 2 /glyph[lscript] ) , where x , x ′ ∈ R d , glyph[lscript] > 0 , and φ ∈ C 2 ( R ) . Let p and q be two bimodal mixture distributions satisfying Assumption 2.1. We define w glyph[star] λ as the optimal mixture weight of q with respect to the entropic regularized KSD distance, i.e., w glyph[star] λ = argmin w ∈ [0 , 1] KSD λ ( P , Q w ) . If E [log( p ( Z L ))] = E [log( p ( Z R ))]

where Z L ∼ Q L and Z R ∼ Q R , it exists λ ∈ R such that w glyph[star] λ = w p .

Proof of Theorem 3.3. We consider the mixture distributions p and q satisfying Assumption 2.1, for µ > 0 and r > 0 , and denote by q L the distribution of the left mode of the mixture q , and similarly, q R is the distribution of the right mode of q . The probability measures Q L and Q R respectively admits the densities q L and q R . As in the proof of Theorem 2.3, we get that

Next, we define Z ∼ Q w , Z L ∼ Q L , and Z R ∼ Q R , and develop the entropic regularization term,

Combining these two results, we have

We recall that p does not depend on w , but only on the fixed weight w p and the two mode distributions. Consequently, KSD 2 λ ( P , Q w ) is a second-order polynomial with respect to w . As for Theorem 2.3, the glyph[negationslash]

coefficient of w 2 is KSD 2 ( P , Q 1 / 2 ) / 4 , and is therefore positive. Then, the polynomial is minimized with respect to w , at w glyph[star] defined by

glyph[negationslash]

Since E [log( p ( Z L ))] -E [log( p ( Z R ))] = 0 by assumption, we get that w glyph[star] λ = w p for lambda defined by

E Proof of Theorem 3.4

Theorem 3.4. Let k p be the Stein kernel associated with the IMQ kernel with glyph[lscript] > 0 , β ∈ (0 , 1) , and c = 1 . For m ≥ 2 , let { x i } i m =1 ⊂ R d be a set of points concentrated at x 0 , a local minimum or saddle point of p , and of empirical measure Q m . Then, we have L-KSD 2 ( P , Q m ) > E [ L-KSD 2 ( P , P m ) ] , if the density at x 0 satisfies p ( x 0 ) < ∆ + p ( x 0 ) / ( E [ ‖ s p ( X ) ‖ 2 2 ] + E [∆ + log p ( X )] ) .

Proof of Theorem 3.4. Following the same approach as in the proof of Theorem 2.4, we have

where X ∼ P , and, with the empirical distribution Q m = 1 m ∑ m i =1 δ ( x i ) , glyph[negationslash]

As { x i } i m =1 are concentrated at a local minimum or saddle point x 0 , the score is null for all particles, as well as the distances between them, and we get

Next, the difference m ( L-KSD 2 ( P , Q m ) -E [ L-KSD 2 ( P , P m )] ) writes

By definition, with

leading to

As x 0 is a stationary point of p , ∂p ( x 0 ) /∂x ( j ) = 0 , and we obtain

Finally, we get

which ensures that L-KSD 2 ( P , Q m ) -E [ L-KSD 2 ( P , P m )] > 0 for m ≥ 2 , provided that

F Proof of Theorem 3.6

Assumption 3.5. Let Q be a probability distribution on R d , such that P is absolutely continuous with respect to Q . Let { Z i } i ∈ N ⊂ R d be a Q -invariant, time-homogeneous Markov chain, generated using a V -uniformly ergodic transition kernel, such that V ( x ) ≥ d P d Q √ 2 βd/glyph[lscript] 2 + ‖ s p ( x ) ‖ 2 2 . Suppose that, for some γ > 0 , sup i ∈ N E [ e γ | log( p ( Z i )) | ] < ∞ , sup i ∈ N E [ e γ ∆ + log p ( Z i ) ] < ∞ ,

Assumption 3.5 is close to the assumption made in Riabiz et al. [2022]. For the last two finite expectations of the assumption, notice that we have just plugged the formula k p ( x , x ) = 2 βd/glyph[lscript] 2 + || s p ( x ) || 2 into the integrability conditions, since this formula is quite straightforward. But in addition to Riabiz et al. [2022], we require two integrability assumptions, one for each regularization term. We give additional insights about these quite complex conditions below. Overall, our assumptions are hardly stronger than those of Riabiz et al. [2022].

The condition E [ e γ | log( p ( Z i )) | ] < ∞ is satisfied if ∫ f ( x ) p ( x ) -γ dx < ∞ , where f is the distribution of Z i (since p ( x ) > 1 only on a compact set and p is continuous). Therefore, there exists γ > 0 satisfying such condition, provided that the tails of density f are not too heavy compared to the tails of p . This has to be verified for all iterations of the Markov Chain, and happens to be a quite mild assumption.

Regarding the condition involving the Laplacian term, we can use the analysis from Riabiz et al. (2022) (Appendix S2.4). It relies on a Lipschitz condition for the score function ∇ log p , to transform the first finite expectation of Assumption 3.5 into a more amenable and practical integrability condition which writes E [ κ ‖ Z i ‖ 2 2 ] < ∞ for some κ ∈ (0 , ∞ ) . In the same spirit, if we add a Lipschitz condition on ∆ + log p , our second additional integrability assumption reduces to a similar amenable condition. Finally, even though Riabiz et al (2022) do not elaborate further on this condition, it can be noticed that it is satisfied if all Z i have subgaussian tails, for example. Besides, one of the main example of distantly dissipative distributions are log-concave functions outside of a compact set. In this case, the Laplacian regularization is null outside of a compact set, since the Laplacian is negative for concave functions. Then, the integrability condition is automatically satisfied for this type of distributions.

Theorem 3.6. Let P be a distantly dissipative probability measure, that admits the density p ∈ C 2 ( R d ) , k p be the Stein kernel associated with the IMQ kernel wh glyph[lscript],c > 0 , β ∈ (0 , 1) . Let { Z i } i ∈ N ⊂ R d be a Markov chain satisfying Assumption 3.5, π be the index sequence of length m n generated by regularized Stein thinning, and Q m n be the empirical measure of { Z π i } m n i =1 . If log( n ) α < m n < n , with any α > 1 , and λ m n = o (log( m n ) /m n ) , then we have almost surely Q m n = ⇒ n →∞ P .

Theorem 3.6 extends Theorem 3 from Riabiz et al. [2022] to regularized Stein Thinning, using Lemmas 3-4 and assuming the following convergence rate of the regularization parameter λ m n = o (log( m n ) /m n ) . Lemma 3 also extends Theorem 1 from Riabiz et al. [2022][Theorem 1], whereas Lemma 4 is a slight modification of Lemma 5 from Riabiz et al. [2022].

Lemma3. Let P be a probability measure on R d that admits density p ∈ C 2 ( R d ) , k p be a reproducing Stein kernel, and { x i } n i =1 ⊂ R d a fixed set of points. If π is an index sequence of length m produced by regularized Stein thinning, then we have for λ > 0 ,

where the weights w glyph[star] are defined as

Lemma 4. Let f be a non-negative function on R d . Consider a sequence of random variables ( X i ) i ∈ N ⊂ R d such that, for some γ > 0 ,

If log( n ) α < m n < n , with any α > 1 , then we have almost surely,

The proofs of Lemmas 3 and 4 are reported at the end of this section. We first proceed with the proof of Theorem 3.6.

Proof of Theorem 3.6. From Lemma 3, we have

Riabiz et al. [2022][Proof of Theorem 3, p.12] showed that the term ( glyph[star] ) converges towards 0 almost surely as n →∞ . For the remaining terms, from Assumption 3.5, we have

and

We can use Lemma 4 with f ( x ) = k p ( x , x ) and f ( x ) = ∆ + log( x ) to deduce that ( glyph[star]glyph[star] ) → 0 and ( glyph[star] glyph[star] glyph[star] ) → 0 , respectively. The remaining term ( glyph[diamondmath] ) can be rewritten as

Using the assumption that λ m n = o (log( m n ) /m n ) and Lemma 4 with f ( x ) = | log( p ( x )) | , we conclude that ( glyph[diamondmath] ) → 0 . It follows that KSD 2 ( 1 m n ∑ m n j =1 δ ( x π j ) ) → 0 almost surely as n → ∞ . Given that p is assumed to be distantly dissipative, we apply Theorem 4 from Chen et al. [2019] to obtain that Q m n ⇒ P almost surely, as n →∞ .

Proof of Lemma 3. We consider an iteration t ∈ { 2 , . . . , m } of regularized Stein thinning, where m ≥ 2 is the final length of the thinned sample. We define a t = t 2 KSD 2 ( P , 1 t ∑ t j =1 δ ( x π j )) and f t = ∑ t j =1 k p ( x π j , · ) . We also denote S 2 1 = max i =1 ,...,n k p ( x i , x i ) + max i =1 ,...,n ∆ + log( p ( x i )) , and S 2 = 2max i =1 ,...,n | log( p ( x i )) | . Using the definition of the squared KSD of an empirical measure, we have

Let x glyph[star] t = arg min y ∈{ x i } n i =1 f t -1 ( y ) . By definition, x π t minimizes the cost function of the regularized Stein thinning algorithm at iteration t , and we have

We combine this last inequality with the first equation to obtain

and then,

By definition, 0 ≤ | log( p ( x i )) | ≤ S 2 and k p ( x glyph[star] t , x glyph[star] t ) + ∆ + log( p ( x glyph[star] t )) -∆ + log( p ( x π t )) ≤ S 2 1 , hence, we have

As in the proof of Riabiz et al. [2022], we have min y ∈{ x i } n i =1 f t -1 ( y ) ≤ √ a t -1 ‖ h glyph[star] ‖ H ( k p ) where h glyph[star] is the element in the RKHS H ( k p ) of the form h glyph[star] = ∑ n i =1 w glyph[star] i k p ( x i , · ) . As a result, a t is bounded as follows:

We then show by induction that where

With such as a result, we will have KSD 2 ( P , 1 t ∑ t j =1 δ ( x π j )) ≤ KSD 2 ( P , ∑ n i =1 w glyph[star] i δ ( x π j )) + C t + λS 2 and obtain the advertised result in Theorem 3 for the last iteration t = m .

For t = 1 , we have a 1 = k p ( x π 1 , x π 1 ) ≤ S 2 1 and thus a 1 ≤ ‖ h glyph[star] ‖ 2 H ( k p ) + C 1 + λS 2 . For a fixed t ≥ 2 , assume that a t -1 ≤ ( t -1) 2 ( ‖ h glyph[star] ‖ 2 H ( k p ) + C t -1 + λS 2 ) where C t -1 = 1 t -1 ( S 2 1 -‖ h glyph[star] ‖ 2 H ( k p ) ) ∑ t -1 j =1 1 j . We then have

where

Using Riabiz et al. [2022][Lemma 1], we have

It follows from Equation (9) that we need R t ≤ 0 , i.e. ,

The above inequality is always satisfied as long as

which is equivalent to where

With the assumption that log( n ) β ≤ m n < n with β > 1 , it is deduced that c 1 ( m 2 n ) → 0 and c 2 ( m n ) → 0 as n →∞ .

G Laplacian Stein Operator

Instead of the standard Langevin operator, we can use T p ( g ) = ∆( pg ) /p , mentioned in Oates et al. [2017, Appendix A.2]. However, such Stein operator introduces similar problems as Pathology I, since the Stein kernel associated with T ′ p ( g ) also has spurious minimum in regions where second derivatives of p vanish, as illustrated below. Therefore, it is more appropriate to taylor a specific Laplacian correction as proposed in this paper, which cannot directly be derived from a Stein kernel, since it is not differentiable. In this appendix, we study the operator T p defined as T p g = ∆( pg ) /p and such that [Oates et al., 2017, Appendix A.2]

and always true by definition of C t . Hence we have shown that a t ≤ t 2 ( ‖ h glyph[star] ‖ 2 H ( k p ) + C t + λS 2 ) . Given that ‖ h glyph[star] ‖ 2 H ( k p ) = KSD 2 ( P , ∑ n i =1 w glyph[star] i δ ( x i )) , we have

where C t ≤ 1+log( t ) t ( max i =1 ,...,n k p ( x i , x i ) + max i =1 ,...,n ∆ + log( p ( x i )) ) (see [Riabiz et al., 2022][Lemma 2]).

Proof of Lemma 4. We follow the proof of Lemma 5 from [Riabiz et al., 2022]. In the last step of the proof, we essentially need to show that

for all g belonging to G , and with Z ∼ P . For each dimension j ∈ { 1 , . . . , d } , let T j p be the operator defined as

for g : R d → R , and where x j is the j -th coordinate of x to simplify notations. The operator T p g can then be rewritten as ( T p g )( x ) = ∑ d j =1 ( T j p g )( x ) . Gorham and Mackey [2017][Proposition 2] establishes the closed-form expression of the KSD in the case of the multidimensional Langevin operator. We generalize the proof of Gorham and Mackey [2017][Proposition 2] for the Laplacian operator ( T p g )( x ) = ∆( p ( x ) g ( x )) /p ( x ) and establish a closed-form expression of the Stein kernel k p : R d × R d → R . For a given kernel function k : R d × R d → R , k ∈ C 2 , 2 , the Stein kernel k p is given by [Gorham and Mackey, 2017]

After a few developments, it is found that

A closed-form expression can then be obtained in the case of, e.g. , an inverse multiquadratic kernel of the form k ( x , y ) = (1 + ‖ x -y ‖ 2 2 /glyph[lscript] 2 ) -1 / 2 where glyph[lscript] denotes the bandwidth. In this case, one has

A closed-form expression of k p can then be obtained by combining Equations (10)-(13). In contrast to the Langevin Stein kernel (see Equation ( 2 ) of the main article), the above Stein kernel involves second-order derivatives of the density p .

We run experiments based on our Example 1 of Gaussian mixtures for this new Stein operator T p g = ∆( pg ) /p . We sequentially set µ = 2 and µ = 5 , with σ = 1 and w = 0 . 5 . Results are displayed in Figure 17. In the left panel with µ = 2 , we see that Pathology II does not occur, as opposed to Figure 2 of the main article with the Langevin Stein operator. However, when µ is set to 5 in the right panel, all particles are concentrated in spurious minimum again. Therefore, introducing higher-order derivatives of the target density through the Stein operator does not seem to be a promising route.

where

3

2

0

0

-1

-2

-2

0

2

0

5

10

C1

C1

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper explores two pathologies associated with KSD thinning: the inability to distinguish mixing weights and concentration on low-probability regions of the target. The authors provide formalisations of when these pathologies can occur (Theorem 2.3 and Theorem 2.4), supported by empirical evidence. To address these issues, they propose the addition of two penalty terms: entropy regularisation and a Laplacian correction. The effectiveness of the proposed regularised KSD approach in mitigating these pathologies is demonstrated through both theoretical studies and numerical experiments.\n\nSTRENGTHS:\n**Novelty**: Although the pathologies of blindness to mixing proportions and spurious concentration to low-probability regions have been acknowledged in the community, this paper stands out by providing formalizations and solutions to these issues. Solutions to Pathology I has been previously proposed in other applications of KSD (e.g., Zhang et al., 2022; Liu et al., 2023), but no solution nor formalisation have been provided regarding Pathology II. The proposed regularisation for Stein thinning represents a novel and principled solution, supported by strong theoretical guarantees.\n\n**Applicability**: The paper thoroughly discusses and provides experimental evidence on the impact of the regularisation parameter, the degree of freedom, on the performance of the proposed approach. Practical guidelines are given to aid in choosing appropriate parameter values in practice.\n\n**Writing**: This paper is well motivated and clearly written. \n\nZhang, M., Key, O., Hayes, P., Barber, D., Paige, B., & Briol, F. X. (2022). Towards healing the blindness of score matching. *arXiv preprint arXiv:2209.07396*.\n\nLiu, X., Duncan, A. B., & Gandy, A. (2023). Using Perturbation to Improve Goodness-of-Fit Tests based on Kernelized Stein Discrepancy. ICML 2023.\n\nWEAKNESSES:\nThis paper is well written overall. I only found a couple of minor areas that could potentially be improved:\n\n**Limitations**: It would be beneficial to include more in-depth discussions on any limitations of the L-KSD approach. E.g., whilst a convergence guarantee was given when $n \\to \\infty$ (Theorem 3.5), for finite $n$ and $m_n$, how prominent is the bias introduced due to the addition of regularisation?\n\n**Clarity**: Whilst this paper is mostly clearly written, some parts deserve elaboration and clarification. E.g.,\n\n1. In Section 3.2, the intuition behind why the entropic regularisation solves Pathology I is not very clearly explained (also see Q1).\n2. In Theorem 3.3, it would be helpful to clarify the precise meaning of “concentrated at $x_0$”.\n3. Also in Theorem 3.3, it might be beneficial to provide an interpretation of the inequality for $p(x_0)$ on L313. Specifically, is this condition satisfied by the mixture of Gaussian example used in the paper?\n4. In Section 4, it would be helpful to clarify the choice of the kernel (and any hyper-parameters) in the MMD used to evaluate the results.\n\nQUESTIONS:\n**Q1. Entropic regularisation**: Could you elaborate on why the entropic regularisation helps alleviate Pathology I? In particular, L242 says “The main idea of this entropic regularization is that $− \\log(p(x))$ takes higher values in modes of smaller probability, and therefore provides the relative mode weight information”. Whilst I agree with the first half of this claim, I am unsure how this sensitivity to low-density regions relates to the ability to detect mis-specification of the mixing weights, and why this regularisation would lead to small values when the mixing weights are correctly specified.\n\n**Q2. Figure 3**: The proportion of particles selected by L-KSD that lie in the left mode, while better complying with the true proportion of $w=0.2$ compared to the result of KSD, still differs from the true value by a statistically significant gap (since the reported estimated proportion is $0.11 \\pm 0.03$). Could you explain whether this gap is expected? Does it mean that the regularised penalty suffers from a significant bias in the finite sample scheme, despite being consistent in the infinite limit?\n\nLIMITATIONS:\nMore discussions on any limitations of the approach should be discussed; see the question in \"Weaknesses\".", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this article the author(s) studied the Stein thining method, an algorithm for post-processing outputs of MCMC based on the kernelized Stein discrepancy (KSD). This article first theoretically analyzed two pathologies of KSD, and then proposed methods to mitigate the two issues by regularizing the KSD objective function. Both theoretical analysis and numerical experiments show that the regularized Stein thinning method has improved the sampling quality.\n\nSTRENGTHS:\nThe blindness of score-based methods to multimodal distributions is a long-standing issue, and this paper presents a careful theoretical anaysis in the context of Stein thining. The author(s) also point out another pathology of KSD caused by spurious minimum of the objective function. Both analyses deepen our understanding of KSD-based methods. The proposed regularzed Stein thinning method also shows promising performance.\n\nWEAKNESSES:\n1. For the theoretical analyses, some questions need to be clarified. See the questions section.\n2. For the numerical experiments, the target densities explored so far are somewhat simplified. I would suggest considering the following additional scenarios:\n    - Sampling from energy-based models, e.g. [1].\n    - Gaussian mixtures with more than two distant modes, e.g. the experiments in [2].\n    - Bayesian neural networks, as an extension of the linear logistic regression.\n\nThese are all \"difficult\" target distributions, and I would not expect regularized Stein thinning solving them perfectly. But exploring at least  some of the cases above would be helpful.\n\n3. If my understanding is correct, the truncated Laplacian operator requires computing the eigen decomposition of the Hessian matrix for every point in the MCMC sample, and this seems to be a huge computational cost. (After author rebuttal: **this point was my misunderstanding**, and the actual computing cost was clarified by the author(s). The truncated Laplacian operator only requires computing the diagonal elements of the Hessian matrix.)\n\n[1] Che, T., Zhang, R., Sohl-Dickstein, J., Larochelle, H., Paull, L., Cao, Y., & Bengio, Y. (2020). Your GAN is secretly an energy-based model and you should use discriminator driven latent sampling. Advances in Neural Information Processing Systems.\n\n[2] Qiu, Y., & Wang, X. (2023). Efficient Multimodal Sampling via Tempered Distribution Flow. Journal of the American Statistical Association.\n\nQUESTIONS:\n1. Both Theorem 2.3 and Theorem 3.2 are based on radial basis function kernels, whereas Theorem 2.4 and Theorem 3.3 are using  inverse multi-quadratic kernels. What considerations is this based on?\n2. Assumption 2.2 seems too hard to verify. Can you give an example for it to hold? Considering Example 1 of this article with only one dimension, is it possible to given an expression for $\\eta$? I would expect $\\eta$ to be a function of $\\mu$ and $w$, and $\\eta\\rightarrow 0$ as $\\mu\\rightarrow\\infty$.\n3. As mentioned in the weakness section, the computing time may be taken into account and reported in the experiments.\n\nLIMITATIONS:\nAfter rebuttal: my previous concern on computational cost is no longer a major problem, and other limitations may include the tuning of $\\lambda$, and some technicalities such as the choice of kernel functions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper proposes a regularized version of Stein thinning for post-processing the output of Markov Chain Monte Carlo algorithms. The regularization addresses two common challenges of Stein thinning: insensitivity to mode proportions and samples concentrating at stationary points. The paper analyzes the above two pathologies under a theoretical framework which is then utilized to justify the proposed regularization technique. Numerical experiments demonstrate gains over standard Stein thinning in both toy examples and in a logistic regression setting.

STRENGTHS:
- The paper is well-written and easy to follow. The structure of the paper is logical.
  
- I found the study of pathologies illuminating and interesting. The theoretical framework appears to be sound.
  
- The proposed technique shows promise both in the toy case and in the simple logistic regression setup.

WEAKNESSES:
- More intuition on the assumptions of the theoretical formulation would be helpful, with some discussion on the limitations of the results with more focus on practical applications of the method.
  
- Extra hyperparameter is introduced, for which a heuristic tuning method is provided, however a systematic approach is missing.

QUESTIONS:
- Why do we see an initial sharp increase of MMD with respect to $d$ followed by gradual decrease in the Gaussian mixture example (Fig. 5) for RST, but strict decrease in MMD for ST?
  
- Is there a systematic way to set $\lambda$ beyond the heuristics provided in the paper?

LIMITATIONS:
Limitations are not discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-21

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe goal of Stein thinning is to post-process the outputs of Markov chain Monte Carlo (MCMC) methods by minimizing the kernelized Stein discrepancy (KSD) between the produced chain and the target distribution. It is useful and has become fast a quite popular method in Bayesian inference because it automatically removes the burn-in period, corrects the bias, and has asymptotic properties of convergence towards the target.\n\nHowever it suffers from two pathologies, namely Pathology I of mode proportion blindness and Pathology II of mode collapse. Pathology I originates from the score insensitivity to mode weights, while Pathology II arises from the over-representation of a few modes. To mitigate these pathologies, the paper proposes a regularization technique called regularized Stein thinning based on adding an entropy and Laplacian term to the KSD.\n\nSTRENGTHS:\nThe paper does a good pedagogic job on the issues related to the optimization of the KSD, and offers a complete picture by providing theoretical negative results about the regular KSD (eg Theorem 2.3, 2.4).\nFor instance Th 2.4 (through Cor 2.5) says one can attain a better KSD than empirical samples of the true target mixture by putting all the points between the modes because this is a region where the score is close to zero.\n\nMoreover the proposed regularised KSD is shown to enjoy good properties regarding Pathologies I and II and has theoretical guarantees and extensive experiments to demonstrate its efficiency.\n\nWEAKNESSES:\nTheorem 2.4 is still limited to the IMQ kernel.\n\nThe Laplacian correction involves second-order derivatives of the target.\n\nTheir regularised KSD introduces an additional hyperparameter lambda that a user should tune, however it is not clear how to fix this parameter in advance (depending on the dimension, the target etc). It is not clear in the theoretical result Theorem 3.2 neither what lambda recovers the true mixture weights and fixes Pathology I. Still, the authors conducted a reasonable number of experiments with different values of lambda.\n\nQUESTIONS:\nDo the authors have any intuition on a general recipee or dependence of lambda on the parameters of the problem?\n\nLIMITATIONS:\nThe authors reasonably discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}