# Dissecting adaptive methods in GANs

Samy Jelassi1, David Dobre2, Arthur Mensch3, Yuanzhi Li4, Gauthier Gidel2,5

(1Princeton University
2 Mila & Université de Montréal
3 Deepmind
4 Carnegie Mellon University
5 Canada CIFAR AI Chair )

###### Abstract

Adaptive methods are a crucial component widely used for training generative
adversarial networks (GANs). While there has been some work to pinpoint the
“marginal value of adaptive methods” in standard tasks, it remains unclear why
they are still critical for GAN training. In this paper, we formally study how
adaptive methods help train GANs; inspired by the grafting method proposed in
Agarwal et al. (2020), we separate the magnitude and direction components of
the Adam updates, and graft them to the direction and magnitude of SGDA
updates respectively. By considering an update rule with the magnitude of the
Adam update and the normalized direction of SGD, we empirically show that the
adaptive magnitude of Adam is key for GAN training. This motivates us to have
a closer look at the class of normalized stochastic gradient descent ascent
(nSGDA) methods in the context of GAN training. We propose a synthetic
theoretical framework to compare the performance of nSGDA and SGDA for GAN
training with neural networks. We prove that in that setting, GANs trained
with nSGDA recover all the modes of the true distribution, whereas the same
networks trained with SGDA (and any learning rate configuration) suffer from
mode collapse. The critical insight in our analysis is that normalizing the
gradients forces the discriminator and generator to be updated at the same
pace. We also experimentally show that for several datasets, Adam’s
performance can be recovered with nSGDA methods.

##  1 Introduction

Adaptive algorithms have become a key component in training modern neural
network architectures in various deep learning tasks. Minimization problems
that arise in natural language processing (Vaswani et al., 2017), fMRI
(Zbontar et al., 2018), or min-max problems such as generative adversarial
networks (GANs) (Goodfellow et al., 2014) almost exclusively use adaptive
methods, and it has been empirically observed that Adam (Kingma & Ba, 2014)
yields a solution with better generalization than stochastic gradient descent
(SGD) in such problems (Choi et al., 2019). Several works have attempted to
explain this phenomenon in the minimization setting. Common explanations are
that adaptive methods train faster (Zhou et al., 2018), escape flat “saddle-
point”–like plateaus faster (Orvieto et al., 2021), or handle heavy-tailed
stochastic gradients better (Zhang et al., 2019; Gorbunov et al., 2022).
However, much less is known about why adaptive methods are so critical for
solving min-max problems such as GANs.

Several previous works attribute this performance to the superior convergence
speed of adaptive methods. For instance, Liu et al. (2019) show that an
adaptive variant of Optimistic Gradient Descent (Daskalakis et al., 2017)
converges faster than SGDA for a class of non-convex, non-concave min-max
problems. However, contrary to the minimization setting, convergence to a
stationary point is not guaranteed, nor is it even a requirement to ensure a
satisfactory GAN performance. Mescheder et al. (2018) empirically shows that
popular architectures such as Wasserstein GANs (WGANs) (Arjovsky et al., 2017)
do not always converge, and yet they produce realistic images. We support this
observation with our own experiments in Section 2 (see Fig. 1(b).) This
finding motivates the central question in this paper: _what factors of Adam
contribute to better quality solutions than SGDA when training GANs?_

In this paper, we investigate why GANs trained with adaptive methods
outperform those trained using stochastic gradient descent ascent (SGDA).
Directly analyzing Adam is challenging due to the highly non-linear nature of
its gradient oracle and its path-dependent update rule. Inspired by the
grafting approach in (Agarwal et al., 2020), we disentangle the adaptive
magnitude and direction of Adam and show evidence that an algorithm made of
the adaptive magnitude of Adam and the direction of SGDA (Ada-nSGDA) recovers
the performance of Adam in GANs. The adaptive magnitude in Adam is thus key
for the performance, and the standard nSGDA direction . Our contributions are
as follows:

  * •

In Section 2, we present the Ada-nSGDA algorithm and the standard normalized
SGDA (nSGDA). We further show that for some architectures and datasets, nSGDA
can be used to model the dynamics of the performance of Adam in GAN training.

  * •

In Section 3, we prove that for a synthetic dataset consisting of two modes, a
model trained with SGDA suffers from mode collapse (producing only a single
type of output), while a model trained with nSGDA does not. This provides an
explanation for why GANs trained with nSGDA outperform those trained with
SGDA.

  * •

In Section 4, we empirically confirm that Ada-nSDGA recovers the performance
of Adam when using different GAN architectures on a wide range of datasets.

Our key theoretical insight is that when using SGDA and any step-size
configuration, either the generator $G$ or discriminator $D$ updates much
faster than its counterpart. By normalizing the gradients as done in nSGDA,
$D$ and $G$ are forced to update at the same speed throughout training. The
consequence is that whenever $D$ learns a mode of the distribution, $G$ learns
it right after, which makes both of them learn all the modes of the
distribution separately.

###  1.1 Related work

Adaptive methods in games optimization. Several works designed adaptive
algorithms and analyzed their convergence to show their benefits relative to
SGDA e.g. in variational inequality problems, Gasnikov et al. (2019);
Antonakopoulos et al. (2019); Bach & Levy (2019); Antonakopoulos et al.
(2020); Liu et al. (2019); Barazandeh et al. (2021). Heusel et al. (2017) show
that Adam locally converges to a Nash equilibrium in the regime where the
step-size of the discriminator is much larger than the one of the generator.
Our work differs as we do not focus on the convergence properties of Adam, but
rather on the fit of the trained model to the true (and not empirical) data
distribution.

Statistical results in GANs. Early works studied whether GANs memorize the
training data or actually learn the distribution (Arora et al., 2017; Liang,
2017; Feizi et al., 2017; Zhang et al., 2017; Arora et al., 2018; Bai et al.,
2018; Dumoulin et al., 2016). Some works explained GAN performance through the
lens of optimization. Lei et al. (2020); Balaji et al. (2021) show that GANs
trained with SGDA converge to a global saddle point when the generator is one-
layer neural network and the discriminator is a specific quadratic/linear
function. Our contribution differs as i) we construct a setting where SGDA
converges to a locally optimal min-max equilibrium but still suffers from mode
collapse, and ii) we have a more challenging setting since we need at least a
degree-3 discriminator to learn the distribution, which is discussed in
Section 3.

Normalized gradient descent. Introduced by Nesterov (1984), normalized
gradient descent has been widely used in minimization problems. Normalizing
the gradient remedies the issue of iterates being stuck in flat regions such
as spurious local minima or saddle points (Hazan et al., 2015; Levy, 2016).
Normalized gradient descent methods outperforms their non-normalized
counterparts in multi-agent coordination (Cortés, 2006) and deep learning
tasks (Cutkosky & Mehta, 2020). Our work considers the min-max setting and
shows that nSGDA outperforms SGDA as it forces discriminator and generator to
update at same rate.

(a) Each circle corresponds to a specific step-size configuration
$\eta_{D}/\eta_{G}$. The best-performing models have step-size ratios between
$10^{-1}$ and $1$, and do not converge. As $\eta_{D}/\eta_{G}$ increases, the
models perform worse but get closer to an equilibrium.

(b) shows that during training, the gradient ratio of a well-performing GAN
approximately stays constant to 1. We also display the images produced by the
model during training.

Figure 1: Gradient ratio against FID score (a) and number of epochs (b)
obtained with DCGAN on CIFAR-10. This ratio is equal to
$\|\mathrm{grad}_{G}^{(t)}\|_{2}/\|\mathrm{grad}_{G}^{(0)}\|_{2}+\|\mathrm{grad}_{D}^{(t)}\|_{2}/\|\mathrm{grad}_{D}^{(0)}\|_{2}$,
where $\mathrm{grad}_{G}^{(t)}$ (resp. $\mathrm{grad}_{D}^{(t)}$) and
$\mathrm{grad}_{G}^{(0)}$ (resp. $\mathrm{grad}_{D}^{(0)}$) are the current
and initial gradients of $G$ (resp. $D$). Note that $\|\cdot\|_{2}$ refers to
the sum of all the parameters norm in a network. For all the plots, the models
are trained for 100 epochs using a batch-size 64. For (b), the results are
averaged over 5 seeds.

###  1.2 Background

Generative adversarial networks. Given a training set sampled from some target
distribution $\mathcal{D}$, a GAN learns to generate new data from this
distribution. The architecture is comprised of two networks: a generator that
maps points in the latent space $\mathcal{D}_{z}$ to samples of the desired
distribution, and a discriminator which evaluates these samples by comparing
them to samples from $\mathcal{D}$. More formally, the generator is a mapping
$G_{\mathcal{V}}\colon\mathbb{R}^{k}\rightarrow\mathbb{R}^{d}$ and the
discriminator is a mapping
$D_{\mathcal{W}}\colon\mathbb{R}^{d}\rightarrow\mathbb{R}$, where
$\mathcal{V}$ and $\mathcal{W}$ are their corresponding parameter sets. To
find the optimal parameters of these two networks, one must solve a min-max
optimization problem of the form

$\displaystyle\min_{\mathcal{V}}\max_{\mathcal{W}}$
$\displaystyle\;\mathbb{E}_{X\sim
p_{data}}[\log(D_{\mathcal{W}}(X))]+\mathbb{E}_{z\sim
p_{z}}[\log(1-D_{\mathcal{W}}(G_{\mathcal{V}}(z)))]:=f(\mathcal{V},\mathcal{W}),$
(GAN)

where $p_{data}$ is the distribution of the training set, $p_{z}$ the latent
distribution, $G_{\mathcal{V}}$ the generator and $D_{\mathcal{W}}$ the
discriminator. Contrary to minimization problems where convergence to a local
minimum is _required_ for high generalization, we empirically verify that most
of the well-performing GANs do not converge to a stationary point.

#### Convergence and performance are decorrelated in GANs.

We support this observation through the following experiment. We train a DCGAN
(Radford et al., 2015) using Adam and set up the step-sizes for $G$ and $D$ as
$\eta_{D},\eta_{G}$, respectively. Note that $D$ is usually trained faster
than $G$ i.e. $\eta_{D}\geq\eta_{G}.$ 1(a) displays the GAN convergence
measured by the ratio of gradient norms and the GAN’s performance measured in
FID score (Heusel et al., 2017). We observe that when $\eta_{D}/\eta_{G}$ is
close to $1$, the algorithm does not converge, and yet, the model produces
high-quality solutions. On the other hand, when $\eta_{D}/\eta_{G}\gg 1$, the
model converges to an equilibrium; a similar statement has been proved by Jin
et al. (2020) and Fiez & Ratliff (2020) in the case of SGDA. However, the
trained GAN produces low-quality solutions at this equilibrium, so simply
comparing the convergence speed of adaptive methods and SGDA cannot explain
the performance obtained with adaptive methods.

SGDA and adaptive methods. The most simple algorithm to solve the min-max
(GAN) is SGDA, which is defined as follows:

$\displaystyle\mathcal{W}^{(t+1)}$
$\displaystyle=\mathcal{W}^{(t)}+\eta_{D}\mathbf{M}_{\mathcal{W},1}^{(t)},\quad\mathcal{V}^{(t+1)}=\mathcal{V}^{(t)}-\eta_{G}\mathbf{M}_{\mathcal{V},1}^{(t)}\,,$
(1)

where $\mathbf{M}_{\mathcal{W},1}^{(t)},\mathbf{M}_{\mathcal{V},1}^{(t)}$ are
the first-order momentum gradients as defined in Algorithm 1. While this
method has been used in the first GANs (Radford et al., 2015), most modern
GANs are trained with adaptive methods such as Adam (Kingma & Ba, 2014).

The definition of this algorithm for game optimizations is given in Algorithm
1. The hyperparameters $\beta_{1},\beta_{2}\in[0,1)$ control the weighting of
the exponential moving average of the first and second-order moments. In many
deep-learning tasks, practitioners have found that setting $\beta_{2}=0.9$
works for most problem settings. It has been empirically observed that having
no momentum (i.e., $\beta_{1}\approx 0$) is optimal for many popular
architectures (Karras et al., 2020; Brock et al., 2018). Thus, in what
follows, we only consider the case where $\beta_{1}=0$.

Optimizers such as Adam (Algorithm 1) are adaptive because they keep updating
step-sizes while training the model. There are two components that contribute
to the update step: the adaptive magnitude
$\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}$and the adaptive direction
$\mathbf{A}_{\mathcal{Y}}^{(t)}/\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}$. The
two components are entangled and it remains unclear how they contribute to the
superior performance of adaptive methods relative to SGDA in GANs.

Input: initial points $\mathcal{W}^{(0)},\mathcal{V}^{(0)}$, step-size
schedules $\\{(\eta_{G}^{(t)},\eta_{D}^{(t)})\\}$ , hyperparameters
$\\{\beta_{1},\beta_{2},\varepsilon\\}$. Initialize
$\mathbf{M}_{\mathcal{W},1}^{(0)}$, $\mathbf{M}_{\mathcal{W},2}^{(0)}$,
$\mathbf{M}_{\mathcal{V},1}^{(0)}$ and $\mathbf{M}_{\mathcal{V},2}^{(0)}$ to
zero.

for $t=0\dots T-1$ do

Receive stochastic gradients
$\mathbf{g}_{\mathcal{W}}^{(t)},\mathbf{g}_{\mathcal{V}}^{(t)}$ evaluated at
${\mathcal{W}}^{(t)}$ and ${\mathcal{V}}^{(t)}$.

Update for $\mathcal{Y}\in\\{\mathcal{W},\mathcal{V}\\}$:
$\mathbf{M}_{\mathcal{Y},1}^{(t+1)}=\beta_{1}\mathbf{M}_{\mathcal{Y},1}^{(t)}+\mathbf{g}_{\mathcal{Y}}^{(t)}$
and
$\mathbf{M}_{\mathcal{Y},2}^{(t+1)}=\beta_{2}\mathbf{M}_{\mathcal{Y},2}^{(t)}+{\mathbf{g}_{\mathcal{Y}}^{(t)}}^{2}.$

Compute gradient oracles for $Y\in\\{V,W\\}$:
$\mathbf{A}_{\mathcal{Y}}^{(t+1)}=\nicefrac{{\mathbf{M}_{\mathcal{Y},1}^{(t+1)}}}{{\sqrt{\mathbf{M}_{\mathcal{Y},2}^{(t+1)}+\varepsilon}}}.$

Update:
$\mathcal{W}^{(t+1)}=\mathcal{W}^{(t)}+\eta_{D}^{(t)}\mathbf{A}_{\mathcal{W}}^{(t+1)},\qquad\mathcal{V}^{(t+1)}=\mathcal{V}^{(t)}-\eta_{G}^{(t)}\mathbf{A}_{\mathcal{V}}^{(t+1)}.$

Algorithm 1 Adam (Kingma & Ba, 2014) for games. All operations on vectors are
element-wise.

##  2 nSGDA as a model to analyze Adam in GANs

In this section, we show that normalized stochastic gradient descent-ascent
(nSGDA) is a suitable proxy to study the learning dynamics of Adam.

To decouple the adaptive magnitude and direction in Adam, we adopt the step-
size grafting approach proposed by Agarwal et al. (2020). At each iteration,
we compute stochastic gradients, pass them to two optimizers
$\mathcal{A}_{1},\mathcal{A}_{2}$ and make a grafted step that combines the
magnitude of $\mathcal{A}_{1}$’s step and direction of $\mathcal{A}_{2}$’s
step. We focus on the optimizer defined by grafting the Adam magnitude onto
the SGDA direction, i.e:

$\displaystyle\mathcal{W}^{(t+1)}$
$\displaystyle=\mathcal{W}^{(t)}+\eta_{D}^{(t)}\|\mathbf{A}_{\mathcal{W}}^{(t)}\|_{2}\frac{\mathbf{g}_{\mathcal{W}}^{(t)}}{\|\mathbf{g}_{\mathcal{W}}^{(t)}\|_{2}+\varepsilon},\quad\mathcal{V}^{(t+1)}=\mathcal{V}^{(t)}-\eta_{G}^{(t)}\|\mathbf{A}_{\mathcal{V}}^{(t)}\|_{2}\frac{\mathbf{g}_{\mathcal{V}}^{(t)}}{\|\mathbf{g}_{\mathcal{V}}^{(t)}\|_{2}+\varepsilon},$
(2)

where $\mathbf{A}_{\mathcal{V}}^{(t)},\mathbf{A}_{\mathcal{W}}^{(t)}$ are the
Adam gradient oracles as in Algorithm 1 and
$\bm{g}_{\mathcal{V}}^{(t)},\bm{g}_{\mathcal{W}}^{(t)}$ the stochastic
gradients. We refer to this algorithm as Ada-nSGDA (combining the Adam
magnitude and SGDA direction). There are two natural implementations for
nSDGA. In the _layer-wise_ version, $\mathcal{Y}^{(t)}$ is a single parameter
group (typically a layer in a neural network), and the updates are applied to
each group. In the _global_ version, $\mathcal{Y}^{(t)}$ contains all of the
model’s weights.

In Fig. 2, we show that Ada-nSGDA and Adam appear to have similar learning
dynamics in terms of the FID score. Both Adam and Ada-nSGDA significantly
outperform SGDA as well as AdaDir, which is the alternate case of (2) where we
instead graft the magnitude of the SGDA update to the direction of the Adam
update. AdaDir diverged after a single step so we omit it in Fig. 2. This
confirms that the critical component of Adam is the _adaptive magnitude_ , and
that the standard update direction recovered by SGDA is sufficient to recover
a good solution. However, from a theoretical perspective, directly analyzing
Ada-nSGDA is difficult due to the adaptive magnitudes
$\|\mathbf{A}_{\mathcal{V}}^{(t)}\|_{2},\|\mathbf{A}_{\mathcal{W}}^{(t)}\|_{2}$.
Therefore, in Section 3, we analyze normalized SGDA (nSGDA) which is Ada-nSGDA
(2) where we omit the adaptive magnitudes. Although nSGDA does not consider
the adaptive magnitude, it still recovers the performance of Adam for some
architectures such as WGAN-GP (Arjovsky et al., 2017) as we show in Fig. 2.
This may come from the fact that the adaptive magnitudes stay within a
constant range and do not fluctuate across time, as shown in Figs. 2(b) and
2(c).

(a)

(b)

(c)

Figure 2:  (a) shows the FID training curve for a WGAN-GP ResNet, averaged
over 5 seeds. We see that Ada-nSGDA and nSGDA have very similar performance to
Adam for a WGAN-GP. (b, c) displays the fluctuations of Ada-nSGDA adaptive
magnitude. We plot the ratio
$\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}/\|\mathbf{A}_{\mathcal{Y}}^{(0)}\|_{2}$
for each discriminator’s (b) and generator’s (c) layers. At early stages, this
ratio barely increases and remains constant after 10 steps.

##  3 Why does nSGDA perform better than SGDA in GANs?

In Section 2, we numerically showed that nSGDA recovers the performance Adam.
Similar to other optimization works in minimization, we use nSGDA as a model
to understand Adam in GANs. Our goal is to construct a dataset and model where
we can prove that a model trained with nSGDA generates samples from the true
training distribution while SGDA fails. To this end, we consider a dataset
where the underlying distribution consists of two modes, defined as vectors
$u_{1},u_{2}$, that are slightly correlated (See Assumption 1) and consider
the standard GANs’ training objective. We show that a GAN trained with SGDA
using any reasonable111Reasonable simply means that the learning rates are
bounded to prevent the training from diverging. step-size configuration
suffers from _mode collapse_ (Theorem 3.1); it only outputs samples from a
single mode which is a weighted average of $u_{1}$ and $u_{2}$. Conversely,
nSGDA-trained GANs learn the two modes separately (Theorem 3.2).

#### Notation

We set the GAN 1-sample loss
$L_{\mathcal{V},\mathcal{W}}^{(t)}(X,z)=\log(D_{\mathcal{W}}^{(t)}(X))+\log(1-D_{\mathcal{W}}^{(t)}(G_{\mathcal{V}}^{(t)}(z))).$
We denote
$\mathbf{g}_{\mathcal{Y}}^{(t)}=\nabla_{\mathcal{Y}}L_{\mathcal{V},\mathcal{W}}^{(t)}(X,z)$
as the 1-sample stochastic gradient. We use the asymptotic complexity
notations when defining the different constants e.g. $\mathrm{poly}(d)$ refers
to any polynomial in the dimension $d$, $\mathrm{polylog}(d)$ to any
polynomial in $\log(d)$, and $o(1)$ to a constant $\ll d$. We denote $a\propto
b$ for vectors $a$ and $b$ in $\mathbb{R}^{d}$ if there is a positive scaling
factor $c>0$ such that $\|a-cb\|_{2}=o(\|b\|_{2})$.

###  3.1 Setting

In this section, we present the setting to sketch our main results in Theorem
3.1 and Theorem 3.2. We first define the distributions for the training set
and latent samples, and specify our GAN model and the algorithms we analyze to
solve (GAN). Note that for many assumptions and theorems below, we present
informal statements which are sufficient to capture the main insights. The
precise statements can be found in Appendix B.

Our synthetic theoretical framework considers a bimodal data distribution with
two correlated modes:

######  Assumption 1 ($p_{data}$ structure).

Let $\gamma=\frac{1}{\mathrm{polylog}(d)}$. We assume that the modes are
correlated. This means that $\langle u_{1},u_{2}\rangle=\gamma>0$ and the
generated data point $X$ is either $X=u_{1}$ or $X=u_{2}.$

Next, we define the latent distribution $p_{z}$ that $G_{\mathcal{V}}$ samples
from and maps to $p_{data}$. Each sample from $p_{z}$ consists of a data-point
$z$ that is a binary-valued vector $z\in\\{0,1\\}^{m_{G}}$, where $m_{G}$ is
the number of neurons in $G_{\mathcal{V}}$, and has non-zero support, i.e.
$\|z\|_{0}\geq 1$. Although the typical choice of a latent distributions in
GANs is either Gaussian or uniform, we choose $p_{z}$ to be a binary
distribution because it models the weights’ distribution of a hidden layer of
a deep generator; Allen-Zhu & Li (2021) argue that the distributions of these
hidden layers are sparse, non-negative, and non-positively correlated. We now
make the following assumptions on the coefficients of $z$:

######  Assumption 2 ($p_{z}$ structure).

Let $z\sim p_{z}$. We assume that with probability $1-o(1)$, there is only one
non-zero entry in $z$. The probability that the entry $i\in[m_{G}]$ is non-
zero is $\Pr[z_{i}=1]=\Theta(1/m_{G}).$

In Assumption 2, the output of $G_{\mathcal{V}}$ is only made of one mode with
probability $1-o(1)$. This avoids summing two of the generator’s neurons,
which may cause mode collapse.

To learn the target distribution $p_{data}$, we use a linear generator
$G_{\mathcal{V}}$ with $m_{G}$ neurons and a non-linear neural network with
$m_{D}$ neurons:

$G_{\mathcal{V}}(z)=Vz=\sum_{i=1}^{m_{G}}v_{i}z_{i}\,,\qquad
D_{\mathcal{W}}(X)=\mathrm{sigmoid}\Big{(}a\sum_{i=1}^{m_{D}}\langle
w_{i},X\rangle^{3}+\frac{b}{\sqrt{d}}\Big{)}.$ (3)

where
$V=[v_{1}^{\top},v_{2}^{\top},\cdots,v_{m_{G}}^{\top}]\in\mathbb{R}^{m_{G}\times
d}$, $z\in\\{0,1\\}^{m_{G}}$,
$W=[w_{1}^{\top},\dots,w_{m_{D}}^{\top}]\in\mathbb{R}^{m_{D}\times d}$, and
$a,b\in\mathbb{R}$. Intuitively, $G_{\mathcal{V}}$ outputs linear combinations
of the modes $v_{i}$. We choose a cubic activation as it is the smallest
monomial degree for the discriminator’s non-linearity that is sufficient for
the generator to recover the modes $u_{1},u_{2}$.222Li & Dou (2020) show that
when using linear or quadratic activations, the generator can fool the
discriminator by only matching the first and second moments of $p_{data}$.

We now state the SGDA and nSGDA algorithms used to solve the GAN training
problem (GAN). For simplicity, we set the batch-size to 1. The resultant
update rules for SGDA and nSGDA are:333In the nSGDA algorithm defined in (2),
the step-sizes were time-dependent. Here, we assume for simplicity that the
step-sizes $\eta_{D},\eta_{G}>0$ are _constant_.

SGDA: at each step $t>0,$ sample $X\sim p_{data}$ and $z\sim p_{z}$ and update
as

$\displaystyle\mathcal{W}^{(t+1)}$
$\displaystyle=\mathcal{W}^{(t)}+\eta_{D}\mathbf{g}_{\mathcal{W}}^{(t)},\quad\mathcal{V}^{(t+1)}=\mathcal{V}^{(t)}-\eta_{G}\mathbf{g}_{\mathcal{V}}^{(t)},$
(4)

nSGDA: at each step $t>0,$ sample $X\sim p_{data}$ and $z\sim p_{z}$ and
update as

$\displaystyle\mathcal{W}^{(t+1)}$
$\displaystyle=\mathcal{W}^{(t)}+\eta_{D}\tfrac{\mathbf{g}_{\mathcal{W}}^{(t)}}{\|\mathbf{g}_{\mathcal{W}}^{(t)}\|_{2}},\quad\mathcal{V}^{(t+1)}=\mathcal{V}^{(t)}-\eta_{G}\tfrac{\mathbf{g}_{\mathcal{V}}^{(t)}}{\|\mathbf{g}_{\mathcal{V}}^{(t)}\|_{2}}.$
(5)

Compared to the versions of SGDA and Ada-nSGDA that we introduced in Section
2, we have the same algorithms except that we set $\beta_{1}=0$ and omit
$\varepsilon$ in (4) and (5). Lastly, we detail how to set the optimization
parameters for SGDA and nSGDA in (4) and (5).

######  Parametrization 3.1 (Informal).

When running SGDA and nSGDA on (GAN), we set:

– Initialization: $b^{(0)}=0,$ and
$a^{(0)},\;w_{i}^{(0)}(i\in[m_{D}]),\;v_{j}^{(0)}(j\in[m_{G}])$ are
initialized with a Gaussian with small variance.

– Number of iterations: we run SGDA for $t\leq T_{0}$ iterations where $T_{0}$
is the first iteration such that the algorithm converges to an approximate
first order local minimum. For nSGDA, we run for
$T_{1}=\tilde{\Theta}(1/\eta_{D})$ iterations.

– Step-sizes: For SGDA, $\eta_{D},\eta_{G}\in(0,\frac{1}{\mathrm{poly}(d)})$
can be arbitrary. For nSGDA, $\eta_{D}\in(0,\frac{1}{\mathrm{poly}(d)}]$, and
$\eta_{G}$ is slightly smaller than $\eta_{D}.$

– Over-parametrization: For SGDA, $m_{D},m_{G}=\mathrm{polylog}(d)$ are
arbitrarily chosen i.e. $m_{D}$ may be larger than $m_{G}$ or the opposite.
For nSGDA, we set $m_{D}=\log(d)$ and $m_{G}=2\log(d).$

Our theorem holds when running SGDA for any (polynomially) possible number of
iterations; after $T_{0}$ steps, the gradient becomes inverse polynomially
small and SGDA essentially stops updating the parameters. Additionally, our
setting allows any step-size configuration for SGDA i.e. larger, smaller, or
equal step-size for $D$ compared to $G$. Note that our choice of step-sizes
for nSGDA is the one used in practice, i.e. $\eta_{D}$ slightly larger than
$\eta_{G}.$

###  3.2 Main results

We state our main results on the performance of models trained using SGDA (4)
and nSGDA (5). We show that nSGDA learns the modes of the distribution
$p_{data}$ while SGDA does not.

######  Theorem 3.1 (Informal).

Consider a training dataset and a latent distribution as described above and
let Assumption 1 and Assumption 2 hold. Let $T_{0}$, $\eta_{G},\eta_{D}$ and
the initialization be as defined in 3.1. Let $t$ be such that $t\leq T_{0}.$
Run SGDA on the GAN problem defined in (GAN) for $t$ iterations with step-
sizes $\eta_{G},\eta_{D}$. Then, with probability at least $1-o(1)$, the
generator outputs for all $z\in\\{0,1\\}^{m_{G}}$:

$\displaystyle
G_{\mathcal{V}}^{(t)}(z)\propto\begin{cases}u_{1}+u_{2}&\text{if
}\eta_{D}\geq\eta_{G}\\\ \xi^{(t)}(z)&\text{otherwise}\end{cases},$ (6)

where $\xi^{(t)}(z)\in\mathbb{R}^{d}$ is some vector that is not correlated to
any of the modes. Formally, $\forall\ell\in[2]$,
$\cos(\xi^{(t)}(z),u_{\ell})=o(1)$ for all $z\in\\{0,1\\}^{m_{G}}$.

A formal proof can be found in Appendix G. Theorem 3.1 indicates that when
training with SGDA and any step-size configuration, the generator either does
not learn the modes at all ($G_{\mathcal{V}}^{(t)}(z)=\xi^{(t)}(z)$) or learns
an average of the modes ($G_{\mathcal{V}}^{(t)}(z)\propto u_{1}+u_{2}$). The
theorem holds _for any_ time $t\leq T_{0}$ which is the iteration where SGDA
converges to an approximate first-order locally optimal min-max equilibrium.
Conversely, nSGDA succeeds in learning the two modes separately:

######  Theorem 3.2 (Informal).

Consider a training dataset and a latent distribution as described above and
let Assumption 1 and Assumption 2 hold. Let $T_{1}$, $\eta_{G},\eta_{D}$ and
the initialization as defined in 3.1. Run nSGDA on (GAN) for $T_{1}$
iterations with step-sizes $\eta_{G},\eta_{D}$. Then, the generator learns
both modes $u_{1},u_{2}$ i.e., for $\ell\in\\{1,2\\}$,

$\displaystyle\Pr{}_{z\sim p_{z}}[G_{\mathcal{V}}^{(T_{1})}(z)\propto
u_{\ell}]\quad\text{is non-negligible}.$ (7)

A formal proof can be found in Appendix I. Theorem 3.2 indicates that when we
train a GAN with nSGDA in the regime where the discriminator updates slightly
faster than the generator (as done in practice), the generator successfully
learns the distribution containing the direction of both modes.

We implement the setting introduced in Subsection 3.1 and validate Theorem 3.1
and Theorem 3.2 in Fig. 3. Fig. 3(a) displays the relative update speed
$\eta\|\mathbf{g}_{\mathcal{Y}}^{(t)}\|_{2}/\|\mathcal{Y}^{(t)}\|_{2}$, where
$\mathcal{Y}$ corresponds to the parameters of either $D$ or $G$. Fig. 3(b)
shows the correlation $\langle
w_{i}^{(t)},u_{\ell}\rangle/\|w_{i}^{(t)}\|_{2}$ between _one_ of $D$’s
neurons and a mode $u_{\ell}$ and Fig. 3(c) the correlation $\langle
v_{j}^{(t)},u_{\ell}\rangle/\|v_{j}^{(t)}\|_{2}$ between $G$’s neurons and
$u_{\ell}$. We discuss the interpretation of these plots to the next section.

### Why does SGDA suffer from mode collapse and nSGDA learn the modes?

(a) Relative gradients

(b) $D$ weight correlation

(c) $\mathcal{V}$ learned via SGDA

(d) $\mathcal{V}$ learned via nSGDA

Figure 3: (a) shows the relative gradient updates for SGDA. $D$ first updates
its weights while $G$ does not move until iteration 20, then $G$ moves its
weights. (b) shows the correlation for one neuron of $D$ (with maximal
correlation to $u_{2}$ at initialization) with the modes $u_{1},u_{2}$ during
the learning process of SGDA. (c, d) shows the correlations of the neurons of
$G$ with the modes when trained with SGDA and nSGDA respectively. This shows
that for SGDA (c), the model ultimately learns the weighted average
$u_{1}+u_{2}.$ For nSGDA, we see from (d) that one of the neurons ($V_{4}$) is
highly correlated with $u_{1}$ and another one ($V_{3}$) is correlated with
$u_{2}.$

We now explain why SGDA suffers from mode collapse, which corresponds to the
case where $\eta_{D}\geq\eta_{G}$. Our explanation relies on the
interpretation of Figs. 3(a), 3(b), and 3(c), and on the updates around
initialization that are defined as followed. There exists $i\in[m_{D}]$ such
that $D$’s update is with high probability

$\displaystyle\mathbb{E}[w_{i}^{(t+1)}|w_{i}^{(t)}]\approx
w_{i}^{(t)}+\eta_{D}\sum_{l=1}^{2}\mathbb{E}[\langle
w_{i}^{(t)},u_{l}\rangle^{2}]u_{l}\,.$ (8)

Thus, the weights of $D$ receive gradients directed by $u_{1}$ and $u_{2}$. On
the other hand, the weights of $G$ at early stages receive gradients directed
by $w_{j}^{(t)}$:

$\displaystyle v_{i}^{(t+1)}\approx v_{i}^{(t)}+\eta_{G}\sum_{j}\langle
v_{i}^{(t)},w_{j}^{(t)}\rangle^{2}w_{j}^{(t)}.$ (9)

We observe that the learning process in Figs. 3(a) & 3(b) has three distinct
phases. In the first phase (iterations 1-20), $D$ learns one of the modes
($u_{1}$ or $u_{2}$) of $p_{data}$ (Fig. 3(b)) and $G$ barely updates its
weights (Fig. 3(a)). In the second phase (iterations 20-40), $D$ learns the
weighted average $u_{1}+u_{2}$ (Fig. 3(b)) while $G$ starts moving its weights
(Fig. 3(a)). In the final phase (iterations 40+), $G$ learns $u_{1}+u_{2}$
(Fig. 3(c)) from $D$. In more detail, the learning process is described as
follows:

Phase 1 : At initialization, $w_{j}^{(0)}$ and $v_{j}^{(0)}$ are small. Assume
w.l.o.g. that $\langle w_{i}^{(0)},u_{2}\rangle>\langle
w_{i}^{(0)},u_{1}\rangle$. Because of the $\langle
w_{i}^{(t)},u_{l}\rangle^{2}$ in front of $u_{2}$ in (8), the parameter
$w_{i}^{(t)}$ gradually grows its correlation with $u_{2}$ (Fig. 3(b)) and
$D$’s gradient norm thus increases (Fig. 3(a)). While $\|w_{j}^{(t)}\|\ll
1\,\forall j$, we have that $v_{i}^{(t)}\approx v_{i}^{(0)}$ (Fig. 3(a)).

Phase 2: $D$ has learned $u_{2}$. Because of the sigmoid in the gradient of
$w_{i}^{(t)}$ (that was negligible during Phase 1) and $\langle
u_{1},u_{2}\rangle=\gamma>0$, $w_{i}^{(t)}$ now mainly receives updates with
direction $u_{2}$. Since $G$ did not update its weights yet, the min-max
problem (GAN) is approximately just a minimization problem with respect to
$D$’s parameters. Since the optimum of such a problem is the weighted average
$u_{1}+u_{2}$, $w_{j}^{(t)}$ slowly converges to this optimum. Meanwhile,
$v_{i}^{(t)}$ start to receive some significant signal (Fig. 3(a)) but mainly
learn the direction $u_{1}+u_{2}$ (Fig. 3(c)), because $w_{j}^{(t)}$ is
aligning with this direction.

Phase 3: The parameters of $G$ only receive gradient directed by
$u_{1}+u_{2}$. The norm of its relative updates stay large and $D$ only
changes its last layer terms (slope $a$ and bias $b$).

In contrast to SGDA, nSGDA ensures that $G$ and $D$ always learn at the same
speed with the updates:

$\displaystyle w_{i}^{(t+1)}\approx w_{i}^{(t)}+\eta_{D}\frac{\langle
w_{i}^{(t)},X\rangle^{2}X}{\|\langle w_{i}^{(t)},X\rangle^{2}X\|_{2}},\text{
and }v_{i}^{(t+1)}\approx v_{i}^{(t)}+\eta_{G}\frac{\sum_{j}\langle
w_{j}^{(t)},v_{i}^{(t)}\rangle^{2}w_{j}^{(t)}}{\|\sum_{j}\langle
w_{j}^{(t)},v_{i}^{(t)}\rangle^{2}w_{j}^{(t)}\|_{2}}$ (10)

No matter how large $\langle w_{i}^{(t)},X\rangle$ is, $G$ still learns at the
same speed with $D$. There is a tight window (iteration $25$, Fig. 3(b)) where
only one neuron of $D$ is aligned with $u_{1}$. This is when $G$ can also
learn to generate $u_{1}$ by “catching up” to $D$ at that point, which avoids
mode collapse.

##  4 Numerical performance of nSGDA

In Section 2, we present the Ada-nSGDA algorithm (2) which corresponds to
“grafting” the Adam magnitude onto the SGDA direction. In Section 3, we
construct a dataset and GAN model where we prove that a GAN trained with nSGDA
can generate examples from the true training distribution, while a GAN trained
with SGDA fails due to mode collapse. We now provide more experiments
comparing nSGDA and Ada-nSGDA with Adam on real GANs and datasets.

(a)

(b)

(c)

(d)

Figure 4: (a, b) are the final FID scores for a ResNet WGAN-GP model trained
for 45k steps on CIFAR-10 and STL-10 respectively. (c, d) are the final FID
scores for a StyleGAN2 model trained for 2600kimgs on FFHQ and LSUN Churches
respectively. WGAN-GP results are obtained over 5 seeds, and StyleGAN2 results
over 3 seeds.

We train a ResNet WGAN with gradient penalty on CIFAR-10 (Krizhevsky et al.,
2009) and STL-10 (Coates et al., 2011) with Adam, Ada-nSDGA, SGDA, as well as
nSGDA with a fixed learning rate as done in Section 3. We use the default
architectures and training parameters specified in Gulrajani et al. (2017)
($\lambda_{GP}=10$, $n_{dis}=5$, learning rate decayed linearly to 0 over 100k
steps). We also train a StyleGAN2 model (Karras et al., 2020) on FFHQ (Karras
et al., 2019) and LSUN Churches (Yu et al., 2016) (both resized to $128\times
128$ pixels) with Adam, Ada-nSGDA, and SGDA. We use the recommended StyleGAN2
hyperparameter configuration for this resolution (batch size = 32,
$\gamma=0.1024$, map depth = 2, channel multiplier = 16384). We use the
Fréchet Inception distance (FID) (Heusel et al., 2017) to quantitatively
assess the performance of the model. For each optimizer, we conduct a coarse
log-space sweep over step sizes and optimize for FID. We train the WGAN-GP
models for 2880 thousand images (kimgs) on CIFAR-10 and STL-10 (45k steps with
a batch size of 64), and the StyleGAN2 models for 2600 kimgs on FFHQ and LSUN
Churches. We average our results over 5 seeds for the WGAN-GP ResNets, and
over 3 seeds for the StyleGAN2 models due to the computational cost associated
with training GANs.

#### WGAN-GP

Figures 4(a) and 4(b) show the FID curves during training for the WGAN-GP
model. We test Adam, Ada-nSGDA, AdaDir (that contains the adaptive direction),
nSGDA and SGDA. The first remark is that Ada-nSGDA recovers the performance of
Adam while AdaDir diverges (NaN loss) for these experiments and is hence
omitted. This means that the adaptive magnitude is key to the performance of
adaptive methods. Additionally, nSGDA not only recovers the performance of
Adam but obtains a final FID of $\sim$2-3 points lower than Ada-nSGDA. Such
performance is possible because the adaptive
magnitude$\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}$ stays within a constant
range and does not fluctuate much (Figs. 2(b), 2(c)) and a constant learning
rate as in nSGDA is enough to recover Adam’s performance. Thus, nSGDA is a
valid model for Adam in the case of WGAN-GP. In contrast, models trained with
SGDA consistently perform significantly worse, with final FID scores $4\times$
larger than Adam.

#### StyleGAN2

Figures 4(c) and 4(d) show the FID curves during training for StyleGAN2.
Similar to WGAN-GP, we observe that Ada-nSGDA recovers the Adam performance
while SGDA performs much worse. However, nSGDA was not able to recover the
performance of Adam in this setting. There are several hypotheses to explain
this:

  * •

We observe in Figs. 5(b), 5(c) that the adaptive magnitude the ratio
$\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}$ stays within a constant range but
varies significantly across time. This violates our theoretical setting of
assuming a constant adaptive magnitude. In contrast,
$\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}$ stays within a constant range and
does not fluctuate across time in the WGAN-GP experiment (Figs. 2(b), 2(c)).

  * •

The GAN problem formulation in the case of StyleGAN2 is very different from
the one in WGAN-GP. Specifically, StyleGAN2 has a drastically different
generator architecture than the ResNet generator used in the WGAN-GP
experiments (utilizing weight demodulation), as well as using adaptive data
augmentation and additional regularizers (such as a path-length
regularization).

Because our theoretical setting does not capture these observations, we do not
expect nSGDA to work in this setting. However in spite of these differences,
Ada-nSGDA performed similarly to Adam indicating that the nSGDA direction that
we theoretically study is still valid in modern real-world GAN architectures.
We further validate our theory with additional experiments on DCGAN (Radford
et al., 2015) which more closely matches our theoretical setting in Appendix
A, and find that nSGDA recovers the performance of Adam as in WGAN-GP.

(a)

(b)

(c)

Figure 5: (a) shows the FID curve for the StyleGAN2 model trained for 2600
kimgs on FFHQ. (b, c) shows the fluctuations of the Ada-nSGDA adaptive
magnitude. Like Fig. 2, we plot the ratio
$\|\mathbf{A}_{\mathcal{Y}}^{(t)}\|_{2}/\|\mathbf{A}_{\mathcal{Y}}^{(0)}\|_{2}$
for each of the discriminator’s (b) and generator’s (c) layers. This ratio
stays in a stable range but fluctuates a lot.

##  5 Conclusion

Our work addresses the question of how adaptive methods improve the GAN
performance. We empirically showed that the adaptive magnitude of Adam is
responsible for the high performance of adaptive methods in GANs. We further
show that the adaptive magnitude is almost constant in some settings (for
instance in the case of WGAN-GP), and nSGDA is able to recover the learning
dynamics of Adam. We constructed a setting where we proved that nSGDA –thanks
to its balanced updates– recovers the modes of the true distribution while
SGDA fails to do it. Our theory provides insights on the effectiveness of
adaptive methods in architectures such as DCGAN or WGAN-GP, however the
limitation of our work is that we do not fully capture all of the components
in modern architectures such as StyleGAN2. An exciting direction would be to
find an optimization model that recovers the same perform for Adam on these
modern GAN architectures.

## References

  * Agarwal et al. (2020) Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang.  Disentangling adaptive gradient methods from learning rates.  _arXiv preprint arXiv:2002.11803_ , 2020.
  * Allen-Zhu & Li (2020) Zeyuan Allen-Zhu and Yuanzhi Li.  Towards understanding ensemble, knowledge distillation and self-distillation in deep learning.  _CoRR_ , abs/2012.09816, 2020.  URL https://arxiv.org/abs/2012.09816.
  * Allen-Zhu & Li (2021) Zeyuan Allen-Zhu and Yuanzhi Li.  Forward super-resolution: How can gans learn hierarchical generative models for real-world distributions.  _arXiv preprint arXiv:2106.02619_ , 2021.
  * Antonakopoulos et al. (2019) Kimon Antonakopoulos, Veronica Belmega, and Panayotis Mertikopoulos.  An adaptive mirror-prox method for variational inequalities with singular operators.  _Advances in Neural Information Processing Systems_ , 32:8455–8465, 2019.
  * Antonakopoulos et al. (2020) Kimon Antonakopoulos, E Veronica Belmega, and Panayotis Mertikopoulos.  Adaptive extra-gradient methods for min-max optimization and games.  _arXiv preprint arXiv:2010.12100_ , 2020.
  * Arjovsky et al. (2017) Martin Arjovsky, Soumith Chintala, and Léon Bottou.  Wasserstein generative adversarial networks.  In _International conference on machine learning_ , pp. 214–223. PMLR, 2017.
  * Arora et al. (2017) Sanjeev Arora, Rong Ge, Yingyu Liang, Tengyu Ma, and Yi Zhang.  Generalization and equilibrium in generative adversarial nets (gans).  In _International Conference on Machine Learning_ , pp. 224–232. PMLR, 2017.
  * Arora et al. (2018) Sanjeev Arora, Andrej Risteski, and Yi Zhang.  Do gans learn the distribution? some theory and empirics.  In _International Conference on Learning Representations_ , 2018.
  * Bach & Levy (2019) Francis Bach and Kfir Y Levy.  A universal algorithm for variational inequalities adaptive to smoothness and noise.  In _Conference on Learning Theory_ , pp. 164–194. PMLR, 2019.
  * Bai et al. (2018) Yu Bai, Tengyu Ma, and Andrej Risteski.  Approximability of discriminators implies diversity in gans.  _arXiv preprint arXiv:1806.10586_ , 2018.
  * Balaji et al. (2021) Yogesh Balaji, Mohammadmahdi Sajedi, Neha Mukund Kalibhat, Mucong Ding, Dominik Stöger, Mahdi Soltanolkotabi, and Soheil Feizi.  Understanding overparameterization in generative adversarial networks.  _arXiv preprint arXiv:2104.05605_ , 2021.
  * Barazandeh et al. (2021) Babak Barazandeh, Davoud Ataee Tarzanagh, and George Michailidis.  Solving a class of non-convex min-max games using adaptive momentum methods.  In _ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 3625–3629. IEEE, 2021.
  * Brock et al. (2018) Andrew Brock, Jeff Donahue, and Karen Simonyan.  Large scale gan training for high fidelity natural image synthesis.  _arXiv preprint arXiv:1809.11096_ , 2018.
  * Choi et al. (2019) Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl.  On empirical comparisons of optimizers for deep learning.  _arXiv preprint arXiv:1910.05446_ , 2019.
  * Coates et al. (2011) Adam Coates, Andrew Ng, and Honglak Lee.  An analysis of single-layer networks in unsupervised feature learning.  In _Proceedings of the fourteenth international conference on artificial intelligence and statistics_ , pp. 215–223. JMLR Workshop and Conference Proceedings, 2011.
  * Cortés (2006) Jorge Cortés.  Finite-time convergent gradient flows with applications to network consensus.  _Automatica_ , 42(11):1993–2000, 2006.
  * Cutkosky & Mehta (2020) Ashok Cutkosky and Harsh Mehta.  Momentum improves normalized sgd.  In _International Conference on Machine Learning_ , pp. 2260–2268. PMLR, 2020.
  * Daskalakis et al. (2017) Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, and Haoyang Zeng.  Training gans with optimism.  _arXiv preprint arXiv:1711.00141_ , 2017.
  * Dumoulin et al. (2016) Vincent Dumoulin, Ishmael Belghazi, Ben Poole, Olivier Mastropietro, Alex Lamb, Martin Arjovsky, and Aaron Courville.  Adversarially learned inference.  _arXiv preprint arXiv:1606.00704_ , 2016.
  * Feizi et al. (2017) Soheil Feizi, Farzan Farnia, Tony Ginart, and David Tse.  Understanding gans: the lqg setting.  _arXiv preprint arXiv:1710.10793_ , 2017.
  * Fiez & Ratliff (2020) Tanner Fiez and Lillian Ratliff.  Gradient descent-ascent provably converges to strict local minmax equilibria with a finite timescale separation.  _arXiv preprint arXiv:2009.14820_ , 2020.
  * Gasnikov et al. (2019) AV Gasnikov, PE Dvurechensky, FS Stonyakin, and AA Titov.  An adaptive proximal method for variational inequalities.  _Computational Mathematics and Mathematical Physics_ , 59(5):836–841, 2019.
  * Goodfellow et al. (2014) Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  _Advances in neural information processing systems_ , 27, 2014.
  * Gorbunov et al. (2022) Eduard Gorbunov, Marina Danilova, David Dobre, Pavel Dvurechensky, Alexander Gasnikov, and Gauthier Gidel.  Clipped stochastic methods for variational inequalities with heavy-tailed noise.  _arXiv preprint arXiv:2206.01095_ , 2022.
  * Gulrajani et al. (2017) Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville.  Improved training of wasserstein gans.  _arXiv preprint arXiv:1704.00028_ , 2017.
  * Hazan et al. (2015) Elad Hazan, Kfir Y Levy, and Shai Shalev-Shwartz.  Beyond convexity: Stochastic quasi-convex optimization.  _arXiv preprint arXiv:1507.02030_ , 2015.
  * Heusel et al. (2017) Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  _Advances in neural information processing systems_ , 30, 2017.
  * Jin et al. (2020) Chi Jin, Praneeth Netrapalli, and Michael Jordan.  What is local optimality in nonconvex-nonconcave minimax optimization?  In _International Conference on Machine Learning_ , pp. 4880–4889. PMLR, 2020.
  * Karras et al. (2019) Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 4401–4410, 2019.
  * Karras et al. (2020) Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila.  Analyzing and improving the image quality of stylegan.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 8110–8119, 2020.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Lei et al. (2020) Qi Lei, Jason Lee, Alex Dimakis, and Constantinos Daskalakis.  Sgd learns one-layer networks in wgans.  In _International Conference on Machine Learning_ , pp. 5799–5808. PMLR, 2020.
  * Levy (2016) Kfir Y Levy.  The power of normalization: Faster evasion of saddle points.  _arXiv preprint arXiv:1611.04831_ , 2016.
  * Li & Dou (2020) Yuanzhi Li and Zehao Dou.  Making method of moments great again?–how can gans learn distributions.  _arXiv preprint arXiv:2003.04033_ , 2020.
  * Liang (2017) Tengyuan Liang.  How well can generative adversarial networks learn densities: A nonparametric view.  _arXiv preprint arXiv:1712.08244_ , 2017.
  * Liu et al. (2019) Mingrui Liu, Youssef Mroueh, Jerret Ross, Wei Zhang, Xiaodong Cui, Payel Das, and Tianbao Yang.  Towards better understanding of adaptive gradient algorithms in generative adversarial nets.  _arXiv preprint arXiv:1912.11940_ , 2019.
  * Mescheder et al. (2018) Lars Mescheder, Andreas Geiger, and Sebastian Nowozin.  Which training methods for gans do actually converge?  In _International conference on machine learning_ , pp. 3481–3490. PMLR, 2018.
  * Nesterov (1984) Y.E. Nesterov.  Minimization methods for nonsmooth convex and quasiconvex functions.  _Econ. Mat. Met._ , 20:519–531, 01 1984.
  * Orvieto et al. (2021) Antonio Orvieto, Jonas Kohler, Dario Pavllo, Thomas Hofmann, and Aurelien Lucchi.  Vanishing curvature and the power of adaptive methods in randomly initialized deep networks, 2021.
  * Radford et al. (2015) Alec Radford, Luke Metz, and Soumith Chintala.  Unsupervised representation learning with deep convolutional generative adversarial networks.  _arXiv preprint arXiv:1511.06434_ , 2015.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Advances in neural information processing systems_ , pp. 5998–6008, 2017.
  * Yu et al. (2016) Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao.  Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop, 2016.
  * Zbontar et al. (2018) Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew J Muckley, Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, et al.  fastmri: An open dataset and benchmarks for accelerated mri.  _arXiv preprint arXiv:1811.08839_ , 2018.
  * Zhang et al. (2019) Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J Reddi, Sanjiv Kumar, and Suvrit Sra.  Why are adaptive methods good for attention models?  _arXiv preprint arXiv:1912.03194_ , 2019.
  * Zhang et al. (2017) Pengchuan Zhang, Qiang Liu, Dengyong Zhou, Tao Xu, and Xiaodong He.  On the discrimination-generalization tradeoff in gans.  _arXiv preprint arXiv:1711.02771_ , 2017.
  * Zhou et al. (2018) Dongruo Zhou, Jinghui Chen, Yuan Cao, Yiqi Tang, Ziyan Yang, and Quanquan Gu.  On the convergence of adaptive gradient methods for nonconvex optimization.  _arXiv preprint arXiv:1808.05671_ , 2018.

##  Appendix A Additional experiments

###  A.1 Experiments with StyleGAN2 and WGAN-GP

In this section, we put additional curves and images produced by WGAN and
StyleGAN2.

(a)

(b)

Figure 6:  (a) is the FID curve of StyleGAN2 on LSUN-Churches and (b) the FID
curve of WGAN on STL-10. These complement the figures of Section 4. Figure 7:
Images generated by a StyleGAN2 model trained with Adam for 2600 kimgs on FFHQ
128. Note that this is not convergence. Figure 8: Images generated by a
StyleGAN2 model trained with Ada-nSDGDA for 2600 kimgs on FFHQ 128. Note that
this is not convergence. Figure 9: Images generated by a StyleGAN2 model
trained with Adam for 2600 kimgs on LSUN Churches 128. Note that this is not
convergence. Figure 10: Images generated by a StyleGAN2 model trained with
Ada-nSDGDA for 2600 kimgs on LSUN Churches 128. Note that this is not
convergence.

###  A.2 Experiments with DCGAN

This section shows that experimental results obtained in Section 4 are also
valid for other architectures such as DCGAN. Indeed, we observe that nSGDA
methods compete with Adam and nSGDA work when the batch size is small. In this
section, lnSGDA refers to the layer-wise nSGDA and gnSGDA to the global nSGDA.

(a) CIFAR-10

(b) LSUN Churches

(c) STL-10

(d) Celeba-HQ

Figure 11: FID scores obtained when training a Resnet WGAN-GP using Adam,
l-nSGDA, g-nSGDA, and SGD on different datasets. In all these datasets,
l-nSGDA, g-nSGDA and Adam perform approximately as well. SGDA performs much
worse.

(a) CIFAR-10

(b) LSUN Churches

(c) STL-10

(d) Celeba-HQ

Figure 12: FID scores obtained when training a DCGAN using Adam, lnSGDA,
gnSGDA and SGD on different datasets. In all these datasets, lnSGDA, gnSGDA
and Adam perform approximately as well. As expected, SGDA performs much worse
than the other optimizers. The models are trained with batch-size 64 –which is
the usual batch-size used for DCGAN.

In this section, we display the images obtained when training the Resnet WGAN-
GP model from Section 4.

(a) Adam

(b) lnSGDA

(c) gnSGDA

Figure 13: CIFAR-10 images generated by a Resnet WGAN-GP model

(a) Adam

(b) lnSGDA

(c) gnSGDA

Figure 14: LSUN-Churches images generated by a Resnet WGAN-GP model

(a) Adam

(b) lnSGDA

(c) gnSGDA

Figure 15: STL-10 images generated by a Resnet WGAN-GP model

(a) Adam

(b) lnSGDA

(c) gnSGDA

Figure 16: Celeba-HQ images generated by a Resnet WGAN-GP model

##  Appendix B Technical statements in the theory section

In this section, we provide the technical version of the statements made in
Section 3.

###  B.1 Setting

The distribution $p_{data}$ we consider is more general than Assumption 1 in
the main paper.

######  Assumption 3 (Data structure).

Let $\gamma=\frac{1}{\mathrm{polylog}(d)}$. The coefficients $s_{1},s_{2}$ and
modes $u_{1},u_{2}$ of the distribution $p_{data}$ respect one of the
following conditions:

  1. 1.

Correlated modes: $\langle u_{1},u_{2}\rangle=\gamma$ and the generated data
point is either $X=u_{1}$ or $X=u_{2}.$

  2. 2.

Correlated coefficients: $\mathbb{P}[s_{1}=s_{2}=1]=\gamma$ and the modes are
orthogonal, $\langle u_{1},u_{2}\rangle=0$.

We now present a more technical version of Assumption 2.

######  Assumption 4 ($p_{z}$ structure).

Let $z\sim p_{z}$. We assume that for $i,j\in[m_{G}]$,

$\displaystyle\Pr[z_{i}=1]$
$\displaystyle=\Theta\left(\frac{1}{m_{G}}\right),\;\Pr[z_{i}=z_{j}=1]$
$\displaystyle=\frac{1}{m_{G}^{2}\mathrm{polylog}(d)}$ (11)

We set $\Pr[z_{i}=z_{j}=1]=\frac{1}{m_{G}^{2}\mathrm{polylog}(d)}$ to ensure
that that the output of the generator is only made of one mode with
probability $1-o(1)$.

In the proof, we actually consider a more complicated version of the
discriminator

$\displaystyle D_{\mathcal{W}}(X)$
$\displaystyle=\mathrm{sigmoid}\left(a\sum_{i\in[m_{D}]}\sigma(\langle
w_{i},X\rangle)+\lambda b\right),\;\text{where }\sigma(z)$
$\displaystyle=\begin{cases}z^{3}&\text{if }|z|\leq\Lambda\\\
3\Lambda^{2}z-2\Lambda^{3}&\text{if }z>\Lambda\\\
3\Lambda^{2}z+2\Lambda^{3}&\text{otherwise}\end{cases},$ (12)

where $\Lambda=d^{0.2}$. $\sigma(\cdot)$ is the truncated degree-3 activation
function—it is thus made Lipschitz, which is only needed in the proof to
_deal_ with the case where the generator is trained much faster than the
discriminator. Note that this latter case is uncommon in practice.

We now present the technical version of 3.1.

######  Parametrization B.1.

When running SGDA and nSGDA on GAN, we set the parameters as

– Initialization:
$b^{(0)}=0,\;a^{(0)}\sim\mathcal{N}\left(0,\frac{1}{m_{D}\mathrm{polylog}(d)}\right),\;w_{i}^{(0)}\sim\mathcal{N}\left(0,\frac{1}{d}\mathbf{I}\right),\;v_{j}^{(0)}\sim\mathcal{N}\left(0,\frac{1}{d^{2}}\mathbf{I}\right)$
for $i\in[m_{D}]$, $j\in[m_{G}].$

– Number of iterations: we run SGDA for $t\leq T_{0}$ iterations where $T_{0}$
is the first iteration such that
$\|\nabla\mathbb{E}[L_{\mathcal{V}^{(T_{0})},\mathcal{W}^{(T_{0})}}(X,z)]\|_{2}\leq
1/\mathrm{poly}(d)$. . For nSGDA, we run for
$T_{1}=\tilde{\Theta}\left(\frac{1}{\eta_{D}}\right)$iterations.

– Step-sizes: For SGDA, $\eta_{D},\eta_{G}\in(0,\frac{1}{\mathrm{poly}(d)})$.
For nSGDA, $\eta_{D}\in(0,\frac{1}{\mathrm{poly}(d)}]$,
$\eta_{G}=\frac{\eta_{D}}{\mathrm{polylog}(d)}.$

– Over-parametrization: For SGDA, $m_{D},m_{G}=\mathrm{polylog}(d)$ are
arbitrarily chosen i.e. $m_{D}$ may be larger than $m_{G}$ or the opposite.
For nSGDA, we set $m_{D}=\log(d)$ and $m_{G}=\log(d)\log\log d$.

Regarding initialization, the discriminator’s weights are sampled from a
standard normal and its bias is set to zero. The weights of the generator are
initialized from a normal with variance $1/d^{2}$ (instead of the $1/d$ in
standard normal). Such a choice is explained as follows. In practice, the
target $X\sim p_{data}$ is a 1D image, thus has entries in $[0,1]^{d}$ and
norm $O(\sqrt{d}).$ Yet, we sample the initial generator’s weights from
$\mathcal{N}(0,\mathbf{I}_{d}/d)$ in this case. In our case, since
$\|u_{i}\|_{2}=1,$ the target $X=s_{1}u_{1}+s_{2}u_{2}$ has norm $O(1)$.
Therefore, we scale down the variance in the normal distribution by a factor
of $1/d$ to match the configuration encountered in practice. Therefore, we
also set $\lambda=\frac{1}{\sqrt{d}\mathrm{polylog}(d)}$ in (12) to ensure
that the weights and the bias in the discriminator learn at the same speed.

Remark: In our theory, we consider the global version of nSGDA;
$\|\mathbf{g}_{\mathcal{W}}^{(t)}\|_{2}$ in the update refers to
$\|\mathbf{g}_{\mathcal{W}}^{(t)}\|_{2}=\|\mathbf{g}_{a}^{(t)}\|_{2}+\|\mathbf{g}_{b}^{(t)}\|_{2}+\|\mathbf{g}_{W}^{(t)}\|_{2}$,
where $\mathbf{g}_{a}^{(t)}$ is the stochastic gradient with respect to $a$,
$\mathbf{g}_{b}^{(t)}$ with respect to $b$ and $\mathbf{g}_{W}^{(t)}$ with
respect to $W.$

###  B.2 Main results

We state now the technical version of Theorem 3.1 and Theorem 3.2.

######  Theorem B.1 (SGDA suffers from mode collapse).

Let $T_{0}$, $\eta_{G},\eta_{D}$ and the initialization as defined in B.1. Let
$t$ be such that $t\leq T_{0}.$ Run SGDA for $t$ iterations with step-sizes
$\eta_{G},\eta_{D}$. Then, with probability at least $1-o(1)$, for all
$z\in\\{0,1\\}^{m_{G}}$, we have:

$\displaystyle
G_{\mathcal{V}}^{(t)}(z)=\alpha^{(t)}(z)(u_{1}+u_{2})+\xi^{(t)}(z),$

where $\alpha^{(t)}(z)\in\mathbb{R}\;\text{ and
}\;\xi^{(t)}(z)\in\mathbb{R}^{d}$ and for all $\ell\in[2]$,
$|\langle\xi^{(t)}(z),u_{\ell}\rangle|=o(1)\|\xi^{(t)}(z)\|_{2}$ for every
$z\in\\{0,1\\}^{m_{G}}$.

In the specific case where
$\eta_{D}=\frac{\sqrt{d}\eta_{G}}{\mathrm{polylog}(d)}$, the model mode
collapses i.e. $\|\xi^{(T_{0})}(z)\|_{2}=o(\alpha^{(T_{0})}(z))$.

Theorem 3.1 indicates that with SGDA and any step-size configuration, the
generator either does not learn the modes at all – when
$\alpha^{(t)}(z)=0,\;G_{\mathcal{V}}^{(t)}(z)=\xi^{(t)}(z)$ – or learns an
average of the modes – when $\alpha^{(t)}(z)\neq
0,\;G_{\mathcal{V}}^{(t)}(z)\approx\alpha^{(t)}(z)(u_{1}+u_{2})$.

######  Theorem B.2 (nSGDA recovers modes separately).

Let $T_{1}$, $\eta_{G},\eta_{D}$ and the initialization as defined in B.1. Run
nSGDA for $T_{1}$ iterations with step-sizes $\eta_{G},\eta_{D}$. Then, the
generator learns both modes $u_{1},u_{2}$ i.e., for $\ell\in\\{1,2\\}$

$\Pr{}_{z\sim
p_{z}}\left(\big{\|}\tfrac{G_{\mathcal{V}}^{(T_{1})}(z)}{\|G_{\mathcal{V}}^{(T_{1})}(z)\|_{2}}-u_{\ell}\big{\|}_{2}=o(1)\right)=\tilde{\Omega}(1).$

Normalized GD vs GD in GAN training

##  Appendix C Notations

Let us also write $\tau_{b}=\lambda$ as the scaling factor of the bias. We can
easily observe that at every step, all of $w_{i}^{(t)}$ and $v_{i}^{(t)}$ lies
in the span of $\\{w_{j}^{(0)},v_{j}^{(0)},u_{1},u_{2}\\}$. Therefore, let us
denote

$w_{i}^{(t)}=\sum_{j\in[m_{D}]}\alpha(w_{i},w_{j},t)\frac{w_{j}^{(0)}}{\|w_{j}^{(0)}\|_{2}}+\sum_{j\in[m_{G}]}\alpha(w_{i},v_{j},t)\frac{v_{j}^{(0)}}{\|v_{j}^{(0)}\|_{2}}+\sum_{j\in[2]}\alpha(w_{i},u_{j},t)\frac{u_{j}}{\|u_{j}\|}$

and $v_{i}^{(t)}$ as $\alpha(v_{i},*,t)$, where $\alpha(*,*,*)\in\mathbb{R}$.

Let us denote

$f(X)=a\left(\sum_{i\in[m_{D}]}\sigma(\langle
w_{i},X\rangle)\right)+\tau_{b}b$

as the function in discriminator without going through sigmoid, and define
$h(X)=\sum_{i\in[m_{D}]}\sigma(\langle w_{i},X\rangle)$.

#### Gradient

The gradient of $L(X,z)$ is given as:

$\nabla_{a}L(X,z)=-\operatorname{Sigmoid}(-f(X))h(X)+\operatorname{Sigmoid}(f(G(z)))h(G(z))$
$\nabla_{b}L(X,z)=-\operatorname{Sigmoid}(-f(X))+\operatorname{Sigmoid}(f(G(z)))$
$\nabla_{w_{i}}L(X,z)=-\operatorname{Sigmoid}(-f(X))a\sigma^{\prime}(\langle
w_{i},X\rangle)X+\operatorname{Sigmoid}(f(G(z)))a\sigma^{\prime}(\langle
w_{i},G(z)\rangle)G(z)$
$\nabla_{v_{i}}L(X,z)=-1_{z_{i}=1}\operatorname{Sigmoid}(f(G(z)))a\sum_{j\in[m_{D}]}\sigma^{\prime}(\langle
w_{i},G(z)\rangle)w_{i}$

We use $a^{(t)},b^{(t)},w_{i}^{(t)},v_{i}^{(t)}$ to denote the value of those
weights at iteration $t$.

We use $a=b\pm c$ for $c\in\mathbb{R}^{*}$ to denote: (1). $a\in[b-c,b+c]$ if
$a,b\in\mathbb{R}$, (2). $\|a-b\|_{2}\leq c$ if $a,b$ are vectors.

For simplicity, we focus on the case when all $\Pr[z_{i}=1]$ are equal. The
other cases can be proved similarly (by replacing the $1/m_{G}$ factor in the
generators update by the exact value of $\Pr[z_{i}=1]$).

##  Appendix D Initialization conditions and three regime of learning

We first show the following Lemma regarding initialization:

Let

$A_{i,\ell}=\frac{1}{2}\sigma^{\prime}(\langle
w_{i}^{(0)},u_{\ell}\rangle)\operatorname{sign}(\langle
w_{i}^{(0)},u_{\ell}\rangle)$

and

$B_{i,j}=\frac{1}{m_{G}}\sigma^{\prime}(\langle
w_{i}^{(0)},v_{j}^{(0)}\rangle)\operatorname{sign}(\langle
w_{i}^{(0)},v_{j}^{(0)}\rangle)$

and

$C_{i,\ell}=\sigma^{\prime}(\langle v_{i}^{(0)},u_{\ell}\rangle)$

Let $A=\max_{i\in[m_{D}],\ell\in[2]}A_{i,\ell}$,
$B=\max_{i\in[m_{D}],j\in[m_{G}]}B_{i,j}$,
$C=\max_{i\in[m_{G}],\ell\in[2]}C_{i,\ell}$, we have: Using a corollary of
Proposition $G.1$ in Allen-Zhu & Li (2020):

######  Lemma D.1.

For every $\eta_{D},\eta_{G}>0$, we have that: with probability at least
$1-o(1)$, we have that:
$A=\frac{\operatorname{polyloglog}(d)}{\sqrt{d}},B=\frac{\operatorname{polyloglog}(d)}{dm_{G}}$.
Moreover, with probability at least $1-o(1)$, one and only one of the
following holds:

  1. 1.

(Discriminator trains too fast):
$\eta_{G}B<\frac{1}{\operatorname{polylog}(d)}\eta_{D}A$;

  2. 2.

(Balanced discriminator and generator):
$\eta_{G}B>\frac{1}{\operatorname{polylog}(d)}\eta_{D}A$,
$\eta_{D}A>\eta_{G}B(1+\frac{1}{\operatorname{polyloglog}(d)})$;

  3. 3.

(Generator trains too fast):
$\eta_{D}A<\eta_{G}B(1-\frac{1}{\operatorname{polyloglog}(d)})$.

This Lemma implies that in case 2,
$\eta_{G}=\tilde{\Theta}(\sqrt{d})\eta_{B}$.

We will show the following induction hypothesis for each case. Intuitively, in
case one we have the following learning process: (too powerful $D$).

  1. 1.

At first $D$ starts to learn, then because of the learning rate of $G$ is too
small, so $D$ just saturate the loss to make the gradient to zero.

In case two we have: (“balanced” $D$ and $G$ but still not enough).

  1. 1.

At first $D$ starts to learn one $u_{j}$ in each of the neuron.

  2. 2.

However, the generator still could not catch up immediate after $D$ learns one
$u_{j}$, so $D$ starts to a mixture of $u_{1},u_{2}$ in its neurons since
$u_{1},u_{2}$ are positively correlated.

  3. 3.

After that $G$ starts to learn, however since $D$ already stuck at the
mixtures of $u_{1},u_{2}$, so $G$ is only able to learn mixtures of
$u_{1},u_{2}$ as well.

In case three we have: (Too powerful $G$)

  1. 1.

$G$ starts to learn without $D$ learning any meaningful signal yet, so $G$
aligns its outputs with the (close to random) weights of $D$ and just pushes
the discriminator to zero. In this case, $G$ simply learns something random to
fool $D$ instead of learning the signals.

Moreover, similar to Lemma D.1, we also have the following condition regarding
the gap between the top one and the second largest one in terms of
correlation:

######  Lemma D.2.

Let

$i_{D},\ell_{D}=\operatorname*{arg\,max}_{i\in[m_{D}],\ell\in[2]}A_{i,\ell}$

Let

$i_{G},j_{G}=\operatorname*{arg\,max}_{i\in[m_{D}],j\in[m_{G}]}B_{i,j}$

Then with probability at least $1-o(1)$ over the random initialization, the
following holds:

$\forall i,\ell\not=i_{D},\ell_{D}:A_{i_{D},\ell_{D}}\geq
A_{i,\ell}\left(1+\frac{1}{\operatorname{polyloglog}(d)}\right)$ $\forall
i,j\not=i_{G},j_{G}:B_{i_{G},j_{G}}\geq
B_{i,j}\left(1+\frac{1}{\operatorname{polyloglog}(d)}\right)$

and

$A=\frac{\Theta(\sqrt{\log\log(d)})}{\sqrt{d}},\quad
B,C=\frac{\Theta(\sqrt{\log\log(d)})}{d}$

For simplicity, we also define $i^{*}=i_{D}$.

##  Appendix E Critical Lemma

The proof heavily relies on the following Lemma about tensor power method,
which is a corollary of Lemma C.19 in Allen-Zhu & Li (2020).

######  Lemma E.1.

For every $\delta\in(0,0.1)$, every $C>10$ , for every sequence of
$x_{t},y_{t}>0$ such that $x_{0}>(1+10\delta)y_{0}$, suppose there is a
sequence of $S_{t}\in[0,C]$ such that for
$\eta\in\left(0,\frac{1}{\operatorname{poly}(C/\delta)}\right)$:

$x_{t+1}\geq x_{t}+\eta S_{t}x_{t}^{2}$ $y_{t}\leq y_{t}+\eta
S_{t}(1+\delta)y_{t}^{2}$

For every $\tau>0$, let $T_{0}$ be the first iteration where $x_{t}>\tau$,
then we must have:

$y_{T_{0}}\leq\frac{y_{0}}{\operatorname{poly}(\delta)}$

Moreover, if all $S_{t}\geq H$ for some $H>0$, then $T_{0}\leq
O\left(\frac{1}{\eta Hx_{0}}\right)$.

Similar to the Lemma above, one can easily show the following auxiliary Lemma:

######  Lemma E.2.

Suppose there are sequences $a_{t},b_{t}\in\mathbb{R}^{d}$ such that
$a_{0},b_{0}>0$ with $a_{0}<0.82b_{0}$. Suppose there exists a sequence of
$C_{t}\in(0,d)$ such that

$a_{t+1}\leq a_{t}-\eta_{D}C_{t}b_{t}$ $b_{t+1}\geq
b_{t}-1.0000001\eta_{D}C_{t}a_{t}$

Then we must have that for every $t\leq T$ where $T$ is the first iteration
such that $a_{T}\leq 0$, then the following holds:

$a_{t}=a_{0}-\Theta\left(\eta\sum_{s\leq t-1}C_{s}\right)$ $\sum_{s\leq
t}|a_{t}C_{t}\eta_{D}|\leq 0.49b_{0}$

Moreover, if in addition that $a_{0}<\frac{1}{C}b_{0}$ for any $C>100$, then
we must have:

$\sum_{s\leq t}|a_{t}C_{t}\eta_{D}|\leq\frac{10}{C}b_{0}$

In the end, we have the following comparison Lemma, whose proof is obvious:

######  Lemma E.3.

Suppose $a_{t},b_{t}>0$ satisfies that $a_{0},b_{0}\leq 1$, and the update of
$a_{t},b_{t}$ is given as: For some values $C>0$ and
$C_{t}\in[0,\operatorname{poly}(d)]$:

$\displaystyle a_{t+1}$ $\displaystyle=a_{t}+\eta_{D}C_{t}$ (13)
$\displaystyle b_{t+1}$
$\displaystyle=b_{t}+\eta_{D}\left[\frac{1}{C},1\right]\times C_{t}$ (14)

Then let $T$ be the first iteration where $a_{T}\geq 2C$, we must have:

$b_{T}\in\left[1,2C+1\right]$

Using this Lemma, we can directly prove the following Lemma:

######  Lemma E.4.

For every
$\eta_{D},\eta_{G}\in\left(0,\frac{1}{\operatorname{poly}(d)}\right]$ such
that $\eta_{G}=\eta_{D}\Gamma$ for $\Gamma=\tilde{\Theta}(\sqrt{d})$, suppose
there are vectors $p_{t},q_{i,t}\in\mathbb{R}^{d}$ ($i\in[m_{G}]$) and a value
$a_{t}\in\mathbb{R},H>0$ satisfies that for a sequence of $H_{i,t}\in[H,1]$
for $i\in[m_{G}]$, $G_{t}=\tilde{\Theta}(\sum_{i\in[m_{G}]}H_{i,t})$, a value
$\tau=\tilde{O}(d^{-0.5})$, and a vector
$\beta_{t}\in\text{span}\\{u_{1},u_{2}\\}$ with $\|\beta_{t}\|_{2}=O(1)$: For
all $i\in[m_{G}]$ and $t\geq 0$:

$\|q_{i,0}\|_{2}=\tilde{\Theta}(d^{-0.49}),\|p_{0}\|_{2}=\log^{\Theta(1)}(d),0<a_{0}\leq
0.819\|p_{0}\|_{2}\quad\frac{\langle
q_{i,0},p_{0}\rangle}{\|q_{i,0}\|_{2},\|p_{i,0}\|_{2}}\geq 1-o(1)$

$p_{t}=p_{t}-\eta_{D}\sum_{i\in[m_{G}]}G_{i,t}a_{t}\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)q_{i,t}+\tilde{O}(\eta_{D}a_{t}\gamma_{t})\beta_{t}$
$a_{t}=a_{t}-\eta_{D}\sum_{i\in[m_{G}]}G_{i,t}\sigma(\langle
p_{t},q_{i,t}\rangle)\pm\tilde{O}(\eta_{D}\gamma_{t})$
$q_{i,t}=\left(q_{i,t}+\eta_{G}H_{i,t}a_{t}\left(\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)+\sum_{j\in[m_{G}]}\gamma_{i,j,t}\sigma^{\prime}(\langle
p_{t},q_{j,t}\rangle)\right)p_{t}\pm\eta_{G}|a_{t}|\tilde{O}\left(\tau\|q_{i,t}\|_{2}\right)^{2}\right)$

In addition, we have: $\gamma_{i,j,t}=\tilde{O}(1)$, and

$\max_{i\in[m_{G}]}\|q_{i,t}\|_{2}\in\left(0,\frac{1}{\operatorname{polylog}(d)}\right]\cup[\operatorname{polylog}(d),+\infty)\implies\forall
i,j\in[m_{G}],H_{i,t}=\tilde{\Theta}(G_{t}),\gamma_{i,j,t}=\tilde{\Theta}(1)$

Then we must have that: let $T$ be the first iteration where $a_{T}\leq 0$, we
have: for every $t\leq T$: there is a scaling factor $\ell_{t}=\Theta(1)$ such
that

$\|p_{t}-\ell_{t}p_{0}\|_{2}\leq
o(1)\|p_{0}\|_{2},\quad\|\Pi_{\text{span}\\{u_{1},u_{2},p_{0}\\}^{\bot}}(p_{t}-p_{0})\|_{2}\leq
d^{-0.6}\|p_{0}\|_{2}$

Moreover, for every $i,j\in[m_{G}]$,
$\|q_{i,t}\|_{2}=\tilde{\Theta}(\|q_{j,t}\|_{2})$ and
$\|q_{i,T}\|_{2}\geq\tilde{\Theta}(\sqrt{\Gamma})$, and as long as
$\max_{i\in[m_{G}]}\|q_{i,t}\|_{2}\geq\operatorname{polylog}(d)$, we have that
$a_{t}\|q_{i,t}\|_{2}\geq\operatorname{polylog}(d)$.

Moreover,

$\|\Pi_{\text{span}\\{u_{1},u_{2},p_{0}\\}^{\bot}}(q_{i,t}-q_{i,0})\|_{2}\leq
d^{-0.6}\|q_{i,t}\|_{2}$

###### proof of Lemma E.4.

For simplicity we consider the case when $H=\tilde{\Omega}(1)$, the other
cases follow similarly.

To proof this result, we maintain the following decomposition of $p_{t}$ and
$q_{i,t}$ as:

$p_{t}=\alpha(t)p_{0}+\beta(t)+\gamma(t)$

Where $\beta(t)\in\text{span}\\{u_{1},u_{2}\\}$ and
$\gamma(t)\bot\text{span}\\{u_{1},u_{2},p_{0}\\}$. Note that
$\alpha(0)=1,\beta(0)=\gamma(0)=0$.

$q_{i,t}=\alpha(i,t)p_{0}+\beta(i,t)+\gamma(i,t)$

Where $\beta(i,t)\in\text{span}{u_{1},u_{2}}$ and
$\gamma(i,t)\bot\text{span}\\{p_{0},u_{1},u_{2}\\}$.

We maintain the following induction hypothesis (which we will prove at the
end): For some $\mu=0.00001$ and $C_{1}=d^{-0.1},C_{2}=d^{-0.6}$, we have:

  1. 1.

Through out the iterations, $\alpha(t)\geq 0.5$ and $\|\beta(t)\|_{2}\leq
0.5(1-\mu)C_{1},\|\gamma(t)\|_{2}\leq 0.5(1-\mu)C_{2}$.

  2. 2.

$\alpha(i,t)\in(0,\tilde{O}\sqrt{\Gamma})$ and $\|\beta(i,t)\|_{2}\leq
C_{1}\alpha(i,t)+\|\beta(i,0)\|_{2},\quad\|\gamma(i,t)\|_{2}\leq
C_{2}\alpha(i,t)+\|\gamma(i,0)\|_{2}$

The induction hypothesis implies that through out the iterations, $\langle
q_{i,t},p_{t}\rangle=\tilde{\Omega}(\|q_{i,t}\|_{2})$.

We can now write down the update of $a_{t},\alpha^{\prime}s,\beta^{\prime}s$
and $\gamma^{\prime}s$ as:

$\displaystyle
a_{t+1}=a_{t}-\eta_{D}\left(\sum_{i\in[m_{G}]}G_{i,t}\sigma(\langle
p_{t},q_{i,t}\rangle)\pm\tilde{O}(1)\right)$ (15)

$\displaystyle\alpha(t+1)=\alpha(t)-\eta_{D}a_{t}\sum_{i\in[m_{G}]}G_{i,t}\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)\alpha(i,t)$ (16)

$\displaystyle\beta(t+1)=\beta(t)-\eta_{D}a_{t}\sum_{i\in[m_{G}]}G_{i,t}\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)\beta(i,t)\pm\tilde{O}(\eta_{D}|a_{t}|)$ (17)

$\displaystyle\gamma(t+1)=\gamma(t)-\eta_{D}a_{t}\sum_{i\in[m_{G}]}G_{i,t}\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)\gamma(i,t)$ (18)

By the induction hypothesis, we know that

$\sigma^{\prime}(\langle
p_{t},q_{j,t}\rangle)\geq\tilde{\Omega}\left(\|q_{j,t}\|_{2}^{2}\times\frac{\Lambda^{2}}{\Gamma}\right)$

Moreover, we have that let $h_{i,t}:=\left(\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)+\sum_{j\in[m_{G}]}\tilde{\Theta}(\sigma^{\prime}(\langle
p_{t},q_{j,t}\rangle))\right)$

$\displaystyle\alpha(i,t+1)=\left(\alpha(i,t)+\eta_{G}H_{i,t}a_{t}h_{i,t}(1\pm\tilde{O}(\tau^{2}\Lambda^{2}/\Gamma)\alpha(t)\right)$
(19)

$\displaystyle\beta(i,t+1)=\left(\beta(i,t)+\eta_{G}H_{i,t}a_{t}h_{i,t}\left(\beta(t)\pm\tilde{O}(\tau^{2}\Lambda^{2}/\Gamma)\right)\right)$
(20)

$\displaystyle\gamma(i,t+1)=\left(\alpha(i,t)+\eta_{G}H_{i,t}a_{t}h_{i,t}\left(\gamma(t)\pm\tilde{O}(\tau^{2}\Lambda^{2}/\Gamma)\right)\right)$
(21)

From these formula, we can easily that as long as (1). $\alpha(t)\geq 0.5$ and
$\|\beta(t)\|_{2}\leq 0.5(1-\mu)C_{1},\|\gamma(t)\|_{2}\leq 0.5(1-\mu)C_{2}$,
(2). $C_{1},C_{2}=\tilde{\Omega}(\tau^{2}\Lambda^{2}/\Gamma)$, we must have
that $\alpha(i,t)>0$ and $\|\beta(i,t)\|_{2}\leq
C_{1}\alpha(i,t)+\|\beta(i,0)\|_{2},\quad\|\gamma(i,t)\|_{2}\leq
C_{2}\alpha(i,t)+\|\gamma(i,0)\|_{2}$. Therefore, it remains to only prove (1)
in the induction hypothesis. Moreover, it is easy to observe that
$\alpha(i,t)=\tilde{\Theta}(\alpha(j,t))$ for all $i,j\in[m_{G}]$ and all $t$.

Now, we divide the update process into two stages:

#### Before all $\|q_{i,t}\|_{2}=\Omega(\Lambda)$. Let’s call these iterations
$[T_{1}]$

Let us consider $T_{i,1}$ such that for all $t\in[T_{i,1}]$ when
$q_{i,t}=O(\Lambda)$ and $a_{t}=\Omega(1)$. In these iterations, by the update
rule, we have

$q_{i,t}=q_{i,t}+\tilde{\Omega}(\eta_{G})\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)p_{t}\pm\tilde{O}(\eta_{G}\tau^{2}\|q_{i,t}\|_{2}^{2})$

By the induction hypothesis, we can simplify the update as:

$\langle q_{i,t},p_{0}\rangle\geq\langle
q_{i,t},p_{0}\rangle+\tilde{\Omega}\left(\eta_{G}\sigma^{\prime}(\langle
q_{i,t},p_{0}\rangle)\right)$

Therefore, we know that
$T_{i,1}\leq\tilde{O}\left(\frac{d^{0.49}}{\eta_{G}}\right)$ and

$\displaystyle\sum_{t\leq T_{i,1}}\sigma^{\prime}(\langle
q_{i,t},p_{0}\rangle),\sum_{t\leq T_{i,1}}\sigma^{\prime}(\langle
q_{i,t},p_{t}\rangle)\leq\tilde{O}\left(\frac{\Lambda}{\eta_{G}}\right)$ (22)

Together with the induction hypothesis, the fact that
$\alpha(i,t)=\tilde{\Theta}(\alpha(j,t))$, the fact that $\sigma(\langle
p_{t},q_{i,t}\rangle)=\tilde{\Theta}(\sigma^{\prime}(\langle
p_{t},q_{i,t}\rangle)\|q_{i,t}\|_{2})$ and update formula Eq equation 15
equation 30 equation 17 equation 18, we know that for all
$t\leq\max\\{T_{i,1}\\}$:

$\displaystyle
a_{t}=a_{0}\pm\tilde{O}\left(\frac{\eta_{D}\Lambda^{2}}{\eta_{G}}\right)=a_{0}\pm\tilde{O}(d^{-0.01})$
(23)
$\displaystyle\alpha(t)=\alpha(0)\pm\tilde{O}\left(\frac{\eta_{D}\Lambda^{2}}{\eta_{G}}\right)=\alpha(0)\pm\tilde{O}(d^{-0.01})$
(24)
$\displaystyle\|\beta(t)\|_{2}\leq\tilde{O}\left(\frac{\eta_{D}\Lambda^{2}}{\eta_{G}}\right)C_{1}+\tilde{O}\left(\frac{\eta_{D}\|\beta(i,0)\|_{2}\Lambda}{\eta_{G}}\right)\leq\tilde{O}(d^{-0.01})C_{1}$
(25)
$\displaystyle\|\gamma(t)\|_{2}\leq\tilde{O}\left(\frac{\eta_{D}\Lambda^{2}}{\eta_{G}}\right)C_{2}+\tilde{O}\left(\frac{\eta_{D}\|\gamma(i,0)\|_{2}\Lambda}{\eta_{G}}\right)\leq\tilde{O}(d^{-0.01})C_{2}$
(26)

#### When all $\|q_{i,t}\|_{2}=\Omega(\Lambda)$:

In this case, since $\|p_{0}\|_{2}=\omega(1)$, we know that $\langle
p_{t},q_{i,t}\rangle=\omega(\Lambda)$, so $\sigma(\langle
p_{t},q_{i,t}\rangle)$ acts on the linear regime, which means that:

$\sigma(\langle p_{t},q_{i,t}\rangle)=(1\pm o(1))3\Lambda^{2}\langle
p_{t},q_{i,t}\rangle,\quad\sigma^{\prime}(\langle p_{t},q_{i,t}\rangle)=(1\pm
o(1))3\Lambda^{2}$

Therefore, we know that:

$\displaystyle a_{t+1}\leq
a_{t}-(1-o(1))\eta_{D}\left(\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\|q_{i,t}\|_{2}\right)\alpha(t)\|p_{0}\|_{2}$
(27)
$\displaystyle\alpha(t+1)\|p_{0}\|_{2}\geq\alpha(t)\|p_{0}\|_{2}-(1+o(1))\eta_{D}\left(\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\|q_{i,t}\|_{2}\right)a_{t}$
(28)

Now, using the fact that $a_{0}\leq 0.819\alpha(0)$ and with Eq equation 23
and Eq equation 24, apply Lemma E.2 we have that until $a_{t}\leq 0$,

$\displaystyle\sum_{t}\eta_{D}\left(\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\|q_{i,t}\|_{2}\right)a_{t}\leq
0.49\|p_{0}\|_{2}$ (29)

Plug in to the update rule:

$\displaystyle\alpha(t+1)$
$\displaystyle=\alpha(t)\pm(1+o(1))\eta_{D}a_{t}\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\alpha(i,t)$
(30)
$\displaystyle=\alpha(t)\pm(1+o(1))\eta_{D}a_{t}\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\frac{\|q_{i,t}\|_{2}}{\|p_{0}\|}$
(31)

$\displaystyle\|\beta(t+1)\|_{2}$
$\displaystyle\leq\|\beta(t)\|_{2}+(1+o(1))\eta_{D}a_{t}\left(\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\beta(i,t)+\tilde{O}(1)\right)$
(32)
$\displaystyle\leq\|\beta(t)\|_{2}+\eta_{D}(1+o(1))a_{t}\left(\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\frac{\|q_{i,t}\|_{2}C_{1}}{\|p_{0}\|_{2}}+\tilde{O}(1)\right)$
(33)
$\displaystyle\leq\|\beta(t)\|_{2}+\eta_{D}(1+o(1))a_{t}\left(\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\frac{\|q_{i,t}\|_{2}C_{1}}{\|p_{0}\|_{2}}\right)$
(34)

$\displaystyle\|\gamma(t+1)\|_{2}$
$\displaystyle\leq\|\gamma(t)\|_{2}+\eta_{D}(1+o(1))a_{t}\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\gamma(i,t)$
(35)
$\displaystyle\leq\|\gamma(t)\|_{2}+\eta_{D}(1+o(1))a_{t}\sum_{i\in[m_{G}]}G_{i,t}3\Lambda^{2}\frac{\|q_{i,t}\|_{2}C_{2}}{\|p_{0}\|_{2}}$
(36)

We directly complete the proof of the induction hypothesis using Eq equation
29.

Now it remains to prove that $\|q_{i,T}\|_{2}=\Omega(\sqrt{\Gamma})$. Compare
the update rule of $q_{i,t}$ and $a_{t}$ we have:

$\displaystyle
a_{t+1}=a_{t}-\tilde{\Theta}(\eta_{D})G_{t}\left(\sum_{i\in[m_{G}]}\Lambda^{2}\|q_{i,t}\|_{2}\right)$
(37)

and

$\displaystyle\sum_{i\in[m_{G}]}\|q_{i,t+1}\|_{2}=\|q_{i,t}\|_{2}+\tilde{\Theta}(\eta_{G})G_{t}\Lambda^{2}a_{t}$
(38)

We can directly conclude that
$\|q_{i,t}\|_{2}\leq\tilde{O}\left(\sqrt{\Gamma}\right)$ and
$\|q_{i,T}\|_{2}=\tilde{\Omega}(\sqrt{\Gamma})$ .

∎

######  Lemma E.5.

For every
$\eta_{D},\eta_{G}\in\left(0,\frac{1}{\operatorname{poly}(d)}\right]$ such
that $\eta_{G}=\eta_{D}\Gamma$ for $\Gamma\geq\tilde{\Omega}(\sqrt{d})$,
suppose for sufficiently large $C=\operatorname{poly}(\log(d)m_{D})$ there are
vectors $\\{q_{i,t}\\}_{i\in[m_{G}]}$, $\\{p_{i}\\}_{i\in[m_{D}]}$ in
$\mathbb{R}^{d}$ such that $\|p_{i}\|_{2}=1,\langle
p_{i},p_{i^{\prime}}\rangle\leq\tilde{O}(1/\sqrt{d})$ for $i,i^{\prime}$,
values $H_{i,j,t},G_{i,t}\in\left[\frac{1}{C^{2}},C^{2}\right]$ and a value
$a_{0}\geq 0$ satisfies that:

$\displaystyle a_{0}$
$\displaystyle=\frac{1}{\operatorname{polylog}(d)},\|q_{j,0}\|_{2}=\tilde{\Theta}(\Lambda);\quad
q_{j,0}=\sum_{i\in[m_{D}]}a_{i}p_{i}+\xi_{j},a_{i}\geq
0,\|\xi_{i}\|_{2}\leq\frac{1}{C}\|q_{j,0}\|_{2}$ (39) $\displaystyle a_{t+1}$
$\displaystyle=a_{t}-\eta_{D}\left(H_{i,j,t}\sum_{i\in[m_{D}],j\in[m_{G}]}\sigma(\langle
p_{i},q_{j,t}\rangle)\right)$ (40) $\displaystyle q_{i,t+1}$
$\displaystyle=q_{i,t}+\eta_{G}a_{t}G_{i,t}\sum_{j\in[m_{D}]}\left(\sigma^{\prime}(\langle
p_{j},q_{i,t}\rangle)\left(p_{j}\pm\frac{1}{C}\right)\right)$ (41)

Then we must have: within
$T=\tilde{O}\left(\frac{\sqrt{\Gamma}}{\eta_{G}}\right)$ many iterations, we
must have that $a_{t}\leq 0$ and
$\max_{j\in[m_{G}]}\|q_{j,T}\|_{2}=\tilde{\Theta}(\sqrt{\Gamma})$. Moreover,
for every $t\leq T$, we have: for every $j\in[m_{G}]$,

$\sum_{i\in[m_{D}]}\sigma(\langle
p_{i},q_{j,t}\rangle)=\Omega\left(\max_{i\in[m_{D}]}\sigma^{\prime}(\langle
p_{i},q_{j,t}\rangle)\|q_{j,t}\|_{2}\right)$

and

$\max_{i\in[m_{D}]}\langle
p_{i},q_{j,t}\rangle\geq\left(1-\frac{1}{C^{0.2}}\right)\|q_{j,t}\|_{2}$

###### Proof of Lemma E.5.

Let us denote $r_{i,t}=\max_{j\in[m_{D}]}\\{\langle p_{j},q_{i,t}\rangle\\}$.

By the update rule, we can easily conclude that:

$r_{i,t+1}=r_{i,t}+\eta_{G}G_{i,t}\left(1-\frac{1}{C^{0.5}}\right)\sigma^{\prime}(r_{i,t})$

On the other hand, let us write
$q_{i,t}=\sum_{j\in[m_{D}]}\alpha_{i,j,t}q_{j}+\xi_{i,t}$, where
$\alpha_{i,j,t}\geq 0$. We know that:

$\displaystyle\|\xi_{i,t+1}\|_{2}\leq\|\xi_{i,t}\|_{2}+\eta_{G}G_{i,t}\frac{m_{D}}{C}\sigma^{\prime}(r_{i,t})$
(42)

By the comparison Lemma E.3 we can easily conclude that for every $t$,

$\|\xi_{i,t}\|_{2}\leq\frac{1}{C^{0.5}}r_{i,t}$

This implies that: there exists values $u_{t}\in[1/C^{2},C^{2}]$ such that

$\displaystyle a_{t+1}=a_{t}-\eta_{D}u_{t}\sum_{i\in[m_{G}]}\sigma(r_{i,t})$
(43)

Comparing this with the update rule of $r_{i,t}$, we know that for every $t$
with $a_{t}\geq 0$, we must have:

$r_{i,t}=\tilde{O}\left(\sqrt{\Gamma}\right),\quad
r_{i,T}=\tilde{\Theta}(\sqrt{\Gamma})$

∎

######  Lemma E.6 (Auxiliary Lemma).

For every $g>0$ we must have: $\operatorname{Sigmoid}(-gx-b)x$ is a decreasing
function of $x$ as long as $gx>1$ and $gx+b>0$.

######  Lemma E.7.

For $a_{t},b_{t},c_{t},d_{t}\in\mathbb{R}^{d}$ be defined as:
$a_{0},c_{0},d_{0}=\frac{1}{\operatorname{polylog}(d)}$, $|b_{t}|\leq O(\log
d)$ and $|b_{t}|\leq\min\\{a_{t}c_{t}^{3},a_{t}d_{t}^{3}\\}$.

$\displaystyle a_{t+1}$
$\displaystyle=a_{t}+\eta_{D}\frac{1}{2}\left(\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(-a_{t}c_{t}^{3}-b_{t})c_{t}^{3}+\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(-a_{t}d_{t}^{3}-b_{t})d_{t}^{3}\right)$
(44) $\displaystyle c_{t+1}$
$\displaystyle=c_{t}+\eta_{D}\frac{3}{2}\left(\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(-a_{t}c_{t}^{3}-b_{t})c_{t}^{2}a_{t}\right)$
(45) $\displaystyle d_{t+1}$
$\displaystyle=d_{t}+\eta_{D}\frac{3}{2}\left(\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(-a_{t}d_{t}^{3}-b_{t})d_{t}^{2}a_{t}\right)$
(46)

Then we have: for every
$t\in\left(\frac{\operatorname{polylog}(d)}{\eta_{D}},\frac{\operatorname{poly}(d)}{\eta_{D}}\right]$,
we must have:

$\displaystyle a_{t}$
$\displaystyle=\sqrt{\frac{2}{3}}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)c_{t}$
(47) $\displaystyle c_{t}$
$\displaystyle=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)d_{t}$ (48)

###### Proof of Lemma E.7.

By the update formula, we can easily conclude that for
$t\leq\frac{\operatorname{poly}(d)}{\eta_{D}}$, we have that
$a_{t},c_{t},d_{t}\in\left[\frac{1}{\operatorname{polylog}(d)},\operatorname{polylog}(d)\right]$.
This implies that for every
$t\in\left[\frac{\operatorname{polylog}(d)}{\eta_{D}},\frac{\operatorname{poly}(d)}{\eta_{D}}\right]$,
we have that

$a_{t}c_{t}^{3},a_{t}d_{t}^{3}\in[1,O(\log d)]$

Apply Lemma E.6 we have that: As long as $a_{t}>3(c_{t}+d_{t})$, we must have
that

$\operatorname{Sigmoid}(-a_{t}c_{t}^{3}-b_{t})c_{t}^{3}+\operatorname{Sigmoid}(-a_{t}d_{t}^{3}-b_{t})d_{t}^{3}<\operatorname{Sigmoid}(-a_{t}c_{t}^{3}-b_{t})c_{t}^{2}a_{t}+\operatorname{Sigmoid}(-a_{t}d_{t}^{3}-b_{t})d_{t}^{2}a_{t}$

This implies that

$\frac{a_{t+1}}{3}-\frac{a_{t}}{3}<c_{t+1}+d_{t+1}-c_{t}-d_{t}$

Note that initially, $a_{0},c_{0},d_{0}=\frac{1}{\operatorname{polylog}(d)}$.
This implies that when $t\geq\frac{\operatorname{polylog}(d)}{\eta_{D}}$, we
must have that $a_{t}\leq 4(c_{t}+d_{t})$, therefore $c_{t}+d_{t}=\Omega(1)$.
Similarly, we can prove that $a_{t}\geq 0.1\min\\{c_{t},d_{t}\\}$.

as long as $c_{t}>d_{t}$, we must have:

$\operatorname{Sigmoid}(-a_{t}c_{t}^{3}-b_{t})c_{t}^{2}a_{t}<\operatorname{Sigmoid}(-a_{t}d_{t}^{3}-b_{t})d_{t}^{2}a_{t}$

Which implies that:

$\displaystyle\frac{c_{t+1}}{1+1/\operatorname{polylog}(d)}-\frac{c_{t}}{1+1/\operatorname{polylog}(d)}<d_{t+1}-d_{t}$
(49)

Note that initially, $c_{0},d_{0}=\frac{1}{\operatorname{polylog}(d)}$ and
when $t\geq\frac{\operatorname{polylog}(d)}{\eta_{D}}$,
$c_{t}+d_{t}=\Omega(1)$. This implies that for every
$t\in\left[\frac{\operatorname{polylog}(d)}{\eta_{D}},\frac{\operatorname{poly}(d)}{\eta_{D}}\right]$,
we have: $c_{t}=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)d_{t}$.
Which also implies that $c_{t},d_{t}\leq O(\log d)$.

Similarly, we can prove the bound for $a_{t}$.

∎

##  Appendix F Induction hypothesis

###  F.1 Case 1: Balanced generator and discriminator

In this section we consider the case 2 in Lemma D.1. Here we give the
induction hypothesis to prove the case of balanced generator and
discriminator, this is the most difficult case and other cases are just simple
modification of this one. Without loss of generality (by symmetry), let us
assume that $a^{(0)}>0$ and $a^{(0)}=\frac{1}{\operatorname{polylog}(d)}$
(this happens with probability $1-o(1)$).

We divide the training into five stages: For a sufficiently large
$C=\operatorname{polylog}(d)$

  1. 1.

Stage 1: Before one of the $\alpha(w_{i},u_{j},t)\geq 1/C$. Call this exact
iteration $T_{B,1}$.

  2. 2.

Stage 2: After $T_{B,1}$, before
$T_{B,2}=T_{B,1}+\frac{1}{\eta_{D}2^{\sqrt{\log(d)}}}$.

  3. 3.

Stage 3: After $T_{B,2}$, before one of the $\alpha(v_{i},u_{j},t)\geq
d^{-0.49}$. Call this exact iteration $T_{B,3}$.

  4. 4.

Stage 4: After $T_{B,3}$, before
$a^{(t)}\leq\tilde{O}\left(\frac{1}{\Lambda^{2}d^{1/4}}\right)$. Call this
exact iteration $T_{B,4}$.

  5. 5.

Stage 5: After $T_{B,4}$, until convergence.

We maintain the following things about $\alpha$ and $a,b$ during each stage:

#### Stage 1

: We maintain: For every $t\leq T_{B,1}$:

  1. 1.

(B.1.0). For all but the $i^{*}\in[m_{D}]$, and for all $j\in[m_{G}]$ (Below
$*$ can be $w_{i^{\prime}},v_{j^{\prime}},u_{\ell}$ for every
$i^{\prime}\in[m_{D}],j^{\prime}\in[m_{G}]$ and $\ell\in[2]$).

$\forall*\not=u_{1},u_{2}:|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{1}{d^{0.9}},\quad|\alpha(w_{i},u_{\ell},t)-\alpha(w_{i},u_{\ell},0)|\leq\frac{C}{\sqrt{d}}$

  2. 2.

(B.1.1). For all $j\in[m_{G}]$:

$|\alpha(v_{j},*,t)-\alpha(v_{j},*,0)|\leq\frac{\operatorname{polyloglog}(d)}{d}$
$|\alpha(v_{j},u_{\ell},t)|\leq\frac{1}{d}$

  3. 3.

(B.1.2). For $i^{*}=i$, we have that: for all $*\not=u_{1},u_{2}$:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{1}{d^{0.9}}$

  4. 4.

(B.1.3). $a$ and $b$ remains nice:

$a^{(t)}\in(1-1/C,1+1/C)a_{0},|b^{(t)}|\leq\frac{1}{C}$

#### Stage 2

: For every $t\in[T_{B,1},T_{B,2}]$.

  1. 1.

(B.1.0), (B.1.1), (B.1.2) still holds.

  2. 2.

(B.2.1): For $i=i^{*}$, we have:

$a^{(t)},\alpha(w_{i},u_{\ell},t)=\tilde{\Theta}(1)$

#### Stage 3

: For every $t\in[T_{B,2},T_{B.3}]$.

  1. 1.

(B.1.0), (B.1.2) still holds.

  2. 2.

(B.3.2): For every $j\in[m_{G}]$: For $*\not=w_{i^{*}},u_{1},u_{2}$, we have:

$|\alpha(v_{j},*,t)-\alpha(v_{j},*,0)|\leq\frac{C^{3}}{\sqrt{d}}\|v_{j}^{(t)}\|_{2}$

and

$\alpha(v_{i},u_{\ell},t)\geq-O\left(\frac{1}{d}\right)$

Moreover, let $\alpha(t):=\max_{j\in[m_{G}],\ell\in[2]}\langle
v_{j}^{(t)},u_{\ell}\rangle$, we have that:

$|\alpha(v_{j},w_{i^{*}},t)|\leq\operatorname{polyloglog}(d)\alpha(t),\quad|\langle
v_{j}^{(t)},u_{\ell}\rangle|\leq O(\alpha(t))$

  3. 3.

(B.3.3): Balanced update: for every $X$,

$\operatorname{Sigmoid}\left(-a^{(t)}\langle
w_{i^{*}}^{(t)},X\rangle^{3}-b^{(t)}\right)\in\left[\frac{1}{\sqrt{d}\operatorname{polylog}(d)},\frac{1}{\operatorname{polylog}(d)}\right]\operatorname{Sigmoid}\left(b^{(t)}\right)$

and

$\operatorname{Sigmoid}\left(-a^{(t)}\langle
w_{i^{*}}^{(t)},u_{1}\rangle^{3}-b^{(t)}\right)=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(-a^{(t)}\langle
w_{i^{*}}^{(t)},u_{2}\rangle^{3}-b^{(t)}\right)$

#### Stage 4

: For every $t\in[T_{B,3},T_{B.4}]$.

  1. 1.

(B.3.1), (B.3.2) still holds.

  2. 2.

(B.4.1) for $i=i^{*}$, we have that for all $*\not=u_{1},u_{2},w_{i}$:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{C}{\sqrt{d}}$

For all $*\in[u_{1},u_{2},w_{i}]$:

$\alpha(w_{i},*,t)=\Theta(\alpha(w_{i},*,T_{B,3}))$

  3. 3.

For every $i,j\in[m_{G}]$:
$\|v_{i}^{(t)}\|_{2}=\tilde{\Theta}(\|v_{j}^{(t)}\|_{2})$ and after
$t=T_{B,4}$, we have that for every $i\in[m_{G}]$,
$\|v_{i}^{(t)}\|_{2}=\tilde{\Theta}(d^{1/4})$.

  4. 4.

$|a^{(t)}|,|b^{(t)}|=O(\log(d))$.

#### Stage 5

: For every $t\in[T_{B,4},T_{0}]$.

  1. 1.

For every $i\in[m_{D}]$,

$|\alpha(w_{i},*,T_{B,4})-\alpha(w_{i},*,t)|\leq d^{-0.1}$

  2. 2.

For every $i\in[m_{G}]$,

$|\alpha(v_{i},*,T_{B,4})-\alpha(v_{i},*,t)|\leq d^{0.2}$

  3. 3.

$|a^{(t)}|\leq\tilde{O}\left(\frac{1}{\Lambda^{2}d^{1/4}}\right)$, and for
every $z$:

$\langle w_{i^{*}}^{(t)},G^{(t)}(z)\rangle=\tilde{\Omega}(d^{1/4})$

###  F.2 Case 2: Generator is dominating

We now consider another case where the generator’s learning rate dominates
that of the discriminator. This corresponds to case 3 in Lemma D.1. In this
case, we divide the learning into four stages: For a sufficiently large
$C=2^{\sqrt{\log d}}$:

  1. 1.

Before $\alpha(v_{j_{G}},w_{i_{G}},t)\geq d^{-0.49}$. Call this iteration
$T_{G,1}$.

  2. 2.

After $T_{G,1}$, before $\alpha(v_{j_{G}},w_{i_{G}},t)\geq\Lambda$. Call this
iteration $T_{G,2}$.

  3. 3.

After iteration $T_{G,2}$, before $a_{t}\leq 0$. Call this iteration
$T_{G,3}$.

  4. 4.

After $T_{G,3}$.

We maintain the following induction hypothesis:

#### Stage 1

: In this stage, we maintain the following induction hypothesis: Let
$\alpha(t):=\alpha(v_{j_{G}},w_{i_{G}},t)$, for every $t\leq T_{G,1}$:

  1. 1.

(G.1.1). For all $i\in[m_{D}]$, and for all $j\in[m_{G}]$:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{C}{\sqrt{d}}$

  2. 2.

(G.1.2). For all $j\in[m_{G}]$, for all $*\not=w_{i_{G}}$:

$|\alpha(v_{j},*,t)-\alpha(v_{j},*,0)|\leq\frac{C}{\sqrt{d}}\alpha(t)$

#### Stage 2

: In this stage, we maintain: for every $t\in[T_{G,1},T_{G,2}]$

  1. 1.

(G.2.1). For every $i\in[m_{D}]$, we have:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{1}{C}$

  2. 2.

(G.2.2). For every $j\in[m_{G}]$, $\alpha(v_{i},w_{i_{G}},t)\geq d^{-0.49}$.

  3. 3.

For every $i\in[m_{G}]$, we have: for every $*\not=w_{i_{G}}$:

$|\alpha(v_{i},*,t)-\alpha(v_{i},*,0)|\leq\frac{2}{C}|\alpha(v_{i},w_{i_{G}},t)|$

#### Stage 3

: In this stage, we maintain: For every $t\in[T_{G,2},T_{G,3}]$:

  1. 1.

(G.2.1), (G.2.2) still holds.

  2. 2.

For every $i\in[m_{G}]$, we have: for every $*=v_{r}$ or $*=u_{\ell}$:

$|\alpha(v_{i},*,t)-\alpha(v_{i},*,0)|\leq\frac{2}{C}\|v_{i}^{(t)}\|_{2}$

#### Stage 4

: In this stage, we maintain: For every $t\in[T_{G,3},T_{1}]$:

  1. 1.

(G.2.1) still holds.

  2. 2.

For every $i\in[m_{G}]$, we have:

$|\alpha(v_{i},*,t)-\alpha(v_{i},*,T_{G,3})|\leq\frac{1}{C}\|v_{i}^{(T_{G,3})}\|_{2}$

  3. 3.

$|\alpha_{t}|=\tilde{O}\left(\frac{1}{\Lambda^{2}\sqrt{\eta_{G}/\eta_{D}}}\right)$,
$\|v_{i}^{(t)}\|_{2}=\tilde{\Theta}(\sqrt{\eta_{G}/\eta_{D}})$, and for all
$z\not=0$,
$\sum_{i\in[m_{D}]}h(G^{(t)}(z))=\tilde{\Theta}(\Lambda^{2}\sqrt{\eta_{G}/\eta_{D}})$.

##  Appendix G Proof of the learning process in balanced case

For simplicity, we are only going to prove the case when $u_{1}\bot u_{2}$ and
$\Pr[s_{1}=s_{2}=1]=\gamma$. The other case can be proved identically.

###  G.1 Stage 1

In this stage, by the induction hypothesis we know that
$\|v_{i}^{(t)}\|_{2}\leq\tilde{O}(1/\sqrt{d})$. Therefore, the update of
$w_{i}^{(t)}$ can be approximate as:

######  Lemma G.1.

For every $t\leq T_{B,1}$, we know that: when the random samples are $(X,z)$:

$\displaystyle
w_{i}^{(t+1)}=w_{i}^{(t)}+\eta_{D}a^{(0)}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{3}{2}\langle
w_{i}^{(t)},X\rangle^{2}X\pm\eta_{D}\tilde{O}\left(\frac{1}{d^{1.5}}\right)$
(50)

Moreover, we have that if $z_{i}=1$:

$\displaystyle v_{i}^{(t+1)}$
$\displaystyle=v_{i}^{(t)}+\eta_{G}a^{(0)}\frac{3}{2}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\sum_{j\in[m_{D}]}\left(\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\frac{1}{C^{0.5}d}\right)^{2}w_{j}^{(t)}$ (51)
$\displaystyle=v_{i}^{(t)}+\eta_{G}a^{(0)}\frac{3}{2}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\sum_{j\in[m_{D}]}\left(\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\frac{1}{C^{0.5}d}\right)^{2}w_{j}^{(0)}\pm\eta_{G}{O}\left(\frac{1}{C^{0.5}d^{2}}\right)$
(52)

Taking expectation of the above Lemma, we can easily conclude that:

$\displaystyle\mathbb{E}[w_{i}^{(t+1)}]$
$\displaystyle=\mathbb{E}[w_{i}^{(t)}]+\eta_{D}a^{(0)}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{3}{4}\left(\langle
w_{i}^{(t)},u_{1}\rangle^{2}u_{1}+\langle
w_{i}^{(t)},u_{2}\rangle^{2}u_{2}+\Theta(\gamma)\langle
w_{i}^{(t)},u_{1}+u_{2}\rangle^{2}(u_{1}+u_{2})\right)$ (53)
$\displaystyle\pm\eta_{D}\tilde{O}\left(\frac{1}{d^{1.5}}\right)$ (54)

and

$\displaystyle\mathbb{E}[\langle w_{j}^{(0)},G^{(t)}(z)\rangle^{2}\mid
z_{i}=1]$ $\displaystyle=\langle w_{j}^{(0)},v^{(t)}_{i}\rangle^{2}\pm
O\left(\frac{1}{m_{G}\operatorname{polylog}(d)}\sum_{i\in[m_{G}]}|\langle
w_{j}^{(0)},v^{(t)}_{i}\rangle|\right)^{2}$ (55)

Therefore, let $\zeta_{t}=\max_{i\in[m_{G}],j\in[m_{D}]}\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle$,
$\Upsilon_{t}=\max_{j\in[m_{D}],\ell\in[2]}\langle
w_{j}^{(t)},u_{\ell}\rangle$, we have that:

$\displaystyle\mathbb{E}[\Upsilon_{t+1}]$
$\displaystyle=\Upsilon+\eta_{D}a^{(0)}\frac{3}{4}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\Upsilon_{t}^{2}$
(56) $\displaystyle\mathbb{E}[\zeta_{t+1}]$
$\displaystyle=\zeta_{t}+\eta_{G}a^{(0)}\frac{3}{2m_{G}}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\zeta_{t}^{2}$
(57)

###### Proof of Lemma G.1.

By the gradient formula, we have:

$\nabla_{w_{i}}L(X,z)=-\operatorname{Sigmoid}(-f(X))a\sigma^{\prime}(\langle
w_{i},X\rangle)X+\operatorname{Sigmoid}(f(G(z)))a\sigma^{\prime}(\langle
w_{i},G(z)\rangle)G(z)$
$\nabla_{v_{i}}L(X,z)=-1_{z_{i}=1}\operatorname{Sigmoid}(f(G(z)))a\sum_{j\in[m_{D}]}\sigma^{\prime}(\langle
w_{i},G(z)\rangle)w_{i}$

At iteration $t$, by induction hypothesis, we have that $a^{(t)}=a^{(0)}(1\pm
1/C)$.

Moreover, by the induction hypothesis again, we have hat
$|f(X)|=\tilde{O}(d^{-1.5})$ and $\|G(z)\|_{2}\leq{O}(d^{-0.5})$. Together
with $\|w_{i}^{(t)}\|_{2}=\tilde{O}(1)$, this implies that

$\|\operatorname{Sigmoid}(f(G(z)))a\sigma^{\prime}(\langle
w_{i},G(z)\rangle)G(z)\|_{2}=\tilde{O}(d^{-1.5})$

This proves the update formula for $w_{i}^{(t)}$. As for $v_{i}$, we observe
that by the induction hypothesis and notice that w.h.p. over the randomness of
initialization, $|\langle v_{i}^{(0)},u_{\ell}\rangle|\leq\frac{\log d}{d}$,
therefore, we can conclude that

$\displaystyle\langle w_{j}^{(t)},G^{(t)}(z)\rangle=\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\tilde{O}\left(\frac{1}{d^{1.35}}\right)\pm
O\left(\frac{\log d}{Cd}\right)=\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\frac{1}{C^{0.5}d}$ (58)

Note that by induction hypothesis,
$\|w_{j}^{(t)}-w_{j}^{(0)}\|_{2}\leq\frac{1}{C}$ and $\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\leq\frac{m_{G}(C^{0.1}+\log d)}{d}$. This
implies that

$\displaystyle\langle w_{j}^{(t)},G^{(t)}(z)\rangle^{2}w_{j}^{(t)}$
$\displaystyle=\left(\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\frac{1}{C^{0.5}d}\right)^{2}w_{j}^{(t)}$ (59)
$\displaystyle=\left(\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\frac{1}{C^{0.5}d}\right)^{2}w_{j}^{(0)}+{O}\left(\frac{1}{C^{0.5}d^{2}}\right)$
(60)

∎

Now, apply Lemma E.4 and the fact that w.p. $1-o(1)$,
$\zeta_{0}=\frac{\operatorname{polyloglog}(d)}{d}$, we have that:

######  Lemma G.2.

$\displaystyle\sum_{t\leq T_{1}}\eta_{G}\zeta_{t}^{2}\leq
O\left(\frac{m_{G}\operatorname{polyloglog}(d)}{a^{(0)}d}\right)$ (61)

In the end, we can show the following Lemma:

######  Lemma G.3.

When $t=T_{B,1}$, we have that: for both $\ell\in[2]$,

$\alpha(w_{i^{*}},u_{\ell},t)=\frac{1}{\operatorname{polylog}(d)}$

###### Proof of Lemma G.3.

By the update formula in Eq equation 53, and the fact that
$\Pr[X=u_{1}+u_{2}]\geq\frac{1}{\operatorname{polylog}(d)}$ and the induction
hypothesis, we know that for $i=i^{*}$, for $t\leq T_{B,1}$ we have that:

$\alpha(w_{i},u_{\ell},t+1)\geq\alpha(w_{i},u_{\ell},t)+\tilde{\Omega}(\eta_{D})\times\left(\alpha(w_{i},u_{3-\ell},t)-\frac{1}{d}\right)^{2}$

This implies that at the end of Stage 1, when
$\alpha(w_{i},u_{3-\ell},t)\geq\frac{1}{C}$, we must have
$\alpha(w_{i},u_{\ell},t)\geq\tilde{\Omega}(1)$ as well. ∎

###  G.2 Stage 2 and Stage 3

At this stage, by the induction hypothesis, we can approximate the function
value as:

$\displaystyle\sum_{i\in[m_{D}]}\sigma\left(\langle
w_{i}^{(t)},X\rangle\right)=\langle
w_{i^{*}}^{(t)},X\rangle^{3}\pm\tilde{O}\left(\frac{1}{d^{1.5}}\right)$ (62)
$\displaystyle\left|\sum_{i\in[m_{D}]}\sigma\left(\langle
w_{i}^{(t)},G^{(t)}(z)\rangle\right)\right|\leq\tilde{O}\left(\|G^{(t)}(z)\|_{2}\right)^{3}\leq\tilde{O}\left(\frac{1}{d^{1.45}}\right)$
(63)

Therefore, at this stage, we can easily approximate the update of
$W_{D}^{(t)}$ as:

######  Lemma G.4.

When the sample is $(X,z)$, we have: for every $t\in(T_{B,1},T_{B,3}]$, the
following holds:

$\displaystyle a^{(t+1)}$
$\displaystyle=a^{(t)}+\eta_{D}\left(1\pm\tilde{O}\left(\frac{1}{d}\right)\right)\operatorname{Sigmoid}\left(-a^{(t)}\langle
w_{i^{*}}^{(t)},X\rangle^{3}-b^{(t)}\right)\langle
w_{i^{*}}^{(t)},X\rangle^{3}$ (64)
$\displaystyle\pm\eta_{D}\tilde{O}\left(\frac{1}{d^{1.45}}\right)\operatorname{Sigmoid}\left(b^{(t)}\right)$
(65) $\displaystyle w_{i}^{(t+1)}$
$\displaystyle=w_{i}^{(t)}+3\eta_{D}\left(1\pm\tilde{O}\left(\frac{1}{d}\right)\right)\operatorname{Sigmoid}\left(-a^{(t)}\langle
w_{i}^{(t)},X\rangle^{3}-b^{(t)}\right)a^{(t)}\langle
w_{i}^{(t)},X\rangle^{2}X$ (66)
$\displaystyle\pm\eta_{D}\tilde{O}\left(\frac{1}{d^{1.45}}\right)\operatorname{Sigmoid}\left(b^{(t)}\right)$
(67) $\displaystyle b^{(t)}$
$\displaystyle=b^{(t)}+\eta_{D}\tau_{b}\left(1\pm\tilde{O}\left(\frac{1}{d}\right)\right)\operatorname{Sigmoid}\left(-a^{(t)}\langle
w_{i^{*}}^{(t)},X\rangle^{3}-b^{(t)}\right)$ (68)
$\displaystyle-\eta_{D}\tau_{b}\left(1\pm\tilde{O}\left(\frac{1}{d^{1.45}}\right)\right)\operatorname{Sigmoid}\left(b^{(t)}\right)$
(69)

Moreover, the update formula also let us bound
$a^{(t)},\alpha(w_{i^{*}},u_{1},t)$ as:

######  Lemma G.5.

Let $\alpha_{t},a_{t}$ be updated as: for $t=T_{B,2}$,
$\alpha_{t}=\alpha(w_{i^{*}},u_{1},t)$ and $a_{t}=a^{(t)}$, such that

$a_{t+1}=a_{t}+\eta_{G}\operatorname{Sigmoid}(-a_{t}\alpha_{t}^{3}-b_{t})\alpha_{t}^{3}$
$\alpha_{t+1}=\alpha_{t}+\frac{3}{2}\eta_{G}a_{t}\operatorname{Sigmoid}(-a_{t}\alpha_{t}^{3}-b_{t})\alpha_{t}^{2}$

Where $b_{t}$ be updated as: for $t=T_{B,2}$, $b_{t}=b^{(t)}$ and update as:

$b_{t+1}=b_{t}-\eta_{D}\tau_{b}\operatorname{Sigmoid}(b_{t})$

Then we have: for every $t\in[T_{B,2},T_{B,3}]$

$a_{t}=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)a^{(t)},\quad\alpha_{t}=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\alpha(w_{i^{*}},u_{1},t)$

$\operatorname{Sigmoid}(b^{(t)})=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(b_{t})$

Moreover, when $t=T_{B,3}$, we have: $a_{t}\leq 0.819\alpha_{t}$.

###### Proof of Lemma G.5.

By the update formula in Lemma G.4 and the bound in induction hypothesis
(B.3.3), we can simplify the update of $a^{(t)},b^{(t)}$ and
$\alpha(w_{i},u_{\ell},t)$ as: for $i=i^{*}$, when $X=u_{\ell}$:

$\displaystyle a^{(t+1)}$
$\displaystyle=a^{(t)}+\eta_{D}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{\ell},t)^{3}-b^{(t)}\right)\alpha(w_{i},u_{\ell},t)^{3}$
(70) $\displaystyle\alpha(w_{i},u_{\ell},t+1)$
$\displaystyle=\alpha(w_{i},u_{\ell},t)$ (71)
$\displaystyle+3\eta_{D}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{\ell},t)^{3}-b^{(t)}\right)a^{(t)}\alpha(w_{i},u_{\ell},t)^{2}$
(72) $\displaystyle b^{(t)}$
$\displaystyle=b^{(t)}-\eta_{D}\tau_{b}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(b^{(t)}\right)$
(73)

B͡y the last inequality, we know that when
$\operatorname{Sigmoid}(b^{(t)})>\left(1+\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(b_{t})$,
then $b^{(t)}$ must be decreasing faster than $b_{t}$, otherwise if
$\operatorname{Sigmoid}(b^{(t)})>\left(1-\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(b_{t})$,
then $b_{t}$ must be decreasing faster than $b^{(t)}$, which proves the bound
of $b^{(t)}$. Moreover, the update formula of $b_{t}$ also gives us that for
every $t\leq\operatorname{poly}(d)$, we have: $|b_{t}|=O(\log d)$. This
implies that for every $Z$:

$\operatorname{Sigmoid}(Z+b^{(t)})=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(Z+b_{t})$

To obtain the bound of $a^{(t)}$ and $\alpha(w_{i}^{*},u_{1},t)$, notice that
when $X=u_{1}+u_{2}$, we have that:

$\operatorname{Sigmoid}\left(-a^{(t)}(\alpha(w_{i},u_{1},t)+\alpha(w_{i},u_{2},t))^{3}-b^{(t)}\right)\leq\min_{\ell\in[2]}\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{\ell},t)^{3}-b^{(t)}\right)$

Therefore, we can conclude:

$\displaystyle\mathbb{E}[a^{(t+1)}]$ $\displaystyle=a^{(t)}$ (74)
$\displaystyle+\eta_{D}\frac{1}{2}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\left(\sum_{\ell\in[2]}\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{\ell},t)^{3}-b_{t}\right)\alpha(w_{i},u_{\ell},t)^{3}\right)$
(75) $\displaystyle\mathbb{E}[\alpha(w_{i},u_{\ell},t+1)]$
$\displaystyle=\alpha(w_{i},u_{\ell},t)$ (76)
$\displaystyle+\frac{3}{2}\eta_{D}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{\ell},t)^{3}-b_{t}\right)a^{(t)}\alpha(w_{i},u_{\ell},t)^{2}$
(77)

Using Lemma G.3 we can conclude that

$\alpha(w_{i^{*}},u_{\ell},T_{B,1})=\frac{1}{\operatorname{polylog}(d)}$

and now apply Lemma E.7, we have: $a^{(t)}=\Theta(\alpha(w_{i^{*}},u_{1},t))$
and

$\alpha(w_{i^{*}},u_{1},t)=\left[\alpha(w_{i^{*}},u_{1},t)\right]\left(1\pm\frac{1}{\operatorname{polyloglog}(d)}\right)$

This implies that:

$\displaystyle\mathbb{E}[a^{(t+1)}]$
$\displaystyle=a^{(t)}+\eta_{D}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\left(\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{1},t)^{3}-b_{t}\right)\alpha(w_{i},u_{\ell},t)^{3}\right)$
(78) $\displaystyle\mathbb{E}[\alpha(w_{i},u_{1},t+1)]$
$\displaystyle=\alpha(w_{i},u_{1},t)$ (79)
$\displaystyle+\frac{3}{2}\eta_{D}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{1},t)^{3}-b_{t}\right)a^{(t)}\alpha(w_{i},u_{1},t)^{2}$
(80)

Apply Lemma E.7 again, we know that when

$a^{(t)}\alpha(w_{i},u_{1},t)^{3}>\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)a_{t}\alpha_{t}^{3}$

We must have that $a^{(t)}\geq a_{t}$ and $\alpha(w_{i},u_{1},t)>\alpha_{t}$.
Therefore, apply Lemma E.6 we know that in this case:

$\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{1},t)^{3}-b_{t}\right)\alpha(w_{i},u_{\ell},t)^{3}\leq\operatorname{Sigmoid}(-a_{t}\alpha_{t}^{3}-b_{t})\alpha_{t}^{3}$

and

$\operatorname{Sigmoid}\left(-a^{(t)}\alpha(w_{i},u_{1},t)^{3}-b_{t}\right)a^{(t)}\alpha(w_{i},u_{1},t)^{2}\leq\operatorname{Sigmoid}(-a_{t}\alpha_{t}^{3}-b_{t})a_{t}\alpha_{t}^{2}$

Combine this with the update rule we can directly complete the proof.

∎

The Lemma G.5 immediately implies that the $\alpha(w_{i^{*}},u_{\ell},t)$ will
be balanced after a while:

######  Lemma G.6.

We have that for every $t\in[T_{B,2},T_{B,3}]$, the following holds:

$\alpha(w_{i^{*}},u_{1},t)=\left[\alpha(w_{i^{*}},u_{1},t)\right]\left(1\pm\frac{1}{\operatorname{polyloglog}(d)}\right)$

and

$\alpha(w_{i^{*}},u_{1},t)\geq\log^{0.1}(d)$

Using Lemma G.6, we also have the Lemma that approximate the update of
$v_{i}^{(t)}$ as:

######  Lemma G.7.

Let us define $\alpha(t):=\max_{j\in[m_{G}],\ell\in[2]}\langle
v_{j}^{(t)},u_{\ell}\rangle$. For every $t\in[T_{B,2},T_{B,3}]$, we have: for
$j\not=i^{*}$:

$\displaystyle\langle w_{j}^{(t)},G^{(t)}(z)\rangle^{2}w_{j}^{(t)}=\langle
w_{j}^{(t)},G^{(t)}(z)\rangle^{2}w_{j}^{(0)}\pm\frac{C^{2}}{d^{0.5}}\langle
w_{j}^{(t)},G^{(t)}(z)\rangle^{2}$ (81)

For $j=i^{*}$:

$\displaystyle\langle w_{j}^{(t)},G^{(t)}(z)\rangle^{2}w_{j}^{(t)}=\langle
w_{j}^{(t)},G^{(t)}(z)\rangle^{2}\left(w_{j}^{(0)}+\alpha(w_{j},u_{1},t)u_{1}+\alpha(w_{j},u_{2},t)u_{2}\right)\pm\frac{C}{d^{0.9}}\langle
w_{j}^{(t)},G^{(t)}(z)\rangle^{2}$ (82)

Now, for $\langle w_{j}^{(t)},G^{(t)}(z)\rangle$ we have: For $j\not=i^{*}$:

$\displaystyle\mathbb{E}_{z}[\langle w_{j}^{(t)},G^{(t)}(z)\rangle^{2}\mid
z_{i}=1]$
$\displaystyle=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\left(\langle
w_{j}^{(0)},v^{(t)}_{i}\rangle\pm\frac{C^{2}}{\sqrt{d}}\alpha(t)\right)^{2}\pm\frac{\alpha(t)^{2}}{\operatorname{polylog}(d)}$
(83)

For $j=i^{*}$:

$\displaystyle\mathbb{E}_{z}[\langle w_{j}^{(t)},G^{(t)}(z)\rangle^{2}\mid
z_{i}=1]$
$\displaystyle=\left(1\pm\frac{1}{\operatorname{polyloglog}(d)}\right)\alpha(w_{i^{*}},u_{1},t)^{2}\left\langle(u_{1}+u_{2}),v^{(t)}_{i}\right\rangle^{2}$
(84)
$\displaystyle\pm\frac{\alpha(t)^{2}}{\operatorname{polyloglog}(d)}\alpha(w_{i^{*}},u_{1},t)^{2}$
(85)

Moreover, the update of Sigmoid can be approximate as:

$\operatorname{Sigmoid}\left(f^{(t)}\left(G^{(t)}(z)\right)\right)=\left(1\pm\tilde{O}\left(\frac{1}{d^{1.45}}\right)\right)\operatorname{Sigmoid}\left(b^{(t)}\right)$

###### Proof of Lemma G.7.

The first half of the lemma regarding $w_{j}^{(t)}$ follows trivially from the
induction hypothesis, we only need to look at $\langle
w_{j}^{(t)},G^{(t)}(z)\rangle$.

We know that for $j\not=i^{*}$, we have that by the induction hypothesis,

$\displaystyle\langle w_{j}^{(t)},G^{(t)}(z)\rangle=\langle
w_{j}^{(0)},G^{(t)}(z)\rangle\pm\tilde{O}\left(\frac{1}{d^{1.35}}\right)\pm
O\left(\frac{Cm_{G}\alpha(t)}{\sqrt{d}}\right)$ (86)

For $j=i^{*}$, we have that

$\displaystyle\langle w_{j}^{(t)},G^{(t)}(z)\rangle$ $\displaystyle=\langle
w_{j}^{(0)},G^{(t)}(z)\rangle+\alpha(w_{j},u_{1},t)\langle
u_{1},G^{(t)}(z)\rangle+\alpha(w_{j},u_{2},t)\langle
u_{2},G^{(t)}(z)\rangle\pm\tilde{O}\left(\frac{1}{d^{1.35}}\right)$ (87)
$\displaystyle=\alpha(w_{j},u_{1},t)\langle
u_{1}+u_{2},G^{(t)}(z)\rangle\pm\frac{1}{\operatorname{polyloglog}(d)}\alpha(w_{j},u_{1},t)\alpha(t)\|z\|_{1}\pm\tilde{O}\left(\frac{1}{d^{1.35}}\right)$
(88) $\displaystyle=\alpha(w_{j},u_{1},t)\langle
u_{1}+u_{2},v_{i}^{(t)}\rangle\pm\tilde{O}\left(\frac{1}{d^{1.35}}\right)\pm
O(\alpha(w_{j},u_{1},t)\alpha(t))(\|z\|_{1}-1)$ (89)

Taking expectation we can complete the proof.

∎

With Eq equation 81 and Eq equation 82 in lemma G.7, together with the
induction hypothesis, we immediately obtain

######  Lemma G.8.

For every $t\in[T_{B,2},T_{B,3}]$, we have that: for every $i\in[m_{G}]$:

$\displaystyle v_{i}^{(t)}$
$\displaystyle=v_{i}^{(T_{B,2})}+\sum_{\ell\in[2]}\alpha_{i,\ell}^{(t)}u_{\ell}+\sum_{j\in[m_{D}]}\beta_{i,j}^{(t)}w_{j}^{(0)}\pm\xi_{i,t}$
(90)

Where $\alpha_{i,\ell}^{(t)},\beta_{i,j}^{(t)}>0$ and
$\alpha_{i,\ell}^{(t)}=(1\pm o(1))\alpha_{i,3-\ell}^{(t)}$;
$\|\xi_{i,t}\|_{2}^{2}\leq\tilde{O}(1/d)\left(\sum_{\ell,j}(\alpha_{i,\ell}^{(t)})^{2}+(\beta_{i,j}^{(t)})^{2}\right)$

We now can immediately control the update of $v_{i}^{(t)}$ using the following
sequence:

######  Lemma G.9.

Let $v_{t}$ be defined as: for $t=T_{B,2}$

$v_{t}=\max_{i\in[m_{G}]}\langle
v_{i}^{(t)},u_{1}+u_{2}\rangle\left(1+\frac{1}{\operatorname{polyloglog}(d)}\right)$

and the update of $v_{t}$ is given as: for $\alpha_{t}$ defined as in Lemma
G.5

$v_{t+1}=v_{t}+\frac{3}{m_{G}}\operatorname{Sigmoid}(b_{t})\alpha_{t}^{2}v_{t}^{2}$

Then we must have: for every $t\in[T_{B,2},T_{B,3}]$:

$\displaystyle\max_{i\in[m_{G}]}\langle v_{i}^{(t)},u_{1}+u_{2}\rangle\leq
v_{t}$ (91)

On the other hand, if for $t=T_{B,2}$,

$v_{t}=\max_{i\in[m_{G}]}\langle
v_{i}^{(t)},u_{1}+u_{2}\rangle\left(1-\frac{1}{\operatorname{polyloglog}(d)}\right)$

Then we must have: for every $t\in[T_{B,2},T_{B,3}]$:

$\displaystyle\max_{i\in[m_{G}]}\langle v_{i}^{(t)},u_{1}+u_{2}\rangle\geq
v_{t}$ (92)

###### Proof of Lemma G.9.

In the setting of Lemma G.7, let us define $beta(t):=\max_{j\in[m_{G}]}\langle
v_{j}^{(t)},u_{1}+u_{2}\rangle$. We have that:

$\displaystyle\beta(t+1)=\beta(t)+\eta_{G}\alpha_{t}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{3}{m_{G}}\beta(t)^{2}\pm\frac{1}{\operatorname{polyloglog}(d)}\alpha(t)^{2}$
(93)

By the induction hypothesis we know that for all $j\in[m_{G}],\ell\in[2]$:

$\displaystyle\langle v_{j}^{(t)},u_{\ell}\rangle\geq\langle
v_{j}^{(0)},u_{\ell}\rangle-O\left(\frac{1}{d}\right)\geq-\frac{\log\log^{2}d}{d}$
(94)

This implies that $\beta(t)\geq\alpha(t)-\frac{\log\log^{2}(d)}{\sqrt{d}}$ and
$\beta(T_{B,2})\geq\frac{1}{d}$,
$\alpha(T_{B,2})\leq\frac{\operatorname{polyloglog}(d)}{\sqrt{d}}$. This
implies that:

$\displaystyle\beta(t+1)=\beta(t)+\eta_{G}\alpha_{t}\left(1\pm\frac{1}{\operatorname{polyloglog}(d)}\right)\frac{3}{m_{G}}\beta(t)^{2}$
(95)

This completes the proof by applying Lemma E.1. ∎

Now, by the comparison Lemma E.1, we know that one of the following event
would happen (depending on the initial value of $v_{t}$ at iteration
$T_{B,2}$):

######  Lemma G.10.

With probability $1-o(1)$, one of the following would happen:

  1. 1.

$T_{B,3}\geq T_{0}$.

  2. 2.

$T_{B,3}<T_{0}$, moreover, at iteration $T_{B,3}$, we have that
$\operatorname{Sigmoid}(b_{t})\geq\frac{1}{\operatorname{polylog}(d)}$.

In the end, we can easily derive an upper bound on the sum of
$\operatorname{Sigmoid}$ as below, which will be used to prove the induction
hypothesis.

######  Lemma G.11.

For every $t\in(T_{B,1},T_{B,3}]$, we have that: for every $X,z$:

$\displaystyle\sum_{t\in(T_{B,1},T_{B,3}]}\eta_{D}\operatorname{Sigmoid}\left(a^{(t)}\langle
w_{i^{*}}^{(t)},X\rangle^{3}+b^{(t)}\right)=\tilde{O}(1)$ (96)
$\displaystyle\sum_{t\in(T_{B,1},T_{B,3}]}\eta_{D}\tau_{b}\left(b^{(t)}\right)\leq\tilde{O}(1)$
(97)

We will also show the following Lemma regarding all the $v_{i}^{(t)}$ at
iteration $T_{B,3}$:

######  Lemma G.12.

For all $i\in[m_{G}]$, if we are in case 2 in Lemma G.10, we have that:

$\displaystyle\langle v_{i}^{(t)},u_{1}\rangle,\langle
v_{i}^{(t)},u_{2}\rangle=\tilde{\Omega}(d^{-0.49})$ (98)

###### Proof of Lemma G.12.

Since with probability at least $1/\operatorname{poly}(d)$, $z_{i}=z_{j}=1$,
so we have: By the update Lemma G.7 of $v$, we know that for all
$j\in[m_{G}]$: Let $\alpha(t)$ be defined as in Lemma G.7:

$\displaystyle\mathbb{E}[\alpha(v_{j},u_{\ell},t+1)]$
$\displaystyle\geq\alpha(v_{j},u_{\ell},t)+\eta_{G}\frac{1}{\operatorname{polylog}(d)}\left(\alpha(t)-\tilde{O}\left(\frac{1}{d}\right)\right)^{2}\pm\tilde{O}\left(\frac{1}{d^{0.8}}\right)\alpha(t)^{2}$
(99) $\displaystyle\mathbb{E}[\alpha(v_{j},u_{\ell},t+1)]$
$\displaystyle\leq\alpha(v_{j},u_{\ell},t)+\eta_{G}\operatorname{polylog}(d)\alpha(t)^{2}$
(100)

By Lemma G.10 we know that $\alpha(t)=\tilde{\Theta}(d^{-0.49})$ at iteration
$t=T_{B,3}$, which implies what we want to prove. ∎

###  G.3 Stage 4 and 5

In Stage 4 we can easily calculate that by induction hypothesis, for every
$i\in[m_{G}]$ and for every $j\in[m_{D}]$, $j\not=i^{*}$:

$|\langle
v_{i}^{(t)},w_{j}^{(t)}\rangle|\leq\tilde{O}\left(\frac{\|v_{i}^{(t)}\|_{2}}{\sqrt{d}}\right)$

Let

$S_{i,t}=\mathbb{E}_{z}\left[\operatorname{Sigmoid}\left(a^{(t)}\sigma\left(\langle
w_{i^{*}}^{(t)},G^{(t)}(z)\rangle\right)+b^{(t)}\right)\mid z_{i}=1\right]$

Note that by induction hypothesis, $|a^{(t)}|,b^{(t)}=O(\log(d))$. Which
implies that as long as
$\max_{i\in[m_{G}]}\|v_{i}^{(t)}\|_{2}\leq\frac{1}{\operatorname{polylog}(d)}$
or for all $i\in[m_{G}]$, $a^{(t)}\sigma^{\prime}(\langle
v_{i}^{(t)},w_{i^{*}}\rangle)\geq\tilde{\Omega}(\log(d))$, we have that: for
all $z,z^{\prime}$ we have that:

$\operatorname{Sigmoid}\left(a^{(t)}\sigma\left(\langle
w_{i^{*}}^{(t)},G^{(t)}(z)\rangle\right)+b^{(t)}\right)=\Theta(1)\times\operatorname{Sigmoid}\left(a^{(t)}\sigma\left(\langle
w_{i^{*}}^{(t)},G^{(t)}(z^{\prime})\rangle\right)+b^{(t)}\right)$

We can immediately obtain the following Lemma:

######  Lemma G.13.

The update of $v_{i}^{(t)}$ is given as:

$\displaystyle\mathbb{E}[v_{i}^{(t+1)}]=v_{i}^{(t)}+\tilde{\Theta}(\eta_{G})a^{(t)}S_{i,t}\left(\left(\sigma^{\prime}(\langle
w_{i^{*}}^{(t)},v_{i}^{(t)}\rangle+\sum_{j\in[m_{G}]}\gamma_{i,j,t}\sigma^{\prime}(\langle
w_{i^{*}}^{(t)},v_{j}^{(t)}\rangle)\right)w_{i}^{(t)}\pm\tilde{O}\left(\frac{\|v_{i}^{(t)}\|_{2}}{\sqrt{d}}\right)^{2}\right)$
(101)

Where $\gamma_{i,j,t}>0$; $\gamma_{i,j,t}=\tilde{\Theta}(1)$ if
$\max_{i\in[m_{G}]}\|v_{i}^{(t)}\|_{2}\leq\frac{1}{\operatorname{polylog}(d)}$
or for all $i\in[m_{G}]$, $a^{(t)}\sigma^{\prime}(\langle
v_{i}^{(t)},w_{i^{*}}\rangle)\geq\tilde{\Omega}(\log(d))$, and
$\gamma_{i,j,t}=\tilde{O}(1)$ otherwise.

Here the additional $\sigma^{\prime}(\langle
w_{i^{*}}^{(t)},v_{j}^{(t)}\rangle)$ part comes from
$\Pr[z_{i},z_{j}=1]=\frac{1}{\operatorname{polylog}(d)}$. The remaining part
of this stage follows from simply apply Lemma E.4.

In stage 5, we bound the update of $a^{(t)},b^{(t)}$ as:

Let

$S_{t}=\mathbb{E}_{z}\left[\operatorname{Sigmoid}\left(a^{(t)}\sigma\left(\langle
w_{i^{*}}^{(t)},G^{(t)}(z)\rangle\right)+b^{(t)}\right)\right]$

In this stage, with the induction hypothesis, we can easily approximate the
sigmoid as:

######  Lemma G.14.

For $t\geq T_{B,4}$, the sigmoid can be approximate as: For every $X,z$

$\operatorname{Sigmoid}(-f^{(t)}(X))=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}(-b^{(t)})$

$\operatorname{Sigmoid}(f^{(t)}(G^{(t)}(z)))=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(a^{(t)}\sigma\left(\langle
w_{i^{*}}^{(t)},G^{(t)}(z)\rangle\right)+b^{(t)}\right)$

Then by the update rule, we can easily conclude that:

######  Lemma G.15.

For $t\geq T_{B,4}$, the update of $a^{(t)},b^{(t)}$ is given as:

$a^{(t+1)}=a^{(t)}+\tilde{O}(\eta_{D})\operatorname{Sigmoid}\left(-b^{(t)}\right)-\tilde{\Omega}(\eta_{D})S_{t}\Lambda^{2}d^{1/4}$

$\mathbb{E}[b^{(t+1)}]=b^{(t)}+\eta_{D}\tau_{b}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\operatorname{Sigmoid}\left(-b^{(t)}\right)-\eta_{D}\tau_{b}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)S_{t}$

This Lemma, together with the induction hypothesis, implies that:

######  Lemma G.16.

We have:

$\displaystyle\sum_{t\geq
T_{B,4}}S_{t}\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}\Lambda^{2}d^{1/4}}\right)$
(102) $\displaystyle\sum_{t\geq
T_{B,4}}\operatorname{Sigmoid}(b^{(t)})\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}}\right)$
(103)

###### Proof of Lemma G.16.

Let us denote
$R=\sum_{t=T_{B,4}}^{T_{0}}\operatorname{Sigmoid}\left(-b^{(t)}\right)$ and
$Q=\sum_{t=T_{B,4}}^{T_{0}}S_{t}$

Sum the update up for $t=T_{B,4}$ to $T_{0}$, we have that:

$\displaystyle a^{(T_{0})}-a^{(T_{B,4})}$
$\displaystyle=\tilde{O}(\eta_{D})R-\tilde{\Omega}(\eta_{D})Q\Lambda^{2}d^{1/4}$
(104) $\displaystyle\mathbb{E}[b^{(T_{0})}]-b^{(T_{B,4})}$
$\displaystyle=\Theta(\eta_{D}\tau_{b})R-\Theta(\eta_{D}\tau_{b})Q$ (105)

By the induction hypothesis that
$|a^{(t)}|\leq\frac{\tilde{O}(1)}{\Lambda^{2}d^{1/4}}$ and
$|b^{t}|=\tilde{O}(1)$, we have that:

$\displaystyle|\tilde{O}(\eta_{D})R-\tilde{\Omega}(\eta_{D})Q\Lambda^{2}d^{1/4}|\leq\frac{\tilde{O}(1)}{\Lambda^{2}d^{1/4}}$
(106)
$\displaystyle|\Theta(\eta_{D}\tau_{b})R-\Theta(\eta_{D}\tau_{b})Q|\leq\tilde{O}(1)$
(107)

Thus, we have:

$\displaystyle\tilde{\Omega}(\eta_{D})Q\Lambda^{2}d^{1/4}$
$\displaystyle\leq\tilde{O}(\eta_{D})R+\frac{\tilde{O}(1)}{\Lambda^{2}d^{1/4}}\leq\tilde{O}(\eta_{D})\left(\frac{1}{\eta_{D}\tau_{b}}+Q\right)+\frac{\tilde{O}(1)}{\Lambda^{2}d^{1/4}}$
(108)

Therefore we have that
$\tilde{\Omega}(\eta_{D})Q\Lambda^{2}d^{1/4}\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}}+\frac{1}{\eta_{D}\Lambda^{2}d^{1/4}}\right)$,
which implies that

$Q\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}\Lambda^{2}d^{1/4}}\right)$

Similarly, we can show that

$\displaystyle R$
$\displaystyle\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}}+Q\right)\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}}+\frac{1}{(\Lambda^{2}d^{1/4})^{2}}+\frac{1}{\Lambda^{2}d^{1/4}}R\right)$
(109)

This implies that $R\leq\tilde{O}\left(\frac{1}{\eta_{D}\tau_{b}}\right)$.

∎

###  G.4 Proof of the induction hypothesis and the final theorem

The final theorem follows immediately from the induction hypothesis ($v$ part)
together with Lemma G.8.

Now it remains to prove the induction hypothesis. We will assume that all the
hypothesises are true until iteration $t$, then we will prove that they are
true at iteration $t+1$.

#### Stage 1

.

To prove the induction hypothesis at Stage 1, for $w$, we have that by Lemma
G.1, we know that: for $*\not=u_{1},u_{2}$,

$\displaystyle|\alpha(w_{i},*,t+1)-\alpha(w_{i},*,t)|\leq\eta_{D}\tilde{O}\left(\frac{1}{d^{1.5}}\right)$
(110)

By $T_{1}\leq{O}\left(\frac{\sqrt{d}}{\eta_{D}a^{(0)}}\right)$ we can conclude
that

$\displaystyle|\alpha(w_{i},*,t+1)-\alpha(w_{i},*,0)|\leq\eta_{D}\tilde{O}\left(\frac{1}{d^{1.5}}\right)\times{O}\left(\frac{\sqrt{d}}{\eta_{D}a^{(0)}}\right)\leq\frac{1}{d^{0.9}}$
(111)

On the $v$ part, again by Lemma G.1, we know that for
$*\notin\\{w_{j}\\}_{j\in[m_{D}]}$:

$\displaystyle|\alpha(v_{i},*,t)-\alpha(v_{i},*,0)|\leq\eta_{G}\left(\frac{1}{C^{0.5}d^{2}}\right)\times{O}\left(\frac{\sqrt{d}}{\eta_{D}a^{(0)}}\right)\leq\frac{1}{d}$
(112)

On the other hand, we know that for $w_{j}$:

$\displaystyle|\alpha(v_{i},w_{j},t+1)-\alpha(v_{i},w_{j},t)|\leq\eta_{G}\left(1+\frac{1}{\operatorname{polylog}(d)}\right)\frac{3}{2m_{G}}\zeta_{t}^{2}$
(113)

Apply Lemma G.2 we complete the proof using Lemma E.3.

As for the $a^{(t)},b^{(t)}$ part, we know that:

$\displaystyle|a^{(t+1)}-a^{(t)}|\leq
O\left(\eta_{D}m_{D}\Upsilon_{t}^{3}\right),|b^{(t)}|\leq\tau_{b}\eta_{D}T_{1}$
(114)

Combine with the update rule of $\Upsilon$ in Eq equation 56, we complete the
proof.

#### Stage 2 and 3

For the $w$ part, we know that by Lemma G.4, we have that for every
$*\not=u_{1},u_{2}$

$\displaystyle|\alpha(w_{i},*,t+1)-\alpha(w_{i},*,t)|\leq\eta_{D}\tilde{O}\left(\frac{1}{d^{1.45}}\right)\operatorname{Sigmoid}(b^{(t)})$
(115)

Now, by Lemma G.11 we have that:

$\sum_{t\in(T_{B,1},T_{B,3}]}\eta_{D}\tau_{b}\left(b^{(t)}\right)\leq\tilde{O}(1)$

This implies that

$|\alpha(w_{i},*,t+1)-\alpha(w_{i},*,T_{B,1})|\leq\eta_{D}\tilde{O}\left(\frac{1}{d^{1.45}}\right)\times\frac{1}{\eta_{b}\eta_{D}}\leq\frac{1}{d^{0.9}}$

For the $v$ part for $t\leq T_{B,2}$, since
$T_{B,2}-T_{B,1}=\tilde{O}(d^{o(1)}/\eta_{D})$, we can easily prove it for
$t\leq T_{B,2}$ as in stage 1. On the other hand, for $t\in(T_{B,2},T_{B,3}]$:
By Lemma G.7 and Lemma G.6, we have that define

$\alpha(t):=\max_{j\in[m_{G}],\ell\in[2]}\langle
v_{j}^{(t)},u_{\ell}\rangle,\quad\beta(t):=\max_{j\in[m_{G}],j^{\prime}\in[m_{G}],j\not=j^{\prime};i\in[m_{D}],i\not=i^{*}}|\langle
v_{j}^{(t)},w_{i}^{(0)}\rangle|+|\langle
v_{j}^{(t)},v_{j^{\prime}}^{(0)}\rangle|$

We have that:

$\displaystyle\mathbb{E}[\alpha(t+1)]\geq\alpha(t)+\eta_{G}\Omega\left(\frac{1}{m_{G}}\right)\operatorname{Sigmoid}(b^{(t)})\alpha(t)^{2}\log^{0.1}(d)$
(116)

and

$\displaystyle\mathbb{E}[\beta(t+1)]\leq\beta(t)+\eta_{G}O\left(\frac{1}{m_{G}}\right)\operatorname{Sigmoid}(b^{(t)})\left(\beta(t)^{2}+\frac{C^{2}}{\sqrt{d}}\alpha(t)^{2}\right)$
(117)

By Lemma D.2 and Lemma E.1 we can show that
$\beta(t)=O\left(\beta(0)+\frac{C^{2}}{\sqrt{d}}\alpha(t)\right)$, which
complete the proof that for all $*\not=w_{i^{*}},u_{1},u_{2}$:

$|\alpha(v_{j},*,t)|\leq\frac{C^{3}}{\sqrt{d}}\|v_{j}^{(t)}\|_{2}$

.

#### Stage 4 and 5

At stage 4 we simply use Lemma E.4, the only remaining part is to show that
$|b^{(t)}|=O(\log(d))$. To see this, we know that by the update formula:

$\nabla_{b}L(X,z)=-\operatorname{Sigmoid}(-f(X))+\operatorname{Sigmoid}(f(G(z)))$

By our induction hypothesis, we know that
$a^{(t)}\left(\sum_{i\in[m_{D}]}\sigma(\langle w_{i}^{(t)},X\rangle)\right)>0$

and $a^{(t)}\left(\sum_{i\in[m_{D}]}\sigma(\langle
w_{i}^{(t)},G(z)\rangle)\right)>0$ . Therefore, $b<O(\log(d))$ is immediate.
Now it remains to show that $b>-O(\log d)$: By the update formula, we have:

$-b^{(t+1)}\leq-b^{(t)}+\eta_{D}\tau_{b}\sum_{i\in[m_{G}]}S_{i,t}$

and by Lemma G.13 and the proof in Lemma E.4, we have that:

$\displaystyle\sum_{i\in[m_{G}]}\langle
v_{i}^{(t+1)},w_{i^{*}}^{(0)}\rangle\geq\sum_{i\in[m_{G}]}\langle
v_{i}^{(t)},w_{i^{*}}^{(0)}\rangle+a^{(t)}\tilde{\Omega}(\eta_{G})\left(\sum_{i\in[m_{G}]}S_{i,t}\right)\left(\sum_{i\in[m_{G}]}\sigma^{\prime}(\langle
v_{i}^{(t)},w_{i^{*}}^{(0)}\rangle)\right)$ (118)

Compare this two updates we can easily obtain that $|b^{(t)}|=O(\log(d))$.

At stage 5, we have that since
$|a^{(t)}|=\tilde{O}\left(\frac{1}{d^{1/4}\Lambda^{2}}\right)$: For every
$j\in[m_{G}],i\in[m_{D}]$

$\displaystyle\|v^{(t+1)}_{j}-v^{(t)}_{j}\|_{2}$
$\displaystyle\leq\tilde{O}(\eta_{G})S_{t}\Lambda^{2}\times\frac{1}{\Lambda^{2}d^{1/4}}$
(119) $\displaystyle\|w^{(t+1)}_{i}-w^{(t)}_{i}\|_{2}$
$\displaystyle\leq\tilde{O}(\eta_{D})\operatorname{Sigmoid}(-b^{(t)})\frac{1}{d^{1/4}\Lambda^{2}}$
(120)

Apply Lemma G.16 we have that:

$\displaystyle\|v^{(t+1)}_{j}-v^{(T_{B,4})}_{j}\|_{2}\leq\tilde{O}(\eta_{G})\times\frac{1}{d^{1/4}}\times\frac{1}{\eta_{D}\tau_{b}\Lambda^{2}d^{1/4}}\leq
d^{0.15}$ (121)
$\displaystyle\|w^{(t+1)}_{i}-w^{(T_{B,4})}_{i}\|_{2}\leq\tilde{O}(\eta_{D})\frac{1}{d^{1/4}\Lambda^{2}}\times\frac{1}{\eta_{D}\tau_{b}}\leq\frac{1}{d^{0.1}}$
(122)

Which proves the induction hypothesis.

##  Appendix H Proof of the learning process in other cases

We now consider other cases, in case 1 of Lemma D.1, the proof is identical to
case 2, the only difference is at Stage 3, we have that $T_{B,3}>T_{0}$.

In case 2, the Stage 1 is identical to the Stage 1, 2, 3 in the balanced case.
For Stage 3, its identical to Stage 4 in the balanced case (the only
difference is to apply Lemma E.5 and the case 2 of Lemma E.2 instead of Lemma
E.4). For Stage 4, its identical to Stage 5 in the balanced case.

At Stage 2, by the induction hypothesis, we know that for $j\not=i_{G}$, we
have that $|\langle
v_{j}^{(t)},w_{j}^{(t)}\rangle|\leq\tilde{O}\left(\frac{1}{C}\|v_{j}^{(t)}\|_{2}\right)$.
Thus, we can approximate the update of $w,v$ as:

$\displaystyle w_{i}^{(t+1)}$
$\displaystyle=w_{i}^{(t)}\pm\tilde{O}(\eta_{D}\sum_{j\in[m_{G}]}\|v_{j}^{(t)}\|_{2}^{2})\Lambda$
(123) $\displaystyle v_{j}^{(t+1)}$
$\displaystyle=v_{j}^{(t)}+\tilde{\Theta}(\eta_{G}\sum_{j\in[m_{G}]}\|v_{j}^{(t)}\|_{2}^{2})w_{i_{G}}^{(0)}\pm\frac{1}{C^{2}}\tilde{\Theta}(\eta_{G}\sum_{j\in[m_{G}]}\|v_{j}^{(t)}\|_{2}^{2})$
(124)

Using the fact that $\eta_{G}\geq\tilde{\Omega}(\sqrt{d})\eta_{D}$ in case 3
we immediately proves the induction hypothesis.

The proof of the theorem follows immediately from the induction hypothesis on
$v$ in this case $v$ only learns noises (linear combinations of
$w_{i}^{(0)}$).

##  Appendix I Normalized SGD

In this section we look at the update of normalized SGD.

Let us define:

$i_{1}^{*}=\operatorname*{arg\,max}_{i\in[m_{D}]}\\{\langle
w_{i}^{(0)},u_{1}\rangle\\}$
$i_{2}^{*}=\operatorname*{arg\,max}_{i\in[m_{D}]}\\{\langle
w_{i}^{(0)},u_{2}\rangle\\}$

Let us define:

$g_{j}^{*}=\operatorname*{arg\,max}_{i\in[m_{D}]}\\{\langle
v_{j}^{(0)},w_{i}^{(0)}\rangle\\}$

Then we first show the following Lemma about initialization:

######  Lemma I.1.

With probability at least $1-o(1)$ over the randomness of the initialization,
the following holds:

  1. 1.

For all $\ell\in[2]$, for all $i\in[m_{D}]$ such that $i\not=i^{*}_{\ell}$, we
have:

$\langle
w_{i_{\ell}^{*}}^{(0)},u_{\ell}\rangle\geq\left(1-\frac{1}{\operatorname{polyloglog}(d)}\right)\langle
w_{i}^{(0)},u_{\ell}\rangle$

  2. 2.

For all $j\in[m_{G}]$, we have that for all $i\in[m_{D}]$ such that
$i\not=g_{j}^{*}$,

$\langle
v_{j}^{(0)},w_{g_{j}^{*}}^{(0)}\rangle\geq\left(1-\frac{1}{\log^{4}(d)}\right)\langle
v_{j}^{(0)},w_{i}^{(0)}\rangle$

  3. 3.

$\\{g_{j}^{*}\\}_{j\in[m_{G}]}=[m_{D}]$.

We now divide the training stage into two: For a sufficiently large
$C=\operatorname{polylog}(d)$, consider the case when
$\eta_{G}=\eta_{D}*C^{-0.6}$.

  1. 1.

Stage 1: When both
$\alpha(w_{i*_{1}},u_{1},t),\alpha(w_{i*_{2}},u_{2},t)\leq\frac{1}{C^{0.95}}$.
Call this iteration $T_{N,1}$.

  2. 2.

Stage 2: After $T_{N,1}$, before $T_{1}$

###  I.1 Induction Hypothesis

We will use the following induction hypothesis: for a

#### Stage 1:

for every $t\leq T_{N,1}$: Let
$\alpha(t):=\max_{\ell\in[2]}\alpha(w_{i_{\ell}^{*}},u_{\ell},t)$,
$\beta(t):=\max_{i\in[m_{G}]}\alpha(v_{i},w_{g_{i}^{*}},t)$.

  1. 1.

Domination: For every $i\in[m_{G}]$, we have:

$|\alpha(v_{i},*,t)-\alpha(v_{i},*,0)|\leq\min\left\\{\frac{1}{C}\alpha(t),\beta(t)\right\\}$

For every $i\in[m_{D}]$, $i\not=i_{1}^{*},i_{2}^{*}$, we have that for
$*\not=u_{1},u_{2}$:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{1}{C}\alpha(t)$

and

$|\alpha(w_{i},u_{1},t)|,|\alpha(w_{i},u_{2},t)|\leq\alpha(t)$

For $i_{1}^{*}$, we have that for every $*\not=u_{1}$,

$|\alpha(w_{i_{1}^{*}},*,t)-\alpha(w_{i_{1}^{*}},*,0)|\leq\frac{1}{C}\alpha(t)$

For $i_{2}^{*}$, we have that for every $*\not=u_{2}$,

$|\alpha(w_{i_{2}^{*}},*,t)-\alpha(w_{i_{2}^{*}},*,0)|\leq\frac{1}{C}\alpha(t)$

  2. 2.

(N.1.2): Growth rate: we have that for every $i\in[m_{D}]$

$\alpha(w_{i_{\ell}^{*}},u_{\ell},t)\in\left(\Omega\left(\frac{1}{m_{D}}\right),1\right)\eta_{D}t$

and for every $i\in[m_{G}]$:

$\alpha(v_{i},w_{g_{i}^{*}},t)\in\left(\Omega\left(\frac{1}{m_{G}^{2}}\right),1\right)\eta_{G}t$

Therefore by our choice of $\eta_{D},\eta_{G}$ we have that $\beta(t)\in
C^{-0.6}\left[\frac{1}{\log^{5}d},\log^{5}d\right]\times\alpha(t)$.

  3. 3.

$a^{(t)}=[0.5,1]a^{(0)},|b^{(t)}|\leq\frac{1}{d^{0.1}}$.

#### Stage 2:

We maintain: For every $t\in[T_{N,1},T_{1}]$:

  1. 1.

(N.1.2) still holds.

  2. 2.

$\alpha(w_{i_{\ell}^{*}},u_{\ell},t)\in\left[\frac{1}{C},\operatorname{polylogloglog}(d)\right]$,
$\beta(t)\in
C^{-0.6}\left[\frac{1}{\log^{5}d},\log^{5}d\right]\times\alpha(t)$,
$a^{(t)}=\Omega(\alpha(w_{i_{\ell}^{*}},u_{\ell},t))$.

  3. 3.

$w_{i}$’s are good: For every $i\not=i_{1}^{*},i_{2}^{*}$, for every $*$:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{1}{C}\alpha(t)$

and for $\ell\in[2]$: for every $*\not=u_{\ell}$, we have:

$|\alpha(w_{i},*,t)-\alpha(w_{i},*,0)|\leq\frac{1}{C}\alpha(t)$

  4. 4.

$v_{i}$’s are good: For every $i\in[m_{G}]$ and every $j\in[m_{D}]$,
$j\not=g_{i}^{*}$:

$\langle v^{(t)}_{i},w_{g_{i}^{*}}^{(t)}\rangle\geq C^{0.9}|\langle
v^{(t)}_{i},w_{j}^{(t)}\rangle|$

and for $g_{i}^{*}\not=i_{\ell}^{*}$, we have that:

$\langle v^{(t)}_{i},w_{g_{i}^{*}}^{(t)}\rangle\geq C^{0.9}|\langle
v^{(t)}_{i},u_{\ell}\rangle|$

For $g_{i}^{*}=i_{\ell}$ , we have that $\langle
v^{(t)}_{i},u_{\ell}\rangle\geq-\frac{1}{C^{0.5}}\beta(t)$

###  I.2 Stage 1 training

With the induction hypothesis, we can show the following Lemma:

######  Lemma I.2.

For $t\leq T_{N,1}$, for
$\varepsilon_{t}:=\frac{(\alpha(t)+d^{-0.5})^{2}}{C^{1.5}}+C^{0.5}(\alpha(t)+d^{-0.5})^{3}$,
when the sample is $X\in\left\\{u_{1},u_{2}\right\\}$, the update of
$w_{i}^{(t)}$ can be approximate as:

$\displaystyle w_{i}^{(t+1)}$
$\displaystyle=w_{i}^{(t)}+\eta_{D}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{\sigma^{\prime}(\langle
w_{i}^{(t)},X\rangle)X\pm\varepsilon_{t}}{\sqrt{\sum_{j\in[m_{D}]}\sigma^{\prime}(\langle
w_{j}^{(t)},X\rangle)^{2}\|X\|_{2}^{2}}}$ (125)

Which can be further simplified as:

$\displaystyle\mathbb{E}[\langle w_{i}^{(t+1)},u_{\ell}\rangle]$
$\displaystyle=\langle
w_{i}^{(t)},u_{\ell}\rangle+\eta_{D}\frac{1}{2}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{\sigma^{\prime}(\langle
w_{i}^{(t)},u_{\ell}\rangle)\pm\varepsilon_{t}}{\sqrt{\sum_{j\in[m_{D}]}\sigma^{\prime}(\langle
w_{j}^{(t)},u_{\ell}\rangle)^{2}}}\pm\eta_{D}\gamma$ (126)

When $z=e_{i}$, the update of $v$ can be approximate as: For
$\delta_{t}:=O\left(\frac{1}{C^{0.94}}(\beta(t)+\frac{1}{d})^{2}\right)$:

$\displaystyle v_{i}^{(t+1)}$
$\displaystyle=v_{i}^{(t)}+\eta_{G}\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{\sum_{j\in[m_{D}]}\sigma^{\prime}(\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle)w_{j}^{(0)}\pm\delta_{t}}{\|\sum_{j\in[m_{D}]}\sigma^{\prime}(\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle)w_{j}^{(0)}\|_{2}}$ (127)

Where we have:

$\langle v_{i}^{(t)},w_{j}^{(t)}\rangle=\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle\pm\frac{1}{C^{0.9}}\beta(t)$

###### Proof of the update Lemma I.2.

By the induction hypothesis, We have that:

$\displaystyle\langle v_{i}^{(t)},w_{j}^{(t)}\rangle$ $\displaystyle=\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle+\langle
v_{i}^{(t)},w_{j}^{(t)}-w_{j}^{(0)}\rangle$ (128) $\displaystyle=\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle+\langle
v_{i}^{(0)},w_{j}^{(t)}-w_{j}^{(0)}\rangle+\langle
v_{i}^{(t)}-v_{i}^{(0)},w_{j}^{(t)}-w_{j}^{(0)}\rangle$ (129)
$\displaystyle=\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle\pm\tilde{O}\left(\frac{1}{\sqrt{d}}\alpha(t)\right)\pm
O\left(\frac{1}{C^{0.94}}\beta(t)\right)$ (130) $\displaystyle=\langle
v_{i}^{(t)},w_{j}^{(0)}\rangle\pm\frac{1}{C^{0.9}}\beta(t)$ (131)

Here we use the fact that $\|w_{j}^{(0)}-w_{j}^{(t)}\|_{2}\leq
O\left(\frac{1}{C^{0.95}}\right)$ from the induction hypothesis.

Consider the update of $w_{i}$, we have that: at stage 1, we must have
$|f(X)|,|f(G(z))|\leq\frac{1}{\operatorname{polylog}(d)}$. Therefore,

$\displaystyle\nabla_{w_{i}}L(X,z)$
$\displaystyle=\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)a^{(t)}\sigma^{\prime}(\langle
w_{i}^{(t)},X\rangle)X-a^{(t)}\sigma^{\prime}(\langle
w_{i}^{(t)},G^{(t)}(z)\rangle)G^{(t)}(z)$ (132)

By the induction hypothesis, we have that by $\beta(t)\leq\alpha(t)$, it holds
that:

$\|\sigma^{\prime}(\langle w_{i}^{(t)},G^{(t)}(z)\rangle)G^{(t)}(z)\|_{2}\leq
O\left(\frac{\alpha(t)}{C}+\frac{1}{C\sqrt{d}}\right)^{2}m_{G}^{2}\times
m_{G}\left(\frac{1}{\sqrt{d}}+\frac{\alpha(t)}{C}\right)\leq\epsilon_{t}$

On the other hand, we must have that when $X=u_{\ell}$, we have

$\displaystyle\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\sigma^{\prime}(\langle
w_{i^{*}_{\ell}}^{(t)},X\rangle)\geq\left(\frac{1}{m_{D}}\alpha(t)+\frac{1}{\sqrt{d}}\right)^{2}\geq\operatorname{polylog}(d)\epsilon_{t}$
(133)

This completes the proof of the $w_{i}$ part. For $v_{i}$ part the proof is
the same using the fact that $\|w_{j}^{(0)}-w_{j}^{(t)}\|_{2}\leq
O\left(\frac{1}{C^{0.95}}\right)$ from the induction hypothesis.

∎

###  I.3 Stage 2 training

In this stage, we can maintain the following simple update rule: For $w_{i}$:

######  Lemma I.3.

For every $t\in(T_{N,1},T_{1}]$, we have that: for every $i\in[m_{D}]$, for
$i=i^{*}_{\ell}$:

$\mathbb{E}[w_{i}^{(t+1)}]=w_{i}^{(t)}+\Theta(\eta_{D})u_{\ell}\pm\eta_{D}\frac{1}{C^{1.501}}\pm\eta_{D}\gamma$

and for $i\not=i^{*}_{1},i^{*}_{2}$,

$\mathbb{E}[w_{i}^{(t+1)}]=w_{i}^{(t)}\pm\eta_{D}\frac{1}{C^{1.501}}\pm\eta_{D}\gamma$

For $v_{i}$:

$\displaystyle\mathbb{E}[v_{i}^{(t+1)}]=v_{i}^{(t)}+\left(1\pm\frac{1}{\operatorname{polylog}(d)}\right)\frac{1}{m_{G}}\eta_{G}\frac{w_{g_{i}^{*}}^{(t)}}{\|w_{g_{i}^{*}}^{(t)}\|_{2}}\pm\eta_{G}\frac{1}{C^{1.5}}$
(134)

###### Proof of Lemma I.3.

This Lemma can be proved identically to Lemma I.2: By the induction
hypothesis, we have

$\displaystyle|\langle w_{i}^{(t)},v_{j}^{(t)}\rangle|\leq\log^{5}\beta(t)$
(135)

Therefore,

$\|\sigma^{\prime}(\langle w_{i}^{(t)},G^{(t)}(z)\rangle)v_{j}^{(t)}\|_{2}\leq
C^{0.01}\beta(t)^{3}\leq\frac{1}{C^{1.51}}\alpha(t)^{2}$

Which implies that:

$\displaystyle w_{i}^{(t+1)}$
$\displaystyle=w_{i}^{(t)}+\eta_{D}\frac{\sigma^{\prime}(\langle
w_{i}^{(t)},X\rangle)X\pm\frac{1}{C^{1.51}\alpha(t)^{2}}}{\sqrt{\sum_{j\in[m_{D}]}(\sigma^{\prime}(\langle
w_{j}^{(t)},X\rangle)^{2}\|X\|_{2}^{2}+\sum_{j\in[m_{D}]}\left(\sigma^{\prime}(\langle
w_{j}^{(t)},X\rangle)^{2}\right)^{3}}}$ (136)