Amortized Projection Optimization for Sliced Wasserstein Generative Models
Abstract
Seeking informative projecting directions has been an important task in utilizing sliced Wasserstein distance in applications.
However, finding these directions usually requires an iterative optimization procedure over the space of projecting directions, which is computationally expensive.
Moreover, the computational issue is even more severe in deep learning applications, where computing the distance between two mini-batch probability measures is repeated several times.
This nested loop has been one of the main challenges that prevent the usage of sliced Wasserstein distances based on good projections in practice.
To address this challenge, we propose to utilize the learning-to-optimize technique or amortized optimization to predict the informative direction of any given two mini-batch probability measures.
To the best of our knowledge, this is the first work that bridges amortized optimization and sliced Wasserstein generative models.
In particular, we derive linear amortized models, generalized linear amortized models, and non-linear amortized models which are corresponding to three types of novel minibatch losses, named amortized sliced Wasserstein .
We demonstrate the favorable performance of the proposed sliced losses in deep generative modeling on standard benchmark datasets 1 .
Introduction
Generative modeling is one of the most important tasks in machine learning and data science.
Leveraging the expressiveness of neural networks in parameterizing the model distribution, deep generative models such as GANs [17], VAEs [23] and diffusion models [19, 54], achieve a significant quality of sampling images.
Despite differences in the way of modeling the model distribution, optimization objectives of training generative models can be written as minimizing a discrepancy D ( · , · ) between data distribution µ and the model distribution ν ϕ with ϕ ∈ Φ , parameter space of neural networks weights, namely, we solve for
ˆ ϕ ∈ arg min ϕ ∈
Φ D ( µ, ν ϕ ) .
For example, Kullback–Leibler divergence is used in VAEs and diffusion models, Jensen–Shannon divergence appears in GANs, and f-divergences are utilized in f-GANs [43].
Because of the complexity of the neural networks ϕ , closed-form optimal solutions to these optimization problems are intractable.
Therefore, gradient-based methods and their stochastic versions are widely used in practice to approximate these solutions.
Recently, optimal transport-based losses, which we denote as D ( · , · ) , are utilized to train generative models due to their training stability, efficiency, and geometrically meaning.
Examples of these models include Wasserstein GAN [3] with the dual form of Wasserstein-1 distance [46], and
OT-GANs [14, 51] with the primal form of Wasserstein distance and with Sinkhorn divergence [8] between mini-batch probability measures.
Although these models considerably improve the generative performance, there have been remained certain problems.
In particular, Wasserstein GAN is reported to fail to approximate the Wasserstein distance [55] while OT-GAN suffers from high computational complexity of Wasserstein distance: O ( m 3 log m ) and its curse of dimensionality: the sample complexity of O ( m − 1 /d ) where m is the number of supports of two mini-batch measures.
The entropic regularization [8] had been proposed to improve the computational complexity of approximating optimal transport to O ( m 2 ) [1, 30, 31, 29] and to remove the curse of dimensionality [34].
However practitioners usually choose to use the slicing (projecting version) of Wasserstein distance [57, 11, 25, 42] due to a fast computational complexity O ( m log m ) and no curse of dimensionality O ( m − 1 / 2 ) .
The distance is known as sliced Wasserstein distance (SW) [4].
Sliced Wasserstein distance is defined as the expected one-dimensional Wasserstein distance between two projected measures over the uniform distribution over the unit sphere.
Due to the intractability of the expectation, Monte Carlo samples from the uniform distribution over the unit sphere are used to approximate the distance.
The number of samples is often called the number of projections and it is denoted as L .
From applications, practitioners observe that sliced Wasserstein distance requires a sufficiently large number of projections L relative to the dimension of data to perform well [25, 11].
Increasing L leads to a linear increase in computational time and memory.
However, when data lie in a low dimensional manifold, several projections are redundant since they collapse projected measures to a Dirac-Delta measure at zero.
There are some attempts to overcome that issue including sampling orthogonal directions [49] and mapping the data to a lower-dimensional space [11].
The most popular approach is to search for the direction that maximizes the projected distance, which is known as max-sliced
Wasserstein distance (Max-SW) [10].
Nevertheless, in the context of deep generative models and deep learning in general, the optimization over the unit sphere requires iterative projected gradient descent methods that can be computationally expensive.
In detail, each gradient-update of the model parameters (neural networks) requires an additional loop for optimization of Max-SW between two mini-batch probability measures.
Therefore, we have two nested optimization loops: the global loop (optimizing model parameters) and the local loop (optimizing projection).
These optimization loops can slow down the training considerably.
Contribution.
To overcome the issue, we propose to leverage learning to learn techniques ( amortized optimization ) to predict the optimal solution of the local projection optimization.
We bridge the literature on amortized optimization and optimal transport by designing amortized models to solve the iterative optimization procedure of finding optimal slices in the sliced Wasserstein generative model.
To the best of our knowledge, this is the first time amortized optimization is used in sliced
Wasserstein literature.
In summary, our main contributions are two-fold:
First, we introduce a novel family of mini-batch sliced Wasserstein losses that utilize amortized models to yield informative projecting directions, named amortized sliced Wasserstein losses ( A -SW).
We specify three types of amortized models: linear amortized, generalized linear amortized, and non-linear amortized models that are corresponding to three mini-batch losses: linear amortized sliced Wasserstein ( LA -SW), generalized linear amortized sliced Wasserstein ( GA -SW), and non-linear amortized sliced Wasserstein ( NA -SW).
Moreover, we discuss some properties of A -SW losses including metricity, complexities, and connection to mini-batch Max-SW.
We then introduce the application of A -SW in generative modeling.
Furthermore, we carry out extensive experiments on standard benchmark datasets including CIFAR10, CelebA, STL10, and CelebAHQ to demonstrate the favorable performance of A -SW in learning generative models.
Finally, we measure the computational speed and memory of A -SW, mini-batch Max-SW, and mini-batch SW to show the efficiency of A -SW.
Organization.
The remainder of the paper is organized as follows.
We first provide background about
Wasserstein distance, sliced Wasserstein distance, max-sliced Wasserstein distance, and amortized optimization in Section 2.
In Section 3, we propose amortized sliced Wasserstein distances and analyze some of their theoretical properties.
The discussion on related works is given in Section 4.
Section 5 contains the application of A -SW to generative models, qualitative experimental results, and quantitative experimental results on standard benchmarks.
In Section 6, we provide a conclusion.
Finally, we defer the proofs of key results and extra materials to the Appendices.
Notation.
For any d ≥ 2 , S d − 1 := { θ ∈ R d | ||
θ ||
22 = 1 } denotes the d dimensional unit hyper-sphere in L 2 norm, and U ( S d − 1 ) is the uniform measure over S d − 1 .
Moreover, δ denotes the Dirac delta function.
For p ≥ 1 , P p ( R d ) is the set of all probability measures on R d that has finite p -moments.
For µ, ν ∈ P p ( R d ) , Π( µ, ν ) := { π ∈ P p ( R d × R d ) | (cid:82) R d π ( x, y ) dx =
ν, (cid:82) R d π ( x, y ) dy = µ } is the set of transportation plans between µ and ν .
For m ≥ 1 , we denotes µ ⊗ m as the product measure which has the supports are the joint vector of m random variables that follows µ .
For a vector
X ∈ R dm , X :=
( x 1 , . . .
, x m ) , P X denotes the empirical measures 1 m (cid:80) mi =
δ x i .
We denote θ♯µ as the push-forward probability measure of µ through the function T θ : R d → R where T θ ( x ) =
θ ⊤ x .
Background
In this section, we first review the definitions of the Wasserstein distance, the sliced Wasserstein distance, and the max-sliced Wasserstein distance.
We then formulate generative models based on the max-sliced Wasserstein distances and review the amortized optimization problem and its application to the max-sliced Wasserstein generative models.
2.1 (Sliced)-Wasserstein Distances
We first define the Wasserstein- p distance [56, 45] between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) as follows: W p ( µ, ν ) := (cid:16) inf π ∈ Π( µ,ν ) (cid:82) R d × R d ∥ x − y ∥
pp dπ ( x, y ) (cid:17) 1 p .
When d = 1 , the Wasserstein distance has a closed form which is W p ( µ, ν )
=
( (cid:82)
[Equation]
where F µ and F ν are the cumulative distribution function (CDF) of µ and ν respectively.
To utilize this closed-form property of Wasserstein distance in one dimension and overcome the curse of dimensionality of Wasserstein distance in high dimension, the sliced Wasserstein distance [4] between µ and ν had been introduced and admitted the following formulation:
SW
[Equation]
[Equation] [Equation]
1 p .
For each θ ∈ S d − 1 , W pp ( θ♯µ, θ♯ν ) can be computed in linear time O ( n log n ) where n is the number of supports of µ and ν .
However, due to the integration over the unit sphere, the sliced Wasserstein distance does not have closed-form expression.
To approximate the intractable expectation, Monte Carlo scheme is used, namely, we draw uniform samples θ 1 , . . .
, θ L ∼ U ( S d − 1 ) from the unit sphere and obtain the following approximation:
[Equation]
1 p .
In practice, L should be chosen to be sufficiently large compared to the dimension d .
It is not appealing since the computational complexity of SW is linear with L .
To reduce projection complexity, max-sliced Wasserstein (Max-SW) is introduced [10] .
In particular, the max-sliced Wasserstein distance between µ and ν is given by:
[Equation]
To solve the optimization problem, a projected gradient descent procedure is used.
We present a simple algorithm in Algorithm 1 in Appendix B.
In practice, practitioners often set a fixed number of gradient updates, e.g., T = 100 .
2.2 Learning Generative Models with Max-Sliced Wasserstein and Amortized Optimization
We now provide an application of (sliced)-Wasserstein distances to generative models settings.
The problem can be seen as the following optimization:
[Equation]
where D ( · , · ) can be Wasserstein distance or SW distance or Max-SW distance.
Despite the recent progress on scaling up Wasserstein distance in terms of the size of supports of probability measures [1, 30], using the original form of Wasserstein distances is still not tractable in real training due to both the memory constraint and time constraint.
In more detail, the number of training samples is often huge, e.g., one million, and the dimension of data is also huge ,e.g., ten thousand.
Therefore, minibatch losses based on Wasserstein distances have been proposed [12, 40, 41].
The corresponding population form of these losses between two probability measures µ and ν is:
[Equation]
where m ≥ 1 is the mini-batch size and D is a Wasserstein metric.
In the generative model context [17], a stochastic gradient of the parameters of interest is utilized to update these parameters, namely,
[Equation]
where k is the number of mini-batches (is often set to 1), and ( X i , Y ϕ i ) is i.i.d sample from µ ⊗ m ⊗ ν ⊗ m ϕ .
The exchangeability between derivatives and expectation, and unbiasedness of the stochastic gradient are proven in [13].
Mini-batch losses are not distances; however, we can derive mini-batch energy distances from them [51].
Learning generative models via max-sliced Wasserstein: As we mentioned in Section 2.1, the maxsliced Wasserstein distance can overcome the curse of dimensionality of the Wasserstein distance and the issues of Monte Carlo samplings in the sliced Wasserstein distance.
Therefore, it is an appealing divergence for learning generative models.
By replacing the Wasserstein metric in equation (3), we arrive at the following formulation of the mini-batch max-sliced Wasserstein loss, which is given by:
[Equation]
Here, we can observe that each pair of mini-batch contains its own optimization problem of finding the "max" slice.
Placing this in the context of iterative training of generative models, we can foresee its expensive computation.
For a better understanding, we present an algorithm for training generative models with mini-batch max-sliced Wasserstein in Algorithm 2 in Appendix B.
In practice, there are some modifications of training generative models with mini-batch Max-SW for dealing with unknown metric space [11].
We defer the details of these modifications in Appendix D.
Amortized optimization: A natural question appears: "How can we avoid the nested loop in minibatch Max-SW due to several local optimization problems?".
In this paper, we propose a practical solution for this problem, which is known as amortized optimization [2].
In amortized optimization, instead of solving all optimization problems independently, an amortized model is trained to predict optimal solutions to all problems.
We now state the adapted definition of amortized models based on that in [52, 2]:
Definition 1 For each context variable x in the context space X , θ ⋆ ( x ) is the solution of the optimization problem θ ⋆ ( x ) = arg
[Equation]
θ ∈ Θ L ( θ, x ) , where Θ is the solution space.
A parametric function f ψ : X → Θ , where ψ ∈ Ψ , is called an amortized model if
[Equation]
The amortized model is trained by the amortized optimization objective which is defined as:
[Equation]
where p ( x ) is a probability measure on X which measures the "importance" of optimization problems.
The amortized model in Definition 1 is sometimes called a fully amortized model for a distinction with the other concept of semi amortized model [2].
The gap between the predicted solution and the optimal solution E x ∼ p ( x ) ||
f ψ ( x ) − θ ⋆ ( x ) || 2 is called the amortization gap.
However, understanding this gap depends on specific configurations of the objective L ( · , x ) , such as convexity and smoothness, which are often non-trivial to obtain in practice.
Amortized Sliced Wasserstein
In this section, we discuss an application of amortized optimization to the mini-batch max-sliced
Wasserstein.
In particular, we first formulate the approach into a novel family of mini-batch losses, named Amortized Sliced Wasserstein .
Each member of this family utilizes an amortized model for predicting informative slicing directions of mini-batch measures.
We then propose several useful amortized models in practice, including the linear model, the generalized linear model, and the non-linear model.
3.1 Amortized Sliced Wasserstein and Amortized Models
We extend the definition of the mini-batch max-sliced Wasserstein in Equation (5) with the usage of an amortized model to obtain the amortized sliced Wasserstein as follows.
Definition 2 Let p ≥ 1 , m ≥ 1 , and µ, ν are two probability measures in P ( R d ) .
Given an amortized model f ψ : R dm × R dm → S d − 1 where ψ ∈ Ψ , the amortized sliced Wasserstein between µ and ν is:
[Equation]
From the definition, we can see that the amortized model maps each pair of mini-batches to the optimal projecting direction on the unit hypersphere between two corresponding mini-batch probability measures.
We have the following result about the symmetry and positivity of A -SW.
Proposition 1 The amortized sliced Wasserstein losses are positive and symmetric.
However, they are not metrics since they do not satisfy the identity property, namely, A -SW ( µ, ν ) =
̸⇐⇒ µ = ν .
Proof of Proposition 1 is in Appendix A.1.
Our next result indicates that we can upper bound the amortized sliced Wasserstein in terms of mini-batch max-sliced Wasserstein.
Proposition 2 Assume that the space Ψ is a compact set and the function f ψ is continuous in terms of ψ .
Then, the amortized sliced Wasserstein are lower-bounds of the mini-batch max-sliced Wasserstein (Equation 5), i.e., A -SW ( µ, ν ) ≤ m-Max-SW ( µ, ν ) for all probability measures µ and ν .
Proof of Proposition 2 is in Appendix A.2.
Parametric forms of the amortized model: Now we define three types of amortized models that we will use in the experiments.
Definition
Given X, Y ∈ R dm , and the one-one "reshape" mapping T : R dm → R d × m , the linear amortized model is defined as:
[Equation]
where w 1 , w 2 ∈ R m , w 0 ∈ R d and ψ = ( w 0 , w 1 , w 2 ) .
In Definition 3, the assumption is that the optimal projecting direction lies on the subspace that is spanned by the basis { x 1 , . . .
, x m , y 1 , . . .
, y m , w 0 } where
X =
( x 1 , . .
x m ) and Y =
( y 1 , . . .
, y m ) .
The computational complexity of this function is O ((2 m + 1) d ) since those of the operators T ( X ) wand T ( Y ) w 2 are O ( md ) while adding the bias w 0 costs an additional computational complexity
O ( d ) .
The number of parameters in linear amortized model is 2 m + d .
To increase the expressiveness of the linear amortized model, we apply some (non-linear) mappings to the inputs X and Y , which results in the generalized linear amortized model as follows.
Definition
Given X, Y ∈ R dm , and the one-one "reshape" mapping T : R dm → R d × m , the generalized linear amortized model is defined as:
[Equation]
where w 1 , w 2 ∈ R m , w 0 ∈ R d , ψ 1 ∈ Ψ 1 , g ψ 1 : R dm → R dm and ψ = ( w 0 , w 1 , w 2 , ψ 1 ) .
In Definition 4, the assumption is that the optimal projecting direction lies on the subspace that is spanned by the basis { x ′ 1 , . . .
, x ′ m , y ′ 1 , . . .
, y ′ m , w 0 } where g ψ 1 ( X ) =
( x ′ 1 , . . .
, x ′ m ) and g ψ 1 ( Y ) =
( y ′ 1 , . . .
, y ′ m ) .
To specify, we let g ψ 1 ( X ) =
( W 2 σ ( W 1 x 1 ) + b 0 , . . .
, W 2 σ ( W 1 x m ) + b 0 ) , where σ ( · ) is the Sigmoid function, W 1 ∈ R d × d , W
2 ∈ R d × d , and b 0 ∈ R d .
Compared to the linear model, the generalized linear model needs additional computations for g ψ ( T ( X )) and g ψ ( T ( Y )) , which are at the order of O (2 m ( d 2 + d )) .
It is because we need to include the complexity for matrix multiplication, e.g., W 1 x 1 that costs O ( d 2 ) , for Sigmoid function that costs O ( d ) , and for adding bias b 0 that costs
O ( d ) .
Therefore, the total computational complexity of the function f ψ is O (4 md 2 + 6 md + d ) while the number of parameters is 2( m + d 2 + d ) .
We finally propose another amortized model where we instead consider some mapping on the function ω
+ T ( X ) ω
1 + T ( Y ) ω 2 in the linear amortized model so as to increase the approximation power of the function f ψ .
Definition
Given X, Y ∈ R dm , and the one-one "reshape" mapping T : R dm → R d × m , the non-linear amortized model is defined as:
[Equation]
where w 1 , w 2 ∈ R m , w 0 ∈ R d , ψ 2 ∈ Ψ 2 , h ψ 2 : R d → R d and ψ = ( w 0 , w 1 , w 2 , ψ 2 ) .
In Definition 5, the assumption is that the optimal projecting direction lies on the image of the function h ψ 2 ( · ) that maps from the subspace spanned by { x 1 , . . .
, x m , y 1 , . . .
, y m , w 0 } where
X =
( x 1 , . .
x m ) and Y =
( y 1 , . . .
, y m ) .
The computational complexity for h ψ 2 ( x ) =
W 4 σ ( W 3 x ))
+ bwhen x ∈ R d , W 3 ∈ R d × d , W 4 ∈ R d × d , and b 0 ∈ R
d is at the order of O (2( d 2 + d )) .
Therefore, the total computational complexity of the function f ψ is O (2 md + 2 d 2 + 3 d ) while the number of parameters is 2( m + d 2 + d ) .
Using amortized models in Definitions 3-5 leads to three amortized sliced Wasserstein losses, which are linear amortized sliced Wasserstein loss ( LA -SW), generalized linear amortized sliced Wasserstein loss ( GA -SW), and non-linear amortized sliced Wasserstein loss ( NA -SW) in turn.
Remark
The parametric forms in Definitions 3-5 are chosen as they are well-known choices for parametric functions.
There are still several other ways of parameterization that can be utilized in practice based on prior knowledge about data, e.g., we can use convolution operator for saving parameters or we can strengthen the dependence between samples via recursive functions.
We leave the design of these amortized models for future work.
3.2 Amortized Sliced Wasserstein Generative Models
Based on the amortized sliced Wasserstein losses, our objective function for training a generative model ν ϕ parametrized by ϕ ∈
Φ now becomes:
[Equation]
Since the above optimization forms a minimax problem, we can use an alternating stochastic gradient descent-ascent algorithm to solve it.
In particular, the stochastic gradients of ϕ and ψ can be estimated from mini-batches ( X 1 , Y ϕ, 1 ) , . . .
, ( X k , Y ϕ,k ) ∼ µ ⊗ m ⊗ ν ⊗ m ϕ as follows:
[Equation]
For more details, we present the procedure in Algorithm 3 in Appendix B.
Computational complexity:
From Algorithm 2 and Algorithm 3 in Appendix B, we can see that training with A -SW can escape the inner while-loop for finding the optimal projecting directions.
In each iteration of the global while-loop, the computational complexity of computing the minibatch Max-SW is O (2 kT 2 ( m log m + dm )) , which is composed by k mini-batches with T 2 loops of the projection to one-dimension operator which costs O (2 dm ) and the computation of the sliced
Wasserstein which costs O (2 m log m ) .
For the mini-batch sliced Wasserstein, the overall computational complexity is O (2 kL ( m log m + dm )) where L is the number of projections.
For LA -SW, the overall computation complexity is O (2 k ( m log m + 3 md + d )) where the extra complexity
O ((2 m + 1) d ) comes from the computation of f ψ ( · ) (see Section 3.1).
Similarly, the computational complexities of GA -SW and NA -SW are respectively O
(2 k ( m log m + 4 md
+ 7 md + d )) and
O (2 k ( m log m + 3 md + 2 d 2 + 3 d )) .
Projection Complexity:
Compared to the sliced Wasserstein, Max-SW reduces the space for projecting directions from O ( L ) to O (1) .
For LA -SW, GA -SW, and NA -SW, the projection complexity is also O (1) .
However, compared to d parameters of Max-SW, LA -SW needs 2 m + d parameters for creating the projecting directions while GA -SW and NA -SW respectively need
O (2( m + d 2 + d )) parameters for producing the directions (see Section 3.1).
Remark
The computational complexities and the projection complexities of GA -SW and NA -SW are based on the specific parameterization that we choose in Section 3.
We would like to recall that these complexities can be reduced by lighter parameterization as in the remark at the end of
Section 3.1.
Related Works
Generalized sliced Wasserstein [24] was introduced by changing the push-forward function from linear T θ ( x )
=
θ ⊤ x to non-linear T θ ( x )
=
g ( θ, x ) for some non-linear function g ( · , · ) .
To cope with the projection complexity of sliced Wasserstein, a biased approximation based on the concentration of
Gaussian projections was proposed in [37].
An implementation technique that utilizes both RAM and
GPUs’ memory for training sliced Wasserstein generative model was introduced in [27].
Augmenting the data to a higher-dimensional space for a better linear separation results in augmented sliced
Wasserstein [6].
Projected Robust Wasserstein (PRW) metrics appeared in [44] that finds the best orthogonal linear projecting operator onto d ′ > 1 dimensional space.
Riemannian optimization techniques for solving PRW were proposed in [28, 20].
We would like to recall that, amortized optimization techniques can be also applied to the case of PRW, max-K-sliced Wasserstein [9], sliced divergences [36], and might be applicable for sliced mutual information [16].
Statistical guarantees of training generative models with sliced Wasserstein were derived in [38].
Amortized optimization was first introduced in the form of amortized variational inference [23, 47].
Several techniques were proposed to improve the usage of amortized variational inference such as using meta sets in [58], using iterative amortized variational inference in [33], using regularization in [53].
Amortized inference was also applied into many applications such as probabilistic reasoning [15], probabilistic programming [48], and structural learning [5].
However, to the best of our knowledge, it is the first time that amortized optimization is used in the literature of optimal transport.
We refer to [2] for a tutorial about the amortized optimization.
Experiments
In this section, we focus on comparing A -SW generative models with SNGAN [35], the sliced
Wasserstein generator [11], and the max-sliced Wasserstein generator [10].
The parameterization of model distribution is based on the neural network architecture of SNGAN [35].
The detail of the training processes of all models is given in Appendix D. For datasets, we choose standard benchmarks such as CIFAR10 (32x32) [26], STL10 (96x96) [7], CelebA (64x64), and CelebAHQ (128x128) [32].
For quantitative comparison, we use the FID score [18] and the Inception score (IS) [50].
We also show some randomly generated images from different models for qualitative comparison.
We give full experimental results in Appendix E.
The detailed settings about architectures, hyperparameters, and evaluation of FID and IS are given in Appendix F.
We would like to recall that all losses that are used in this section are in their mini-batch version.
We first demonstrate the quality of using A -SW in the training generative model compared to the baseline SNGAN, and other mini-batch sliced Wasserstein variants.
Then, we investigate the convergence of generative models trained by different losses including the standard SNGAN’s loss, mini-batch SW, mini-batch Max-SW, and A -SW by looking at their FID scores and IS scores over training epochs of their best settings.
After that, we compare models qualitatively by showing their randomly generated images.
Finally, we report the training speed (number of training iterations per second) and the training memory (megabytes) of all settings of all training losses.
Summary of FID and IS scores: We show FID scores and IS scores of all models at the last training step on all datasets in Table 1.
For SW and Max-SW, we select the best setting of hyperparameters for each score.
In particular, we search for the best setting of the number of projections L ∈ { 1 , 100 , 1000 , 10000 } .
Also, we do a grid search on two hyperparameters of Max-SW, namely, the slice maximum number of iterations T 2 ∈ { 1 , 10 , 100 } and the slice learning rate η 2 ∈ { 0 .
The detailed FID scores and IS scores for all settings are reported in Table
[Figure]
models converge faster.
[Table]
in Appendix E.
For amortized models
[Table]
we fix the slice learning rate η 2 = 0 .
From Table
[Table]
the best amortized model provides lower FID scores and IS scores than SNGAN, SW
[Table]
and Max-SW on all datasets of multiple image resolutions.
We would like to recall that, SNGAN is reported to be better than WGAN [3] in [35].
Furthermore, the best generative models trained by A -SW are better than models trained with SNGAN, SW, and Max-SW.
Interestingly, the LA -SW performs consistently well compared to other members of A -SW.
Also, we observe that Max-SW performs worse than both A -SW and SW.
This might be because the local optimization of Max-SW gets stuck at some bad optimum.
However, we would like to recall that Max-SW is still better than SW with L = 1 (see
Table 3 in Appendix E).
It emphasizes the benefit of searching for a good direction for projecting.
FID and IS scores over training epochs: We show the values of FID scores and Inception scores over epochs on CIFAR10, CelebA, STL10, and CelebA-HQ in Figure 1.
According to the figures in Figure 1, we observe that using SW and A -SW helps the generative models converge faster than
SNGAN.
Moreover, FID lines of A -SW are usually under the lines of other losses and the IS lines of
A -SW are usually above the lines of others.
Therefore, A -SW losses including LA -SW, GA -SW, and NA -SW can improve the convergence of training generative models.
Generated images: We show generated images on CIFAR10, CelebA, STL10 from SNGAN, and
LA -SW in Figure 2 as a qualitative comparison.
The generated images on CelebAHQ and the generated images of Max-SW, GA -SW, and NA -SW are given in Appendix E. From these images, we observe that the quality of generated images is consistent with the FID scores and the IS scores.
Therefore, it reinforces the benefits of using A -SW to train generative models.
Again, we would like to recall that all generated images are completely random without cherry-picking.
Computational time and memory: We report the number of training iterations per second and the memory in megabytes (MB) in Table 2.
We would like to recall that reported numbers are under some errors due to the state of the computational device.
From the table, we see that LA -SW is comparable to Max-SW and SW ( L = 1) about the computational memory and the computational
[Figure]
[Table]
time.
More importantly, LA -SW is faster and consumes less memory than SW ( L ≥ 100 ) and
Max-SW ( T 2 ≥ 10 ).
Compared to SNGAN, SW variants increase the demand for memory and computation slightly.
From LA -SW to GA -SW and NA -SW, the computational time is slower slightly; however, we need between 800 to 2100 MB of memory in extra.
Again, the additional memory depends on the chosen parameterization (see Section 3).
From this table, we can see that using sliced Wasserstein models gives better generative quality than SNGAN but it also costs more computational time and memory.
Among sliced Wasserstein variants, LA -SW is the best option since it costs the least additional memory and time while it gives consistently good results.
We refer to
Section 3 for discussion of the time and projection complexities of A -SW.
Conclusion
We propose using amortized optimization for speeding up the training of generative models that are based on mini-batch sliced Wasserstein with projection optimization.
We introduce three types of amortized models, including the linear, generalized, and non-linear amortized models, for predicting optimal projecting directions between all pairs of mini-batch probability measures.
Moreover, using three types of amortized models leads to three corresponding mini-batch losses which are the linear amortized sliced Wasserstein, the generalized linear amortized sliced Wasserstein, and the non-linear amortized sliced Wasserstein.
We then show that these losses can improve the result of training deep generative models in both training speed and generative performance.
Acknowledgements
NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of
Machine Learning.
Supplement
Amortized for
Sliced Wasserstein Generative Models"
In this supplement, we first collect some proofs in Appendix A. We then present algorithms that are mentioned in the main text in Section B. Next, we discuss the training detail of generative models with different mini-batch losses in Appendix D. Moreover, we present detailed results on the deep generative model in
Appendix E. Next, we report the experimental settings including neural network architectures, and hyper-parameter choices in
Appendix F. Finally, we discuss the potential impacts of our works in Appendix G.
A Proofs
In this appendix, we provide proofs for main results in the main text.
A.1 Proof of Proposition
Recall that, the definition of A -SW ( µ, ν ) is as follows:
[Equation]
For the symmetric property of the amortized sliced Wasserstein, we have
[Equation]
where the second equality is because of the symmetry of Wasserstein distance, the third equality is due to the symmetry of f ψ ( X, Y )
(see forms of f ψ ( X, Y ) in Section 3).
The positiveness of A -SW comes directly from the non-negativity of the Wasserstein distance.
To prove that A -SW violates the identity, we use a counter example where µ = ν = 12 δ x 1 + 12 δ x 2 ( x 1 ̸ = x 2 ).
In this example, there exists a pair of mini-batches X =
( x 1 , x 1 ) and Y =
( x 2 , x 2 ) .
We choose f ψ ( X, Y ) =
x 1 + x 2 || x 1 + x 2 || 2 , then f ψ ( X, Y )
♯P X ̸
= f ψ ( X, Y ) ♯P Y which implies
W p ( f ψ ( X, Y ) ♯P X , f ψ ( X, Y ) ♯P Y ) > 0 .
Since A -SW defines on the maximum value of ψ ∈ Ψ ,
A -SW ( µ, ν ) ≥ W p ( f ψ ( X, Y ) ♯P X , f ψ ( X, Y ) ♯P Y ) > 0 .
A.2 Proof of Proposition
Since the function f ψ is continuous in terms of ψ , it indicates that the function
[Equation]
is continuous in terms of ψ .
Furthermore, as the parameter space Ψ is compact, there exist ψ ∗ ∈ arg
[Equation] [Equation]
( f
[Equation]
( X, Y
[Equation] [Equation]
ψ
[Equation]
X, Y ) ♯P Y )] .
Then, we have
[Equation]
As a consequence, we obtain the conclusion of the proposition.
B Algorithms
We collect algorithms that are mentioned in the main paper.
The projected gradient ascent algorithm for Max-SW is given in Algorithm 1.
After that, the training algorithm for generative models with mini-batch Max-SW is presented in Algorithm 2.
Finally, the algorithm for training generative models with A -SW is in Algorithm 3.
[Table]
Input: Probability measures: µ, ν , learning rate η , max number of iterations T .
Initialize θ while θ not converge or reach T do θ = θ + η · ∇ θ W p ( θ♯µ, θ♯ν ) θ
= θ ||
θ
|| 2 end while
Return: θ
Input: Data probability measure µ learning maximum number of iterations T slice maximum number iterations
Initialize ϕ , the model probability measure ν ϕ while ϕ not converge or reach T 1 do ∇ ϕ =
Sample ( X 1 , Y ϕ, 1 ) , . . .
, ( X k , Y ϕ,k ) ∼ µ ⊗ m ⊗ ν ⊗ m ϕ for i
= 1 to k do while θ not converge or reach T 2 do θ
[Equation]
) θ =
θ ||
θ
|| 2 end while
[Equation]
end for ϕ
[Equation]
end while
Return: ϕ, ν ϕ
C Amortized Projected Robust Wasserstein
We first recall the definition of projected robust Wasserstein (PRW) distance [44].
Given two probability measures µ, ν ∈ P p ( R d ) , the projected robust Wasserstein distance between µ and ν is
Input: Data probability measure µ , model learning rate η 1 , amortized learning rate η 2 , maximum number of iterations while ϕ, ψ not converge or reach T do ∇
[Equation]
Sample ( X 1 , Y
[Equation]
[Equation] [Equation]
( X
[Equation]
, Y
[Equation] [Equation]
) ∼
[Equation]
ϕ for i = 1 to k do
[Equation]
) end for ϕ = ϕ −
[Equation]
∇
[Equation]
ψ
= ψ + η 2 · ∇ ψ end while
Return: ϕ, ϕ defined as:
[Equation]
(14) where V k ( R d ) :=
{ U ∈ R d × k | U ⊤ U
= I k } is the Stefel Manifold.
PRW can be seen as the generalization of Max-SW since PRW with k = 1 is equivalent to Max-SW.
Similar to Max-SW, the optimization of PRW is solved by using projected gradient ascent.
The detailed of the algorithm is given in Algorithm 4.
We would like to recall that other methods of optimization have also been used to solved PRW such as Riemannian optimization [28], block coordinate descent [21].
However, in this paper, we consider the original and simplest method which is projected gradient ascent.
In deep learning and large-scale applications, the mini-batch loss version of PRW is used, that is defined as follow:
[Equation]
Amortized Projected Robust Wasserstein loss: We define Amortized Projected Rubust Wasserstein loss as follow:
Definition 6 Let p ≥ 1 , m ≥ 1 , and µ, ν are two probability measures in P ( R d ) .
Given an amortized model f ψ : R dm × R dm
→
V k ( R d ) where ψ ∈ Ψ , the amortized projected robust
Wasserstein between µ and ν is:
[Equation]
Similar to the case of A -SW,
A -PRW is symmetric, positive, and is a lowerbound of PRW.
Also,
A -PRW is not a metric since it does not satisfy the identity property.
Amortized models: Similar to the case of A -SW, we can derive linear model, generalized linear model, and non-linear amortized model.
The only change is that the model gives k output vectors instead of 1 vector.
Definition 7 Given X, Y ∈ R dm , and the one-one "reshape" maping T : R dm → R d × m , the linear projected amortized model is defined as:
[Equation]
(17) where W 1 , W 2 ∈ R m × k , W 0 ∈ R d × k , and
Proj V k
( R d ) return the Q matrix in QR decomposition.
The definitions of the generalized linear projected amortized model and non-linear projected amortized model are straight-forward from the definitions of generalized linear model and non-linear model in
A -SW.
Input: Probability measures: µ, ν , learning rate η , max number of iterations T .
Initialize U while U not converge or reach T do
U =
U + η · ∇ U W p ( U♯µ, U♯ν )
Q, R = QR ( U ) (QR decomposition)
U =
Q end while
Return: θ
D Training Generative Models
In this section, we review the parameterization of training losses of generative models.
Parametrization: We first discuss the parametrization of the model distribution ν ϕ .
In particular, ν ϕ is a pushforward probability measure that is created by pushing a unit multivariate Gaussian ( ϵ ) through a neural network G ϕ that maps from the realization of the noise to the data space.
The detail of the architecture of G ϕ is given in Appendix F. For training both SNGAN and generative models of SW, Max-SW, and A -SW, we need a second neural network T β that maps from data space to a single scalar.
The second neural network is called Discriminator in SNGAN or Feature encoder in the others.
However, the architecture of the second neural network is the same for all models (see
Appendix F).
For the better distinction between training objectives of SNGAN and the objectives of the others, we denote T β 1 is the sub neural network of T β that maps from the data space to a feature space (output of the last Resnet block), and T β 2 that maps from the feature space (image of T β 1 ) to a single scalar.
More precisely, T β
= T β 2 ◦ T β 1 .
Again, we specify T β 1 and T β 1 in Appendix F.
Training SNGAN:
Let µ is theta data probability measure, these two optimization problems are done alternatively in training SNGAN:
[Equation]
Training SW, Max-SW, and A -SW:
For training these models, we adapt the framework in [11] to
SNGAN, namely, we use these two objectives:
[Equation]
where the function ˜ T β 1 ,β 2 =
[ T β 1 ( x ) , T β 2 ( T β 1 ( x ))] which is the concatenation vector of T β 1 ( x ) and T β 2 ( T β 1 ( x )) , D is one of the mini-batch SW, the mini-batch Max-SW (see Equation 5), and A -SW (see Definition 2).
This technique is an application of metric learning since L p norm is not meaningful on the space of natural images.
This observation is mentioned in previous works [11, 14, 55, 39].
Other settings: The information about the mini-batch size, the learning rate, the optimizer, the number of iterations, and so on, are given in Appendix F.
E Full Experimental Results
Detailed FID scores and Inception scores: We first show the detailed FID scores and IS scores of all settings in Table 3.
From the table, we can see that the quality of the SW depends on the number of projections.
Namely, a higher number of projections often leads to better performance.
For Max-SW, we obverse that increasing the number of iterations T 2 might not lead to a lower FID score and a higher IS score.
The reason might be that the optimization gets stuck at some local optima.
For the choice of the learning rate η 2 , we do not see any superior setting for Max-SW.
Generated Images: We show generated images from SW, GA -SW, and NA -SW on CIFAR10,
CelebA, and STL10 in Figure 3.
The generated images from Max-SW on CIFAR10, CelebA, and
[Figure] [Figure]
STL10 are given in Figure 4.
The generated images from SNGAN and LA -SW are given in Figure 5.
The generated images from SW, Max-SW, GA -SW, and NA -SW on CelebA-HQ are presented in
Figure 6.
Again, we observe consistent quality results compared to the quantitative results of FID scores and Inception scores.
[Figure]
[Figure]
Results on Amortized PRW: We present the result of training generative models on CIfAR with mini-batch PRW loss and amortized PRW losses in Table 4.
For both PRW and A -PRW,
[Table]
[Table]
we set the learning rate for U is 0.01.
We choose the best result from PRW with the number of gradient updates in { 10 , 100 } while we only update the amortized model once for A -PRW.
We observe that A -PRW gives better FID and IS than PRW for all choice of k ∈ { 2 , 4 , 16 } .
Moreover, linear amortized projected model gives the best result among amortized models.
When k = 16 , the non-linear amortized model suffers from numerical error when using QR decomposition, hence, we cannot provide the result for it.
Overall, the result on PRW strengthen the claim that using amortized optimization for deep generative models with (sliced) projected Wasserstein can improve the result
F Experimental Settings
Neural network architectures: We present the neural network architectures on CIFAR10 in Table 5,
CelebA in Table 6, STL10 in Table 7, and CelebA-HQ in Table 8.
In summary, we use directly the architectures from https://github.com/GongXinyuu/sngan.pytorch .
Hyper-parameters: For CIFAR10, CelebA, and CelebA-HQ, we set the training iterations towhile we set it to 100000 in STL10.
We update T β 1 and T β 2 every iterations while we update G ϕ each 5 iterations.
The mini-batch size m is set to 128 on CIFAR10 and CelebA, is set to 32 on STL10, is set to 16 on CelebA-HQ.
The learning rate of G ϕ , T β 1 , and T β 2 is set to 0 . 0002 .
The optimizers for all optimization problems are Adam [22] with ( β 1 , β 2 ) =
FID scores and Inception scores: For these two scores, we calculate them based on 50000 random samples from trained models.
For FID scores, the statistics of datasets are calculated on all training samples.
[Table]
) 128 → 4 × 4 × 256 , dense linear
ResBlock up
ResBlock up ResBlock up
ResBlock up ResBlock up
BN, ReLU, 3 × 3 conv, 3 Tanh
[Table] [Table]
Input: ϵ ∈ R 128 ∼ N (
[Table]
1) 128 → 3 × 3 × 256 , dense , linear
ResBlock up ResBlock up
ResBlock up ResBlock up
ResBlock up
BN, ReLU, 3 × 3 conv, 3 Tanh
[Table]
ResBlock
[Table]
ResBlock
[Table]
ResBlock down
ResBlock
[Table] [Table]
[Table]
Input: ϵ ∈ R 128 ∼ N (
[Table]
1) 128 → 4 × 4 × 256 , dense , linear
ResBlock up
ResBlock up ResBlock up
ResBlock up ResBlock up
BN, ReLU, 3 × 3 conv, 3 Tanh
[Table]
ResBlock down
ResBlock down ResBlock down
ResBlock ResBlock
ResBlock
[Table]
ReLU
Global sum pooling 128 →
Spectral normalization
G Potential Impact and Limitations
Potential Impact: This work improves training generative models with sliced Wasserstein by using amortized optimization.
Moreover, amortized sliced Wasserstein losses can be applied to various applications such as generative models, domain adaptation, and approximate inference, adversarial attack, and so on.
Due to its widely used potential, it can be used as a component in some applications that do not have a good purpose.
For example, some examples are creating images of people without permission, attacking machine learning systems, and so on.
Limitations:
In the paper, we have not been able to investigate the amortization gaps of the proposed amortized models since the connection of the optima of Max-SW to the supports of two probability measures has not been well-understand yet.
Moreover, the design of amortized models requires more engineering to achieve better performance since there is no inductive bias for designing them at the moment.
The hardness in designing amortized models is that we need to trade-off between the performance and computational efficiency.
We will leave these questions to future work.