SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Clément Bénard $^{1}$  Brian Staber $^{1}$  Sébastien Da Veiga $^{2}$
<sup>1</sup> Safran Tech, Digital Sciences & Technologies, 78114 Magny-Les-Hameaux, France
2 ENSAI, CREST, F-35000 Rennes, France
{clement.benard, brian.staber}@{safrangroup.com}
sebastien.da-veiga@ensai.fr
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.
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 post-processing 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  $\tilde{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
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 kernel-based discrepancy measure  $D(\mathbb{P}, \mathbb{Q}_m)$  between the empirical distribution  $\mathbb{Q}_m$  of a subsample of the MCMC output of size  $m$ , and the target distribution  $\mathbb{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  $\mathbb{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  $\mathbb{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 from 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  $\mathbb{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  $\mathbb{P}$  and  $\mathbb{Q}$  on  $\mathbb{R}^d$  with  $d\geq 1$ , whenever  $\mathbb{P}$  admits a continuously differentiable density  $p$ , and the normalization constant of  $p$  is not tractable. Let  $k:\mathbb{R}^d\times \mathbb{R}^d\to \mathbb{R}$  be a positive semi-definite kernel and let  $\mathcal{H}(k)$  be the associated reproducing kernel Hilbert space (RKHS) with inner product  $\langle \cdot ,\cdot \rangle_{\mathcal{H}(k)}$  and norm  $\| \cdot \|_{\mathcal{H}(k)}$ . Kernelized Stein discrepancy belongs to the family of maximum mean discrepancies (MMD) [Gretton et al., 2006] defined as
$$
\operatorname {M M D} _ {k} (\mathbb {P}, \mathbb {Q}) = \sup  _ {\| f \| _ {\mathcal {H} (k)} \leq 1} | \mathbb {E} [ f (\mathbf {X}) ] - \mathbb {E} [ f (\mathbf {Z}) ] |, \tag {1}
$$
where  $\mathbf{X} \sim \mathbb{P}$ ,  $\mathbf{Z} \sim \mathbb{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  $\mathbb{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  $\mathcal{T}_p$  that maps functions  $g: \mathbb{R}^d \to \mathbb{R}^d$  to real-valued functions such that  $\mathbb{E}[\mathcal{T}_p g(\mathbf{X})] = 0$ , with  $\mathbf{X} \sim \mathbb{P}$ , for all  $g$  in  $\mathcal{G}(k) = \{g: \mathbb{R}^d \to \mathbb{R}^d: \sum_{i=1}^{d} \|g_i\|_{\mathcal{H}(k)}^2 \leq 1\}$ . The probability measure  $\mathbb{P}$  on  $\mathbb{R}^d$  is assumed to admit a continuously differentiable Lebesgue density  $p \in \mathcal{C}^1(\mathbb{R}^d)$ , such that  $\mathbb{E}[\|\nabla \log p(\mathbf{X})\|_2^2] < \infty$ . The Stein discrepancy is then defined as  $\mathrm{SD}(\mathbb{P}, \mathbb{Q}) = \sup_{g \in \mathcal{G}(k)} |\mathbb{E}[(\mathcal{T}_p g)(\mathbf{Z})]|$ , where  $\mathbf{Z} \sim \mathbb{Q}$ . If the Stein operator  $\mathcal{T}_p$  is chosen as the Langevin operator  $(\mathcal{T}_p g)(\mathbf{x}) = \langle g(\mathbf{x}), \nabla \log p(\mathbf{x}) \rangle + \langle \nabla, g(\mathbf{x}) \rangle$ , then Stein's discrepancy has a closed-form expression known as kernelized Stein discrepancy [Chwialkowski et al., 2016, Liu et al., 2016],  $\mathrm{KSD}^2(\mathbb{P}, \mathbb{Q}) = \mathbb{E}[k_p(\mathbf{Z}, \mathbf{Z}')]$ , where  $\mathbf{Z} \sim \mathbb{Q}, \mathbf{Z}' \sim \mathbb{Q}$ , and  $k_p$  denotes the Langevin Stein kernel defined from the score function  $s_p(\mathbf{x}) = \nabla \log p(\mathbf{x})$  for  $\mathbf{x}, \mathbf{x}' \in \mathbb{R}^d$ , as
$$
\begin{array}{l} k _ {p} (\mathbf {x}, \mathbf {x} ^ {\prime}) = \langle \nabla_ {\mathbf {x}}, \nabla_ {\mathbf {x} ^ {\prime}} k (\mathbf {x}, \mathbf {x} ^ {\prime}) \rangle + \langle s _ {p} (\mathbf {x}), \nabla_ {\mathbf {x} ^ {\prime}} k (\mathbf {x}, \mathbf {x} ^ {\prime}) \rangle \\ + \langle s _ {p} (\mathbf {x} ^ {\prime}), \nabla_ {\mathbf {x}} k (\mathbf {x}, \mathbf {x} ^ {\prime}) \rangle + \langle s _ {p} (\mathbf {x}), s _ {p} (\mathbf {x} ^ {\prime}) \rangle k (\mathbf {x}, \mathbf {x} ^ {\prime}). \tag {2} \\ \end{array}
$$
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  $\mathbb{P}$ . Gorham and Mackey [2017] also established
convergence guaranties when the kernel  $k$  is chosen as the inverse multi-quadratic (IMQ) kernel function  $k(\mathbf{x},\mathbf{x}^{\prime}) = (c + \| \mathbf{x} - \mathbf{x}^{\prime}\|_{\Gamma}^{2})^{-\beta}$  with  $c > 0, \beta \in (0,1)$ , the positive definite matrix  $\Gamma$  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.
Definition 1.1 (Distant dissipativity Gorham and Mackey [2017]). The density  $p \in \mathcal{C}^1(\mathbb{R}^d)$  is distantly dissipative if  $\lim_{r \to \infty} \inf \kappa(r) > 0$ , where  $\kappa(r) = \inf \left\{-2\frac{\langle s_p(\mathbf{x}) - s_p(\mathbf{y}), \mathbf{x} - \mathbf{y}\rangle}{\|\mathbf{x} - \mathbf{y}\|_2^2} : \| \mathbf{x} - \mathbf{y}\|_2 = r\right\}$ .
**Stein thinning algorithm.** Let  $\mathbb{P}$  be a target probability measure that admits density  $p$ , and let  $\{\mathbf{x}_i\}_{i=1}^n \subset \mathbb{R}^d$  be a MCMC output. The Stein thinning algorithm [Riabiz et al., 2022] selects  $m \leq n$  particles  $\mathbf{x}_{\pi_1}, \ldots, \mathbf{x}_{\pi_m}$  by greedily minimizing the kernelized Stein discrepancy. Given  $t - 1 < m$  particles  $\mathbf{x}_{\pi_1}, \ldots, \mathbf{x}_{\pi_{t-1}}$ , the  $t$ -th particle is defined as
$$
\pi_ {t} \in \operatorname * {a r g m i n} _ {i \in \{1, \dots , n \}} k _ {p} (\mathbf {x} _ {i}, \mathbf {x} _ {i}) + 2 \sum_ {j = 1} ^ {t - 1} k _ {p} \bigl (\mathbf {x} _ {\pi_ {j}}, \mathbf {x} _ {i} \bigr),
$$
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  $\ell$  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  $\ell$  varies. Furthermore, the sample quality output by Stein thinning varies in an erratic fashion with respect to  $\ell$ , making the design of heuristic procedures for the choice of  $\ell$  notoriously difficult. Following the literature recommendations [Riabiz et al., 2022], we use the median heuristic to set  $\ell$  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  $(-\mu, \mathbf{0}_{d-1})$  and  $(\mu, \mathbf{0}_{d-1})$ , of weights  $w$  and  $1 - w$ , and of variance  $\sigma^2 \mathbf{I}_d$ . The initial particles  $\{\mathbf{x}_i\}_{i=1}^n$  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. We set  $\mu = 3$  and  $\sigma = 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,
Figure 1: Illustration of Pathology I with the Gaussian mixture of Example 1 ( $d = 2$ ,  $\mu = 3$ ,  $\sigma = 1$ ,  $w = 0.2$ ,  $n = 3000$ ,  $m = 300$ ). Initial particles are in black, and the Stein thinning output is red.
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$ .
Assumption 2.1 (Distant bimodal mixture distributions). Let  $p$  and  $q$  be two mixture distributions in  $\mathbb{R}^d$ , made of two modes centered in  $(-\mu, \mathbf{0}_{d-1})$  and  $(\mu, \mathbf{0}_{d-1})$ , with  $\mu > 0$ . The distribution of each mode of  $p \in \mathcal{C}^1(\mathbb{R}^d)$  has  $\mathbb{R}^d$  as support, whereas each mode distribution of  $q$  have a compact support, included in a ball of radius  $r > 0$ , with  $r < \mu$ . The left mode of  $p$  has weight  $w_p \neq 1/2$ , and the right mode has weight  $1 - w_p$ . Similarly,  $w$  and  $1 - w$  are the mode weights of  $q$ . Let  $\mathbb{Q}_L$  and  $\mathbb{Q}_R$  be the probability measures that respectively admit the density of the left and right modes of  $q$ , and  $\mathbb{P}$  and  $\mathbb{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  $\eta \in (0,1)$ , we have  $|\mathrm{KSD}^2 (\mathbb{P},\mathbb{Q}_L) / \mathrm{KSD}^2 (\mathbb{P},\mathbb{Q}_R) - 1| < \eta$ .
Theorem 2.3. Let  $k_p$  be the Stein kernel associated with the radial kernel  $k(\mathbf{x},\mathbf{x}') = \phi (\| \mathbf{x} - \mathbf{x}'\| _2 / \ell)$ , where  $\mathbf{x},\mathbf{x}'\in \mathbb{R}^d$ ,  $\ell >0$ , and  $\phi \in C^2 (\mathbb{R})$ , such that  $\phi (z)\to 0$ ,  $\phi '(z)\to 0$ , and  $\phi ''(z)\to 0$  for  $z\rightarrow \infty$ . Let  $p$  and  $q$  be two bimodal mixture distributions satisfying Assumptions 2.1 and 2.2, for any  $\eta \in (0,1)$ . We define  $w^{\star}$  as the optimal mixture weight of  $q$  with respect to the KSD distance, i.e.,  $w^{\star} = \operatorname *{argmin}_{w\in [0,1]}\mathrm{KSD}(\mathbb{P},\mathbb{Q}_w)$ . Then, for  $\mu$  large enough, we have  $|w^{\star} - \frac{1}{2}| < \frac{\eta}{2(1 - \eta)}$ .
Theorem 2.3, proved in Appendix B, states that the weight  $w^{\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  $\mathcal{M}_{s_0}$ , the region of the input space where the score norm is lower than the threshold  $s_0 \geq 0$ , formally  $\mathcal{M}_{s_0} = \{\mathbf{x} \in$
Figure 2: Illustration of Pathology II for the Gaussian mixture of Example 1 ( $d = 2$ ,  $\mu = 2$ ,  $\sigma = 1$ ,  $w = 0.5$ ,  $n = 3000$ ,  $m = 300$ ): many particles are selected around the line  $x^{(1)} = 0$  (left panel), because of the squared first component of the score  $s_p(\mathbf{x})$  along  $x^{(1)}$  (for  $x^{(2)} = 0$  in the right panel).
$\mathbb{R}^d: \| s_p(\mathbf{x}) \|_2 \leq s_0$ . We also introduce an independent and identically distributed (iid) sample  $\mathbf{X}_1, \ldots, \mathbf{X}_m$  of  $\mathbb{P}$ , with  $\mathbb{P}_m$  the associated empirical measure for a positive integer  $m$ , and  $X^{(j)}$  the  $j$ -th component of  $\mathbf{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  $\ell$  in the Stein thinning algorithm.
Theorem 2.4 (KSD spurious minimum). Let  $k_p$  be the Stein kernel associated with the IMQ kernel with  $\ell > 0$ ,  $\beta \in (0,1)$ , and  $c = 1$ . Let  $\{\mathbf{x}_i\}_{i=1}^m \subset \mathcal{M}_{s_0} = \{\mathbf{x} \in \mathbb{R}^d : \| s_p(\mathbf{x}) \|_2 \leq s_0\}$  be a fixed set of points of empirical measure  $\mathbb{Q}_m = \frac{1}{m} \sum_{i=1}^m \delta(\mathbf{x}_i)$ , with  $s_0 \geq 0$  and  $m \geq 2$ . We have  $\mathrm{KSD}^2(\mathbb{P}, \mathbb{Q}_m) < \mathbb{E}[\mathrm{KSD}^2(\mathbb{P}, \mathbb{P}_m)]$ , if the score threshold  $s_0$  and the sample size  $m$  are small enough to satisfy  $m < 1 + (\mathbb{E}[|s_p(\mathbf{X})|_2^2] - s_0^2) / (2\beta d / \ell^2 + 2\beta s_0 / \ell + s_0^2)$ .
Corollary 2.5 (Low KSD samples at density minimum). Let  $k_{p}$  be the Stein kernel associated with the IMQ kernel with  $\ell > 0$ ,  $\beta \in (0,1)$ , and  $c = 1$ . Let  $p$  be a density with at least one local minimum or saddle point. For  $m \geq 2$ , if  $\{\mathbf{x}_i\}_{i=1}^m \subset \mathbb{R}^d$  is a set of points, all located at local minimum or saddle points of  $p$ , then we have  $\mathrm{KSD}^2(\mathbb{P},\mathbb{Q}_m) < \mathbb{E}[\mathrm{KSD}^2(\mathbb{P},\mathbb{P}_m)]$ , if  $m < 1 + \frac{\ell^2}{2\beta d}\mathbb{E}[\|s_p(\mathbf{X})\|_2^2]$ .