Title: A MEAN-FIELD THEORY FOR KERNEL ALIGNMENT WITH RANDOM FEATURES IN GENERATIVE ADVERSE-RIAL NETWORKS

Abstract: We propose a novel supervised learning method to optimize the kernel in maximum mean discrepancy generative adversarial networks (MMD GANs). Specifically, we characterize a distributionally robust optimization problem to compute a good distribution for the random feature model of Rahimi and Recht to approximate a good kernel function. Due to the fact that the distributional optimization is infinite dimensional, we consider a Monte-Carlo sample average approximation (SAA) to obtain a more tractable finite dimensional optimization problem. We subsequently leverage a particle stochastic gradient descent (SGD) method to solve the derived finite dimensional optimization problem. Based on a mean-field analysis, we then prove that the empirical distribution of the interactive particles system at each iteration of the SGD follows the path of the gradient descent flow on the Wasserstein manifold. We also establish the non-asymptotic consistency of the finite sample estimator. Our empirical evaluation on synthetic data-set as well as MNIST and CIFAR-10 benchmark data-sets indicates that our proposed MMD GAN model with kernel learning indeed attains higher inception scores well as Frèchet inception distances and generates better images compared to the generative moment matching network (GMMN) and MMD GAN with untrained kernels.

1. INTRODUCTION:

A fundamental and long-standing problem in unsupervised learning systems is to capture the underlying distribution of data. While deep generative models such as Boltzmann machines Salakhutdinov & Hinton (2009) and auto-encoding variational Bayes Kingma & Welling (2013) accomplish this task to some extent, they are inadequate for many intractable probabilistic computations that arise in maximum likelihood estimation. Moreover, in many machine learning tasks such as caption generation Xu et al. (2015) , the main objective is to obtain new samples rather than to accurately estimate the underlying data distribution. Generative adverserial network (GAN) Goodfellow et al. (2014) provides a framework to directly draw new samples without estimating data distribution. It consists of a deep feedforward network to generate new samples from a base distribution (e.g. Gaussian distribution), and a discriminator network to accept or reject the generated samples. However, training GAN requires finding a Nash equilibrium of a non-convex minimax game with continuous, high-dimensional parameters. Consequently, it is highly unstable and prone to miss modes Salimans et al. (2016) ; Che et al. (2016) . To obtain more stable models, the generative moment matching networks (GMMNs) Li et al. (2015) are proposed, wherein instead of training a discriminator network, a non-parametric statistical hypothesis test is performed to accept or reject the generated samples via the computation of the kernel maximum mean discrepancy Gretton et al. (2007) . While leveraging a statistical test simplifies the loss function for training GMMN, in practice, the diversity of generated samples by GMMN is highly sensitive to the choice of the kernel. Thus, to improve the sampling performance, the kernel function also needs to be jointly optimized with the generator. Rather than i optimizing the kernel directly, the MMD GAN model Li et al. (2017) is proposed in which an embedding function is optimized in conjunction with a fixed user-defined kernel (e.g. RBF Gaussian kernel). However, there are no theoretical guarantees that the user-defined kernel is the 'right' kernel for embedded features.

Contributions.:

To address the kernel model selection problem in MMD GAN Li et al. (2017) , in this paper we put forth a novel framework to learn a good kernel function from training data. Our kernel learning approach is based on a distributional optimization problem to learn a good distribution for the random feature model of Rahimi and Recht Rahimi & Recht (2008; 2009) to approximate the kernel. Since optimization with respect to the distribution of random features is infinite dimensional, we consider a Monte Carlo approximation to obtain a more tractable finite dimensional optimization problem with respect to the samples of the distribution. We then use a particle stochastic gradient descent (SGD) to solve the approximated finite dimensional optimization problem. We provide a theoretical guarantee for the consistency of the finite sample-average approximations. Based on a mean-field analysis, we also show the consistency of the proposed particle SGD. In particular, we show that when the number of particles tends to infinity, the empirical distribution of the particles in SGD follows the path of the gradient descent flow of the distributional optimization problem on the Wasserstein manifold.

2. PRELIMINARIES OF MMD GANS:

Assume we are given data {v i } n i=1 that are sampled from an unknown distribution P V with the support V. In many unsupervised tasks, we wish to attain new samples from the distribution P V without directly estimating it. Generative Adversarial Network (GAN) Goodfellow et al. (2014) provides such a framework. In vanilla GAN, a deep network G(•; W ) parameterized by W ∈ W is trained as a generator to transform the samples Z ∼ P Z , Z ∈ Z from a user-defined distribution P Z (e.g. Gaussian) into a new sample G(Z; W ) ∼ P W , such that the distributions P W and P V are close. In addition, a discriminator network D(•; δ) parameterized by δ ∈ ∆ is also trained to reject or accept the generated samples as a realization of the data distribution. The training of the generator and discriminator networks is then accomplished via solving a minimax optimization problem as below (1)

In the high dimensional settings, the generator trained via the min-max program of equation 1 can potentially collapse to a single mode of distribution where it always emits the same point Che et al. (2016) . To overcome this shortcoming, other adversarial generative models are proposed in the literature, which propose to modify or replace the discriminator network by a statistical two-sample test based on the notion of the maximum mean discrepancy which is defined below:

Definition 2.1. (MAXIMUM MEAN DISCREPANCY GRETTON ET AL. ( 2007)) Let (X , d) be a metric space, F be a class of functions f : X → IR, and P, Q ∈ B(X ) be two probability measures from the set of all Borel probability measures B(X ) on X . The maximum mean discrepancy (MMD) between the distributions P and Q with respect to the function class F is defined below

EQUATION

Different choices of the function class F in equation 2 yield different adversarial models such as Wasserstein GANs (WGAN) Arjovsky et al. (2017) , f -GANs Nowozin et al. (2016) , and GMMN and MMD GAN Li et al. (2017; 2015) . In the latter two cases, the function class F def = {f : X → IR : f H X ≤ 1}, where H X is a reproducing kernel Hilbert space (RKHS) of functions with a kernel K : X × X → IR, denoted by (H X , K). Then, the squared MMD loss in equation 2 as a measure of the distance between the distributions P V and P W has the following expression

EQUATION

where X = V ∪ W. Instead of training the generator via solving the minimax optimization in equation 1, the MMD GAN model of Li et al. (2015) proposes to optimize the discrepancy between two distributions via optimization of an embedding function ι : IR D → IR p , p ≤ D, i.e., min

EQUATION

where k : IR p × IR p → IR is a user-defined fixed kernel. In Li et al. (2015) , the proposal for the kernel k : IR p × IR p → IR is a mixture of the Gaussians,

EQUATION

where the bandwidth parameters σ 1 , • • • , σ m > 0 are manually selected. Nevertheless, in practice there is no guarantee that the user-defined kernel k(ι(x), q(y)) can capture the structure of the embedded features ι(x).

3. PROPOSED APPROACH: KERNEL LEARNING WITH RANDOM FEATURES FOR MMD GANS:

In this section, we first expound our kernel learning approach. Then, we describe a novel MMD GAN model based on the proposed kernel learning approach.

3.1. ROBUST DISTRIBUTIONAL OPTIMIZATION FOR KERNEL LEARNING:

To address the kernel model selection issue in MMD GAN Li et al. (2017) , we consider a kernel optimization scheme with random features Rahimi & Recht (2008; 2009) . Let ϕ : IR d × IR D → [-1, 1] denotes the explicit feature maps and µ ∈ M(IR D ) denotes a probability measure from the space of probability measures M(IR D ) on IR D . The kernel function is characterized via the explicit feature maps using the following integral equation 

Here, P is the set of probability distributions corresponding to a kernel class K. In the sequel, we consider P to be the distribution ball of radius R as below

EQUATION

where µ 0 is a user-defined base distribution, and To establish the proof, we consider the W p (•, •) is the p-Wasserstein distance defined as below

EQUATION

)

where the infimum is taken with respect to all couplings π of the measures µ, µ 0 ∈ M(IR D ), and Π(µ, µ 0 ) is the set of all such couplings with the marginals µ and µ 0 .

The kernel MMD loss function in equation 8 is defined with respect to the unknown distributions of the data-set P V and the model P W . Therefore, we construct an unbiased estimator for the MMD loss function in equation 8 based on the training samples. To describe the estimator, sample the labels from a uniform distribution y 1 , • • • , y n ∼ i.i.d. Uniform{-1, +1}, where we assume that the number of positive and negative labels are balanced. In particular, consider the set of positive labels

I = {i ∈ {1, 2, • • • , n} : y i = +1}, and negative labels J = {1, 2, • • • , n}/I, where their cardinality is |I| = |J | = n 2 .

We consider the following assignment of labels: • Positive class labels: If y i = +1, sample the corresponding feature map from data-distribution

x i = v i ∼ P V . • Negative class labels: If y i = -1, sample from the corresponding feature map from the generated distribution x i = G(Z i , W ) ∼ P W , Z i ∼ P Z .

By this construction, the joint distribution of features and labels P Y,X has the marginals P X|Y =+1 = P V , and P X|Y =-1 = P W . Moreover, the following statistic, known as the kernel alignment in the literature (see, e.g., Sinha & Duchi (2016) ; Cortes et al. (2012) ), is an unbiased estimator of the MMD loss in equation 8,

EQUATION

See Appendix C.1 for the related proof. The kernel alignment in equation 11 can also be viewed through the lens of the risk minimization

min W ∈W inf µ∈P MMD α µ P V , P W def = 8 n(n -1)α 1≤i<j≤n (αy i y j -K µ (x i , x j )) 2 (12a) = 8 n(n -1)α 1≤i<j≤n (αy i y j -IE µ [ϕ(x i ; ξ)ϕ(x j ; ξ)]) 2 . (12b)

Here, α > 0 is a scaling factor that determines the separation between feature vectors, and K * def = αyy T is the ideal kernel that provides the maximal separation between the feature vectors over the training data-set, i.e., K * (x i , x j ) = α when features have identical labels y i = y j , and K * (x i , x j ) = -α otherwise. Upon expansion of the risk function in equation 12, it can be easily shown that it reduces to the kernel alignment in equation 11 when α → +∞. Intuitively, the risk minimization in equation 12 gives a feature space in which pairwise distances are similar to those in the output space Y = {-1, +1}.

3.2. SAA FOR DISTRIBUTIONAL OPTIMIZATION:

The distributional optimization problem in equation 8 is infinite dimensional, and thus cannot be solved directly. To obtain a tractable optimization problem, instead of optimizing with respect to the distribution µ of random features, we optimize the i.i.d. samples (particles)

ξ 1 , • • • , ξ N ∼ i.i.d. µ

generated from the distribution. The empirical distribution of these particles is accordingly defined as follows

EQUATION

where δ(•) is the Dirac's delta function concentrated at zero. In practice, the optimization problem in equation 12 is solved via the Monte-Carlo sample average approximation of the objective function,

EQUATION

where

P N def = IB N R ( µ N 0 ) = µ N ∈ M(IR D ) : W p ( µ N , µ N 0 ) ≤ R , and µ N 0 is the empirical measure associated with the initial samples ξ 1 0 , • • • , ξ N 0 ∼ i.i.d. µ 0 .

The empirical objective function in equation 14 can be optimized with respect to the samples ξ1 , • • • , ξ N using the particle stochastic gradient descent. For the optimization problem in equation 14, the (projected) stochastic gradient descent (SGD) takes the following recursive form, 1

EQUATION

for k = 1, 2, • • • , N , where (y m , x m ), ( y m , x m ) ∼ i.i.d P x,y and η ∈ IR >0 denotes the learning rate of the algorithm, and the initial particles are

ξ 1 0 , • • • , ξ N 0 ∼ i.i.d. µ 0 .

At each iteration of the SGD dynamic in equation 15, a feasible solution for the inner optimization of the empirical risk function in equation 14 is generated via the empirical measure

EQUATION

Indeed, we prove in Section 4 that for an appropriate choice of the learning rate η > 0, the empirical measure in equation 16 remains inside the distribution ball µ N m ∈ P N for all m ∈ [0, N T ] ∩ IN, and is thus a feasible solution for the empirical risk minimization equation 14 (see Corollary 4.2.1 in Section 4).

3.3. PROPOSED MMD GAN WITH KERNEL LEARNING:

In Algorithm 1, we describe the proposed method MMD GAN model with the kernel learning approach described earlier. Algorithm 1 has an inner loop for the kernel training and an outer loop for training the generator, where we employ RMSprop Tieleman & Hinton (2012) . Our proposed MMD GAN model is distinguished from MMD GAN of Li et al. (2017) in that we learn a good kernel function in equation 17 of the inner loop instead of optimizing the embedding function that is implemented by an auto-encoder. However, we mention that our kernel learning approach is compatible with the auto-encoder implementation of Li et al. (2017) for the dimensionality reduction of features (and particles). In the case of including an auto-encoder, the inner loop in Algorithm 1 must be modified to add an additional step for training the auto-encoder. However, to convey the main ideas more clearly, the training step of the auto-encoder is omitted from Algorithm 1.

3.4. COMPUTATIONAL COMPLEXITY ANALYSIS:

Sampling the labels y, y ∼ i.i.d. Uniform{-1, +1} require O(1) complexity, while sampling x| y = +1, x|y = +1 ∼ P V and x| y = +1, x|y = -1 ∼ P W has a complexity of O(d). The computation of the stochastic gradient step in equation 17 requires computing the sum

1 N N i=1 ϕ(x; ξ k )ϕ( x; ξ k ).

This can be done as a separate preprocessing step prior to executing the SGD in equation 17, and requires preparing the vectors

ϕ def = [ϕ(x; ξ 1 ), • • • , ϕ(x; ξ N )] and ϕ def = [ϕ( x; ξ 1 ), • • • , ϕ( x; ξ N )].

Using the random feature model of Rahimi and Recht Rahimi & Recht (2008; 2009) , where ϕ(x; ξ) = √ 2 cos(x T ξ + b). Here b ∼ Uniform[-π, +π], the complexity of computing the vectors ϕ and ϕ is of the order O(N d) on a single processor. However, this construction is trivially parallelizable. Furthermore, computation can be sped up even further for certain distributions µ 0 . For example, the Fastfood technique can approximate ϕ and φ in O(N log(d))

Algorithm 1 MMD GAN with a supervised kernel learning Method (Monte-Carlo Approach)

Inputs: The learning rates η, η > 0 , the number of iterations of discriminator per generator update T ∈ N, the batch-size n, the number of random features N ∈ N. Regularization parameter α > 0. while W has not converged do for t = 1, 2, • • • , T do Sample the labels y, y ∼i.i.d Uniform{-1, 1}.

Sample the features x|y = +1 ∼ P V , and x|y = -1 ∼ P W . Similarly, x| y = +1 ∼ P V , and x| y = -1 ∼ P W . For all k = 1, 2, • • • , N , update the particles,

EQUATION

end for Sample a balanced minibatch of labels {yi} n i=1 ∼i.i.d. Uniform{-1, +1}. Sample the minibatch {x} n i=1 such that xi|yi = +1 ∼ P V , and xi|yi = -

EQUATION

)

end while time for the Gaussian kernel Le et al. (2013) . Updating each particle in equation 17, involves the computation of the gradient

∇ ξ (ϕ(ξ; ξ k )ϕ( ξ; ξ k ))) which is O(d).

Thus, the complexity of one iteration of SGD for all the particles is O(N d).

Overall, one step of the kernel learning has a complexity of O(N d). On the other hand, to attain

ε-suboptimal solution max k=1,2,••• ,N ξ k -ξ k * ≤ ε, the SGD requires has the sample complexity O log( 1 ε )

. Consequently, the computational complexity of the kernel learning is of the order of O(N d log( 1 ε )). To compare this complexity with that of MMD GAN is of the order O B 2 log( 1 ε ) , where B is the batch size for approximation of the population MMD, and is the number of kernel mixtures.

4. CONSISTENCY AND A MEAN-FIELD ANALYSIS:

In this section, we provide theoretical guarantees for the consistency of various approximations we made to optimize the population MMD loss function in equation 8. We defer the proofs of the following theoretical results to AppendixC. The main assumptions ((A.1),(A.2), and (A.3)) underlying our theoretical results are also stated in the same section.

Consistency of finite-sample estimate::

In this part, we prove that the solution to finite sample optimization problem in equation 14 approaches its population optimum in equation 8 as the number of data points as well as the number of random feature samples tends to infinity. (19a)

EQUATION

respectively. Then, with the probability of (at least) 1 -3 over the training data samples {(x i , y i )} n i=1 and the random feature samples {ξ k 0 } N k=1 , the following non-asymptotic bound holds

MMD µ * [P V , P W * ] -MMD µ N * [P V , P W N * ] (20) ≤ L 2 (d + 2) N ln 1 2 2 8 N diam 2 (X ) + 2 max c 1 L 2 n ln 1 2 4 , c 2 RL 4 n 2 ln 4e L 4 9 + 8L 2 α ,

where c 1 = 3 1 4 × 2 4 , and c 2 = 9 × 2 11 .

The proof of Theorem 4.1 is presented in Appendix C.1.

Notice that there are three key parameters involved in the upper bound of Theorem 4.1. Namely, the number of training samples n, the number of random feature samples N , and the regularization parameter α. The upper bound in equation 20 thus shows that when n, N, α → +∞, the solution obtained from solving the empirical risk minimization in equation 12 yields a MMD population value tending to the optimal value of the distributional optimization in equation 8.

Consistency of particle SGD for solving distributional optimization. The consistency result of Theorem 4.1 is concerned with the MMD value of the optimal empirical measure µ N * (ξ) =

1 N N k=1 δ(ξ -ξ k * )

of the empirical risk minimization equation 14. In practice, the particle SGD is executed for a few iterations and its values are used as an estimate for

(ξ 1 * , • • • , ξ N * ).

Consequently, it is desirable to establish a similar consistency type result for the particle SGD estimates

(ξ 1 m , • • • , ξ N m )

at the m-th iteration. To reach this objective, we define the scaled empirical measure as follows

EQUATION

At any time t, the scaled empirical measure µ N t is a random element, and thus (µ N t ) 0≤t≤T is a measured-valued stochastic process. Therefore, we characterize the evolution of its Lebesgue density p N t (ξ) def = µ N t (dξ)/dξ in the following theorem: Theorem 4.2. (MCKEAN-VLASOV MEAN-FIELD PDE) Suppose conditions (A.1)-(A.3) of Appendix C are satisfied. Further, suppose that the Radon-Nikodyme derivative q 0 (ξ) = µ 0 (dξ)/dξ exists. Then, there exists a unique solution (p * t (ξ)) 0≤t≤T to the following non-linear partial differential equation

EQUATION

)

Moreover, the measure-valued process {(µ N t ) 0≤t≤T } N ∈IN defined in equation 21 converges (weakly) to the unique solution µ * t (ξ) = p * t (ξ)dξ as the number of particles tend to infinity N → ∞. 

∈ P N = { µ N ∈ M(IR D ) : W p ( µ N , µ N 0 ) ≤ R} for all m ∈ [0, N T ] ∩ IN,

with the probability of (at least) 1 -δ.

Let us make two remarks about the PDE in equation 22. First, the seminal works of Otto Otto (2001), and Jordan, et al. Jordan et al. (1998) establishes a deep connection between the McKean-Vlasov type PDEs specified in equation 22 and the gradient flow on the Wasserstein manifolds. More specifically, the PDE equation in equation 22 can be thought of as the minimization of the energy functional

EQUATION

using the following gradient flow dynamics

EQUATION

where

grad pt E(p t (ξ)) = ∇ ξ • (p t (ξ)∇ ξ R α (p t (ξ)))

is the Riemannian gradient of R α (µ t (ξ)) with respect to the metric of the Wasserstein manifold . This shows that when the number of particles in particle SGD equation 15 tends to infinity (N → +∞), their empirical distribution follows a gradient descent path for minimization of the population version (with respect to data samples) of the distributional risk optimization in equation 12. In this sense, the particle SGD is the 'consistent' approximation algorithm for solving the distributional optimization.

4.1. RELATED WORKS:

The mean-field description of SGD dynamics has been studied in several prior works for different information processing tasks. 

EQUATION

where µ 0 is a user defined base distribution, and h ∈ H is a function that transforms the base distribution µ 0 into a distribution µ that provides a better kernel. Therefore, the work of Li, et al Li et al. (2019) implicitly optimizes the distribution of random features via transforming a random variable with a function. In contrast, the proposed distributional optimization framework in this paper optimizes the distribution of random feature explicitly, via optimizing their empirical measures. Perhaps more importantly from a practical perspective is the fact that our kernel learning approach does not require the user-defined function class H. Moreover, our particle SGD method in equation 15 obviates tuning hyper-parameters related to the implicit kernel learning method such as the gradient penalty factor and the variance constraint factor (denoted by λ GP and λ h , respectively, in Algorithm 1 of Li et al. (2019) ).

5. EMPIRICAL EVALUATION:



5.1. SYNTHETIC DATA-SET:

Due to the space limitation, the experiments on the synthetic data are deferred to Appendix A.

5.2. PERFORMANCE ON BENCHMARK DATASETS:

We evaluate our kernel learning approach on large-scale benchmark data-sets. We train our MMD GAN model on two distinct types of datasets, namely on MNIST LeCun et al. (1998) and CIFAR-10 LeCun et al. (1998) , where the size of training instances are 60 × 10 3 and 50 × 10 3 , respectively. All the generated samples are from a fixed noise random vectors and are not singled out.

Implementation and hyper-parameters. We implement Algorithm 1 as well as MMD GAN Li et al. (2017) in Pytorch using NVIDIA Titan V100 32GB graphics processing units (GPUs). The source code of Algorithm 1 is built upon the code of Li et al. (2017) , and retains the auto-encoder implementation. In particular, we use a sequential training of the auto-encoder and kernel as explained in the Synthetic data in Section A of Supplementary. For a fair comparison, our hyper-parameters are adjusted as in Li et al. (2017) , i.e., the learning rate of 0.00005 is considered for RMSProp Tieleman & Hinton (2012) . Moreover, the batch-size for training the generator and auto-encoder is n = 64. The learning rate of particle SGD is tuned to η = 10. Random Feature Maps. To approximate the kernel, we use the the random feature model of Rahimi and Recht Rahimi & Recht (2008; 2009) , where ϕ(

x; ξ) = √ 2 cos(x T ξ + b). Here b ∼ Uniform{-π, +π} is a random bias term. Practical considerations. When data-samples {V i } ∈ IR d are high dimensional (as in CIFAR- 10), the particles ξ 1 , • • • , ξ N ∈ IR D , D = d in SGD equation 15 are also high-dimensional.

To reduce the dimensionality of the particles, we apply an auto-encoder architecture similar to Li et al. (2017) , and train our kernel on top of learned embedded features. More specifically, in our simulations, we train an auto-encoder where the dimensionality of the latent space is h = 10 for MNIST, and h = 128 (thus D = d = 128) for CIFAR-10. Therefore, the particles ξ 1 , • • • , ξ N in subsequent kernel training phase have the dimension of D = 10, and D = 128, respectively. Choice of the scaling parameter α. There is a trade-off in the choice of α. While for large values of α, the kernel is better able to separate data-samples from generated samples, in practice, it slows down the convergence of particle SGD. This is due to the fact that the coupling between the particle dynamics in equation 15 decrease as α increase. The scaling factor is set to be α = 1 in all the following experiments. Qualitative comparison. We now show that without the bandwidth tuning for Gaussian kernels and using the particle SGD to learn the kernel, we can attain better visual results on benchmark data-sets. In Figure 1 , we show the generated samples on CIFAR-10 and MNIST data-sets, using our Algorithm 1, MMD GAN Li et al. (2017) with a mixed and homogeneous Gaussian RBF kernels, and GMMN Li et al. (2015) . 2017) with a mixture RBF Gaussian kernel κ(x, y) = 5 k=1 κ σ k (x, y), where σ k ∈ {1, 2, 4, 8, 16} are the bandwidths of the Gaussian kernels that are fine tuned and optimized. We observe that our MMD GAN with automatic kernel learning visually attains similar results to MMD GAN Li et al. (2017) which requires manual tuning of the hyper-parameters. In Figure 1 (c), we show the MMD GAN result with a single kernel RBF Gaussian kernel whose bandwidth is manually tuned at σ = 16. Lastly, in Figure 1(d) , we show the samples from GMMN Li et al. (2015) which does not exploit an auto-encoder or kernel training. Clearly, GMMN yield a poor results compared to other methods due to high dimensionality of features, as well as the lack of an efficient method to train the kernel.

On MNIST data-set in Figure 1 ,(e)-(h), the difference between our method and MMD GAN Li et al. (2017) is visually more pronounced. We observe that without a manual tuning of the kernel bandwidth and by using the particle SGD equation 15 to optimize the kernel, we attain better generated images in Figure 1 (e), compared to MMD GAN with mixed RBF Gaussian kernel and manual bandwidth tuning in Figure 1 (f). Moreover, using a single RBF Gaussian kernel yields a poor result regardless of the choice of its bandwidth. The generated images from GMMN is also shown in Figure 1 Quantitivative comparison. To quantitatively measure the quality and diversity of generated samples, we compute the inception score (IS) Salimans et al. (2016) as well as Frèchet Inception Distance (FID) Heusel et al. (2017) on CIFAR-10 images. Intuitively, the inception score is used for GANs to measure samples quality and diversity. Accordingly, for generative models that are collapsed into a single mode of distribution, the inception score is relatively low. The FID improves on IS by actually comparing the statistics of generated samples to real samples, instead of evaluating generated samples independently.

(h). (a) (b) (c) (d) (e) (f) (g) (h)

In Table 1 , we report the quantitative measures for different MMD GAN model using different scoring metric. Note that in To reduce the dimensionality of data, we consider the embedding ι :

IR d → IR p , x → ι(x) = Σx,

where Σ ∈ IR p×d and p < d. In this case, the distribution of the embedded features are P X|Y =+1 = N(0, (1 + λ)ΣΣ T ), and

P X|Y =-1 = N(0, (1 -λ)ΣΣ T ).

Note that λ ∈ [0, 1] is a parameter that determines the separation of distributions. In particular, the Kullback-Leibler divergece of the two multi-variate Gaussian distributions is controlled by λ ∈ [0, 1],

EQUATION

In Figure 2 , we show the distributions of i.i.d. samples from the distributions P V and P W for different choices of variance parameter of λ = 0.1, λ = 0.5, and λ = 0.9. Notice that for larger λ the divergence is reduced and thus performing the two-sample test is more difficult. From Figure 2 , we clearly observe that for large values of λ, the data-points from the two distributions P V and P W have a large overlap and conducting a statistical test to distinguish between these two distributions is more challenging.

A.0.1 KERNEL LEARNING APPROACH:

Figure 4 depicts our two-phase kernel learning procedure which we also employed in our implementations of Algorithm 1 on benchmark data-sets of Section 5.2 in the main text. The kernel learning approach consists of training the auto-encoder and the kernel optimization sequentially, i.e., sup

EQUATION

)

where the function class is defined

Q def = {ι(z) = Σz, Σ ∈ IR D×d }, and (K µ N • ι)(x 1 , x 2 ) = K µ N (ι(x 1 ), ι(x 2 ))

. Now, we consider a two-phase optimization procedure:

• Phase (I): we fix the kernel function, and optimize the auto-encoder to compute a covariance matrix Σ for dimensionality reduction • Phase (II): we optimize the kernel based on the learned embedded features. This two-phase procedure significantly improves the computational complexity of SGD as it reduces the dimensionality of random feature samples ξ ∈ IR D , D d. When the kernel function K is fixed, optimizing the auto-encoder is equivalent to the kernel learning step of Li et al. (2017) . 

A.0.2 STATISTICAL HYPOTHESIS TESTING WITH THE KERNEL MMD:

EQUATION

used to distinguish between these hypotheses:

• Null hypothesis H 0 : P V = P W (thus λ = 0), • Alternative hypothesis H 1 : P V = P W (thus λ > 0).

To perform hypothesis testing via the kernel MMD, we require that H X is a universal RKHS, defined on a compact metric space X . Universality requires that the kernel K(•, •) be continuous and, H X be dense in C(X ). Under these conditions, the following theorem establishes that the kernel MMD is indeed a metric: Theorem A.1. (METRIZABLITY OF THE RKHS) Let F denotes a unit ball in a universal RKHS H X defined on a compact metric space X with the associated continuous kernel K(•, •). Then, the kernel MMD is a metric in the sense that MMD K [P V , P W ] = 0 if and only if

P V = P W . To design a test, let µ N m (ξ) = 1 N N k=1 δ(ξ -ξ k m )

denotes the solution of SGD in equation 15 for solving the optimization problem. Consider the following MMD estimator consisting of two U -statistics and an empirical function

EQUATION

)

Given the samples {V i } m i=1 and {W i } n i=1 , we design a test statistic as below

EQUATION

where τ ∈ IR is a threshold. Notice that the unbiased MMD estimator of equation 28 can be negative despite the fact that the population MMD is non-negative. Consequently, negative values for the statistical threshold τ equation 29 are admissible. In the following simulations, we only consider non-negative values for the threshold τ .

A Type I error is made when H 0 is rejected based on the observed samples, despite the null hypothesis having generated the data. Conversely, a Type II error occurs when H 0 is accepted despite the alternative hypothesis H 1 being true. The significance level of a test is an upper bound on the probability of a Type I error: this is a design parameter of the test which must be set in advance, and is used to determine the threshold to which we compare the test statistic. The power of a test is the probability of rejecting the null hypothesis H 0 when it is indeed incorrect. In particular,

EQUATION

)

In this sense, the statistical power controls the probability of making Type II errors.

A.0.3 EMPIRICAL RESULTS:

In Figure 3 , we show the evolution of the empirical measure µ N m (ξ) of SGD particles by plotting the 2D histogram of the particles

ξ 1 m , • • • , ξ N m ∈ IR D at different iterations of SGD (D = d).

Clearly, starting with a uniform distribution in 3(a), the empirical measure seemingly evolves into a Gaussian measure in Figure 3(d) . The evolution to a Gaussian distribution demonstrates that the RBF Gaussian kernel corresponding to a Gaussian distribution for the random features indeed provides a good kernel function for the underlying hypothesis test with Gaussian distributions.

In Figure 4 , we evaluate the power of the test for 100 trials of hypothesis test using the test statistics of equation 29. To obtain the result, we used an autoencoder to reduce the dimension from d = 100 to p = 50. Clearly, for the trained kernel in Panel (a) of Figure 4 , the threshold τ for which Power = 1 increases after learning the kernel via the two phase procedure described earlier. In comparison, in Panel (b), we observe that training an auto-encoder only with a fixed standard Gaussian kernel K(x, y) = e -x-y 2 2 attains lower thresholds compared to our two-phase procedure. In Panel (c), we demonstrate the case of a fixed Gaussian kernel without an auto-encoder. In this case, the threshold is significantly lower due to the large dimensionality of the data.

From Figure 4 , we also observe that interestingly, the phase transition in the statistical threshold τ is not sensitive to the parameter λ. This phenomenon can be justified by the fact that the kernel learning indeed improves the MMD value more significantly for smaller values of λ (i.e., more difficult hypothesis testing problems) than larger values of λ. See Figure 5 . The inception score for LSUN bedroom data set is 3.860 ± 0.0423 using a single Gaussian kernel with the bandwidth of σ = 16, and 4.064 ± 0.061 using our proposed method in Algorithm 1.

C PROOFS OF MAIN THEORETICAL RESULTS:

Before we delve into proofs, we state the main assumptions underlying our theoretical results:

Assumptions:

(A.1) The feature space X = V ∪ W ⊂ IR d is compact with a finite diameter diam(X ) < ∞, where V = support(P V ) and W = support(P W ) are the supports of the distributions P V and P W respectively. (A.

2) The feature maps are bounded and Lipchitz almost everywhere (a.e.) ξ ∈ IR D . In particular,

sup x∈X |ϕ(x; ξ)| < L 0 , sup x∈X ∇ ξ ϕ(x; ξ) 2 ≤ L 1 , and sup ξ∈IR D ∇ x ϕ(x; ξ) < L 2 . Let L def = max{L 0 , L 1 , L 2 } < +∞. (A.3) Let µ N 0 (ξ) def = 1 N N k=1 δ(ξ -ξ k 0 ) denotes the empirical measure for the initial particles ξ 1 0 , • • • , ξ N 0 .

We assume that µ N 0 (ξ) converges (weakly) to a deterministic measure µ 0 ∈ M(IR D ). Furthermore, we assume the limiting measure µ 0 is absolutely continuous w.r.t. Lebesgue measure and has a compact support support(µ 0 ) = Ξ ⊂ IR D .

Notation: We denote vectors by lower case bold letters, e.g.

x = (x 1 , • • • , x n ) ∈ IR n ,

and matrices by the upper case bold letters, e.g.,

M = [M ij ] ∈ IR n×m . The Frobenius norm of a matrix is denoted by M F = n i=1 m j=1 |M ij | 2 . Let IB r (x) def = {y ∈ IR d : y -x 2 ≤ r}

denote the Euclidean ball of radius r centered at x. For a given metric space X , Let C b (IR d ) denote the space of bounded and continuous functions on X equipped with the usual supremum norm

EQUATION

Further, C k b (X ) the space of all functions in C b (X ) whose partial derivatives up to order k are bounded and continuous, and C k c (X ) the space of functions whose partial derivatives up to order k are continuous with compact support.

We denote the class of the integrable functions f with f (t) ≥ 0 a.e., on 0 ≤ t ≤ T by L 1 + [0, T ]. Similarly, L ∞ + [0, T ] will denote the essentially bounded functions with f (t) ≥ 0 almost everywhere. For a given metric space X , we denote the Borel σ-algebra by B(X ). For a Borel set B ∈ B(X ), the measure value of the set B with respect to the measure is given by µ(B). The space of finite non-negative measures defined on X is denoted by M(X ). The Dirac measure with the unit mass at x ∈ X is denoted by δ(x). For any measure µ ∈ M(X ) and any bounded function f ∈ C b (X ), we define

EQUATION

The space M(X ) is equipped with the weak topology, i.e., a (random) sequence {µ N } N ∈IN converges weakly to a deterministic measure µ ∈ M(X ) if and only if µ N , f → µ, f for all f ∈ C b (X ). We denote the weak convergence by µ N t weakly → µ. Notice that when X is Polish, then M(X ) equipped with the weak topology is also Polish. 3 For a Polish space X , let D X ([0, T ]) denotes the Skorokhod space of the cádlág functions that take values in X defined on [0, T ]. We assume that D X ([0, T ]) is equipped with the Skorokhod's J 1 -topology Billingsley (2013) , which in that case D X ([0, T ]) is also a Polish space.

We use asymptotic notations throughout the paper. We use the standard asymptotic notation for sequences. If a n and b n are positive sequences, then

a n = O(b n ) means that lim sup n→∞ a n /b n < ∞, whereas a n = Ω(b n ) means that lim inf n→∞ a n /b n > 0. Furthermore, a n = O(b n ) implies a n = O(b n poly log(b n )). Moreover a n = o(b n ) means that lim n→∞ a n /b n = 0 and a n = ω(b n ) means that lim n→∞ a n /b n = ∞. Lastly, we have a n = Θ(b n ) if a n = O(b n ) and a n = Ω(b n ).

Finally, for positive a, b > 0, denote a b if a/b is at most some universal constant. Definition C.1. (ORLICZ NORM) The Young-Orlicz modulus is a convex non-decreasing function ψ : IR + → IR + such that ψ(0) = 0 and ψ(x) → ∞ when x → ∞. Accordingly, the Orlicz norm of an integrable random variable X with respect to the modulus ψ is defined as

EQUATION

)

In the sequel, we consider the Orlicz modulus ψ ν (x) def = (x ν ) -1 . Accordingly, the cases of • ψ2 and • ψ1 norms are called the sub-Gaussian and the sub-exponential norms and have the following alternative definitions: Definition C.2. (SUB-GAUSSIAN NORM) The sub-Gaussian norm of a random variable Z, denoted by Z ψ2 , is defined as

EQUATION

)

For a random vector Z ∈ IR n , its sub-Gaussian norm is defined as follows

EQUATION

)

Definition C.3. (SUB-EXPONENTIAL NORM) The sub-exponential norm of a random variable Z, denoted by Z ψ1 , is defined as follows

EQUATION

)

For a random vector Z ∈ IR n , its sub-exponential norm is defined below

EQUATION

C.1 PROOF OF THEOREM 4.1

By the triangle inequality, we have that

EQUATION

where the terms A i , i = 1, 2, 3, 4 are defined as follows

A 1 def = MMD µ * [P V , P W * ] -min W ∈W sup µ∈P MMD µ [P V , P W ] A 2 def = min W ∈W sup µ∈P MMD µ [P V , P W ] -min W ∈W sup µ N ∈P N MMD µ N P V , P W A 3 def = min W ∈W sup µ N ∈P N MMD µ N P V , P W -MMD µ N * P V , P W A 4 def = MMD µ N * P V , P W N * -MMD µ N * P V , P W N * .

In the sequel, we compute an upper bound for each term on the right hand side of equation 38:

Upper bound on A 1 :

First, notice that the squared kernel MMD loss in equation 4 can be characterized in terms of class labels and features defined in Section 3.1 as follows

EQUATION

To see this equivalence, we first rewrite the right hand side of equation 39 as follows

IE P ⊗2 y,x [y yK µ (x, x)] = IP{y = +1}IP{ y = +1}IE x, x∼P ⊗2 x|y=+1 [K µ (x, x)] + IP{y = -1}IP{ y = -1}IE x, x∼P ⊗2 x|y=-1 [K µ (x, x)] -IP{y = -1}IP{ y = +1}IE x∼P x|y=-1 , x∼P x|y=+1 [K µ (x, x)] -IP{y = +1}IP{ y = -1}IE x∼P x|y=+1 , x∼P x|y=-1 [K µ (x, x)].

(40) Now, recall from Section 3.1 that P x|y=+1 = P V , and P x|y=-1 = P W by construction of the labels and random features. Moreover, y, y ∼ i.i.d. Uniform{-1, +1}, and thus IP{y = -1} = IP{y = +1} = 1 2 . Therefore, from equation 40, we derive

IE P ⊗2 y,x [y yK µ (x, x)] = 1 4 IE P ⊗2 V [K µ (x; x)] + 1 4 IE P ⊗2 W [K µ (x; x)] - 1 2 IE P V ,P W [K µ (x; x)] = 1 4 MMD µ [P V , P W ].

For any given W ∈ W, we have that

sup µ∈P MMD µ [P V , P W ] -sup µ∈P MMD µ [P V , P W ] ≤ sup µ∈P MMD µ [P V , P W ] -MMD µ [P V , P W ] = 4 sup µ∈P 1 n(n -1) i =j y i y j K µ (x i , x j ) -IE P ⊗2 y,x [y yK µ (x, x)] = 4 sup µ∈P IE µ [E n (ξ)] ≤ 4 sup µ∈P IE µ [E n (ξ)] ,

xx where the error term is defined using the random features 

EQUATION

EQUATION

)

provided that Ψ is upper semi-continuous in ξ.

Under the strong duality of Theorem C.4, we obtain that

EQUATION

)

In the sequel, let p = 2. The Moreau's envelope Parikh & Boyd (2014) of a function f : X → IR is defined as follows

EQUATION

where β > 0 is the regularization parameter. When the function f is differentiable, the following lemma can be established: Lemma C.5. (MOREAU'S ENVELOPE OF DIFFERENTIABLE FUNCTIONS) Suppose the function f : X → IR is differentiable. Then, the Moreau's envelope defined in equation 44 has the following upper bound and lower bounds

EQUATION

)

In particular, when f is L f -Lipschitz, we have

EQUATION

)

The proof is presented in Appendix D.1. Now, we return to Equation equation 43. We leverage the lower bound on Moreau's envelope in equation 45 of Lemma C.5 as follows

EQUATION

) Let ζ * = ζ * (ξ, s) = arg sup ζ∈IR D ∇E n (1 -s)ξ + sζ 2 .

Then, applying the union bound in conjunction with Inequality equation 47 yields

EQUATION

Now, we state the following lemma: Lemma C.6. (TAIL BOUNDS FOR THE FINITE SAMPLE ESTIMATION ERROR) Consider the estimation error E n defined in equation 41. Then, the following statements hold:

• Z = ∇E n (ξ) 2 2 is a sub-exponential random variable with the Orlicz norm of Z ψ1 ≤ 9×2 9 ×L 4 n 2

for every ξ ∈ IR D . Moreover,

EQUATION

• E n (ξ) is zero-mean sub-Gaussian random variable with the Orlicz norm of E n (ξ) ψ2 ≤ 16 √ 3L 4 n

for every ξ ∈ IR D . Moreover,

EQUATION

)

The proof of Lemma C.6 is presented in Appendix C.2. Now, we leverage the probability bounds equation 49 and equation 50 of Lemma C.6 to upper bound the terms on the right hand side of equation 48 as below

EQUATION

where the last inequality comes from the basic inequality a + b ≤ 2 max{a, b}. Therefore, with the probability of at least 1 -, we have that

EQUATION

for all W ∈ W. 

EQUATION

)

with the probability of (at least) 1 -.

Upper bound on A 2 ::

To establish the upper bound on A 2 , we recall that

EQUATION

Therefore,

sup µ∈P MMD µ [P V , P W ] -sup µ N ∈P N MMD µ N [P V , P W ] ≤ sup µ∈P IE µ [ϕ(x i ; ξ)ϕ(x j ; ξ)] -sup µ N ∈P N IE µ N [ϕ(x i ; ξ)ϕ(x j ; ξ)] . (56)

Here, the last inequality is due to Theorem C.4 and the following duality results hold

sup µ∈P IEµ[ϕ(x; ξ)ϕ( x; ξ)] = inf λ≥0 λR 2 - IR D inf ζ∈IR D {λ ξ -ζ 2 2 -ϕ(xi; ζ)ϕ(xj; ζ)}µ0(dξ) sup µ N ∈P N IE µ N [ϕ(xi; ξ)ϕ(xj; ξ)] = inf λ≥0 λR 2 - 1 N N k=1 inf ζ∈IR D {λ ξ k 0 -ζ 2 2 -ϕ(xi; ζ)ϕ(xj; ζ)} .

Now, in the sequel, we establish a uniform concentration result for the following function

T λ (x, x) : IR N ×D → IR (ξ 1 0 , • • • , ξ N 0 ) → T λ (x, x) (ξ 1 0 , • • • , ξ N 0 ) = 1 N N k=1 M 1 2λ -ϕ(x,•)ϕ( x,•) (ξ k 0 ) - IR D M 1 2λ -ϕ(x,•)ϕ( x,•) (ξ)µ0(dξ).

Then, from equation 56 we have

EQUATION

)

We now closely follow the argument of Rahimi & Recht (2008) to establish a uniform concentration result with respect to the data points x, x ∈ X . In particular, consider an -net cover of X ⊂ IR d .

Then, we require N = 4diam(X ) d balls of the radius > 0, e.g., see (Pollard, 1990, Lemma 4 .1, Section 4). Let Z = {z 1 , • • • , z N } ⊂ X denotes the center of the covering net. Now, let

EQUATION

)

Without loss of generality suppose M

EQUATION

where 

T λ (zi,zj ) (ξ 1 0 , • • • , ξ k 0 , • • • , ξ N 0 ) -T λ (zi,zj ) (ξ 1 0 , • • • , ξ k 0 , • • • , ξ N 0 ) ≤ 2L 2 N .

From McDiarmid's Martingale inequality McDiarmid (1989) and the union bound, we obtain that

EQUATION

for all λ ≥ 0. Now, consider arbitrary points (x, x) ∈ X × X . Let the center of the balls containing those points be z i , z j ∈ Z, i.e., x ∈ IB ε (z i ) and x ∈ IB ε (z j ) for some z i , z j ∈ Z. Then, by the triangle inequality, we have that

EQUATION

where

L T = L T (ξ 1 0 , • • • , ξ N 0 ) def = sup x, x∈X ∇ x T λ (x, x) (ξ 1 0 , • • • , ξ N 0 )

2 is the Lipschitz constant of the mapping T . Note that the Lipschitz constant L T is a random variable with respect to the random feature samples ξ

0 , • • • , ξ N . Let (x * , x * ) def = arg sup x, x∈X ∇ x T λ (x, x) (ξ 1 0 , • • • , ξ N 0 ) 2 .

We compute an upper bound on the second moment of the random variable L T as follows

IEµ 0 L 2 T = IEµ 0 ∇xT λ (x * , x * ) (ξ 1 0 , • • • , ξ N 0 ) 2 2 = IEµ 0   1 N N k=1 ∇xM 1 2λ -ϕ(x * ;•)ϕ( x * ;•) (ξ k 0 ) - IR D ∇xM 1 2λ -ϕ(x * ;•)ϕ( x * ;•) (ξ)µ0(dξ) 2 2   = IEµ 0   1 N N k=1 ∇xM 1 2λ -ϕ(x * ;•)ϕ( x * ;•) (ξ k 0 ) 2 2   -IEµ 0 IR D ∇xM 1 2λ -ϕ(x * ;•)ϕ( x * ;•) (ξ)µ0(dξ) 2 2 ≤ 1 N 2 IEµ 0   N k=1 ∇xM 1 2λ -ϕ(x * ;•)ϕ( x * ;•) (ξ k 0 ) 2 2   .

We further proceed using the triangle inequality as well as the basic inequality (a

EQUATION

)

To proceed from equation 62, we leverage the following lemma: Lemma C.8. (MOREAU'S ENVELOP OF PARAMETRIC FUNCTIONS) Consider the parametric function f : X × Θ → IR and the associated Moreau's envelope for a given θ ∈ Θ ⊂ IR d :

EQUATION

Furthermore, define the proximal operator as follows

EQUATION

Then, Moreau's envelope has the following upper bound

EQUATION

)

The proof is presented in Appendix D.2.

Equipped with Inequality equation 65 of Lemma C.8, we now compute an upper bound on the right hand side of equation 62 as follows

EQUATION

where the last inequality is due to (A.2).

Invoking Markov's inequality now yields:

EQUATION

)

Now, using the union bound, for every arbitrary pair of data points (x, x) ∈ X × X the following inequality holds

IP |T λ (x, x) (ξ 1 0 , • • • , ξ N 0 )| ≥ δ ≤ IP |T λ (zi,zj ) (ξ 1 0 , • • • , ξ N 0 )| ≥ δ/2 + IP |T λ (x, x) (ξ 1 0 , • • • , ξ N 0 ) -T λ (zi,zj ) (ξ 1 0 , • • • , ξ N 0 )| ≥ δ/2 (68) ≤ 2 δ 2 L 4 + 4diam(X ) d • - N δ 2 L 2 .

Following the proposal of Rahimi & Recht (2008) , we select = (κ 1 /κ 2 )

1 d+2 , where κ 1 def = (4diam(X )) d • e -2N λδ 2 2L 2 λ+L 4 and κ 2 def = (2/δ) 2 L 4 . Then, IP sup x, x∈X |T λ (x, x) (ξ 1 0 , • • • , ξ N 0 )| ≥ δ ≤ 2 8 L 2 diam(X ) δ 2 • - N δ 2 L 2 (d + 2) .

Thus, with the probability of at least 1 -, the following inequality holds

EQUATION

)

where W(•) is the Lambert W -function. 4 Since W(x) ≤ ln(x) for x > e, we can rewrite the upper bound in terms of elementary functions

EQUATION

provided that N is sufficiently large and/or is sufficiently small so that 2 8 N diam 2 (X ) ≥ e. Plugging Inequality equation 70 in equation 57 now results in the following inequality

EQUATION

for all W ∈ W. Employingequation 53 from Lemma C.7 now yields the following upper bound

EQUATION

)

Upper bound on A 3 :

Recall that the solution of the empirical risk function of equation 14 is denoted by

EQUATION

)

We also define the solution of the empirical kernel alignment as follows

EQUATION

)

Due to the optimality of the empirical measure µ N * for the inner optimization in equation 73, the following inequality holds

EQUATION

Upon expansion of MMD α µ N P V , P W N *

, and after rearranging the terms in equation 75, we arrive at

EQUATION

where the last inequality is due to the fact that ϕ ∞ < L by (A.1). Now, due to optimality of W N for the outer optimization problem in equation 74, we have

EQUATION

)

Putting together Inequalities equation 76 and equation 77 yields

EQUATION

)

Similarly, due to the optimality of the empirical measure µ N for the optimization in equation 74 we have that

EQUATION

)

Combining equation 78 and equation 79 then yields

EQUATION

)

Upper bound on A 4 :

The upper bound on A 4 can be obtained exactly the same way as A 1 . Specifically, from equation 52 it follows directly that

EQUATION

)

Now, plugging the derived upper bounds in A 1 -A 4 in equation 38 and employing the union bound completes the proof.

C.2 PROOF OF THEOREM 4.2

The proof has three main ingredients and follows the standard procedure in the literature, see, e.g., Wang et al. (2017) ; Luo & Mattingly (2017) . In the first step, we identify the mean-field limit of the particle SGD in equation 15. In the second step, we prove the convergence of the measured-valued process {(µ N t ) 0≤t≤T } to the mean-field solution by establishing the pre-compactness of Sokhorhod space. Lastly, we prove the uniqueness of the mean-field solution of the particle SGD.

Step 1-Identification of the scaling limit: First, we identify the weak limit of converging subsequences via the action of the empirical measure

µ N m (ξ) = 1 N N k=1 δ(ξ -ξ k m ) on a test function f ∈ C 3 b (IR D ).

In particular, we use the standard techniques of computing the scaling limits from Luo & Mattingly (2017) .

Recall that the action of an empirical measure on a bounded function is defined as follows

EQUATION

)

We analyze the evolution of the empirical measure µ N m via its action on a test function

f ∈ C 3 b (IR D ). Using Taylor's expansion, we obtain f, µ N m+1 -f, µ N m = f, µ N m+1 -f, ν N m+1 = 1 N N k=1 f (ξ k m+1 ) -f (ξ k m ) = 1 N N k=1 ∇f (ξ k m )(ξ k m+1 -ξ k m ) T + R N m .

where R N m is a remainder term defined as follows

EQUATION

where

ξ k def = ( ξ k (1), • • • , ξ k (p)), and ξ k (i) ∈ [ξ k m (i), ξ k m+1 (i)], for i = 1, 2, • • • , p. Plugging the difference term (ξ k m+1 -ξ k m ) from the SGD equation in equation 15 results in f, µ N m+1 -f, µ N m (85) = η N 2 α N k=1 ∇f (ξ k m ) • 1 N N =1 ϕ(xm; ξ m )ϕ( xm; ξ m ) -αym ym ∇ ξ ϕ(xm; ξ k m )ϕ( xm; ξ k m ) + R N m .

Now, we define the drift and Martingale terms as follows

D N m def = η N α X ×Y ϕ(x, ξ)ϕ( x, ξ), µ N m -αy y (86a) × ∇f (ξ)(ϕ( x; ξ)∇ ξ ϕ(x; ξ) + ϕ(x; ξ)∇ ξ ϕ( x; ξ)), µ N m dP ⊗2

x,y ((x, y), ( x, y))

M N m def = η N α ϕ(xm, ξ)ϕ( xm, ξ), µ N m -αym ym (86b) × ∇f (ξ)(ϕ( xm; ξ)∇ ξ ϕ(xm; ξ) + ϕ( xm; ξ)∇ ξ ϕ(xm; ξ)), µ N m -D N m .

respectively. Using the definitions of D N m and M N m in equation 86a-equation 86b, we recast Equation equation 85 as follows

EQUATION

Summation over = 0, 1, 2 • • • , m -1 and using the telescopic sum yields

EQUATION

)

We also define the following continuous embedding of the drift, martingale, and the remainder terms as follows

EQUATION

The scaled empirical measure µ N t def = µ N N t then can be written as follows

EQUATION

)

Since the drift process (D N t ) 0≤t≤T is a piecewise cádlág process, we have

EQUATION

where the functional R[µ s ] is defined as follows

R[µ s ] def = η α X ×Y ( ϕ(x, ξ)ϕ( x, ξ), µ s -αy y) (92) × ∇f (ξ)(ϕ( x; ξ)∇ ξ ϕ(x; ξ) + ϕ(x; ξ)∇ ξ ϕ( x; ξ)) T , µ s P ⊗2

x,y ((dx, d x), (dy, d y)).

Therefore, the expression in equation 90 can be rewritten as follows

EQUATION

In the following lemma, we prove that the remainder term sup 0≤t≤T |R N t | vanishes in probabilistic sense as the number of particles tends to infinity N → ∞: Lemma C.9. (LARGE N -LIMIT OF THE REMAINDER PROCESS) Consider the remainder process (R N t ) 0≤t≤T defined via scaling in equation 84-equation 89c. Then, there exists a constant C 0 > 0 such that

EQUATION

)

and thus lim sup

N →∞ sup 0≤t≤T |R N t | = 0.

Proof. The proof is relegated to Appendix D.5.

We can also prove a similar result for the process defined by the remainder term: Lemma C.10. (LARGE N -LIMIT OF THE MARTINGALE PROCESS) Consider the Martingale process (M N t ) 0≤t≤T defined via scaling in equation 86b-equation 89b. Then, for some constant C 1 > 0, the following inequality holds

EQUATION

)

In particular, with the probability of at least 1 -ρ, we have

EQUATION

)

and thus lim sup N →∞ sup 0≤t≤T |M N t | = 0 almost surely.

Proof. The proof is deferred to Appendix D.6. Now, using the results of Lemmata C.10 and C.9 in conjunction with equation 93 yields the following mean-field equation as N → ∞,

µ t , f = µ 0 , f + η α t 0 X ×Y ( ϕ(x, ξ)ϕ( x, ξ), µ s -αy y) (97) × ∇f (ξ)(ϕ( x; ξ)∇ ξ ϕ(x; ξ) + ϕ(x; ξ)∇ ξ ϕ( x; ξ)), µ s P ⊗2

x,y ((dx, d x), (dy, d y)) ds.

Notice that he mean-field equation in equation 97 is in the weak form. When the Lebesgue density p t (ξ) = dµ t /dξ exists, the McKean-Vlasov PDE in equation 22 can be readily obtained from equation 97.

Step 2: Pre-compactness of the Skorkhod space: To establish our results in this part of the proof, we need a definition and a theorem:

Definition C.11. (TIGHTNESS) A set A of probability measures on a metric space S is tight if there exists a compact subset S 0 ⊂ S such that

EQUATION

for all ε > 0. A sequence {X N } N ∈IN of random elements of the metric space S is tight if there exists a compact subset S 0 ⊂ S such that

EQUATION

for all ε > 0, and all N ∈ IN.

Now, to show the tightness of the measured valued process (µ N t ) 0≤t≤T , we must verify Jakubowski's criterion (Jakubowski, 1986, Thm. 1) :

Theorem C.12. (JAKUBOWSKI'S CRITERION (JAKUBOWSKI, 1986, THM. 1) 

) A sequence of measured-valued process {(ζ N t ) 0≤t≤T } N ∈IN is tight in D M(IR D ) ([0, T ]

) if and only if the following two conditions are satisfied:

(J.1) For each T > 0 and γ > 0, there exists a compact set U T,γ such that lim

EQUATION

)

This condition is referred to as the compact-containment condition.

(J.2) There exists a family H of real-valued functions H : M(IR D ) → IR that separates points in M(IR D ) and is closed under addition such that for every H ∈ H, the sequence

{(H(ξ N t )) 0≤t≤T } N ∈IN is tight in D IR ([0, T ]).

To establish (J1), we closely follow the proof of (Giesecke et al., 2013, Lemma 6.1.) . In particular, for each L > 0, we define S L = [0, B] p . Then, S B ⊂ IR p is compact, and for each t ≥ 0, and N ∈ IN, we have

EQUATION

EQUATION

EQUATION

where (a) follows from Markov's inequality, and (b) follows from the upper bound on the norm of the particles in equation 177 of Appendix D. We now define the following set

EQUATION

We let U T,γ = U B , where U B is the completion of the set U B . By definition, U T,γ is a compact subset of M(IR D ). Now, we have

EQUATION

)

Now, since

EQUATION

this implies that for any γ > 0, there exists a B > 0, such that

EQUATION

This completes the proof of (J.1). To verify (J.2), we consider the following class of functions

EQUATION

By definition, every function H ∈ H is continuous with respect to the weak topology of M(IR D ) and further the class of functions H separate points in M(IR D ) and is closed under addition. Now, we state the following sufficient conditions to establish (J.2). The statement of the theorem is due to (Robert, 2013 , Thm. C.9): Theorem C.13. (TIGHTNESS IN D IR ([0, T ]), (ROBERT, 2013, THM. C.9)) A sequence

{(Z N t ) 0≤t≤T } N ∈IN is tight in D IR ([0, T ]

) iff for any δ > 0, we have (T.1) There exists > 0, such that

EQUATION

for all N ∈ IN.

(T.2) For any ρ > 0, there exists σ > 0 such that

EQUATION

This completes the tightness proof of the of the laws of the measured-valued process {(µ N t ) 0≤t≤T } N ∈IN . Now, we verify the condition (J.2) by showing that the sufficient conditions (T.1) and (T.2) hold for function values {(H(µ N t )) 0≤t≤T } N ∈IN , where H ∈ H and H is defined in Eq. equation 108. Now, condition (T.1) is readily verified since

EQUATION

where in the last step, we used the fact that f ∈ C 3 b (IR D ), and hence, f ∞ ≤ b. Thus, IP(H(µ N 0 ) ≥ b) = 0 for all N ∈ IN, and the condition (T.1) is satisfied. Now, consider the condition (T.2). From Equation equation 93, and with 0 ≤ t 1 < t 2 ≤ T we have

EQUATION

To bound the first term, recall the definition of R[µ s ] from equation 92. The following chain of inequalities holds,

EQUATION

Let I : IR D → IR, I(ξ) = 1 denotes the identity function. Notice that I, µ s = IR D µ s (ds) = 1.

From equation 115, we proceed as follows

EQUATION

where the last inequality is due to (A.1). Therefore,

EQUATION

where

s 0 def = 2η α (L 2 + α)L 2 C 1 .

Consider the middle term of equation 114. Using the definition of the martingale term in equation 89b, we obtain that

EQUATION

In Equation of Section D, we have proved the following concentration bound

EQUATION

Now, recall the alternative definition of the sub-Gaussian random variables:

EQUATION

EQUATION

)

We enumerate a few standard consequences of sub-Gaussianity Boucheron et al. (2013) . If X i are independent and σ 2 i -sub-Gaussian, then

EQUATION

)

Now, it is clear from equation 119 andthat M N m is sub-Gaussian random variable with a zero mean, and with the parameter

σ 2 m = 4mL 4 η 2 C 2 1 (L 2 + α) 2 N 2 α 2

. Therefore,

N t2 = N t1 M N is sub-Gaussian with the parameter σ 2 (t 1 , t 2 ) def = 2L 4 η 2 C 2 1 (L 2 + α) 2 N 2 α 2 ( N t 1 -N t 2 + 1)( N t 1 + N t 2 ).

Consequently, from Inequality equation 118 and the concentration inequality in equation 121, we have

EQUATION

where

(t * 1 , t * 2 ) def = arg sup t1,t2≤T,|t1-t2|≤ρ N t2 = N t1 M N .

We first compute a bound for the last term of equation 114 using the definition of the scaled term R N t from equation 89c. We have

EQUATION

where (a) follows from the upper bound in equation 178 of Section D, and in (b) we define s

1 def = C 0 N (ηL 2 + (L 4 /α)).

Putting together equation 117, equation 122, and equation 126, we conclude from Inequality equation 114 that

IP sup t 1 ,t 2 ≤T,|t 1 -t 2 |≤ρ |H(µ N t 1 ) -H(µ N t 1 )| ≥ σ ≤ IP sup t 1 ,t 2 ≤T,|t 1 -t 2 |≤ρ |M N t 1 -M N t 2 | + (s0 + s1)ρ ≥ σ ≤ 2 - α 2 (σ -(s0 + s1)ρ) 2 4L 4 η 2 C 2 1 (L 2 + α) 2 (ρ + 1)T .

Therefore, condition (T.2) is also satisfied. Since the sufficient conditions (T.1) and (T.2) are satisfied, the condition (J.2) is satisfied. This completes the tightness proof of the measured-valued sequence {µ N t } N ∈IN . Now, we prove its convergence to a mean-field solution (µ * t ) 0≤t≤T . Theorem C.15. (PROKHOROV'S THEOREM PROKHOROV ( 1956)) A subset of probability measures on a complete separable metric space is tight if and only if it is pre-compact.

According to Theorem C.15, the tightness of the Skorkhod Space D M(IR D ) ([0, T ]) implies its precompactness which in turn implies the existence of a converging sub-sequence

{(µ N t ) 0≤t≤T } N k of {µ N t } N ∈IN . Notice that {(µ N t ) 0≤t≤T } N k

is a stochastic process defined on the Skorkhod space. Therefore, let π N k denotes the law of the converging sub-sequence {(µ N t ) 0≤t≤T } N k . By definition, π N k is an element of the measure space M(D [0,T ] (M(IR D ))). In the sequel, we closely follow the argument of (Wang et al., 2017, Proposition 4) to show that the limiting measure π ∞ is a Dirac's delta function concentrated at a mean-field solution µ * t ∈ D [0,T ] (M(IR D )). We define the following functional

EQUATION

We compute the expectation of the functional F t with respect to π N k . We then have

EQUATION

Now, from Equation equation 93, we have that

EQUATION

Plugging equation 129 in equation 128 gives

EQUATION

where the last equality is due to the bounds in equation 94 and equation 95 of Lemmata C.9 and C.10, respectively. Taking the limit of N → ∞ from equation 130 yields

EQUATION

It can be shown that the functional F t [•] is continuous and bounded. Therefore, due the weak convergence of the sequence {π N k } N k ∈IN to π ∞ , equation 131 implies that

EQUATION

)

Since the identity equation 132 holds for all bounded test functions f ∈ C 3 b (IR D ) and for all t ∈ (0, T ], it follows that π ∞ is a Dirac's delta function concentrated at a solution (µ * t ) 0≤t≤T of the mean-field equation.

Step 3: Uniqueness of a mean-field solution: Before we establish the uniqueness result we make two remarks:

First, we make it clear that from the compact-containment condition (J.1) of Jakubowski's criterion in Theorem C.12, the support of the measured-valued process (µ N t ) 0≤t≤T = ( µ N N t ) 0≤t≤T is compact for all 0 ≤ t ≤ T . Moreover, in Step 2 of the proof, we established that the measure valued process (µ N t ) 0≤t≤T converges weakly to a mean-field solution as the number of particles tends to infinity (i.e., N → ∞). Thus, all the possible solutions of the mean-field equation also have compact supports. Let Ξ ⊂ IR D denotes a compact set containing the supports of all such solutions at 0 ≤ t ≤ T . In the sequel, it suffices to establish the uniqueness of the mean-field solution for the test functions with a compact domain, i.e., let f ∈ C 3 b ( Ξ).

Second, for all bounded continuous test functions f ∈ C 3 b ( Ξ), the operator f → µ t , f is a linear operator with µ t (IR D ) = 1. Hence, from Riesz-Markov-Kakutani representation theorem Rudin (1987); Varadarajan (1958) by assuming µ t ∈ M(IR D ) , existence of unique operator implies f → f, µ t implies the existence of the unique probability measure µ t . Now, we equip the measure space M(IR D ) with the following norm

EQUATION

)

Given an initial measure µ 0 , we next prove that there exists at most one mean-field model solution by showing that there exists at most one real valued process µ t , f corresponding to the mean-field model. Suppose (µ * ,1 t ) 0≤t≤T , (µ * ,2 t ) 0≤t≤T are two solutions satisfying the mean-field equations equation 97 with the initial distributions µ 1 0 , µ 2 0 ∈ M(IR D ), respectively. For any test function

f ∈ C 3 b ( Ξ) we have that µ * ,1 t -µ * .2 t , f = µ 1 0 -µ 2 0 , f + η α t 0 X ×Y ϕ(x, ξ)ϕ( x, ξ), µ * ,1 s -µ * ,2 s -αy y (134) × ∇f (ξ)(∇ ξ (ϕ( x; ξ)ϕ(x; ξ))) T , µ * ,1 s -µ * ,2 s P ⊗2

x,y ((dz, d z) ds.

We bound the first term on the right side of Equation equation 134 as follows

EQUATION

where used the definition of the norm • on the measure space M(IR D ) from equation 133.

xxxvi Under review as a conference paper at ICLR 2020 Furthermore, let

EQUATION

where in the last inequality, we used the fact that ∇f (ξ) ≤ C 1 since the test function is threetimes continuously differentiable f ∈ C 3 b ( Ξ) on a compact support. Similarly, we have

EQUATION

)

Putting together the inequalities in equation 136,equation 137, and equation 138 yield

EQUATION

The above inequality holds for all bounded functions f ∈ C 3 b ( Ξ). Thus, by taking the supremum with respect to f we obtain 

EQUATION

≤ b µ 1 0 -µ 2 0 + L 2 C 1 η t 0 µ * ,1 s -µ * ,2 s ds + L 4 C 1 η α t 0 µ * ,1 s -µ * ,

EQUATION

where c 0 ∈ L ∞ [0, T ] is non-decreasing, and

EQUATION

It follows that for a.e. t ∈ [0, T ]

EQUATION

)

We now apply the extended Gronewall's Inequality equation 142 with p = 1, c 0 

(t) = b µ 1 0 - µ 2 0 , 0 ≤ t ≤ T , c 1 (t) = ηL 2 C 1 , 0 ≤ t ≤ T , c 2 (t) = ηL 4 C1 α , 0 ≤ t ≤ T ,

EQUATION

Thus, starting from an initial measure µ 1 0 = µ 2 0 = µ 0 , there exists at most one solution for the mean-field model equations equation 97.

C.3 PROOF OF COROLLARY 4.2.1:

To establish the proof, we recall from equation 88 that

EQUATION

for all f ∈ C b (IR D ). Recall the definition of the total variation distance TV(•, •) on a metric space (X , d).

EQUATION

The total variation distance admits the following variational form TV(µ, ν) = sup 

EQUATION

EQUATION

)

Based on the upper bound equation 178 on the remainder term, we have

EQUATION

for some constant C 0 > 0. Moreover, from the concentration inequality equation 189, we also have that with the probability of at least 1 -δ, the following inequality holds

EQUATION

Lastly, recall the definition of the drift term in equation 86a. By carrying out a similar bounding method leading to equation 116, it can be shown that

EQUATION

By plugging equation 150, equation 151, and equation 152 into equation 149, we derive that

EQUATION

with the probability of 1 -δ. We now leverage the following lemma: xxxviii Lemma C.17. (BOUNDED EQUIVALENCE OF THE WASSERSTEIN AND TOTAL VARIATION DIS-TANCES, SINGH & P ÓCZOS (2018)) Suppose (X , d) is a metric space, and suppose µ and ν are Borel probability measures on X with countable support; i.e., there exists a countable set X ⊆ X such that µ(X ) = ν(X ) = 1. Then, for any p ≥ 1, we have Sep(X )(2TV(µ, ν))

EQUATION

where Diam(X ) Consider the metric space (IR D , • 2 ). Note that the empirical measures

µ N m , µ N 0 have a countable support X = {ξ k m } N k=1 ∪ {ξ k 0 } N k=1 ⊂ IR D .

Therefore, using the upper bounds in equation 154 of Lemma C.17 and 153, we conclude that when the step-size is of the order

EQUATION

) then W p ( µ N m , µ N 0 ) ≤ R for all m ∈ [0, N T ] ∩ IN.

D PROOFS OF AUXILIARY RESULTS:

D.1 PROOF OF LEMMA C.5

The upper bound follows trivially by letting x = y in the optimization problem equation 44.

Now, consider the lower bound. Define the function g :

[0, 1] → IR, t → g(t) = f (y + t(x -y)).

Then, when f is differentiable, we have g (t) = xy, ∇f (y + t(x -y)) . In addition, g(0) = f (y), and g(1) = f (x). Based on the basic identity g(1) = g(0) + 1 0 g (s)ds, we derive

EQUATION

where the last step is due to the Cauchy-Schwarz inequality. Using Inequality equation 156 yields the following lower bound on Moreau's envelope

EQUATION

where (a) is due to Jensen's inequality.

D.2 PROOF OF LEMMA C.8:

Let u ∈ IR d denotes an arbitrary unit vector u 2 = 1. From the definition of the gradient of a function, we have that

EQUATION

where

EQUATION

We now have

EQUATION

where the inequality in (a) follows by letting y = Prox f (•;θ) (x) in the optimization problem. Now, recall that

EQUATION

Substitution of equation 161 in equation 160 yields

M f (•;θ+δu) (x) ≤ M f (•;θ) (x) + δ ∇ θ f (Prox f (•;θ) (x); θ), u + O(δ 2 ). (162) Hence, D u [M f (•;θ) (x)] ≤ ∇ θ f (Prox f (•;θ) (x); θ), u .

From equation 158 and by using Cauchy-Schwarz inequality, we compute the following bound on the inner product of the gradient with the unit vectors u ∈ IR d , u 2 = 1,

EQUATION

)

Since the preceding upper bound holds for all the unit vectors u ∈ IR d , we let u =

∇ θ M β f (•;θ) (x) ∇ θ M β f (•;θ) (x) 2

to get Inequality equation 65.

D.3 PROOF OF LEMMA C.6:

Let z ∈ S d-1 denotes an arbitrary vector on the sphere. Define the random variable

EQUATION

xl Under review as a conference paper at ICLR 2020

Clearly, IE Px,y [Q z ] = 0. Now, let ( y m , x m ) ∈ Y × X , 1 ≤ m ≤ n.

By repeated application of the triangle inequality, we obtain that

EQUATION

where the last inequality is due to assumption (A.2) and the fact that z 2 = 1 for z ∈ S d-1 . In particular, to derive Inequality equation 165, we employed the following upper bounds

|ϕ(x i ; ξ)| ≤ L, |y m ϕ(x m ; ξ) -y m ϕ( x m ; ξ)| ≤ |ϕ(x m ; ξ)| + |ϕ( x m ; ξ)| ≤ 2L, ∇ ξ ϕ(x i ; ξ) 2 ≤ L, y m ∇ ξ ϕ(x m ; ξ) -y m ∇ ξ ϕ( x m ; ξ) 2 ≤ ∇ ξ ϕ(x m ; ξ) 2 + ∇ ξ ϕ( x m ; ξ) 2 ≤ 2L.

Using McDiarmid Martingale's inequality McDiarmid (1989) then gives us

EQUATION

for x ≥ 0. Now, for every p ∈ IN, the 2p-th moment of the random variable Q z is given by

EQUATION

where (a) is due to the concentration bound in equation 166. Now Therefore, 

IE Q 2 z ((y 1 , x 1 ), • • • , (y n , x n ))/γ 2 = ∞ p=0 1 p!γ 2p IE φ 2p z ((y 1 , x 1 ), • • • , (y n , x n )) = 1 + 2 p∈IN 16L 4 nγ 2p = 2 1 -(16L 4 /nγ) 2 -1. For γ = 16 √ 3L 4 /n, we obtain IE Q 2 z ((y 1 , x 1 ), • • • , (y n , x n ))/γ 2 ≤ 2. Therefore, Q z ψ2 = z, ∇E n (ξ) ψ2 ≤ 16 √ 3L 4 /n

≥ 0. Then, Z 2 2 ψ ν 2 ≤ 2 • 3 2 ν • β 2 .

Using Lemma D.1, we now have that ∇E n (ξ) 2 2 ψ1 ≤ 4608L 4 /n 2 for every ξ ∈ IR D . Applying the exponential Chebyshev's inequality with β = 4608L 4 /n 2 yields

IP IR D 1 0 ∇E n ((1 -s)ξ + sζ * ) 2 2 -IE x,y [ ∇E n ((1 -s)ξ + sζ * ) 2 2 ] µ 0 (dξ) ≥ δ ≤ e -n 2 δ 4608L 4 IE x,y e n 2 4608L 4 IR D 1 0 ∇En((1-s)ξ+sζ * ) 2 2 -IEx,y[ ∇En((1-s)ξ+sζ * ) 2 2 ] dsµ0(dξ) (a) ≤ e -n 2 δ 4608L 4 IR D 1 0 IE x,y e n 2 4608L 4 (| ∇En((1-s)ξ+sζ * ) 2 2 -IEx,y[ ∇En((1-s)ξ+sζ * ) 2 2 ])| dsµ 0 (dξ) (b) ≤ 2e -n 2 δ 4608L 4 ,

where (a) follows by Jensen's inequality, and (b) follows from the fact that

EQUATION

due to Definition C.1. Therefore,

EQUATION

It now remains to compute an upper bound on the expectation

IE x,y [ ∇E n ((1 -s)ξ + sζ * ) 2 2 ]

. But this readily follows from equation 167 by letting p = 1 and z = ∇En((1-s)ξ+sζ * ) ∇En((1-s)ξ+sζ * ) 2 as follows 

EQUATION

EQUATION

Therefore,

EQUATION

Similarly, it can be shown that min

EQUATION

Combining equation 172 and equation 173 yields the desired inequality.

D.5 PROOF OF LEMMA C.10:

We recall the expression of the remainder term {R N m } 0≤m≤N T from equation 84. For each 0 ≤ m ≤ N T , N ∈ IN, we can bound the absolute value of the remainder term as follows

EQUATION

Next, we characterize a bound on the difference term ξ k m+1 -ξ k m 2 . To attain this goal, we use the iterations of the particle SGD in Equation equation 15. We have that

EQUATION

where in (a), we used the fact that ϕ ∞ < L and ∇ ξ ϕ(x, ξ) 2 < L due to (A.1), and y m , y m ∈ {-1, 1}. Plugging the last inequality in equation 175 yields

EQUATION

We next compute an upper bound on the Frobenious norm of the Hessian matrix ∇ 2 f ( ξ k ). To this end, we first show that there exists a compact 

EQUATION

For the difference terms, we derive that

EQUATION

where the last step follows from the fact that ϕ ∞ ≤ L due to (A.1). Similarly, we obtain that 

EQUATION

Thus, (M N m ) m∈[0,N T ]∩IN is a Martingale process with bounded difference. From the Azuma-Hoeffding inequality it follows that

EQUATION

)

xlv Under review as a conference paper at ICLR 2020 Therefore, since M N t = M N N t , we have

EQUATION

)

Then,

EQUATION

)

where the inequality follows from equation 190.

By Doob's Martingale inequality Doob (1953) , the following inequality holds

EQUATION

)

In particular, with the probability of at least 1 -ρ, we have

EQUATION

E CHAOTICITY AND PROPAGATION OF CHAOS IN PARTICLE SGD:

In this appendix, we establish the so called 'propagation of chaos' property of particle SGD. We now establish the so called 'propagation of chaos' property of particle SGD. At a high level, the propagation of chaos means that when the number of samples {ξ k } N k=1 tends to infinity (N → +∞), their dynamics are decoupled. Definition E.1. (EXCHANGABLITY) Let ν be a probability measure on a Polish space S and. For N ∈ IN, we say that ν ⊗N is an exchangeable probability measure on the product space S n if it is invariant under the permutation π def = (π(1), • • • , π(N )) of indices. In particular,

EQUATION

for all Borel subsets B ∈ B(S n ).

An interpretation of the exchangablity condition equation 195 can be provided via De Finetti's representation theorem which states that the joint distribution of an infinitely exchangeable sequence of random variables is as if a random parameter were drawn from some distribution and then the random variables in question were independent and identically distributed, conditioned on that parameter.

Next, we review the mathematical definition of chaoticity, as well as the propagation of chaos in the product measure spaces: Definition E.2. (CHAOTICITY) Suppose ν ⊗N is exchangeable. Then, the sequence {ν ⊗N } N ∈IN is ν-chaotic if, for any natural number ∈ IN and any test function

EQUATION

For the first term on the right side of equation 202 we have lim sup

N →∞ IE k=1 µ N t , f k - k=1 µ * t , f k (a)

≤ lim sup

N →∞ IE k=1 µ N t , f k - k=1 µ * t , f k (b)

≤ IE lim sup

EQUATION

where Further, by the definition of the empirical measure µ N t we obtain that

EQUATION

Therefore, subtracting equation 204 and equation 205 yields

EQUATION

)

Hence, lim sup

EQUATION

)

Combining equation 203-equation 208 yields the desired result.

xlviii

:

To avoid clutter in our subsequent analysis, the normalization factor 16 n(n-1) of the gradient is omitted by modifying the step-size η.

The notion of the weak convergence of a sequence of empirical measures is formally defined in Supplementary.

A topological space is Polish if it is homeomorphic to a complete, separable metric space.

Recall that the lambert W -function is the inverse of the function f (W ) = W e W . xxvi

Figures and Tables:

[D(X; δ)] + IE P Z [log(1 -D(G(Z; W ); δ))].

µ (x, y) = IE µ [ϕ(x; ξ)ϕ(y; ξ)] = Ξ ϕ(x; ξ)ϕ(y; ξ)µ(dξ). (7) Let MMD µ [P V , P W ] def = MMD Kµ [P V , P W ].Then, the kernel optimization problem in can be formulated as a distribution optimization for random features, i.e, min W ∈W sup µ∈P MMD µ [P V , P W ].

Theorem 4.1. (NON-ASYMPTOTIC CONSISTENCY OF FINITE-SAMPLE ESTIMATOR) Suppose conditions (A.1)-(A.3) of Appendix C are satisfied. Consider the distribution balls P and P N that are defined with respect to the 2-Wasserstein distance (p = 2). Furthermore, consider the optimal MMD values of the population optimization and its finite sample estimate (W * , µ * ) P V , P W ].

Due to the mean-field analysis of Theorem 4.2, we can prove that the empirical measure µ N m of the particles in SGD dynamic equation 15 remains inside the feasible distribution ball P N : Corollary 4.2.1. Consider the learning rate η = O R p T √ N T log(2/δ) for the SGD in equation 15. Then, the empirical measure µ N m of the particles remains inside the distributional ball µ N m

Figure 1(a) shows the samples from Algorithm 1, Figure 1(b) shows the samples from MMD GAN Li et al. (2017) with a mixture RBF Gaussian kernel κ(x, y) =

Figure 1: Sample generated images using CIFAR-10 (top row), and MNIST (bottom row) data-sets.Panels (a)-(e): Proposed MMD GAN with an automatic kernel selection via the particle SGD (Algorithm 1), Panels (b)-(f): MMD GAN Li et al. (2017) with an auto-encoder for dimensionality reduction in conjunction with a mixed RBF Gaussian kernel whose bandwidths are manually tuned, Panels (c)-(g): MMD GAN in Li et al. (2017) with a single RBF Gaussian kernel with an auto-encoder for dimensionality reduction in conjunction with a single RBF Gaussian kernel whose bandwidth is manually tuned, Panel (d)-(g): GMMN without an auto-encoder Li et al. (2015).

Figure 2: Visualization of data-points from the synthetic data-set P V = N(0, (1 + λ)I d×d ) and P W = N(0, (1 -λ)I d×d ) for d = 2. Panel (a): λ = 0.1, Panel (b): λ = 0.5, and Panel (c): λ = 0.9.

Figure 3: The evolution of the empirical measure µ N m (ξ) = 1 N

Figure 4: The statistical power versus the threshold τ for the binary hypothesis testing via the unbiased estimator of the kernel MMD. The parameters for this simulations are λ ∈ {0.1, 0.5, 0.9}, d = 100, n + m = 100, p = 50. Panel (a): Trained kernel using the two-phase procedure with the particle SGD in equation 15 and an auto-encoder, Panel (b): Trained kernel with an auto-encoder and a fixed Gaussian kernel with the bandwidth σ = 1, Panel (c): Untrained kernel without an auto-encoder.

Figure 6: The sample generated images from CelebA (top row), and LSUN bedroom data-sets (bottom row). Panels (a): Proposed MMD GAN with an automatic kernel selection via the particle SGD (Algorithm 1), Panels (b): MMD GAN Li et al. (2017) with an auto-encoder for dimensionality reduction in conjunction with a mixed RBF Gaussian kernel whose bandwidths are manually tuned, Panels (c): MMD GAN in Li et al. (2017) with a single RBF Gaussian kernel with an auto-encoder for dimensionality reduction in conjunction with a single RBF Gaussian kernel whose bandwidth is manually tuned.

) Now, we invoke the following strong duality theorem Gao & Kleywegt (2016): Theorem C.4. (STRONG DUALITY FOR ROBUST OPTIMIZATION, (GAO & KLEYWEGT, 2016, THEOREM 1)) Consider the general metric space (Ξ, d), and any normal distribution ν ∈ M(Ξ), where M(Ξ) is the set of Borel probability measures on Ξ. Then, sup µ∈M(Ξ)

Lemma C.7. (DISTANCE BETWEEN MINIMA OF ADJACENT FUNCTIONS) Let Ψ(W ) : W → IR and Φ(W ) : W → IR. Further, suppose Ψ(W ) -Φ(W ) ∞ ≤ δ for some δ > 0. Then, min W ∈W Ψ(W ) -min W ∈W Φ(W ) ≤ δ. (53) See Appendix D.4 for the proof. Let Ψ(W ) def = sup µ∈P MMD µ [P V , P W ], and Φ(W ) def = sup µ∈P MMD µ [P V , P W ]. Then, from Inequality equation 53, we have the following upper bound on A 1

(a) follows by letting ζ = ξ k 0 in the first optimization problem, (b) follows by using the fact that -λ ζξ k 0 2 is non-positive for any ζ ∈ IR D and can be dropped, and (c) follows from Assumption (A.2). Now, plugging the upper bound in equation 59 into equation 58 yields

Now, using the variation form and by taking the supremum of equation 146 with respect to the functions from the function class F c def = {f ∈ C 1/2 (IR D )}, we obtain the following upper bound on the TV distance

sup x,y∈X d(x, y), and Sep(X ) def = inf x =y∈X d(x, y).

for all z ∈ S n-1 and ξ ∈ IR D . Consequently, by the definition xli Under review as a conference paper at ICLR 2020 of the sub-Gaussian random vector in equation 35 of Definition C.2, we have ∇E n (ξ) ψ2 ≤ 16 √ 3L 4 /n for every ξ ∈ IR D . We invoke the following lemma proved by the first author in (Khuzani & Li, 2017, Lemma 16): Lemma D.1. (THE ORLICZ NORM OF THE SQUARED VECTOR NORMS, (KHUZANI & LI, 2017, LEMMA 16)) Consider the zero-mean random vector Z satisfying Z ψν ≤ β for every ν

Plugging the expectation upper bound of equation 170 into equation 169 completes the proof of the first part of Lemma C.6. The second part of Lemma C.6 follows by a similar approach and we thus omit its proof. xlii Under review as a conference paper at ICLR 2020 D.4 PROOF OF LEMMA C.7 Let W * def = arg min W ∈W Ψ(W ) and W def = arg min W ∈W Φ(W ). Then, since |Ψ(W ) -Φ(W )| ≤ δ for all W ∈ W, we have that

set C ⊂ IR D such that ξ k m ∈ C for all k = 1, 2, • • • , N and all m ∈ [0, N T ] ∩ IN. For each k = 1, 2, • • • , N , from Inequality equation 175 we obtain that c 0 for some constant c 0 > 0 since the initial samples ξ 1 0 , • • • , ξ N 0 are drawn from the measure µ 0 whose support support(µ 0 ) = Ξ is assumed to be compact by(A.3). From bound in equation 177, it thus follows that ξ k m 2 < C for some constantC > 0, for all m ∈ [0, N T ] ∩ IN. Now, recall that ξ k = ( ξ k (1), • • • , ξ k (p)), where ξ k (i) ∈ [ξ k m (i), ξ k m+1 (i)], i = 1, 2, • • • , m + 1, for i = 1, 2, • • • , p. Therefore, ξ k ∈ C. Since all the test function f ∈ C 3 b (IR 3) are three-times continuously differentiable, it follows that there exists a constant C 0 def = C 0 (T ) > 0 such that sup ξ∈C ∇ 2 f ( ξ) F < C 0 . From Inequality equation 176, it follows that the definition of the scaled term R N t from equation 89c. Using the Inequality equation 178 as well as the definition of R N t OF LEMMA C.10Let F m-1 = σ((x k , y k ) 0≤k≤m-1 , ( x k , y k ) 0≤k≤m-1) denotes the σ-algebra generated by the samples up to time m -1. We define F -1 def = ∅. Further, define the following random variable∆ N m def = ϕ(x m , ξ)ϕ( x m , ξ), µ N m -αy m y m × ∇f (ξ)∇ ξ (ϕ( x m ; ξ)ϕ(x m ; ξ)), µ N m . (180) Notice that 1 N IE[∆ N m |F m-1 ] = D N m .We now rewrite the martingale term in equation 86b in term of∆ By construction of M N m in equation 181, it is a Martingale IE[M N m |F m-1 ] = M N m-1. We now prove that M N m has also bounded difference. To do so, we define the shorthand notations a N m def = ϕ(x m , ξ)ϕ( x m , ξ), µ N m -αy m y m , ∇f (ξ)(∇ ξ (ϕ( x m ; ξ)ϕ(x m ; ξ))) T , µ N m .

|b N m | = | ∇f (ξ)(∇ ξ (ϕ( x m ; ξ)ϕ(x m ; ξ))) T , µ N m m ; ξ k m )| • ∇f (ξ k m )(∇ ξ ϕ( x m ; ξ k m )) m ; ξ k m )| • |∇f (ξ k m )(∇ ξ ϕ( x m ; ξ k m )) T |and (c) follows from (A.1), and (b) follows from the Cauchy-Schwarz inequality. From Inequality equation 177 and the ensuing disucssion in Appendix D.5, we recall that ξ k m 2 < C for some constant and for all m ∈ [0, N T ] ∩ IN, and k = 1, 2, • • • , N . For the two times continuously test function f ∈ C 3 b (IR D ), it then follows that |∇f (ξ k m ) 2 ≤ C 1 for some constant C 1 > 0. The following bound can now be computed from equation 186, bounds on |a N m | and |b N m | from equation 185-equation 187 into equation 184 we obtain that

(a)  is by Jensen's inequality, (b) is by Fatou's lemma, (c) follows from the basic inequalityN i=1 a i -N i=1 b i ≤ N i=1 |a i -b i | for |a i |, |b i | ≤ 1, i = 1, 2, • • • , N , as well as the fact that µ * t , f k ≤ b and µ N t , f k ≤ b for all k = 1, 2, • • • , N due to the boundedness of the test functions f 1 , • • • , f ∈ C 3 b (IR D ),and (d) follows from the weak convergence µ N t weakly → µ * t to the mean-field solution equation 97. Now, consider the second term on the right hand side of equation 202. Due to the exchangeability of the initial states (ξ k 0 ) 1≤k≤N , the law of the random variables (ξ k 0 ) 1≤k≤N is also exchangeable. Therefore, we obtain that , N ) is the set of all permutations of numbers selected from {1, 2, • • • , N }. Notice that the right hand side of equation 204 is the symmetrized version of the left hand side equation 205.

Comparison

Table content: <table><tr><td>Method</td><td>FID (↓)</td><td>IS (↑)</td></tr><tr><td>MMD GAN (Gaussian) Li et al. (2017)</td><td>67.244 ± 0.134</td><td>5.608±0.051</td></tr><tr><td>MMD GAN (Mixture Gaussian) Li et al. (2017)</td><td>67.129 ± 0.148</td><td>5.850±0.055</td></tr><tr><td>SGD Alg. 1</td><td colspan="2">65.059 ± 0.153 5.97 ± 0.046</td></tr><tr><td>Real Data</td><td>0</td><td>11.237±0.116</td></tr></table>

and u s = µ * ,1 s -µ * ,2

Table content: <table><tr><td/><td>s</td><td>. In</td></tr><tr><td>this case, it is easy to see that M 1 = 2bηT L 4 C1 α we obtain that</td><td colspan="2">. Consequently, from equation 140 and equation 144,</td></tr></table>

Peer review (iclr2020):
This paper proposes to learn a kernel for training MMD-GAN by optimizing over the probability distribution that defines the kernel by means of random features. This is unlike the usual setting of MMD-GAN where the kernel is parametrized by composing a fixed top-kernel with a discriminator network that is optimized during training. The main motivation for this approach is to avoid having to *manually* fix some parameters of the top-level kernel like the bandwidth of the RBF kernel. They provide an algorithm to achieve such optimization in probability space along with some consistency results and perform experiments on MNIST and Cifar10 to demonstrate empirically the advantage of such an approach over those that fix the top-level kernel. Theory: Theorem 4.1 provides a convergence result of an oracle finite-sample estimator: that is the one obtained by exactly solving the optimization problem in 19b. In that case, they show the consistency of the proposed estimator. The result is somehow expected but the proof relies on nice duality results for measures and is very technical. The clarity of the proof could be improved: - Currently, the structure of the main proof mixes direct lemma (lemma B.7) with less obvious ones (lemma B.6). Also, some concepts are introduced in the main proof but not necessary for its understanding: The notion of Orlicz norm is in introduced in Definition B.3 on the fly to state lemma B.6, but only equations 50 and 51 are used in this lemma which does not make use of the notion of Orlicz norm at all. Theorem 4.1 doesn*t say anything about the consistency of the algorithm itself. To partly address this, the authors show in theorem 4.2 that as the number of particles grows the empirical process converges to a McKean Vlasov PDE (equation 22). This means that the proposed algorithm is approximating some gradient flow in metric space (25). - However, this gradient flow is a non-convex optimization problem and there is no guarantee that a global solution is reached. Recent work provides cases when global convergence occurs [Chizat2018] but it is not the case in general. Some further clarification about the connection between Thm 4.1 and Thm 4.2 would be therefore useful. Thm 4.2 is also curious in the sense that the process defined by equation 16, which is noisy since it relies on one sample from the data, would converge towards (22) which is a Mc-valsov equation with a drift only (no diffusion or other noise). What happened to the noise coming from sampling from P_v and P_w in equation 16? Wouldn*t there be some sort of diffusion term as in [Hu2018]? More generally it would be nice to have a discussion of the assumptions and results in the paper as they seem to rely on methods that people in the machine learning community are not totally familiar with. Experiments: The experiments are not convincing for several reasons: - The comparison with the other methods is somehow unfair since the bandwidth is manually tuned for the competing methods. It is easy to adaptively learn the bandwidth as well: in this case, it will be just an additional parameter of a discriminator network. This was done in [Arbel2018] where a single gaussian kernel is used and a regularization of the critic allows to learn the bandwidth without manual tuning. Does the proposed method offer an additional advantage compared to those? - In practice and for a scaling parameter alpha=1, isn*t the algorithm strictly equivalent to considering an MMD-GAN with a dot product kernel and a discriminator given by the feature phi(x,zeta)? - Mnist and Cifar10 are somehow very simple, what would happen on more complicated datasets (CelebA or imagenet)? - I also think there is an ablation that is missing: If the auto-encoder also needs to be optimized then does it also help to optimize over the particles as well or is optimizing the auto-encoder discriminator enough to achieve a similar performance? In other words, does optimizing the auto-encoder compensate for the need to learn the distribution mu? Of course, this would depend on how the auto-encoder is parametrized but I don*t see why it wouldn*t in many cases. Overall, I*m not convinced that the proposed approach would lead to any substantial improvement for MMD-GAN in practice, and the experiments are not really convincing as they are now. However, I find the theoretical results interesting and might be used to better analyse the dynamics of GAN*s. But as the paper is currently framed, it is hard to put these theoretical results in perspective.