Stein’s Identity with General Covariance

Why the Covariance Matrix Matters

For a scalar normal variable, integration by parts links a centered observation to a derivative of the function evaluated at that observation. In a multivariate normal model, the same principle holds, but the covariance matrix determines how each coordinate interacts with derivatives in every coordinate. When the coordinates are correlated, the resulting identity is not generally a sum involving only derivatives in matching coordinates.

This tutorial derives the identity with explicit indices, then checks it on correlated Gaussian examples. Keeping the Jacobian convention visible is important: it is what makes the compact trace formula agree with the componentwise calculation.

Definition: Let \(X=(X_1,\ldots,X_d)^\mathsf{T}\sim N(\mu,\Sigma)\), where \(\mu\in\mathbb{R}^d\) and \(\Sigma\) is a symmetric positive-definite \(d\times d\) covariance matrix. For a differentiable vector-valued function \(g=(g_1,\ldots,g_d)^\mathsf{T}\), use the Jacobian convention \[ (\nabla g(x))_{ij}=\frac{\partial g_i(x)}{\partial x_j}. \] Thus rows index components of \(g\), and columns index coordinates of \(x\).

We use a simple sufficient set of regularity conditions. Assume each \(g_i\) is continuously differentiable on \(\mathbb{R}^d\), and that for every \(i,j,k\), \(E[|g_i(X)|]\), \(E[|(X_j-\mu_j)g_i(X)|]\), and \(E[|\partial g_i(X)/\partial x_k|]\) are finite. These assumptions ensure the expectations in the identity exist and allow the integration-by-parts argument below to be justified with smooth cutoffs.

The Componentwise Identity

The density of \(X\) is \[ \varphi_{\mu,\Sigma}(x) =(2\pi)^{-d/2}|\Sigma|^{-1/2} \exp\!\left\{-\frac12(x-\mu)^\mathsf{T}\Sigma^{-1}(x-\mu)\right\}. \] Differentiating its logarithm gives \[ \frac{\partial}{\partial x_k}\varphi_{\mu,\Sigma}(x) =-\big[\Sigma^{-1}(x-\mu)\big]_k\varphi_{\mu,\Sigma}(x). \] Equivalently, \(x_j-\mu_j=-\sum_k\Sigma_{jk}\,\partial\log\varphi_{\mu,\Sigma}(x)/\partial x_k\). This score identity turns a centered coordinate into a weighted combination of derivatives with respect to all coordinates.

Theorem: Under the conditions above, for every pair of indices \(i,j\), \[ E[(X_j-\mu_j)g_i(X)] =\sum_{k=1}^d\Sigma_{jk}\,E\!\left[\frac{\partial g_i(X)}{\partial x_k}\right]. \] Consequently, \[ E[(X-\mu)^\mathsf{T}g(X)] =E[\operatorname{tr}(\Sigma\nabla g(X))]. \]

Proof. First take a continuously differentiable cutoff function \(\eta_R\) that equals \(1\) on the ball of radius \(R\), vanishes outside the ball of radius \(2R\), and satisfies \(|\partial\eta_R/\partial x_k|\leq C/R\) for a constant \(C\) independent of \(R\). The function \(g_i\eta_R\) has compact support, so ordinary integration by parts has no boundary term. Using the derivative of the Gaussian density,

\begin{aligned} \int (x_j-\mu_j)g_i(x)\eta_R(x)\varphi_{\mu,\Sigma}(x)\,dx &=-\sum_{k=1}^d\Sigma_{jk} \int g_i(x)\eta_R(x)\frac{\partial}{\partial x_k} \varphi_{\mu,\Sigma}(x)\,dx\\ &=\sum_{k=1}^d\Sigma_{jk} \int \frac{\partial}{\partial x_k} \big(g_i(x)\eta_R(x)\big)\varphi_{\mu,\Sigma}(x)\,dx\\ &=\sum_{k=1}^d\Sigma_{jk} \int \left( \frac{\partial g_i(x)}{\partial x_k}\eta_R(x) +g_i(x)\frac{\partial\eta_R(x)}{\partial x_k} \right)\varphi_{\mu,\Sigma}(x)\,dx. \end{aligned}

As \(R\to\infty\), the left side converges to \(E[(X_j-\mu_j)g_i(X)]\) by dominated convergence, using the assumed absolute integrability. In the first term on the final line, dominated convergence applies to the derivative of \(g_i\). The remaining cutoff term has absolute value at most \(C E[|g_i(X)|]/R\), apart from the fixed covariance factor, so it tends to zero. This proves the componentwise identity.

Now sum that identity over \(i=j\). The left side is \(E[(X-\mu)^\mathsf{T}g(X)]\), and the right side is \[ \sum_{i=1}^d\sum_{k=1}^d\Sigma_{ik} E\!\left[\frac{\partial g_i(X)}{\partial x_k}\right]. \] By the stated Jacobian convention and symmetry of \(\Sigma\), this double sum is \(E[\operatorname{tr}(\Sigma\nabla g(X))]\): indeed, \(\operatorname{tr}(\Sigma\nabla g)=\sum_{a,b}\Sigma_{ab}\,\partial g_b/\partial x_a\), and relabeling indices gives the displayed sum. \(\square\)

Reading the Trace When Coordinates Are Correlated

Expanding the right side makes the role of covariance explicit: \[ \operatorname{tr}(\Sigma\nabla g(X)) =\sum_{i=1}^d\sum_{k=1}^d \Sigma_{ik}\frac{\partial g_i(X)}{\partial x_k}. \] A diagonal covariance matrix leaves only matching-coordinate derivatives, each weighted by that coordinate’s variance. Nonzero off-diagonal covariances allow cross-coordinate derivative terms. A particular such contribution is zero if its covariance is zero or if the corresponding derivative has zero expectation; correlation alone does not guarantee a nonzero contribution.

The vector version for a scalar function is often useful alongside the theorem.

Corollary: If \(h:\mathbb{R}^d\to\mathbb{R}\) is continuously differentiable and satisfies the corresponding integrability conditions, then \[ E[(X-\mu)h(X)]=\Sigma E[\nabla h(X)], \] where \(\nabla h\) is the column vector of partial derivatives.

To verify the corollary, apply the componentwise identity with \(g_i\) replaced by the scalar function \(h\), for each coordinate \(j\). The \(j\)th entry on the left is \(E[(X_j-\mu_j)h(X)]\); the \(j\)th entry on the right is \(\sum_k\Sigma_{jk}E[\partial h(X)/\partial x_k]\), exactly the \(j\)th entry of \(\Sigma E[\nabla h(X)]\). Thus the vector equality follows component by component. \(\square\)

Worked Applications

Worked Example: An Affine Function with Correlated Coordinates

Let \(X\sim N(\mu,\Sigma)\) in two dimensions, with \[ \Sigma=\begin{pmatrix}4&1\\1&3\end{pmatrix}, \qquad g(x)=Ax,\qquad A=\begin{pmatrix}1&2\\-1&0\end{pmatrix}. \] The covariance matrix is positive definite because its first leading principal minor is \(4>0\) and its determinant is \(12-1=11>0\). Find both sides of the identity.

The Jacobian is \(\nabla g=A\), so the right side is \(\operatorname{tr}(\Sigma A)\). Multiplication gives \[ \Sigma A= \begin{pmatrix}4&1\\1&3\end{pmatrix} \begin{pmatrix}1&2\\-1&0\end{pmatrix} = \begin{pmatrix}3&8\\-2&2\end{pmatrix}, \qquad \operatorname{tr}(\Sigma A)=3+2=5. \] For the left side, write \(X=\mu+Z\), where \(E[Z]=0\) and \(E[ZZ^\mathsf{T}]=\Sigma\). Then \[ E[Z^\mathsf{T}AX] =E[Z^\mathsf{T}A\mu]+E[Z^\mathsf{T}AZ] =0+\operatorname{tr}(A\Sigma). \] The first term is zero since \(E[Z]=0\). For the second, expanding coordinates gives \(E[Z^\mathsf{T}AZ]=\sum_{i,j}A_{ij}E[Z_iZ_j]=\sum_{i,j}A_{ij}\Sigma_{ij}=\operatorname{tr}(A\Sigma)\), using symmetry of \(\Sigma\). Here \[ A\Sigma= \begin{pmatrix}6&7\\-4&-1\end{pmatrix}, \qquad \operatorname{tr}(A\Sigma)=6-1=5. \] Both sides equal \(5\). The answer does not depend on \(\mu\), because the centered vector \(Z\) has mean zero.

Worked Example: A Cross-Coordinate Derivative

Let \(X=(X_1,X_2)^\mathsf{T}\sim N(\mu,\Sigma)\), where \[ \Sigma=\begin{pmatrix}2&0.6\\0.6&1\end{pmatrix}, \qquad g(x)=\begin{pmatrix}x_2\\0\end{pmatrix}. \] The covariance matrix is positive definite since its determinant is \(2-0.36=1.64>0\). Evaluate both sides.

The left side is \[ E[(X_1-\mu_1)X_2+(X_2-\mu_2)0] =\operatorname{Cov}(X_1,X_2)+\mu_2E[X_1-\mu_1] =0.6+0=0.6. \] The Jacobian and its covariance-weighted trace are \[ \nabla g(x)=\begin{pmatrix}0&1\\0&0\end{pmatrix}, \qquad \Sigma\nabla g(x)= \begin{pmatrix}0&2\\0&0.6\end{pmatrix}, \qquad \operatorname{tr}(\Sigma\nabla g(x))=0.6. \] The identity therefore gives \(0.6=0.6\). The contribution comes from the off-diagonal covariance multiplying \(\partial g_1/\partial x_2=1\); no derivative of \(g_2\) contributes in this example.

Worked Example: A Nonlinear Function of a Correlated Pair

Suppose \(X\sim N(\mu,\Sigma)\) with \(\mu=(1,2)^\mathsf{T}\) and \[ \Sigma=\begin{pmatrix}4&1\\1&3\end{pmatrix}, \qquad g(x)=\begin{pmatrix}x_2^2\\0\end{pmatrix}. \] The determinant of \(\Sigma\) is \(11>0\), so this is a nonsingular Gaussian model. Compute the two sides.

For the left side, \[ E[(X-\mu)^\mathsf{T}g(X)] =E[(X_1-1)X_2^2]. \] Write \(X_2=2+W\), where \(W=X_2-2\). Then \[ E[(X_1-1)X_2^2] =E[(X_1-1)(4+4W+W^2)]. \] The constant term has expectation zero. The linear term is \(4E[(X_1-1)W]=4\operatorname{Cov}(X_1,X_2)=4\). The final term is zero: a centered bivariate normal distribution is unchanged by \((X_1-1,W)\mapsto(-(X_1-1),-W)\), while \((X_1-1)W^2\) changes sign. Hence the left side is \(4\).

For the right side, \(\partial g_1/\partial x_2=2x_2\), and all other entries of the Jacobian are zero. Thus \[ \operatorname{tr}(\Sigma\nabla g(X)) =\Sigma_{12}(2X_2)=2X_2, \qquad E[\operatorname{tr}(\Sigma\nabla g(X))]=2E[X_2]=2(2)=4. \] The two sides agree. In general, this particular calculation gives \(2\operatorname{Cov}(X_1,X_2)\mu_2\); the numerical choices make that value \(2(1)(2)=4\).

Common Mistake and Practical Interpretation

A frequent error is to use the independent-coordinate formula after noticing only that \(X\) is Gaussian. For a general covariance matrix, replacing the trace by \(\sum_i\Sigma_{ii}\,\partial g_i/\partial x_i\) discards the off-diagonal terms. In the second example, that incorrect diagonal-only expression is zero, even though the correct identity gives \(0.6\).

The correction is to expand \(\operatorname{tr}(\Sigma\nabla g)\) using the Jacobian convention before simplifying. Off-diagonal covariances make cross-coordinate derivative terms possible, but they do not force those terms to be nonzero: a term vanishes if its covariance is zero, or if the corresponding derivative has zero expectation. This distinction matters in applications of Stein’s identity because it separates the structure permitted by the covariance matrix from the actual contribution of a chosen function.

The general-covariance identity is the appropriate integration-by-parts tool whenever Gaussian observations have dependent coordinates. Its componentwise form is often easiest to apply; its trace form is compact once the indices have been checked.

Check Your Understanding

Use the componentwise formula and keep the row-and-column convention for the Jacobian in view.

For a differentiable vector function \(g\), write the \((i,j)\) componentwise identity and identify which covariance entries multiply derivatives of \(g_i\).
If \(\Sigma\) is diagonal, simplify \(\operatorname{tr}(\Sigma\nabla g(X))\). What changes when \(\Sigma\) has nonzero off-diagonal entries?
In the second worked example, which derivative and covariance entry produce the nonzero trace?
Can a cross-coordinate term be zero even when its covariance is nonzero? Give both conditions under which that term vanishes.
For the affine function \(g(x)=Ax\), explain why the left side is \(\operatorname{tr}(A\Sigma)\), and why this equals \(\operatorname{tr}(\Sigma A)\).