Title: SGD LEARNS ONE-LAYER NETWORKS IN WGANS

Abstract: Generative adversarial networks (GANs) are a widely used framework for learning generative models. Wasserstein GANs (WGANs), one of the most successful variants of GANs, require solving a minmax optimization problem to global optimality, but are in practice successfully trained using stochastic gradient descent-ascent. In this paper, we show that, when the generator is a one-layer network, stochastic gradient descent-ascent converges to a global solution with polynomial time and sample complexity.

1. INTRODUCTION:

Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) are a prominent framework for learning generative models of complex, real-world distributions given samples from these distributions. GANs and their variants have been successfully applied to numerous datasets and tasks, including image-to-image translation (Isola et al., 2017) , image super-resolution (Ledig et al., 2017) , domain adaptation (Tzeng et al., 2017) , probabilistic inference (Dumoulin et al., 2016) , compressed sensing (Bora et al., 2017 ) and many more. These advances owe in part to the success of Wasserstein GANs (WGANs) (Arjovsky et al., 2017; Gulrajani et al., 2017) , leveraging the neural net induced integral probability metric to better measure the difference between a target and a generated distribution.

Along with the afore-described empirical successes, there have been theoretical studies of the statistical properties of GANs-see e.g. (Zhang et al., 2018; Arora et al., 2017; 2018; Bai et al., 2018; Dumoulin et al., 2016) and their references. These works have shown that, with an appropriate design of the generator and discriminator, the global optimum of the WGAN objective identifies the target distribution with low sample complexity.

On the algorithmic front, prior work has focused on the stability and convergence properties of gradient descent-ascent (GDA) and its variants in GAN training and more general min-max optimization problems; see e.g. (Nagarajan & Kolter, 2017; Heusel et al., 2017; Mescheder et al., 2017; 2018; Daskalakis et al., 2017; Daskalakis & Panageas, 2018a; b; Gidel et al., 2019; Liang & Stokes, 2019; Mokhtari et al., 2019; Jin et al., 2019; Lin et al., 2019) and their references. It is known that, even in min-max optimization problems with convex-concave objectives, GDA may fail to compute the min-max solution and may even exhibit divergent behavior. Hence, these works have studied conditions under which GDA converges to a globally optimal solution under a convex-concave objective, or different types of locally optimal solutions under nonconvex-concave or nonconvex-nonconcave objectives. They have also identified variants of GDA with better stability properties in both theory and practice, most notably those using negative momentum.

In the context of GAN training, Feizi et al. (2017) show that for WGANs with a linear generator and quadratic discriminator GDA succeeds in learning a Gaussian using polynomially many samples in the dimension. In the same vein, we are the first to our knowledge to study the global convergence properties of stochastic GDA in the GAN setting, and establishing such guarantees for non-linear generators. In particular, we study the WGAN formulation for learning a single-layer generative model with some reasonable choices of activations including tanh, sigmoid and leaky ReLU.

Our contributions. For WGAN with a one-layer generator network using an activation from a large family of functions and a quadratic discriminator, we show that stochastic gradient descent-ascent learns a target distribution using polynomial time and samples, under the assumption that the target distribution is realizable in the architecture of the generator. This is achieved by a) analysis of the dynamics of stochastic gradient-descent to show it attains a global optimum of the minmax problem, and b) appropriate design of the discriminator to ensure a parametric O( 1 √ n ) statistical rate (Zhang et al., 2018; Bai et al., 2018) . Related Work. We briefly review relevant results in GAN training and learning generative models:

-Optimization viewpoint. For standard GANs and WGANs with appropriate regularization, Nagarajan & Kolter (2017) , Mescheder et al. (2017) and Heusel et al. (2017) establish sufficient conditions to achieve local convergence and stability properties for GAN training. At the equilibrium point, if the Jacobian of the associated gradient vector field has only eigenvalues with negative real-part at the equilibrium point, GAN training is verified to converge locally for small enough learning rates. A follow-up paper by (Mescheder et al., 2018) shows the necessity of these conditions by identifying a prototypical counterexample that is not always locally convergent with gradient descent based GAN optimization. However, the lack of global convergence prevents the analysis to provide any guarantees of learning the real distribution.

The work of (Feizi et al., 2017) described above has similar goals as our paper, namely understanding the convergence properties of basic dynamics in simple WGAN formulations. However, they only consider linear generators, which restrict the WGAN model to learning a Gaussian. Our work goes a step further, considering WGANs whose generators are one-layer neural networks with a broad selection of activations. We show that with a proper gradient-based algorithm, we can still recover the ground truth parameters of the underlying distribution.

More broadly, WGANs typically result in nonconvex-nonconcave min-max optimization problems. In these problems, a global min-max solution may not exist, and there are various notions of local min-max solutions, namely local min-local max solutions Daskalakis & Panageas (2018b) , and local min solutions of the max objective Jin et al. (2019) , the latter being guaranteed to exist under mild conditions. In fact, Lin et al. (2019) show that GDA is able to find stationary points of the max objective in nonconvex-concave objectives. Given that GDA may not even converge for convexconcave objectives, another line of work has studied variants of GDA that exhibit global convergence to the min-max solution Daskalakis et al. (2017) ; Daskalakis & Panageas (2018a); Gidel et al. (2019) ; Liang & Stokes (2019) ; Mokhtari et al. (2019) , which is established for GDA variants that add negative momentum to the dynamics. While the convergence of GDA with negative momentum is shown in convex-concave settings, there is experimental evidence supporting that it improves GAN training (Daskalakis et al., 2017; Gidel et al., 2019) .

-Statistical viewpoint. Several works have studied the issue of mode collapse. One might doubt the ability of GANs to actually learn the distribution vs just memorize the training data (Arora et al., 2017; 2018; Dumoulin et al., 2016) . Some corresponding cures have been proposed. For instance, Zhang et al. (2018) ; Bai et al. (2018) show for specific generators combined with appropriate parametric discriminator design, WGANs can attain parametric statistical rates, avoiding the exponential in dimension sample complexity (Liang, 2018; Bai et al., 2018; Feizi et al., 2017) . Wu et al. (2019) provides an algorithm to learn the distribution of a single-layer ReLU generator network. While our conclusion appears similar, our focus is very different. Our paper targets understanding when a WGAN formulation trained with stochastic GDA can learn in polynomial time and sample complexity. Their work instead relies on a specifically tailored algorithm for learning truncated normal distributions Daskalakis et al. (2018) .

Recent work of:



2. PRELIMINARIES:

We consider GAN formulations for learning a generator G A : R k → R d of the form z → x = φ(Az), where A is a d × k parameter matrix and φ some activation function. We consider discriminators D v : R d → R or D V : R d → R that are linear or quadratic forms respectively for the different purposes of learning the marginals or the joint distribution. We assume latent variables z are sampled from the normal N (0, I k×k ), where I k×k denotes the identity matrix of size k. The real/target distribution outputs samples x ∼ D = G A * (N (0, I k0×k0 )), for some ground truth parameters A * , where A * is d × k 0 , and we take k ≥ k 0 for enough expressivity, taking k = d when k 0 is unknown.

The Wasserstain GAN under our choice of generator and discriminator is naturally formulated as:

min A∈R d×k max v∈R d f (A, v) ≡ E x∼D D v (x) -E z∼N (0,I k×k ) D v (G A (z)) . 1

We use a i to denote the i-th row vector of A. We sometimes omit the 2 subscript, using x to denote the 2-norm of vector x, and X to denote the spectral norm of matrix X. S n ⊂ R n×n represents all the symmetric matrices of dimension n × n. We use Df (X 0 )[B] to denote the directional derivative of function f at point X 0 with direction B: Df

(X 0 )[B] = lim t→0 f (X0+tB)-f (X0) t .

3 WARM-UP: LEARNING THE MARGINAL DISTRIBUTIONS As a warm-up, we ask whether a simple linear discriminator is sufficient for the purposes of learning the marginal distributions of all coordinates of D. Notice that in our setting, the i-th output of the generator is φ(x) where x ∼ N (0, a i 2 ), and is thus solely determined by a i 2 . With a linear discriminator D v (x) = v x, our minimax game becomes:

min A∈R d×k max v∈R d f 1 (A, v) ≡ E x∼D v x -E z∼N (0,I k×k ) v φ(Az) .

(1)

Notice that when the activation φ is an odd function, such as the tanh activation, the symmetric property of the Gaussian distribution ensures that E x∼D [v x] = 0, hence the linear discriminator in f 1 reveals no information about A * . Therefore specifically for odd activations (or odd plus a constant activations), we instead use an adjusted rectified linear discriminator

D v (x) ≡ v R(x -C)

to enforce some bias, where C = 1 2 (φ(x) + φ(-x)) for all x, and R denotes the ReLU activation. Formally, we slightly modify our loss function as:

f1 (A, v) ≡ E x∼D v R(x -C) -E z∼N (0,I k×k ) v R(φ(Az) -C) .

(2)

We will show that we can learn each marginal of D if the activation function φ satisfies the following. Assumption 1. The activation function φ satisfies either one of the following: 1. φ is an odd function plus constant, and φ is monotone increasing; 2. The even component of φ, i.e. 1 2 (φ(x)+φ(-x)), is positive and monotone increasing on x ∈ [0, ∞). Remark 1. All common activation functions like (Leaky) ReLU, tanh or sigmoid function satisfy Assumption 1. Lemma 1. Suppose A * = 0. Consider f 1 with activation that satisfies Assumption 1.2 and f1 with activation that satisfies Assumption 1.1. The stationary points of such f 1 and f1 yield parameters A satisfying

a i = a * i , ∀i ∈ [d].

To bound the capacity of the discriminator, similar to the Lipschitz constraint in WGAN, we regularize the discriminator. For the regularized formulation we have: Theorem 1. In the same setting as Lemma 1, alternating gradient descent-ascent with proper learning rates on

min A max v {f 1 (A, v) -v 2 /2} or respectively min A max v { f1 (A, v) -v 2 /2} recovers A such that a i = a * i , ∀i ∈ [d].

All the proofs of the paper can be found in the appendix. We show that all local min-max points in the sense of (Jin et al., 2019) of the original problem are global min-max points and recover the correct norm of a * i , ∀i. Notice for the source data distribution

x = (x 1 , x 2 , • • • x d ) ∼ D with activation φ, the marginal distribution of each x i follows φ(N (0, a * i 2 ))

and is determined by a * i . Therefore we have learned the marginal distribution for each entry i. It remains to learn the joint distribution.

4. LEARNING THE JOINT DISTRIBUTION:

In the previous section, we utilize a (rectified) linear discriminator, such that each coordinate v i interacts with the i-th random variable. With the (rectified) linear discriminator, WGAN learns the correct a i , for all i. However, since there's no interaction between different coordinates of the random vector, we do not expect to learn the joint distribution with a linear discriminator.

To proceed, a natural idea is to use a quadratic discriminator D V (x) := x V x = xx , V to enforce component interactions. Similar to the previous section, we study the regularized version:

EQUATION

where

f 2 (A, V ) = E x∼D D V (x) -E z∼N (0,I k×k ) D V (φ(Az)) = E x∼D xx -E z∼N (0,I k×k ) φ(Az)φ(Az) , V .

By adding a regularizer on V and explicitly maximizing over V :

g(A) ≡ max V f 2 (A, V ) - 1 2 V 2 F = 1 2 E x∼D xx -E z∼N (0,I k×k ) φ(Az)φ(Az) 2 F .

In the next subsection, we first focus on analyzing the second-order stationary points of g, then we establish that gradient descent ascent converges to second-order stationary points of g .

4.1. GLOBAL CONVERGENCE FOR OPTIMIZING THE GENERATING PARAMETERS:

We first assume that both A and A * have unit row vectors, and then extend to general case since we already know how to learn the row norms from Section 3. To explicitly compute g(A), we rely on the property of Hermite polynomials. Since normalized Hermite polynomials {h i } ∞ i=0 forms an orthonomal basis in the functional space, we rewrite the activation function as φ(x) = ∞ i=0 σ i h i , where σ i is the i-th Hermite coefficient. We use the following claim:

Claim 1 ((Ge et al., 2017) Claim 4.2). Let φ be a function from R to R such that φ ∈ L 2 (R, e -x 2 /2 ),

and let its Hermite expansion be φ = ∞ i=1 σ i h i . Then, for any unit vectors u, v ∈ R d , we have that

E x∼N (0,I d×d ) φ(u x)φ(v x) = ∞ i=0 σ 2 i (u v) i .

Therefore we could compute the value of f 2 explicitly using the Hermite polynomial expansion:

f 2 (A, V ) = ∞ i=0 σ 2 i (A * (A * ) ) •i -(AA ) •i , V .

Here X •i is the Hadamard power operation where (X •i ) jk = (X jk ) i . Therefore we have:

g(A) = 1 2 ∞ i=0 σ 2 i (A * (A * ) ) •i -(AA ) •i 2 F

We reparametrize with Z = AA and define g

(Z) = g(A) with individual component functions gjk (z) ≡ 1 2 ( ∞ i=0 σ 2 i ((z * jk ) i -z i )) 2 . Accordingly z * jk = a * j , a * k is the (j, k)-th component of the ground truth covariance matrix A * (A * ) .

Assumption 2. The activation function φ is an odd function plus constant. In other words, its Hermite expansion φ = ∞ i=0 σ i h i satisfies σ i = 0 for even i ≥ 2. Additionally we assume σ 1 = 0. Remark 2. Common activations like tanh and sigmoid satisfy Assumption 2. Lemma 2. For activations including leaky ReLU and functions satisfying Assumption 2, g(Z) has a unique stationary point where Z = A * (A * ) .

Notice g(Z) = jk gjk (z jk ) is separable across z jk , where each gjk is a polynomial scalar function. Lemma 2 comes from the fact that the only zero point for g jk is z jk = z * jk , for odd activation φ and leaky ReLU. Then we migrate this good property to the original problem we want to solve: Problem 1. We optimize over function g when a * i = 1, ∀i:

min A    g(A) = 1 2 ∞ i=0 σ 2 i (A * (A * ) ) •i -(AA ) •i 2 F    s.t. a i a i = 1, ∀i.

Existing work Journée et al. (2008) connects g(Z) to the optimization over factorized version for g(A) (g(A) ≡ g(AA )). Specifically, when k = d, all second-order stationary points for g(A) are first-order stationary points for g(Z). Though g is not convex, we are able to show that its first-order stationary points are global optima when the generator is sufficiently expressive, i.e., k ≥ k 0 . In reality we won't know the latent dimension k 0 , therefore we just choose k = d for simplicity. We make the following conclusion: Theorem 2. For activations including leaky ReLU and functions satisfying Assumption 2, when k = d, all second-order KKT points for problem 1 are its global minimum. Therefore alternating projected gradient descent-ascent on Eqn.

(3) converges to A : AA = A * (A * ) .

The extension for non-unit vectors is straightforward, and we defer the analysis to the Appendix.

5. FINITE SAMPLE ANALYSIS:

Algorithm 1 Online stochastic gradient descent ascent on WGAN 1: Input: n training samples:

x 1 , x 2 , • • • x n , where each x i ∼ φ(A * z), z ∼ N (0, I k×k )

, learning rate for generating parameters η, number of iterations T . 2: Random initialize generating matrix A (0) . 3:

for t = 1, 2, • • • , T do 4: Generate m latent variables z (t) 1 , z (t) 2 , • • • , z (t)

m ∼ N (0, I k×k ) for the generator. The empirical function becomes

f (t) m,n (A, V ) = 1 m m i=1 φ(Az (t) i )φ(Az (t) i ) - 1 n n i=1 x i x i , V - 1 2 V 2 5:

Gradient ascent on V with optimal step-size η V = 1:

V (t) ← V (t) -η V ∇ V f (t) m,n (A (t-1) , V (t-1)

).

6::

Sample noise e uniformly from unit sphere 7:

Projected Gradient Descent on A, with constraints C = {A|(AA ) ii = (A * A * ) ii } :

A (t) ← Proj C (A (t-1) -η(∇ A f (t) m,n (A (t-1) , V (t)

) + e)). 8: end for 9: Output:

A (T ) (A (T ) )

In this section, we consider analyzing Algorithm 1, i.e., gradient descent ascent on the following:

EQUATION

)

Notice in each iteration, gradient ascent with step-size 1 finds the optimal solution for V . By Danskin's theorem (Danskin, 2012), our min-max optimization is essentially gradient descent over

g(t) m,n (A) ≡ max V f (t) m,n (A, V ) = 1 2 1 m m i=1 φ(Az (t) i )φ(Az (t) i ) -1 n n i=1 x i x i 2 F with a batch of samples {z (t) i }, i.e., stochastic gradient descent for f n (A) ≡ E zi∼N (0,I k×k ),∀i∈[m] [g m,n (A)].

Therefore to bound the difference between f n (A) and the population risk g(A), we analyze the sample complexity required on the observation side (x i ∼ D, i ∈ [n]) and the mini-batch size required on the learning part (φ(Az j ), z j ∼ N (0, I k×k ), j ∈ [m]). We will show that with large enough n, m, the algorithm specified in Algorithm 1 that optimizes over the empirical risk will yield the ground truth covariance matrix with high probability.

Our proof sketch is roughly as follows:

1. With high probability, projected stochastic gradient descent finds a second order stationary point Â of f n (•) as shown in Theorem 31 of (Ge et al., 2015) .

2. For sufficiently large m, our empirical objective, though a biased estimator of the population risk g(•), achieves good -approximation to the population risk on both the gradient and Hessian (Lemmas 4&5). Therefore Â is also an O( )-approximate second order stationary point (SOSP) for the population risk g(A).

3. We show that any -SOSP Â for g(A) yields an O( )-first order stationary point (FOSP) Ẑ ≡ Â Â for the semi-definite programming on g(Z) (Lemma 6).

4. We show that any O( )-FOSP of function g(Z) induces at most O( ) absolute error compared to the ground truth covariance matrix Z * = A * (A * ) (Lemma 7).

5.1. OBSERVATION SAMPLE COMPLEXITY:

For simplicity, we assume the activation and its gradient satisfy Lipschitz continuous, and let the Lipschitz constants be 1 w.l.o.g.: Assumption 3. Assume the activation is 1-Lipschitz and 1-smooth.

To estimate observation sample complexity, we will bound the gradient and Hessian for the population risk and empirical risk on the observation samples:

EQUATION

g n (A) ≡ 1 2 1 n n i=1 x i x i -E z∼N (0,I k×k ) φ(Az)φ(Az) 2 F . Claim 2. ∇g(A) -∇g n (A) = 2E z diag(φ (Az))(X -X n )φ(Az)z , where X = E x∼D [xx ], and X n = 1 n n i=1 x i x i .

The directional derivative with arbitrary direction B is:

D∇g(A)[B] -D∇g n (A)[B] = 2E z diag(φ (Az))(X n -X)φ (Az) • (Bz)z + 2E z diag(φ (Az) • (Bz))(X n -X)φ(Az)z

Lemma 3. Suppose the activation satisfies Assumption 3.

Pr[ X -X n ≤ X ] ≥ 1 -δ, for n ≥ Θ(d/ 2 log 2 (1/δ)) 2 .

Lemma 4. Suppose the activation satisfies Assumption 2&3. With samples

n ≥ Θ(d/ 2 log 2 (1/δ)), ∇g(A) -∇g n (A) 2 ≤ O( d A 2 ) with probability 1 -δ. Meanwhile, D∇g(A)[B] - D∇g n (A)[B] 2 ≤ O( d 3/2 A 2 B 2 )

with probability 1 -δ.

5.2. BOUNDING MINI-BATCH SIZE:

Normally for empirical risk for supervised learning, the mini-batch size can be arbitrarily small since the estimator of the gradient is unbiased. However in the WGAN setting, notice for each iteration, we randomly sample a batch of random variables

{z i } i∈[m]

, and obtain a gradient of

gm,n (A) ≡ 1 2 1 n n i=1 x i x i -1 m m j=1 φ(Az j )φ(Az j ) 2 F

, in Algorithm 1. However, the finite sum is inside the Frobenius norm and the gradient on each mini-batch may no longer be an unbiased estimator for our target g n

(A) = 1 2 1 n n i=1 x i x i -E z φ(Az)φ(Az) 2 F .

In other words, we conduct stochastic gradient descent over the function f (A) ≡ E z gm,n (A).

Therefore we just need to analyze the gradient error between this f (A) and g n (A) (i.e. gm,n is almost an unbiased estimator of g n ). Finally with the concentration bound derived in last section, we get the error bound between f (A) and g(A). Lemma 5. The empirical risk gm,n is almost an unbiased estimator of g n . Specifically, the expected function

f (A) = E zi∼N (0,I k×k ),i∈[m] [g m,n ] satisfies: ∇f (A) -∇g n (A) ≤ O( 1 m A 3 d 2 ).

For arbitrary direction matrix B,

D∇f (A)[B] -D∇g n (A)[B] ≤ O( 1 m B A 3 d 5/2 ).

In summary, we conduct concentration bound over the observation samples and mini-batch sizes, and show the gradient of f (A) that Algorithm 1 is optimizing over has close gradient and Hessian with the population risk g(A). Therefore a second-order stationary point (SOSP) for f (A) (that our algorithm is guaranteed to achieve) is also an approximated SOSP for g(A). Next we show such a point also yield an approximated first-order stationary point of the reparametrized function g(Z) ≡ g(A), ∀Z = AA .

5.3. RELATION ON APPROXIMATE OPTIMALITY:

In this section, we establish the relationship between g and g. We present the general form of our target Problem 1:

min A∈R d×k g(A) ≡ g(AA ) (5) s.t. Tr(A X i A) = y i , X i ∈ S, y i ∈ R, i = 1, • • • , n.

Similar to the previous section, the stationary property might not be obvious on the original problem. Instead, we could look at the re-parametrized version as:

min Z∈S g(Z) (6) s.t. Tr(X i Z) = y i , X i ∈ S, y i ∈ R, i = 1, • • • , n, Z 0, Definition 1. A matrix A ∈ R d×k

is called an -approximate second-order stationary point ( -SOSP) of Eqn. ( 5) if there exists a vector λ such that:

   Tr(A X i A) = y i , i ∈ [n] (∇ Z g(AA ) - n i=1 λ i X i )ã j ≤ ãj , {ã j } j span the column space of A Tr(B D∇ A L(A, λ)[B]) ≥ -B 2 , ∀B s.t. Tr(B X i A) = 0

Here L(A, λ) is the Lagrangian form g(AA ) -n i=1 λ i (Tr(A X i A) -y i ). Specifically, when = 0 the above definition is exactly the second-order KKT condition for optimizing (5). Next we present the approximate first-order KKT condition for (6): Definition 2. A symmetric matrix Z ∈ S n is an -approximate first order stationary point of function (6) ( -FOSP) if and only if there exist a vector σ ∈ R m and a symmetric matrix S ∈ S such that the following holds:

         Tr(X i Z) = y i , i ∈ [n] Z 0, S -I, S ãj ≤ ãj , {ã j } j span the column space of Z S = ∇ Z g(Z) - n i=1 σ i X i . Lemma 6. Let latent dimension k = d.

For an -SOSP of function ( 5) with A and λ, it infers an -FOSP of function ( 6) with Z, σ and S that satisfies: Z = AA , σ = λ and S = ∇ Z g(AA )i λ i X i . Now it remains to show an -FOSP of g(Z) indeed yields a good approximation for the ground truth parameter matrix. Lemma 7. If Z is an -FOSP of function ( 6), then Z -Z * F ≤ O( ). Here Z * = A * (A * ) is the optimal solution for function (6).

Together with the previous arguments, we finally achieve our main theorem on connecting the recovery guarantees with the sample complexity and batch size3 : Theorem 3. For arbitrary δ < 1, , given small enough learning rate η < 1/poly(d, 1/ , log(1/δ)), let sample size n ≥ Θ(d 5 / 2 log 2 (1/δ)), batch size m ≥ O(d 5 / ), for large enough T =poly(1/η, 1/ , d, log(1/δ)), the output of Algorithm 1 satisfies A (T ) (A (T ) ) -Z * F ≤ O( ) with probability 1 -δ, under Assumptions 2 & 3 and k = d. 

6. SIMULATIONS:

In this section, we provide simple experimental results to validate the performance of stochastic gradient descent ascent and provide experimental support for our theory.

We focus on Algorithm 1 that targets to recover the parameter matrix. We conduct a thorough empirical studies on three joint factors that might affect the performance: the number of observed samples m (we set n = m as in general GAN training algorithms), the different choices of activation function φ, and the output dimension d. In Figure 1 we plot the relative error for parameter estimation decrease over the increasing sample complexity. We fix the hidden dimension k = 2, and vary the output dimension over {3, 5, 7} and sample complexity over {500, 1000, 2000, 5000, 10000}. Reported values are averaged from 20 runs and we show the standard deviation with the corresponding colored shadow. Clearly the recovery error decreases with higher sample complexity and smaller output dimension. To visually demonstrate the learning process, we also include a simple comparison for different φ: i.e. leaky ReLU and tanh activations, when k = 1 and d = 2. We set the ground truth covariance matrix to be [1, 1; 1, 1], and therefore a valid result should be [1, 1] or [-1, -1]. From Figure 2 we could see that for both leaky ReLU and tanh, the stochastic gradient descent ascent performs similarly with exact recovery of the ground truth parameters.

7. CONCLUSION:

We analyze the convergence of stochastic gradient descent ascent for Wasserstein GAN on learning a single layer generator network. We show that stochastic gradient descent ascent algorithm attains the global min-max point, and provably recovers the parameters of the network with absolute error measured in Frobenius norm, from Θ(d 5 / 2 ) i.i.d samples.

A OMITTED PROOF FOR LEARNING THE DISTRIBUTION A.1 STATIONARY POINT FOR MATCHING FIRST MOMENT:

Proof of Lemma 1. To start with, we consider odd-plus-constant monotone increasing activations. Notice that by proposing a rectified linear discriminator, we have essentially modified the activation function as φ := R(φ -C), where C = 1 2 (φ(x) + φ(-x)) is the constant bias term of φ. Observe that we can rewrite the objective f1 for this case as follows:

f 1 (A, v) = E z∼N (0,I k 0 ×k 0 ) v φ(A * z) -E z∼N (0,I k×k ) v φ(Az).

Moreover, notice that φ is positive and increasing on its support which is [0, +∞). Now let us consider the other case in our statement where φ has a positive and monotone increasing even component in [0, +∞). In this case, let us take:

φ(x) = φ(x) + φ(-x), x ≥ 0 0, o.w.

Because of the symmetry of the Gaussian distribution, we can rewrite the objective function for this case as follows:

f 1 (A, v) = E z∼N (0,I k 0 ×k 0 ) v φ(A * z) -E z∼N (0,I k×k ) v φ(Az).

Moreover, notice that φ is positive and increasing on its support which is [0, +∞).

To conclude, in both cases, the optimization objective can be written as follows, where φ satisfies Assumption 1.2 and is only non-zero on [0, +∞).

f 1 (A, v) = E z∼N (0,I k 0 ×k 0 ) v φ(A * z) -E z∼N (0,I k×k ) v φ(Az).

The stationary points of the above objective satisfy:

∇ v f 1 (A, v) = E z∼N (0,I k 0 ×k 0 ) φ(A * z) -E z∼N (0,I k×k ) φ(Az) = 0, ∇ aj f 1 (A, v) = -E z∼N (0,I k×k ) v j φ (a j z)z = 0.

We focus on the gradient over v. To achieve ∇ v f 1 (A, v) = 0, the stationary point satisfies:

EQUATION

To recap, for activations φ that follow Assumption 1, in both cases we have written the necessary condition on stationary point to be Eqn. ( 7), where φ is defined differently for odd or non-odd activations, but in both cases it is positive and monotone increasing on its support [0, ∞). We then argue the only solution for Eqn. ( 7) satisfies a j = a * j , ∀j. This follows directly from the following claim:

Claim 3. The function h(α) := E x∼N (0,α 2 ) f (x), α > 0 is a monotone increasing function if f is positive and monotone increasing on its support [0, ∞).

We could see from Claim 3 that the LHS and RHS of Eqn. ( 7) is simply h( a j ) and h( a * j ) for each j. Now that h is an monotone increasing function, the unique solution for h( a j ) = h( a * j ) is to match the norm: a j = a * j , ∀j.

Proof of Claim 3.

h(α) = E x∼N (0,α 2 ) f (x) = ∞ 0 f (x)e -x 2 2α 2 dx y:=x/α = ∞ 0 αf (αy)e -y 2 2 dy = E y∼N (0,1) αf (αy). Notice h (α) = E x∼N (0,1) [αxf (αx) + f (αx)].

Since f , f , and α > 0, and we only care about the support of f where x is also positive, therefore h is always positive and h is monotone increasing.

To sum up, at stationary point where ∇f 1 (A, v) = 0, we have ∀i, a * i = a i .

A.2 PROOF OF THEOREM 1:

Proof of Theorem 1. We will take optimal gradient ascent steps with learning rate 1 on the discriminator side v, hence the function we will actually be optimizing over becomes (using the notation for φ from section A.1):

h(A) = max v f 1 (A, v) = 1 2 E z∼N (0,I k 0 ×k 0 ) φ(A * z) -E z∼N (0,I k×k ) φ(Az) 2 .

We just want to verify that there's no spurious local minimum for h(A). Notice there's no interaction between each row vector of A.

Therefore we instead look at each

h i := 1 2 E z∼N (0,I k 0 ×k 0 ) φ((a * i ) z) -E z∼N (0,I k×k ) φ(a i z) 2 for each i. Now ∇h i (a i ) = -E z∼N (0,I k 0 ×k 0 ) φ((a * i ) z) -E z∼N (0,I k×k ) φ(a i z) (E z∼N (0,I k×k ) z φ (a i z)).

Due to the symmetry of the Gaussian, we take a i = ae 1 , where a = a i . It is easy to see that checking whether E z∼N (0,I k×k ) z φ (a i z) = 0 is equivalent to checking whether E z1∼N (0,1) z 1 φ (az 1 ) = 0.

Recall that φ is supported on [0, +∞) and it is monotonically increasing on its support. Hence, E z1∼N (0,1) z 1 φ (az 1 ) = 0 unless a = 0. Hence, suppose a i = 0, ∀i.

Then ∇ A h(A) = 0 iff h(A) = 0, i.e. E z∼N (0,I k 0 ×k 0 ) φ(A * z) = E z∼N (0,I k×k ) φ(Az).

Therefore all stationary points of h(A) are global minima where E z∼N (0,I k 0 ×k 0 ) φ(A * z) = E z∼N (0,I k×k ) φ(Az) and according to Lemma 1, this only happens when a i = a * i , ∀i ∈ [d].

A.3 STATIONARY POINTS FOR WGAN WITH QUADRATIC DISCRIMINATOR:

Proof of Lemma 2. To study the stationary point for g(Z) = jk gjk (z jk ), we look at individual

gjk (z) ≡ 1 2 ( ∞ i=0 σ 2 i ((z * jk ) i -z i ))

2 . Notice for odd-plus-constant activations, σ i is zero for even i > 0. Recall our assumption in Lemma 2 also requires that σ 1 = 0. Since the analysis is invariance to the position of the matrix Z, we simplify the notation here and essentially want to study the stationary point for f (a) = 1 2 ( i odd σ 2 i (a i -b i )) 2 for some constant b and σ i , where σ 1 = 04 .

f (a) = i odd σ 2 i (a i -b i ) i odd iσ 2 i a i-1 = (a -b)   σ 2 1 + i≥3 odd σ 2 i a i -b i a -b     σ 2 1 + i≥3 odd iσ 2 i a i-1   = (a -b)(I)(II).

Notice now f (a) = 0 ⇔ a = b. This is because the polynomial f (a) is factorized to a -b and two factors I and II that are always positive. Notice here we use a i -b i a-b to denote i j=0 a j b i-j , which is always nonnegative. This is simply because a i -b i always shares the same sign as a -b when i is odd. Therefore I=σ

2 1 + i≥3 odd σ 2 i a i -b i a-b > 0, ∀a.

Meanwhile, since a i-1 is always nonnegative for each odd i, we have II= σ 2 1 + i≥3 odd iσ 2 i a i-1 is also always positive for any a.

Next, for activation like ReLU, loss gjk (Daniely et al., 2016) . Therefore h (-1) = 0 for any z * jk . This fact prevents us from getting the same conclusion for ReLU.

(z) = 1 2 (h(z) -h(z * jk )) 2 , where h(x) = 1 π ( √ 1 -x 2 + (π - cos -1 (x))x)

However, for leaky ReLU with coefficient of leakage α ∈ (0, 1), φ(x) = max{x, αx} = (1α)ReLU(x) + αx.

We have

E z∼N (0,I k×k φ(a i z)φ(a j z) =(1 -α) 2 E z ReLU(a i z)ReLU(a j z) + (1 -α)αE z ReLU(a i z)a j z + (1 -α)αE z a i zReLU(a j z) + α 2 E z a i za j z =(1 -α) 2 h(a i a j ) + αa i a j Therefore for leaky ReLU gjk (z) = 1 2 ((1 -α) 2 (h(z) -h(z jk * )) + α(z -z * jk )) 2 , and g jk (z) = ((1-α) 2 (h(z)-h(z jk * ))+α(z -z * jk ))((1-α) 2 h (z)+α)

. Now with α > 0, (1-α) 2 h (z)+α ≥ α for all z and gjk (z) = 0 ⇔ z = z * jk . To sum up, for odd activations and leaky ReLU, since each gjk (z) only has stationary point of z = z * jk , the stationary point Z of g(Z) = jk gjk also satisfy Z = Z * = A * (A * ) .

Proof of Theorem 2. Instead of directly looking at the second-order stationary point of Problem 1, we look at the following problem on its reparametrized version:

Problem 2. min Z    g(Z) = 1 2 ∞ i=0 σ 2 i (Z * ) •i -Z •i 2 F    s.t. z ii = 1, ∀i. Z 0.

Here Z * = A * (A * ) and satisfies z * ii = 1, ∀i.

Compared to function g in the original problem 1, it satisfies that g(AA ) ≡ g(A).

A matrix Z satisfies the first-order stationary point for Problem 2 if there exists a vector σ such that:

         z ii = 1, Z 0, S 0, SZ = 0, S = ∇ Z g(Z) -diag(σ).

Proof. Denote the normalized variables x = x/α, ŷ = y/β. Let l = min{m, n}.

E[h m (x)h n (y)] =E[h α m (x)h β n (ŷ)] = m 2 i=0 n 2 j=0 η m,i α η n,j β E[h m-2i (x)h n-2j (ŷ)] = m 2 i=0 n 2 j=0 η m,i α η n,j β δ (m-2i),(n-2j) ρ n-2j (Lemma ??) = l 2 i=0 η l,i α η l,i β ρ l-2i if m ≡ n (mod 2) 0 o.w.

.

Now the population risk becomes:

g(A) = 1 2 E x∼D xx -E z∼N (0,I k×k ) φ(Az)φ(Az) 2 = 1 2 i,j∈[d] E z∼N (0,I k 0 ×k 0 ) φ((a * i ) z)φ((a * j ) z) -E z∼N (0,I k×k ) φ(a i z)φ(a j z) 2 ≡ 1 2 i,j gij (z ij ).

To simplify the notation, for a specific i, j pair, we write x = a i z/α, α = a i and ŷ = a j z/β, where β = a j . Namely we have (

x, ŷ) ∼ N (0, [[1, ρ]; [ρ, 1]]), where ρ = cos a i , a j . Again, recall φ(αx) = k odd σ i h i (αx) = k odd σ i h α i (x). E z∼N (0,I k×k ) [φ(αx)φ(β ŷ)] =E m odd σ m h α m (x) n odd σ n h β n (ŷ) = m,n odd σ m σ n E S [h α m (x)h β n (ŷ)] = m odd σ m n≤m odd σ n n 2 k=0 η n,k α η n,k β ρ n-2k

Therefore we could write out explicitly the coefficient for each term ρ k , k odd, as:

c k = n≥k odd σ n η n, n-k 2 α η n, n-k 2 β ( m≥n σ m ). We have gij (z ij ) = ( k odd c k z k ij -k odd c k (z * ij ) k ) 2

. Now suppose σ i to have the same sign, and

α i ≥ 1, ∀ or α i ≤ 1, ∀i, each coefficient c i ≥ 0.

Therefore still the only stationary point for g(Z) is Z * .

B OMITTED PROOFS FOR SAMPLE COMPLEXITY B.1 OMITTED PROOFS FOR RELATION ON APPROXIMATE STATIONARY POINTS:

Proof of Lemma 6. We first review what we want to prove. For a matrix A that satisfies -approximate SOSP for Eqn. (5), we define

EQUATION

We just want to show Z := AA , σ := λ, and S := S A satisfies the conditions for -FOSP of Eqn. (6). Therefore, by going over the conditions, its easy to tell that all other conditions automatically apply and it remains to show S A -I.

By noting that ∇:

EQUATION

(From Eqn. ( 9) we have Tr(B XiA) = 0)

Notice that A ∈ R d×k and we have chosen k = d for simplicity. We first argue when A is rankdeficient, i.e. rank(A) < k. There exists some vector v ∈ R k such that Av = 0. Now for any vector b ∈ R d , let B = bv . Therefore AB = Avb = 0. From (10) we further have:

1 2 Tr(B D A ∇ A L(A, λ)[B]) = Tr(B S A B) + Tr(AB D A ∇ Z g(AA )[B]) = Tr(vb S A bv ) = v 2 b S A b ≥ -/2 B 2 F (from (9)) = -/2 v 2 b 2

Therefore from the last three rows we have b S A b ≥ -/2 b 2 for any b, i.e. S A -/2I d×d . On the other hand, when A is full rank, the column space of A is the entire R d vector space, and therefore S A -I d×d directly follows from the second line of the -SOSP definition.

B.2 DETAILED CALCULATIONS:

Recall the population risk

EQUATION

F . Write the empirical risk on observations as: 

g n (A) ≡ 1 2 1 n n i=1 x i x i -E z∼N (0,I k×k ) φ(Az)φ(Az) 2 F . Claim 4. ∇g(A) -∇g n (A) = 2E z diag(φ (Az))(X -X n )φ(Az)z , where X = E x∼D [xx ], and X n = 1 n n i=1 x i x i . Proof. ∇g(A) -∇g n (A) = ∇(g(A) -g n (A)) = 1 2 ∇ X -X n , X + X n -2E z∼N (0,I k×k ) φ(Az)φ(Az) =∇ X n -X, E z∼N (0,I k×k ) φ(Az)φ(Az) Now write S(A) = φ(Az)φ(Az) . [S(A + ∆A) -S(A)] ij =φ(a i z + ∆a i z)φ(a j z + ∆a j z) -φ(a i z)φ(a j z) =φ (a i z)∆a i zφ(a j z) + φ (a j z)∆a j zφ(a i z) + O( ∆A 2 )

B.3 OMITTED PROOFS FOR OBSERVATION SAMPLE COMPLEXITY:

Proof of Lemma 3. For each x i = φ(Az i ), z i ∼ N (0, I k×k ). Each coordinate |x i,j | = |φ(a j z i )| ≤ |a j z i | since φ is 1-Lipschitz. 6 . Without loss of generality we assumed a j = 1, ∀j, therefore a j z ∼ N (0, I k×k ). For all i ∈ [n], j ∈ [d] |x i,j | ≤ log(nd/δ) with probability 1 -δ.

Then by matrix concentration inequality ((Vershynin, 2010) Corollary 5.52), we have with probability 1-δ: (1-)X X n (1+ )X if n ≥ Ω(d/ 2 log 2 (nd/δ)). Therefore set n = Θ(d/ 2 log 2 (1/δ)) will suffice.

Proof of Lemma 4.

X ij = E z∼N (0,I k×k ) φ(a i z)φ(a j z)

= 0 i = j E[φ 2 (a i z)] ≤ 2 π i = j Therefore X 2 ≤ 2

π . Together with Lemma 3, X -X n ≤ 2 π w.p 1 -δ. Recall ∇g(A) -∇g n (A) = 2E z diag(φ (Az))(X -X n )φ(Az)z := 2E z G(z), where G(z) is defined as diag(φ (Az))(X -X n )φ(Az)z . We have G(z) ≤ A z 2 X -X n .

∇g(A) -∇g n (A)

2 = 2 E z [G(z)] ≤ 2E z G(z) ≤ 2E z A z 2 X -X n ≤ 2 A 2 π E z z 2 = 2 A d 2 π

For the directional derivative, we make the concentration bound in a similar way. Denote D(z) = diag(φ (Az))(X n -X)φ (Az) • (Bz)z + diag(φ (Az) • (Bz))(X n -X)φ(Az)z . On the other hand, our target function is:

g n (A) ≡ 1 2 1 n n i=1

x i x i -E z∼N (0,I k×k ) φ(Az)φ(Az)

2 F = 1 2 E S [S] 2 F -E S [S], X n + 1 2 X n 2 F

Therefore E S gm,n (A) -g n (A) = 1 2m (E S S(A) 2 F -E S S(A) 2 F ). Claim 6. 

:

We will replace v by V ∈ R d×d when necessary.

Θ hides log factors of d for simplicity.

The exact error bound comes from the fact that when diagonal terms of AA are fixed, A 2 = O( √ d).

The zero component has been cancelled out.

Throughout the analysis for low rank optimization inJournée et al. (2008), they require function g(Z) to be convex. However, by carefully scrutinizing the proof, one could see that this condition is not required in building the connection of first-order and second-order stationary points of g(A) and g(Z). For more cautious readers, we also show a relaxed version in the next section, where the equivalence of SOSP of g and FOSP of g is a special case of it.

For simplicity, we analyze as if φ(0) = 0 w.o.l.g. throughout this section, since the bias term is canceled out in the observation side with φ(A * z) and the learning side with φ(Az).

Figures and Tables:

Figure 1: Recovery error ( AA -Z * F ) with different observed sample sizes n and output dimension d.

Figure 2: Comparisons of different performance with leakyReLU and tanh activations. Same color starts from the same starting point. For both cases, parameters always converge to true covariance matrix. Each arrow indicates the progress of 500 iteration steps.

∆A) -S(A)] i: =φ (a i z)∆a i zφ(Az) + (φ (Az) • ∆Az) φ(a i z) + O( ∆A 2 ) Therefore S(A + ∆A) -S(A) = diag(φ (Az))∆Azφ(Az) + φ(Az)z ∆A diag(φ (Az)).(11)Andg(A + ∆A) -g n (A + ∆A) -(g(A) -g n (A)) = X n -X, E z [S(A + ∆A) -S(A)] =E z X n -X, diag(φ (Az))∆Azφ(Az) + φ(Az)z ∆A diag(φ (Az)) =2E z diag(φ (Az))(X n -X)φ(Az)z , ∆A .

we have ∇g(A) -∇g n (A) = 2E z diag(φ (Az))(X n -X)φ(Az)z .Claim 5. For arbitrary matrix B, the directional derivative of ∇g(A) -∇g n (A) with direction B is:D A ∇g(A)[B] -D A ∇g n (A)[B] = 2E z diag(φ (Az))(X n -X)φ (Az) • (Bz)z +2E z diag(φ (Az) • (Bz))(X n -X)φ(Az)z Proof. g(A + tB) =2E z diag(φ (Az + tBz))(X n -X)φ(Az + tBz)z =2E z diag(φ (Az) + t(Bz) • φ (Az))(X n -X)(φ(Az) + tφ (Az) • (Bz))z + O(t 2 ) Therefore lim t→0 g(A + tB) -g(A) t =2E z diag(φ (Az))(X n -X)φ (Az) • (B z)z + 2E z diag(φ (Az) • (Bz))(X n -X)φ(Az)z

(z) ≤ X n -X 2 B z 2 (1 + z A ). Therefore D A ∇g(A)[B] -D A ∇g n (A)[B] ≤ O( d 3/2 A B ) with probability 1 -δ.B.4 OMITTED PROOFS ON BOUNDING MINI-BATCH SIZERecall gmWrite S j (A) ≡ φ(Az j )φ(Az j ) . Then we have

∇E S gm,n (A)-∇g n (A) = 2 m E z diag(φ (Az))S(A)φ(Az)zdiag(φ (Az))E S [S(A)]φ(Az)z . Proof. ∇E S gm,n -∇g n , ∆A =E S gm,n (A + ∆A) + g n (A + ∆A) -(E S gm,n (A) + g n (A)) + O( ∆A 2 ) = 1 2m E S S(A + ∆A) 2 F -E S S(A) 2 F -E S S(A + ∆A) 2 F + E S S(A) 2 F + O( ∆A 2 ) = 1 m (E S S(A), S(A + ∆A) -S(A) -E S [S(A)], E S [S(A + ∆A) -S(A)]) + O( ∆A 2 ) = 1 m E z S(A), diag(φ (Az))∆Azφ(Az) -E S [S(A)], E z diag(φ(Az))∆Azφ(Az) + O( ∆A 2 ) (from Eqn. (11) and symmetry of S) = 2 m E z diag(φ (Az))S(A)φ(Az)zdiag(φ (Az))E S [S(A)]φ(Az)z , ∆A + O( ∆A 2 )

Peer review (iclr2020):
I have read the authors response. In the response the authors clarified the contributions of this paper. I agree with the authors that the analysis of gradient descent-ascent is a difficult problem, and the optimization results given in this paper is a contribution of importance. Because of this I have improved my score. However, I do not agree with the authors that studying quadratic discriminators instead of more complicated ones should be considered as a contribution instead of drawback. In my opinion, as long as the focus is on WGAN, results involving standard neural networks are still more desired compared with the results in this submission. For example, similar results for a neural network discriminator might be even more impactful, because the optimization problem is even more difficult. Therefore I still consider the simple discriminator and generator as a weak point of this paper. ====================================================================================================== This paper studies the training of WGANs with stochastic gradient descent. The authors show that for one-layer generator network and quadratic discriminator, if the target distribution is modeled by a teacher network same as the generator, then stochastic gradient descent-ascent can learn this target distribution in polynomial time. The authors also provide sample complexity results. The paper is well-written and the theoretical analysis seems to be valid and complete. However, I think the WGANs studied in this paper are simplified too much that the analysis can no longer capture the true nature of WGAN training. First, the paper only studies linear and quadratic discriminators. This is not very consistent with the original intuition of WGAN, which is to use the worst Lipschitz continuous neural network to approximate the worst function in the set of all Lipschitz continuous functions in the definition of Wasserstein distance. When the discriminator is as simple as linear or quadratic functions, there is pretty much no “Wasserstein” in the optimization problem. Moreover, the claim that SGD learns one-layer networks can be very misleading. In fact what is a “one-layer” neural network? - if the authors meant “two-layer network” or “single hidden layer network”, then this is not true. Because as far as I can tell, the model is much more difficult than the model . The former is a standard single hidden layer network which is non-convex, while the latter is essentially a linear model especially when phi is known. - if the authors meant “a linear model with elementwise monotonic transform”, then I would like to suggest that a more appropriate name should be used to avoid unnecessary confusion. As previously mentioned, the discriminators are too simple to approximate the Wasserstein distance, and therefore in general it should not be possible to guarantee recovery of the true data distribution. However, in this paper it is still shown that certain true distributions can be learned. This is due to the extremely simplified true model. In fact, even if the activation function is unknown, it seems that one can still learn well (for example, by Kendall’s tau).