Stein's Equation for Multivariate Gaussian Approximation

Why the Covariance Matrix Appears in Stein’s Equation

For a Gaussian target, integration by parts turns expectations involving the random vector into expectations involving derivatives of a test function. In several dimensions, the covariance matrix determines exactly how those derivatives are combined. The resulting differential operator is the multivariate Gaussian version of the normal Stein operator.

Fix a dimension \(d\in\mathbb{N}\) and a symmetric positive semidefinite \(d\times d\) matrix \(\Sigma\). Let \(Z\) have the centered Gaussian distribution \(N_d(0,\Sigma)\). For a twice differentiable function \(f:\mathbb{R}^d\to\mathbb{R}\), write \(\nabla f\) for its gradient and \(\nabla^2 f\) for its Hessian. The matrix contraction in the operator is \(\operatorname{tr}(\Sigma\nabla^2 f)\).

Definition: The centered Gaussian Stein operator with covariance \(\Sigma\) is \(\mathcal{A}_{\Sigma}f(x)=\operatorname{tr}(\Sigma\nabla^2 f(x))-x\cdot\nabla f(x)\). For a test function \(h\), its Stein equation is \(\mathcal{A}_{\Sigma}f(x)=h(x)-\mathbb{E}h(Z)\).

The equation is centered at the target expectation: taking expectations under \(Z\) should make its right-hand side zero. The integration-by-parts identity below explains why the operator has this particular form.

Gaussian Integration by Parts

The covariance identity holds even if \(\Sigma\) is singular. Let \(G\) be a standard Gaussian vector in \(\mathbb{R}^d\), and take a symmetric square root \(A=\Sigma^{1/2}\), so that \(Z=AG\) has covariance \(\Sigma\). For a differentiable function \(g\) with sufficient integrability, standard one-dimensional Gaussian integration by parts applied to each coordinate of \(G\) gives the following result.

Theorem: If \(g:\mathbb{R}^d\to\mathbb{R}\) is continuously differentiable and the terms below are integrable, then, for each \(i\), \(\mathbb{E}[Z_i g(Z)]=\sum_{j=1}^d\Sigma_{ij}\mathbb{E}[\partial_j g(Z)]\).

Proof. Write \(Z=AG\). For each \(i\), standard Gaussian integration by parts gives \[ \mathbb{E}[Z_i g(Z)] =\sum_{k=1}^d A_{ik}\mathbb{E}[G_k g(AG)] =\sum_{k=1}^d\sum_{j=1}^d A_{ik}A_{jk}\mathbb{E}[\partial_j g(AG)]. \] Since \(\sum_k A_{ik}A_{jk}=\Sigma_{ij}\), this is the asserted identity. The calculation uses only the representation \(Z=AG\), so it does not require \(\Sigma\) to be invertible. \(\square\)

Apply the identity with \(g=\partial_i f\), then sum over \(i\). The left side becomes \(\mathbb{E}[Z\cdot\nabla f(Z)]\), while the right side becomes the covariance contraction with the Hessian.

\mathbb{E}[Z\cdot\nabla f(Z)] = \mathbb{E}\!\left[\operatorname{tr}(\Sigma\nabla^2 f(Z))\right].

Corollary: Under the integrability and differentiability conditions needed for the identity, \(\mathbb{E}[\mathcal{A}_{\Sigma}f(Z)]=0\). Thus every solution of the Stein equation has the required target expectation identity.

Worked Example: A Linear Test Function

Worked Example: The Stein solution for a linear function

Let \(h(x)=a\cdot x\), where \(a\in\mathbb{R}^d\). Since \(Z\) is centered, \(\mathbb{E}h(Z)=0\). Set \(f(x)=-a\cdot x\). Its Hessian is zero and its gradient is \(-a\), so

\mathcal{A}_{\Sigma}f(x) =0-x\cdot(-a) =a\cdot x =h(x)-\mathbb{E}h(Z).

Therefore \(f(x)=-a\cdot x\) solves the equation. In dimension one with variance \(1\), taking \(h(x)=x\) gives \(f(x)=-x\); direct substitution gives \(f''(x)-xf'(x)=0-x(-1)=x\).

Worked Example: A Correlated Two-Dimensional Target

Worked Example: The operator on a product

Let \(\Sigma=\begin{pmatrix}1&\rho\\ \rho&1\end{pmatrix}\), where \(-1\leq\rho\leq1\), and let \(f(x_1,x_2)=x_1x_2\). The gradient is \((x_2,x_1)\), and the Hessian is \(\begin{pmatrix}0&1\\1&0\end{pmatrix}\). Hence

\operatorname{tr}(\Sigma\nabla^2 f)=2\rho, \qquad x\cdot\nabla f=x_1x_2+x_2x_1=2x_1x_2,

and therefore \(\mathcal{A}_{\Sigma}f(x)=2\rho-2x_1x_2\). Because \(\mathbb{E}[Z_1Z_2]=\rho\), the expectation of this expression at \(Z\) is \(2\rho-2\rho=0\), as required. In particular, when \(\rho=1\) the operator value is \(2-2x_1x_2\), not \(2-x_1x_2\): both terms in the drift must be included.

Worked Example: A Quadratic Test Function

Worked Example: A centered quadratic

Let \(B\) be a symmetric \(d\times d\) matrix and set \(h(x)=x^{\mathsf T}Bx\). The Gaussian second-moment identity gives \(\mathbb{E}h(Z)=\operatorname{tr}(B\Sigma)\). Choose \(f(x)=-\tfrac12 x^{\mathsf T}Bx\). Then \(\nabla f(x)=-Bx\) and \(\nabla^2 f(x)=-B\), so

\mathcal{A}_{\Sigma}f(x) =-\operatorname{tr}(\Sigma B)+x^{\mathsf T}Bx =h(x)-\mathbb{E}h(Z).

This verifies the equation directly, including its covariance term and its sign. The calculation uses the symmetry of \(B\), which ensures the stated gradient and Hessian.

Constructing Solutions with the Ornstein–Uhlenbeck Semigroup

The preceding examples have polynomial solutions. A useful general construction comes from the Gaussian Ornstein–Uhlenbeck semigroup. For \(t\geq0\), define

P_t h(x)=\mathbb{E}\!\left[ h\!\left(e^{-t}x+\sqrt{1-e^{-2t}}\,Z\right) \right].

As \(t\) grows, the coefficient of \(x\) tends to zero and the argument approaches a draw from the target Gaussian. If \(h\) is twice continuously differentiable with bounded first and second derivatives, define

f_h(x)=-\int_0^\infty\bigl(P_t h(x)-\mathbb{E}h(Z)\bigr)\,dt.

The integral is finite for each fixed \(x\). Indeed, using the same \(Z\) in the two terms and the bounded gradient of \(h\), the absolute difference inside the integral is at most \[ \|\nabla h\|_\infty \left(e^{-t}|x|+\left|1-\sqrt{1-e^{-2t}}\right|\mathbb{E}|Z|\right). \] Both terms are integrable over \(t\in[0,\infty)\).

Theorem: For such an \(h\), the function \(f_h\) solves \(\mathcal{A}_{\Sigma}f_h=h-\mathbb{E}h(Z)\).

Proof. Put \(q=e^{-t}\), \(s=1-e^{-2t}\), and \(Y=qx+\sqrt{s}\,Z\). Differentiating the semigroup gives \(\nabla P_t h(x)=q\,\mathbb{E}\nabla h(Y)\) and \(\nabla^2P_t h(x)=q^2\mathbb{E}\nabla^2 h(Y)\). Differentiating in \(t\), and applying Gaussian integration by parts to the term involving \(Z\), gives \[ \partial_tP_t h(x) =-q\,x\cdot\mathbb{E}\nabla h(Y) +q^2\mathbb{E}\!\left[\operatorname{tr}(\Sigma\nabla^2h(Y))\right] =\mathcal{A}_{\Sigma}P_t h(x). \] Integrating this identity from \(0\) to \(T\) yields \[ \mathcal{A}_{\Sigma}\!\left(-\int_0^T(P_t h-\mathbb{E}h(Z))\,dt\right) =h(x)-P_T h(x). \] As \(T\to\infty\), \(P_T h(x)\to\mathbb{E}h(Z)\). The derivative formulas and bounded derivatives justify passing the operator through the integral defining \(f_h\); thus the displayed identity converges to \(\mathcal{A}_{\Sigma}f_h(x)=h(x)-\mathbb{E}h(Z)\). \(\square\)

Which Derivatives Matter for Wasserstein Bounds?

For a test function \(h\), the Stein equation rewrites a difference of expectations as

\mathbb{E}h(X)-\mathbb{E}h(Z) = \mathbb{E}\!\left[ \operatorname{tr}(\Sigma\nabla^2 f_h(X)) -X\cdot\nabla f_h(X) \right],

whenever \(X\) has the required integrability and \(f_h\) is a solution. This identity identifies the derivative terms that an approximation argument must control: the gradient appears in the drift, while the Hessian appears in the covariance term. In coupling and exchangeable-pair arguments, a Taylor expansion commonly produces a remainder controlled by a Hessian bound.

For smooth tests with bounded Hessian, the semigroup also gives a direct bound. If \(\|\nabla^2h(x)\|_{\mathrm{op}}\leq M\) for every \(x\), then the Hessian formula for \(P_t h\) implies \(\|\nabla^2P_t h(x)\|_{\mathrm{op}}\leq e^{-2t}M\). Integrating in \(t\) gives the following result.

Derivative bound: Under the preceding smoothness assumptions, \(\|\nabla^2 f_h(x)\|_{\mathrm{op}}\leq M/2\) for every \(x\).

Proof. Differentiate the integral defining \(f_h\) twice and use the semigroup Hessian formula: \[ \|\nabla^2 f_h(x)\|_{\mathrm{op}} \leq\int_0^\infty e^{-2t}M\,dt =\frac{M}{2}. \] The integral evaluates to \(1/2\), proving the bound. \(\square\)

Wasserstein distance uses a supremum over Lipschitz test functions, which need not have bounded classical second derivatives. For the standard Gaussian target \(\Sigma=I\), a useful bound for smooth \(L\)-Lipschitz \(h\) follows by differentiating the Gaussian smoothing and applying integration by parts: \[ \partial_{ij}P_t h(x) =\frac{e^{-2t}}{\sqrt{1-e^{-2t}}} \,\mathbb{E}\!\left[Z_j\,\partial_i h\!\left(e^{-t}x+\sqrt{1-e^{-2t}}\,Z\right)\right]. \] For unit vectors \(u,v\), the absolute value of the corresponding Hessian bilinear form is at most \[ \frac{e^{-2t}}{\sqrt{1-e^{-2t}}}\,L\,\mathbb{E}|v\cdot Z| \leq \frac{e^{-2t}}{\sqrt{1-e^{-2t}}}\,L\sqrt{\frac{2}{\pi}}. \] Integrating over \(t\) gives \(\|\nabla^2 f_h(x)\|_{\mathrm{op}}\leq L\sqrt{2/\pi}\), since \(\int_0^\infty e^{-2t}(1-e^{-2t})^{-1/2}\,dt=1\). This illustrates why Hessian estimates for Stein solutions are central to Wasserstein bounds; nonsmooth Lipschitz tests can be handled by smooth approximation.

A Common Pitfall: Mixing Covariance Conventions

The operator here is written in the original coordinates of \(X\), with target covariance \(\Sigma\). Its diffusion term is \(\operatorname{tr}(\Sigma\nabla^2 f)\), and its drift is \(-x\cdot\nabla f\). Replacing \(\Sigma\) by the identity without also changing coordinates gives the wrong equation when the target is not standard Gaussian. A quick check is to apply the proposed operator to a quadratic: its Gaussian expectation must be zero.

A second useful check is sign and centering. The right side is \(h-\mathbb{E}h(Z)\), not its negative. For \(h(x)=a\cdot x\), the solution \(f(x)=-a\cdot x\) verifies the sign immediately. For \(h(x)=x_1x_2\) under a correlated target, both drift contributions must be included, so the drift is \(2x_1x_2\).

Check Your Understanding

Use the operator and examples above to check that you can move between the covariance identity, the Stein equation, and derivative bounds.

What integration-by-parts identity gives the covariance term in \(\mathcal{A}_{\Sigma}\)?
For \(h(x)=a\cdot x\), verify the Stein equation using \(f(x)=-a\cdot x\).
For the correlated two-dimensional example, why is the drift of \(f(x)=x_1x_2\) equal to \(2x_1x_2\)?
What Hessian bound follows from the semigroup when \(h\) has Hessian operator norm at most \(M\)?
Which derivatives of the Stein solution occur in the expectation identity used for Wasserstein approximation?