#  Neural Stein critics with staged $L^{2}$-regularization

Matthew Repasky  H. Milton Stewart School of Industrial and Systems
Engineering, Georgia Institute of Technology Xiuyuan Cheng Email:
<EMAIL_ADDRESS>Department of Mathematics, Duke University  Yao Xie
H. Milton Stewart School of Industrial and Systems Engineering, Georgia
Institute of Technology

###### Abstract

Learning to differentiate model distributions from observed data is a
fundamental problem in statistics and machine learning, and high-dimensional
data remains a challenging setting for such problems. Metrics that quantify
the disparity in probability distributions, such as the Stein discrepancy,
play an important role in high-dimensional statistical testing. In this paper,
we investigate the role of $L^{2}$ regularization in training a neural network
Stein critic so as to distinguish between data sampled from an unknown
probability distribution and a nominal model distribution. Making a connection
to the Neural Tangent Kernel (NTK) theory, we develop a novel staging
procedure for the weight of regularization over training time, which leverages
the advantages of highly-regularized training at early times. Theoretically,
we prove the approximation of the training dynamic by the kernel optimization,
namely the “lazy training”, when the $L^{2}$ regularization weight is large,
and training on $n$ samples converge at a rate of ${O}(n^{-1/2})$ up to a log
factor. The result guarantees learning the optimal critic assuming sufficient
alignment with the leading eigen-modes of the zero-time NTK. The benefit of
the staged $L^{2}$ regularization is demonstrated on simulated high
dimensional data and an application to evaluating generative models of image
data.

Keywords: Stein Discrepancy, Goodness-of-fit Test, Neural Tangent Kernel, Lazy
Training, Generative Models

##  1 Introduction

Understanding the discrepancy between probability distributions is a central
problem in machine learning and statistics. In training generative models,
learning to minimize such a discrepancy can be used to construct a probability
density model given observed data, such as in the case of generative
adversarial networks (GANs) trained to minimize $f$-divergences [29],
Wasserstein GANs [2], and score matching techniques [21, 31]. Generally, GANs
and other generative models require discriminative critics to distinguish
between data and a distribution [13]; such critics have the ability to tell
the location where the model and data distributions differ as well as the
magnitude of departure. Recent developments in the training of generative
models have facilitated advancements in out-of-distribution detection [15], in
which such models learn to predict the higher likelihood for in-distribution
samples. There exists a wide array of integral probability metrics that
quantify distances on probability distributions [33], including the Stein
discrepancy [14]. In particular, the computation of the Stein discrepancy only
requires knowledge of the score function (the gradient of the log density) of
the model distribution. Thus, it avoids the need to integrate the normalizing
constant in high dimensions. This makes the Stein discrepancy useful for
evaluating some deep models like energy-based models.

Implicit in the minimization of the discrepancy between a model distribution
and observed data is the concept of goodness-of-fit (GoF). In the GoF problem
and the closely related two-sample test problem, the goal of the analysis is
to approximate and estimate the discrepancy between two probability
distributions. Integral probability metrics are widely used for such problems.
For example, Reproducing Kernel Hilbert space (RKHS) Maximum Mean Discrepancy
(MMD) [17], a kernel-based approach, is used for two-sample testing among
other testing tasks. Kernels parameterized by deep neural networks have been
adopted recently in [26] to improve the testing power. For GoF tests, where
the task is to detect the departure of an unknown distribution of observed
data from a model distribution, methods using the Stein discrepancy metrics
have been developed. The Stein discrepancy is also calculated using kernel
methods [27, 28, 8] and more recently using deep neural network-aided
techniques [16]. We provide more background information related to the Stein
discrepancy and its role in GoF testing in Section 2.

In machine learning, a wide array of modern generative model architectures
exist. Energy-based models (EBMs) are a particularly useful subset of
generative models. Such models can be described by an energy function which
describes a probability density up to a normalizing constant [36]. While such
models provide flexibility in representing a probability density, the
normalizing constant (which requires an integration over the energy function
to compute) is required to compute the likelihood of data given the model. The
Stein discrepancy provides a metric for evaluating EBMs without knowledge of
this normalization constant [16]. Another popular class of generative models
is flow-based models [9, 10]. Flow-based modeling approaches, such as RealNVP
and Glow [24], provide reversible and efficient transformations representing
complex distributions, yielding simple log-likelihood computation. In Section
5.2, we outline our approach for evaluating generative EBM models using neural
Stein critics.

In this paper, we introduce a method for learning the Stein discrepancy via a
novel staged regularization strategy when training neural network Stein
critics. We consider the $L^{2}$-regularization of the neural Stein critic,
which has been adopted in past studies of neural network Stein discrepancy
[19, 16]. Another motivation to use $L^{2}$-regularization is due to the fact
that the (population) objective of the $L^{2}$-regularized Stein discrepancy
is equivalent to the mean-squared error between the trained network critic and
the optimal one (up to an additive constant); see (11). It has also been shown
previously that the Stein discrepancy evaluated at the optimal critic under
$L^{2}$ regularization reveals the Fisher divergence [19]. We analyze the role
of the regularization strength parameter in $L^{2}$ neural Stein methods,
emphasizing its impact on neural network optimization, which was overlooked in
previous studies. On the practical side, our study shows the benefit of
softening the impact of this regularization over the course of training,
yielding critics which fit more quickly at early times, followed by stable
convergence with weaker regularization. An example is shown in Figure 1:
(A)-(C) illustrate that the target critic changes in magnitude throughout
training as the weight of regularization is decreased, and (D)-(F) shows the
rough approximation to the optimum at early times followed by more nuanced
changes in the later stages of training.

The proposed staging of $L^{2}$-regularization is closely connected to the so-
called “lazy training” phenomenon of neural networks [7], which approximately
solves a kernel regression problem corresponding to the Neural Tangent Kernel
(NTK) [22] at the early training times. Theoretically, we prove the kernel
learning dynamic of $L^{2}$ neural Stein methods when the regularization
weight is large - the main observation is that the role played by the
regularization weight parameter is equivalent to a scaling of the neural
network function, which itself leads to kernel learning in the case of strong
regularization. This theoretical result motivates the usage of a large penalty
weight in the beginning phase of training before decreasing it to a smaller
value later.

For GoF problems, the trained neural Stein critic provides model comparison
capabilities that assess the accuracy of a model’s approximation of the true
distribution, allowing for identifying the locations of distribution departure
in observed data. This naturally leads to applications for GoF testing and
evaluation of generative models. In summary, the contributions of the current
work are as follows:

  1. 1.

We introduce a new method for training neural Stein critics, which
incorporates a staging of the weight of the $L^{2}$ regularization over the
process of mini-batch training.

  2. 2.

We prove the NTK kernel learning (lazy-training dynamic) of neural Stein
critic training with large $L^{2}$ regularization weight, providing a
theoretical justification of the benefit of using strong $L^{2}$
regularization in the early training phase. The analysis reveals a convergence
at the rate of $O(n^{-1/2})$ up to a log factor, $n$ being the sample size,
when training with finite-sample empirical loss.

  3. 3.

The advantage of the proposed neural Stein method is demonstrated in
experiments, exhibiting improvements over fixed-regularization neural Stein
critics and the kernelized Stein discrepancy on simulated data. The neural
Stein critic is applied to evaluating generative models of image data.

###  1.1 Related works

The Stein discrepancy has been widely used in various problems in machine
learning, including measuring the quality of sampling algorithms [14],
evaluating generative models by diffusion kernel Stein discrepancy which
unifies score matching with minimum Stein discrepancy estimators [5], GoF
testing, among others. For GoF testing, a kernel Stein discrepancy (KSD)
approach has allowed for closed-form computation of the discrepancy metric
[27, 8]. Similar metrics have been used in the GoF setting, such as the finite
set Stein discrepancy (FSSD), which behaves as the KSD but can be computed in
linear time [23]. Our work leverages neural networks, which potentially have
large expressiveness, to parameterize the critic function space and studies
the influence of $L^{2}$ regularization from a training dynamic point of view.

Recent studies have developed alternatives to kernelized approaches to
computing the Stein discrepancy. Using neural networks to learn Stein critic
functions, [19] applied neural Stein in the training of high-quality samplers
from un-normalized densities. The neural network Stein critic has also been
applied to the GoF hypothesis test settings to evaluate EBMs [16]. These
methods impose a boundedness constraint on the $L^{2}$ norm of the functions
represented by the neural network using a regularization term added to the
training objective. The optimal critic associated with this method yields a
Stein discrepancy equivalent to the Fisher divergence and provides an
additional benefit in that the critic can be used as a diagnostic tool to
identify regions of poor fit [19]. The staging scheme we introduce in this
work is motivated by the observed connection between large $L^{2}$ penalty and
lazy training. In practice, it yields an improvement upon past techniques used
to learn the $L^{2}$-penalized critic.

Many methods have been developed to train and evaluate generative models
without knowledge of the likelihood of a model. Early works used the method of
Score Matching, which minimizes the difference in score function between the
data and model distributions using a proxy objective [21]. Methods building on
this approach are known as score-based methods. Score-based methods include
approaches that can estimate the normalizing constant for computation of
likelihoods, as in the case of Noise-Contrastive Estimation [18], and can
conduct score matching using deep networks with robust samplers, as in the
case of Noise Conditional Score Networks with Langevin sampling [31]. Our
approach potentially provides a more efficient training scheme to obtain
discriminative Stein discrepancy critics, leading to a metric representing the
discrepancy between distributions. We also experimentally observe that the
trained Stein critic function indicates the differential regions between
distributions of high dimensional data.

The Stein divergence can be interpreted as a divergence to measure the
discrepancy between two distributions. The estimation error analysis of neural
network function class-based divergence measure has been studied in [4, 40,
32], where the global optimizer within the neural network function class on
the empirical loss is assumed without addressing the optimization guarantee.
To incorporate the analysis of neural network training dynamics, the current
paper utilizes theoretical understandings developed by the NTK theory [22, 11,
3]. Particularly, we utilize ideas surrounding the “lazy training” phenomenon
of neural networks, suggesting the loss through training decays rapidly with a
relatively small change to the parameters of the network model, which results
in a kernel regression optimization dynamic [7]. For neural network hypothesis
testing, the NTK learning was used in computing neural network MMD for two-
sample testing [6]. As for when neural network training falls under the lazy
training regime, [7] showed that it depends on a choice of scaling of the
network mapping. In this paper, we adopt a staging of the weight of $L^{2}$
regularization, which can utilize the NTK kernel learning regime at early
periods of training, and also, in practice, go beyond kernel learning at later
phases.

###  1.2 Notation

Notations in this paper are mostly standard, with a few clarifications as
follows: We use $\partial_{\rm x}$ with subscript x to denote partial
derivative, e.g., $\partial_{\theta}f(x,\theta)$ means
$\frac{\partial}{\partial\theta}f(x,\theta)$. We use $\mathbb{E}_{x\sim p}$ to
denote integral over the measure $p(x)dx$, that is, $\mathbb{E}_{x\sim
p}f(x)=\int_{\mathcal{X}}f(x)p(x)dx$. We may omit the variable as a short-hand
notation of integrals, that is, we write $\int_{\mathcal{X}}f$ for
$\int_{\mathcal{X}}f(x)dx$. The symbol $\cdot$ stands for vector-vector inner-
product (also used in the divergence operator $\nabla\cdot$), and $\circ$
stands for matrix-vector multiplication.

##  2 Background

We begin by providing necessary preliminaries of the Stein Discrepancy,
$L^{2}$ Stein critics, and GoF testing.

###  2.1 Stein Discrepancy

The recent works on Stein discrepancy in machine learning are grounded in the
theory of Stein’s operator, the study of which dates back to statistical
literature of the 1970s [34]. Let $\mathcal{X}$ be a domain in
$\mathbb{R}^{d}$, and we consider probability densities on $\mathcal{X}$.
Given a density $q$ on $\mathcal{X}$, for a sufficiently regular vector field
${\bf f}:\mathcal{X}\to\mathbb{R}^{d}$, the Stein operator $T_{q}$ [34, 1]
applied to ${\bf f}$ is defined as

$T_{q}{\bf f}(x):={\bf s}_{q}(x)\cdot{\bf f}(x)+\nabla\cdot{\bf f}(x),$ (1)

where ${\bf s}_{q}$ is the score function of $q$, defined as

${\bf s}_{q}:={\nabla q}/{q}=\nabla\log q.$ (2)

Note that the divergence of ${\bf f}$, denoted by $\nabla\cdot{\bf f}(x)$, is
the same as the trace of the Jacobian of ${\bf f}$ evaluated at $x$, and
throughout this paper $\cdot$ stands for vector inner-product.

Given another probability density $p$ on $\mathcal{X}$ and a bounded function
class $\mathcal{F}$ of sufficient regularity, the Stein discrepancy [14] which
measures the difference between $p$ and $q$ is defined as

${\rm SD}_{\mathcal{F}}(p,q):=\underset{{\bf f}\in\mathcal{F}}{\sup}\,{\rm
SD}[{\bf f}],\quad{\rm SD}[{\bf f}]:=\mathbb{E}_{x\sim p}T_{q}{\bf f}(x).$ (3)

In this paper, we call ${\bf f}$ the “critic”, and we call ${\rm SD}[{\bf f}]$
the Stein discrepancy evaluated at the critic ${\bf f}$. We further consider
the class of ${\bf f}$ that satisfy a mild boundary condition, namely
$\mathcal{F}_{0}(p):=\\{{\bf f},\,p{\bf f}|_{\partial X}=0\\}$. (When $X$ is
unbounded, $|_{\partial X}$ is understood as the infinity boundary condition.
When $X$ is bounded, sufficient regularity of $\partial X$ is assumed so that
the divergence theorem can apply). When $p=q$, one can verify that for any
${\bf f}\in\mathcal{F}_{0}(p)$ and sufficiently regular, ${\rm SD}[{\bf
f}]=0$, which is known as the Stein’s identity. The other direction, namely
for some $\mathcal{F}\subset\mathcal{F}_{0}(p)$, ${\rm
SD}_{\mathcal{F}}(p,q)=0\Rightarrow p=q$, is also established under certain
conditions [35]. This explains the name “critic” because the ${\bf f}$ which
makes ${\rm SD}[{\bf f}]$ significantly large can be viewed as a test function
(vector field) that indicates the difference between $p$ and $q$.

To proceed, we introduce the following assumption on the densities $p$ and
$q$:

######  Assumption 1.

The two densities $p$ and $q$ are supported and non-vanishing on $\mathcal{X}$
(vanishing at the boundary of $\mathcal{X}$), differentiable on $\mathcal{X}$,
and the score functions ${\bf s}_{p}$ and ${\bf s}_{q}$ are in
$\mathcal{F}_{0}(p)$.

When the function class $\mathcal{F}$ is set to be the unit ball in an RKHS,
the definition (3) yields the kernelized Stein discrepancy (KSD) [27, 8]. A
possible limitation of the KSD approach is the sampling complexity and
computational scalability in high dimension. In this work, we consider critics
of regularized $L^{2}$ norm to be parametrized by neural networks as
previously studied in [19, 16], which have potential advantages in model
expressiveness and computation.

###  2.2  $L^{2}$ Stein critic

This paper focuses on when the critic ${\bf f}$ in Stein discrepancy (3) is at
least squared integrable on $(\mathcal{X},p(x)dx)$. Consider the critic
function ${\bf f}:\mathcal{X}\to\mathbb{R}^{d}$, and the $d$ coordinates of
${\bf f}$ are denoted as ${\bf f}(x)=(f_{1}(x),\cdots,f_{d}(x))^{T}$. We first
introduce notations of the $L^{2}$ space of vector fields. Define the inner-
product and $L^{2}$-norm of vector fields on $(\mathcal{X},p(x)dx)$ as the
following: for ${\bf v},{\bf w}:\mathcal{X}\to\mathbb{R}^{d}$, let

$\langle{\bf v},{\bf w}\rangle_{p}:=\int_{\mathcal{X}}{\bf v}(x)\cdot{\bf
w}(x)p(x)dx,$ (4)

and then the $L^{2}$ norm is defined as

$\|{\bf v}\|_{p}^{2}:=\langle{\bf v},{\bf v}\rangle_{p}.$ (5)

We stress that, throughout this paper, subscript p in the norm means 2-norm
weighted by the measure $p(x)dx$. We denote all critics ${\bf
f}:\mathcal{X}\to\mathbb{R}^{d}$ such that $\|{\bf f}\|_{p}^{2}<\infty$ the
space of $L^{2}(p):=L^{2}(\mathcal{X},p(x)dx)$.

The Stein discrepancy over the class of $L^{2}$ critics with bounded $L^{2}$
norm is defined, for some $r>0$, as

${\rm SD}_{r}(p,q)=\underset{{\bf f}\in\mathcal{F}_{0}(p),\,\|{\bf
f}\|_{p}\leq r}{\sup}{\rm SD}[{\bf f}],$ (6)

where ${\rm SD}[{\bf f}]$ is defined as in (3). Define

${\bf f}^{*}:={\bf s}_{q}-{\bf s}_{p},$ (7)

the following lemma characterizes the solution of (6), and the proof (left to
Appendix A) also derives a useful equality when ${\bf f}$ and ${\bf f}^{*}$
are in $L^{2}(p)\cap\mathcal{F}_{0}(p)$, see also e.g. [27, Lemma 2.3]:

${\rm SD}[{\bf f}]=\mathbb{E}_{x\sim p}T_{q}{\bf f}(x)=\langle{\bf f}^{*},{\bf
f}\rangle_{p}.$ (8)

######  Lemma 2.1.

For any $r>0$, suppose ${\bf s}_{q}$ and ${\bf s}_{p}$ are in $L^{2}(p)$ and
$\mathcal{F}_{0}(p)$, then ${\rm SD}_{r}(p,q)=r\|{\bf f}^{*}\|_{p}<\infty$
and, if $\|{\bf f}^{*}\|_{p}>0$, the supremum of (6) is achieved at ${\bf
f}=({r}/{\|{\bf f}^{*}\|_{p}}){\bf f}^{*}$.

Lemma 2.1 implies that if $\|{\bf f}^{*}\|_{p}=0$, then ${\rm SD}_{r}(p,q)=0$.
Thus, we take the following assumption, which restricts the case where the
$L^{2}$ Stein critic can achieve a positive discrepancy when $q\neq p$.

######  Assumption 2.

The score functions ${\bf s}_{p}$ and ${\bf s}_{q}$ are in $L^{2}(p)$, and
when $q\neq p$, $\|{\bf s}_{q}-{\bf s}_{p}\|_{p}>0$.

###  2.3 Goodness-of-Fit tests

In GoF testing, we are presented with samples $X=\\{x_{i}\\}_{i}$ drawn from
an unknown distribution $p$, and we wish to assess whether this sample is
likely to have come from the model distribution $q$. That is, we may define
the null and alternative hypotheses as follows:

$H_{0}:p=q,\quad\quad H_{1}:p\neq q.$ (9)

A GoF test is conducted using a test statistic $\hat{T}=\hat{T}(X)$ computed
using observed samples in $X$. After specifying a number $t_{\rm thresh}$,
which is called the “test threshold”, the null hypothesis is rejected if
$\hat{T}>t_{\rm thresh}$. There are different approaches to specifying a test
threshold, for example, if prior knowledge of the distribution of $\hat{T}$
under $H_{0}$ is available, it can be used to choose $t_{\rm thresh}$. In the
experiments of this work, we compute $t_{\rm thresh}$ using a bootstrap
strategy from samples from the model distribution $q$.

The selection of $t_{\rm thresh}$ is to control the Type-I error, defined as
$\Pr[\hat{T}>t_{\rm thresh}]$ under $H_{0}$. The randomness of $\Pr$ is with
respect to data $X$. The goal is to guarantee that $\Pr[\hat{T}>t_{\rm
thresh}]\leq\alpha$, which is called the the “significance level” (typically,
$\alpha=0.05$). The Type-II error measures the probability that the null
hypothesis is improperly accepted as true, that is, $\Pr[\hat{T}\leq t_{\rm
thresh}]$ under $H_{1}$. Finally, the “test power” is defined as one minus the
Type-II error. For the application to GoF testing, the current work develops a
test statistic computed using a Stein critic parameterized and learned by a
neural network, computed from a training-testing split of the dataset. More
details will be introduced in Section 5.1.

##  3 Method

###  3.1 Neural $L^{2}$ Stein critic

Replacing the $L^{2}$-norm constraint in (6) to be a regularization term leads
to the following minimization over a certain class of ${\bf f}$ (inside
$L^{2}(p)\cap\mathcal{F}_{0}(p)$) as

$\mathcal{L}_{\lambda}[{\bf f}]:=-{\rm SD}[{\bf f}]+\frac{\lambda}{2}\|{\bf
f}\|_{p}^{2}=\mathbb{E}_{x\sim p}\left(-T_{q}{\bf
f}(x)+\frac{\lambda}{2}\|{\bf f}(x)\|^{2}\right),$ (10)

where $\lambda>0$ is the penalty weight of the $L^{2}$ regularization. Under
Assumption 1-2, the equality (8) gives that

$\displaystyle\mathcal{L}_{\lambda}[{\bf f}]$ $\displaystyle=-\langle{\bf
f}^{*},{\bf f}\rangle_{p}+\frac{\lambda}{2}\langle{\bf f},{\bf
f}\rangle_{p}=\frac{1}{2\lambda}(\|\lambda{\bf f}-{\bf f}^{*}\|_{p}^{2}-\|{\bf
f}^{*}\|_{p}^{2}).$ (11)

Since $\|{\bf f}^{*}\|_{p}^{2}$ is a constant independent of ${\bf f}$, (11)
immediately gives that $\mathcal{L}_{\lambda}[{\bf f}]$ is minimized at

${\bf f}_{\lambda}^{*}:=\frac{{\bf f}^{*}}{\lambda}=\frac{1}{\lambda}({\bf
s}_{q}-{\bf s}_{p}),$ (12)

see also Theorem 4.1 of [19]. The expression (12) reveals an apparent issue:
if one only considers global functional minimization, then the choice of
$\lambda$ plays no role but contributes a scalar normalization of the optimal
critic, and consequently does not affect the learned Stein critic in practice,
e.g., in computing test statistics. A central observation of the current paper
is that the choice of $\lambda$ plays a role in optimization, specifically, in
training neural network parameterized Stein critics.

Following [19, 16], we parameterize the critic by a neural network mapping
${\bf f}(x,\theta)$ parameterized by $\theta$, and we assume that ${\bf
f}(\cdot,\theta)\in L^{2}(p)$ for any $\theta$ being considered. We denote
this ${\bf f}(\cdot,\theta)$ the “neural Stein critic”. The population loss of
$\theta$ follows from (10) as

${L}_{\lambda}(\theta)=\mathcal{L}_{\lambda}[{\bf f}(\cdot,\theta)].$ (13)

The neural Stein critic is trained by minimizing the empirical version of
${L}_{\lambda}(\theta)$ computed from finite i.i.d. samples $x_{i}\sim p$,
i.e., given $n_{\rm tr}$ training samples, the empirical loss is defined as

$\hat{L}_{\lambda}(\theta):=\frac{1}{n_{\rm tr}}\sum_{i=1}^{n_{\rm
tr}}(-T_{q}{\bf f}(x_{i},\theta)+\frac{\lambda}{2}\|{\bf
f}(x_{i},\theta)\|^{2}).$

######  Remark 3.1 ($\lambda$ as scaling parameter).

We note a “$\lambda$-scaling” view of the $L^{2}$-regularized objective
${L}_{\lambda}(\theta)$, namely, it is equivalent to a scaling of the
parameterized function ${\bf f}(x,\theta)$ by $\lambda$. Specifically,
consider $\mathcal{L}_{1}[{\bf f}]$ as the “scaling-free” ($\lambda$-agnostic)
functional minimizing objective, then by definition, we have that

$\mathcal{L}_{\lambda}[{\bf f}]=\frac{1}{\lambda}\mathcal{L}_{1}[\lambda{\bf
f}].$ (14)

This $\lambda$-scaling view of the $L^{2}$ Stein discrepancy allows us to
connect to the lazy training of neural networks [7], which indicates that,
with large values of $\lambda$, the early stage of the training of the network
approximates the optimization of a kernel solution. We detail the theoretical
proof of such an approximation in Section 4. This view motivates our proposed
usage of large $\lambda$ in early training, which will be detailed in Section
3.2.

We close the current subsection by introducing a few notations to clarify the
$\lambda$-scaling indicated by (12). Suppose the neural Stein critic ${\bf
f}(\cdot,\theta)$ trained by minimizing $\hat{L}_{\lambda}(\theta)$
approximates the minimizer of ${\mathcal{L}}_{\lambda}$, namely the “optimal
critic” ${\bf f}_{\lambda}^{*}$, then both ${\bf f}(\cdot,\theta)$ and the
value of ${\rm SD}[{\bf f}(\cdot,\theta)]$ would scale like $1/\lambda$. We
call ${\bf f}^{*}$ defined in (7) the “scaleless optimal critic function”. The
expression (12) also suggests that, if the neural Stein critic successfully
approximates the optimal, we would expect

$\lambda{\bf f}(\cdot,\theta)\approx{\bf f}^{*},$

which will be confirmed by the theory in Section 4. We thus call $\lambda{\bf
f}(\cdot,\theta)$ the “scaleless neural Stein critic”.

Figure 1:  Visualization of staged $L^{2}$ regularization throughout training.
The two distributions $p$ and $q$ are Gaussian mixtures in 1D, and the
scaleless optimal critic ${\bf f}^{*}$ has explicit expression (A.5), see the
details of the distributions and optimization in Appendix B.2. In (A)-(C), the
trained neural Stein critic ${\bf f}(\cdot,\theta)$ is compared to the optimal
critic ${\bf f}^{*}_{\lambda}$ defined in (12). The optimal critic function
changes scale through training as the weight of regularization is decreased.
In (D)-(F), the scaleless neural Stein critic $\lambda{\bf f}(\cdot,\theta)$
is compared to the scaleless optimal critic ${\bf f}^{*}$ defined in (7). The
neural Stein critic roughly approximates the optimal critic at early times,
followed by fine-tuned changes in the later epochs of training.

###  3.2 Staged-$\lambda$ regularization in training

In practice, the neural network is optimized via the $\lambda$-dependent loss
$L_{\lambda}$. We propose a staging of the regularization weight $\lambda$
that begins with a large value followed by a gradual decrease until some
terminal value. Specifically, we adopt a log-linear staged regularization
scheme by which we decrease $\lambda$ via the application of a multiplicative
factor $\beta<1$ after a certain interval of training time measured by the
number of batches $B_{w}$. In addition to $B_{w}$, this staging scheme uses
three parameters: the initial weight $\lambda_{\rm init}$, the decay factor
$\beta\in(0,1)$, and the terminal weight $\lambda_{\rm term}$ (at which point
the decay ceases). We denote the discrete-time regularization staging with
these parameters as $\Lambda(B_{i};\lambda_{\rm init},\lambda_{\rm
term},\beta)$, where $B_{i}=i\cdot B_{w}$ stands for the $i$-th interval of
batches, $i\in\mathbb{N}$, and $\lambda$ will be set to the value on the
$i$-th interval. $i=0$ refers to the period before training has begun such
that $\Lambda(B_{0})=\lambda_{\rm init}$ and, for $i>0$,

$\Lambda\left(B_{i};\lambda_{\rm init},\lambda_{\rm
term},\beta\right)=\text{max}\left\\{\lambda_{\rm
init}\cdot\beta^{i},\lambda_{\rm term}\right\\}.$ (15)

That is, when $i$ increments of $B_{w}$ number of batches have occurred,
$\lambda_{\rm init}$ is annealed by a factor of $\beta^{i}$. $\beta$ can be
any positive number in $(0,1)$ and we typically choose $\beta$ to be about
0.7$\sim$ 0.9.

The beginning of the training with a large $\lambda$ is motivated by the
analysis in Section 4, which suggests that with large $\lambda$ at early times
of training, the training of the neural Stein critic can be approximately
understood from the perspective of kernel regression optimization, rapidly
reaching its best approximation in $\sim 1/\lambda$ time (up to a log factor),
see Theorem 4.6 (and Theorem 4.10 for the analysis with finite-sample loss).
In many cases, the kernel learning solution may not be sufficient to
approximate the optimal critic ${\bf f}^{*}$, and this calls for going beyond
the NTK kernel learning in training neural networks, which remains a
theoretically challenging problem. On the other hand, the proper
regularization strength $\lambda$, depending on the problem (the distributions
$p$ and $q$ and the sample size), may be much smaller than the initial large
$\lambda$, which we set to be $\lambda_{\rm term}$. We empirically observed
that gradually reducing the $\lambda$ in later training phases will allow the
neural Stein critic to progressively fit to ${\bf f}^{*}$, as illustrated in
Figure 1. Meanwhile, we also observed that using large $\lambda$ at the
beginning phase of training and then gradually tuning down to small $\lambda$
performs better than using small and fixed $\lambda$ throughout, see Section
6. These empirical results suggest the benefit of fully exploiting the kernel
learning (by using large $\lambda$) in the early training phase and annealing
to small $\lambda$ in the later phase of training.

To further study other possible staging methods, we investigated an adaptive
scheme of annealing $\lambda$ over training time by monitoring the validation
error, and the details are provided in Appendix B.5. On simulated high
dimensional Gaussian mixture data, the adaptive staging gives performance
similar to the scheme (15), and the resulting trajectory of $\lambda$ also
resembles the exponential decay as in (15), see Figure A.4. Thus we think the
scheme (15) can be a typical heuristic choice of the staging, and other
choices of annealing $\lambda$ are also possible.

###  3.3 Evaluation and validation metrics

We introduce the mean-squared error metric to quantitatively evaluate how the
learned neural Stein critic ${\bf f}(x,\theta)$ approximates the optimal
critic. In our experiments on synthetic data, we can use the knowledge of
${\bf f}^{*}$ (which requires knowledge of ${\bf s}_{p}$) and compute an
estimator of ${\rm MSE}_{q}:=\mathbb{E}_{x\sim q}\left\|\lambda{\bf
f}(x,\theta)-{\bf f}^{*}(x)\right\|^{2}$ by the sample average on a set of
$n_{\rm te}$ test samples $x_{i}\sim q$, that is,

$\widehat{\rm MSE}_{q}[{\bf f}(\cdot,\theta)]=\frac{1}{n_{\rm
te}}\sum_{i=1}^{n_{\rm te}}\left\|\lambda{\bf f}(x_{i},\theta)-{\bf
f}^{*}(x_{i})\right\|^{2}.$ (16)

Similarly, if additional testing or validation samples from data distribution
$p$ are available, we can compute an estimator of ${\rm
MSE}_{p}:=\mathbb{E}_{x\sim p}\left\|\lambda{\bf f}(x,\theta)-{\bf
f}^{*}(x)\right\|^{2}$ by the sample average denoted as $\widehat{\rm
MSE}_{p}$.

We note that the relationship between ${\rm MSE}_{p}$ and
$\mathcal{L}_{\lambda}[{\bf f}(\cdot,\theta)]$ as in (11) leads to an
estimator of ${\rm MSE}_{p}$ (up to a constant) that can be computed without
the knowledge of ${\bf f}^{*}$ as follows. Define ${\rm MSE}^{(m)}_{p}[{\bf
f}(\cdot,\theta)]:=2\lambda\mathcal{L}_{\lambda}[{\bf
f}(\cdot,\theta)]=2\lambda\mathbb{E}_{x\sim p}\left(-T_{q}{\bf
f}(x,\theta)+\frac{\lambda}{2}\|{\bf f}(x,\theta)\|^{2}\right)$, and (11)
gives that $2\lambda\mathcal{L}_{\lambda}[{\bf f}(\cdot,\theta)]={\rm
MSE}_{p}[{\bf f}(\cdot,\theta)]-\|{\bf f}^{*}\|_{p}^{2}$. Thus, we can use the
sample-average estimator of ${\rm MSE}^{(m)}_{p}$ to estimate ${\rm MSE}_{p}$
up to the unknown constant $\|{\bf f}^{*}\|_{p}^{2}$. That is, given $n_{\rm
val}$ validation samples $x_{i}\sim p$, the estimator can be computed as

$\widehat{\rm MSE}_{p}^{(m)}[{\bf f}(\cdot,\theta)]:=\frac{2\lambda}{n_{\rm
val}}\sum_{i=1}^{n_{\rm val}}\left(-T_{q}{\bf
f}(x_{i},\theta)+\frac{\lambda}{2}\|{\bf f}(x_{i},\theta)\|^{2}\right).$ (17)

The superscript (m) stands for “monitor”, and in experiments we will use
$\widehat{\rm MSE}_{p}^{(m)}$ to monitor the training progress of neural Stein
critic by a stand-alone validation set.

##  4 Theory of lazy training

In this section, we show the theoretical connection between choosing large
$L^{2}$-regularization weight $\lambda$ at the beginning of training and “lazy
learning” - referring to when the training dynamic resembles a kernel learning
determined by the model initialization [7]. The main result, Theorem 4.6 (for
population loss training) and Theorem 4.10 (for empirical loss training),
prove the kernel learning of neural Stein critic at large $\lambda$ and
theoretically support the benefit of using large $\lambda$ in the beginning
phase of training.

To recall the setup, we assume the score function of model distribution ${\bf
s}_{q}$ is accessible, and the score of data distribution ${\bf s}_{p}$ is
unknown. We aim to train a neural network critic to infer the unknown
scaleless optimal critic ${\bf f}^{*}$ from data. All the proofs are in
Section 8 and Appendix A.

###  4.1 Evolution of network critic under gradient descent

Throughout this section, we derive the continuous-time optimization dynamics
of gradient descent (GD). The continuous-time GD dynamic reveals the small
learning rate limit of the discrete-time GD, and the analysis may be extended
to minibatch-based Stochastic Gradient Descent (SGD) method [12] used in
practice. We start from training with the population loss, and the finite-
sample analysis will be given in Section 4.4.

Consider a neural Stein critic ${\bf f}(\cdot,\theta)$ parameterized by
$\theta$ which maps from $\mathcal{X}\subset\mathbb{R}^{d}$ to
$\mathbb{R}^{d}$. We assume $\theta\in\Theta$, which is some bounded set in
$\mathbb{R}^{M_{\Theta}}$, where $M_{\Theta}$ is the total number of trainable
parameters. The notation $\langle\cdot,\cdot\rangle_{\Theta}$ stands for the
inner-product in $\mathbb{R}^{M_{\Theta}}$, and $\|\cdot\|_{\Theta}$ the
Euclidean norm. The subscript Θ may be omitted when there is no confusion.

######  Assumption 3.

The network function ${\bf f}(x,\theta)$ is differentiable on
$\mathcal{X}\times\Theta$, and for any $\theta\in\Theta$, ${\bf
f}(\cdot,\theta)\in L^{2}(p)\cap\mathcal{F}_{0}(p)$.

The boundedness of $\Theta$ is valid in most applications, and our theory will
restrict $\theta$ inside a Euclidean ball in $\mathbb{R}^{M_{\Theta}}$, see
more in Assumption 4. Recall that for regularization weight $\lambda>0$, the
training objective $L_{\lambda}(\theta)$ is defined as in (13). Suppose the
neural network parameter $\theta$ evolving over training time is denoted as
$\theta(t)$ for $t>0$. The GD dynamic is defined by the ordinary differential
equation

$\dot{\theta}(t)=-{\partial_{\theta}}L_{\lambda}(\theta(t)),$ (18)

starting from some initial value of $\theta(0)$. The following lemma gives the
expression of (18).

######  Lemma 4.1.

For $\lambda>0$, the GD dynamic of $\theta(t)$ of minimizing
$L_{\lambda}(\theta)$ can be written as

$\dot{\theta}(t)=-\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta(t))\cdot\big{(}\lambda{\bf f}(x,\theta(t))-{\bf f}^{*}(x)\big{)}.$
(19)

Next, we derive the evolution of the network critic over time. We start by
defining

${\bf u}(x,t):={\bf f}(x,\theta(t)),$ (20)

and by chain rule, we have that

${\partial_{t}}{\bf u}(x,t)=\langle\partial_{\theta}{\bf
f}(x,\theta(t)),\dot{\theta}(t)\rangle_{\Theta}.$ (21)

Combining Lemma 4.1 and (21) leads to the evolution equation of ${\bf u}(x,t)$
to be derived in the next lemma. To proceed, we introduce the definition of
the finite-time (matrix) neural tangent kernel (NTK) ${\bf
K}_{t}(x,x^{\prime})$ as

$[{\bf
K}_{t}(x,x^{\prime})]_{ij}=\langle\partial_{\theta}f_{i}(x,\theta(t)),\partial_{\theta}f_{j}(x^{\prime},\theta(t))\rangle_{\Theta},\quad
i,j=1,\cdots,d,$ (22)

where $f_{i}$ denotes the $i$-th coordinate of ${\bf f}$. With the notation of
${\bf K}_{t}(x,x^{\prime})$ we have the following lemma.

######  Lemma 4.2.

The dynamic of ${\bf u}(x,t)$ follows that

${\partial_{t}}{\bf u}(x,t)=-\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{t}(x,x^{\prime})\circ\left(\lambda{\bf u}(x^{\prime},t)-{\bf
f}^{*}(x^{\prime})\right).$ (23)

The evolution equation (23) is exact for ${\bf u}(x,t)$, but the kernel ${\bf
K}_{t}(x,x^{\prime})$ changes with $t$ as the training progresses. Next in
Section 4.2, we will replace ${\bf K}_{t}(x,x^{\prime})$ with the zero-time
kernel throughout time in the evolution equation, resulting in the more
analyzable kernel learning solution.

###  4.2 Kernel learning with zero-time NTK

Theoretical studies of NTK to analyze neural network learning [22, 11, 3]
typically use two approximations: (i) zero-time approximation, namely ${\bf
K}_{t}(x,x^{\prime})\approx{\bf K}_{0}(x,x^{\prime})$ where both ${\bf K}_{t}$
and ${\bf K}_{0}$ are finite-width NTK, and (ii) infinite-width approximation,
namely at initialization, ${\bf K}_{0}(x,x^{\prime})\approx{\bf
K}_{0}^{(\infty)}(x,x^{\prime})$ as the widths of hidden layers increase,
where ${\bf K}_{0}^{(\infty)}$ is the limiting kernel at infinite width. As
has been pointed out in [7], the reduction to kernel learning in “lazy
training” does not necessarily require model over-parameterization \-
corresponding to large width (number of neurons) in a neural network - but can
be a consequence of a scaling of the network function. Here we show the same
phenomenon where the scaling factor is the $L^{2}$ regularization parameter
$\lambda$, see (14). That is, we show the approximation (i) only and prove
lazy training for finite-width neural networks.

We derive the property of the zero-time NTK kernel learning in this
subsection, and prove the validity of the approximation (i) in the next
subsection. To proceed, consider the kernel ${\bf K}_{t}(x,x^{\prime})$
defined in (22) at time zero, which can be written as

$\left[{\bf
K}_{0}(x,x^{\prime})\right]_{ij}=\langle\partial_{\theta}f_{i}(x,\theta(0)),\partial_{\theta}f_{j}(x^{\prime},\theta(0))\rangle_{\Theta},\quad
i,j=1,\cdots,d.$ (24)

The kernel ${\bf K}_{0}(x,x^{\prime})$ only depends on the initial network
weights $\theta(0)$, which is usually random and independent from the data
samples. We call ${\bf K}_{0}(x,x^{\prime})$ the zero-time finite-width
(matrix) NTK. Following the NTK analysis of neural network training, we will
show in Section 4.3 that the evolution dynamic of the network function ${\bf
u}(x,t)$ can be approximated by that of a kernel regression optimization \-
the so-called “lazy-training” dynamic - which is expressed by replacing the
finite-time NTK with the zero-time NTK. For the dynamic in (23), the lazy-
training dynamic counterpart is defined by the evolution of another solution
$\bar{\bf u}(x,t)$ by replacing the kernel ${\bf K}_{t}$ with ${\bf K}_{0}$ in
(23) starting from the same initial value, namely, $\bar{\bf u}(x,0)={\bf
u}(x,0)$ and

${\partial_{t}}\bar{\bf u}(x,t)=-\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{0}(x,x^{\prime})\circ\left(\lambda\bar{\bf u}(x^{\prime},t)-{\bf
f}^{*}(x^{\prime})\right).$ (25)

For simplicity, one assumes that at initialization, the network function is
zero mapping, that is, both ${\bf u}(x,0)$ and $\bar{\bf u}(x,0)$ are zero.
The argument generalizes to when the initial network function is small in
magnitude, cf. Remark 4.1, which can be obtained, e.g., by initializing neural
network weights with small values.

To analyze the dynamic of (25), we introduce the eigen-decomposition of the
kernel in the next lemma.

######  Lemma 4.3.

Suppose $\|\partial_{\theta}f_{i}(x,\theta(0))\|_{\Theta}$ for $i=1,\cdots,d$
are squared integrable on $(\mathcal{X},p(x)dx)$. The kernel ${\bf
K}_{0}(x,x^{\prime})$ on $(\mathcal{X},p(x)dx)$ has a finite collection of $M$
eigen-functions ${\bf v}_{k}:\mathcal{X}\to\mathbb{R}^{d}$, $k=1,2,\cdots$,
associated with positive eigenvalues, in the sense that

$\int_{\mathcal{X}}{\bf K}_{0}(x,x^{\prime})\circ{\bf
v}_{k}(x^{\prime})p(x^{\prime})dx^{\prime}=\mu_{k}{\bf v}_{k}(x),$ (26)

where $\mu_{1}\geq\cdots\geq\mu_{M}>0$. The eigen-functions are ortho-normal
in $L^{2}(p)$, namely, $\langle{\bf v}_{k},{\bf
v}_{l}\rangle_{p}=\delta_{kl}$, and for any ${\bf v}$ orthogonal to
${span}\\{{\bf v}_{1},\cdots,{\bf v}_{M}\\}$, $\int_{\mathcal{X}}{\bf
K}_{0}(x,x^{\prime})\circ{\bf v}(x^{\prime})p(x^{\prime})dx^{\prime}=0$.

The square integrability of $\partial_{\theta}f_{i}(x,\theta(0))$ can be
guaranteed by certain boundedness condition on $\partial_{\theta}f$, see more
in Assumption 4 below. The finite rank of the kernel, as shown in the proof,
is due to the fact that we use a neural network of finite width. Below, we
will assume that the optimal critic ${\bf f}^{*}$ can be efficiently expressed
by the span of finite many leading eigen-functions ${\bf v}_{k}$. To show that
the kernel spectrum is expressive enough to approximate an ${\bf f}^{*}$, one
may combine our analysis with NTK approximation (ii): the expressiveness of
the limiting kernel ${\bf K}_{0}^{(\infty)}$ at infinite width can be
theoretically characterized in certain settings, meaning that its eigen-
functions ${\bf v}^{(\infty)}_{k}$ collectively can span a rich functional
space. By the approximation ${\bf K}_{0}\approx{\bf K}_{0}^{(\infty)}$ in
spectrum, the expressiveness of the span of ${\bf v}_{k}$ can also be shown.
Such an extension is postponed here.

By Lemma 4.3, the eigen-functions $\\{{\bf v}_{1},\cdots,{\bf v}_{M}\\}$ form
an ortho-normal set with respect to the inner-product
$\langle\cdot,\cdot\rangle_{p}$. Thus for any integer $m\leq M$, the scaleless
optimal critic ${\bf f}^{*}\in L^{2}(p)$ (by Assumption 2) can be orthogonally
decomposed into two parts ${\bf f}_{1}^{*}$ and ${\bf f}_{2}^{*}$ such that

${\bf f}^{*}={\bf f}_{1}^{*}+{\bf f}_{2}^{*},\quad{\bf
f}_{1}^{*}\in\text{span}\\{{\bf v}_{1},\cdots,{\bf v}_{m}\\},\quad{\bf
f}_{2}^{*}\in\text{span}\\{{\bf v}_{1},\cdots,{\bf v}_{m}\\}^{\perp}.$

Making use of the orthogonal decomposition and by assuming that ${\bf f}^{*}$
has a significant projection on the eigen-space of the kernel ${\bf
K}_{0}(x,x^{\prime})$ associated with large eigenvalues, the following
proposition derives the optimization guarantee of $\bar{\bf u}(x,t)$.

######  Proposition 4.4.

Under the condition of Lemma 4.3 and notations as therein, suppose for
$\delta>0$ and some integer $m\leq M$,
$\mu_{1}\geq\cdots\geq\mu_{m}\geq\delta>0$. Let ${\bf f}^{*}={\bf
f}_{1}^{*}+{\bf f}_{2}^{*}$ be the orthogonal decomposition as above. Then,
for $\lambda>0$, starting from $\bar{\bf u}(x,0)=0$, for all $t>0$,

$\|\lambda\bar{\bf u}(\cdot,t)-{\bf f}^{*}\|_{p}\leq e^{-t\lambda\delta}\|{\bf
f}_{1}^{*}\|_{p}+\|{\bf f}_{2}^{*}\|_{p}.$ (27)

In particular, if for some $0<\epsilon<1$, $\|{\bf
f}_{2}^{*}\|_{p}\leq\epsilon\|{\bf f}^{*}\|_{p}$, then we have

$\|\lambda\bar{\bf u}(\cdot,t)-{\bf f}^{*}\|_{p}\leq 2\epsilon\|{\bf
f}^{*}\|_{p},\quad\text{when}\quad
t\geq\frac{1}{\lambda}\frac{\log(1/\epsilon)}{\delta}.$ (28)

###  4.3 Approximation by lazy-training dynamic

The network function ${\bf f}(x,\theta)$ maps from $\mathcal{X}\times\Theta$
to $\mathbb{R}^{d}$, and $\partial_{\theta}{\bf f}(x,\theta)$ is a
$d$-by-$M_{\Theta}$ matrix. We denote by $\|\cdot\|$ the vector 2-norm and the
matrix operator norm. We denote by $B_{r}$ the open Euclidean ball of radius
$r$ (centered at the origin) in $\mathbb{R}^{M_{\Theta}}$.

######  Assumption 4.

Under Assumption 3, there are positive constants $r$ and $L_{1},L_{2}$ such
that $B_{r}\subset\Theta$ and

(C1) For any $\theta\in B_{r}$, $\sup_{x\in\mathcal{X}}\|\partial_{\theta}{\bf
f}(x,\theta)\|\leq L_{1}$.

(C2) For any $\theta_{1},\theta_{2}\in B_{r}$,
$\sup_{x\in\mathcal{X}}\|\partial_{\theta}{\bf
f}(x,\theta_{1})-\partial_{\theta}{\bf f}(x,\theta_{2})\|\leq
L_{2}\|\theta_{1}-\theta_{2}\|$.

######  Proposition 4.5.

Under Assumptions 1-4, suppose $\theta(0)\in B_{r/2}$, and ${\bf
u}(x,0)=\bar{\bf u}(x,0)=0$, then for any $\lambda>0$,

(i) For any $t\geq 0$ such that $B_{r/2+\sqrt{{t}/{(2\lambda)}}\|{\bf
f}^{*}\|_{p}}\subset\Theta$,

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

(ii) For any $t\leq\frac{1}{2}(\frac{r}{\|{\bf f}^{*}\|_{p}})^{2}\lambda$
(recall that $\|{\bf f}^{*}\|_{p}>0$ by Assumption 2),

$\|\lambda{\bf u}(\cdot,t)-\lambda\bar{\bf
u}(\cdot,t)\|_{p}\leq\frac{4\sqrt{2}}{3}L_{1}L_{2}\sqrt{\lambda}t^{3/2}\|{\bf
f}^{*}\|_{p}^{2}.$ (30)

While the idea of proving Proposition 4.5 largely follows that of Theorem 2.2
in [7], we adopt a slightly improved analysis. Specifically, when we choose
$t\sim 1/\lambda$, both bounds in (29) and (30) reduce to $O(1/\lambda)$,
which echoes Theorem 2.2 in [7]. (Note that our normalization of objective
multiplies another factor of $\lambda$, and thus our time $t$ corresponds to
$T=t\lambda$ for time $T$ in [7].) Here we would like to derive the
approximation up to time $t\sim\log(1/\epsilon)/\lambda$, corresponding to
time $T\sim\log(1/\epsilon)$ in [7] instead of $O(1)$ time. Technically, our
analysis also improves the bounding constant in Theorem 2.2 of [7] by removing
a factor of $(e^{CT}-1)$, which will become a factor of $\epsilon^{-C}$ when
$T\sim\log(1/\epsilon)$. Thus our improvement is important to apply to the
case when $\epsilon$ is small. The improvement is by the special property of
mean-squared loss, which is equivalent to the neural Stein minimizing loss
$\mathcal{L}_{\lambda}$ up to a constant, cf. (11). We include a proof in
Section 8 for completeness.

We are ready to derive the main theorem of this section by combining
Propositions 4.4 and 4.5.

######  Theorem 4.6.

Under Assumptions 1-4, let the decomposition of ${\bf f}^{*}$ into ${\bf
f}_{1}^{*}$ and ${\bf f}_{2}^{*}$ be as in Proposition (4.4) and satisfy the
conditions therein and for some $0<\epsilon<1$, $\|{\bf
f}_{2}^{*}\|_{p}\leq\epsilon\|{\bf f}^{*}\|_{p}$. Suppose $\theta(0)\in
B_{r/2}$ and ${\bf u}(x,0)=0$, then when
$\lambda>(\frac{2\log(1/\epsilon)}{\delta})^{1/2}\frac{\|{\bf
f}^{*}\|_{p}}{r}$, for

$t=\frac{t_{0}}{\lambda}\frac{\log(1/\epsilon)}{\delta},\quad 1\leq
t_{0}\leq(\frac{r}{\|{\bf
f}^{*}\|_{p}})^{2}\frac{\delta}{2\log(1/\epsilon)}\lambda^{2},$ (31)

we have

$\|\lambda{\bf u}(\cdot,t)-{\bf
f}^{*}\|_{p}\leq\left(2\epsilon+\frac{C_{1}}{\lambda}\left(\frac{t_{0}\log(1/\epsilon)}{\delta}\right)^{3/2}\right)\|{\bf
f}^{*}\|_{p},$ (32)

where $C_{1}:=\frac{4\sqrt{2}}{3}L_{1}L_{2}\|{\bf f}^{*}\|_{p}$ is a constant
determined by ${\bf f}^{*}$ and Assumption 4(C1)(C2).

The needed upper bound of $t$, which $\sim\lambda$ is technical (to ensure
that $\theta(t)$ stays inside $B_{r}$). The theorem suggests that when the
scaleless optimal critic ${\bf f}^{*}$ can be represented using the leading
eigen-modes of the zero-time NTK up to an $O(\epsilon)$ residual, training the
neural Stein critic for $\sim\log(1/\epsilon)/\lambda$ time achieves an
approximation of ${\bf f}^{*}$ with a relative error of
$O(\epsilon,1/\lambda)$ up to a factor involving $\log(1/\epsilon)$. The
theoretical bound in Theorem 4.6 does not depend on data dimension $d$ or the
domain $\mathcal{X}$ explicitly, however, such dependence are indirect through
the constants $L_{1}$ and $L_{2}$ in (C1)(C2). The same applies to the finite-
sample analysis in Theorem 4.10 with respect to the constants in Assumptions 4
and 5.

######  Remark 4.1 (Small initialization).

The result extends to when the initial network function ${\bf u}(x,0)$ is non-
zero by satisfying $\|\lambda{\bf u}(x,0)\|_{p}\leq\epsilon\|{\bf
f}^{*}\|_{p}$. By considering the evolution of $\bar{\bf u}(x,t)$ starting
from $\bar{\bf u}(x,0)={\bf u}(x,0)$, one can extend Proposition 4.5 where the
bounds in (29) and (30) are multiplied by $O(1)$ constant factors (due to that
$\|\lambda{\bf u}(x,0)-{\bf f}^{*}\|_{p}\leq(1+\epsilon)\|{\bf
f}^{*}\|_{p}\leq 2\|{\bf f}^{*}\|_{p}$ when using the argument in (52)).
Proposition 4.4 also extends by considering the evolution equation (51) from
${\bf w}(x,0)=\lambda{\bf u}(x,t)-{\bf f}^{*}(x)$, and then $\lambda{\bf
u}(x,t)$ contributes to another $\epsilon\|{\bf f}^{*}\|_{p}$ in the bound
(28).

###  4.4 Training with finite samples

In this subsection, we extend the analysis to training using empirical
training loss $\hat{L}_{\lambda}(\theta)$ with $n$ training samples, which can
be written as

$\hat{L}_{\lambda}(\theta)=\mathbb{E}_{x\sim\hat{p}}\left(-T_{q}{\bf
f}(x,\theta)+\frac{\lambda}{2}\|{\bf f}(x,\theta)\|^{2}\right),$ (33)

where $\mathbb{E}_{x\sim\hat{p}}$ denotes the sample average over i.i.d.
samples $x_{i}\sim p$, i.e.,
$\mathbb{E}_{x\sim\hat{p}}g(x)=\frac{1}{n}\sum_{i=1}^{n}g(x_{i})$. Again,
using continuous-time evolution, the GD dynamic of $\hat{\theta}(t)$ is
defined by

$\dot{\hat{\theta}}(t)=-{\partial_{\theta}}\hat{L}_{\lambda}(\hat{\theta}(t)),\quad\hat{\theta}(0)=\theta(0),$
(34)

where $\theta(0)$ is some random initialization of the parameters. We define

$\hat{\bf u}(x,t):={\bf f}(x,\hat{\theta}(t)),$

assume zero-initialization $\hat{\bf u}(x,0)=0$, and similarly as in (21) we
have

${\partial_{t}}\hat{\bf u}(x,t)=\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\dot{\hat{\theta}}(t)\rangle_{\Theta}.$ (35)

As the counterpart of (22), we introduce the finite-time empirical NTK as

$[\hat{\bf
K}_{t}(x,x^{\prime})]_{ij}:=\langle\partial_{\theta}f_{i}(x,\hat{\theta}(t)),\partial_{\theta}f_{j}(x^{\prime},\hat{\theta}(t))\rangle_{\Theta},\quad
i,j=1,\cdots,d,$

which we also denote as

$\hat{\bf K}_{t}(x,x^{\prime})=\langle\partial_{\theta}{\bf
f}(x,\hat{\theta}(t)),\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))\rangle_{\Theta}=\partial_{\theta}{\bf
f}(x,\hat{\theta}(t))\partial_{\theta}{\bf
f}(x^{\prime},\hat{\theta}(t))^{T},\quad\partial_{\theta}{\bf
f}(x,\theta)\in\mathbb{R}^{d\times M_{\Theta}}.$ (36)

Since $\hat{\theta}(0)=\theta(0)$, we have that

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

which is a kernel matrix independent from training data, and this fact is
important for our analysis. Using the definition of $\hat{\bf
K}_{t}(x,x^{\prime})$, the dynamic (35) has the following equivalent form.

######  Lemma 4.7.

The dynamic of $\hat{\bf u}(x,t)$ follows that

${\partial_{t}}\hat{\bf
u}(x,t)=-\mathbb{E}_{x^{\prime}\sim\hat{p}}\left(\hat{\bf
K}_{t}(x,x^{\prime})\circ\left(\lambda\hat{\bf u}(x^{\prime},t)-{\bf
s}_{q}(x^{\prime})\right)-\nabla_{x^{\prime}}\cdot\hat{\bf
K}_{t}(x,x^{\prime})\right).$ (38)

Our analysis will compare the kernel $\hat{\bf K}_{t}(x,x^{\prime})$ to ${\bf
K}_{0}(x,x^{\prime})$, which allows to compare $\hat{\bf u}(x,t)$ to $\bar{\bf
u}(x,t)$ where we will also need to control the error by replacing
$\mathbb{E}_{x^{\prime}\sim{p}}$ with $\mathbb{E}_{x^{\prime}\sim\hat{p}}$.
The kernel comparison relies on showing that $\|\hat{\theta}(t)-\theta(0)\|$
is small, which we derive in the next lemma after introducing additional
technical assumptions on the score functions ${\bf s}_{p}$, ${\bf s}_{q}$, the
function ${\bf f}(x,\theta)$ and its derivatives. For a set
$B\in\mathbb{R}^{d}$, we denote by $\bar{B}$ the closure of the set.

######  Assumption 5.

Suppose ${\bf f}(x,\theta)$ is $C^{2}$ on $\mathcal{X}\times\Theta$. For the
$B_{r}$ as in Assumption 4,

(C3) There is $L_{3}>0$ such that, for any $\theta\in B_{r}$,
$\sup_{x\in\mathcal{X}}\|\nabla_{x}\cdot\partial_{\theta}{\bf
f}(x,\theta)\|\leq L_{3}$.

(C4) There is $L_{4}>0$ such that, for any $\theta_{1},\theta_{2}\in B_{r}$,
$\sup_{x\in\mathcal{X}}\|\nabla_{x}\cdot\partial_{\theta}{\bf
f}(x,\theta_{1})-\nabla_{x}\cdot\partial_{\theta}{\bf f}(x,\theta_{2})\|\leq
L_{4}\|\theta_{1}-\theta_{2}\|$.

(C5) There are positive constants $b_{(0)}$, $b_{(1)}$ and $b_{p}$ such that,
for any $\theta\in B_{r}$, $\sup_{x\in\mathcal{X}}\|{\bf f}(x,\theta)\|\leq
b_{(0)}$, $\sup_{x\in\mathcal{X}}|\nabla_{x}\cdot{\bf f}(x,\theta)|\leq
b_{(1)}$, and $\sup_{x\in\mathcal{X}}|{\bf f}(x,\theta)\cdot{\bf
s}_{p}(x)|\leq b_{(0)}b_{p}$. There are positive constants $C_{\mathcal{F}}$
and $b_{\mathcal{F}}$ and a positive $\gamma\leq 1/2$, such that, w.p. $\geq
1-2n^{-10}$,

$\max\\{\sup_{\theta\in\overline{B}_{r}}|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\nabla_{x}\cdot{\bf
f}(x,\theta)|,\,\sup_{\theta\in\overline{B}_{r}}|(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
f}(x,\theta)\cdot{\bf
s}_{p}(x)|\\}\leq\frac{C_{\mathcal{F}}}{n^{\gamma}}+b_{\mathcal{F}}\sqrt{\frac{\log
n}{n}}.$ (39)

(C6) There is a constant $b_{q}>0$ such that $\|{\bf s}_{q}\|_{p}\leq b_{q}$
and, for any $\theta\in B_{r}$, $\sup_{x\in\mathcal{X}}|{\bf
f}(x,\theta)\cdot{\bf s}_{q}(x)|\leq b_{(0)}b_{q}$, and
$\sup_{x\in\mathcal{X}}\|\partial_{\theta}{\bf f}(x,\theta)^{T}{\bf
s}_{q}(x)\|\leq L_{1}b_{q}$. The random variables $\|{\bf s}_{q}(x)\|^{2}$ and
$\|{\bf f}^{*}(x)\|^{2}$ with $x\sim p$ are sub-exponential. Specifically, the
constant $b_{q}$ and another constant $b_{2}>0$ satisfy that, when $n$ is
large (s.t. $\sqrt{{\log n}/{n}}$ is less than a constant possibly depending
on $b_{2}$ and $b_{q}$), w.p. $\geq 1-n^{-10}$,
$(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|{\bf
s}_{q}(x)\|^{2}\leq\sqrt{20}b_{q}^{2}\sqrt{{\log n}/{n}}$; w.p. $\geq
1-n^{-10}$, $(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|{\bf
f}^{*}(x)\|^{2}\leq\sqrt{20}b_{2}\sqrt{{\log n}/{n}}$.

In the below, we adopt big-O notation to facilitate exposition and $\tilde{O}$
stands for the involvement of a log factor. The constant dependence can be
tracked in the proof. We derive a non-asymptotic result which holds at a
sufficiently large finite sample size $n$.

######  Remark 4.2 (Uniform-law and rates).

The condition (C5) gives the standard uniform-law bounds which can be derived
by Rademacher complexity of the relevant function classes over
$\theta\in\overline{B}_{r}$, see e.g. [39], where we treat ${\bf s}_{p}$ as a
fixed bounded function on $\mathcal{X}$. The $C_{\mathcal{F}}n^{-\gamma}$ term
corresponds to the Rademacher complexity which is bounded by the covering
complexity of the function class, and the exponent $\gamma$ is determined by
the covering number bound that usually involves the dimensionality of the
domain and the regularity of the function class. The constant
$b_{\mathcal{F}}$ corresponds to the boundedness of the functions, namely
$b_{(1)}$ and $b_{(0)}b_{p}$. We note that while (C5) gives an overall
$\tilde{O}(n^{-{\gamma}})$ bound, it will only be used in the middle-step
analysis (Lemma 4.8) and our final finite-sample bound in Theorem 4.10
achieves the parametric rate of $\tilde{O}(n^{-1/2})$.

######  Remark 4.3 (Boundedness condition).

In (C5)(C6), the uniform boundedness of ${\bf f}(x,\theta)\cdot{\bf
s}_{p}(x)$, ${\bf f}(x,\theta)\cdot{\bf s}_{q}(x)$, and
$\|\partial_{\theta}{\bf f}(x,\theta)^{T}{\bf s}_{q}(x)\|$ on $\mathcal{X}$
can be fulfilled if ${\bf f}(x,\theta)$ and $\partial_{\theta}{\bf
f}(x,\theta)$ vanish sufficiently fast when approaching $\partial\mathcal{X}$,
and thus allowing the score functions ${\bf s}_{p}$ and ${\bf s}_{q}$ to be
potentially unbounded in $L^{\infty}(\mathcal{X})$. The score functions still
need to be in $L^{2}(p)$ and the tails cannot be too heavy to guarantee that
$\|{\bf s}_{q}(x)\|^{2}$ and $\|{\bf f}^{*}(x)\|^{2}$ are sub-exponential. The
Bernstein-type concentration bound of in (C6) follows from the standard
property of sub-exponential random variable, see e.g. [39, Chapter 2].

######  Lemma 4.8.

Under Assumptions 1-5, suppose $\theta(0)\in B_{r/2}$, and $\hat{\bf
u}(x,0)=0$. Then for any $\lambda>0$, there is positive integer $n_{\lambda}$
s.t. when $n>n_{\lambda}$, under a good event which happens w.p. $\geq
1-3n^{-10}$, for any $t\geq 0$ s.t. $\|{\bf
f}^{*}\|_{p}\sqrt{{t}/{\lambda}}\leq{r}/{2}$,
$\|\hat{\theta}(t)-{\theta}(0)\|\leq\sqrt{{t}/{\lambda}}\|{\bf f}^{*}\|_{p}$
and $\hat{\theta}(t)\in B_{r}$.

Based on the lemma, we derive the comparison of $\hat{\bf u}(x,t)$ with
$\bar{\bf u}(x,t)$ in the following proposition.

######  Proposition 4.9.

Under the same assumption as in Lemma 4.8, $\hat{\bf u}(x,0)=\bar{\bf
u}(x,0)=0$. For $\lambda>0$ and any $t\leq(\frac{r}{2\|{\bf
f}^{*}\|_{p}})^{2}\lambda$, when $n$ is sufficiently large (depending on
$\lambda,t,\|{\bf f}^{*}\|_{p},\log M_{\Theta}$), under the intersection of
the good event in Lemma 4.8 and another good event which happens w.p.$\geq
1-4n^{-10}$,

$\|\lambda\hat{\bf u}(\cdot,t)-\lambda\bar{\bf u}(\cdot,t)\|_{p}\leq
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)$ (40)

where $C_{2}$ is an $O(1)$ constant bounded by a multiple of $(1+b_{q}+\|{\bf
f}^{*}\|_{p})$ and the constant factor depends on $L_{1},L_{2},L_{3},L_{4}$.

######  Remark 4.4 (Largeness of $n$ and the relation to $\|{\bf
f}^{*}\|_{p}$).

The largeness requirement of $n$ depends on $\lambda,t,\|{\bf
f}^{*}\|_{p},\log M_{\Theta}$, and specifically,
$n>\max\\{n_{\lambda},n_{5},n_{\lambda,t}\\}$ which are defined in
(69)(73)(91) respectively: $n_{\lambda}$ depends on $\lambda$, $\|{\bf
f}^{*}\|_{p}$ and constants $b_{2}$, $b_{q}$, $C_{\mathcal{F}}$,
$b_{\mathcal{F}}$, and when $\|{\bf f}^{*}\|_{p}$ is small it calls for larger
$n$, specifically $\|{\bf f}^{*}\|_{p}^{2}>\lambda n^{-\gamma}$ as the leading
term, up to constant; $n_{5}$ depends on $\log M_{\Theta}$; $n_{\lambda,t}$
depends on $\lambda t$ and $\log M_{\Theta}$. The construction of the constant
$C_{2}$ can be found in the proof. The first term in the r.h.s. of (40) is the
analog of (30) in Proposition 4.5(ii). When $\|{\bf f}^{*}\|_{p}$ is small,
$C_{2}$ stays bounded (as long as $b_{q}$ is bounded) and the first term in
(40) is proportional to $\|{\bf f}^{*}\|_{p}$. The second term of the order
$\tilde{O}(n^{-1/2})$ is due to the finite-sample training, and when $\|{\bf
f}^{*}\|_{p}$ is small, it suggests that larger $n$ is needed so as to make
the second term balance with the first term.

The main theorem for finite-sample lazy training follows by combining
Propositions 4.4 and 4.9.

######  Theorem 4.10.

Under Assumptions 1-5, suppose the decomposition of ${\bf f}^{*}$ satisfies
the same condition as in Theorem 4.6 with respect to $0<\epsilon<1$,
$\theta(0)\in B_{r/2}$, and $\hat{\bf u}(x,0)=0$. For
$\lambda>2(\frac{\log(1/\epsilon)}{\delta})^{1/2}\frac{\|{\bf
f}^{*}\|_{p}}{r}$, let

$t=\frac{t_{0}}{\lambda}\frac{\log(1/\epsilon)}{\delta},\quad
1\leq{t_{0}}\leq(\frac{r}{\|{\bf
f}^{*}\|_{p}})^{2}\frac{\delta}{4\log(1/\epsilon)}\lambda^{2}.$ (41)

Then, when $n$ is sufficiently large (depending on $\lambda,\|{\bf
f}^{*}\|_{p},\log M_{\Theta}$ and $t_{0}\frac{\log(1/\epsilon)}{\delta}$),
w.p.$\geq 1-7n^{-10}$,

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

where $C_{2}$ is as in Proposition 4.9 and $\kappa_{1}$, $\kappa_{2}$ are
constant factors involving powers of $t_{0}\frac{\log(1/\epsilon)}{\delta}$.

The result in Theorems 4.6 and 4.10 suggests that by using a larger $\lambda$
at the beginning of the training process, the training can achieve the NTK
kernel learning solution more rapidly. In practice, staying with large
$\lambda$ too long would lead to the worsening of the model, and we propose
the gradual annealing scheme of $\lambda$ as described in Section 3.2 so as to
combine the benefits of both large $\lambda$ in the beginning and small
$\lambda$ in the later phases of training. The theoretical benefit of using
large $\lambda$ at the beginning phase of training is supported by experiments
in Section 6.

##  5 Applications to testing and model evaluation

###  5.1 Goodness-of-Fit (GoF) testing

In a GoF test, we are given $n_{\rm sample}$ data samples $x_{i}\sim p$ and a
model distribution $q$, and assume we can sample from $q$ as well as access
its score function ${\bf s}_{q}$. To apply the neural Stein test, we conduct a
training-test split of the samples, where the two splits have $n_{\rm tr}$ and
$n_{\rm GoF}$ samples respectively, and $n_{\rm sample}=n_{\rm tr}+n_{\rm
GoF}$. We first train a neural Stein critic ${\bf f}(x,\theta)$ from the
training split $\\{x_{i}^{\rm tr}\\}_{i=1}^{n_{\rm tr}}$, and then we compute
the following test statistic on the test split $\\{x_{i}\\}_{i=1}^{n_{\rm
GoF}}$

$\hat{T}=\frac{1}{n_{\rm GoF}}\sum_{i=1}^{n_{\rm GoF}}T_{q}{\bf
f}(x_{i},\theta),$ (43)

which can be viewed as a sample-average estimator of ${\rm SD}[{\bf
f}(\cdot,\theta)]$ as defined in (3).

To assess the null hypothesis as in (9), we will adopt a bootstrap strategy to
compute the test threshold $t_{\rm thresh}$ by drawing samples from $q$, to be
detailed in Section 5.1.1. We also derive GoF test consistency analysis in
Section 5.1.2.

####  5.1.1 Bootstrap strategy to compute test threshold

The bootstrap strategy draws independent samples $y_{i}\sim q$ to simulate the
distribution of the test statistic under $H_{0}$. We denote the test statistic
as $\hat{T}_{\rm null}$, which can be computed from a set of samples
$\\{y_{i}\\}_{i=1}^{n_{\rm GoF}}$ as $\hat{T}_{\rm null}=\frac{1}{n_{\rm
GoF}}\sum_{i=1}^{n_{\rm GoF}}T_{q}{\bf f}(y_{i},\theta)$. We will set $t_{\rm
thresh}$ as the $(1-\alpha)$ quantile of the distribution of $\hat{T}_{\rm
null}$. To simulate the distribution, one can compute $\hat{T}_{\rm null}$ in
$n_{\rm boot}$ independent replicas, and then set $t_{\rm thresh}$ to be the
quantile of the empirical distribution. This means that the $\hat{T}_{\rm
null}$ in each replica is computed from a “fresh” set of $n_{\rm GoF}$ samples
from $q$. We call these independent copies of $\hat{T}_{\rm null}$ the “fresh
null statistics”. Note that this would require evaluating the trained neural
network on $n_{\rm boot}n_{\rm GoF}$ many samples $y_{i}$ in total, which can
be significant since $n_{\rm boot}$ is usually a few hundred (we use $n_{\rm
boot}=500$ in all experiments). To accelerate computation, we propose an
“efficient” bootstrap procedure, which begins by drawing $n_{\rm pool}$
samples from $q$, $n_{\rm pool}=r_{\rm pool}\cdot n_{\rm GoF}$. The trained
network is evaluated on the $n_{\rm pool}$ samples of $y_{i}$ to obtain the
values of $T_{q}{\bf f}(y_{i},\theta)$, and then we compute $n_{\rm boot}$
many times of the $n_{\rm GoF}$-sample average by drawing from the pool with
replacement. We call the set of values of $\hat{T}_{\rm null}$ computed this
way the “efficient null statistics”. We observed that setting $r_{\rm
pool}=50$ is usually sufficient to render a null statistic distribution that
resembles that of the fresh statistics. With $n_{\rm boot}=500$, this yields
about ten times speedup in the computation of the bootstrap.

To illustrate the validity of the bootstrap strategy, we apply the method to a
critic trained on 1D Gaussian mixture data as described in Appendix B.2. In
this case, a partially trained critic, displayed in Figure 2(A), is used to
compute the test statistics. We set $n_{\rm GoF}=100$, and to better
illustrate the empirical distribution of the test statistics we use $n_{\rm
boot}=10,000$ replicas. Three sets of 10,000 test statistics are computed: 1)
using $n_{\rm GoF}$ fresh samples from $p$ to compute each $\hat{T}$, 2) using
$n_{\rm GoF}$ fresh samples from $q$ to compute $\hat{T}_{\rm null}$, and 3)
efficient bootstrap null statistics computed from a pre-generated pool with
$r_{\rm pool}=50$. The empirical distributions are visualized by histograms in
Figure 2(B). The plot shows a clear disparity between the distribution of test
statistics under $H_{1}$ and the two distributions of test statistics under
$H_{0}$, and the validity of the efficient bootstrap procedure is demonstrated
by the similarity of the histograms of the fresh and efficient null
statistics.

Figure 2:  In (A), a scaleless neural Stein critic function $\lambda{\bf
f}(\cdot,\theta)$ which is trained for a few epochs is compared to the
scaleless optimal critic ${\bf f}^{*}$. This trained neural critic may be
further improved by more training time, as shown in Figure 1. (B) The trained
critic, though not fully fitting the ${\bf f}^{*}$, still provides a GoF test
with significant power: the distribution of the $\hat{T}_{\rm null}$ under
$H_{0}$ is centered around zero, while the distribution of $\hat{T}$ under
$H_{1}$ has a non-zero mean around 5, which is significantly beyond the right
95% quantile of the $\hat{T}_{\rm null}$ distribution. The distribution of
$\hat{T}_{\rm null}$ (blue) and that approximately computed from the efficient
bootstrap (green) are close. See Section 5.1.1 for the details of the test
statistic computation. The details of the 1D Gaussian mixture distributions
are given in Appendix B.2.

####  5.1.2 Theoretical guarantee of test power

Let the trained neural Stein critic be $\hat{\bf f}(x):={\bf
f}(x,\hat{\theta}(t))$. We first observe the asymptotic normality of the test
statistics $\hat{T}$ and $\hat{T}_{\rm null}$: by definition, they are sample
averages over i.i.d. test samples $x_{i}\sim p$ ($y_{i}\sim q$) independent
from the training split. Thus, conditioning on training from $n_{\rm tr}$
samples, the test statistic is an independent sum averaged over i.i.d. random
variables

$\xi_{i}:=T_{q}\hat{\bf f}(x_{i})={\bf s}_{q}\cdot\hat{\bf
f}(x_{i})+\nabla\cdot\hat{\bf f}(x_{i}),\quad i=1,\cdots,n_{\rm GoF}.$

By Assumption 5(C5)(C6) (and that $\hat{\theta}(t)\in B_{r}$ by Lemma 4.8),
$\xi_{i}$ are uniformly bounded (either evaluated on $x_{i}$ or $y_{i}$ in
$\mathcal{X}$). Thus, by Central Limit Theorem, as $n_{\rm GoF}\to\infty$,

  * •

$\sqrt{n_{\rm GoF}}\hat{T}$ converges in distribution to
$\mathcal{N}(\mu_{1},\sigma_{1}^{2})$ where $\mu_{1}=\mathbb{E}_{x\sim
p}T_{q}\hat{\bf f}(x)$, and $\sigma_{1}^{2}={\rm Var}_{x\sim p}(T_{q}\hat{\bf
f}(x))$,

  * •

$\sqrt{n_{\rm GoF}}\hat{T}_{\rm null}$ converges in distribution to
$\mathcal{N}(\mu_{0},\sigma_{0}^{2})$ where $\mu_{0}=\mathbb{E}_{y\sim
q}T_{q}\hat{\bf f}(y)=0$, and $\sigma_{0}^{2}={\rm Var}_{y\sim
q}(T_{q}\hat{\bf f}(y))$.

The closeness to normal density is typically observed when $n_{\rm GoF}$ is a
few hundred, see Figure 2(B) where $n_{\rm GoF}=100$. The uniform boundedness
of $\xi_{i}$ implies that $\sigma_{1}^{2}$ and $\sigma_{0}^{2}$ are at most
$O(1)$ constants. This indicates that as long as the training obtains a critic
$\hat{\bf f}(x)$ that makes $\mu_{1}=\mathbb{E}_{x\sim p}T_{q}\hat{\bf
f}(x)>0$, then the GoF test can successfully reject the null when $n_{\rm
GoF}$ is large enough.

We also derive the finite-sample test power guarantee of the neural Stein GoF
test in the following corollary by incorporating the learning guarantee in
Section 4, and the proof is left to Section 8.

######  Corollary 5.1.

Consider a significance level $\alpha$ and target Type-II error $\beta$,
$0<\alpha,\beta<1$. Suppose the conditions in Theorem 4.10 are satisfied with
$\epsilon<1/2$, $\lambda>0$ as required and
$2\epsilon+C_{2}\kappa_{1}\lambda^{-1}<0.8$, and $n_{\rm tr}$ is large enough
for the good event (call it $E_{\rm tr}$) to hold,
$C_{2}\kappa_{2}\sqrt{{(\log n_{\rm tr}+\log M_{\Theta})}/{n_{\rm
tr}}}<0.1\|{\bf f}^{*}\|_{p}$, and $7n_{\rm
tr}^{-10}<0.1\min\\{\alpha,\beta\\}$. Then the GoF test using $\hat{\bf
f}(x)={\bf f}(x,\hat{\theta}(t))$ achieves a significance level $\alpha$ and a
test power at least $1-\beta$, if

$\sqrt{2}(b_{(0)}b_{q}+b_{(1)})\left(1+\sqrt{\log\frac{1}{\alpha}}+\sqrt{\log\frac{1}{\beta}}\right)\lambda
n_{\rm GoF}^{-1/2}<0.1\|{\bf f}^{*}\|_{p}^{2}.$

In particular, the test power $\to 1$ as $n_{\rm tr},n_{\rm GoF}\to\infty$.

###  5.2 Evaluation of EBM generative models

We consider the application of the Stein discrepancy in evaluating generative
models. The model evaluation problem is to detect how the model density $q$
(by a given generative model) is different from the unknown data density $p$.
While this can be formulated as a GoF testing, in machine learning
applications it is of interest where and how much the two densities differ
rather than rejecting or accepting the null hypothesis. To this end, we note
that the trained Stein critic can be used as an indicator to reveal where $p$
and $q$ locally differ.

In the case of EBMs, the model probability density takes the form as
$q(x)={\exp(-E_{\phi}(x))}/{Z}$, where $Z$ is the normalizing constant and the
real-valued energy function $E_{\phi}(x)$ is parameterized by $\phi$, which is
another model such as a neural network (so the set of parameters $\phi$
differs from the neural Stein critic parametrization $\theta$). The score
function of $q$, therefore, equals the gradient of the energy function, i.e.,

${\bf s}_{q}=-\nabla E_{\phi}(x),$ (44)

which can be computed from the parameterized form of $E_{\phi}(x)$. In the
case that the energy function is represented by a deep generative neural
network, the gradient (44) can be computed by back-propagation, which is
compatible with the auto-differentiation implementation of widely-used deep
network platforms. In practice, the evaluation metric of a trained EBM can be
computed on a holdout validation dataset.

Below we give the expression of the score function for Gaussian-Bernoulli
Restricted Boltzmann Machines (RBMs) [16], which is a specific type of EBM.
The energy of an RBM with latent Bernoulli variable $h$ is defined as
$E(x,h|B,b,c)=-\frac{1}{2}x^{\text{T}}Bh-b^{\text{T}}x-c^{\text{T}}h+\frac{1}{2}\|x\|^{2}$.
Therefore, the score function has the expression ${\bf
s}_{q}(x)=b-x+B\cdot{\rm tanh}(B^{\text{T}}x+c)$. In Section 6.3, we evaluate
Gaussian-Bernoulli RBMs using a Stein discrepancy test computed via neural
Stein critic functions, and we also compare neural Stein critics trained using
different regularization strategies.

##  6 Experiment

In this section, we present numerical experiments applying the proposed neural
Stein method on differentiating a data distribution $p$ (from which we have
access to a set of data samples) and a model distribution $q$ (of which the
score function is assumed to be known). We compare the proposed neural Stein
method with staged regularization to that with fixed $L^{2}$ regularization,
as well as to a kernel Stein method previously developed in the literature.

In Sections 6.1 and 6.2, we consider a set of Gaussian mixture models for both
$p$ and $q$. In Section 6.3, the data are sampled from the MNIST handwritten
digits dataset [25], and the model distribution is a Gaussian-Bernoulli RBM
neural network model. Codes to reproduce the results in this section can be
found at the following repository:
https://github.com/mrepasky3/Staged_L2_Neural_Stein_Critics.

###  6.1 Gaussian mixture data

####  6.1.1 Simulated datasets

We compare the performance of a variety of neural Stein critics in the
scenario in which the distributions are bimodal, $d-$dimensional Gaussian
mixtures, following an example studied in [26]. The model $q$ has equally-
weighted components with means $\mu_{1}=\mathbf{0}_{d}$ and
$\mu_{2}=0.5\times\mathbf{1}_{d}$, both having identity covariance. The data
distribution $p$ has the following form:

$p=\frac{1}{2}\mathcal{N}\left(\mu_{1},\begin{bmatrix}1&\rho_{1}&\mathbf{0}_{d-2}^{\text{T}}\\\
\rho_{1}&1&\mathbf{0}_{d-2}^{\text{T}}\\\
\mathbf{0}_{d-2}&\mathbf{0}_{d-2}&\mathbf{I}_{d-2}\end{bmatrix}\right)+\frac{1}{2}\mathcal{N}\left(\mu_{2},\begin{bmatrix}\omega^{2}&\omega\rho_{2}&\mathbf{0}_{d-2}^{\text{T}}\\\
\omega\rho_{2}&1&\mathbf{0}_{d-2}^{\text{T}}\\\
\mathbf{0}_{d-2}&\mathbf{0}_{d-2}&\mathbf{I}_{d-2}\end{bmatrix}\right),$ (45)

where $\rho_{1}=-\rho_{2}$ represents the covariance shift with respect to the
model distribution. We also introduce a parameter $\omega$, which scales the
covariance matrix of the second component of the mixture. We examine this
scenario in three settings of increasingly high dimensions. In each setting of
this section, we fix the parameters $\rho_{1}=0.5$ and $\omega=0.8$.

####  6.1.2 Experimental setup

##### Neural network training.

All neural Stein critics trained are two-hidden-layer MLPs with Swish
activation [30], where each hidden layer includes 512 hidden units. The
weights of the linear layers of the models are initialized using the standard
PyTorch weight initialization, and the biases of the linear layers are
initialized as zero. The Adam optimizer (using the default momentum parameters
by PyTorch) was used for network optimization. The critics are trained using
$n_{\rm tr}=$2,000 samples from the data distribution $p$; the learning rate
is fixed at $10^{-3}$, and the batch size is 200 samples. All models are
trained for 60 epochs. To compute the divergence $\nabla\cdot{\bf
f}(x,\theta)$ in dimension greater than two, we use Hutchinson’s unbiased
estimator of the trace Jacobian [20] following [16]. In these experiments, the
staging of $\lambda$ according to the staged regularization strategies occurs
every $B_{w}=10$ batches, which is equivalent to the end of every epoch. For
each choice of regularization strategy, we train 10 neural Stein critic
network replicas. We compare the neural Stein critics trained using
fixed-$\lambda$ regularization vs. staged regularization schemes.

##### Computation of MSE.

Having knowledge of the score functions of both $p$ and $q$, we are able to
compute the $\widehat{\rm MSE}_{q}$ (16) using $n_{\rm te}=20,000$ samples
from the model $q$. Over the course of training, we also computed
$\widehat{\rm MSE}^{(m)}_{p}$ (17) using $n_{\rm val}=1,000$ samples from the
data distribution $p$. This monitor is used to select the “best” model over
the course of training, where the model with the lowest $\widehat{\rm
MSE}^{(m)}_{p}$ value is selected.

##### Computation of test power.

After a trained neural critic is obtained, we perform the GoF hypothesis
testing as described in Section 5.1, including computing the test threshold
$t_{\rm thresh}$ by the efficient bootstrap strategy. We set the significance
level $\alpha=0.05$, the number of bootstrap $n_{\rm boot}=500$, and the
efficient bootstrap ratio number $r_{\rm pool}=50$. We use $n_{\rm GoF}$
specific for each data distribution example, see below.

By definition, the test power is the probability at which the null hypothesis
$H_{0}:p=q$ is correctly rejected. For each trained neural critic from a given
training split of $p$ samples, we estimate the test power of this critic
empirically by conducting $n_{\rm run}$ times of the GoF tests (including
independent realizations of the test split of $p$ samples and computing the
$t_{\rm thresh}$ by bootstrap) and counting the frequency when the null
hypothesis is rejected. Because the training also contributes to the
randomness of the quality of the GoF test, we repeat the procedure for $n_{\rm
replica}$ replicas (including training the model on an independent realization
of the training split in each replica), and report the mean and standard
deviation of the estimated test power from the replicas. We use $n_{\rm
run}=500$ and $n_{\rm replica}=10$ in our experiments.

##### Setup per example.

We introduce the specifics of the experimental setup for Gaussian mixtures in
2D, 10D, and 25D below.

  * (a)

2-dimensional mixture. The selected fixed-$\lambda$ regularization strategies
in 2 dimensions are $\lambda\in\\{1\times 10^{-3},1\times 10^{-2},1\times
10^{-1},1\times 10^{0}\\}$. We compare these fixed schemes to a few staged
regularization strategies. Using the notation of Equation (15), these are the
$\Lambda(1\times 10^{0},5\times 10^{-2},0.95)$ and $\Lambda(1\times
10^{0},5\times 10^{-2},0.90)$ staged regularization schemes. In 2D, the number
of test samples for the GoF testing is $n_{\rm GoF}=75$.

  * (b)

10-dimensional mixture. For this setting, we examine fixed
$\lambda\in\\{2.5\times 10^{-4},1\times 10^{-3},4\times 10^{-3},1.6\times
10^{-2},6.4\times 10^{-2},2.56\times 10^{-1},1.024\times 10^{0}\\}$. We
analyze the staging schemes $\Lambda(5\times 10^{-1},1\times 10^{-3},0.80)$
and $\Lambda(5\times 10^{-1},1\times 10^{-3},0.85)$. For the GoF power
analysis, we use $n_{\rm GoF}=200$ test samples from the data distribution
$p$.

  * (c)

25-dimensional mixture. The fixed regularization weights are selected as
$\lambda\in\\{2.5\times 10^{-4},1\times 10^{-3},4\times 10^{-3},1.6\times
10^{-2},6.4\times 10^{-2}\\}$. The staging schemes analyzed in this case are
$\Lambda(4\times 10^{-1},5\times 10^{-4},0.80)$, $\Lambda(4\times
10^{-1},5\times 10^{-4},0.85)$, and $\Lambda(4\times 10^{-1},5\times
10^{-4},0.90)$. We use $n_{\rm GoF}=500$ test samples from the data
distribution $p$ for the GoF tests in 25D.

Figure 3:  $\widehat{\rm MSE}_{q}$ results for the Gaussian mixture model
settings in 2 dimensions (A), in 10 dimensions (B), and in 25 dimensions (C).
In each case, 10 models are trained for 60 epochs using each regularization
strategy; the mean and standard deviation are visualized over these 10 models.
The legends indicate the value of $\lambda$ corresponding to each curve. We
use $\Lambda(1\times 10^{0},5\times 10^{-2},\cdot)$ staging in 2D,
$\Lambda(5\times 10^{-1},1\times 10^{-3},\cdot)$ in 10D, and $\Lambda(4\times
10^{-1},5\times 10^{-4},\cdot)$ in 25D. The advantage of staged regularization
becomes more pronounced with increasing dimension. Figure 4:  A visualization
of the scaleless neural Stein critic $\lambda{\bf f}(\cdot,\theta)$ throughout
training, where we use the $\Lambda(1\times 10^{0},5\times 10^{-2},0.90)$
staging strategy on the 2-dimensional Gaussian mixture data. In (A), the
critic quickly approximates some regions in the first epoch, whereas in (B)
and (C) the critic makes smaller adjustments throughout the domain, creating
an approximation of the scaleless optimal critic function ${\bf f}^{*}$ (7)
pictured in (D). Figure 5:  Testing power results for the Gaussian mixture
model settings in 2 dimensions (A), in 10 dimensions (B), and in 25 dimensions
(C). The mean power for fixed-$\lambda$ settings is plotted with standard
deviation error bars, and the highest-average-power staged setting is plotted
as a (mean) line with shaded standard deviation in each setting, where the
dashed line spans the range of $\lambda$ for each staging strategy.

####  6.1.3 Results

2-dimensional mixture. The value of $\widehat{\rm MSE}_{q}$ computed via
Equation (16) throughout training for all 2D critics can be seen in Figure
3(A), where the mean and standard deviation are plotted across the 10 networks
for each staging strategy. We begin by examining the training behavior for
fixed regularization weight $\lambda$. Rapid descent in $\widehat{\rm
MSE}_{q}$ is observed for large $\lambda$ early in training while small
$\lambda$ make slower and steadier progress. The staged regularization
strategies exploit this rapid descent at early times, followed by steady late-
stage training. However, note that the advantage of staging in 2 dimensions
only yields a marginal benefit in MSE. In Figure 4, we visualize the critic
vector field plotted through training for the $\Lambda(1\times 10^{0},5\times
10^{-2},0.90)$ staging scheme. Rapid fitting of some regions can be seen at
the beginning of training in Figure 4(A), followed by more nuanced adjustments
in Figures 4(B) and 4(C), resulting in a good fit to the theoretically optimal
critic function in Figure 4(D). Furthermore, consider the test power results
displayed in Figure 5(A). While the staged regularization cases have among the
lowest $\widehat{\rm MSE}_{q}$ values, similar testing power is achieved by
the fixed-$\lambda$ settings. For more detail regarding the 2-dimensional
experiment, see Table A.1 in Appendix B.1, which shows the average GoF
hypothesis testing power at the approximate “best” training epoch as chosen by
finding the average lowest monitor $\widehat{\rm MSE}_{p}^{(m)}$ calculated
using Equation (17). The table also displays the average $\widehat{\rm
MSE}_{q}$ as calculated by Equation (16) at the “best” epoch for each
regularization scheme.

10-dimensional mixture. As in 2D, the $\widehat{\rm MSE}_{q}$ is plotted for
critics over the course of training in Figure 3(B). The curves corresponding
to some fixed $\lambda$ strategies, such as $\lambda=6.4\times 10^{-2}$,
$2.56\times 10^{-1}$, and $1.024\times 10^{0}$ are omitted since they quickly
diverge in $\widehat{\rm MSE}_{q}$. In higher dimensions, the difference in
using a wider range of regularization weights $\lambda$ is more pronounced.
Larger values result in networks that rapidly fit a (relatively) poor
representation of the optimal critic ${\bf f}^{*}_{\lambda}$, followed by
dramatic overfitting. Smaller choices of fixed $\lambda$ delay both phenomena,
fitting to a critic which has a lower value of $\widehat{\rm MSE}_{q}$ at
best. The staged regularization strategies exploit both types of training
dynamic, descending in $\widehat{\rm MSE}_{q}$ more rapidly than most
fixed-$\lambda$ strategies while achieving lower value (and hence better fit)
than any fixed strategy, on average. Examining the results in Figure 5(B), we
find that the power of the $\Lambda(5\times 10^{-1},1\times 10^{-3},0.80)$
regularization strategy exceeds (on average) that of all the fixed-$\lambda$
strategies. More detail related to the networks obtained via validation can be
found in Table A.2 in Appendix B.1.

25-dimensional mixture. As in the previous settings, the $\widehat{\rm
MSE}_{q}$ computed via Equation (16) is plotted in Figure 3(C) for the 25D
regularization strategies. As is the case in 10D, the $\widehat{\rm MSE}_{q}$
of some regularization strategies are relatively high and are therefore
omitted from Figure 3. The observed trend of increasing the dimension from 2
to 10 is further exemplified by the increase to 25 dimensions. The larger
choices of fixed $\lambda$ result in networks that quickly obtain a poor fit
of the optimal critic, followed by overfitting. The smaller choices of fixed
$\lambda$ yield more stable $\widehat{\rm MSE}_{q}$ curves. Combining these
dynamics in our staging strategies, the $\widehat{\rm MSE}_{q}$ performance
gap between fixed- and staged-$\lambda$ regularization dramatically increases,
where the staged strategies substantially outperform the fixed strategies.
Figure 5(C) further corroborates this finding. The GoF hypothesis test power
of the staged critics is dramatically higher than any fixed-$\lambda$ training
strategy. See Table A.3 in Appendix B.1 for more detail pertaining to trained
networks in 25D.

On the simulated Gaussian mixture data, we find that staging the
regularization of the neural Stein critic throughout training yields greater
benefit as the dimension increases from 2 to 25. The scaleless neural Stein
critic $\lambda{\bf f}(\cdot,\theta)$ rapidly fits the scaleless optimal
critic ${\bf f}^{*}$ (7) at early times when $\lambda$ is large, followed by
stable convergence to a low-$\widehat{\rm MSE}_{q}$ critic throughout
training. Furthermore, we find that the GoF hypothesis testing power follows a
similar trend as that of the $\widehat{\rm MSE}_{q}$, such that staging
$\lambda$ yields an increase in power, especially in higher-dimension.

###  6.2 Comparison to kernel Stein discrepancy

Using the simulated Gaussian mixture data as in Section 6.1, we compare the
testing power of the neural Stein critic GoF hypothesis test to that of KSD.
To do so, we perform the GoF hypothesis test by computing the KSD test
statistic outlined in [27, 8], whereby the Stein discrepancy is computed using
a critic restricted to an RKHS. Following common practice in the literature,
we construct this RKHS to be defined by a radial basis function (RBF) with a
bandwidth equal to the median of the data Euclidean distances in a given GoF
test. We also compare our method to a KSD test with RBF bandwidth which is
selected to maximize the power. To compute the KSD test, we adopt the
implementation of [8]. Further details of the KSD GoF test, including
selecting the most effective bandwidth, are given in Appendix B.3, with GoF
testing power for a range of bandwidths displayed in Figure A.1.

It would be natural to ask how the neural network test compares with other
traditional parametric and non-parametric tests. There are cases where the
generalized likelihood ratio (GLR) test can be computed from parametric models
and would be the optimal test. When a parametric model is not available,
traditional non-parametric tests like Kolmogorov-Smirnov may be restricted to
low-dimensional data. We thus focus on comparison with the KSD test which is a
kernel-based non-parametric test generally applicable to high dimensional
data.

####  6.2.1 Experimental setup

Using the model distribution $q$ and data distribution $p$ as defined in
Section 6.1.1, we construct a two-component Gaussian mixture in 50D to compare
the power and computation time of the neural Stein discrepancy GoF hypothesis
test (Section 5.1) to the KSD test (Appendix B.3). The model distribution $q$
remains as the isotropic, two-component Gaussian mixture, while the data
distribution $p$ has the form of (45) with covariance shift $\rho_{1}=0.5$ and
covariance scaling $\omega=0.5$.

##### Computation of test power.

For all neural Stein critic GoF tests, a total of $n_{\rm tr}+n_{\rm
GoF}=n_{\rm sample}$ samples from $p$ are used, where $n_{\rm sample}$ takes
on values in $\\{100,200,300,500,1000,1500,2500\\}$. In each case, $n_{\rm
tr}$ and $n_{\rm GoF}$ are equal to $n_{\rm sample}/2$. The 50%/50%
training/testing split was chosen by comparing the testing power for various
training/testing splits. Our findings indicate a range of training/testing
splits exists for which the testing power performance is comparable. While the
training duration contributes most significantly to the computation time, the
test achieves high power once the training split reaches 50% of the sample
size. That is, the 50%/50% train/test split is a generic split choice in this
range of high power splits, and therefore we use this split in the results
outlined in Section 6.2.2 and in Figure 6. The details of this finding can be
found in Appendix B.4, Figure A.2, and an application to a simpler setting is
displayed in Figure A.3. Furthermore, the training data are partitioned into
an 80%/20% train/validation split. The testing procedure and test power
computation are as described in Section 6.1.2. Specifically, the significance
level $\alpha=0.05$ in all tests, the number of bootstraps $n_{\rm boot}=500$,
and the efficient bootstrap ratio is $r_{\rm pool}=50$.

The KSD tests [8] are computed over the same range of samples size. The KSD
test is conducted with an RBF kernel using the median data distance heuristic
for bandwidth, and we also examine the selected bandwidth chosen to maximize
the test power. To compute $t_{\rm thresh}$, the KSD test uses a “wild
bootstrap” procedure [8]. We provide details of bandwidth selection and wild
bootstrap in Appendix B.3. The number of wild bootstrap samples used for KSD
is equal to the number of bootstrap samples used for the neural Stein test,
namely $n_{\rm boot}=500$.

The test power is computed using $n_{\rm run}=400$ and $n_{\rm replica}=5$ for
each KSD and neural Stein GoF test.

##### Comparison of test time.

We compare the time taken to perform one KSD GoF test vs. training a neural
Stein critic and performing a neural GoF test. All reported times are elapsed
time on a laptop with an Intel Core i7-1165G7 processor with 16 GB of RAM. We
record and report the computation times of the tests over a range of sample
sizes.

For the neural test, we measure the duration time of training followed by the
computation time of a single test statistic and the computation time of the
bootstrap, averaged over $n_{\rm replica}=5$ replicas with $n_{\rm run}=10$
and $n_{\rm boot}=500$.

To determine the computation time of the KSD test, we measure the duration of
10 tests using the best-selected bandwidth, which can be broken down into the
time taken to compute an individual test statistic and the computation time of
the wild bootstrap, averaged over $n_{\rm replica}=5$, $n_{\rm run}=10$, and
$n_{\rm boot}=500$.

##### Neural network training.

The neural Stein critics are learned using two-hidden-layer (512 nodes) MLPs
with Swish activation, initialized using PyTorch standard initialization with
biases set to 0. We use the Adam optimizer with a learning rate set to
$5\times 10^{-3}$, minibatch size equal to $40$, and the critics are trained
for 60 epochs. The $\Lambda(4\times 10^{-1},5\times 10^{-4},0.85)$ staging is
used in the training of the neural Stein critics, where the frequency of
$\lambda$ staging, $B_{w}$ batches, is again chosen to be equivalent to the
number of training batches per epoch. The networks are trained using $n_{\rm
tr}$ number of training samples from the data distribution $p$, which are
split into an 80%/20% train/validation split for model selection. As in
Section 6.1.2, the networks are selected to minimize the monitor $\widehat{\rm
MSE}^{(m)}_{p}$ of Equation (17).

Figure 6:  Comparison between neural Stein critic GoF test and KSD test [8] in
the 50D setting described in 6.2.1. We perform $n_{\rm run}=400$ number of GoF
tests to compute power with $n_{\rm replica}=5$ replicas. All results use
$n_{\rm boot}=500$ bootstrap samples. The KSD is computed using an RBF kernel
with median data distance heuristic bandwidth and bandwidth selected to
maximize power (see Appendix B.3 for details). In each plot, the x-axis
indicates the overall sample size $n_{\rm sample}$ provided to each method,
where the neural Stein GoF test uses half for training and a half for
computation of the test statistic (refer to Figure A.2 for a comparison of
neural Stein tests with other choices of train/test split), and KSD uses the
entire sample for computation of the test statistic. In (A), the average test
power and its standard deviation error bar are shown, computed over the 5
models in each case. Figure (B) shows the average duration of network training
plus an individual GoF test (computation of test statistic and bootstrap) for
the neural Stein test and the average duration of a single test (test
statistic and bootstrap) for the KSD test. Finally, (C) shows the breakdown
between the average computation time required to compute a test statistic
(including training for the neural network test) and the time required to
compute the bootstrap. See Section 6.2.1 for details of the GoF test power and
recording of computation time.

####  6.2.2  Results

The result of the testing power comparison on the 50D data can be found in
Figure 6. First, examining the comparison in power between the methods in
Figure 6(A), the neural Stein critic GoF test (solid black line) remains more
powerful than the KSD GoF tests for all sample sizes. The KSD GoF test using
the median data distances heuristic bandwidth (solid red line) does not
achieve comparable power to the staged-regularization neural Stein critic
approach, even for as many as 2500 samples. While the KSD test using the best-
selected bandwidth (dashed red line) achieves much higher power than the
heuristic approach, the power of this method still falls below that of the
neural Stein GoF test. Next, Figure 6(B) highlights the quadratic time
complexity of the KSD GoF test, resulting in a dramatic increase in
computation time that surpasses the overall time to train a neural Stein
critic and compute a test for samples sizes larger than 1,000. Furthermore,
the breakdown into test statistic computation time and bootstrap computation
time in Figure 6(C) reveals that the time to compute an individual test
statistic for the neural Stein approach becomes more time-efficient than KSD
when $n_{\rm sample}$ is just larger than 1500. While the test statistic
computation and wild bootstrap of KSD contribute substantially to the overall
computation time of KSD, the neural Stein GoF test computation time is
dominated by the training period.

Our findings indicate that the neural Stein critic clearly outperforms KSD in
this setting for a larger sample size, with higher power and lower computation
time. This result highlights the deeper expressivity of the $L^{2}$ function
space compared to kernel methods in addition to the benefit of the linear time
complexity of the neural Stein critic GoF test, as opposed to the quadratic
complexity of KSD.

###  6.3 MNIST handwritten digits data

The results of Section 6.1 indicate that the staging of regularization when
training neural Stein critics yields greater benefit in higher dimensions.
Therefore, we extend to a real-data example in an even higher dimension: the
MNIST handwritten digits dataset [25]. We compare the fixed and staged
regularization strategies to train critics that discriminate between an RBM
and a mixture model of MNIST digits.

####  6.3.1 Authentic and synthetic MNIST data

To construct the model distribution $q$ for the MNIST setting, we follow the
approach of [16]. We use a Gaussian-Bernoulli RBM that models the MNIST data
distribution, trained using a learned neural Stein critic to minimize the
Stein discrepancy between the true MNIST density and the RBM. We declare the
model distribution $q$ to be this 728-dimensional RBM. The data distribution
$p$ is a mixture model composed of 97% the RBM and 3% true digits “1” from the
MNIST dataset. Therefore, any disparity in the distributions of $p$ and $q$
are caused by this infusion of digits from MNIST into $p$.

####  6.3.2 Learned neural Stein critics

In addition to being a realistic setting, training a neural Stein critic using
MNIST digits will allow us to better interpret the discrepancy. Of course,
since we do not have access to the “true” score function $p$, we cannot
calculate the $\widehat{\rm MSE}_{q}$ for the trained critics using Equation
(16). Furthermore, the computation of $\widehat{\rm MSE}^{(m)}_{p}$ using
validation data from $p$ via (17) becomes less accurate in high dimension, as
the method would require a large amount of validation data to be an accurate
representation of the population ${\rm MSE}_{p}$. Therefore, we introduce an
additional validation metric to evaluate the fit of the neural Stein critic
${\bf f}(\cdot,\theta)$. First, in the language of the GoF test introduced in
Section 5.1, we denote the quantity computed by applying the Stein operator
with respect to $q$ on neural Stein critic ${\bf f}(\cdot,\theta)$ evaluated
at a sample $x\in\mathcal{X}$ as the “critic witness” of the sample $x$:

$w(x{,\theta})=T_{q}{\bf f}(x{,\theta}).$ (46)

Intuitively, since ${\bf f}(\cdot,\theta)$ is trained on samples from $p$ (the
data distribution) by maximizing the Stein discrepancy, the value of
$w(x{,\theta})$ represents the magnitude of the difference between
distributions $p$ and $q$ at $x\in\mathcal{X}$. Evaluating the critic witness
at ${n_{\rm GoF}}$ samples $x_{i}\sim q$, under the central limit theorem
(CLT) assumption, random variables $w(x_{i}{,\theta})$ have a (centered)
normal distribution with standard deviation $\sigma(w)/\sqrt{{n_{\rm GoF}}}$
when ${n_{\rm GoF}}$ is large, where $\sigma(w)$ is the standard deviation of
$w(x_{i}{,\theta})$.

Note that the test statistic (43) is the mean of $w(x_{i}{,\theta})$ computed
over testing data $x_{i}\sim p$. As an assessment for the GoF testing power
for the neural Stein critic function ${\bf f}{(\cdot,\theta)}$ (which is
expensive to compute in such a high dimension), we may compare the mean and
variance of $w(\cdot{,\theta})$ applied to a testing dataset sampled from $p$
and to a “null” dataset drawn from $q$, both of size ${n_{\rm GoF}}$:

$\hat{P}=\frac{{\bar{w}}_{p}}{\sigma(w_{p})+\sigma(w_{q})},\quad{\bar{w}}_{p}=\frac{1}{{n_{\rm
GoF}}}\sum_{i=1}^{{n_{\rm
GoF}}}w(x_{i}{,\theta}),\quad\sigma(w_{p})=\frac{1}{{n_{\rm
GoF}}}\sqrt{\sum_{i=1}^{{n_{\rm GoF}}}(w(x_{i}{,\theta})-{\bar{w}}_{p})^{2}}$
(47)

This quantity acts to reflect the capability of the neural Stein critic to
differentiate between the distributions in the GoF hypothesis testing setting
described in Section 5.1. In addition to the $\hat{P}$ metric from Equation
(47), we may also apply the Stein discrepancy evaluated at the scaleless
neural Stein critic, i.e., the ${\rm SD}[\lambda{\bf f}(\cdot,\theta)]$ (3),
to the holdout dataset from the data distribution as an evaluation metric for
the models as described in Section 5.2.

####  6.3.3 Experimental setup

We again train 2-hidden-layer MLP’s with Swish activation, where each hidden
layer is composed of 512 hidden units, using the default Adam optimizer
parameters by PyTorch. The critics observe 2,000 training samples from $p$,
training on mini-batches of size 100 with a learning rate $10^{-3}$. Each
model is trained for 25 epochs. We consider fixed $\lambda\in\\{1\times
10^{-3},1\times 10^{-2},1\times 10^{-1},1\times 10^{0},2\times 10^{0}\\}$ and
staging scheme $\Lambda(5\times 10^{-1},1\times 10^{-3},0.90)$. The frequency
of updates via the staging strategies is $B_{w}=20$ batches. We fit 10 critics
per regularization strategy. For each critic, we compute the validation SD and
the power metric (47) throughout training using ${n_{\rm GoF}}=1,000$ samples
from $p$ and the same number of “null” samples from $q$.

In addition to assessing the proxy for the test power of the neural Stein
critic in the GoF test via Equation (47), we examine the interpretability of
the critic as a diagnostic tool for anomalous observations. By Equation (7),
the scaleless optimal critic captures the difference in the score of the data
distribution and model distribution. Therefore, a trained neural Stein critic
can indicate which samples in a validation dataset represent the largest
departure from the distribution $q$. We isolate such samples by identifying
samples with high critic witness value (46). We do so by both visualizing the
images in a holdout validation dataset sampled from $p$ which have a high
$w(\cdot{,\theta})$ value, in addition to plotting a heatmap of
$w(\cdot{,\theta})$ reduced using a t-SNE embedding [38] applied to the entire
validation dataset.

Figure 7: In (A), we visualize the $\hat{P}$ metric from Equation (47)
throughout training for various regularization strategies when training neural
Stein critics using the MNIST mixture dataset. To compute $\hat{P}$, the model
is applied to a validation dataset of 1,000 samples. In each case, 10 models
are trained for 25 epochs using each regularization strategy; the mean and
standard deviation are visualized over these 10 models. In (B), a distribution
of null statistics (43) are plotted alongside a statistic calculated over an
${n_{\rm GoF}}=100$ sample testing set from $p$, computed using a
$\Lambda(5\times 10^{-1},1\times 10^{-3},0.90)$ staged critic. In (C), the
Stein discrepancy evaluated at the scaleless neural Stein critic $\lambda{\bf
f}(\cdot,\theta)$ through training is visualized, which is computed using the
same validation datasets of 1,000 samples used to compute $\hat{P}$ in (A).

####  6.3.4 Results

The power approximation using Equation (47) is plotted in Figure 7(A) for each
regularization strategy, where the means and standard deviations are
calculated using the ten models for each regularization scheme. While the
staging strategy does not exhibit such an advantage as in the high-dimension
Gaussian mixture data, staging provides a more rapid increase in the
validation metric in the early training period, yielding a final model of
comparable performance to the fixed-$\lambda$ strategies. In Figure 7(B), we
observe that the test statistic (43) exhibits clear separation from its
bootstrapped ($n_{\rm boot}=500$) null distribution, even in the case when the
number of test samples is relatively small (in this case, ${n_{\rm
GoF}}=100$). While Figure 7(B) shows this distribution for the
$\Lambda(5\times 10^{-1},1\times 10^{-3},0.90)$ staged regularization
strategy, this holds for fixed-$\lambda$ training as well. Finally, the Stein
discrepancy evaluated at the scaleless neural Stein critics $\lambda{\bf
f}(\cdot,\theta)$ applied to the holdout dataset from $p$ is visualized
throughout training for each model in Figure 7(C). This result is similar to
the result from Figure 7(A) in that the staged approach performs comparably to
the best fixed-$\lambda$ regularization strategies.

Figure 8: Embedding via t-SNE of the critic witness function (46) scaled by
the regularization weight $\lambda$ using a $\Lambda(5\times 10^{-1},1\times
10^{-3},0.90)$ staged, regularized neural Stein critic applied to a validation
dataset of 6,000 samples from the data distribution $p$ used to evaluate the
diagnostic capacity of the neural Stein critic. After 25 epochs of training,
the critic was selected to produce these critic witness values. In (A), the
black points represent the portion of the validation set coming from the RBM,
while the red points represent the portion that is true digits 1 from MNIST.
In (B), the points with high critic witness value are more darkly colored. In
(C), the 12 images in the validation set with the highest critic witness value
are shown, applying the critic trained by the staging strategy. Similarly, (D)
shows the 12 images in the validation set with the lowest critic witness
value.

For a more direct understanding of how the neural Stein critics perform,
consider instances of validation data sampled from $p$ in which the critic
witness (46) value is very large. These samples indicate the largest deviation
between the model distribution $q$ and the data distribution $p$. In the
staged $\Lambda(5\times 10^{-1},1\times 10^{-3},0.90)$ setting, we visualize
the critic witness applied to a validation set of 6,000 samples from $p$ by
reducing the images to a two-dimensional embedding via t-SNE. In Figure 8(A),
we observe the embedding of the validation data into this space, where the
true MNIST points are highlighted in red. In Figure 8(B), the true MNIST
digits are found to have a larger critic witness value than those of RBM
samples in the validation dataset.

Furthermore, in the $\Lambda(5\times 10^{-1},1\times 10^{-3},0.90)$ case,
visualizing the images in the validation set from $p$ which have the highest
critic witness value in Figure 8(C) and those which have the lowest critic
witness value in Figure 8(D), we find that this approach correctly identifies
true digits one from MNIST as anomalous while accepting those generated by the
RBM as normal. In the case of the fixed-$\lambda$ regularization strategies,
it seems that all methods do well to identify the true digits “1” in samples
from the data distribution $p$.

##  7 Discussion

We have developed a novel training approach for learning neural Stein critics
by starting with strong $L^{2}$ regularization and progressively decreasing
the regularization weight over the course of training. The advantage of staged
regularization is empirically observed in experiments, especially for high-
dimension data. In all the experiments, it is observed that critics trained
using larger regularization weights at the beginning of training enjoy a more
rapid approximation of the target function, including in the task of detecting
distribution differences between authentic and synthetic MNIST data. These
experimental findings are consistent with the NTK lazy-training phenomenon
that happens with large regularization strength at the beginning phase of
training, which we theoretically prove. We apply the neural Stein method to
GoF tests for which we derive a theoretical guarantee of test power.

The work can be extended in several future directions. Theoretically, within
the lazy-training framework, it would be interesting to characterize the
expressiveness of the NTK kernel (the assumption of Proposition 4.4). One can
also explore a more advanced analysis of the neural network training dynamic,
including the SGD training using mini-batches and going beyond kernel
learning. A related question is to theoretically analyze the later training
stage, for example, to prove the convergence guarantee with decreasing
$\lambda$, and to find the theoretically optimal annealing strategy of
$\lambda$ in later stages. A possible way is to track the change of the NTK
kernel as time evolves, which remains a challenge in studying NTK theory. For
the algorithm, one can further investigate the optimal staging scheme and
explore other types of regularization of the neural Stein critic than the
$L^{2}$-regularization. For the application to GoF testing, here we consider
the typical setting for the Goodness-of-Fit test, which is a simple null
hypothesis, i.e., we test whether or not the data follows a particular
distribution/model; a potentially interesting direction is extending to cases
where the null hypothesis is composite, i.e., comparing the data to a set of
$K$ target distributions. An interesting strategy would be training $K$ neural
networks jointly with possible sharing parameters so as to reduce model size
and computation. Finally, further applications can be conducted on other
modern generative modeling approaches in the same manner as the Gaussian-
Bernoulli RBM, including normalizing flow architectures.

##  8 Proofs

###  8.1 Proofs in Sections 4.1-4.3

###### Proof of Lemma 4.3.

Introduce the notation

$\left[{\bf K}_{0}(x,x^{\prime})\right]_{ij}=K((x,i),(x^{\prime},j)),$

where $K(z,z^{\prime})$ is a positive semi-definite (PSD) kernel defined on
the space of $\mathcal{Z}=\mathcal{X}\times[d]$, that is,

$z=(x,i),\quad x\in\mathcal{X},\quad i\in[d],$

with $dz$ being the induced product measure. The kernel $K$ is PSD due to the
definition (24). One can also verify that $K$ is Hilbert-Schmidt because

$\displaystyle\int_{\mathcal{Z}}\int_{\mathcal{Z}}K(z,z^{\prime})^{2}dzdz^{\prime}=\sum_{i,,j=1}^{d}\int_{\mathcal{X}}\int_{\mathcal{X}}\left[{\bf
K}_{0}(x,x^{\prime})\right]_{ij}^{2}p(x)p(x^{\prime})dxdx^{\prime}$
$\displaystyle~{}~{}~{}=\sum_{i,,j=1}^{d}\int_{\mathcal{X}}\int_{\mathcal{X}}\langle\partial_{\theta}f_{i}(x,\theta(0)),\partial_{\theta}f_{j}(x^{\prime},\theta(0))\rangle_{\Theta}^{2}p(x)p(x^{\prime})dxdx^{\prime}$
$\displaystyle~{}~{}~{}\leq\left(\sum_{i=1}^{d}\int_{\mathcal{X}}\|\partial_{\theta}f_{i}(x,\theta(0))\|_{\Theta}^{2}p(x)dx\right)^{2}<\infty,\quad\text{(by
Cauchy-Schwarz)}$ (48)

and the integrability follows by the $L^{2}$ integrability of
$\partial_{\theta}f_{i}(x,\theta(0))$ in $(\mathcal{X},p(x)dx)$ assumed in the
condition of the lemma.

As a result, the spectral theorem implies that $K(z,z^{\prime})$ has discrete
spectrum $\mu_{k}$, $k=1,2,\cdots$ which decreases to 0, each $\mu_{k}$ is
associated with an eigenfunction $v_{k}(z)$, and
$\\{{v}_{k}\\}_{k=1}^{\infty}$ form an orthonormal basis on $\mathcal{Z}$.
This means that

$\sum_{j=1}^{d}\int_{\mathcal{X}}\left[{\bf
K}_{0}(x,x^{\prime})\right]_{ij}v_{k}(x^{\prime},j)p(x^{\prime})dx^{\prime}=\mu_{k}v_{k}(x,i),\quad\forall
x\in\mathcal{X},\,i\in[d],$

and for any $k,l=1,2,\cdots$,

$\sum_{i=1}^{d}\int_{\mathcal{X}}v_{k}(x,i)v_{l}(x,i)p(x)dx=\delta_{kl}.$

At last, because the neural network has a finite width, $\Theta$ is in a
finite-dimensional Euclidean space of dimensionality $M_{\Theta}$. By (24),
the kernel $K$ has a finite rank at most $M_{\Theta}$. Let the rank of $K$ be
$M$, then $\mu_{1}\geq\cdots\geq\mu_{M}>0$, and the other eigenvalues are
zero. Defining ${\bf v}_{k}(x)$ by $(v_{k}(x,i))_{i=1}^{d}$ finishes the
proof. ∎

###### Proof of Proposition 4.4.

As shown in the proof of Lemma 4.3, there are ortho-normal basis $\\{{\bf
v}_{k}\\}_{k=1}^{\infty}$ of $\mathcal{X}$ with respect to
$\langle\cdot,\cdot\rangle_{p}$, where the first $M$ many consist of eigen-
functions of ${\bf K}_{0}(x,x^{\prime})$. Using the ortho-normal basis, ${\bf
f}^{*}$ has the following expansion with coefficients $c_{k}\in\mathbb{R}$,

${\bf f}^{*}=\sum_{k=1}^{\infty}c_{k}{\bf v}_{k},\quad\|{\bf
f}^{*}\|^{2}_{p}=\sum_{k=1}^{\infty}c_{k}^{2}<\infty.$ (49)

The uniqueness of orthogonal decomposition gives that

${\bf f}_{1}^{*}=\sum_{k=1}^{m}c_{k}{\bf v}_{k},\quad{\bf
f}_{2}^{*}=\sum_{k=m+1}^{\infty}c_{k}{\bf v}_{k}.$ (50)

To prove the proposition, it suffices to prove (27), and then (28) follows by
that

$\|{\bf f}_{1}^{*}\|_{p}\leq\|{\bf f}^{*}\|_{p}$

which follows from (50), and that $e^{-t\delta\lambda}\leq\epsilon$ when
$t\geq{\log(1/\epsilon)}/{(\delta\lambda)}$.

To prove (27): By (25), and define

${\bf w}(x,t):=\lambda\bar{\bf u}(x,t)-{\bf f}^{*}(x),$

we have

${\partial_{t}}{\bf w}(x,t)=-\lambda\int_{\mathcal{X}}{\bf
K}_{0}(x,x^{\prime})\circ{\bf w}(x^{\prime},t)p(x^{\prime})dx^{\prime},$ (51)

and, by that $\bar{\bf u}(x,0)=0$, we have

${\bf w}(x,0)=-{\bf f}^{*}(x).$

Because ${\bf w}(x,0)\in L^{2}(p)$, the evolution equation (51) implies that

${\bf w}(x,t)=\sum_{k=1}^{\infty}b_{k}(t){\bf
v}_{k}(x),\quad\dot{b}_{k}(t)=\begin{cases}-\lambda\mu_{k}b_{k}(t),&k\leq M\\\
0,&k>M\end{cases},\quad b_{k}(0)=-c_{k}.$

Using the notation $\mu_{l}=0$ for $l>M$, we have

$b_{k}(t)=-c_{k}e^{-t\lambda\mu_{k}},\quad k=1,2,\cdots$

This gives that that for any $t>0$,

$-{\bf w}(x,t)=\sum_{k=1}^{\infty}e^{-t\lambda\mu_{k}}c_{k}{\bf
v}_{k}(x)=\sum_{k=1}^{m}e^{-t\lambda\mu_{k}}c_{k}{\bf
v}_{k}(x)+\sum_{k=m+1}^{\infty}e^{-t\lambda\mu_{k}}c_{k}{\bf v}_{k}(x)=:①+②.$

Because $\mu_{k}\geq\delta$ for $k=1,\cdots,m$,

$\|①\|_{p}^{2}=\sum_{k=1}^{m}e^{-2t\lambda\mu_{k}}c_{k}^{2}\leq
e^{-2t\lambda\delta}\sum_{k=1}^{m}c_{k}^{2}=e^{-2t\lambda\delta}\|{\bf
f}_{1}^{*}\|_{p}^{2},$

where the last equality is by (50). In addition,

$\|②\|_{p}^{2}=\sum_{k=m+1}^{\infty}e^{-2t\lambda\mu_{k}}c_{k}^{2}\leq\sum_{k=m+1}^{\infty}c_{k}^{2}{=}\|{\bf
f}_{2}^{*}\|_{p}^{2}.$

Putting together, this gives that

$\|{\bf w}(x,t)\|_{p}\leq\|①\|_{p}+\|②\|_{p}\leq e^{-t\lambda\delta}\|{\bf
f}_{1}^{*}\|_{p}+\|{\bf f}_{2}^{*}\|_{p},$

which proves (27). ∎

###### Proof of Proposition 4.5.

Proof of (i): Recall that
$\dot{\theta}(t)=-{\partial_{\theta}}L_{\lambda}(\theta(t))$ as in (18), and
then by the chain rule,

$\frac{d}{dt}L_{\lambda}(\theta(t))=\langle{\partial_{\theta}}L_{\lambda}(\theta(t)),\dot{\theta}(t)\rangle_{\Theta}=-\|{\partial_{\theta}}L_{\lambda}(\theta(t))\|_{\Theta}^{2}\leq
0,$

which means that $L_{\lambda}(\theta(t))$ is monotonically decreasing over
time. Together with (11), we have that $\forall t\geq 0$,

$\displaystyle\int_{0}^{t}\|\dot{\theta}(s)\|_{\Theta}^{2}ds$
$\displaystyle=\int_{0}^{t}\|{\partial_{\theta}}L_{\lambda}(\theta(s))\|_{\Theta}^{2}ds=L_{\lambda}(\theta(0))-L_{\lambda}(\theta(t))=\mathcal{L}_{\lambda}[{\bf
u}(\cdot,0)]-\mathcal{L}_{\lambda}[{\bf u}(\cdot,t)]$
$\displaystyle=\frac{1}{2\lambda}\left(\|\lambda{\bf u}(\cdot,0)-{\bf
f}^{*}\|_{p}^{2}-\|\lambda{\bf u}(\cdot,t)-{\bf
f}^{*}\|_{p}^{2}\right)\leq\frac{1}{2\lambda}\|\lambda{\bf u}(\cdot,0)-{\bf
f}^{*}\|_{p}^{2}=\frac{1}{2\lambda}\|{\bf f}^{*}\|_{p}^{2}.$ (52)

In the last equality, we have used the condition that ${\bf u}(x,0)=0$. (52)
bounds the change of $\theta(t)$ by Cauchy-Schwarz, that is,

$\|\theta(t)-\theta(0)\|\leq\int_{0}^{t}\|\dot{\theta}(s)\|ds\leq\sqrt{t}\left(\int_{0}^{t}\|\dot{\theta}(s)\|^{2}ds\right)^{1/2}\leq\sqrt{t}\left(\frac{1}{2\lambda}\|{\bf
f}^{*}\|_{p}^{2}\right)^{1/2},$

which proves (29). Finally, by that $\|\theta(0)\|<r/2$, we have
$\|\theta(t)\|<r/2+\sqrt{{t}/{(2\lambda)}}\|{\bf f}^{*}\|_{p}$, and then the
upper bound of $t$ in the condition of (i) ensures that $\theta(t)\in\Theta$.
Thus, for all such $t$, ${\bf u}(\cdot,t)={\bf f}(\cdot,\theta(t))$ is well-
defined and ${\bf f}(\cdot,\theta(t))\in L^{2}(p)$ by Assumption 4, so that
the involved integrals in the loss $L_{\lambda}(\theta(t))$ are all finite.

Proof of (ii): We first note that

$\theta(t)\in B_{r},\quad\forall 0\leq t\leq\frac{1}{2}(\frac{r}{\|{\bf
f}^{*}\|_{p}})^{2}\lambda.$ (53)

This is because $r/2+\sqrt{{t}/{(2\lambda)}}\|{\bf f}^{*}\|_{p}\leq r$ and
$B_{r}\subset\Theta$, thus the condition needed for $t$ by (i) is satisfied,
and then (29) gives that $\|\theta(t)-\theta(0)\|\leq r/2$ for the range of
$t$ being considered. The claim (53) then follows together with the assumption
that $\|\theta(0)\|<r/2$.

In the below, we may omit the dependence on $t$ in notation when there is no
confusion, e.g., we write $\lambda{\bf u}(x^{\prime},t)-{\bf
f}^{*}(x^{\prime})$ as $(\lambda{\bf u}-{\bf f}^{*})(x^{\prime})$. To prove
(30), recall the evolution equations of ${\bf u}$ and $\bar{\bf u}$ as in (23)
and (25),

$\begin{cases}{\partial_{t}}{\bf u}(x,t)=-\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{t}(x,x^{\prime})\circ\left(\lambda{\bf u}-{\bf
f}^{*}\right)(x^{\prime})\\\ {\partial_{t}}\bar{\bf
u}(x,t)=-\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{0}(x,x^{\prime})\circ\left(\lambda\bar{\bf u}-{\bf
f}^{*}\right)(x^{\prime})\\\ {\bf u}(x,0)=\bar{\bf u}(x,0)=0\end{cases}$ (54)

which gives that $({\bf u}-\bar{\bf u})|_{t=0}=0$ and

$\displaystyle\partial_{t}({\bf u}-\bar{\bf u})$
$\displaystyle=-\mathbb{E}_{x^{\prime}\sim p}({\bf K}_{t}-{\bf
K}_{0})(x,x^{\prime})\circ\left(\lambda{\bf u}-{\bf
f}^{*}\right)(x^{\prime})-\lambda\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{0}(x,x^{\prime})\circ\left({\bf u}-\bar{\bf u}\right)(x^{\prime}).$ (55)

By taking $\langle\cdot,\cdot\rangle_{p}$ inner-product with ${\bf u}-\bar{\bf
u}$ on both sides, and that ${\bf K}_{0}$ has the expression (24) and thus is
PSD, we have that

$\displaystyle\frac{d}{dt}\frac{1}{2}\|{\bf u}-\bar{\bf u}\|_{p}^{2}$
$\displaystyle\leq\langle{\bf u}-\bar{\bf u},-\mathbb{E}_{x^{\prime}\sim
p}({\bf K}_{t}-{\bf K}_{0})(\cdot,x^{\prime})\circ\left(\lambda{\bf u}-{\bf
f}^{*}\right)(x^{\prime})\rangle_{p}$ $\displaystyle\leq\|{\bf u}-\bar{\bf
u}\|_{p}\|{\bf K}_{t}-{\bf K}_{0}\|_{p}\|\lambda{\bf u}-{\bf f}^{*}\|_{p},$
(56)

where $\|{\bf K}_{t}-{\bf K}_{0}\|_{p}$ stands for the operator norm of the
kernel integral operator in $L^{2}(p)$. We claim that

$\displaystyle\text{(Claim 1)}\quad\|\lambda{\bf u}(\cdot,t)-{\bf
f}^{*}\|_{p}\leq\|{\bf f}^{*}\|_{p},\quad\forall t\geq 0,$ (57)
$\displaystyle\text{(Claim 2)}\quad\|{\bf K}_{t}-{\bf K}_{0}\|_{p}\leq
L_{1}L_{2}\sqrt{\frac{2t}{\lambda}}\|{\bf f}^{*}\|_{p},\quad\forall
t\leq\frac{1}{2}(\frac{r}{\|{\bf f}^{*}\|_{p}})^{2}\lambda.$ (58)

If both claims hold, then (56) continues as

$\frac{d}{dt}\frac{1}{2}\|{\bf u}-\bar{\bf u}\|_{p}^{2}\leq
L_{1}L_{2}\sqrt{\frac{2t}{\lambda}}\|{\bf f}^{*}\|_{p}^{2}\|{\bf u}-\bar{\bf
u}\|_{p},\quad\forall t\leq\frac{1}{2}(\frac{r}{\|{\bf
f}^{*}\|_{p}})^{2}\lambda.$ (59)

We define

$\Delta(t):=\|\lambda{\bf u}(\cdot,t)-\lambda\bar{\bf u}(\cdot,t)\|_{p}\geq
0,\quad\Delta(0)=0,$

and $\Delta(t)$ is continuous. For fixed $t\leq\frac{1}{2}(\frac{r}{\|{\bf
f}^{*}\|_{p}})^{2}\lambda$, suppose

$\sup_{s\in[0,t]}\Delta(s)=\Delta(t^{*}),\quad\text{for some $t^{*}\leq t$},$

then by (59),

$\displaystyle\frac{1}{2}\Delta(t^{*})^{2}$
$\displaystyle=\lambda^{2}\int_{0}^{t^{*}}(\frac{d}{dt}\frac{1}{2}\|({\bf
u}-\bar{\bf
u})(\cdot,s)\|_{p}^{2})ds\leq\lambda\int_{0}^{t^{*}}L_{1}L_{2}\sqrt{\frac{2s}{\lambda}}\|{\bf
f}^{*}\|_{p}^{2}\Delta(s)ds$ $\displaystyle\leq
L_{1}L_{2}\sqrt{2\lambda}\|{\bf
f}^{*}\|_{p}^{2}\Delta(t^{*})\int_{0}^{t^{*}}\sqrt{s}ds$

this gives that

$\Delta(t^{*})\leq\frac{4\sqrt{2}}{3}L_{1}L_{2}\sqrt{\lambda}\|{\bf
f}^{*}\|_{p}^{2}(t^{*})^{3/2},$

which proves (ii) by that $\Delta(t)\leq\Delta(t^{*})$ and $t^{*}\leq t$.

To finish the proof, it remains to show (57) and (58).

Proof of (57): Define ${\bf w}(x,t):=\lambda{\bf u}(x,t)-{\bf f}^{*}(x)$, By
(54),

$\partial_{t}{\bf w}(x,t)=-\lambda\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{t}(x,x^{\prime})\circ{\bf w}(x^{\prime}),$

and thus

$\frac{d}{dt}\frac{1}{2}\|{\bf w}(\cdot,t)\|_{p}^{2}=-\lambda\langle{\bf
w},\mathbb{E}_{x^{\prime}\sim p}{\bf K}_{t}(\cdot,x^{\prime})\circ{\bf
w}(x^{\prime})\rangle_{p}.$

Denoting the $i$-th entry of ${\bf w}$ as $w_{i}$, by (22), one can verify
that

$\displaystyle\langle{\bf w},\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{t}(\cdot,x^{\prime})\circ{\bf w}(x^{\prime})\rangle_{p}=\mathbb{E}_{x\sim
p}\mathbb{E}_{x^{\prime}\sim
p}\sum_{i,j=1}^{d}\langle\partial_{\theta}f_{i}(x,\theta(t)),\partial_{\theta}f_{j}(x^{\prime},\theta(t))\rangle_{\Theta}w_{i}(x)w_{j}(x^{\prime})$
$\displaystyle~{}~{}~{}=\left\langle\sum_{i=1}^{d}\mathbb{E}_{x\sim
p}\partial_{\theta}f_{i}(x,\theta(t))w_{i}(x),\sum_{j=1}^{d}\mathbb{E}_{x^{\prime}\sim
p}\partial_{\theta}f_{j}(x^{\prime},\theta(t))w_{j}(x^{\prime})\right\rangle_{\Theta}=\|\mathbb{E}_{x\sim
p}\partial_{\theta}{\bf f}(x,\theta(t))\cdot{\bf w}(x)\|_{\Theta}^{2}.$ (60)

This shows that $\frac{d}{dt}\frac{1}{2}\|{\bf w}(\cdot,t)\|_{p}^{2}\leq 0$,
that is, $\|{\bf w}(\cdot,t)\|_{p}^{2}$ monotonically decreases over time.
Thus for any $t\geq 0$,

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

due to that ${\bf u}(x,0)=0$.

Proof of (58): It suffices to show that for any ${\bf v}\in L^{2}(p)$,

$\langle{\bf v},\mathbb{E}_{x^{\prime}\sim p}({\bf K}_{t}-{\bf
K}_{0})(\cdot,x^{\prime})\circ{\bf v}(x^{\prime})\rangle_{p}\leq
L_{1}L_{2}\sqrt{\frac{2t}{\lambda}}\|{\bf f}^{*}\|_{p}\|{\bf v}\|_{p}^{2}.$
(61)

Denote the $i$-th entry of ${\bf v}$ as $v_{i}$, and, similarly as in (60) and
by (22) and (24), one has

$\displaystyle\langle{\bf v},\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{t}(\cdot,x^{\prime})\circ{\bf v}(x^{\prime})\rangle_{p}$
$\displaystyle=\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta(t))\cdot{\bf v}(x)\|_{\Theta}^{2},$ $\displaystyle\langle{\bf
v},\mathbb{E}_{x^{\prime}\sim p}{\bf K}_{0}(\cdot,x^{\prime})\circ{\bf
v}(x^{\prime})\rangle_{p}$ $\displaystyle=\|\mathbb{E}_{x\sim
p}\partial_{\theta}{\bf f}(x,\theta(0))\cdot{\bf v}(x)\|_{\Theta}^{2}.$

Subtracting the two gives that

$\displaystyle\langle{\bf v},\mathbb{E}_{x^{\prime}\sim p}({\bf K}_{t}-{\bf
K}_{0})(\cdot,x^{\prime})\circ{\bf
v}(x^{\prime})\rangle_{p}=\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta(t))\cdot{\bf v}(x)\|_{\Theta}^{2}-\|\mathbb{E}_{x\sim
p}\partial_{\theta}{\bf f}(x,\theta(0))\cdot{\bf v}(x)\|_{\Theta}^{2}$
$\displaystyle~{}~{}~{}=(\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta(t))\cdot{\bf v}(x)\|_{\Theta}+\|\mathbb{E}_{x\sim
p}\partial_{\theta}{\bf f}(x,\theta(0))\cdot{\bf v}(x)\|_{\Theta})$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}(\|\mathbb{E}_{x\sim
p}\partial_{\theta}{\bf f}(x,\theta(t))\cdot{\bf
v}(x)\|_{\Theta}-\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta(0))\cdot{\bf v}(x)\|_{\Theta}).$ (62)

Because $\theta(t)\in B_{r}$, cf. (53),

$\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf f}(x,\theta(t))\cdot{\bf
v}(x)\|\leq\mathbb{E}_{x\sim p}\|\partial_{\theta}{\bf f}(x,\theta(t))\|\|{\bf
v}(x)\|\leq L_{1}\mathbb{E}_{x\sim p}\|{\bf v}(x)\|\leq L_{1}\|{\bf v}\|_{p},$

where the second inequality is by (C1). We have the same upper bound of
$\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf f}(x,\theta(0))\cdot{\bf
v}(x)\|_{\Theta}$ because $\theta(0)\in B_{r}$. Meanwhile, triangle inequality
gives that

$\displaystyle\|\mathbb{E}_{x\sim p}\partial_{\theta}{\bf
f}(x,\theta(t))\cdot{\bf v}(x)\|_{\Theta}-\|\mathbb{E}_{x\sim
p}\partial_{\theta}{\bf f}(x,\theta(0))\cdot{\bf v}(x)\|_{\Theta}$
$\displaystyle~{}~{}~{}\leq\|\mathbb{E}_{x\sim p}(\partial_{\theta}{\bf
f}(x,\theta(t))-\partial_{\theta}{\bf f}(x,\theta(0)))\cdot{\bf v}(x)\|$
$\displaystyle~{}~{}~{}\leq\mathbb{E}_{x\sim p}\|\partial_{\theta}{\bf
f}(x,\theta(t))-\partial_{\theta}{\bf f}(x,\theta(0))\|\|{\bf v}(x)\|$
$\displaystyle~{}~{}~{}\leq L_{2}\|\theta(t)-\theta(0)\|\mathbb{E}_{x\sim
p}\|{\bf v}(x)\|\quad\text{(by (C2) and both $\theta(t)$ and $\theta(0)$ are
in $B_{r}$)}$ $\displaystyle~{}~{}~{}\leq L_{2}\|\theta(t)-\theta(0)\|\|{\bf
v}\|_{p}.$ (63)

By (i) which is proved, the r.h.s. of (63) is upper-bounded by
$L_{2}\sqrt{{t}/{(2\lambda)}}\|{\bf f}^{*}\|_{p}\|{\bf v}\|_{p}$. Putting back
to (62), we have

$\text{(r.h.s. of \eqref{eq:v(Kt-K0)v-bound})}\leq
2L_{1}L_{2}\sqrt{{t}/{(2\lambda)}}\|{\bf f}^{*}\|_{p}\|{\bf v}\|_{p}^{2}$

which proves (61). ∎

###### Proof of Theorem 4.6.

The condition on the largeness of $\lambda$ guarantees that the range of
$t_{0}$ in (31) is non-empty, and (31) ensures that
$\frac{\log(1/\epsilon)}{\delta}\frac{1}{\lambda}\leq
t\leq\frac{1}{2}(\frac{r}{\|{\bf f}^{*}\|_{p}})^{2}\lambda$. For this range of
$t$, Proposition 4.5 applies to give (30). Meanwhile, Proposition 4.4 requires
that $\partial_{\theta}{\bf f}(\cdot,\theta(0))$ are in $L^{2}(p)$, which
follows from the uniform boundedness of $\partial_{\theta}{\bf
f}(\cdot,\theta(0))$ on $\mathcal{X}$ by Assumption 4(C1) due to that
$\theta(0)\in B_{r}$. The condition
$t\geq\frac{\log(1/\epsilon)}{\delta}\frac{1}{\lambda}$ allows Proposition 4.4
to apply to bound $\|\lambda\bar{\bf u}(\cdot,t)-{\bf f}^{*}\|_{p}$ as in (28)
for the range of $t$ being considered. Putting together (30) and (28), by
triangle inequality, we have

$\|\lambda{\bf u}(x,t)-{\bf f}^{*}\|_{p}\leq 2\epsilon\|{\bf
f}^{*}\|_{p}+\frac{4\sqrt{2}}{3}L_{1}L_{2}\sqrt{\lambda}t^{3/2}\|{\bf
f}^{*}\|_{p}^{2}$

and by definition
$\sqrt{\lambda}t^{3/2}=\lambda^{-1}(t_{0}{\log(1/\epsilon)}/{\delta})^{3/2}$,
which proves (32). ∎

###  8.2 Proofs in Section 4.4

###### Proof of Lemma 4.8.

Similarly as in the proof of Proposition 4.5(i), we have

$\frac{d}{dt}\hat{L}_{\lambda}(\hat{\theta}(t))=\langle{\partial_{\theta}}\hat{L}_{\lambda}(\hat{\theta}(t)),\dot{\hat{\theta}}(t)\rangle_{\Theta}=-\|{\partial_{\theta}}\hat{L}_{\lambda}(\hat{\theta}(t))\|_{\Theta}^{2}\leq
0,$

and for any $t\geq 0$,

$\displaystyle\int_{0}^{t}\|\dot{\hat{\theta}}(s)\|_{\Theta}^{2}ds$
$\displaystyle=\int_{0}^{t}\|{\partial_{\theta}}\hat{L}_{\lambda}(\hat{\theta}(s))\|_{\Theta}^{2}ds=\hat{L}_{\lambda}(\theta(0))-\hat{L}_{\lambda}(\hat{\theta}(t)).$
(64)

By the definition (33), because ${\bf f}(x,\theta(0))=\hat{\bf u}(x,0)=0$,
$\hat{L}_{\lambda}(\theta(0))=0$; Meanwhile, let $\hat{\bf f}(x):={\bf
f}(x,\hat{\theta}(t))$,

$\hat{L}_{\lambda}(\hat{\theta}(t))=\mathbb{E}_{x\sim\hat{p}}\left(\frac{\lambda}{2}\|\hat{\bf
f}(x)\|^{2}-{\bf s}_{q}(x)\cdot\hat{\bf
f}(x)\right)-\mathbb{E}_{x\sim\hat{p}}\nabla_{x}\cdot\hat{\bf f}(x).$ (65)

If the last term is with $\mathbb{E}_{x\sim{p}}$, we have that (by that
$\hat{\bf f}(x)$ is in $L^{2}(p)\cap\mathcal{F}_{0}(p)$ by Assumption 3)

$-\mathbb{E}_{x\sim{p}}\nabla_{x}\cdot\hat{\bf
f}(x)=\mathbb{E}_{x\sim{p}}\hat{\bf f}(x)\cdot{\bf
s}_{p}(x)=\mathbb{E}_{x\sim\hat{p}}\hat{\bf f}\cdot{\bf
s}_{p}-(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\hat{\bf f}\cdot{\bf
s}_{p},$

and thus

$\displaystyle-\mathbb{E}_{x\sim\hat{p}}\nabla_{x}\cdot\hat{\bf f}(x)$
$\displaystyle=-\mathbb{E}_{x\sim{p}}\nabla_{x}\cdot\hat{\bf
f}(x)-(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\nabla_{x}\cdot\hat{\bf
f}(x)$ $\displaystyle=\mathbb{E}_{x\sim\hat{p}}\hat{\bf f}\cdot{\bf
s}_{p}-(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\hat{\bf f}\cdot{\bf
s}_{p}-(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\nabla_{x}\cdot\hat{\bf
f}(x)$ $\displaystyle=\mathbb{E}_{x\sim\hat{p}}\hat{\bf f}\cdot{\bf
s}_{p}-①-②,$ (66)

where

$①:=(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}}){\bf
f}(x,\hat{\theta}(t))\cdot{\bf
s}_{p}(x),\quad②:=(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\nabla_{x}\cdot{\bf
f}(x,\hat{\theta}(t)).$

We then have, recalling that ${\bf f}^{*}={\bf s}_{q}-{\bf s}_{p}$,

$\displaystyle-\hat{L}_{\lambda}(\hat{\theta}(t))$
$\displaystyle=-\mathbb{E}_{x\sim\hat{p}}\left(\frac{\lambda}{2}\|\hat{\bf
f}\|^{2}-{\bf s}_{q}\cdot\hat{\bf f}+{\bf s}_{p}\cdot\hat{\bf f}\right)+①+②$
$\displaystyle=-\mathbb{E}_{x\sim\hat{p}}\left(\frac{\|\lambda\hat{\bf f}-{\bf
f}^{*}\|^{2}}{2\lambda}-\frac{\|{\bf f}^{*}\|^{2}}{2\lambda}\right)+①+②$
$\displaystyle\leq\frac{\mathbb{E}_{x\sim\hat{p}}\|{\bf
f}^{*}\|^{2}}{2\lambda}+①+②$ $\displaystyle=\frac{\|{\bf
f}^{*}\|_{p}^{2}}{2\lambda}+③+①+②,\quad③:=\frac{(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|{\bf
f}^{*}(x)\|^{2}}{2\lambda}.$ (67)

The difference $(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|{\bf
f}^{*}(x)\|^{2}$ is the deviation of an independent sum sample average from
its expectation. By Assumption 5(C6), there is an integer $n_{6}$ (possibly
depending on constants $b_{2}$, $b_{q}$) s.t. when $n>n_{6}$, under a good
event $E_{1}$ which happens w.p. $\geq 1-n^{-10}$,

$(\mathbb{E}_{x\sim\hat{p}}-\mathbb{E}_{x\sim{p}})\|{\bf
f}^{*}(x)\|^{2}\leq\sqrt{20}b_{2}\sqrt{\frac{\log n}{n}}.$

As a result, $③\leq\frac{\sqrt{20}b_{2}}{2\lambda}\sqrt{\frac{\log n}{n}}$. As
for $①$ and $②$, we know that if $\hat{\theta}(t)\in{B}_{r}$, then under the
good event in (C5), called $E_{2}$,

$\max\\{|①|,\,|②|\\}\leq\frac{C_{\mathcal{F}}}{n^{\gamma}}+b_{\mathcal{F}}\sqrt{\frac{\log
n}{n}}.$ (68)

However, we have not shown $\hat{\theta}(t)\in B_{r}$ yet. We now let
$n_{\lambda}$ such that when $n>n_{\lambda}$,

$n>n_{6},\quad{\sqrt{20}b_{2}}\sqrt{\frac{\log
n}{n}}+4\lambda(\frac{C_{\mathcal{F}}}{n^{\gamma}}+b_{\mathcal{F}}\sqrt{\frac{\log
n}{n}})<{\|{\bf f}^{*}\|_{p}^{2}}.$ (69)

Since $\|{\bf f}^{*}\|_{p}^{2}>0$ (by Assumption 2) is a fixed constant, for
any $\lambda>0$, the l.h.s. is
$\tilde{O}(n^{-1/2})+\lambda\tilde{O}(n^{-\gamma})$ and thus will be less than
the r.h.s. when $n$ is large enough. We also let $t_{\rm max}$ be defined s.t.
$\|{\bf f}^{*}\|_{p}\sqrt{{t_{\rm max}}/{\lambda}}={r}/{2}$. Consider being
under the intersection of good events $E_{1}$ and $E_{2}$ which happens w.p.
$\geq 1-3n^{-10}$, we claim that, for any $t\leq t_{\rm max}$,
$\hat{\theta}(t)\in B_{r}$.

We prove this by contradiction. If not, then by the continuity of
$\hat{\theta}(t)$ over time, there must be a $0<t_{1}\leq t_{\rm max}$ s.t.
$\|\hat{\theta}(t_{1})\|=r$. By that $\|\theta(0)\|<r/2$, we have

$\|\hat{\theta}(t_{1})-{\theta}(0)\|>r/2.$ (70)

Meanwhile, by (64)(67),

$\int_{0}^{t_{1}}\|\dot{\hat{\theta}}(s)\|_{\Theta}^{2}ds=-\hat{L}_{\lambda}(\hat{\theta}(t_{1}))\leq\frac{\|{\bf
f}^{*}\|_{p}^{2}}{2\lambda}+①+②+③$

where $①$ and $②$ are taking value at $t=t_{1}$. Under $E_{1}$, we already
have $③\leq\frac{\sqrt{20}b_{2}}{2\lambda}\sqrt{\frac{\log n}{n}}$; now
$\hat{\theta}(t_{1})\in\overline{B}_{r}$, applying (C5) we know that under
$E_{2}$ the bound (68) holds. By the definition of $n_{\lambda}$ in (69), we
have that

$①+②+③<\frac{\|{\bf f}^{*}\|_{p}^{2}}{2\lambda},$

and as a result,

$\int_{0}^{t_{1}}\|\dot{\hat{\theta}}(s)\|_{\Theta}^{2}ds<\frac{\|{\bf
f}^{*}\|_{p}^{2}}{\lambda}.$

Then, by Cauchy-Schwarz,

$\|\hat{\theta}(t_{1})-\theta(0)\|\leq\int_{0}^{t_{1}}\|\dot{\hat{\theta}}(s)\|_{\Theta}ds\leq\sqrt{t_{1}}(\int_{0}^{t_{1}}\|\dot{\hat{\theta}}(s)\|_{\Theta}^{2}ds)^{1/2}<\sqrt{\frac{t_{1}}{\lambda}}\|{\bf
f}^{*}\|_{p}\leq\sqrt{\frac{t_{\rm max}}{\lambda}}\|{\bf
f}^{*}\|_{p}=\frac{r}{2},$

which means that $\|\hat{\theta}(t_{1})-\theta(0)\|<r/2$ and this contradicts
with (70).

Now we have shown that for $\hat{\theta}(t)\in B_{r}$ for any $t\leq t_{\rm
max}$. Applying the same argument to bound
$\int_{0}^{t}\|\dot{\hat{\theta}}(s)\|_{\Theta}^{2}ds$ for any such $t$, and
in particular the upper bound of $①+②+③<\frac{\|{\bf
f}^{*}\|_{p}^{2}}{2\lambda}$ holds, one can verify that

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

and this finishes the proof of the lemma. ∎

###### Proof of Proposition 4.9.

Recall the evolution equation of $\bar{\bf u}(x,t)$ in (25), which can be
written as

${\partial_{t}}\bar{\bf u}(x,t)=-\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}))-\mathbb{E}_{x^{\prime}\sim p}{\bf
K}_{0}(x,x^{\prime})\circ{\bf s}_{p}(x^{\prime}).$

Because $\|\partial_{\theta}{\bf f}(x,\theta)\|$ is uniformly bounded by
$L_{1}$ by Assumption 4(C1), each column of the $d$-by-$M_{\Theta}$ matrix
$\partial_{\theta}{\bf f}(x,\theta(0))$, which can be viewed as a vector field
on $\mathcal{X}$, is uniformly bounded and thus in