Partial derivative of expectation and Stein's lemma

Question

Currently, I am reading a paper about the Gaussian Process in Neural Network [1]. In the solution of the main result in this paper, the author applied Stein's lemma and claimed an equation about the partial derivative of expectation as below ([1], page. 68),

Let $\Psi(w,\tau) = \mathbb{E}_{z \sim \mathcal{N}(0,1)}\psi(g_1,\dots,g_{m-1},w+\tau z)$. By Stein's Lemma, $\Psi(w,\tau)$ is differentiable in $w$, and $$\partial_{w}\Psi(w,\tau) = \tau^{-1}\mathbb{E}_{z \sim \mathcal{N}(0,1)}[z\psi(g_1,\dots,g_{m-1},w+\tau z)],$$

Stein's Lemma: For jointly Gaussian random variable $Z_1$ and $Z_2$ with zero mean, and any function $\phi:\mathcal{R} \rightarrow \mathcal{R}$, where $\mathbb{E}[\phi'(Z_1)]$ and $\mathbb{E}[Z_1\phi(Z_2)]$ exists, we have \begin{equation} \mathbb{E}[Z_1 \phi(Z_2)] = Cov(Z_1,Z_2) \mathbb{E}[\phi'(Z_2)], \end{equation} where $Cov(Z_1,Z_2)$ is the covariance of $Z_1$ and $Z_2$.

As I understand, the author applied Stein's lemma for $z$ and $\phi(g_1,\dots,g_{m-1},w+\tau z)$ to obtain $$\mathbb{E}[z\phi(g_1,\dots,g_{m-1},w+\tau z)] = \tau \mathbb{E}[\phi'(g_1,\dots,g_{m-1},w+\tau z)].$$

However, I couldn't figure out why the author can obtain the equation of the partial derivative in terms of $w$ as above. I am very appreciative if someone explains it to me.

[1]. https://arxiv.org/abs/1910.12478

Answer 1

Stein's lemma (with $Z_1=Z_2$) gives $$\mathbb{E}[z\psi(g_1,\dots,g_{m-1},w+\tau z)] =\text{cov}\,(z,z)\mathbb{E}[\frac{d}{dz}\psi(g_1,\dots,g_{m-1},w+\tau z)]$$ $$\qquad=\tau\frac{d}{dw}\mathbb{E}[\psi(g_1,\dots,g_{m-1},w+\tau z)]\equiv \tau\partial_w\Psi(w,\tau),$$ since $\text{cov}\,(z,z)=1$ and $df(w+\tau z)/dz=\tau df(w+\tau z)/dw$.