# Large scale global optimization of ultra-high dimensional non-convex
landscapes based on generative neural networks

Jiaqi Jiang,1 Jonathan A. Fan, 1

###### Abstract

We present a non-convex optimization algorithm metaheuristic, based on the
training of a deep generative network, which enables effective searching
within continuous, ultra-high dimensional landscapes. During network training,
populations of sampled local gradients are utilized within a customized loss
function to evolve the network output distribution function towards one peaked
at high performing optima. The deep network architecture is tailored to
support progressive growth over the course of training, which allows the
algorithm to manage the curse of dimensionality characteristic of high
dimensional landscapes. We apply our concept to a range of standard
optimization problems with dimensions as high as one thousand and show that
our method performs better with fewer functional evaluations compared to
state-of-the-art algorithm benchmarks. We also discuss the role of deep
network over-parameterization, loss function engineering, and proper network
architecture selection in optimization, and why the required batch size of
sampled local gradients is independent of problem dimension. These concepts
form the foundation for a new class of algorithms that utilize customizable
and expressive deep generative networks to solve non-convex optimization
problems.

## Introduction

High dimensional, non-convex optimization problems are pervasive in many
scientific and engineering domains, including computational materials science
(Chen et al. 2017; Jørgensen et al. 2018), electromagnetics (Piggott et al.
2015; Sell et al. 2017), circuits design (Held et al. 2011), process
engineering (Sieniutycz and Jezowski 2009), and systems biology (Balsa-Canto
et al. 2012). These problems are known to be very difficult to solve because
they are NP-hard, and algorithms aiming to definitively search for the global
optimum, such as branch and bound methods, cannot practically scale to high
dimensional systems. As such, various algorithm heuristics have been
developed, ranging from evolutionary metaheuristics to Bayesian optimization
(Pelikan et al. 1999; Snoek et al. 2015), which use judicious sampling of the
landscape to identify high performing optima. In all cases, it remains
challenging to apply these algorithms to ultra-high dimensional spaces with
dimensions of hundreds to thousands due to the curse of dimensionality.

The explosion of interest and research in deep neural networks over the last
decade has presented new opportunities in optimization, as the process of
training a deep network involves solving a high dimensional optimization
problem. To this end, gradient-based optimization metaheuristics termed global
topology optimization networks (GLOnets) (Jiang and Fan 2020, 2019) were
recently proposed that use the training of a deep generative network to
perform non-convex optimization. The concept applies to optimization problems
where $\mathbf{x}$ is a $d$-dimensional variable and the goal is to maximize
the smoothly varying, non-convex objective function $f(\mathbf{x})$. To run
the metaheuristic, the generative network is first initialized so that it
outputs a distribution of $\mathbf{x}$ values that spans the full optimization
landscape. Over the course of network training, this distribution is sampled,
$f(\mathbf{x})$ and local gradients are computed for these sampled points, and
these values are incorporated into a customized loss function and
backpropagated to evolve and narrow the distribution around high performing
optima. Initial demonstrations indicate that GLOnets can perform better than
standard gradient-based optimizers and global search heuristics for various
non-convex optimization problems. However it is unable to extend to high
dimensional problems in its current form, and the lack of interpretability
with this black box algorithm has made it difficult to understand if and how
it can to adapt to more general problems, including high dimensional problems.

In this Article, we introduce the progressive growing GLOnet (PG-GLOnet) in
which optimization within an ultra-high dimensional non-convex landscape is
mediated through the training of a progressive growing deep generative
network. Our tailoring of the network architecture for this optimization task
serves to incorporate knowledge and assumptions about the optimization
landscape into the metaheuristic, which is a requirement for tractably
navigating ultra-high dimensional landscapes. We also explain how the
algorithm works to smoothen the design landscape, how evaluation of the loss
function serves as a gradient estimation calculation, and why the number of
required functional evaluations is independent of problem dimension. With
standard benchmarking test functions, we show that our concept performs better
than state-of-the-art algorithms with fewer functional evaluations for one
thousand dimensional problems. We anticipate that the customization of network
architectures within the GLOnets framework will seed new connections between
deep learning and optimization.

## Progressive Growing GLOnets Algorithm and Benchmarking

Figure 1: (a) Framework of the GLOnet algorithm. A deep generative network
produces a distribution of design variables $\mathbf{x}$ and the distribution
is narrowed around high performing optima by backpropagation. (b) PG-GLOnet
generator architecture. A set of growing blocks are implemented and activated
over the course of network training to enable optimization in very high-
dimensional landscapes.

The PG-GLOnet concept builds on the foundation of the original GLOnet
algorithm, which we briefly review here. The optimization problem to be solved
with GLOnets can be written in the following form:

$\max_{\mathbf{x}}f(\mathbf{x})$ (1)

where $f(\mathbf{x})$ is a non-convex, continuous objective function with
feasible gradients. With GLOnets, this optimization problem is indirectly
solved through the training of a general neural network (Figure 1a), where the
input is a $d$-dimensional random variable $\mathbf{z}$ with a standard normal
distribution and the output is a distribution of $\mathbf{x}$’s. The generator
therefore serves to map $\mathbf{z}$ onto $\mathbf{x}=G(\mathbf{z};\phi)$ with
a distribution $P(\mathbf{x};\phi)$, where $\mathbf{\phi}$ denotes the
trainable neural network parameters. The optimization objective for the
generator is defined as:

$L=\max_{\mathbf{\phi}}\mathop{\mathbb{E}}_{\mathbf{x}\sim
P(\mathbf{x};\phi)}\exp\left[\frac{f(\mathbf{x})}{T}\right]$ (2)

The distribution that maximizes this expected value is a delta function
centered at the global optimum, and as such, an ideally trained generator will
produce a narrow distribution centered at the global optimum, thereby solving
the original optimization problem. The use of the exponential function and the
hyperparameter $T$ in the optimization objective further enhance the valuation
of the global optimum, and more generally high performing optima, in the
design space.

Generator training is consistent with conventional deep learning training
methods: gradients of the objective function with respect to network
parameters, $\nabla_{\phi}\mathbb{E}f$, are calculated through
backpropagation, and they are used to iteratively optimize $\mathbf{\phi}$
using standard gradient-based methods. In practice, the objective function is
approximated by a batch of $M$ samples. $P(\mathbf{x};\phi)$, on the other
hand, is typically implicit and cannot be directly sampled. To circumvent this
issue, we draw $M$ samples $\\{\mathbf{z}^{(m)}\\}_{m=1}^{M}$ from the
standard normal distribution, transform them to
$\\{\mathbf{x}^{(m)}\\}_{m=1}^{M}$, and then approximate $L$ and its gradient
$\nabla_{\phi}L$ with respect to network parameters $\phi$:

$L\approx\frac{1}{M}\sum_{m=1}^{M}\exp\left[\frac{f(\mathbf{x}^{(m)})}{T}\right]$
(3)
$\nabla_{\phi}L\approx\frac{1}{M}\sum_{m=1}^{M}\frac{1}{T}\exp\left[\frac{f(\mathbf{x}^{(m)})}{T}\right]\nabla_{\mathbf{x}}f\cdot
D_{\phi}\mathbf{x}^{(m)}$ (4)

$\nabla_{\mathbf{x}}f=[\frac{\partial f}{\partial x_{1}},\frac{\partial
f}{\partial x_{2}},\dots,\frac{\partial f}{\partial x_{d}}]$ are the gradients
of $f(\mathbf{x})$ and
$D_{\phi}\mathbf{x}=\frac{\partial(x_{1},x_{2},\dots)}{\partial(\phi_{1},\phi_{2},...)}$
is the Jacobian matrix. Evaluation of $f(\mathbf{x})$ is usually performed by
a numerical simulator and the gradient of $f(\mathbf{x})$ can be calculated
explicitly or by auto-differentiation for analytic expressions, or by the
adjoint variables method (AVM).

In the initial conception of GLOnet, which we term FC-GLOnet, the generative
network was a fully connected deep network and was capable of effectively
addressing optimization problems with a modest number of dimensions. However,
it was found to be ineffective at optimizing within very high dimensional
landscapes due to the curse of dimensionality, which makes a direct search for
the global optimum within a full, high dimensional landscape an intractable
proposition. We therefore propose the PG-GLOnet, which utilizes a generative
network that outputs a distribution that gradually grows from a coarse, low
dimensional space to a fine, high dimensional space. By tailoring the network
architecture in this way, we regularize the optimization process to take place
over differing degrees of optimization landscape smoothing, enabling our
search process to be computationally efficient and tractable.

The PG-GLOnet generator architecture is shown in Figure 1b. The progressive
growth concept is inspired by progressively growing GANs (Karras et al. 2017)
that have been developed in the computer vision community to process images
with increasing spatial resolution during network training. The input to the
network is a $D$-dimensional random vector $\mathbf{x}^{0}$, and its dimension
is much smaller than that of $\mathbf{x}$. With $L$ growing blocks, the
network simultaneously transforms and increases the dimensionality of the
input vector, and its output is a $2^{L}D$ dimensional vector $\mathbf{x}^{L}$
that matches the dimensionality of $\mathbf{x}$.

In each growing block, the input vector dimension is doubled in two ways, by
direct upsampling and by a linear transform. The resulting outputs are
combined together and further transformed using a non-linear activation
function:

$\mathbf{x}^{out}_{2d\times
1}=q\left((1-\alpha)\begin{pmatrix}\mathbf{x}^{in}_{d\times 1}\\\
\mathbf{x}^{in}_{d\times 1}\end{pmatrix}+\alpha\ A_{2d\times
d}\cdot\mathbf{x}^{in}_{d\times 1}\right)$ (5)

$A_{2d\times d}$ are trainable parameters in the linear transformation branch,
$q(\cdot)$ is a non-linear activation function, and $\alpha$ is a
hyperparameter that is manually tuned over the course of optimization.

Initially, $\alpha$’s for all of the growing blocks in the network are set to
$0$, such that the vector outputted by each block has the same effective
dimensionality as its input vector. The network output $\mathbf{x}^{L}$
therefore has an effective dimensionality that matches the dimensionality of
the input $\mathbf{x}^{0}$. As $\alpha$ is increased for a particular growing
block, its output vector becomes dominated by its linear transformation
branch, as opposed to its upsampling branch, and it has an effective
dimensionality that exceeds and eventually doubles that of the growing block
input vector. The effective dimensionality of $\mathbf{x}^{L}$ therefore
arises from the aggregation of effective dimensionality increases from all
growing blocks. To control the effective dimensionality of $\mathbf{x}^{L}$
over the course of PG-GLOnet training, $\alpha$ is manually changed from 0 to
1 sequentially from the left to right blocks (bottom of Figure 1b). At the end
of PG-GLOnet training, $\alpha$ is $1$ for all growing blocks and the
effective dimensionality of $\mathbf{x}^{L}$ matches that of $\mathbf{x}$.

Figure 2: Benchmark results for PG-GLOnet. (a,b) Benchmark of PG-GLOnet with
ADAM for a modified Rastrigin test function that can be tuned from convex to
non-convex. (a) Functional evaluations required by both algorithms for a two
dimensional Rastrigin test function with differing $\rho$. Plots of the test
function for $\rho=0$ and $\rho=10$ are shown on the left. (b) Optimal values
achieved by PG-GLOnet and ADAM for a ten dimensional Rastrigin test function
with differing $\rho$. (c,d) Benchmark of PG-GLOnet with CMA-ES for modified
(c) Rastrigin and (d) Schwefel test functions. Both plots show the required
number of functional evaluations required to find the global optimum as a
function of test function dimension.

To evaluate the efficacy of PG-GLOnet in solving high dimensional non-convex
optimization problems, we perform a series of benchmark numerical experiments
where we optimize a set of standard test functions with PG-GLOnet and other
established algorithms. In the first set of experiments, we consider a testing
function that can be tuned from a convex to non-convex function and compare
PG-GLOnet with ADAM, a well known momentum-based gradient descent algorithm
that is typically more effective than gradient descent. ADAM is a local
optimization algorithm and performs well on convex objective functions but can
get trapped within local optima for non-convex functions. Our test function is
a modified Rastrigin function defined as follows:

$f(\mathbf{x};\rho)=\rho d+\sum_{i=1}^{d}[x_{i}^{2}-\rho\cos(2\pi x_{i})]$ (6)

$\rho$ is a hyperparameter that specifies the amplitude of the sinusoidal
modulation within the function. When $\rho=0$,
$f(\mathbf{x};\rho)=\sum_{i=1}^{d}x_{i}^{2}$ and is a convex function. As
$\rho$ increases, more local optima emerge and these optima become separated
by larger magnitude barriers.

We first consider the computational cost required by ADAM and PG-GLOnet to
find the global optimum of a two dimensional modified Rastrigin function as a
function of $\rho$. For ADAM, we run 10000 optimizations for 200 iterations
with random starting points, and for PG-GLOnet, we run the algorithm 10 times
with a batch size of 20 for 200 total iterations. In both cases, the
algorithms terminate early when they output results within $10^{-3}$ of the
global optimum, and computational cost is quantified as the average number of
function evaluations required to find the global optimum. The results are
summarized in Figure 2a and indicate that for convex or nearly convex
optimization landscapes, ADAM is more efficient at finding the global optimum.
This efficiency arises because ADAM is a specially tailored local optimizer
that is well suited for these types of problems, while PG-GLOnet always
requires relatively large batch sizes and more iterations to converge. As
$\rho$ increases, orders-of-magnitude more ADAM evaluations are required to
search for the global optimum due to trapping within local optima in the
design landscape. The computational cost for PG-GLOnet, on the other hand,
does not increase nearly as rapidly due to its ability to navigate non-convex
landscapes and is ten times more efficient than ADAM for $\rho$ greater than
$3$.

We also perform benchmarks between ADAM and PG-GLOnet for a ten dimensional
problem. Due to the inability for ADAM to converge to the global optimum in
non-convex, high dimensional landscapes, we perform this benchmark differently
and compare the best optimal value found by ADAM and PG-GLOnet given the same
amount of computational resources. Here, we run ADAM for 200 iterations with
20 random starting points and PG-GLOnet for 200 iterations with a batch size
of 20. We run these benchmark experiments ten times and average the best
values from each experiment, and the results are reported in Figure 2b. We
find that the PG-GLOnet is able to consistently find solutions at or near the
global optimum for all values of $\rho$, but the local optimizer gets
progressively worse as $\rho$ increases.

In our next set of benchmark experiments, we compare PG-GLOnet with the
covariance matrix adaptation evolution strategy (CMA-ES), which is an
established evolutionary algorithm used to perform population-based global
searching of an optimization landscape. Compared to ADAM, it is more suitable
for performing non-convex optimization. We consider two standard non-convex
testing functions with lots of local optima, the Rastrigin and Schwefel
functions (defined in the Appendix).

Plots in Figures 2c and 2d show the average number of function evaluations
required to find the global optimum as a function of problem dimension $d$.
The computational cost of CMA-ES increases exponentially as the problem
dimension becomes larger, indicating the intractability of applying this
algorithm to ultra-high dimensional problems. For the Schwefel function, we
limited our CMA-ES benchmarking experiments to a problem dimension of 20 due
to this scaling trend. PG-GLOnet, on the other hand, has a relatively small
computational cost that is not sensitive to the dimension. In fact, the same
neural network architecture and batch size is used for all problems. A more
detailed discussion as to the origins of problem dimension and batch size
decoupling is provided in the Discussion section.

Finally, we benchmark PG-GLOnet with state-of-art algorithms on testing
functions proposed by the CEC’2013 Special Session and Competition on Large-
Scale Global Optimization (LSGO) (Li et al. 2013). We consider the six non-
convex benchmark functions from the competition, which involve variations and
combinations of the Rastrigin and Ackely functions and are defined in the
Appendix. These benchmark functions were designed to incorporate a number of
challenging features for optimization, including:

  1. 1.

High dimensions. The design space of a optimization problem grows
exponentially as the dimension of design variables increases. These benchmark
functions utilize one thousand dimensional landscapes.

  2. 2.

Functions with non-separable subcomponents. The whole design variable is
decomposed into several subcomponents and dimensions within each subcomponent
are strongly coupled together.

  3. 3.

Imbalance in the contribution of subcomponents. The contribution of a
subcomponent is magnified or dampened by a coefficient.

  4. 4.

Non-linear transformations to the base functions. Three transformations are
applied to break the symmetry and introduce some irregularity on the
landscape: (1) Ill-conditioning (2) Irregularities (3) Symmetry breaking.

Table 1: Optimization results from conjugate gradient, CC-RDG3, DGSC, FC-GLOnet followed by local gradient descent, and PG-GLOnet followed by local gradient descent, as applied to 1000-dimensional benchmark functions. |  CG |  CC-RDG3 |  DGSC |  FC-GLOnet +Local Refinement |  PG-GLOnet +Local Refinement
---|---|---|---|---|---
$f_{1}$ |  (3.65 $\pm$ 0.07)e+04 |  (2.36 $\pm$ 0.11)e+03 |  (7.15 $\pm$ 1.61)e+02 |  (1.74 $\pm$ 0.01)e+03 |  (1.85 $\pm$ 0.63)e+03
$f_{2}$ |  (1.98 $\pm$ 0.00)e+01 |  (2.04 $\pm$ 0.00)e+01 |  (2.07 $\pm$ 0.00)e+01 |  (1.98 $\pm$ 0.00)e+01 |  (4.38 $\pm$ 0.75)e-03
$f_{3}$ |  (2.08 $\pm$ 0.18)e+07 |  (2.27 $\pm$ 0.30)e+06 |  (3.27 $\pm$ 0.66)e+06 |  (8.76 $\pm$ 2.99)e+05 |  (5.05 $\pm$ 0.86)e+05
$f_{4}$ |  (9.76 $\pm$ 0.01)e+05 |  (9.96 $\pm$ 0.00)e+05 |  (1.06 $\pm$ 0.00)e+06 |  (9.94 $\pm$ 0.01)e+05 |  (2.95 $\pm$ 0.82)e+02
$f_{5}$ |  (1.15 $\pm$ 0.97)e+09 |  (1.45 $\pm$ 0.32)e+08 |  (1.79 $\pm$ 0.53)e+08 |  (4.22 $\pm$ 0.77)e+07 |  (1.17 $\pm$ 0.49)e+08
$f_{6}$ |  (8.84 $\pm$ 0.00)e+07 |  (9.11 $\pm$ 0.14)e+07 |  (9.38 $\pm$ 0.03)e+07 |  (9.03 $\pm$ 0.01)e+07 |  (4.66 $\pm$ 0.79)e+04

To globally search these landscapes for the global optimum, we perform a two
step optimization procedure. First, we run PG-GLOnet for each benchmark
function for 200 iterations and a batch size of 100, from which our generative
network outputs a narrow distribution of $\mathbf{x}$’s in promising regions
of the optimization landscape. We then sample this distribution 100 times and
perform local gradient descent on each of these design variables for an
additional 200 iterations. The best function values found by PG-GLOnet plus
local gradient descent are reported in Table 1, together with results produced
from FC-GLOnet plus local gradient descent, local conjugate gradient descent,
and two state-of-art non-convex optimization algorithms that were the best
performing algorithms in the most recent LSGO contest: CC-RDG3, which is a
divide-and-conquer method (Sun et al. 2019), and DGSC, which is a differential
group method utilizing spectral clustering (Li et al. 2019). We observe that
PG-GLOnet with local gradient descent refinement is able to significantly
outperform the other algorithms for the majority of test functions. In
addition, the total computational cost of the two step optimization procedure
is only $4\times 10^{4}$ function evaluations, while CC-RDG3 and DGSC require
$3\times 10^{6}$ function evaluations.

## Discussion

We discuss the origins of the efficiency and efficacy of PG-GLOnet in solving
ultra-high dimensional non-convex optimization problems. First, we examine how
the generic GLOnet algorithm operates and why it is able to effectively
utilize a gradient-based strategy to solve non-convex optimization problems.
Second, we examine the role of the progressive growing generative network
architecture in PG-GLOnet in solving ultra-high dimensional problems. By
understanding the relationship between network architecture and optimization
procedure, we elucidate built-in assumptions used by PG-GLOnet in its search
for the global optimum.

With the generic GLOnet algorithm, the original optimization problem cited in
Equation 1 is reframed as a related problem (Equation 2) that addresses a
transformed, smoothened optimization landscape. The key concepts that produce
this landscape transformation and enable effective gradient-based optimization
are outlined in Figure 3a and are: 1) distribution optimization, where the
original problem involving the optimization of $\mathbf{x}$ is transformed to
a problem involving the optimization of parameters within a simple
distribution $P(\mathbf{x})$; 2) exponential transformation, where the
objective function is exponentially weighted; 3) over-parametrization, where
the distribution $P(\mathbf{x})$ is now parameterized by a neural network with
hundreds to thousands of weights; and 4) gradient estimation, where gradients
that specify the evolution of the continuous distribution $P(\mathbf{x})$ are
accurately computed through discrete samplings of $\mathbf{z}$.

Figure 3: Conceptualization of the GLOnets optimization platform. (a)
Visualization of key concepts that enable effective gradient-based
optimization within a non-convex landscape, including: 1) transforming the
optimization problem to the optimization of parameters within a distribution;
2) exponential weighing of the objective function; 3) over-parameterization of
the distribution function; and 4) Effective gradient estimation during the
network training procedure. b) Role of the progressively growing PG-GLOnet
network architecture in optimization, visualized for a single growing block
applied to a two-dimensional problem. When $\alpha$ is zero, PG-GLOnet
searches within a one dimensional slice of the two dimensional landscape. As
$\alpha$ increases, the effective dimensionality of the PG-GLOnet output
distribution increases and enables searching of more of the landscape. Upon
the completion of PG-GLOnet training, the generator output distribution
collapses to the global optimum.

Distribution optimization. With the concept of distribution optimization, the
original problem of searching for an optimal $\mathbf{x}$ is recast as a
population-based search in which parameters within a distribution function are
optimized, thereby enabling a search for the global optimum in a smoother and
higher dimensional optimization landscape. This concept is shared by other
population-based optimization algorithms, such as CMA-ES. To visualize the
concept, we consider a non-convex one-dimensional function $f(\mathbf{x})$
plotted as a blue line in the leftmost figure in Figure 3a. The objective is
to maximize $f(\mathbf{x})$, and the function contains multiple local maxima
separated by deep valleys. It is easy for optimization algorithms,
particularly gradient-based algorithms, to get trapped in the local optima.
For example, if gradient descent optimization is used and is initialized at
the yellow dot position, the algorithm will converge to the local optimum
delineated by the red dot. With this approach, multiple independent gradient
descent optimizations with random starting points are needed to increase the
possibility of finding the global optimum. For these problems, gradient-free
optimization heuristics are often employed, which can reduce the chances of
trapping within suboptimal maxima but which introduce a more stochastic nature
to the search process.

However, if we consider the optimization of a distribution function that
interacts with the global optimization landscape, local information at
different parts of the landscape can be aggregated and collectively utilized
to evolve this distribution in a manner that reduces issues of trapping within
suboptimal maxima. Formally, we transform the optimization variable
$\mathbf{x}$ to parameters within the distribution $P(\mathbf{x})$, and the
globally optimal distribution is one that is narrowly peaked around the global
optimum. Distribution functions can be explicitly parameterized in many ways.
As a simple illustrative example that builds on our discussion of the one-
dimensional $f(\mathbf{x})$, we consider the one-dimensional Gaussian
distribution denoted as $P(\mathbf{x;\mu,\sigma})$, shown as the red curve in
the leftmost figure in Figure 3a. $\mu$ and $\sigma$ refer to mean and
standard deviation, respectively.

With a Gaussian distribution function, the objective function now becomes
transformed to the expected value of $f(\mathbf{x})$ as a function of
$(\mu,\sigma)$: $\mathop{\mathbb{E}}_{\mathbf{x}\sim
P(\mathbf{x};\mu,\sigma)}f(\mathbf{x})$. As this new optimization landscape is
a function of two distribution parameters, $\mu$ and $\sigma$, it is two
dimensional. We can directly visualize this new landscape by evaluating $\int
f(\mathbf{x})P(\mathbf{x};\mu,\sigma)d\mathbf{x}$ for all values of
$(\mu,\sigma)$, and the result is summarized in the second figure from the
left in Figure 3a. The horizontal line section at the bottom of the contour
plot, where $\sigma$ equals zero, is the original one-dimensional
$f(\mathbf{x})$ with multiple optima. As $\sigma$ increases to finite values
above zero, the landscape becomes smoother. Mathematically, horizontal line
sections for finite sigma are calculated by convolving $f(\mathbf{x})$ with
the Gaussian function, producing a Gaussian blur that leads to smoothening.
This smoothened landscape facilitates gradient-based optimization of
$(\mu,\sigma)$ when the distribution is initialized to large $\sigma$ values,
and the final optimized distributions converge to the original $f(\mathbf{x})$
space at the bottom of the plot. However, while this two-dimensional landscape
is smoother than the original $f(\mathbf{x})$, there remain multiple
distribution parameter initializations for which the gradient-based optimizer
converges to suboptimal maxima.

Exponential transformation. To further smoothen the optimization landscape and
enhance the presence of the global optimum, we perform an exponential
transformation of the objective function. Mathematically, the objective
function for the distribution optimization problem becomes:
$\mathop{\mathbb{E}}_{\mathbf{x}\sim
P(\mathbf{x};\mu,\sigma)}\exp\left[\frac{f(\mathbf{x})}{T}\right]$. The
temperature term $T$ modulates the impact of the global optimum on the
optimization landscape such that low $T$ produces strong landscape modulation
by the global optimum. For our one-dimensional $f(\mathbf{x})$ example, the
exponentially transformed landscape is plotted in the second figure from the
left in Figure 3a and shows that the local optima has faded out, such that
gradient-based optimization within this landscape is more likely to converge
to the global optimum.

The choice of $T$ depends on the scale of $f(\mathbf{x})$. Consider
$f(\mathbf{x})$ that is linearly normalized to span $(0,1)$. Such
normalization can be typically achieved based on prior knowledge about the
upper and lower bound of $f(\mathbf{x})$. If we want to amplify
$f(\mathbf{x})$ for $f(\mathbf{x})>f_{d}$ and minimize $f(\mathbf{x})$ for
$f(\mathbf{x})<f_{d}$, where $f_{d}$ is a division point between 0 and 1, the
temperature is chosen to be $T=f_{d}/\log(1+f_{d})$. For example, if $f_{d}$
is chosen to be the golden ratio, then the temperature is roughly $T=1.3$. In
practice, the selection of $f_{d}$ is problem specific, and $T$ can be treated
as a hyperparameter that can be manually tuned around 1 for tailoring to a
particular problem.

Over-parameterization. To further enhance the ability for GLOnet to
efficiently and reliably converge to the global optimum, we next consider the
concept of over-parameterization in which the distribution $P(\mathbf{x})$ is
now a neural network parameterized by weights $\mathbf{\phi}$. The objective
function then becomes: $\mathop{\mathbb{E}}_{\mathbf{x}\sim
P(\mathbf{x};\phi)}\exp\left[\frac{f(\mathbf{x})}{T}\right]$. Our use of a
neural network is inspired by the fact that deep network training involves the
solving of an extremely high dimensional non-convex optimization problem, that
the convergence of the neural network is typically insensitive to
initialization, and that good neural network parameters can be found using
backpropagation.

The underlying mathematical principles outlining why gradient descent is so
effective for deep network training have been revealed to some extent by
computer scientists in recent years. (Choromanska et al. 2015; Arora, Cohen,
and Hazan 2018; Draxler et al. 2018) First, the parameter space of deep
networks is a high-dimensional manifold, such that most local optima are
equivalently good and the probability of converging to a bad optimum during
training decreases quickly with network size. Second, these equivalently high
performing local optima originate from neural network over-parameterization,
which builds in redundancy in the optimization landscape that speeds up and
stabilizes the gradient-based optimization process.

To understand how this applies to GLOnet, we revisit our one-dimensional
$f(\mathbf{x})$ landscape in which local optima are separated by deep
barriers. When the optimization landscape is transformed using
$P(\mathbf{x,\phi})$, it frames the optimization problem in a very high
dimensional landscape, as the dimensionality of $\mathbf{\phi}$ is much higher
than $\mathbf{x}$. Solutions to the optimization problem therefore reside in a
high-dimensional manifold, such that many different $\mathbf{\phi}$’s serve as
high performing local optima. Additionally, local optima in $f(\mathbf{x})$
are no longer separated by deep barriers but are instead connected by pathways
with low to no barriers in our transformed high dimensional landscape,
mitigating trapping within these local optima during gradient-based
optimization. The high dimensional landscape representing the transformed
$f(\mathbf{x})$ is visualized as a two-dimensional projection in the rightmost
plot in Figure 3a. The global optimum is now a connected band in the
optimization landscape, as opposed to a single point in $f(\mathbf{x})$, and
there are fewer energy barriers preventing gradients from converging to the
global optimum, enabling gradient descent optimization to be more robust and
faster. We note that neural network depth and expressivity play a large role
in determining the practical impact of over-parameterization on optimization,
and as a demonstration, we compare the performance of GLOnets based on linear
and deep non-linear networks in the Appendix.

Gradient estimation. A critical feature to maximizing the performance of
GLOnet is ensuring that gradients used to evolve $P(\mathbf{x})$, which are
approximated using a finite batch of samples, are sufficiently accurate. There
are two methods for gradient estimation that can be used for GLOnets. The
first is to use a score function gradient estimator, which utilizes the
evaluated derivatives of the probability distribution $P(\mathbf{x};\phi)$ and
$f(\mathbf{x})$. This method for estimation requires explicit evaluation of
derivatives to $P(\mathbf{x};\phi)$ but only an implicit evaluation of
$\nabla_{\mathbf{x}}f$. The second is to use a pathwise gradient estimator,
which relies on knowing the explicit derivatives of $f(\mathbf{x})$ but for
which the probability distribution $P(\mathbf{x};\phi)$ can be implicit.
Empirically, we find for GLOnet that the pathwise gradient estimator more
consistently produces smaller gradient error compared with the score function
gradient estimator, and we therefore implement the pathwise gradient estimator
in Equation 4. (Xu et al. 2019; Mohamed et al. 2020)

The pathwise gradient estimator is based on the principle of Monte Carlo
estimation, such that the estimation error decreases with the inverse square
root of batch size. Importantly, this estimation error is independent of
dimension. As a result, GLOnet and specifically PG-GLOnet are able to operate
for batch sizes that are independent of problem dimension, as demonstrated in
Figures 2c and 2d. This scaling of problem dimension without a required
scaling in the number of functional evaluations allows PG-GLOnet to readily
scale and address the 1000-dimensional problems in Table 1 with modest
computational resources.

Progressive growth. Direct searching within a high dimensional, non-convex
landscape is an intractable problem. In the case of FC-GLOnet, which utilizes
all of the features above, including distribution optimization and over-
parameterization, the algorithm is still not effective in directly searching
high dimensional landscapes (Table 1). With PG-GLOnet, the progressive growing
architecture regularizes the optimization procedure to search first within a
relatively coarse, low dimensional representation of the optimization
landscape, followed by relatively local searching within increasingly higher
dimensional landscape representations. This hierarchical increase of landscape
dimensionality directly corresponds to the serial toggling of $\alpha$ within
the series of growing blocks in the generator. As such, the optimization
landscape is evolved over the course of PG-GLOnet training in a manner that
maintains the tractability of the optimization problem.

To further visualize the relationship between generative network architecture
and optimization search procedure, we consider a non-convex two-dimensional
landscape shown in Figure 3b. The generative network contains a single growing
block, and the toggling of $\alpha$ from zero to one modulates the effective
dimensionality of the generator output from one to two. Initially, $\alpha$ is
zero and the vector outputted by the generator has the same effective
dimensionality as its input vector and is one. The optimization landscape
being searched is therefore a diagonal line within the two-dimensional
landscape (Figure 3b, left-most plot), and with optimal solutions near the
center of the line, the outputted generator distribution (red coloring in
plot) narrows towards this region. As $\alpha$ is increased, the generator
output vector becomes dominated by its linear transformation branch, as
opposed to its upsampling branch, and it has an effective dimensionality that
increases and eventually doubles. In our PG-GLOnet visualization, this
increase in effective dimensionality corresponds to a broadening of the
optimization landscape being searched, and the outputted generator
distribution widens relative to the diagonal line. Upon the completion of
network growth, the PG-GLOnet distribution converges to the global optimum.

The success of PG-GLOnet is therefore predicated on the ability for the
outputted distribution of the generative network to be narrowed down to
smaller but more promising regions of a coarse optimization landscape, prior
to increasing the landscape dimensionality and adding more degrees of freedom
to the problem. This concept therefore works particularly well for problems
where optima within a low dimensional analogue of the optimization landscape
help to inform of the presence and position of optima within the high
dimensional landscape. This regularization of the optimization procedure also
indicates that for problems where optima within coarse variants of the
optimization landscape do not inform the position of the global optimum, PG-
GLOnet will not work well.

In summary, we present a general global optimization algorithm metaheuristic
based on progressive growing deep generative neural networks termed PG-GLOnet.
Unlike other population-based algorithms, PG-GLOnet uses gradient-based
optimization to evolve an expressive, complex distribution in the optimization
landscape to one centered around promising optima. This complex distribution,
parameterized using the deep network framework, utilizes loss function
engineering and over-parameterization to facilitate effective gradient-based
searching. PG-GLOnet is particularly well suited to address ultra-high
dimensional problems because the required batch size is independent of problem
dimension and the progressively growing network architecture facilitates a
hierarchical search process within a landscape with progressively growing
effective dimensionality. This use of a hierarchical search strategy also
provides bounds as to the types of problems and landscapes that are suited for
PG-GLOnet optimization. We anticipate that further research in the tailoring
of application-specific generative network architectures to particular
optimization landscapes will enable the GLOnet platform to extend and adapt to
an even wider range of non-convex, high dimensional optimization problems.

## References

  * Arora, Cohen, and Hazan (2018) Arora, S.; Cohen, N.; and Hazan, E. 2018.  On the optimization of deep networks: Implicit acceleration by overparameterization.  In _International Conference on Machine Learning_ , 244–253. PMLR.
  * Balsa-Canto et al. (2012) Balsa-Canto, E.; Banga, J. R.; Egea, J. A.; Fernandez-Villaverde, A.; and de Hijas-Liste, G. 2012.  Global optimization in systems biology: stochastic methods and their applications.  In _Advances in Systems Biology_ , 409–424. Springer.
  * Chen et al. (2017) Chen, Z.; Jia, W.; Jiang, X.; Li, S.-S.; and Wang, L.-W. 2017.  SGO: A fast engine for ab initio atomic structure global optimization by differential evolution.  _Computer Physics Communications_ , 219: 35–44.
  * Choromanska et al. (2015) Choromanska, A.; Henaff, M.; Mathieu, M.; Arous, G. B.; and LeCun, Y. 2015.  The loss surfaces of multilayer networks.  In _Artificial intelligence and statistics_ , 192–204. PMLR.
  * Draxler et al. (2018) Draxler, F.; Veschgini, K.; Salmhofer, M.; and Hamprecht, F. 2018.  Essentially no barriers in neural network energy landscape.  In _International conference on machine learning_ , 1309–1318. PMLR.
  * Held et al. (2011) Held, S.; Korte, B.; Rautenbach, D.; and Vygen, J. 2011.  Combinatorial optimization in VLSI design.  _Combinatorial Optimization_ , 33–96.
  * Jiang and Fan (2019) Jiang, J.; and Fan, J. A. 2019.  Global optimization of dielectric metasurfaces using a physics-driven neural network.  _Nano letters_ , 19(8): 5366–5372.
  * Jiang and Fan (2020) Jiang, J.; and Fan, J. A. 2020.  Simulator-based training of generative neural networks for the inverse design of metasurfaces.  _Nanophotonics_ , 9(5): 1059–1069.
  * Jørgensen et al. (2018) Jørgensen, M. S.; Larsen, U. F.; Jacobsen, K. W.; and Hammer, B. 2018.  Exploration versus exploitation in global atomistic structure optimization.  _The Journal of Physical Chemistry A_ , 122(5): 1504–1509.
  * Karras et al. (2017) Karras, T.; Aila, T.; Laine, S.; and Lehtinen, J. 2017.  Progressive growing of gans for improved quality, stability, and variation.  _arXiv preprint arXiv:1710.10196_.
  * Li et al. (2019) Li, L.; Fang, W.; Wang, Q.; and Sun, J. 2019.  Differential grouping with spectral clustering for large scale global optimization.  In _2019 IEEE Congress on Evolutionary Computation (CEC)_ , 334–341. IEEE.
  * Li et al. (2013) Li, X.; Tang, K.; Omidvar, M. N.; Yang, Z.; Qin, K.; and China, H. 2013.  Benchmark functions for the CEC 2013 special session and competition on large-scale global optimization.  _gene_ , 7(33): 8.
  * Mohamed et al. (2020) Mohamed, S.; Rosca, M.; Figurnov, M.; and Mnih, A. 2020.  Monte Carlo Gradient Estimation in Machine Learning.  _J. Mach. Learn. Res._ , 21(132): 1–62.
  * Pelikan et al. (1999) Pelikan, M.; Goldberg, D. E.; Cantú-Paz, E.; et al. 1999.  BOA: The Bayesian optimization algorithm.  In _Proceedings of the genetic and evolutionary computation conference GECCO-99_ , volume 1, 525–532. Citeseer.
  * Piggott et al. (2015) Piggott, A. Y.; Lu, J.; Lagoudakis, K. G.; Petykiewicz, J.; Babinec, T. M.; and Vučković, J. 2015.  Inverse design and demonstration of a compact and broadband on-chip wavelength demultiplexer.  _Nature Photonics_ , 9(6): 374–377.
  * Sell et al. (2017) Sell, D.; Yang, J.; Doshay, S.; Yang, R.; and Fan, J. A. 2017.  Large-angle, multifunctional metagratings based on freeform multimode geometries.  _Nano letters_ , 17(6): 3752–3757.
  * Sieniutycz and Jezowski (2009) Sieniutycz, S.; and Jezowski, J. 2009.  _Energy optimization in process systems_.  Elsevier.
  * Snoek et al. (2015) Snoek, J.; Rippel, O.; Swersky, K.; Kiros, R.; Satish, N.; Sundaram, N.; Patwary, M.; Prabhat, M.; and Adams, R. 2015.  Scalable bayesian optimization using deep neural networks.  In _International conference on machine learning_ , 2171–2180. PMLR.
  * Sun et al. (2019) Sun, Y.; Li, X.; Ernst, A.; and Omidvar, M. N. 2019.  Decomposition for large-scale optimization problems with overlapping components.  In _2019 IEEE congress on evolutionary computation (CEC)_ , 326–333. IEEE.
  * Xu et al. (2019) Xu, M.; Quiroz, M.; Kohn, R.; and Sisson, S. A. 2019.  Variance reduction properties of the reparameterization trick.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , 2711–2720. PMLR.


$\|\mu_{i}-\mu_{i}^{*}\|\leq\frac{c_{r}}{4}\sqrt{\log d}$ for all $i\in[K]$,
then for any distinct $i,j\in[K]$, we have

$\displaystyle\mathinner{\\!\left\lVert\mathbb{E}_{X}[w_{i}(X,\mu)(1-w_{i}(X,\mu))(X-\mu_{i})(X-\mu_{i})^{\top}]\right\rVert}_{\mathsf{op}}$
$\displaystyle\leq\frac{d^{2}c_{r}^{2}}{d^{c_{r}^{2}/1000}}$
$\displaystyle\mathinner{\\!\left\lVert\mathbb{E}_{X}[w_{i}(X,\theta)w_{j}(x,\theta)(X-\mu_{i})(X-\mu_{j})^{\top}]\right\rVert}_{\mathsf{op}}$
$\displaystyle\leq\frac{d^{2}c_{r}^{2}}{d^{c_{r}^{2}/1000}}$

###  E.2 Closeness between population gradient descent and empirical gradient
descent

In this section, we show that the population gradient descent on the DDPM
objective is close to the empirical gradient descent for mixtures of $K$
Gaussians.

######  Lemma E.7.

For any $\varepsilon$ that is $\Theta(\frac{1}{\mathrm{poly}(d)})$ and noise
scale $t>t^{\prime}$ where $t^{\prime}\lesssim 1$, the empirical estimate of
gradient descent update on the DDPM objective with the number of samples
$n>n^{\prime}$ concentrates well to the population gradient descent update
where $n^{\prime}=O(\frac{K^{4}d^{5}B^{6}}{\varepsilon^{2}})$. More
specifically, the following inequality holds with probability at least
$1-\exp(-d^{0.99})$:

$\displaystyle\mathinner{\\!\left\lVert\nabla_{\mu_{1,t}}\Big{(}\frac{1}{n}\sum_{i=1}^{n}L_{t}(s_{\theta_{t}}(x_{i,0},z_{i,t}))\Big{)}-\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})\right\rVert}\leq\varepsilon.$

###### Proof.

Recall that the population gradient is given by

$\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})=\mathbb{E}\Big{[}\frac{1}{2}\nabla_{\mu_{1,t}}\mathinner{\\!\left\lVert
s_{\theta_{t}}(X_{t})\right\rVert}^{2}+\frac{\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}}{\beta_{t}}\Big{]}\,,$
(E.11)

where

$\displaystyle\mathbb{E}\Big{[}\frac{1}{2}\nabla_{\mu_{1,t}}\mathinner{\\!\left\lVert
s_{\theta_{t}}(X_{t})\right\rVert}^{2}\Big{]}$
$\displaystyle=\mathbb{E}\bigg{[}\Big{(}w_{1,t}(X_{t})(X_{t}-\mu_{1,t})\mu_{1,t}^{\top}+w_{1,t}(X_{t})\cdot\mathrm{Id}$
(E.12) $\displaystyle\hskip 28.45274pt-
w_{1,t}(X_{t})\sum_{i=1}^{K}w_{i,t}(X_{t})(X_{t}-\mu_{1,t})\mu_{i,t}^{\top}\Big{)}\cdot\sum_{i=1}^{K}\big{(}w_{i,t}(X_{t})\mu_{i,t}-X_{t}\big{)}\bigg{]}\,,$
(E.13)

and

$\displaystyle\mathbb{E}\Big{[}\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}\Big{]}$
$\displaystyle=\mathbb{E}\Big{[}\Big{(}w_{1,t}(X_{t})(X_{t}-\mu_{1,t})\mu_{1,t}^{\top}Z_{t}$
(E.14)
$\displaystyle\quad\quad\qquad+w_{1,t}(X_{t})Z_{t}-w_{1,t}(X_{t})\sum_{i=1}^{K}w_{i,t}(X_{t})(X_{t}-\mu_{1,t})\mu_{i,t}^{\top}Z_{t}\Big{)}\Big{]}\,.$
(E.15)

We will prove that the sample estimate of each coordinate in Eq. (E.15)
concentrates well around the expectation. We will prove the concentration of
the first coordinate and a similar analysis holds for other coordinates. For
the rest of the proof, we use $\tilde{x}_{t}$ to denote the first coordinate
of $X_{t}$ and $\tilde{\mu}_{i,t}$ to indicate the first coordinate
$\mu_{i,t}$. For any random variable $Y\in\mathbb{R}$, we use
$\|Y\|_{\psi_{1}}$ to denote the sub-exponential norm of $Y$ and
$\|Y\|_{\psi_{2}}$ to denote the sub-gaussian norm of $Y$ (See lemma B.1 for
details). Using properties of a sub-Gaussian random variable from Lemma B.1,
we get

$\displaystyle\Big{\|}\sum_{j=1}^{K}w_{1,t}(X_{t})w_{j,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\mu_{1,t}^{\top}\mu_{j,t}\Big{\|}_{\psi_{2}}$
(E.16) $\displaystyle\lesssim$
$\displaystyle\;\sum_{j=1}^{K}\Big{\|}w_{1,t}(X_{t})w_{j,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\mu_{1,t}^{\top}\mu_{j,t}\Big{\|}_{\psi_{2}}$
(Using sum of sub-Gaussian random variables property in Lemma B.1)
$\displaystyle\lesssim$
$\displaystyle\;\sum_{j=1}^{K}\Big{\|}w_{1,t}(X_{t})w_{j,t}(X_{t})\mu_{1,t}^{\top}\mu_{j,t}z\Big{\|}_{\psi_{2}}+\Big{\|}w_{1,t}(X_{t})w_{j,t}(X_{t})\mu_{1,t}^{\top}\mu_{j,t}(\tau-\tilde{\mu}_{1,t})\Big{\|}_{\psi_{2}}$
(E.17) $\displaystyle\lesssim$ $\displaystyle\;KB^{2}+KB^{3}\lesssim KB^{3},$
(E.18)

where the third inequality follows by writing $\tilde{x}_{t}=z+\tau$ where
$z\sim\mathcal{N}(0,1)$ and $\tau$ is a random variable that takes
$\tilde{\mu}_{i,t}^{*}$ for every $i\in[K]$ with probability $\frac{1}{K}$.
The fourth inequality follows from the sub-Gaussian property of a bounded
random variable and the product of a sub-Gaussian random variable with bounded
random variable property in Lemma B.1. Using the sum of sub-Gaussian random
variable property in Lemma B.1, we have

$\Big{\|}\sum_{i=1}^{K}w_{1,t}(X_{t})w_{i,t}(X)\tilde{\mu}_{i,t}\Big{\|}_{\psi_{2}}\lesssim\sum_{i=1}^{K}\|w_{1,t}(X_{t})w_{i,t}(X)\tilde{\mu}_{i,t}\|_{\psi_{2}}\lesssim
KB.$ (E.19)

Using properties of the sub-Gaussian random variable from Lemma B.1 in a
similar way of Eq. (E.18), we have

$\displaystyle\Big{\|}\sum_{i=1}^{K}\sum_{j=1}^{K}w_{1,t}(X_{t})w_{i,t}(X_{t})w_{j,t}(X_{t})\mu_{i,t}^{\top}\mu_{j,t}(\tilde{x}_{t}-\tilde{\mu}_{1,t})\Big{\|}_{\psi_{2}}$
(E.20) $\displaystyle\leq$
$\displaystyle\;\sum_{i=1}^{K}\sum_{j=1}^{K}\Big{\|}w_{1,t}(X_{t})w_{i,t}(X_{t})w_{j,t}(X_{t})\mu_{i,t}^{\top}\mu_{j,t}(\tilde{x}_{t}-\tilde{\mu}_{1,t})\Big{\|}_{\psi_{2}}$
(E.21) $\displaystyle\leq$
$\displaystyle\;\sum_{i=1}^{K}\sum_{j=1}^{K}\Big{\|}w_{1,t}(X_{t})w_{i,t}(X_{t})w_{j,t}(X_{t})\mu_{i,t}^{\top}\mu_{j,t}z\Big{\|}_{\psi_{2}}+\Big{\|}w_{1,t}(X_{t})w_{i,t}(X_{t})w_{j,t}(X_{t})\mu_{i,t}^{\top}\mu_{j,t}(\tau-\tilde{\mu}_{i,t})\Big{\|}_{\psi_{2}}$
(E.22) $\displaystyle\leq$
$\displaystyle\;K^{2}B^{2}+K^{2}B^{3}\lesssim\;K^{2}B^{3}$ (E.23)

We know that
$\|w_{1,t}(X_{t})\mu_{1,t}^{\top}X_{t}\|_{\psi_{2}}\leq\|\sum_{i=1}^{d}\mu_{1,t}(i)X_{t}(i)\|_{\psi_{2}}\lesssim
dB^{2}$ and $\|\tilde{x}_{t}-\tilde{\mu}_{1,t}\|_{\psi_{2}}\lesssim B$. Using
the fact that the product of two sub-Gaussian random variables is a sub-
exponential random variable, we have

$\displaystyle\|w_{1,t}(X_{t})\mu_{1,t}^{\top}X_{t}(\tilde{x}_{t}-\tilde{\mu}_{1,t})\|_{\psi_{1}}$
$\displaystyle\leq\|\tilde{x}_{t}-\tilde{\mu}_{1,t}\|_{\psi_{2}}\|w_{1,t}(X_{t})\mu_{1,t}^{\top}X_{t}\|_{\psi_{2}}\lesssim
dB^{3}$ (E.24)

The sub-gaussian norm of $w_{1,t}(X_{t})\tilde{x}_{t}$ term in the gradient is
given by

$\displaystyle\|w_{1,t}(X_{t})\tilde{x}_{t}\|_{\psi_{2}}\leq\|X_{t}\|_{\psi_{2}}\lesssim\|Z\|_{\psi_{2}}+\|\tau\|_{\psi_{2}}\lesssim
B$ (E.25)

Using the property that the product of two sub-Gaussian random variables is a
sub-exponential random variable, we obtain

$\displaystyle\Big{\|}w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\Big{(}\sum_{i=1}^{K}w_{i,t}(X_{t})\mu_{i,t}^{\top}X_{t}\Big{)}\Big{\|}_{\psi_{1}}$
(E.26)
$\displaystyle\lesssim\|w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\|_{\psi_{2}}\Big{\|}\Big{(}\sum_{i=1}^{K}w_{i,t}(X_{t})\mu_{i,t}^{\top}X_{t}\Big{)}\Big{\|}_{\psi_{2}}$
(E.27) $\displaystyle\lesssim KdB^{3}$ (E.28)

For any random variable $Y$, we know that
$\|X\|_{\psi_{1}}\leq\|X\|_{\psi_{2}}$. Therefore, combining Eq. (E.18),
(E.19), (E.23), (E.24), (E.25) and (E.28), we have

$\displaystyle\|[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}s_{\theta_{t}}(X_{t})]_{1}-\mathbb{E}[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}s_{\theta_{t}}(X_{t})]_{1}\|_{\psi_{1}}$
$\displaystyle\lesssim\|[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}s_{\theta_{t}}(X_{t})]_{1}\|_{\psi_{1}}$
(E.29) $\displaystyle\lesssim K^{2}dB^{3}$ (E.30)

Now, we shift our focus on obtaining the sub-exponential norm of
$\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}$. Using
$\|w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\|_{\psi_{2}}\lesssim B$ and
$\|\mu_{1,t}^{\top}Z_{t}\|_{\psi_{2}}\lesssim dB$, we obtain

$\displaystyle\|w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\mu_{1,t}^{\top}Z_{t}\|_{\psi_{1}}\leq\|w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\|_{\psi_{2}}\|\mu_{1,t}^{\top}Z_{t}\|_{\psi_{2}}\lesssim
dB^{2}$ (E.31)

Using Lemma B.1, we have
$\|w_{1,t}(X_{t})z_{t}\|_{\psi_{2}}\leq\|z_{t}\|_{\psi_{2}}\lesssim 1$. For
the last term, we have

$\displaystyle\Big{\|}w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\sum_{i=1}^{K}w_{i,t}(X_{t})\mu_{i,t}^{\top}Z_{t}\Big{\|}_{\psi_{1}}$
$\displaystyle\leq\|w_{1,t}(X_{t})(\tilde{x}_{t}-\tilde{\mu}_{1,t})\|_{\psi_{2}}\Big{\|}\sum_{i=1}^{K}w_{i,t}(X_{t})\mu_{i,t}^{\top}Z_{t}\Big{\|}_{\psi_{2}}$
(E.32) $\displaystyle\lesssim KdB^{2}$ (E.33)

Combining Eq. (E.31), (E.33), we have

$\displaystyle\Big{\|}\frac{[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}]_{1}}{\beta_{t}}-\frac{\mathbb{E}[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}]_{1}}{\beta_{t}}\Big{\|}_{\psi_{1}}\lesssim\Big{\|}\frac{[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}]_{1}}{\beta_{t}}\Big{\|}_{\psi_{1}}\lesssim\frac{KdB^{2}}{\beta_{t}},$
(E.34)

where $[\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}]_{1}$ denotes the
first coordinate of $\nabla_{\mu_{1,t}}s_{\theta_{t}}(X_{t})^{\top}Z_{t}$.
Combining Eq. (E.30) and Eq. (E.34), we have

$\displaystyle\Big{\|}[\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}}(X_{t}))]_{1}-[\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})]_{1}\Big{\|}_{\psi_{1}}\lesssim\frac{K^{2}dB^{3}}{\beta_{t}}$

For each i.i.d. sample $x_{i,t}$, the term
$[\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}}(x_{i,t}))]_{1}-[\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})]_{1}$
is also independent and identically distributed. Therefore, using Lemma B.3,
for any $\varepsilon$ that is $\Theta(\frac{1}{\text{poly}(d)})$, we have

$\displaystyle\Pr\Big{[}\Big{|}\frac{1}{n}\sum_{i=1}^{n}[\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}}(x_{i,t}))]_{1}-[\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})]_{1}\Big{|}\geq\varepsilon\Big{]}\leq
2\exp\Big{(}-\frac{n\varepsilon^{2}\beta_{t}^{2}}{K^{4}d^{2}B^{6}}\Big{)}.$

A similar analysis will give the concentration for each coordinate. Using the
union bound and rescaling $\varepsilon$ as $\frac{\varepsilon}{d}$, with
probability at least
$1-2d\exp\Big{(}-\frac{n\varepsilon^{2}\beta_{t}^{2}}{K^{4}d^{4}B^{6}}\Big{)}$,
we have

$\displaystyle\mathinner{\\!\left\lVert\nabla_{\mu_{1,t}}\Big{(}\frac{1}{n}\sum_{i=1}^{n}L_{t}(s_{\theta_{t}}(x_{i,t}))\Big{)}-\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})\right\rVert}\leq\varepsilon$

Note that for any $t=\Omega(1)$, $\beta_{t}\geq c$ for some constant $c$.
Therefore, choosing $n$ provided in the Lemma E.7 statement, we obtain the
result. ∎

###  E.3 Proof of Theorem E.1

###### Proof of Theorem E.1.

For any training iteration $h$, assume that parameters $\theta_{t}^{(h)}$ are
such that
$\mathinner{\\!\left\lVert\mu_{i,t}^{(h)}-\mu_{i,t}^{*}\right\rVert}\leq\frac{c_{r}}{4}\sqrt{\log
d}$ we can write the update on the DDPM objective as follows:

$\displaystyle\|\mu_{1,t}^{(h+1)}-\mu_{1,t}^{*}\|=$
$\displaystyle\;\Big{\|}\mu_{1,t}^{(h)}-\eta\nabla\Big{(}\frac{1}{n}\sum_{i=1}^{n}L_{t}(s_{\theta_{t}^{(h)}}(x_{i,0},z_{i,t}))\Big{)}-\mu_{1,t}^{*}\Big{\|}$
$\displaystyle\leq$
$\displaystyle\;\big{\|}\mu_{1,t}^{(h)}+\eta\,\mathbb{E}[w_{1,t}(X_{t})(X_{t}-\mu_{1,t}^{(h)})]-\mu_{1,t}^{*}\big{\|}$
$\displaystyle\quad\quad+\eta\Big{\|}\left(-\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})\right)-\mathbb{E}[w_{1,t}(X_{t})(X_{t}-\mu_{1,t}^{(h)})]\Big{\|}$
$\displaystyle\quad\quad+\eta\Big{\|}\left(\nabla_{\mu_{1,t}}L_{t}(s_{\theta_{t}})\right)-\nabla_{\mu_{1,t}}\Big{(}\frac{1}{n}\sum_{i=1}^{n}L_{t}(s_{\theta_{t}^{(h)}}(x_{i,0},z_{i,t}))\Big{)}\Big{\|}\,.$

Using Lemma E.4, Lemma E.7 and Theorem 3.2 from [SN21], for any
$\eta\in(0,K)$, we have

$\displaystyle\|\mu_{1,t}^{(h+1)}-\mu_{1,t}^{*}\|\leq$
$\displaystyle\;\left(1-\frac{3\eta}{8K}\right)\|\mu_{1,t}^{(h)}-\mu_{1,t}^{*}\|+\frac{\eta
K^{2}B^{2}}{d^{\frac{c_{r}^{2}}{4000}}}+\eta\varepsilon.$

Choosing $\eta=\frac{2K}{3}$, $c_{r}$ to be sufficiently large constant and
$\varepsilon$ to be $\Theta(\frac{1}{\text{poly}(d)})$, we have

$\displaystyle\|\mu_{1,t}^{(h+1)}-\mu_{1,t}^{*}\|\leq\frac{3}{4}\|\mu_{1,t}^{(h)}-\mu_{1,t}^{*}\|+\varepsilon$

By assumption 15, $\|\mu_{1,t}^{(0)}-\mu_{1,t}^{*}\|\leq O(\sqrt{\log d})$ and
therefore, choosing $H$ to be $\Omega(\log(\frac{\log d}{\varepsilon}))$, we
obtain the result. ∎

##  Appendix F Additional proofs

###  F.1 Proof of Lemma C.2

###### Proof of Lemma C.2.

By calculating the negative gradient of the DDPM objective in Eq. (13), we
obtain

$\displaystyle-\nabla_{\mu_{t}}L_{t}(s_{\mu_{t}})$
$\displaystyle=-\mathbb{E}_{X_{0},Z_{t}}[(\tanh(\mu_{t}^{\top}X_{t})I+\tanh^{\prime}(\mu_{t}^{\top}X_{t})X_{t}\mu_{t}^{\top})(s_{\mu_{t}}(X_{t})+\frac{Z_{t}}{\beta_{t}})]$
(F.1)
$\displaystyle=-\mathbb{E}[(\tanh(\mu_{t}^{\top}X_{t})I+\tanh^{\prime}(\mu_{t}^{\top}X_{t})X_{t}\mu_{t}^{\top})(\tanh(\mu_{t}^{\top}X_{t})\mu_{t}-X_{t}+\frac{Z_{t}}{\beta_{t}})]$
$\displaystyle=\mathbb{E}[-\tanh^{2}(\mu_{t}^{\top}X_{t})\mu_{t}-\tanh(\mu_{t}^{\top}X_{t})\tanh^{\prime}(\mu_{t}^{\top}X_{t})X_{t}\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}+\tanh(\mu_{t}^{\top}X_{t})X_{t}$
$\displaystyle\quad+\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}X_{t}X_{t}-\tanh(\mu_{t}^{\top}X_{t})\frac{Z_{t}}{\beta_{t}}-\tanh^{\prime}(\mu_{t}^{\top}X_{t})X_{t}\mu_{t}^{\top}\frac{Z_{t}}{\beta_{t}}]$

By simplifying the gradient terms involving $Z_{t}$ by the Stein’s identity as
in Lemma F.1 and plugging it back in the gradient, we obtain

$\displaystyle-\nabla_{\mu_{t}}L_{t}(s_{\mu_{t}})$
$\displaystyle=\mathbb{E}\Big{[}\left(\tanh(\mu_{t}^{\top}X_{t})-\tanh(\mu_{t}^{\top}X_{t})\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}+\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}X_{t}\right)X_{t}\Big{]}$
$\displaystyle\quad-\mu_{t}-\mathbb{E}\mathinner{\left[\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}X_{t}\right]}-\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}$
$\displaystyle=\mathbb{E}\Big{[}\left(\tanh(\mu_{t}^{\top}X_{t})-0.5\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}+\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}X_{t}\right)X_{t}\Big{]}$
$\displaystyle\quad-\mu_{t}-\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}$

Observe that
$\left(\tanh(\mu^{\top}x)-\frac{1}{2}\tanh^{\prime\prime}(\mu^{\top}x)\mathinner{\\!\left\lVert\mu\right\rVert}^{2}+\tanh^{\prime}(\mu^{\top}x)\mu^{\top}x\right)x$
and $\tanh^{\prime}(\mu^{\top}x)$ are even functions and $X_{t}$ is a
symmetric distribution, therefore, for any even function $f$, we can write
$\mathbb{E}_{X_{t}}[f(X_{t})]=\frac{1}{2}\mathbb{E}_{X_{t}\sim\mathcal{N}(\mu_{t}^{*},\mathrm{Id})}[f(X_{t})]+\frac{1}{2}\mathbb{E}_{X_{t}\sim\mathcal{N}(-\mu_{t}^{*},I)}[f(X_{t})]=\mathbb{E}_{X_{t}\sim\mathcal{N}(\mu_{t}^{*},\mathrm{Id})}[f(X_{t})]$.
Applying this property of the even function on the gradient update, we obtain
the result. ∎

######  Lemma F.1.

When random variable $X_{t}=\alpha_{t}X_{0}+\beta_{t}Z_{t}$ where
$Z_{t}\sim\mathcal{N}(0,I),\alpha_{t}=\exp(-t)$ and
$\beta_{t}=\sqrt{1-\exp(-2t)}$, then for any $t>0$, the following two
equations hold.

$\displaystyle\mathbb{E}_{X_{0},Z_{t}}$
$\displaystyle\Big{[}\tanh(\mu_{t}^{\top}X_{t})\frac{Z_{t}}{\beta_{t}}+\tanh^{2}(\mu_{t}^{\top}X_{t})\mu_{t}\Big{]}=\mu_{t}$
$\displaystyle\mathbb{E}_{X_{0},Z_{t}}$
$\displaystyle\Big{[}\tanh^{\prime}(\mu_{t}^{\top}X_{t})\frac{\mu_{t}^{\top}Z_{t}}{\beta_{t}}X_{t}\Big{]}=\mathbb{E}_{X_{0},Z_{t}}\mathinner{\left[\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}X_{t}+\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}$

###### Proof.

Applying Stein’s lemma on the first term, we get the first equation of the
statement in the Lemma.

$\displaystyle\mathbb{E}_{X_{0},Z_{t}}\mathinner{\left[\tanh(\mu_{t}^{\top}X_{t})\frac{Z_{t}}{\beta_{t}}\right]}$
$\displaystyle=\mathbb{E}_{X_{0},Z_{t}}\mathinner{\left[\tanh(\mu_{t}^{\top}(\alpha_{t}X_{0}+\beta_{t}Z_{t}))\frac{Z_{t}}{\beta_{t}}\right]}=\mathbb{E}_{X_{0},Z_{t}}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}$
$\displaystyle=\mathbb{E}_{X_{0},Z_{t}}\mathinner{\left[\left(1-\tanh^{2}(\mu_{t}^{\top}X_{t})\right)\mu_{t}\right]}$

For the second term, we have

$\displaystyle\mathbb{E}$
$\displaystyle\Big{[}\tanh^{\prime}(\mu_{t}^{\top}X_{t})\frac{\mu_{t}^{\top}Z_{t}}{\beta_{t}}X_{t}\Big{]}=\mathbb{E}\Big{[}\tanh^{\prime}(\mu_{t}^{\top}X_{t})\frac{\mu_{t}^{\top}Z_{t}}{\beta_{t}}\alpha_{t}X_{0}\Big{]}+\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}Z_{t}Z_{t}\right]}$
$\displaystyle=\sum_{i=1}^{d}\mathbb{E}\Big{[}\alpha_{t}X_{0}\tanh^{\prime}(\mu_{t}^{\top}X_{t})\frac{\mu_{t}(i)Z_{t}(i)}{\beta_{t}}\Big{]}+\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}+\mathbb{E}\mathinner{\left[\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}Z_{t}\beta_{t}\mu_{t}\right]}$
$\displaystyle=\sum_{i=1}^{d}\mathbb{E}\Big{[}\alpha_{t}X_{0}\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mu_{t}(i)\mu_{t}(i)\Big{]}+\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}+\mathbb{E}\mathinner{\left[\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}Z_{t}\beta_{t}\mu_{t}\right]}$

where the second equality follows from the Stein’s lemma on the
$\mathbb{E}[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}Z_{t}Z_{t}]$ and
the last equality follows from the Stein’s lemma on
$\mathbb{E}[\alpha_{t}X_{0}\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mu_{t}(i)Z_{t}(i)]$.
Applying Stein’s inequality on the
$\mathbb{E}\mathinner{\left[\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mu_{t}^{\top}Z_{t}\beta_{t}\mu_{t}\right]}$,
we obtain

$\displaystyle=\mathbb{E}\mathinner{\left[\alpha_{t}X_{0}\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}\right]}+\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}+\sum_{i=1}^{d}\beta_{t}\mu_{t}\mathbb{E}\mathinner{\left[\tanh^{\prime\prime\prime}(\mu_{t}^{\top}X_{t})\mu_{t}(i)\beta_{t}\mu_{t}(i)\right]}$
$\displaystyle=\mathbb{E}\mathinner{\left[X_{t}\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}\right]}-\mathbb{E}\mathinner{\left[\beta_{t}Z_{t}\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}\right]}+\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}$
$\displaystyle\quad+\beta_{t}^{2}\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}\mu_{t}\mathbb{E}\mathinner{\left[\tanh^{\prime\prime\prime}(\mu_{t}^{\top}X_{t})\right]}$
$\displaystyle=\mathbb{E}\mathinner{\left[X_{t}\tanh^{\prime\prime}(\mu_{t}^{\top}X_{t})\mathinner{\\!\left\lVert\mu_{t}\right\rVert}^{2}\right]}+\mathbb{E}\mathinner{\left[\tanh^{\prime}(\mu_{t}^{\top}X_{t})\mu_{t}\right]}.$

∎

###  F.2 Proof of Lemma C.8

###### Proof of Lemma C.8.

Recall that the gradient update for any $\mu_{t}^{*}$ is given by

$\displaystyle-\nabla_{\mu^{*}_{t}}L_{t}(s_{\mu^{*}_{t}})$
$\displaystyle=G(\mu_{t}^{*},\mu_{t}^{*})+\eta\mathbb{E}_{x\sim\mathcal{N}(\mu_{t}^{*},\mathrm{Id})}[\tanh(\mu^{*\top}_{t}x)x]-\eta\mu_{t}^{*}$
(F.2)

We know that
$\mathbb{E}_{x\sim\mathcal{N}(\mu_{t}^{*},\mathrm{Id})}[\tanh(\mu^{*\top}_{t}x)x]=\mu_{t}^{*}$
(Eq.(2.1) of [DTZ17]) and $\nabla_{\mu^{*}_{t}}L_{t}(s_{\mu^{*}_{t}})=0$
because $\mu^{*}_{t}$ is a stationary point of the regression objective of
diffusion model. This implies that $G(\mu^{*}_{t},\mu^{*}_{t})=0$ for any
$\mu^{*}_{t}$.

Note that this proof only talks about 1D case therefore, for the purpose of
this proof, we use $a$ to denote $\mu$ and $b$ to denote $\mu^{*}$. In 1D,
using Mean value theorem, we have

$\displaystyle\frac{G(a,b)-G(a,a)}{b-a}$
$\displaystyle=\frac{dG(a,\xi)}{d\xi}\text{ for some $\xi\in[a,b]$ (if
$a<b$)}$ (F.3)

Using the fact that $G(a,a)=0$ in Eq. (F.3), we have

$\displaystyle\mathinner{\\!\left\lvert G(a,b)\right\rvert}$
$\displaystyle=\mathinner{\\!\left\lvert\frac{dG(a,\xi)}{d\xi}\right\rvert}\mathinner{\\!\left\lvert
b-a\right\rvert}$

Observe that it suffices to prove
$\mathinner{\\!\left\lvert\frac{dG(a,\xi)}{d\xi}\right\rvert}\leq 0.01$ to
obtain the lemma. By computing the gradient of $G$, we obtain

$\displaystyle\frac{dG(a,\xi)}{d\xi}=\eta\mathbb{E}_{x\sim\mathcal{N}(\xi,1)}\Big{[}2\tanh^{\prime}(ax)ax+\tanh^{\prime\prime}(ax)\left(\frac{-3a^{2}}{2}+a^{2}x^{2}\right)-\frac{1}{2}a^{3}x\tanh^{\prime\prime\prime}(ax)\Big{]}$

For the first term, we have

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(\xi,I)}[\tanh^{\prime}(ax)ax]$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\int_{-\infty}^{\infty}\tanh^{\prime}(ax)axe^{-\frac{(x-\xi)^{2}}{2}}dx$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\int_{0}^{\infty}\tanh^{\prime}(ax)ax\left(e^{-\frac{(x-\xi)^{2}}{2}}-e^{-\frac{(x+\xi)^{2}}{2}}\right)dx$
$\displaystyle\leq\frac{1}{\sqrt{2\pi}}\int_{0}^{\infty}e^{-ax}axe^{-\frac{(x-\xi)^{2}}{2}}dx$
$\displaystyle\leq\frac{ae^{\frac{a^{2}-2a\xi}{2}}}{\sqrt{2\pi}}\int_{0}^{\infty}xe^{-\frac{(x-\xi+a)^{2}}{2}}dx$
$\displaystyle\leq
ae^{\frac{a^{2}-2a\xi}{2}}(\sqrt{\frac{2}{\pi}}e^{-\frac{(\xi-a)^{2}}{2}}+(\xi-a)\text{erf}\left(\frac{\xi-a}{\sqrt{2}}\right))$
$\displaystyle\leq
ae^{-\frac{\xi^{2}}{2}}+a\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-2a(\xi-a)-a^{2}}{2}}$

Using Lemma 1 of [DTZ17], we know that
$\mathbb{E}_{x\sim\mathcal{N}(\xi,I)}[\tanh^{\prime}(ax)ax]>0$. Therefore, we
have

$\displaystyle\mathinner{\\!\left\lvert\mathbb{E}_{x\sim\mathcal{N}(\xi,I)}[\tanh^{\prime}(ax)ax]\right\rvert}\leq
ae^{-\frac{\xi^{2}}{2}}+a\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-2a(\xi-a)-a^{2}}{2}}$
(F.4)

For the second term, we have

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(\xi,1)}[\tanh^{\prime\prime}(ax)(-\frac{3a^{2}}{2}+a^{2}x^{2})]$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\int_{0}^{\infty}a^{2}\tanh^{\prime\prime}(ax)(-\frac{3}{2}+x^{2})\left(\exp(-\frac{(x-\xi)^{2}}{2})-\exp(-\frac{(x+\xi)^{2}}{2})\right)dx$
$\displaystyle\leq\frac{1}{\sqrt{2\pi}}\int_{0}^{\sqrt{\frac{3}{2}}}a^{2}e^{-2ax}(\frac{3}{2}-x^{2})\exp(-\frac{(x-\xi)^{2}}{2})dx$
$\displaystyle\leq\frac{3}{\sqrt{2\pi}}a^{2}\exp(-\frac{a^{2}}{16})$

Assuming $a\geq\sqrt{6}$, then when $\xi\geq a\geq\sqrt{6}$, we have
$\exp(-\frac{(x-\xi)^{2}}{2})\leq\exp(-\frac{a^{2}}{4})$ and when $\xi\leq a$,
using $\xi\geq\frac{3a}{4}$, we have
$\exp(-\frac{(x-\xi)^{2}}{2})\leq\exp(-\frac{a^{2}}{16})$. For the lower
bound, we have

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(\xi,1)}[\tanh^{\prime\prime}(ax)(-\frac{3a^{2}}{2}+a^{2}x^{2})]$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\int_{0}^{\infty}\tanh^{\prime\prime}(ax)(-\frac{3a^{2}}{2}+a^{2}x^{2})\left(\exp(-\frac{(x-\xi)^{2}}{2})-\exp(-\frac{(x+\xi)^{2}}{2})\right)dx$
$\displaystyle\geq\frac{1}{\sqrt{2\pi}}\int_{\sqrt{\frac{3}{2}}}^{\infty}\tanh^{\prime\prime}(ax)(-\frac{3a^{2}}{2}+a^{2}x^{2})\left(\exp(-\frac{(x-\xi)^{2}}{2})-\exp(-\frac{(x+\xi)^{2}}{2})\right)dx$
$\displaystyle\geq\frac{1}{\sqrt{2\pi}}\int_{\sqrt{\frac{3}{2}}}^{\infty}\tanh^{\prime\prime}(ax)a^{2}x^{2}\left(\exp(-\frac{(x-\xi)^{2}}{2})-\exp(-\frac{(x+\xi)^{2}}{2})\right)dx$
$\displaystyle\geq-\frac{8a^{2}}{\sqrt{2\pi}}\int_{\sqrt{\frac{3}{2}}}^{\infty}e^{-2ax}x^{2}\left(\exp(-\frac{(x-\xi)^{2}}{2})-\exp(-\frac{(x+\xi)^{2}}{2})\right)dx$
$\displaystyle\geq-\frac{8a^{2}e^{-\sqrt{6}a}}{\sqrt{2\pi}}\int_{\sqrt{\frac{3}{2}}}^{\infty}x^{2}\exp(-\frac{(x-\xi)^{2}}{2})dx\geq-8a^{2}e^{-\sqrt{6}a}$

Using upper bound and lower bound, we have

$\displaystyle\mathinner{\\!\left\lvert\mathbb{E}_{x\sim\mathcal{N}(\xi,1)}[\tanh^{\prime\prime}(ax)a^{2}(-\frac{3}{2}+x^{2})]\right\rvert}\leq
8a^{2}e^{-\sqrt{6}a}$

For the third term, we have

$\displaystyle\Big{|}\mathbb{E}_{x\sim\mathcal{N}(\xi,1)}[\frac{a^{3}x}{2}\tanh^{\prime\prime\prime}(ax)]\Big{|}$
$\displaystyle=\bigg{|}\frac{1}{32\sqrt{2\pi}}\int_{0}^{\infty}a^{3}x\sigma(2ax)(1-\sigma(2ax))\left(1-6\sigma(2ax)(1-\sigma(2ax))\right)\bigg{(}\exp\bigg{(}-\frac{(x-\xi)^{2}}{2}\bigg{)}$
$\displaystyle\quad\quad-\exp\bigg{(}-\frac{(x+\xi)^{2}}{2}\bigg{)}\bigg{)}dx\bigg{|}$
$\displaystyle\leq\bigg{|}\frac{3a^{3}}{16\sqrt{2\pi}}\int_{0}^{\infty}x\sigma^{2}(2ax)(1-\sigma(2ax))^{2}\bigg{(}\exp\bigg{(}-\frac{(x-\xi)^{2}}{2}\bigg{)}-\exp\bigg{(}-\frac{(x+\xi)^{2}}{2}\bigg{)}\bigg{)}dx\bigg{|}$
$\displaystyle\leq\frac{3a^{3}}{16\sqrt{2\pi}}\int_{0}^{\infty}xe^{-ax}\exp\bigg{(}-\frac{(x-\xi)^{2}}{2}\bigg{)}dx$
$\displaystyle\leq\frac{a^{3}}{10}e^{-\frac{\xi^{2}}{2}}+\frac{a^{3}}{10}\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-2a(\xi-a)-a^{2}}{2}}\,.$

We can lower bound the third term as follows:

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(\xi,1)}[\frac{a^{3}x}{2}\tanh^{\prime\prime\prime}(ax)]$
$\displaystyle\geq\frac{1}{2\sqrt{2\pi}}\int_{0}^{c}a^{3}x\tanh^{\prime\prime\prime}(ax)\bigg{(}\exp\bigg{(}-\frac{(x+\xi)^{2}}{2}\bigg{)}-\exp\bigg{(}-\frac{(x-\xi)^{2}}{2}\bigg{)}\bigg{)}dx$
$\displaystyle\geq\frac{a^{3}}{2\sqrt{2\pi}}\int_{0}^{c}x\exp\bigg{(}-\frac{(x-\xi)^{2}}{2}\bigg{)}\left(\exp\left(-2\xi
x\right)-1\right)dx$
$\displaystyle\geq-\frac{a^{3}\xi}{\sqrt{2\pi}}\int_{0}^{c}x^{2}\exp\bigg{(}-\frac{(x-\xi)^{2}}{2}\bigg{)}dx\geq-\frac{\xi\exp(-\frac{\xi^{2}}{4})}{\sqrt{2\pi}}$

Using all the bounds, we have

$\displaystyle\mathinner{\\!\left\lvert\frac{dG(a,\xi)}{d\xi}\right\rvert}$
$\displaystyle\leq\frac{a^{3}}{10}e^{-\frac{\xi^{2}}{2}}+\frac{a^{3}}{10}\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-2a(\xi-a)-a^{2}}{2}}+8a^{2}e^{-\sqrt{6}a}+ae^{-\frac{\xi^{2}}{2}}+a\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-2a(\xi-a)-a^{2}}{2}}$

When $\xi\geq a$ and $a\geq c$ for some sufficiently large constant $c$ (for
example, $c=25$), then, we have

$\displaystyle\mathinner{\\!\left\lvert\frac{dG(a,\xi)}{d\xi}\right\rvert}$
$\displaystyle\leq\frac{a^{3}}{10}e^{-\frac{a^{2}}{2}}+\frac{a^{3}}{10}\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-a^{2}}{2}}+8a^{2}e^{-\sqrt{6}a}+ae^{-\frac{a^{2}}{2}}+a\mathinner{\\!\left\lvert\xi-a\right\rvert}e^{\frac{-a^{2}}{2}}\leq
0.01$

When $\frac{3a}{4}\leq\xi\leq a$ and $a>c$ for sufficiently large constant $c$
(for example, $c=25$), we have

$\displaystyle\mathinner{\\!\left\lvert\frac{dG(a,\xi)}{d\xi}\right\rvert}$
$\displaystyle\leq\frac{a^{3}}{10}e^{-\frac{9a^{2}}{32}}+\frac{a^{4}}{40}e^{\frac{-a^{2}}{4}}+8a^{2}e^{-\sqrt{6}a}+ae^{-\frac{a^{2}}{2}}+\frac{a^{2}}{4}e^{\frac{-a^{2}}{4}}\leq
0.01$

Pluggint the bound on $|\frac{dG(a,\xi)}{d\xi}|$ in Eq. (F.2), we obtain the
final result. ∎

###  F.3 Proof of Lemma C.10

###### Proof of Lemma C.10.

We will prove this by induction. For $h=0$, this is true because the algorithm
initializes the gradient descent on the low noise regime with the output of
gradient descent on the high noise regime, and the output is guaranteed to
have $\langle\hat{\mu}_{t}^{(0)},\hat{\mu}_{t}^{*}\rangle$ to be $\Omega(1)$
and by assumption $\|\mu_{t}^{*}\|>c^{\prime}$, therefore
$\|\mu_{t}^{(0)}\|\in[c,\frac{4\langle\hat{\mu}_{t}^{(0)},\mu_{t}^{*}\rangle}{3}]$.

Suppose
$\|\mu_{t}^{(h)}\|\in[c,\frac{4\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle}{3}]$,
then we know that
$\|\mu_{t}^{(h+1)}-\mu_{t}^{*}\|<\|\mu_{t}^{(h)}-\mu_{t}^{*}\|$. To prove
$\|\mu_{t}^{(h+1)}\|\in[c,\frac{4\langle\hat{\mu}_{t}^{(h+1)},\mu_{t}^{*}\rangle}{3}]$,
first we will prove that
$\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(r+1)}\rangle\in[c,\frac{6\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle}{5}]$.
Note that the update in the direction of $\langle\hat{\mu}_{t},\mu_{t}\rangle$
works like 1D. Therefore, we have a contraction for it as follows.

$\displaystyle\mathinner{\\!\left\lvert\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle-\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle\right\rvert}<\mathinner{\\!\left\lvert\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h)}\rangle-\langle\hat{\mu}_{t},\mu_{t}^{*}\rangle\right\rvert}$

If $\|\mu_{t}^{(h)}\|\leq\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle$, then
using Lemma F.4, we know
$\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle\leq\frac{6\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle}{5}$
and
$\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle\geq\|\mu_{t}^{(h)}\|\geq c$
because of the contraction. If
$\|\mu_{t}^{(h)}\|\geq\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle$ and
$\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle\geq\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle$,
then $\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle\leq\|\mu_{t}^{(h)}\|$
because of the contraction. If
$\|\mu_{t}^{(h)}\|\geq\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle$ and
$\langle\hat{\mu}_{t}^{(h+1)},\mu_{t}^{(h)}\rangle\leq\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle$,
then using
$\langle\hat{\mu}_{t}^{(h+1)},\mu_{t}^{(h)}\rangle\geq\|\mu_{t}^{(h)}\|-\mathinner{\\!\left\lvert
U(\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h)}\rangle,\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle)\right\rvert}\geq\frac{4\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle}{5}\geq\frac{4\langle\hat{\mu}_{t}^{(0)},\mu_{t}^{*}\rangle}{5}\geq
c$ from Lemma F.2, we get the result that
$\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle\in[c,\frac{6\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle}{5}]$.
Now, using Lemma F.2, we get

$\displaystyle\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{(h+1)}\rangle\in[c,\frac{6\langle\hat{\mu}_{t}^{(h)},\mu_{t}^{*}\rangle}{5}]$
$\displaystyle\implies\|\mu_{t}^{(h+1)}\|\in\big{[}\frac{c}{\cos\alpha_{h}},\frac{6\|\mu_{t}^{*}\|\cos\beta_{h}}{5\cos\alpha_{h}}\big{]}$
$\displaystyle\implies\|\mu_{t}^{(h+1)}\|\in\big{[}c,\frac{4\|\mu_{t}^{*}\|\cos\beta_{h+1}}{3}\big{]}$
$\displaystyle\implies\|\mu_{t}^{(h+1)}\|\in\big{[}c,\frac{4\langle\hat{\mu}_{t}^{(h+1)},\mu_{t}^{*}\rangle}{3}\big{]}\qed$

######  Lemma F.2.

Suppose the angle between $\mu^{(r)}$ and $\mu^{*}$ is $\beta_{r}$ and
$\alpha_{r}$ is the angle between $\mu^{(r)}$ and $\mu^{(r+1)}$ and assume the
contraction is true at time $r$. Assume that $\beta_{0}\in(0,\frac{\pi}{2})$.
Then:

$\alpha_{r}\in(0,\pi/2)\;\;\forall
r\qquad\text{and}\qquad\cos\beta_{r}\leq\cos\beta_{r+1}$ (F.5)

which implies that

$\displaystyle\cos\beta_{r}\leq\cos\beta_{r+1}\;\;\forall
r\implies\langle\hat{\mu}^{(r)},\mu^{*}\rangle\geq\langle\hat{\mu}^{(0)},\mu^{*}\rangle$

###### Proof.

First, we will prove that if $\beta_{r}\in(0,\frac{\pi}{2})$ and
$\|\mu^{(r)}\|\in[c,\frac{4\langle\hat{\mu}_{t}^{(r)},\mu_{t}^{*}\rangle}{3}]$,
then $\alpha_{r}\in(0,\beta_{r})$ for any $r$. We denote $\alpha_{r}>0$ if
$\mu^{(r)}$ moves towards $\mu^{(r)\perp}$ and hence towards $\mu^{*}$. The
following simple observation of
$\langle\hat{\mu}^{(r)\perp},\mu^{(r+1)}\rangle\geq 0$ proves that
$\alpha_{r}>0$.

$\displaystyle\langle\hat{\mu}^{(r)\perp},\mu^{(r+1)}\rangle$
$\displaystyle=\mathbb{E}_{x\sim\mathcal{N}(\mu^{*},1)}\Big{[}\eta\big{(}\tanh(\mu^{(r)\top}x)-\frac{1}{2}\tanh^{\prime\prime}(\mu^{(r)\top}x)\|\mu^{(r)}\|^{2}+\tanh^{\prime}(\mu^{(r)\top}x)\mu^{(r)\top}x\big{)}\cdot\langle\hat{\mu}^{(r)\perp},x\rangle\Big{]}$
$\displaystyle=\mathbb{E}_{x\sim\mathcal{N}(0,1)}\Big{[}\eta\Big{(}\tanh(\mu^{(r)\top}(x+\mu^{*}))-\frac{1}{2}\tanh^{\prime\prime}(\mu^{(r)\top}(x+\mu^{*}))\|\mu^{(r)}\|^{2}$
$\displaystyle\quad\quad\quad+\tanh^{\prime}(\mu^{(r)\top}(x+\mu^{*}))\mu^{(r)\top}(x+\mu^{*})\Big{)}\cdot\langle\hat{\mu}^{(r)\perp},(x+\mu^{*})\rangle\Big{]}$
$\displaystyle=\mathbb{E}_{\alpha_{1},\alpha_{2}\sim\mathcal{N}(\langle\hat{\mu}^{(r)},\mu^{*}\rangle,1)}\Big{[}\eta\Big{(}\tanh(\|\mu^{(r)}\|\alpha_{1})-\frac{1}{2}\tanh^{\prime\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|^{2}$
$\displaystyle\quad\quad\quad+\tanh^{\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|\alpha_{1}\Big{)}(\alpha_{2}+\langle\hat{\mu}^{(r)\perp},\mu^{*}\rangle)\Big{]}$
$\displaystyle=\mathbb{E}_{\alpha_{1},\alpha_{2}\sim\mathcal{N}(\langle\hat{\mu}^{(r)},\mu^{*}\rangle,1)}\Big{[}\eta\Big{(}\tanh(\|\mu^{(r)}\|\alpha_{1})-\frac{1}{2}\tanh^{\prime\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|^{2}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad+\tanh^{\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|\alpha_{1}\Big{)}\cdot\langle\hat{\mu}^{(r)\perp},\mu^{*}\rangle\Big{]}>0\,,$

where in the last step we used the fact that
$\langle\hat{\mu}^{(r)},\mu^{*}\rangle>0$ and
$\langle\hat{\mu}^{(r)\perp},\mu^{*}\rangle>0$.

Now, we will prove that $\cot\alpha_{r}>\cot\beta_{r}$ which will prove that
$\alpha_{r}\in(0,\beta_{r})$. Note that

$\displaystyle\cot\alpha_{r}$
$\displaystyle=\frac{\langle\hat{\mu}^{(r)},\mu^{(r+1)}\rangle}{\langle\hat{\mu}^{(r)\perp},\mu^{(r+1)}\rangle}\hskip
28.45274pt\text{where}$
$\displaystyle\langle\hat{\mu}^{(r)},\mu^{(r+1)}\rangle=$
$\displaystyle\;(1-\eta)\|\mu^{(r)}\|+\eta\mathbb{E}_{\alpha_{1}\sim\mathcal{N}(\hat{\mu}^{(r)\top}\mu^{*},1)}[\tanh(\|\mu^{(r)}\|\alpha_{1})\alpha_{1}]$
$\displaystyle+\eta\mathbb{E}_{\alpha_{1}\sim\mathcal{N}(\hat{\mu}^{(r)\top}\mu^{*},1)}[-\frac{1}{2}\tanh^{\prime\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|^{2}\alpha_{1}+\tanh^{\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|\alpha_{1}^{2}$
$\displaystyle-\tanh^{\prime}(\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|]$
$\displaystyle\langle\hat{\mu}^{(r)\perp},\mu^{(r+1)}\rangle=$
$\displaystyle\;\eta\langle\hat{\mu}^{(r)\perp},\mu^{*}\rangle\mathbb{E}_{\alpha_{1}\sim\mathcal{N}(\hat{\mu}^{(r)\top}\mu^{*},1)}[\tanh(\|\mu^{(r)}\|\alpha_{1})-\frac{1}{2}\tanh^{\prime\prime}(\;\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|^{2}$
$\displaystyle+\tanh^{\prime}(\;\|\mu^{(r)}\|\alpha_{1})\|\mu^{(r)}\|\alpha_{1}]$
and
$\displaystyle\cot\beta_{r}=\frac{\langle\hat{\mu}^{(r)},\mu^{*}\rangle}{\langle\hat{\mu}^{(r)\perp},\mu^{*}\rangle}$

Observe the fact that to prove $\frac{a+c^{\prime}}{b+c}-\frac{a}{b}>0$, it is
sufficient to prove $c^{\prime}>\frac{ac}{b}$ for $b,c>0$. Using this
observation, to prove $\cot\alpha_{r}>\cot\beta_{r}$, it is sufficient to
prove

$\displaystyle\Big{(}1-\eta-\eta\mathbb{E}[\tanh^{\prime}(\|\mu^{(r)}\|x)]\Big{)}\|\mu^{(r)}\|+\eta\mathbb{E}_{x}\Big{[}-\frac{1}{2}\tanh^{\prime\prime}(\|\mu^{(r)}\|x)\|\mu^{(r)}\|^{2}(x-\langle\hat{\mu}^{(r)},\mu^{*}\rangle)$
$\displaystyle+\tanh^{\prime}(\|\mu^{(r)}\|x)(x^{2}-\langle\hat{\mu}^{(r)},\mu^{*}\rangle
x)+\tanh(\|\mu^{(r)}\|x)(x-\langle\hat{\mu}^{(r)},\mu^{*}\rangle)\Big{]}>0,$

where the expectation is wrt $\mathcal{N}(\langle\mu^{(r)},\mu^{*}\rangle,1)$.
Lemma F.3 shows that this is indeed true. ∎

######  Lemma F.3.

For any $\eta=\frac{1}{20}$, assuming $a\in[30,\frac{4b}{3}]$, we have

$\displaystyle(1-\eta-\eta\mathbb{E}_{x\sim\mathcal{N}(b,1)}[\tanh^{\prime}(ax)])a$
$\displaystyle+\eta\,\mathbb{E}_{x\sim\mathcal{N}(b,1)}\Big{[}-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}(x-b)\tanh^{\prime}(ax)(x^{2}-bx)+\tanh(ax)(x-b)\Big{]}>0\,.$

###### Proof.

First, we will find the upper bound on
$\mathbb{E}[\tanh^{\prime\prime}(ax)(x-b)]$.

$\displaystyle\mathbb{E}[\tanh^{\prime\prime}(ax)(x-b)]$
$\displaystyle=\int_{-\infty}^{\infty}\tanh^{\prime\prime}(ax)(x-b)\exp\Big{(}-\frac{(x-b)^{2}}{2}\Big{)}dx$
$\displaystyle\leq\int_{0}^{b}\tanh^{\prime\prime}(ax)(x-b)\exp\Big{(}-\frac{(x-b)^{2}}{2}\Big{)}dx$
$\displaystyle\leq\int_{0}^{b}\tanh^{\prime\prime}(ax)x\exp\Big{(}-\frac{(x-b)^{2}}{2}\Big{)}dx$
$\displaystyle\leq\int_{0}^{b}\exp(-ax)x\exp\Big{(}-\frac{(x-b)^{2}}{2}\Big{)}dx$
$\displaystyle\leq\exp\Big{(}\frac{a^{2}-2ab}{2}\Big{)}\int_{0}^{b}x\exp\Big{(}-\frac{(x-b)^{2}+2a(x-b)+a^{2}}{2}\Big{)}dx$
$\displaystyle\leq\exp(\frac{a^{2}-2ab}{2})\int_{0}^{\infty}x\Big{[}\exp\Big{(}-\frac{(x-b+a)^{2}}{2}\Big{)}+\exp\Big{(}-\frac{(x+b-a)^{2}}{2}\Big{)}\Big{]}dx$
$\displaystyle\leq\exp(-b^{2}/2)+\mathinner{\\!\left\lvert
a-b\right\rvert}\cdot\exp\Big{(}\frac{a^{2}-2ab}{2}\Big{)}\,.$

Now, for the second term, we have

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(b,1)}[\tanh^{\prime}(ax)(x^{2}-bx)]$
$\displaystyle=\int_{-\infty}^{\infty}\tanh^{\prime}(ax)x(x-b)\exp\Big{(}-\frac{(x-b)^{2}}{2}\Big{)}dx$
$\displaystyle\geq-b\int_{0}^{b}xe^{-ax}\exp\Big{(}-\frac{(x-b)^{2}}{2}\Big{)}dx$
$\displaystyle\geq-b\exp\Big{(}\frac{a^{2}-2ab}{2}\Big{)}\int_{0}^{\infty}x\Big{[}\exp\Big{(}-\frac{(x-b+a)^{2}}{2}\Big{)}+\exp\Big{(}-\frac{(x+b-a)^{2}}{2}\Big{)}\Big{]}dx$
$\displaystyle\geq-b\exp(-b^{2}/2)-b\mathinner{\\!\left\lvert
a-b\right\rvert}\cdot\exp\Big{(}\frac{a^{2}-2ab}{2}\Big{)}$

We can rewrite the last term as
$\mathbb{E}_{x\sim\mathcal{N}(0,1)}[\tanh(a(x+b))x]$. Using the fact that
$\tanh(a(x+b))>\tanh(a(-x+b))$, we get that
$\mathbb{E}_{x\sim\mathcal{N}(0,1)}[\tanh(a(x+b))x]>0$. Finally, using the
upper bound on $\mathbb{E}[\tanh^{\prime}(ax)]$, we get the following lower
bound.

$\displaystyle(1-\eta-\eta\,\mathbb{E}_{x\sim\mathcal{N}(b,1)}[\tanh^{\prime}(ax)])\,a+\eta\mathbb{E}_{x\sim\mathcal{N}(b,1)}\Big{[}-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}(x-b)+\tanh^{\prime}(ax)(x^{2}-bx)\Big{]}$
$\displaystyle\geq\frac{a}{20}(19-4e^{\frac{a^{2}-2ab}{2}})+\frac{1}{20}\Big{(}-\frac{a^{2}}{2}\Big{[}\exp(-b^{2}/2)+\mathinner{\\!\left\lvert
a-b\right\rvert}\exp\Big{(}\frac{a^{2}-2ab}{2}\Big{)}\Big{]}$
$\displaystyle\qquad\qquad\qquad\qquad\qquad\qquad\quad-b\exp(-b^{2}/2)-b\mathinner{\\!\left\lvert
a-b\right\rvert}\exp(\frac{a^{2}-2ab}{2})\Big{)}\geq 1\,.\qed$

######  Lemma F.4.

For any $a,b>0$ and $a\in[30,\frac{4b}{3}]$, the following holds. Define

$U(a,b)\triangleq\eta\mathbb{E}_{x\sim\mathcal{N}(b,1)}\Big{[}\Big{(}\tanh(ax)-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}+\tanh^{\prime}(ax)ax\Big{)}x\Big{]}-\eta\mathbb{E}_{x\sim\mathcal{N}(b,1)}\mathinner{\left[\tanh^{\prime}(ax)a\right]}-\eta
a\,.$

When the learning rate $\eta=\frac{1}{20}$, is given by, we have

$\displaystyle\mathinner{\\!\left\lvert U(a,b)\right\rvert}\leq\frac{a+b}{10}$

###### Proof.

We upper bound each term in $U(a,b)$ and they apply triangle inequality to get
the result. We start with
$|\mathbb{E}_{x\sim\mathcal{N}(b,1)}\mathinner{\left[\tanh^{\prime\prime}(ax)a^{2}x\right]}|$:

$\displaystyle-\mathbb{E}_{x\sim\mathcal{N}(b,1)}\mathinner{\left[\tanh^{\prime\prime}(ax)a^{2}x\right]}$
$\displaystyle=\frac{a^{2}}{8\sqrt{2\pi}}\int_{0}^{\infty}x\sigma(2ax)(1-\sigma(2ax))(2\sigma(2ax)-1)\left(e^{-\frac{(x-b)^{2}}{2}}+e^{-\frac{(x+b)^{2}}{2}}\right)dx$
$\displaystyle\leq\frac{a^{2}}{4\sqrt{2\pi}}\int_{0}^{\infty}xe^{-2ax}e^{-\frac{(x-b)^{2}}{2}}dx$
$\displaystyle\leq\frac{a^{2}}{4\sqrt{2\pi}}\int_{0}^{\infty}e^{-ax}xe^{-\frac{(x-b)^{2}}{2}}dx$
$\displaystyle\leq\frac{a^{2}}{2}e^{-\frac{b^{2}}{2}}+\frac{a^{2}}{2}\mathinner{\\!\left\lvert
b-a\right\rvert}e^{-\frac{-2a(b-a)-a^{2}}{2}}$
$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(b,1)}[\tanh^{\prime}(ax)ax^{2}]$
$\displaystyle=\frac{1}{\sqrt{2\pi}}\int_{0}^{\infty}\tanh^{\prime}(ax)ax^{2}\left(e^{-\frac{(x-b)^{2}}{2}}+e^{-\frac{(x+b)^{2}}{2}}\right)dx$
$\displaystyle\leq a\int_{0}^{\infty}e^{-ax}x^{2}e^{-\frac{(x-b)^{2}}{2}}dx$
$\displaystyle\leq
ae^{\frac{a^{2}-2ab}{2}}\int_{0}^{\infty}x^{2}e^{-\frac{(x-b+a)^{2}}{2}}dx$
$\displaystyle\leq 2a(a-b)^{2}e^{\frac{a^{2}-2ab}{2}}$
$\displaystyle-\mathbb{E}_{x\sim\mathcal{N}(b,1)}[a\tanh^{\prime}(ax)]$
$\displaystyle=-\frac{a}{\sqrt{2\pi}}\int_{0}^{\infty}\tanh^{\prime}(ax)\left(e^{-\frac{(x-b)^{2}}{2}}+e^{-\frac{(x+b)^{2}}{2}}\right)dx$
$\displaystyle\geq-a\int_{0}^{\infty}e^{-ax}e^{-\frac{(x-b)^{2}}{2}}dx$
$\displaystyle\geq-
ae^{\frac{a^{2}-2ab}{2}}\int_{0}^{\infty}e^{-\frac{(x-b+a)^{2}}{2}}dx$
$\displaystyle\geq-4ae^{\frac{a^{2}-2ab}{2}}\,.$

Now, using the fact that $\tanh^{\prime}(x)$ and $-\tanh^{\prime\prime}(x)x$
are always positive, we have the following upper bound.

$\displaystyle\mathinner{\\!\left\lvert U(a,b)\right\rvert}$
$\displaystyle\leq\eta\,\Big{|}\mathbb{E}_{x\sim\mathcal{N}(b,1)}\Big{[}\big{(}\tanh(ax)-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}+\tanh^{\prime}(ax)ax\big{)}\cdot
x\Big{]}\Big{|}$ $\displaystyle\qquad\qquad+\eta\mathinner{\\!\left\lvert
a\right\rvert}+\eta\,\big{|}-\mathbb{E}_{x\sim\mathcal{N}(b,I)}\mathinner{\left[\tanh^{\prime}(ax)a\right]}\big{|}$
$\displaystyle\leq\eta\Big{(}2b+a+\frac{a^{2}}{2}e^{-\frac{b^{2}}{2}}+\frac{a^{2}}{2}\mathinner{\\!\left\lvert
b-a\right\rvert}e^{\frac{-2a(b-a)-a^{2}}{2}}+2a(b-a)^{2}e^{\frac{a^{2}-2ab}{2}}+2ae^{\frac{a^{2}-2ab}{2}}\Big{)}$

If $b\geq a$ and $a\geq 30$, then we have

$\displaystyle\mathinner{\\!\left\lvert U(a,b)\right\rvert}$
$\displaystyle\leq\eta\left(2b+a+0.1\right)$

If $b\leq a\leq\frac{4b}{3}$ and $a\geq 30$, then

$\displaystyle\mathinner{\\!\left\lvert U(a,b)\right\rvert}$
$\displaystyle\leq\eta\left(2b+a+0.1\right)$

Using $\eta=1/20$ and for any $a>30$, we have

$\displaystyle\mathinner{\\!\left\lvert
U(a,b)\right\rvert}\leq\frac{a+b}{10}.$

∎

###  F.4 Additional proofs for mixtures of two Gaussians

######  Lemma F.5.

Suppose $a,b>0$ satisfy $a\in[30,\frac{4b}{3}]$, then the following inequality
holds:

$\displaystyle|\mathbb{E}_{x\sim\mathcal{N}(b,1)}[-0.5\tanh^{\prime\prime}(ax)a^{2}+\tanh^{\prime}(ax)ax]|\leq
0.01$

###### Proof.

We first show that
$\mathbb{E}_{x\sim\mathcal{N}(b,1)}[-0.5\tanh^{\prime\prime}(ax)a^{2}]>0$ for
any $a,b>0$.

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(b,1)}[-0.5\tanh^{\prime\prime}(ax)a^{2}]$
$\displaystyle=-0.5a^{2}\int_{-\infty}^{\infty}\tanh^{\prime\prime}(ax)\exp(-0.5(x-b)^{2})dx$
$\displaystyle=-0.5$ $\displaystyle
a^{2}\int_{0}^{\infty}\tanh^{\prime\prime}(ax)(\exp(-0.5(x-b)^{2})-\exp(-0.5(x+b)^{2}))dx>0$

where the last inequality follows from
$\exp(-0.5(x-b)^{2})>\exp(-0.5(x+b)^{2})$ and $\tanh^{\prime\prime}(ax)<0$ for
$x>0$. We can upper bound
$\mathbb{E}_{x\sim\mathcal{N}(b,1)}[-0.5\tanh^{\prime\prime}(ax)a^{2}]$ as
follows:

$\displaystyle\mathbb{E}_{x\sim\mathcal{N}(b,1)}[-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}]$
$\displaystyle\leq-\frac{1}{2}a^{2}\int_{0}^{\infty}\tanh^{\prime\prime}(ax)\exp(-\frac{1}{2}(x-b)^{2})dx$
$\displaystyle\leq
a^{2}\int_{0}^{\infty}\exp(-ax)\exp(-\frac{1}{2}(x-b)^{2})dx$
$\displaystyle\leq
a^{2}\exp(\frac{1}{2}(a^{2}-2ab))\int_{0}^{\infty}\exp(-\frac{1}{2}(x-b+a)^{2})dx$
$\displaystyle\leq a^{2}\exp(\frac{1}{2}(a^{2}-2ab))$

When $a\leq b$, by writing $a^{2}-2ab=-2a(b-a)-a^{2}\leq-a^{2}$, we have
$\mathbb{E}[-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}]\leq 0.005$ for $a\geq
30$. When $a\in[b,\frac{4b}{3}]$, $a^{2}-2ab=\leq-\frac{2b^{2}}{9}$, we have
$|\mathbb{E}[-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}]|\leq 0.005$. Similar
to the
$\mathbb{E}_{x\sim\mathcal{N}(b,1)}[-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}]$,
we prove $\mathbb{E}_{x\sim\mathcal{N}(b,1)}[\tanh^{\prime}(ax)ax]>0$ and
$\mathbb{E}_{x\sim\mathcal{N}(b,1)}[\tanh^{\prime}(ax)ax]<0.005$. Combining
bounds for $|\mathbb{E}[\tanh^{\prime}(ax)ax]|$ and
$|\mathbb{E}[-\frac{1}{2}\tanh^{\prime\prime}(ax)a^{2}]|$ using triangle
inequality, we obtain the result. ∎