$L^{2}(p)\cap\mathcal{F}_{0}(p)$, and then

$\mathbb{E}_{x^{\prime}\sim p}\partial_{\theta}{\bf
f}(x^{\prime},\theta(0))\cdot{\bf
s}_{p}(x^{\prime})=-\mathbb{E}_{x^{\prime}\sim
p}\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf f}(x^{\prime},\theta(0)).$

As a result, we have

$\mathbb{E}_{x^{\prime}\sim p}{\bf K}_{0}(x,x^{\prime})\circ{\bf
s}_{p}(x^{\prime})=-\mathbb{E}_{x^{\prime}\sim p}\nabla_{x^{\prime}}\cdot{\bf
K}_{0}(x,x^{\prime}).$

By comparing to (38), we have

$\displaystyle{\partial_{t}}(\hat{\bf u}-\bar{\bf u})(x,t)$
$\displaystyle=-\mathbb{E}_{x^{\prime}\sim\hat{p}}\hat{\bf
K}_{t}(x,x^{\prime})\circ\left(\lambda\hat{\bf u}(x^{\prime},t)-{\bf
s}_{q}(x^{\prime})\right)+\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{0}(x,x^{\prime})\circ(\lambda\bar{\bf u}(x^{\prime},t)-{\bf
s}_{q}(x^{\prime}))$
$\displaystyle~{}~{}~{}~{}~{}+\mathbb{E}_{x^{\prime}\sim\hat{p}}\nabla_{x^{\prime}}\cdot\hat{\bf
K}_{t}(x,x^{\prime})-\mathbb{E}_{x^{\prime}\sim p}\nabla_{x^{\prime}}\cdot{\bf
K}_{0}(x,x^{\prime}),$ $\displaystyle=①+②+③+④+⑤,$ (71)

where

$\displaystyle①$ $\displaystyle:=-\mathbb{E}_{x^{\prime}\sim\hat{p}}\hat{\bf
K}_{t}(x,x^{\prime})\circ(\lambda\hat{\bf u}(x^{\prime},t)-\lambda\bar{\bf
u}(x^{\prime},t))$ $\displaystyle②$
$\displaystyle:=-\mathbb{E}_{x^{\prime}\sim\hat{p}}(\hat{\bf
K}_{t}(x,x^{\prime})-{\bf K}_{0}(x,x^{\prime}))\circ(\lambda\bar{\bf
u}(x^{\prime},t)-{\bf s}_{q}(x^{\prime}))$ $\displaystyle③$
$\displaystyle:=-(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p}){\bf K}_{0}(x,x^{\prime})\circ(\lambda\bar{\bf u}(x^{\prime},t)-{\bf
s}_{q}(x^{\prime}))$ $\displaystyle④$
$\displaystyle:=\mathbb{E}_{x^{\prime}\sim\hat{p}}(\nabla_{x^{\prime}}\cdot\hat{\bf
K}_{t}(x,x^{\prime})-\nabla_{x^{\prime}}\cdot{\bf K}_{0}(x,x^{\prime}))$
$\displaystyle⑤$
$\displaystyle:=(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim{p}})\nabla_{x^{\prime}}\cdot{\bf
K}_{0}(x,x^{\prime}).$

We will analyze each of the five terms respectively toward deriving a bound of
$\|\lambda(\hat{\bf u}-\bar{\bf u})(\cdot,t)\|_{p}$ as has been done in the
proof of Proposition 4.5(ii). We restrict to when $t\leq(\frac{r}{2\|{\bf
f}^{*}\|_{p}})^{2}\lambda$ by default, and consider being under the good event
in Lemma 4.8 assuming that $n>n_{\lambda}$ (by (69), this ensures that $n$ is
greater than the large $n$ threshold integer $n_{6}$ required by (C6)), where
we have $\hat{\theta}(t)\in B_{r}$ and

$\|\hat{\theta}(t)-{\theta}(0)\|\leq\sqrt{\frac{t}{\lambda}}\|{\bf
f}^{*}\|_{p}.$ (72)

In the below, we may omit the dependence on $t$ in notation when there is no
confusion, e.g., we write $\lambda\hat{\bf u}(x^{\prime},t)-{\bf
s}_{q}(x^{\prime})$ as $(\lambda\hat{\bf u}-{\bf s}_{q})(x^{\prime})$.

Bound of $\|⑤\|$: By definition, $⑤=\langle\partial_{\theta}{\bf
f}(x,{\theta}(0)),(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim{p}})\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\rangle_{\Theta}$, and thus

$\|⑤\|\leq\|\partial_{\theta}{\bf
f}(x,{\theta}(0))\|\|(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim{p}})\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\|.$

By (C1), $\|\partial_{\theta}{\bf f}(x,{\theta}(0))\|\leq L_{1}$; To bound
$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\nabla_{x}\cdot\partial_{\theta}{\bf
f}(x,{\theta}(0))\|$, we utilize the matrix Bernstein inequality (see, e.g.,
[37, Theorem 6.1.1], reproduced as Lemma A.2): The matrix
$\nabla_{x}\cdot\partial_{\theta}{\bf f}(x_{i},{\theta}(0))$ is of size
$1$-by-$M_{\Theta}$, and the operator norm is bounded by $L_{3}$ by Assumption
5(C3). By Lemma A.2, let $n_{5}>2$ be an integer s.t. $n>n_{5}$ implies

$(10\log n+\log(2M_{\Theta}))/n<(3/2)^{2},$ (73)

then when $n>n_{5}$, there is a good event $E_{5}$ that happens w.p. $\geq
1-n^{-10}$ under which

$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\nabla_{x}\cdot\partial_{\theta}{\bf
f}(x,{\theta}(0))\|\leq 4\sqrt{11}L_{3}\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}.$

As a result, when $n>n_{5}$ and under $E_{5}$,

$\|⑤\|\leq 4\sqrt{11}L_{1}L_{3}\sqrt{\frac{\log n+\log M_{\Theta}}{n}}.$ (74)

Bound of $\|④\|$: By definition,

$\displaystyle\nabla_{x^{\prime}}\cdot\hat{\bf
K}_{t}(x,x^{\prime})-\nabla_{x^{\prime}}\cdot{\bf K}_{0}(x,x^{\prime})$
$\displaystyle=\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\rangle_{\Theta}$
$\displaystyle~{}~{}~{}-\langle\partial_{\theta}{\bf
f}(x,{\theta}(0)),\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\rangle_{\Theta}$
$\displaystyle=\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t))-\partial_{\theta}{\bf
f}(x,{\theta}(0)),\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\rangle_{\Theta}$
$\displaystyle~{}~{}~{}+\langle\partial_{\theta}{\bf
f}(x,{\theta}(0)),\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))-\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\rangle_{\Theta}.$

By that $\hat{\theta}(0),\,\hat{\theta}(t)\in B_{r}$ and (C1)(C2)(C3)(C4),

$\|\nabla_{x^{\prime}}\cdot\hat{\bf
K}_{t}(x,x^{\prime})-\nabla_{x^{\prime}}\cdot{\bf
K}_{0}(x,x^{\prime})\|\leq(L_{2}L_{3}+L_{1}L_{4})\|\hat{\theta}(t)-{\theta}(0)\|.$

Together with (72), we have that

$\|④\|\leq(L_{2}L_{3}+L_{1}L_{4})\sqrt{{t}/{\lambda}}\|{\bf f}^{*}\|_{p}.$
(75)

Bound of $\|③\|$: By definition, $③=-\langle\partial_{\theta}{\bf
f}(x,{\theta}(0)),(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}(\lambda\bar{\bf
u}(x^{\prime},t)-{\bf s}_{q}(x^{\prime}))\rangle_{\Theta}$, thus

$\displaystyle\|③\|$ $\displaystyle\leq\|\partial_{\theta}{\bf
f}(x,{\theta}(0))\|\|(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}(\lambda\bar{\bf
u}(x^{\prime},t)-{\bf s}_{q}(x^{\prime}))\|\leq L_{1}\|ⓐ_{3}-ⓑ_{3}\|,$

where

$\displaystyle ⓐ_{3}$
$\displaystyle:=(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}\lambda\bar{\bf
u}(x^{\prime},t),$ $\displaystyle ⓑ_{3}$
$\displaystyle:=(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}{\bf
s}_{q}(x^{\prime}).$

For $ⓑ_{3}$, we need to bound
$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\|$. At each $x$, ${\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))$ is a $1$-by-$M_{\Theta}$
vector, and the norm is bounded by $L_{1}b_{q}$ uniformly on $\mathcal{X}$ by
(C6). Applying Lemma A.2 and consider when $n>n_{5}$ which implies $(10\log
n+\log(1+M_{\Theta}))/n<(3/2)^{2}$, there is a good event $E_{b}$ that happens
w.p. $\geq 1-n^{-10}$ under which

$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf
f}(x,{\theta}(0))\|<2\sqrt{11}L_{1}b_{q}\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}.$ (76)

This shows that when $n>n_{5}$ and under $E_{b}$, $\|ⓑ_{3}\|$ is bounded by
the r.h.s. of (76).

For $ⓐ_{3}$, we will bound the deviation to be $\tilde{O}(n^{-1/2})$ uniformly
for all $t$. The observation is that though $\bar{\bf u}(x,t)$ changes over
$t$, the concentration argument of the sample average only involves the kernel
${\bf K}_{0}(x,x^{\prime})$. Specifically, by definition,

$\displaystyle\lambda\bar{\bf u}(x,t)=-\lambda\int_{0}^{t}\mathbb{E}_{y\sim
p}[{\bf K}_{0}(x,y)](\lambda\bar{\bf u}(y,s)-{\bf f}^{*}(y))ds,$ (77)

and then

$ⓐ_{3}=-\lambda\int_{0}^{t}\Big{[}(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\Big{]}\mathbb{E}_{y\sim p}\partial_{\theta}{\bf
f}(y,{\theta}(0))^{T}(\lambda\bar{\bf u}(y,s)-{\bf f}^{*}(y))ds.$

Note that $\|\mathbb{E}_{y\sim p}\partial_{\theta}{\bf
f}(y,{\theta}(0))^{T}(\lambda\bar{\bf u}(y,s)-{\bf f}^{*}(y))\|\leq
L_{1}\mathbb{E}_{y\sim p}\|\lambda\bar{\bf u}(y,s)-{\bf f}^{*}(y)\|$ by (C1),
and $\mathbb{E}_{y\sim p}\|\lambda\bar{\bf u}(y,s)-{\bf
f}^{*}(y)\|\leq\|\lambda\bar{\bf u}(\cdot,s)-{\bf f}^{*}\|_{p}$. The latter
can be bounded by $\|{\bf f}^{*}\|_{p}$ using the same argument as in (Claim
1) (57) proved in the proof of Proposition 4.5, that is, for any $t\geq 0$,

$\|\lambda\bar{\bf u}(\cdot,t)-{\bf f}^{*}\|_{p}\leq\|\lambda\bar{\bf
u}(\cdot,0)-{\bf f}^{*}\|_{p}=\|{\bf f}^{*}\|_{p}.$ (78)

Putting together, this implies that

$\displaystyle\|ⓐ_{3}\|$ $\displaystyle\leq\lambda
L_{1}\int_{0}^{t}\Big{\|}(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\Big{\|}\mathbb{E}_{y\sim p}\|\lambda\bar{\bf
u}(y,s)-{\bf f}^{*}(y)\|ds$ $\displaystyle\leq\lambda tL_{1}\|{\bf
f}^{*}\|_{p}\Big{\|}(\mathbb{E}_{x^{\prime}\sim\hat{p}}-\mathbb{E}_{x^{\prime}\sim
p})\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\Big{\|}.$

Next, we bound
$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\|$ by matrix
Bernstein inequality Lemma A.2: The matrix $\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))$ is of size
$M_{\Theta}$-by-$M_{\Theta}$, and its operator norm is bounded by $L_{1}^{2}$
by (C1). When $n>n_{5}$ as defined by (73), Lemma A.2 gives that w.p. $\geq
1-n^{-10}$,

$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\|\leq
2L_{1}^{2}\sqrt{\frac{10\log
n+\log(2M_{\Theta})}{n}}<2\sqrt{11}L_{1}^{2}\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}.$ (79)

We call the above good event $E_{a}$. Thus, when $n>n_{5}$ and under $E_{a}$,

$\|ⓐ_{3}\|\leq 2\sqrt{11}(\lambda t)L_{1}^{3}\|{\bf
f}^{*}\|_{p}\sqrt{\frac{\log n+\log M_{\Theta}}{n}}.$

Putting together, we have that when $n>n_{5}$ and under $E_{a}\cap E_{b}$,

$\|③\|\leq L_{1}(\|ⓐ_{3}\|+\|ⓑ_{3}\|)\leq 2\sqrt{11}L_{1}^{2}((\lambda
t)L_{1}^{2}\|{\bf f}^{*}\|_{p}+b_{q})\sqrt{\frac{\log n+\log M_{\Theta}}{n}}.$
(80)

Bound of $\|②\|$: By definition,

$\displaystyle\hat{\bf K}_{t}(x,x^{\prime})-{\bf K}_{0}(x,x^{\prime})$
$\displaystyle=\partial_{\theta}{\bf
f}(x,\hat{\theta}(t))\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))^{T}-\partial_{\theta}{\bf
f}(x,{\theta}(0))\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))^{T}.$

Since ${\theta}(0),\,\hat{\theta}(t)\in B_{r}$, by (C1)(C2), $\forall
x,x^{\prime}\in\mathcal{X}$,

$\displaystyle\|\hat{\bf K}_{t}(x,x^{\prime})-{\bf K}_{0}(x,x^{\prime})\|$
$\displaystyle\leq\|\partial_{\theta}{\bf
f}(x,\hat{\theta}(t))\|\|\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))-\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))\|$ $\displaystyle~{}~{}~{}+\|\partial_{\theta}{\bf
f}(x,\hat{\theta}(t))-\partial_{\theta}{\bf
f}(x,{\theta}(0))\|\|\partial_{\theta}{\bf f}(x^{\prime},{\theta}(0))\|$
$\displaystyle\leq 2L_{1}L_{2}\|\hat{\theta}(t)-{\theta}(0)\|$
$\displaystyle\leq 2L_{1}L_{2}\sqrt{\frac{t}{\lambda}}\|{\bf f}^{*}\|_{p},$

where we used (72) in the last inequality. This gives that

$\|②\|\leq 2L_{1}L_{2}\sqrt{\frac{t}{\lambda}}\|{\bf
f}^{*}\|_{p}\mathbb{E}_{x^{\prime}\sim\hat{p}}\|\lambda\bar{\bf
u}(x^{\prime},t)-{\bf s}_{q}(x^{\prime}))\|.$ (81)

We will bound $\mathbb{E}_{x^{\prime}\sim\hat{p}}\|\lambda\bar{\bf
u}(x^{\prime},t)-{\bf s}_{q}(x^{\prime}))\|$ to be $O(1)$. First, note that we
have

$\displaystyle\|\lambda\bar{\bf u}(\cdot,t)-{\bf s}_{q}\|_{p}$
$\displaystyle\leq\|\lambda\bar{\bf u}(\cdot,t)-({\bf s}_{q}-{\bf
s}_{p})\|_{p}+\|{\bf s}_{p}\|_{p}$ $\displaystyle\leq\|{\bf
f}^{*}\|_{p}+\|{\bf s}_{p}\|_{p}\quad\text{(by \eqref{eq:claim1-baru})}$
$\displaystyle\leq 2\|{\bf f}^{*}\|_{p}+b_{q}.\quad\text{(by $\|{\bf
s}_{p}\|_{p}\leq\|{\bf f}^{*}\|_{p}+\|{\bf s}_{q}\|_{p}$ and (C6))}$ (82)

We claim that when $n$ is large enough and under a high probability good
event, to be specified below,

$\text{(Claim)}\quad(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|\lambda\bar{\bf
u}(x,t)-{\bf s}_{q}(x)\|^{2}\leq\left(L_{1}^{2}\|{\bf
f}^{*}\|_{p}+b_{q}\right)^{2}.$ (83)

Together with (82), this will imply that

$\displaystyle\mathbb{E}_{x\sim\hat{p}}\|\lambda\bar{\bf u}(x,t)-{\bf
s}_{q}(x))\|$ $\displaystyle\leq(\mathbb{E}_{x\sim\hat{p}}\|\lambda\bar{\bf
u}(x,t)-{\bf s}_{q}(x)\|^{2})^{1/2}$ $\displaystyle\leq{(2+L_{1}^{2})\|{\bf
f}^{*}\|_{p}+2b_{q}}:=C_{2,2}.$ (84)

Then (81) gives

$\|②\|\leq 2L_{1}L_{2}C_{2,2}\sqrt{{t}/{\lambda}}\|{\bf f}^{*}\|_{p}.$ (85)

To prove the claim (83), we will bound
$(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|\lambda\bar{\bf
u}(x,t)-{\bf s}_{q}(x)\|^{2}$ to be $\tilde{O}(n^{-1/2})$ uniformly for all
$t$ similarly as in the analysis of $ⓐ_{3}$ above, where we leverage the
expression (77) of $\bar{\bf u}(x,t)$ via ${\bf K}_{0}(x,x^{\prime})$. Based
on the expression, we have

$\displaystyle\|\lambda\bar{\bf u}(x,t)\|^{2}$
$\displaystyle=\lambda^{2}\int_{0}^{t}\int_{0}^{t}\mathbb{E}_{x_{1}^{\prime}\sim
p}\mathbb{E}_{x_{2}^{\prime}\sim p}(\lambda\bar{\bf
u}(x_{1}^{\prime},s_{1})-{\bf f}^{*}(x_{1}^{\prime}))^{T}$
$\displaystyle~{}~{}~{}[{\bf K}_{0}(x,x_{1}^{\prime})^{T}{\bf
K}_{0}(x,x_{2}^{\prime})](\lambda\bar{\bf u}(x_{2}^{\prime},s_{2})-{\bf
f}^{*}(x_{2}^{\prime}))ds_{1}ds_{2}.$

Meanwhile, $\|\lambda\bar{\bf u}(x,t)-{\bf s}_{q}(x)\|^{2}=\|\lambda\bar{\bf
u}(x,t)\|^{2}-2{\bf s}_{q}(x)\cdot\lambda\bar{\bf u}(x,t)+\|{\bf
s}_{q}(x)\|^{2}$, and then

$(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|\lambda\bar{\bf
u}(x,t)-{\bf s}_{q}(x)\|^{2}=ⓐ_{2}-2ⓑ_{2}+ⓒ_{2},$

where

$\displaystyle ⓐ_{2}$
$\displaystyle:=\lambda^{2}\int_{0}^{t}\int_{0}^{t}\mathbb{E}_{x_{1}^{\prime}\sim
p}\mathbb{E}_{x_{2}^{\prime}\sim p}(\lambda\bar{\bf
u}(x_{1}^{\prime},s_{1})-{\bf f}^{*}(x_{1}^{\prime}))^{T}$
$\displaystyle~{}~{}~{}~{}~{}~{}\Big{[}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
K}_{0}(x,x_{1}^{\prime})^{T}{\bf
K}_{0}(x,x_{2}^{\prime})\Big{]}(\lambda\bar{\bf u}(x_{2}^{\prime},s_{2})-{\bf
f}^{*}(x_{2}^{\prime}))ds_{1}ds_{2}$
$\displaystyle=\lambda^{2}\int_{0}^{t}\int_{0}^{t}\mathbb{E}_{x_{1}^{\prime}\sim
p}\mathbb{E}_{x_{2}^{\prime}\sim p}(\lambda\bar{\bf
u}(x_{1}^{\prime},s_{1})-{\bf f}^{*}(x_{1}^{\prime}))^{T}\partial_{\theta}{\bf
f}(x_{1}^{\prime},{\theta}(0))$
$\displaystyle~{}~{}~{}~{}~{}~{}\Big{[}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf
f}(x,{\theta}(0))\Big{]}\partial_{\theta}{\bf
f}(x_{2}^{\prime},{\theta}(0))^{T}(\lambda\bar{\bf
u}(x_{2}^{\prime},s_{2})-{\bf f}^{*}(x_{2}^{\prime}))ds_{1}ds_{2},$

and, (C1) implies that $\|\partial_{\theta}{\bf
f}(x_{1}^{\prime},{\theta}(0))\|,\,\|\partial_{\theta}{\bf
f}(x_{2}^{\prime},{\theta}(0))\|\leq L_{1}$,

$\displaystyle|ⓐ_{2}|$
$\displaystyle\leq\lambda^{2}L_{1}^{2}\Big{\|}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf
f}(x,{\theta}(0))\Big{\|}\left(\int_{0}^{t}\mathbb{E}_{x_{1}^{\prime}\sim
p}\|\lambda\bar{\bf u}(x_{1}^{\prime},s_{1})-{\bf
f}^{*}(x_{1}^{\prime})\|ds_{1}\right)^{2}$ $\displaystyle\leq(\lambda
t)^{2}L_{1}^{2}\|{\bf
f}^{*}\|_{p}^{2}\Big{\|}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\Big{\|},$ (86)

where the last inequality used that $\mathbb{E}_{x\sim p}\|\lambda\bar{\bf
u}(x,s)-{\bf f}^{*}(x)\|\leq\|\lambda\bar{\bf u}(\cdot,s)-{\bf
f}^{*}\|_{p}\leq\|{\bf f}^{*}\|_{p}$ by (78); Meanwhile,

$\displaystyle ⓑ_{2}$
$\displaystyle:=-\lambda\int_{0}^{t}\mathbb{E}_{x^{\prime}\sim
p}\Big{[}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}{\bf K}_{0}(x,x^{\prime})\Big{]}(\lambda\bar{\bf
u}(x^{\prime},s)-{\bf f}^{*}(x^{\prime}))ds$
$\displaystyle=-\lambda\int_{0}^{t}\mathbb{E}_{x^{\prime}\sim
p}\Big{[}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf
f}(x,{\theta}(0))\Big{]}\partial_{\theta}{\bf
f}(x^{\prime},{\theta}(0))^{T}(\lambda\bar{\bf u}(x^{\prime},s)-{\bf
f}^{*}(x^{\prime}))ds,$

and, similarly,

$\displaystyle|ⓑ_{2}|$ $\displaystyle\leq\lambda
L_{1}\Big{\|}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf
f}(x,{\theta}(0))\Big{\|}\int_{0}^{t}\mathbb{E}_{x^{\prime}\sim
p}\|\lambda\bar{\bf u}(x^{\prime},s)-{\bf f}^{*}(x^{\prime})\|ds$
$\displaystyle\leq(\lambda t)L_{1}\|{\bf
f}^{*}\|_{p}\Big{\|}(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\Big{\|};$ (87)

Finally,

$ⓒ_{2}:=(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|{\bf
s}_{q}(x)\|^{2}.$

To bound $|ⓐ_{2}|$, $|ⓑ_{2}|$, and $ⓒ_{2}$, we use the concentration argument
respectively.

For $ⓐ_{2}$, recall that we have bounded
$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\partial_{\theta}{\bf
f}(x,{\theta}(0))^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\|$ in (79). This
means that when $n>n_{5}$ and under $E_{a}$,

$|ⓐ_{2}|\leq 2\sqrt{11}(\lambda t)^{2}L_{1}^{4}\|{\bf
f}^{*}\|_{p}^{2}\sqrt{\frac{\log n+\log M_{\Theta}}{n}}.$ (88)

For $ⓑ_{2}$, we have bounded
$\|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
s}_{q}(x)^{T}\partial_{\theta}{\bf f}(x,{\theta}(0))\|$ in (76). Thus, when
$n>n_{5}$ and under $E_{b}$, (87) gives that

$|ⓑ_{2}|\leq 2\sqrt{11}(\lambda t)L_{1}^{2}\|{\bf
f}^{*}\|_{p}b_{q}\sqrt{\frac{\log n+\log M_{\Theta}}{n}}.$ (89)

For $ⓒ_{2}$, by Assumption 5(C6) and that $n>n_{\lambda}$ ensures $n>n_{6}$,
there is a good event $E_{c}$ that happens w.p. $\geq 1-n^{-10}$ under which

$ⓒ_{2}\leq\sqrt{20}b_{q}^{2}\sqrt{\frac{\log n}{n}}.$ (90)

Putting together (88)(89)(90), when $n>\max\\{n_{\lambda},n_{5}\\}$ and under
the intersection of $E_{a}$, $E_{b}$ and $E_{c}$, we have

$\displaystyle(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|\lambda\bar{\bf
u}(x,t)-{\bf s}_{q}(x)\|^{2}$ $\displaystyle\leq|ⓐ_{2}|+2|ⓑ_{2}|+ⓒ_{2}$
$\displaystyle\leq 2\sqrt{11}\left((\lambda t)L_{1}^{2}\|{\bf
f}^{*}\|_{p}+b_{q}\right)^{2}\sqrt{\frac{\log n+\log M_{\Theta}}{n}}.$

Let $n_{\lambda,t}$ be the integer s.t. when $n>n_{\lambda,t}$,

$2\sqrt{11}\max\\{(\lambda t)^{2},1\\}\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}<1,$ (91)

then we have shown that the claim (83) holds when
$n>\max\\{n_{\lambda},n_{5},n_{\lambda,t}\\}$ and under $E_{a}\cap E_{b}\cap
E_{c}$.

Analysis of (71): We handle ① based on the proved bounds of the other four
terms in (71). Define ${II⃝}(x,t):=②+③+④+⑤$. Collecting (74)(75)(80)(85), we
have that when $n>\max\\{n_{\lambda},n_{5},n_{\lambda,t}\\}$ and under the
intersection (of the good event in Lemma 4.8 and) $E_{a}\cap E_{b}\cap
E_{c}\cap E_{5}$,

$\displaystyle\|II⃝(x,t)\|$
$\displaystyle\leq(2L_{1}L_{2}C_{2,2}+L_{2}L_{3}+L_{1}L_{4})\sqrt{{t}/{\lambda}}\|{\bf
f}^{*}\|_{p}$ $\displaystyle~{}~{}~{}+2\sqrt{11}(\lambda tL_{1}^{4}\|{\bf
f}^{*}\|_{p}+L_{1}^{2}b_{q}+2L_{1}L_{3})\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}$ $\displaystyle\leq C_{2,1}\left(\sqrt{{t}/{\lambda}}\|{\bf
f}^{*}\|_{p}+(\lambda t+1)\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}\right)=:g(t),$ (92)

where

$C_{2,1}:=\max\\{2L_{1}L_{2}C_{2,2}+L_{2}L_{3}+L_{1}L_{4},\,2\sqrt{11}L_{1}^{4}\|{\bf
f}^{*}\|_{p},\,2\sqrt{11}(L_{1}^{2}b_{q}+2L_{1}L_{3})\\},$ (93)

and $C_{2,2}$ is defined as in (84). Back to (71), we have

${\partial_{t}}(\hat{\bf u}-\bar{\bf
u})(x,t)=①+II⃝(x,t)=-\mathbb{E}_{x^{\prime}\sim\hat{p}}\hat{\bf
K}_{t}(x,x^{\prime})\circ\lambda(\hat{\bf u}-\bar{\bf
u})(x^{\prime},t)+II⃝(x,t).$ (94)

We first derive a bound of

$\|\lambda(\hat{\bf u}-\bar{\bf
u})(\cdot,t)\|_{\hat{p}}:=(\mathbb{E}_{x\sim\hat{p}}\|\lambda(\hat{\bf
u}-\bar{\bf u})(x,t)\|^{2})^{1/2},$

namely the mean squared error on training samples. By (94),

$\displaystyle\frac{d}{dt}\frac{1}{2}\|\lambda(\hat{\bf u}-\bar{\bf
u})(\cdot,t)\|_{\hat{p}}^{2}$
$\displaystyle=\lambda^{2}\mathbb{E}_{x\sim\hat{p}}(\hat{\bf u}-\bar{\bf
u})(x,t)^{T}{\partial_{t}}(\hat{\bf u}-\bar{\bf u})(x,t)$
$\displaystyle\leq\lambda^{2}\mathbb{E}_{x\sim\hat{p}}(\hat{\bf u}-\bar{\bf
u})(x,t)^{T}II⃝(x,t)$
$\displaystyle\leq\lambda^{2}\mathbb{E}_{x\sim\hat{p}}\|(\hat{\bf u}-\bar{\bf
u})(x,t)\|\|II⃝(x,t)\|$
$\displaystyle\leq\lambda^{2}g(t)\mathbb{E}_{x\sim\hat{p}}\|(\hat{\bf
u}-\bar{\bf u})(x,t)\|$ $\displaystyle\leq\lambda g(t)\|\lambda(\hat{\bf
u}-\bar{\bf u})(\cdot,t)\|_{\hat{p}},$

where in the first inequality we used that

$\mathbb{E}_{x\sim\hat{p}}(\hat{\bf u}-\bar{\bf
u})(x,t)^{T}①=-\lambda\mathbb{E}_{x\sim\hat{p}}\mathbb{E}_{x^{\prime}\sim\hat{p}}(\hat{\bf
u}-\bar{\bf u})(x,t)^{T}[\hat{\bf K}_{t}(x,x^{\prime})](\hat{\bf u}-\bar{\bf
u})(x^{\prime},t)\leq 0.$

Using the same argument as in the proof of Proposition 4.5(ii) based on (59),
we have that for any $t$ being considered,

$\displaystyle\|\lambda(\hat{\bf u}-\bar{\bf u})(\cdot,t)\|_{\hat{p}}$
$\displaystyle\leq 2\lambda\int_{0}^{t}g(s)ds$ $\displaystyle\leq
2\lambda\int_{0}^{t}g(t)ds\quad\text{($g(t)$ is monotone in $t$)}$
$\displaystyle=2\lambda tg(t).$ (95)

Next, we use (95) to bound $\|\lambda(\hat{\bf u}-\bar{\bf
u})(\cdot,t)\|_{{p}}$. Again from (94), we have

$\displaystyle\frac{d}{dt}\frac{1}{2}\|\lambda(\hat{\bf u}-\bar{\bf
u})(\cdot,t)\|_{{p}}^{2}=\lambda^{2}\mathbb{E}_{x\sim{p}}(\hat{\bf u}-\bar{\bf
u})(x,t)^{T}{\partial_{t}}(\hat{\bf u}-\bar{\bf u})(x,t)$
$\displaystyle~{}~{}~{}=-\lambda^{2}\mathbb{E}_{x\sim{p}}\mathbb{E}_{x^{\prime}\sim\hat{p}}(\hat{\bf
u}-\bar{\bf u})(x,t)^{T}[\hat{\bf K}_{t}(x,x^{\prime})]\lambda(\hat{\bf
u}-\bar{\bf u})(x^{\prime},t)+\lambda^{2}\mathbb{E}_{x\sim{p}}(\hat{\bf
u}-\bar{\bf u})(x,t)^{T}II⃝(x,t)$ $\displaystyle~{}~{}~{}\leq\lambda
L_{1}^{2}\mathbb{E}_{x\sim{p}}\|\lambda(\hat{\bf u}-\bar{\bf
u})(x,t)\|\mathbb{E}_{x^{\prime}\sim\hat{p}}\|\lambda(\hat{\bf u}-\bar{\bf
u})(x^{\prime},t)\|+\lambda\mathbb{E}_{x\sim{p}}\|\lambda(\hat{\bf u}-\bar{\bf
u})(x,t)\|\|II⃝(x,t)\|$ $\displaystyle~{}~{}~{}\leq\lambda
L_{1}^{2}\|\lambda(\hat{\bf u}-\bar{\bf u})(\cdot,t)\|_{{p}}\|\lambda(\hat{\bf
u}-\bar{\bf u})(\cdot,t)\|_{\hat{p}}+\lambda g(t)\|\lambda(\hat{\bf
u}-\bar{\bf u})(\cdot,t)\|_{{p}}$
$\displaystyle~{}~{}~{}\leq\lambda(L_{1}^{2}2\lambda
t+1)g(t)\|\lambda(\hat{\bf u}-\bar{\bf u})(\cdot,t)\|_{{p}}\quad\text{(by
\eqref{eq:hatp-hatu-baru-bound})}$

where in the first inequality we used the upper bound $\|\hat{\bf
K}_{t}(x,x^{\prime})\|=\|\partial_{\theta}{\bf
f}(x,\hat{\theta}(t))\partial_{\theta}{\bf f}(x,\hat{\theta}(t))^{T}\|\leq
L_{1}^{2}$ by (C1). Using the argument by integrating over $t$ again, we have

$\displaystyle\|\lambda(\hat{\bf u}-\bar{\bf u})(\cdot,t)\|_{{p}}$
$\displaystyle\leq 2\lambda\int_{0}^{t}(2L_{1}^{2}\lambda s+1)g(s)ds$
$\displaystyle\leq 2\lambda t(2L_{1}^{2}\lambda
t+1)g(t)\quad\text{($(2L_{1}^{2}\lambda t+1)g(t)$ is monotone in $t$)}$

Inserting the definition of $g(t)$ in (92), we have

$\displaystyle\|\lambda(\hat{\bf u}-\bar{\bf u})(\cdot,t)\|_{{p}}$
$\displaystyle\leq 2C_{2,1}(1+2L_{1}^{2}\lambda
t)\left(\lambda^{1/2}t^{3/2}\|{\bf f}^{*}\|_{p}+(\lambda t+1)\lambda
t\sqrt{\frac{\log n+\log M_{\Theta}}{n}}\right),$

which proves (40) by defining $C_{2}:=2C_{2,1}\max\\{1,2L_{1}^{2}\\}$. By the
definition of $C_{2,1}$ in (93) and $C_{2,2}$ in (84),

$C_{2}\leq c(L_{1},\cdots,L_{4})(1+b_{q}+\|{\bf f}^{*}\|_{p}),$

where $c(L_{1},\cdots,L_{4})$ an $O(1)$ constant determined by
$L_{1},L_{2},L_{3},L_{4}$. ∎

###### Proof of Theorem 4.10.

Similarly as in the proof of Theorem 4.6, the required largeness of $\lambda$
guarantees that the range of $t_{0}$ in (41) is non-empty. Proposition 4.4
applies same as before, and $t_{0}\geq 1$ ensures that the bound of
$\|\lambda\bar{\bf u}(\cdot,t)-{\bf f}^{*}\|_{p}$ as in (28) holds.

The upper bound of $t_{0}$ in (41) ensures that $t$ falls in the needed range
by Proposition 4.9; We consider large $n$ as required by Proposition 4.9, and,
by Remark 4.4, the three thresholds $n_{\lambda},n_{5},n_{\lambda,t}$ are
determined by $\lambda$, $\|{\bf f}^{*}\|_{p}$, $\log M_{\Theta}$ and $\lambda
t=t_{0}\frac{\log(1/\epsilon)}{\delta}$. We also consider being under the
intersection of the good events in Proposition 4.9, which happens w.p. $\geq
1-7n^{-10}$, then (40) holds.

Putting together (28) and (40), by triangle inequality, we have

$\|\lambda\hat{\bf u}(x,t)-{\bf f}^{*}\|_{p}\leq 2\epsilon\|{\bf
f}^{*}\|_{p}+C_{2}(1+\lambda t)\left(\lambda^{1/2}t^{3/2}\|{\bf
f}^{*}\|_{p}+(1+\lambda t)\lambda t\sqrt{\frac{\log n+\log
M_{\Theta}}{n}}\right).$

Then (42) follows by the equality $\lambda
t=t_{0}\frac{\log(1/\epsilon)}{\delta}$ and defining

$\displaystyle\kappa_{1}:=(1+t_{0}\frac{\log(1/\epsilon)}{\delta})\left(t_{0}\frac{\log(1/\epsilon)}{\delta}\right)^{3/2},\quad\kappa_{2}:=(1+t_{0}\frac{\log(1/\epsilon)}{\delta})^{2}t_{0}\frac{\log(1/\epsilon)}{\delta}.$

∎

###  8.3 Proof in Section 5.1

###### Proof of Corollary 5.1.

Applying Theorem 4.10 gives that, under $E_{\rm tr}$, $\|\lambda\hat{\bf
f}-{\bf f}^{*}\|_{p}$ is upper bounded by the r.h.s. of (42) where setting
$n=n_{\rm tr}$. The condition of the corollary further implies that

$\|\lambda\hat{\bf f}-{\bf f}^{*}\|_{p}<0.9\|{\bf f}^{*}\|_{p}.$

Then,

$\langle\lambda\hat{\bf f},{\bf f}^{*}\rangle-\|{\bf
f}^{*}\|_{p}^{2}=\langle\lambda\hat{\bf f}-{\bf f}^{*},{\bf
f}^{*}\rangle\geq-\|\lambda\hat{\bf f}-{\bf f}^{*}\|_{p}\|{\bf
f}^{*}\|_{p}>-0.9\|{\bf f}^{*}\|_{p}^{2},$

which gives that $\langle\lambda\hat{\bf f},{\bf f}^{*}\rangle>0.1\|{\bf
f}^{*}\|_{p}^{2}$. Meanwhile, by that $\hat{\bf f}\in\mathcal{F}_{0}(p)\cap
L^{2}(p)$ (Assumption 3), (8) gives that

$\mathbb{E}_{x\sim p}T_{q}\hat{\bf f}(x)=\langle{\bf f}^{*},\hat{\bf
f}\rangle_{p}>\frac{0.1}{\lambda}\|{\bf f}^{*}\|_{p}^{2}.$ (96)

Next, we bound the deviation of $\hat{T}$ and $\hat{T}_{\rm null}$ by
concentration. By definition,

$\hat{T}=\mathbb{E}_{x\sim\hat{p}}T_{q}\hat{\bf
f}(x),\quad\hat{p}:=\frac{1}{n_{\rm GoF}}\sum_{i=1}^{n_{\rm
GoF}}\delta_{x_{i}}.$

Recall that $\hat{\theta}(t)\in B_{r}$ by Lemma 4.8, then, by Assumption
5(C5)(C6),

$|T_{q}\hat{\bf f}(x_{i})|\leq|{\bf s}_{q}\cdot\hat{\bf
f}(x_{i})|+|\nabla\cdot\hat{\bf f}(x_{i})|\leq b_{(0)}b_{q}+b_{(1)}=:b.$

By the Hoeffding’s inequality, Lemma A.1, for any $t>0$, denote $n_{\rm GoF}$
by $n$,

$\Pr[\mathbb{E}_{x\sim\hat{p}}T_{q}\hat{\bf
f}(x)-\mathbb{E}_{x\sim{p}}T_{q}\hat{\bf f}(x)\leq-t|E_{\rm
tr}]\leq\exp\\{-\frac{nt^{2}}{2b^{2}}\\}.$

By union bound, and the condition that $7n_{\rm
tr}^{-10}<0.1\min\\{\alpha,\beta\\}$,

$\Pr[\mathbb{E}_{x\sim\hat{p}}T_{q}\hat{\bf
f}(x)-\mathbb{E}_{x\sim{p}}T_{q}\hat{\bf
f}(x)\leq-t]\leq\exp\\{-\frac{nt^{2}}{2b^{2}}\\}+7n_{\rm
tr}^{-10}<\exp\\{-\frac{nt^{2}}{2b^{2}}\\}+0.1\beta.$

This gives that w.p. $\geq 1-\beta$,

$\hat{T}\geq\mathbb{E}_{x\sim{p}}T_{q}\hat{\bf f}(x)-t_{1},\quad
t_{1}:=\sqrt{2}b\sqrt{\frac{\log\frac{1}{0.9\beta}}{n_{\rm GoF}}}.$

Similarly,

$\hat{T}_{\rm null}=\mathbb{E}_{y\sim\hat{q}}T_{q}\hat{\bf
f}(y),\quad\hat{q}:=\frac{1}{n_{\rm GoF}}\sum_{i=1}^{n_{\rm
GoF}}\delta_{y_{i}},$

and w.p. $\geq 1-\alpha$,

$\hat{T}_{\rm null}\leq t_{2},\quad
t_{2}:=\sqrt{2}b\sqrt{\frac{\log\frac{1}{0.9\alpha}}{n_{\rm GoF}}}.$

We set $t_{2}$ to be the test threshold, combined with (96), we know that the
target significance level and power can be achieved if
$t_{1}+t_{2}\leq\frac{0.1}{\lambda}\|{\bf f}^{*}\|_{p}^{2}$. Together with
that $\sqrt{\log\frac{1}{0.9\alpha}}\leq 0.5+\sqrt{\log\frac{1}{\alpha}}$ and
similarly for $\sqrt{\log\frac{1}{0.9\beta}}$, this proves the corollary. ∎

## Acknowledgement

The work is supported by NSF DMS-2134037. M.R. and Y.X. are partially
supported by an NSF CAREER CCF-1650913, and NSF DMS-2134037, CMMI-2015787,
CMMI-2112533, DMS-1938106, and DMS-1830210. X.C. is also partially supported
by NSF DMS-1818945, NIH R01GM131642 and the Alfred P. Sloan Foundation.

## References

  * [1] Andreas Anastasiou, Alessandro Barp, François-Xavier Briol, Bruno Ebner, Robert E Gaunt, Fatemeh Ghaderinezhad, Jackson Gorham, Arthur Gretton, Christophe Ley, Qiang Liu, et al.  Stein’s method meets computational statistics: A review of some recent developments.  arXiv preprint arXiv:2105.03481, 2021.
  * [2] Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein generative adversarial networks.  In International conference on machine learning, pages 214–223. PMLR, 2017.
  * [3] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  Advances in Neural Information Processing Systems, 32, 2019.
  * [4] Sanjeev Arora, Rong Ge, Yingyu Liang, Tengyu Ma, and Yi Zhang.  Generalization and equilibrium in generative adversarial nets (gans).  In International Conference on Machine Learning, pages 224–232. PMLR, 2017.
  * [5] Alessandro Barp, Francois-Xavier Briol, Andrew Duncan, Mark Girolami, and Lester Mackey.  Minimum stein discrepancy estimators.  Advances in Neural Information Processing Systems, 32, 2019.
  * [6] Xiuyuan Cheng and Yao Xie.  Neural tangent kernel maximum mean discrepancy.  Advances in Neural Information Processing Systems, 34, 2021.
  * [7] Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  Advances in Neural Information Processing Systems, 32, 2019.
  * [8] Kacper Chwialkowski, Heiko Strathmann, and Arthur Gretton.  A kernel test of goodness of fit.  In International conference on machine learning, pages 2606–2615. PMLR, 2016.
  * [9] Laurent Dinh, David Krueger, and Yoshua Bengio.  Nice: Non-linear independent components estimation.  arXiv preprint arXiv:1410.8516, 2014.
  * [10] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp.  arXiv preprint arXiv:1605.08803, 2016.
  * [11] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In International Conference on Learning Representations, 2018.
  * [12] Ian Goodfellow, Yoshua Bengio, and Aaron Courville.  Deep learning.  MIT press, 2016.
  * [13] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  Advances in neural information processing systems, 27, 2014.
  * [14] Jackson Gorham and Lester Mackey.  Measuring sample quality with stein’s method.  Advances in Neural Information Processing Systems, 28, 2015.
  * [15] Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky.  Your classifier is secretly an energy based model and you should treat it like one.  arXiv preprint arXiv:1912.03263, 2019.
  * [16] Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, and Richard Zemel.  Learning the stein discrepancy for training and evaluating energy-based models without sampling.  In International Conference on Machine Learning, pages 3732–3747. PMLR, 2020.
  * [17] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola.  A kernel two-sample test.  The Journal of Machine Learning Research, 13(1):723–773, 2012.
  * [18] Michael Gutmann and Aapo Hyvärinen.  Noise-contrastive estimation: A new estimation principle for unnormalized statistical models.  In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 297–304. JMLR Workshop and Conference Proceedings, 2010.
  * [19] Tianyang Hu, Zixiang Chen, Hanxi Sun, Jincheng Bai, Mao Ye, and Guang Cheng.  Stein neural sampler.  arXiv preprint arXiv:1810.03545, 2018.
  * [20] Michael F Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  Communications in Statistics-Simulation and Computation, 18(3):1059–1076, 1989.
  * [21] Aapo Hyvärinen and Peter Dayan.  Estimation of non-normalized statistical models by score matching.  Journal of Machine Learning Research, 6(4), 2005.
  * [22] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  Advances in neural information processing systems, 31, 2018.
  * [23] Wittawat Jitkrittum, Wenkai Xu, Zoltán Szabó, Kenji Fukumizu, and Arthur Gretton.  A linear-time kernel goodness-of-fit test.  Advances in Neural Information Processing Systems, 30, 2017.
  * [24] Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  Advances in neural information processing systems, 31, 2018.
  * [25] Yann LeCun, Corinna Cortes, and CJ Burges.  Mnist handwritten digit database, 2010.
  * [26] Feng Liu, Wenkai Xu, Jie Lu, Guangquan Zhang, Arthur Gretton, and Danica J Sutherland.  Learning deep kernels for non-parametric two-sample tests.  In International conference on machine learning, pages 6316–6326. PMLR, 2020.
  * [27] Qiang Liu, Jason Lee, and Michael Jordan.  A kernelized stein discrepancy for goodness-of-fit tests.  In International conference on machine learning, pages 276–284. PMLR, 2016.
  * [28] Qiang Liu and Dilin Wang.  Stein variational gradient descent: A general purpose bayesian inference algorithm.  Advances in neural information processing systems, 29, 2016.
  * [29] Sebastian Nowozin, Botond Cseke, and Ryota Tomioka.  f-gan: Training generative neural samplers using variational divergence minimization.  Advances in neural information processing systems, 29, 2016.
  * [30] Prajit Ramachandran, Barret Zoph, and Quoc V Le.  Searching for activation functions.  arXiv preprint arXiv:1710.05941, 2017.
  * [31] Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  Advances in Neural Information Processing Systems, 32, 2019.
  * [32] Sreejith Sreekumar and Ziv Goldfeld.  Neural estimation of statistical divergences.  Journal of machine learning research, 23(126), 2022.
  * [33] Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet.  On integral probability metrics,$\backslash$phi-divergences and binary classification.  arXiv preprint arXiv:0901.2698, 2009.
  * [34] Charles Stein.  A bound for the error in the normal approximation to the distribution of a sum of dependent random variables.  In Proceedings of the sixth Berkeley symposium on mathematical statistics and probability, volume 2: Probability theory, volume 6, pages 583–603. University of California Press, 1972.
  * [35] Charles Stein, Persi Diaconis, Susan Holmes, and Gesine Reinert.  Use of exchangeable pairs in the analysis of simulations.  Lecture Notes-Monograph Series, pages 1–26, 2004.
  * [36] Yee Whye Teh, Max Welling, Simon Osindero, and Geoffrey E Hinton.  Energy-based models for sparse overcomplete representations.  Journal of Machine Learning Research, 4(Dec):1235–1260, 2003.
  * [37] Joel A Tropp et al.  An introduction to matrix concentration inequalities.  Foundations and Trends® in Machine Learning, 8(1-2):1–230, 2015.
  * [38] Laurens Van der Maaten and Geoffrey Hinton.  Visualizing data using t-sne.  Journal of machine learning research, 9(11), 2008.
  * [39] Martin J Wainwright.  High-dimensional statistics: A non-asymptotic viewpoint, volume 48.  Cambridge university press, 2019.
  * [40] Pengchuan Zhang, Qiang Liu, Dengyong Zhou, Tao Xu, and Xiaodong He.  On the discrimination-generalization tradeoff in gans.  In International Conference on Learning Representations, 2018.

##  Appendix A Additional technical lemmas and proofs

###  A.1 Proofs of technical lemmas

###### Proof of Lemma 2.1.

For any ${\bf f}\in\mathcal{F}_{0}(p)$, because $p{\bf
f}|_{\partial\mathcal{X}}=0$, we have

$\int_{\mathcal{X}}p\nabla\cdot{\bf f}=-\int_{\mathcal{X}}\nabla p\cdot{\bf
f}$

Then for any ${\bf f}\in L^{2}(p)\cap\mathcal{F}_{0}(p)$,

$\displaystyle\mathbb{E}_{x\sim p}T_{q}{\bf f}(x)$
$\displaystyle=\int_{\mathcal{X}}p({\bf s}_{q}\cdot{\bf f}+\nabla\cdot{\bf
f})=\int_{\mathcal{X}}p{\bf s}_{q}\cdot{\bf f}-\int_{\mathcal{X}}\nabla
p\cdot{\bf f}$ $\displaystyle=\int_{\mathcal{X}}p({\bf s}_{q}-{\bf
s}_{p})\cdot{\bf f}=\mathbb{E}_{x\sim p}{\bf f}^{*}\cdot{\bf f}=\langle{\bf
f}^{*},{\bf f}\rangle_{p}.$ (A.1)

Note that the integrals are well-defined because both ${\bf s}_{p}$ and ${\bf
s}_{q}$ are in $L^{2}(p)$. By (A.1), whenever $\|{\bf f}\|_{p}\leq r$,

$\mathbb{E}_{x\sim p}T_{q}{\bf f}(x)\leq\|{\bf f}^{*}\|_{p}\|{\bf f}\|_{p}\leq
r\|{\bf f}^{*}\|_{p},$ (A.2)

where $\|{\bf f}^{*}\|_{p}<\infty$ because ${\bf s}_{q}$ and ${\bf s}_{p}$ are
in $L^{2}(p)$. This proves that ${\rm SD}_{r}(p,q)$ is upper bounded by
$r\|{\bf f}^{*}\|_{p}$ and thus is finite. To show that ${\rm
SD}_{r}(p,q)=r\|{\bf f}^{*}\|_{p}$, first when ${\bf f}^{*}=0$ in $L^{2}(p)$
then ${\rm SD}_{r}(p,q)=0$ by (A.2). If $\|{\bf f}^{*}\|_{p}>0$, by property
of the inner product, the ${\bf f}$ that is a positive multiple of ${\bf
f}^{*}$ and has $\|{\bf f}\|_{p}=r$ achieves the supremum of value $r\|{\bf
f}^{*}\|_{p}$. Because ${\bf f}^{*}\in\mathcal{F}_{0}(p)$, this maximizer
${\bf f}=({r}/{\|{\bf f}^{*}\|_{p}}){\bf f}^{*}$ is also in
$\mathcal{F}_{0}(p)$. ∎

###### Proof of Lemma 4.1.

By the definition of $L_{\lambda}$ in (13) and (11), the latter applies to
${\bf f}={\bf f}(\cdot,\theta)$ because ${\bf f}(\cdot,\theta)$ is in
$L^{2}(p)\cap\mathcal{F}_{0}(p)$ by Assumption 3,

$L_{\lambda}(\theta)=-\langle{\bf f}^{*},{\bf
f}(\cdot,\theta)\rangle_{p}+\frac{\lambda}{2}\|{\bf
f}(\cdot,\theta)\|_{p}^{2}=\mathbb{E}_{x\sim p}\left(-{\bf f}^{*}(x)\cdot{\bf
f}(x,\theta)+\frac{\lambda}{2}\|{\bf f}(x,\theta)\|^{2}\right).$

Then we have that

$\displaystyle\frac{\partial}{\partial\theta}L_{\lambda}(\theta)$
$\displaystyle=\mathbb{E}_{x\sim p}\big{(}-{\bf
f}^{*}(x)\cdot\partial_{\theta}{\bf f}(x,\theta)+\lambda{\bf
f}(x,\theta)\cdot\partial_{\theta}{\bf f}(x,\theta)\big{)}$
$\displaystyle=\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta)\cdot\big{(}\lambda{\bf f}(x,\theta)-{\bf f}^{*}(x)\big{)}.$

This proves (19) according to the GD dynamic defined in (18). ∎

###### Proof of Lemma 4.2.

Inserting (19) to (21) gives that

$\frac{\partial}{\partial t}{\bf u}(x,t)=-\langle\partial_{\theta}{\bf
f}(x,\theta(t)),\mathbb{E}_{x^{\prime}\sim p}\partial_{\theta}{\bf
f}(x^{\prime},\theta(t))\cdot\big{(}\lambda{\bf u}(x^{\prime},t)-{\bf
f}^{*}(x^{\prime})\big{)}\rangle_{\Theta},$ (A.3)

where we used the definition (20). The lemma then follows by the definition
(22) and the linearity of inner-product $\langle\cdot,\cdot\rangle_{\Theta}$.
∎

###### Proof of Lemma 4.7.

By (33),

$\partial_{\theta}\hat{L}_{\lambda}(\theta)=\mathbb{E}_{x\sim\hat{p}}\left(\partial_{\theta}{\bf
f}(x,\theta)\cdot(\lambda{\bf f}(x,\theta)-{\bf
s}_{q}(x))-\nabla_{x}\cdot\partial_{\theta}{\bf f}(x,\theta)\right),$ (A.4)

and combined with (34)(35), we have

$\displaystyle{\partial_{t}}\hat{\bf u}(x,t)$
$\displaystyle=-\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\partial_{\theta}\hat{L}_{\lambda}(\hat{\theta}(t))\rangle_{\Theta}$
$\displaystyle=-\left\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\mathbb{E}_{x^{\prime}\sim\hat{p}}\left(\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\cdot(\lambda\hat{\bf u}(x^{\prime},t)-{\bf
s}_{q}(x^{\prime}))-\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\right)\right\rangle_{\Theta}$
$\displaystyle=-\mathbb{E}_{x^{\prime}\sim\hat{p}}\left(\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\rangle_{\Theta}\circ(\lambda\hat{\bf
u}(x^{\prime},t)-{\bf s}_{q}(x^{\prime}))-\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\nabla_{x^{\prime}}\cdot\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\rangle_{\Theta}\right).$

This proves the lemma by the definition of $\hat{\bf K}_{t}(x,x^{\prime})$ in
(36). ∎

###  A.2 Concentration lemmas

The following lemma can be derived from Classical Hoeffding’s inequality.

######  Lemma A.1 (Hoeffding).

Suppose $X_{1},\cdots,X_{n}$ are independent random variables, $|X_{i}|\leq
L$. Then for any $t\geq 0$,

$\Pr[\frac{1}{n}\sum_{i=1}^{n}(X_{i}-\mathbb{E}X_{i})\geq
t],\,\Pr[\frac{1}{n}\sum_{i=1}^{n}(X_{i}-\mathbb{E}X_{i})\leq-t],\leq\exp\\{-\frac{nt^{2}}{2L^{2}}\\}.$

The following lemma can be directly derived from [37, Theorem 6.1.1].

######  Lemma A.2 (Matrix Bernstein [37]).

Let $X_{i}$ be a sequence of $n$ independent, random, real-valued matrices of
size $d_{1}$-by-$d_{2}$. Assume that $\mathbb{E}X_{i}=0$ and $\|X_{i}\|\leq L$
for each $i$, and $\nu>0$ be such that

$\|\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}X_{i}X_{i}^{T}\|,\,\|\frac{1}{n}\sum_{i=1}^{n}\mathbb{E}X_{i}^{T}X_{i}\|\leq\nu.$

Then, for any $t\geq 0$,

$\Pr[\|\frac{1}{n}\sum_{i=1}^{n}X_{i}\|\geq
t]\leq(d_{1}+d_{2})\exp\\{-\frac{nt^{2}}{2(\nu+Lt/3)}\\}.$

##  Appendix B Additional experiments

###  B.1 Supplementary Tables

Regularization ($\lambda$) | Mean Power | Epoch |  Mean Selected $\widehat{\rm MSE}_{q}$ | Mean Stopping Epoch
---|---|---|---|---
$1\mathrm{e}-3$ |  0.814 $\pm 0.018$ | 60 | 0.041 | 59.5
$1\mathrm{e}-2$ |  0.808 $\pm 0.027$ | 50 | 0.027 | 49.4
$1\mathrm{e}-1$ |  0.825 $\pm 0.040$ | 30 | 0.021 | 27.8
$1\mathrm{e}-0$ |  0.835 $\pm 0.025$ | 20 | 0.025 | 21.8
$\Lambda(1\mathrm{e}0,5\mathrm{e}-2,0.90)$ |  0.839 $\pm 0.044$ | 30 | 0.023 | 31.5
$\Lambda(1\mathrm{e}0,5\mathrm{e}-2,0.95)$ |  0.838 $\pm 0.029$ | 30 | 0.020 | 34.1
Table A.1:  For the mixture data in 2D as in (45) of Section 6.1.2, we consider the approximate epoch at which the monitor $\widehat{\rm MSE}^{(m)}_{p}$ (17) is minimized. At this epoch of training, we display the average GoF testing power ($n_{\rm GoF}=75$) and $\widehat{\rm MSE}_{q}$ (16) over $n_{\rm replica}=10$ models for each regularization strategy. Regularization ($\lambda$) | Mean Power | Epoch |  Mean Selected $\widehat{\rm MSE}_{q}$ | Mean Stopping Epoch
---|---|---|---|---
$2.500\mathrm{e}-4$ |  0.839 $\pm 0.034$ | 50 | 0.120 | 52.4
$1.000\mathrm{e}-3$ |  0.894 $\pm 0.027$ | 50 | 0.105 | 50.1
$4.000\mathrm{e}-3$ |  0.910 $\pm 0.014$ | 45 | 0.114 | 46.0
$1.600\mathrm{e}-2$ |  0.893 $\pm 0.025$ | 30 | 0.127 | 34.1
$6.400\mathrm{e}-2$ |  0.905 $\pm 0.013$ | 20 | 0.128 | 21.5
$2.560\mathrm{e}-1$ |  0.878 $\pm 0.019$ | 12 | 0.132 | 12.8
$1.024\mathrm{e}-0$ |  0.824 $\pm 0.037$ | 12 | 0.172 | 12.4
$\Lambda(5\mathrm{e}-1,1\mathrm{e}-3,0.80)$ |  0.928 $\pm 0.007$ | 50 | 0.088 | 49.8
$\Lambda(5\mathrm{e}-1,1\mathrm{e}-3,0.85)$ |  0.920 $\pm 0.021$ | 45 | 0.095 | 40.1
Table A.2: For the mixture data in 10D as in (45), we consider the epoch that minimizes monitor $\widehat{\rm MSE}^{(m)}_{p}$, at which point we display the average testing power (${n_{\rm GoF}}=200$) and $\widehat{\rm MSE}_{q}$ over $n_{\rm replica}=10$ models for each regularization strategy. Regularization ($\lambda$) | Mean Power | Epoch |  Mean Selected $\widehat{\rm MSE}_{q}$ | Mean Stopping Epoch
---|---|---|---|---
$2.500\mathrm{e}-4$ |  0.862 $\pm 0.103$ | 45 | 0.260 | 37.9
$1.000\mathrm{e}-3$ |  0.962 $\pm 0.021$ | 45 | 0.262 | 42.1
$4.000\mathrm{e}-3$ |  0.925 $\pm 0.040$ | 25 | 0.276 | 23.9
$1.600\mathrm{e}-2$ |  0.444 $\pm 0.111$ | 10 | 0.301 | 9.3
$6.400\mathrm{e}-2$ |  0.243 $\pm 0.050$ | 4 | 0.310 | 3.8
$\Lambda(4\mathrm{e}-1,5\mathrm{e}-4,0.80)$ |  0.998 $\pm 0.002$ | 45 | 0.208 | 46.2
$\Lambda(4\mathrm{e}-1,5\mathrm{e}-4,0.85)$ |  1.000 $\pm 0.001$ | 50 | 0.190 | 49.1
$\Lambda(4\mathrm{e}-1,5\mathrm{e}-4,0.90)$ |  0.995 $\pm 0.002$ | 55 | 0.202 | 54.4
Table A.3: For the mixture data in 25D as in (45), we consider the epoch that
minimizes monitor $\widehat{\rm MSE}^{(m)}_{p}$, at which point we display the
average testing power (${n_{\rm GoF}}=500$) and $\widehat{\rm MSE}_{q}$ over
$n_{\rm replica}=10$ models for each regularization strategy.

###  B.2 1D Gaussian mixture

Consider a setting in which $q$ is a 1-dimensional Gaussian mixture of two
equally-weighted, unit variance components at positions -1 and 1. Let $p$ be a
similar mixture, however at positions -0.8 and 1, and the component at 1 has
variance 0.25. Therefore, the score of $q$ has the form:

$\displaystyle{\bf s}_{q}(x)=\nabla q(x)/q(x)=$
$\displaystyle\left(\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{(x+1)^{2}}{2}\right)+\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{(x-1)^{2}}{2}\right)\right)^{-1}\times$
$\displaystyle\left(-\frac{x+1}{\sqrt{2\pi}}\exp\left(-\frac{(x+1)^{2}}{2}\right)-\frac{x-1}{\sqrt{2\pi}}\exp\left(-\frac{(x-1)^{2}}{2}\right)\right),$

where ${\bf s}_{p}$ has similar form. Therefore, the scaleless optimal critic
${\bf f}^{*}={\bf s}_{q}-{\bf s}_{p}$ (7) has analytical form:

$\displaystyle{\bf f}^{*}(x)=$
$\displaystyle\left(\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{(x+1)^{2}}{2}\right)+\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{(x-1)^{2}}{2}\right)\right)^{-1}\times$
(A.5)
$\displaystyle\left(-\frac{x+1}{\sqrt{2\pi}}\exp\left(-\frac{(x+1)^{2}}{2}\right)-\frac{x-1}{\sqrt{2\pi}}\exp\left(-\frac{(x-1)^{2}}{2}\right)\right)-$
$\displaystyle\left(\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{(x+0.8)^{2}}{2}\right)+\frac{1}{(0.5)\sqrt{2\pi}}\exp\left(-\frac{(x-1)^{2}}{2\times(0.5)^{2}}\right)\right)^{-1}\times$
$\displaystyle\left(-\frac{x+1}{\sqrt{2\pi}}\exp\left(-\frac{(x+0.8)^{2}}{2}\right)-\frac{x-1}{(0.5)^{3}\sqrt{2\pi}}\exp\left(-\frac{(x-1)^{2}}{2\times(0.5)^{2}}\right)\right),$

which can be computed in a similar analytical manner in higher dimension given
distributions of the form (45).

A neural Stein critic (MLP of 2 hidden layers with 512-node width) is trained
using staged $\Lambda(1,10^{-3},0.9)$ regularization in this setting given
1,000 training samples from $q$ with batch size 200, learning rate $10^{-3}$,
and $B_{w}=5$.

Figure A.1:  Comparison between the testing power of the KSD GoF test in the
50D setting described in Section 6.2.1 for a range of bandwidth choices with
$n_{\rm sample}=1500$. The horizontal axis indicates the scaling factor
$\delta$ of the kernel bandwidth (see Appendix B.3) and the vertical axis
displays the test power. The red line and error bars display the mean and
standard deviation of the testing power for each choice of kernel bandwidth,
and the black line and shaded region indicate the mean and standard deviation
of the testing power of the neural Stein GoF test in this setting.

###  B.3 Kernel Stein discrepancy goodness-of-fit test

As in Section 5.1, consider the setting in which we are provided with
probability distribution $q$ supported on $\mathcal{X}\subset\mathbb{R}^{d}$
from which we can sample, and consider the scenario wherein we are provided
with a finite sample of data $x_{i}\sim p$ of ${n_{\rm GoF}}$ samples. As an
alternative to learning the Stein critic function using a neural network for
estimation of the test statistic in (43), one may define the function space
$\mathcal{F}$ to be an RKHS defined by a kernel denoted $k(\cdot,\cdot)$. In
this case, the KSD admits computation in closed form by the following
relation:

${\rm KSD}=\mathbb{E}_{x,x^{\prime}\sim p}[u_{q}(x,x^{\prime})],$ (A.6)

where

$\displaystyle u_{q}(x,x^{\prime})=$ $\displaystyle\
\mathbf{s}_{q}(x)^{\text{T}}k(x,x^{\prime})\mathbf{s}_{q}(x^{\prime})+\mathbf{s}_{q}(x)^{\text{T}}\nabla_{x^{\prime}}k(x,x^{\prime})$
$\displaystyle+\nabla_{x}k(x,x^{\prime})^{\text{T}}\mathbf{s}_{q}(x^{\prime})+{\rm
tr}(\nabla_{x,x^{\prime}}k(x,x^{\prime})).$ (A.7)

Using this formulation, [8] defines a quadratic-time $V$-statistic which
estimates the KSD as follows:

$\hat{V}_{q}=\frac{1}{n_{\rm GoF}^{2}}\sum_{i,j=1}^{{n_{\rm
GoF}}}u_{q}(x_{i},x_{j}).$ (A.8)

Given the test statistic (A.8), a “wild bootstrap” is used to approximate the
distribution of the test statistic under the null hypothesis, which is
outlined in [8]. Given the test set of ${n_{\rm GoF}}$ samples $x_{i}\sim p$,
the null hypothesis may be rejected if (A.8) computed using these $x_{i}$
exceeds the ($1-\alpha$) quantile of the wild bootstrapped distribution under
the null hypothesis (where, again, $\alpha$ is selected to tune the Type-I
error). The power is then estimated using $n_{\rm run}$ number of such GoF
tests.

Figure A.2:  Comparison between neural Stein critic GoF test in the 50D
setting described in Appendix B.4 for various proportions of training/testing
split, where the testing procedure and power/time computation is otherwise the
same as that in Section 6.2.1. In all cases, the total number of points is
$n_{\rm sample}=500$. The power and duration of the KSD test using the
selected bandwidth according to Appendix B.3 for $n_{\rm sample}=500$ is shown
in red in each figure. The x-axis of each plot refers to the percentage of
these points which are allocated to training for the neural network test. The
remainder are used for computation of the neural Stein test statistic. The
y-axes of (A), (B), and (C) correspond to those of Figure 6.

Following the approaches of [8, 27], we use an RBF kernel to define the RKHS
used for KSD in our experiments. This kernel has the following form:

$k(x,x^{\prime})=\exp\left(-\frac{\|x-x^{\prime}\|_{2}^{2}}{2\sigma^{2}}\right)=\exp\left(-\gamma\|x-x^{\prime}\|_{2}^{2}\right),$
(A.9)

where $\sigma$ is the kernel bandwidth and $\gamma=1/(2\sigma^{2})$. A
standard heuristic we employ is the choice of the median Euclidean distance
between the data considered as the bandwidth $\sigma$ of this kernel. However,
we also conduct an experiment to select a bandwidth which has the capacity to
achieve higher testing power when used in the KSD GoF test. We consider
bandwidths which are scaled versions of the median data distances heuristic
bandwidth. That is, we compute the testing power for bandwidth factors
$\delta\in\\{2^{-6},2^{-5},2^{-4},2^{-3},2^{-2},2^{-1},2^{0},2^{1},2^{2}\\}$,
where the scaling follows $\gamma^{\prime}=1/(2\delta\sigma^{2})$ for $\sigma$
equal to the median Euclidean distance between data samples. To do so, we
consider the same model and data distributions as in Section 6.2.1. We then
fix $n_{\rm sample}=1500$, and compute the KSD test power from $n_{\rm
run}=400$ GoF tests using $n_{\rm boot}=500$ bootstrapped samples in each
test. This is conducted using $n_{\rm replica}=5$ replicas to generate a mean
and standard deviation of power computed for each choice of bandwidth.
Finally, we compare this to the GoF testing power computed from $n_{\rm
replica}=5$ neural Stein critics trained using $n_{\rm sample}=1500$ samples
split into a $50\%/50\%$ train/test split (where $20\%$ of the train split is
used for validation). The result can be seen in Figure A.1. This shows that
the factor of $\delta=2^{-4}$ maximizes power in this case. Therefore, this
factor is chosen to create the “best” selected bandwidth for the experiments
in Section 6.2.1. This figure also displays that the KSD GoF test power is
still less than that of the neural Stein GoF test, despite better choices of
bandwidth.

Figure A.3:  Comparison between neural Stein critic GoF test and KSD test in
the 50D setting described in Appendix B.4. Other than the change in
distribution perturbation parameters, the experiment is identical to the
experiments described in Section 6.2.1. The sample size $n_{\rm sample}$
provided to each method is indicated on the x axis. The neural method uses
half of these data for training and the other half for testing, and the KSD
uses all $n_{\rm sample}$ samples for test statistic computation. The mean and
standard deviation of the testing power are plotted in each case, which is
computed over 5 replicas.

###  B.4 Comparison over different training-test splits

In Section 6.2.1, we outline a procedure for comparing the neural Stein test
to the KSD GoF test by considering varying overall sample size $n_{\rm
sample}$, training the neural network using half of these samples and
conducting tests using the other half. This section analyzes the power and
running time of the neural Stein GoF hypothesis testing procedure as outlined
in Section 5.1 for a variety of training/testing splits. Using the notation of
Section 6.1.1 for the model and data distributions, let the model distribution
$q$ be an isotropic, two-component Gaussian mixture in $\mathbb{R}^{50}$ with
means $\mu_{1}=\mathbf{0}_{d}$ and $\mu_{2}=0.5\times\mathbf{1}_{d}$, and let
the data distribution $p$ have the form of (45) with covariance shift
$\rho_{1}=0.99$ and scale $\omega=0.05$. In this section, we fix $n_{\rm
sample}=500$ and consider other choices of training/testing split, comparing
to the KSD with “best” selected bandwidth (see Appendix B.3) for $n_{\rm
sample}=500$. That is, we partition the 500 points into splits whereby the
number of training samples takes on a proportion in the range of
$\left\\{10\%,20\%,30\%,40\%,50\%,60\%,70\%,80\%,90\%\right\\}$ of $n_{\rm
sample}$. The testing procedure is otherwise identical to that of Section
6.2.1. The result can be seen in Figure A.2. Figure A.2(A) shows that there
exists a range of splits for which the power is nearly one and that since the
computation time is dominated by the training, as seen in Figures A.2(B) and
(C), a 50%/50% split is a good generic choice of the split. Therefore, this
split is used in the experiments outlined in Section 6.2.1 and displayed in
Figure 6.

We also consider the comparison of testing power in this case for the neural
Stein critic GoF test and the KSD test with the median data distances
heuristic RBF kernel bandwidth and the KSD with bandwidth selected to maximize
power according to the result outlined in Appendix B.3. The methodology for
comparing the power between these tests is identical to that of Section 6.2.1.
That is, for each sample size, the train/test split is $50\%/50\%$, with
$20\%$ of the training partition dedicated to validation. The power is then
computed in each case using $n_{\rm run}=400$ GoF tests, $n_{\rm boot}=500$
bootstrapped samples in each individual test, efficient bootstrap ratio
$r_{\rm pool}=50$, and $n_{\rm replica}=5$. The result can be found in Figure
A.3, which shows that the power of the neural Stein critic GoF test dominates
that of both KSD test methods initially. However, the neural Stein test
achieves near unit power by the time $n_{\rm sample}=500$, and the KSD with
selected bandwidth achieves unit power by the time $n_{\rm sample}=1,000$.
This increase in power compared to Figure 6(A) is due to the fact that the
perturbation between the model and data distribution in this scenario is
larger in this setting than in the setting outlined in Section 6.2.1.

###  B.5 Adaptive Staging

Figure A.4:  The $\widehat{\rm MSE}_{p}$ (A) and $\widehat{\rm MSE}_{q}$ (B)
of adaptive staged regularization strategies $\Lambda_{a}$ and
$\tilde{\Lambda}_{a}$ compared to that of nonadaptive staged regularization
$\Lambda$ and fixed-$\lambda$ strategies. The trajectories of $\lambda$ over
the course of training in the staged regularization strategies is depicted in
(C). See Appendix B.5 for a description of the experiment setup, which is
largely the same as the experiments described in Section 6.1.2. The mean MSE
over 10 models is represented as a solid line for each regularization
strategy, while the standard deviation is represented as the shaded region.
Similarly, for the adaptive staging strategies, the mean value of $\lambda$
for each stage is represented by the solid line, while the standard deviation
is reflected by the shaded region.

In Section 3.2, we specified a method for staging the weight of regularization
$\lambda$ via an exponential decay over the course of optimization. Here we
develop an adaptive procedure which stages $\lambda$ according to the monitor
$\widehat{\rm MSE}^{(m)}_{p}$ (17), and compare with the non-adaptive staging
in (15).

The adaptive staging method decides to stage the weight of regularization down
to a smaller value when the monitored $\widehat{\rm MSE}^{(m)}_{p}$ is
observed to increase. Specifically, we examine two adaptive staging approaches

  * •

The strategy, denoted as $\Lambda_{a}(\lambda_{\rm init},\lambda_{\rm
term},\beta)$, begins with $\lambda=\lambda_{\rm init}$. Over the course of
training, $\widehat{\rm MSE}^{(m)}_{p}$ is computed after every $B_{w}=10$
mini-batches of updates given $n_{\rm val}$ samples from $p$ (in these
examples, the validation samples are re-drawn from $p$ for each computation).
If $\widehat{\rm MSE}^{(m)}_{p}$ is at any observation greater than the
previous observation, the weight of regularization $\lambda$ is multiplied by
the factor $\beta<1$. This process is repeated until $\lambda=\lambda_{\rm
term}$, at which point $\lambda$ is fixed for the remainder of training. We
require that $\widehat{\rm MSE}^{(m)}_{p}$ decrease at least once with each
stage of $\lambda$, so that $\lambda$ may not stage down until first an
improvement followed by a decrease in performance in MSE is observed.

  * •

The staging, denoted by $\tilde{\Lambda}_{a}$, uses $\widehat{\rm MSE}_{p}$
instead of $\widehat{\rm MSE}^{(m)}_{p}$ as the monitor function to determine
the staging. For simulated high dimensional Gaussian mixture data, the
computation of $\widehat{\rm MSE}_{p}$ uses the prior information of the
scaleless optimal critic ${\bf f}^{*}$.

To assess the performance of the adaptive staging strategy, we conduct an
experiment using the Gaussian mixture data of Section 6.1.1 in 25D with
$\omega=0.8$ and $\rho_{1}=0.5$. The neural Stein critic architecture and
initialization are as in Section 6.1.2. We again use the Adam optimizer with
default parameters, with learning rate $10^{-3}$ and batch size 200 samples.
In this setting, the networks are trained using 4,000 samples from $p$ for 60
epochs each. For each staging strategy, 10 critics are trained. The
$\widehat{\rm MSE}^{(m)}_{p}$ and $\widehat{\rm MSE}_{p}$ are computed each
using $n_{\rm val}=20,000$ samples from $p$ and the $\widehat{\rm MSE}_{q}$ is
computed using $n_{\rm te}=20,000$ samples from $q$. We compare the staging to
fixed regularization strategies with $\lambda\in\\{1\times 10^{-3},4\times
10^{-3},1.6\times 10^{-2}\\}$. The staging strategies used are
$\Lambda(3\times 10^{-1},5\times 10^{-4},0.85)$, $\Lambda_{a}(3\times
10^{-1},5\times 10^{-4},0.70)$, and $\tilde{\Lambda}_{a}(3\times
10^{-1},5\times 10^{-4},0.70)$, where an update is made according to the
nonadaptive staging strategy every $B_{w}=10$ batches.

The $\widehat{\rm MSE}_{p}$ and $\widehat{\rm MSE}_{q}$ in the adaptive
staging comparison experiments are shown in Figure A.4(A) and (B),
respectively. As in Section 6.1.3, all staging strategies outperform the
fixed-$\lambda$ strategies in 25D. Furthermore, the performances of the
adaptive strategies $\Lambda_{a}$ and $\tilde{\Lambda}_{a}$ and the non-
adaptive strategy $\Lambda$ are similar, with overall minimal MSE comparable
across all three staging strategies. In addition, a comparison of the
trajectory of $\lambda$ over the course of training via the three staged
regularization strategies is displayed in Figure A.4(C), indicating that the
adaptive staging strategies exhibit a roughly exponential annealing of
$\lambda$ over the course of training. This result justifies the rationale of
the heuristic, nonadaptive staging approach proposed in Section 3.2, which is
computationally less expensive than the adaptive staging which relies on the
repetitive computation of the monitor validation MSE.