Amortized Projection Optimization for Sliced Wasserstein Generative Models
Anonymous Author(s)
AffiliationAddress
Abstract
Seeking informative projecting directions has been an important task in utilizingsliced Wasserstein distance in applications.
However, finding these directionsusually requires an iterative optimization procedure over the space of projectingdirections, which is computationally expensive.
Moreover, the computationalissue is even more severe in deep learning applications, where computing thedistance between two mini-batch probability measures is repeated several times.
This nested loop has been one of the main challenges that prevent the usage ofsliced Wasserstein distances based on good projections in practice.
To address thischallenge, we propose to utilize the learning-to-optimize technique or amortizedoptimization to predict the informative direction of any given two mini-batchprobability measures.
To the best of our knowledge, this is the first work that bridgesamortized optimization and sliced Wasserstein generative models.
In particular,we derive linear amortized models, generalized linear amortized models, andnon-linear amortized models which are corresponding to three types of novel minibatch losses, named amortized sliced Wasserstein .
We demonstrate the favorableperformance of the proposed sliced losses in deep generative modeling on standardbenchmark datasets.
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 generativemodels such as GANs [17], VAEs [22], and diffusion models [19, 52], achieve a significant quality ofsampling 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 datadistribution µ and the model distribution ν ϕ with ϕ ∈ Φ , parameter space of neural networks weights,namely, we solve for ˆ ϕ ∈ arg min ϕ ∈ Φ D ( µ, ν ϕ ) .
For example, Kullback–Leibler divergence is usedin VAEs and diffusion models, Jensen–Shannon divergence appears in GANs, and f-divergences areutilized in f-GANs [40].
Because of the complexity of the neural networks ϕ , closed-form optimalsolutions to these optimization problems are intractable.
Therefore, gradient-based methods and theirstochastic 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 ofthese models include Wasserstein GAN [3] with the dual form of Wasserstein-1 distance [43], and
OT-GANs [14, 48] 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 reportedto fail to approximate the Wasserstein distance [53] while OT-GAN suffers from high computationalcomplexity 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 entropicregularization [8] had been proposed to improve the computational complexity of approximatingoptimal transport to O ( m 2 ) [1, 29, 30, 28] and to remove the curse of dimensionality [33].
Howeverpractitioners usually choose to use the slicing (projecting version) of Wasserstein distance [55, 11, 24]due to a fast computational complexity O ( m log m ) and no curse of dimensionality O ( m − 1 / 2 ) .
Thedistance is known as sliced Wasserstein distance (SW) [4].
Sliced Wasserstein is defined as theexpected one-dimensional Wasserstein distance between two projected measures over the uniformdistribution over the unit sphere.
Due to the intractability of the expectation, Monte Carlo samplesfrom the uniform distribution over the unit sphere are used to approximate the distance.
The numberof 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 largenumber of projections
L relative to the dimension of data to perform well [24, 11].
Increasing L leadsto a linear increase in computational time and memory.
However, when data lie in a low dimensionalmanifold, several projections are redundant since they collapse projected measures to a Dirac-Deltameasure at zero.
There are some attempts to overcome that issue including sampling orthogonaldirections [46] and mapping the data to a lower-dimensional space [11].
The most popular approachis 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 anddeep learning in general, the optimization over the unit sphere requires iterative projected gradientdescent methods that can be computationally expensive.
In detail, each gradient-update of the modelparameters (neural networks) requires an additional loop for optimization of Max-SW between twomini-batch probability measures.
Therefore, we have two nested optimization loops: the global loop(optimizing model parameters) and the local loop (optimizing projection).
These optimization loopscan slow down the training considerably.
Contribution.
To overcome the issue, we propose to leverage learning to learn techniques ( amortizedoptimization ) to predict the optimal solution of the local projection optimization.
We bridge theliterature on amortized optimization and optimal transport by designing amortized models to solvethe iterative optimization procedure of finding optimal slices in the sliced Wasserstein generativemodel.
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:1.
First, we introduce a novel family of mini-batch sliced Wasserstein losses that utilize amortizedmodels to yield informative projecting directions, named amortized sliced Wasserstein ( A -SW).
We specify three types of amortized models: linear amortized, generalized linear amortized, andnon-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-linearamortized sliced Wasserstein ( NA -SW).
Moreover, we discuss some properties of A -SW lossesincluding metricity, complexities, and connection to mini-batch Max-SW.2.
We then introduce the application of A -SW in generative modeling.
Furthermore, we carry outextensive 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 amortizedoptimization in Section 2.
In Section 3, we propose amortized sliced Wasserstein distances andanalyze 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 make some conclusion.
Finally, we defer the proofs of key results and extra materials in the Appendices.
Notation.
For any d ≥ 2 , S d − 1 := { θ ∈ R d | || θ || 22 = 1 } denotes the d dimensional unit hyper-spherein L 2 norm, and U ( S d − 1 ) is the uniform measure over S d − 1 .
Moreover, δ denotes the Dirac deltafunction.
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 theset of transportation plans between µ and ν .
For m ≥ 1 , we denotes µ ⊗ m as the product measurewhich has the support is m random variables 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 probabilitymeasure 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 Wassersteindistance, and the max-sliced Wasserstein distance.
We then formulate generative models based on themax-sliced Wasserstein distances and review the amortized optimization problem and its applicationto the max-sliced Wasserstein generative models.
2.1 (Sliced)-Wasserstein Distances
We first define the Wasserstein- p distance [54, 42] 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
[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 overcomethe curse of dimensionality of Wasserstein distance in high dimension, the sliced Wassersteindistance [4] between µ and ν had been introduced and admitted the following formulation:
[Equation]
1 p .
For each θ ∈ S d − 1 , W pp ( θ♯µ, θ♯ν ) can be computedin linear time O ( n log n ) where n is the number of supports of µ and ν .
However, due to theintegration 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 drawuniform 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 largecompared to the dimension d .
It is not appealing since the computational complexity of SW is linearwith L .
To reduce projection complexity, max-sliced Wasserstein (Max-SW) is introduced [10] .
Inparticular, 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 asimple algorithm in Algorithm 1 in Appendix B.
In practice, practitioners often set a fixed number ofgradient 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.
Theproblem can be seen as the following optimization:
[Equation]
where D ( · , · ) can be Wasserstein distance or SW distance or Max-SW distance.
Despite the recentprogress on scaling up Wasserstein distance in terms of the size of supports of probability measures [1,29], using the original form of Wasserstein distances is still not tractable in real training due to boththe memory constraint and time constraint.
In more detail, the number of training samples isoften huge,e.g., one million, and the dimension of data is also large,e.g., ten thousand.
Therefore,mini-batch losses based on Wasserstein distances have been proposed [51, 12, 13, 38, 39].
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 toupdate 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 samples fromµ ⊗ m ⊗ ν ⊗ m ϕ .
The exchangeability between derivatives and expectation, and unbiasedness of thestochastic gradient are proven in [13].
Mini-batch losses are not distances; however, we can derivemini-batch energy distances from them [48].
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 andthe issues of Monte Carlo samplings in the sliced Wasserstein distance.
Therefore, it is an appealingdivergence for learning generative models.
By replacing the Wasserstein metric in equation (3), wearrive at the following formulation of the mini-batch max-sliced Wasserstein, which is given by:
[Equation]
Here, we can observe that each pair of mini-batch contains its own optimization problem of findingthe "max" slice.
Placing this in the context of iterative training of generative models, we can foreseeits expensive computation.
For a better understanding, we present an algorithm for training generativemodels with mini-batch max-sliced Wasserstein in Algorithm 2 in Appendix B.
In practice, thereare some modifications of training generative models with mini-batch Max-SW for dealing withunknown metric space [11].
We defer the details of these modifications in Appendix C.
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 practicalsolution 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 predictoptimal solutions to all problems.
We now state the adapted definition of amortized models based onthat in [49, 2]:
[-Deﬁnition 
-]Definition 1 For each context variable x in the context space X , θ ⋆ ( x ) is the solution of the optimization problem θ ⋆ ( x ) =
arg min θ ∈ Θ L ( θ, x ) , where Θ is the solution space.
A parametric functionf ψ : 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 distinctionwith the other concept of semi amortized model [2].
The gap between the predicted solution and theoptimal solution E x ∼ p ( x ) ||
f ψ ( x ) − θ ⋆ ( x ) || 2 is called the amortization gap.
However, understandingthis 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 forpredicting informative slicing directions of mini-batch measures.
We then propose several usefulamortized models in practice, including the linear model, the generalized linear model, and thenon-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 ofan amortized model to obtain the amortized sliced Wasserstein as follows.
Definition 2 Let p ≥ 1 , m ≥ 1 , and µ, ν are two probability measure in P ( R d ) .
Given an amortizedmodel 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 optimalprojecting direction on the unit hypersphere between two corresponding mini-batch probabilitymeasures.
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, theyare 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 theamortized sliced Wasserstein in terms of mini-batch max-sliced Wasserstein.
Proposition {+
+}
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 wewill use in the experiments.
Definition
Given X, Y ∈ R dm , and the one-one "reshape" mapping T : R dm → R d × m , the linearamortized 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 isspanned 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 ) w 1and 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) mappingsto 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 , thegeneralized linear amortized model is defined as:
[Equation]
(10) 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 thatis spanned by the basis { x ′ 1 , . . .
, x ′ m , y ′ 1 , . . .
, y ′ m } 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, thegeneralized linear model needs additional computations for g ψ ( T ( X )) and g ψ ( T ( Y )) , which are atthe 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 ofthe function f ψ .
Definition
Given X, Y ∈ R dm , and the one-one mapping T : R dm → R d × m , the non-linearamortized model is defined as:
[Equation]
11) 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 functionh ψ 2 ( · ) that maps from the subspace spanned by { x 1 , . . .
, x m , y 1 , . . .
y m } where X = ( x 1 , . . .
, x m )and Y =
( y 1 , . . .
, y m ) .
The computational complexity for h ψ 2 ( x ) =
W 4 σ ( W 3 x ))
+ b {+
 whenx ∈+} 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 totalcomputational complexity of the function f ψ is O (2 md + 2 d 2 + 3 d ) while the number of parametersis 2( m + d 2 + d ) .
Using amortized models in Definitions 3-5 leads to three corresponding amortized sliced Wasserstein, which are linear amortized sliced Wasserstein ( LA -SW), generalized linear amortized sliced Wasserstein ( GA -SW), and non-linear amortized sliced Wasserstein ( NA -SW) in turn.
Remark
The parametric forms in Definitions 3-5 are chosen as they are well-known choices forparametric functions.
There are still several other ways of parameterization that can be utilizedin practice based on prior knowledge about data, e.g., we can use convolution operator for savingparameters or we can strengthen the dependence between samples via recursive functions.
We leavethe 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 generativemodel ν ϕ parametrized by ϕ ∈ Φ now becomes:
[Equation]
Since the above optimization forms a minimax problem, we can use an alternating stochastic gradientdescent-ascent algorithm to solve it.
In particular, the stochastic gradients of ϕ and ψ can be estimatedfrom 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 thattraining 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 loopsof 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 computationalcomplexities 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 forprojecting directions from O ( L ) to O (1) .
For LA -SW, GA -SW, and NA -SW, the projectioncomplexity is also O (1) .
However, compared to d parameters of Max-SW, LA -SW needs 2 m + dparameters 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 -SWare based on the specific parameterization that we choose in Section 3.
We would like to recallthat these complexities can be reduced by lighter parameterization as in the remark at the end of
Section 3.1.
Related Works
Generalized sliced Wasserstein [23] was introduced by changing the push-forward function fromlinear T θ ( x )
= θ ⊤ x to non-linear T θ ( x ) =
g ( θ, x ) for some non-linear function g ( · , · ) .
To cope withthe projection complexity of sliced Wasserstein, a biased approximation based on the concentrationof Gaussian projections was proposed in [36].
An implementation technique that utilizes two typesof memories for training sliced Wasserstein generative model was introduced in [26].
Augmentingthe data to a higher-dimensional space for a better linear separation results in augmented sliced
Wasserstein [6].
Projected Robust Wasserstein (PRW) metrics appeared in [41] that finds the bestorthogonal linear projecting operator onto d ′ > 1 dimensional space.
Riemannian optimizationtechniques for solving PRW were proposed in [27, 20].
We would like to recall that, amortizedoptimization techniques can be also applied to the case of PRW, max-K-sliced Wasserstein [9], sliceddivergences [35], and might be applicable for sliced mutual information [16].
Statistical guaranteesof training generative models with sliced Wasserstein were derived in [37].
Amortized optimization was first introduced in the form of amortized variational inference [22, 44].
Several techniques were proposed to improve the usage of amortized variational inference such asusing meta sets in [56], using iterative amortized variational inference in [32], using regularizationin [50].
Amortized inference was also applied into many applications such as probabilistic reasoning [15], probabilistic programming [45], and structural learning [5].
However, to the best of ourknowledge, 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 [34], the sliced
Wasserstein generator [11], and the max-sliced Wasserstein generator [10].
The parameterization ofmodel distribution is based on the neural network architecture of SNGAN [34].
The detail of thetraining processes of all models is given in Appendix C. For datasets, we choose standard benchmarkssuch as CIFAR10 (32x32) [25], STL10 (96x96) [7], CelebA (64x64), and CelebAHQ (128x128) [31].
For quantitative comparison, we use the FID score [18] and the Inception score (IS) [47].
We alsoshow some randomly generated images from different models for qualitative comparison.
We givefull experimental results in Appendix D.
The detailed settings about architectures, hyperparameters,and evaluation of FID and IS are given in
Appendix E. We would like to recall that all losses that areused in this section are in their mini-batch version.
We first demonstrate the quality of using A -SW in the training generative model compared tothe baseline SNGAN, and other mini-batch sliced Wasserstein variants.
Then, we investigate theconvergence 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 overtraining epochs of their best settings.
After that, we compare models qualitatively by showing theirrandomly generated images.
Finally, we report the training speed (number of training iterations persecond) 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 lasttraining 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 3in Appendix D. For amortized models, we fix the slice learning rate η 2 = 0 .
From Table 1, the bestamortized model provides lower FID scores and IS scores than SNGAN, SW, and Max-SW on all
[Figure]
models converge faster.
[Table]
datasets of multiple image resolutions.
We would like to recall that, SNGAN is reported to be betterthan WGAN [3] in [34].
Furthermore, the best generative models trained by A -SW are better thanmodels trained with SNGAN, SW, and Max-SW.
Interestingly, the LA -SW performs consistentlywell compared to other members of A -SW.
Also, we observe that Max-SW performs worse thanboth A -SW and SW.
This might be because the local optimization of Max-SW gets stuck at somebad optimum.
However, we would like to recall that Max-SW is still better than SW with L = 1 (see
Table 3 in Appendix D).
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 scoresover epochs on CIFAR10, CelebA, STL10, and CelebA-HQ in Figure 1.
According to the figuresin 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 thegenerated images of Max-SW, GA -SW, and NA -SW are given in Appendix D. 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 liketo 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 thememory in megabytes (MB) in Table 2.
We would like to recall that reported numbers are undersome errors due to the state of the computational device.
From the table, we see that LA -SW iscomparable to Max-SW and SW ( L = 1) about the computational memory and the computationaltime.
More importantly, LA -SW is faster and consumes less memory than SW ( L ≥ 100 ) and
[Figure]
[Table]
Max-SW ( T 2 ≥ 10 ).
Compared to SNGAN, SW variants increase the demand for memory andcomputation slightly.
From LA -SW to GA -SW and NA -SW, the computational time is slowerslightly; however, we need between 800 to 2100 MB of memory in extra.
Again, the additionalmemory depends on the chosen parameterization (see Section 3).
From this table, we can see thatusing sliced Wasserstein models gives better generative quality than SNGAN but it also costs morecomputational time and memory.
Among sliced Wasserstein variants, LA -SW is the best option sinceit 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 arebased on mini-batch sliced Wasserstein with projection optimization.
We introduce three types ofamortized models, including the linear, generalized, and non-linear amortized models, for predictingoptimal projecting directions between all pairs of mini-batch probability measures.
Moreover, usingthree types of amortized models leads to three corresponding mini-batch losses which are the linearamortized sliced Wasserstein, the generalized linear amortized sliced Wasserstein, and the non-linearamortized sliced Wasserstein.
We then show that these losses can improve the result of training deepgenerative models in both training speed and generative performance.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
c) Did you discuss any potential negative societal impacts of your work?
[Yes]
d) Have you read the ethics review guidelines and ensured that your paper conforms
to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experi-
ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models ) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]