The Polynomial Stein Discrepancy for Assessing Moment Convergence

Thank you for your time reviewing our manuscript and for your valuable comments and suggestions. We will create a log/log plot, replace "Theorem 3.1" with "Corollary 3.1." and mention that property of the graph Stein discrepancy.

(Lack of) confidence about power of PSD in the wild

We have highlighted in the paper the strong theoretical performance on Gaussian targets and the strong empirical performance of PSD on non-Gaussian targets. During this review period, we have also highlighted issues with PSD for distributions where moments do not exist (e.g. Cauchy). However, the development of theoretical results for the non-Gaussian setting represents a substantial undertaking that would be better considered in a future paper.

When to use PSD over KSD

Broadly we suggest using PSD when one wishes to obtain high statistical power for detecting discrepancies in moments and the target is "close to Gaussian" (with checks that we have expanded on elsewhere). Fast but biased sampling methods like SG-MCMC, where the discrepancy is often in the variance and the target is often close to Gaussian, are good applications that we highlight in the paper. The linear-time complexity of PSD is particularly helpful for large $n$.

How do practitioners know the posterior is Gaussian?

We propose a diagnostic plot in the discussion. If the gradients are available in analytic form, rather than indirectly through automatic differentiation procedures, then one can also determine which expectations are being tracked by analysing the system of linear equations in equation (14) of Appendix B. We do this for the Rosenbrock target in Appendix C.6, and we are able to determine exactly which expectations are being tracked. We will suggest this technique in the discussion to diagnose situations where the PSD might not behave as expected.

Simpler ways to learn the moments if the target is Gaussian

Estimating the first $r$ moments need not be our only goal. However, checking for discrepancies in moments can be important in practice. For example, in the common application of SG-MCMC algorithms, poor step-sizes can lead to under- or over-estimated variance. Selecting the step-size based on matching the second-order moment can lead to a reasonable posterior approximation, as explained in the paper.

Moment-matching guarantees for a larger class of target distributions

We can guarantee PSD is assessing moments for Gaussian targets because $\nabla \log p(x)$ is a linear function in that context, so the Stein operator applied to an $r$th order polynomial is itself an $r$th order polynomial. In this sense, the closer $\nabla \log p(x)$ is to a linear function, the closer the PSD is to assessing discrepancies in the first moments.

Sensitivity of PSD to mean-shift

We are investigating this further by observing the sensitivity of PSD to transformations. The moment-tracking property of PSD will continue to hold due to similar reasoning as Appendix D.

Topology of PSD and properties when $PSD(P, Q_n)\to 0$

While higher order PSD becomes increasingly complex, we can explicitly consider the case of $P=N(\mu_P, \Sigma_P)$ and $r = 1$. Here, $PSD(P,Q_n) = ||n^{-1}\sum_{i=1}^n\Sigma_P^{-1}(\mu_P-X_i) ||_2 = ||\Sigma_P^{-1}(\mu_P-\bar{X}) ||_2$ is measuring the Mahalanobis distance between $\mu_P$ and the sample mean $\bar{X}$. From the law of large numbers as $n\rightarrow \infty$, $\bar{X}\rightarrow \mu_Q$ and consequently, $PSD(P,Q)\rightarrow \Sigma_P^{-1}(\mu_P-\mu_Q)$ is zero if and only if the moments of $P$ and $Q$ match. Assuming iid draws from $Q$, the CLT gives $\sqrt{n}\Sigma_P^{-1}(\bar{X}_Q-\mu_q)\sim N(0, \Sigma_P^{-1}\Sigma_Q\Sigma_P^{-T})$ so $PSD = ||Y +\Sigma_P^{-1}(\mu_q-\mu_P)||_2$ where $Y\sim N(0, n^{-1}\Sigma_P^{-1}\Sigma_Q\Sigma_p^{-T})$. The PSD behaves as the square root of a noncentral $\chi^2$ distribution.

This result for $r=1$ serves as an illustrative example for iid samples and Gaussian targets. More practically useful theory, including for correlated samples, would be substantially more challenging and we will mention it as future work in the paper.

Simpler formulation for identifying specific moments where discrepancies occur

If the target is Gaussian with diagonal covariance, then for the form of PSD we consider, the individual terms in the PSD correspond exactly to the moments of interest, as explained in the paper. This makes identification of problematic moments simpler than for alternative polynomial kernels.

Benefits of other spanning polynomials?

Unfortunately the benefits of orthogonal polynomials such as Chebyshev polynomials would likely be lost with PSD. It is not clear that one obtains a Chebyshev polynomial after applying a Langevin Stein operator to a Chebyshev polynomial.

Thank you for your time reviewing our manuscript and for your valuable comments and suggestions.

On the BvM limit

Proposition 3.2 explicitly states the Gaussianity assumption without mentioning the BvM limit, but we agree that this can be clarified in other parts of the text. We will reword the text in the abstract and the introduction of the paper. We will replace the text "we prove that it detects differences in the first $r$ moments in the Bernstein-von Mises limit" in the abstract with "we prove that it detects differences in the first $r$ moments for Gaussian targets."

We opt to keep the text about the BvM as a remark after our theoretical result. This is relevant to the practical application of the method because a major application of KSD and related methods is in fast but biased Monte Carlo methods, such as SG-MCMC, where posterior distributions are often close to Gaussian. Bardenet et al (2017) highlight this in their review paper on MCMC for tall data, stating ``we emphasize that we have only been able so far to propose subsampling-based methods which display good performance in scenarios where the Bernstein-von Mises approximation of the target posterior distribution is excellent. It remains an open challenge to develop such methods in scenarios where the Bernstein-von Mises approximation is poor."

Method is based on comparing high moments

As mentioned in Section 3.3, lower-order moments are often of interest and they can be where discrepancies are most likely to appear. We give the example of SG-MCMC, where large step-sizes lead to over-estimation of the posterior variance and small step-sizes can lead to under-estimation of the posterior variance. We therefore do not necessarily agree that the method is based on comparing high moments.

Rosenbrock target context

We will explain that it is challenging because there is high probability mass in a narrow, curved region with complex dependencies and different scales across dimensions. We will add references to the original work (Rosenbrock, 1960), the probabilistic version (Goodman and Weare, 2010) and a paper talking about its complexity (Pagani, 2022).

Gaunt (2020) and relations to theoretical underpinnings

Thank you for the reference. Since the polynomial functions are not Lipschitz, we cannot directly compute the solution of Stein's equation and establish bounds on the derivative. The suggested paper seems to describe a technique to establish the required bounds for unbounded test functions which seems suitable for the PSD case. However, it is not clear how to establish this for our formulation of PSD. Further, the main goal of the paper seems to focus on establishing the closeness of the distribution of $g(W)$ to the distribution of $g(Z)$, where $W$ is a sequence of random variables approximating a standard Gaussian. We are unsure of the relevance of this continuous mapping theorem to our problem - what would be a possible interpretation of the function $g$?

Initial results on the convergence of PSD for a simple setting are given in response to Reviewer JxPo.

Relationship to Xu et al (2022)

The resampling method is indeed similar to the method we have considered. Following Liu et al (2016), we are able to establish the asymptotic exactness of our bootstrap procedure. Theoretical guarantees for resampling under a normal approximation could be a valuable reference to extend the method. Upon our preliminary review of that paper, we feel that the results in Appendix $B.1$ are of interest and we will cite this as a potential extension for the bootstrap procedure. Please let us know if you were referencing a different theorem or section of the paper.

What if moments do not exist?

We believe this could affect things in two ways. First, we might not have the critical zero expectation property for the Stein operator applied to the polynomial, hence we might not expect that PSD $= 0$ when $P = Q$. Second, we will not be assessing moments exactly because the gradients of the log target will not be linear in $x$. We will mention limitations around moments not existing in the discussion.

Dependence on Langevin Stein operator

The advantage of using the Langevin Stein operator is that when it is applied to monomials for a Gaussian target then we obtain monomials back, and therefore we are tracking convergence in moments. This may not hold true for other Stein operators. Nevertheless, applying diffusion Stein operators and other Stein operators, such as those considered in Kanagawa et al (2022), to different types of polynomials is an interesting question and can be studied further in future work. We will mention this in the discussion.

Anonymous Github site

The full set of code has been provided as a zipped folder in the supplementary material for (optional) review. We would prefer to use the first author's Github account if the paper is accepted to encourage appropriate attribution.

Thank you for your time reviewing our manuscript and for your valuable comments and suggestions. We will correct the typographical error.

What if I wanted to test against a specific set of functions other than those given in the paper? Is there any other example of this generating an interesting set of statistics?

It is possible to use functions other than polynomials, but the challenge with doing this is that the application of the Langevin Stein operator, which is required for the zero expectation property, modifies the form of the function. This means we may not obtain interpretable forms for other functions like we do for polynomial functions and Gaussian distributions. As an example, if one were to consider $g(x) = \sin(x)$, then after the application of a second-order Langevin Stein operator for a one-dimensional unit Gaussian target, the function would become $\mathcal{A}_x^{(2)} g(x) = -\sin(x) - x \cos(x)$. A discrepancy based on $g(x) = \sin(x)$ would therefore determine whether $\mathbb{E}_Q[\sin(x) + x \cos(x)] = \mathbb{E}_P[\sin(x) + x \cos(x)]$, which is not the same as the desired function $\sin(x)$.

For instance, if I wanted to test against the moments after applying some transform to my random variable, could I possibly concoct interesting tests this way?

This is an interesting question. Appendix D shows that applying an invertible linear transformation does not change the fact that, in the Gaussian case, the test is for the moments up to order $r$ of the original distribution. However, the statistical power can be affected by the choice of transformation. In particular, such a transformation may improve power when the posterior variances differ substantially in scale. It would be interesting to consider other transformations as well, and we will mention this in the discussion.

Following up on this and a comment by Reviewer JxPo on the effect of a mean shift on the PSD statistic, we aim to empirically investigate the effect of a mean shift during the review period.

My main concern is that the paper could have benefited from further exploration of some of the ideas contained therein.

We have further explored and commented on several topics as part of this review. We would be happy to consider further exploring other aspects of the paper if you would like to provide further details.

Thank you for your time reviewing our manuscript and for your valuable comments and suggestions.

We will correct the typographical issues and the appendix as suggested.

When does KSD outperform PSD? In Figure 1, there are no such examples, and in Table 1 setting $r>1$ is sufficient to match KSD. What are some practical examples of non-moment-based perturbations that KSD can detect but not PSD?

Probability distributions are completely determined by their moment generating functions, provided it exists (a necessary condition is that all moments are finite). Therefore, we can potentially expect the KSD to outperform PSD for target distributions lacking well-defined moments (e.g. those with heavy tails). We comment on this further in response to Reviewer 8qrD. We can also expect KSD to outperform PSD of order $r$ when the discrepancies lie in moments higher than the $r$th order moment. For example, KSD outperforms PSD with $r=1,2,3$ for the student-t and Laplace examples in Figure 1 since the discrepancy is in the kurtosis ($4$th order moment). We will add details of when we can expect KSD to outperform PSD to the discussion.