Amortized Projection Optimization for Sliced Wasserstein Generative Models
Authors: Khai Nguyen
Nhat Ho
Venue: NeurIPS 2022 Accept (2022)
Keywords: Sliced Wasserstein
Optimal Transport
Amortized Optimization
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 \textit{learning-to-optimize} technique or \textit{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 mini-batch losses, named \emph{amortized sliced Wasserstein}. We demonstrate the favorable performance of the proposed sliced losses in deep generative modeling on standard benchmark datasets.

Amortized Projection Optimization for Sliced Wasserstein Generative Models
Khai Nguyen
Nhat Ho
Department of Statistics and Data Sciences
Department of Statistics and Data Sciences The University of Texas at Austin Austin, TX 78712 minhnhat@utexas. edu
The University of Texas Austin Austin, TX 78712
khainb@utexas.cet
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 the main challenges that prevent the usage sliced Wasserstein distances based good projections in practice. To address this challenge, we propose to utilize the learning-to-on to-optimize technique or amortized optimization to predict the informative direction of any given two bini-batch probability measures. To the best 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 corresponding three types of novel mini- batch losses, named amortized sliced Wasserstein. We demonstrate the favorable performance of the proposed sliced losses in deep generative modeling on standard benchmark datasets
Introduction
Generative modeling is one of the most important tasks in machine learning and data science. Lever- aging the expressiveness of neural networks in parameterizing the model distribution, deep generative models such as GANs 17 VAEs [23], and diffusion models [I] 54 achieve a significant quality of sampling images. Despite differences in the way of modeling the model distribution. optimization ob- jectives of training generative models can be written as minimizing a discrepancy ((..)) between data distribution f and the model distribution Vs with & Q parameter space of neural networks weights, namely, we solve for è E arg mingee D(u, M.). For example, Kullback-Leibler divergence used in VAΣs and diffusion models, Jensen-Shannon divergence appears in GANs, and f-divergences are utilized in f-GANs
3 Because of the complexity of the neural networks d,, closed-form optimal solutions to these optimization problems 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 gener- ative models due their training stability, efficiency, and geometrically meaning. Examples of these models include Wasserstein GAN B with the dual form of Wasserstein-1 distance [46], and
OT-GANs [14 5II with the primal form of Wasserstein distance and with Sinkhomn divergence
'Code for the paper is published at AmortizedSW
www./OoT/TTO. Austin Data--aiateeneeeeeee Science-Group)
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

between mini- probability measures. Although these models considerably improve the gen- erative 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: ((m logm) and its curse of dimensionality: the sample complexity of (m -1/d) where m is the number of supports of two mini-batch measures.
The entropic regularization [ had been proposed to improve the computational complexity of and to remove the curse of dimensional- However practitioners usually choose to use the slicing (projecting version) of Wasserstein
approximating optimal transport to O(m²) 1 3C 31 29]
ity
R4
distance
25 2 due to a fast computational complexity O(m log and no curse of di- mensionality (( 1/2). The distance is known as sliced Wasserstein distance (SW) 4 Sliced
5
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 used to approximate the distance. The number of samples is often called the number of projections and is denoted as L.
From applications, practitioners observe that sliced Wasserstein distance requires sufficiently large number of projections I. relative to the dimension of data to perform well [25 11. Increasing 7. leads to 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 Dirac-Delta measure at zero, There are some attempts to overcome that issue including sampling orthogonal directions 4 and mapping the data to lower-dimensional space [ЛI. 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) [IO]. 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 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:
1. First, we introduce a novel family of mini-batch sliced Wasserstein losses that utilize amor- tized models to yield informative projecting directions, named amortized sliced Wasserstein losses (A-SW). We specify three types of amortized models: linear amortized, general- ized linear amortized, and non-linear amortized models that are corresponding three mini-batch losses: linear amortized sliced Wasserstein (C.A-SW), generalized linear amor- tized sliced Wasserstein (GA-SW), and non-linear amortized sliced Wasserstein (WW). Moreover, we discuss some properties of A-SW losses including metricity, complexities, and connection mini-batch Max-SW.
2.
We then introduce the application of .A-SW in generative modeling. Furthermore, we carry out extensive experiments on standard benchmark datasets including CIFARI0 10, CelebA STL10, and CelebAHQ to demonstrate the favorable performance of AA-SW in learning generative models. Finally, we measure the computational speed and memory of A-SW, mini-batch Max-SW, and bani-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 B we propose amortized sliced Wasserstein distances and analyze some of their theoretical properties. The discussion on related works is given in Section 4 Section 5coontanns the application of A-SW to generative models, qualitative experimental results, and quantitative experimental results on standard benchmarks. In Section 0 Finally, we defer the proofs of key results and extra materials to the Appendices.
we provide a conclusion.

---

Notation, For any d 2 2, 8 := 10 E R 1 = 1} denotes the d dimensional unit hyper-sphere in L2 norm, and (((S-1) ) is the uniform measure over g Moreover, denotes the Dirac delta function. For p 2 1, P.(P) is the set of all probability measures on R that has finite p-moments. For U, (R [[( v) := { E Pp [R* x IR") fgz ग (I. y)dz = v, fga ((, y)dy = μ is the set of transportation plans between I and v. For m 2 1, we denotes f y as the product measure which has the supports are the joint vector of m random variables that follows μ.. For vector X E IRd X == (21....,TTm) Px denotes the empirical measures Li=1 62,- We denote f1 as the push-forward probability measure of jt through the function Ta R² + R where To(x) = 0 x.
2 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)-Wasserstee Distances
We first define the Wasserstein-p distance
56 45 between two probability measures E PP((R))
and V E (IR) as follows: Wp((,, v) == inf xel((a))))) FRuxR || =ll,dd(<, v) # When d = 1,
the Wasserstein distance has closed form which is Wp (P, v) (oo |F (=) dz) where F₁. and F, are the cumulative distribution function (CDF) of μ. and v 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 SWp(H, v) (J Wz( #fv))d) $ For each 9 E g (2 Btv) can be computed in linear time Onn log n) where 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 expres- sion. To approximate the intractable expectation, Monte Carlo scheme is used, namely, we draw uniform samples 01....,0 U((d-1) from the unit sphere and obtain the following approximation:
SWp(H, v) =e E Wg #/)))) %
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 [[C] In particular, the max- max-sliced Wasserstein distance between μ and v is given by:
Max-SW v) == max (Ptju 0jv). 6534-1
(1)
To solve the optimization problem, projected gradient descent procedure is used. We present simple algorithm in Algorithm In practice, practitioners often set fixed number of gradient updates, e.g., T = 100.
2.2 Learning Generative Models with Max-Sliced Wasserstein and Amortized Optimization
We now provide application of (sliced)- Wasserrsen distances generative models settings. The problem can be seen as the following optimization:
min D(p, Us), Φ
where -(..)) :) 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 [I] 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, mini-
batch losses based on Wasserstein distances have been proposed 112 40 AT. The corresponding population form of these losses between two probability measures μ and is:
D(p,v) == Exyyr &m D(Px. Py),
(3)
3

---

Algorithm Max-sliced Wasserstein distance
Input: Probability measures: μl, v, learning rate 7) max number of iterations T.
Initialize
while not converge or reach T do
0 ++1 1 (0ijh 8tv)
|
end while Return:
Algorithm 2 Training generative models with mini-batch max-sliced Wasserstein loss
Input:
Data probability measure μ4, model learning rate 7/, slice learning rate 12,, model maximum number of iterations Ti, slice maximum number of iterations T2, number of mini-batches (is often set to 1).
Initialize Φ. the model probability measure V
while e not converge or reach 7i do
V 00
Sample (X1, Yo.1)...... (Xx: ko,k) ~ a @ &m for = 1 to do
while 0 not converge reach T2 do
0 = 0 + 72- VoW,(0:P5 0f
0 Talla
end while
V V + WOW,
end for Φ = Φ
71 V
end while Return: 0,
where 2 1 is the mini-batch size and 'D is Wasserstein metric.
In the generative model context [17, stochastic gradient of the parameters of interest utilized to update these parameters, namely,
1 - V.D(P Py,.); Ko
V D(D(
(4)
where & is the number of mini-batches (is often set to 1), and (X,, Yo,) is i.i.d sample from am @ 3m The exchangeability between derivatives and expectation, and unbiasedness of the stochastic gradient are proven in [3. 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 max- sliced 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, is an appealing divergence for learning generative models. By replacing the Wasserstein metric in equation S we arrive at the following formulation of the mini-batch max-sliced Wasserstein loss, which is given by:
m--Ma--SW v) = Ex.y~aemSuom
max DES
(,(0Pxx0:PP))
(5)
Here, we can observe that each pair 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 In practice, there are some modifications of training generative models with mini batch Max-SW for dealing with unknown metric space [ We defer the details of these modifications in Appendix C
Amortized optimization: natural question appears: "How can we avoid the nested loop in mini- batch Max-SW due to several local optimization problems?". In this paper, we propose 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 trained predict optimal solutions to all problems. We now state the adapted definition of amortized models based on
that in 52 2
Definition For each context variable the context space X, 0- (x) is the solution of the opti- mization problem 0* (x) arg mingge L(0, v), where & is the solution space. A parametric :x O where vb M, called an amortized model if
fa(x) (x), Vx X.
The amortized model is trained by the amortized optimization objective which is defined as:
min Bewppat) (()(())
where p(z) is a probability measure on X which measures the "importance" of optimization problems.
The amortized model in Definition is sometimes called 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 Ez~p(=)) (z) 0* (æ lll is called the amortization gap. However, understanding this gap depends on specific configurations of the objective Z)),)), such as convexity and smoothness, which are often non-trivial to obtain in practice.
3 Amortized Sliced Wasserstein
In this section, we discuss 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 amortized model for predicting informative slicing directions of mini-batch measures. 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 with the usage of an amortized model to obtain the amortized sliced Wasserstein as follows.
Definition Let p 1, m 2 1, and μ. are two probability measures P(TR)). Given an amortized
model fo Rdm x Rdm gd-1 where vp E the amortized sliced Wasserstein between y and v is:
A-SW(H, := max (xX,) Sy (f (X, NEW
(8)
From the definition, we can see that the amortized model maps each pair of mini-batcher to the optimal probability
projecting direction the unit hypersphere between two corresponding measures. We have the following result about the symmetry and positivity of A-SS.
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(u, v) 0 0
Proof of Proposition is in Appendix A.I 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 Vi is compact set and the function f is continuous in terms of of the mini-bated max-sliced Wasserstein (Equation 5 i.e., A-SW(p, v) & mn-Max-SW(pp, v) for all probability measures P and v.
Jb. Then. the amortized sliced Wasserstein are
Proof of Proposition is in Appendix A.2
Parametric forms the amortized model: Now we define three types of amortized models that we will use in the experiments.
5

---

Definition 3 Given X.Y E Rdd, and the one-one "reshape" mapping T Rdm - Rdxm, the linear amortized model is defined as:
wo + [[ + w2
(xX = 1|w0 + T(X) (X)w) + T(Y))al22
where W1, W2 E R wo E Rd and t = (wo, W1, w2).
In Definition S the assumption is that the optimal projecting direction lies on the subspace that is
spanned by the basis {.'....,, 2m, V1...., ym Wo} where X = (#1...., .mm) and Y = (y1 Bmm)
The computational complexity of this function is ( + 1)d) since those of the operators T T(X) and T(Y) w2 are ((md) while adding the bias wo costs an additional computational complexity (d). The number of parameters in linear amortized model is 2m + d.
To increase the expressiveness of the linear amortized model, we apply some (non-linean mappings to the inputs X and Y, which results in the generalized linear amortized model as follows.
Definition 4 Given X,Y E Rdm, and the one-one "reshape" mapping T Rdm - Rdxm, the generalized linear amortized model is defined as:
(X, == wo + Tig (X)) + T(o,, (Y))
(10)
1wo + T (X)) + (9w, (Y) )2)llll
where wi, w2 E R wp E Rdd, 1/1 E 1 9v1 Rdd + Rm and = (wo, w1, W2, vh))
In Definition 4 the assumption is that the optimal projecting direction lies on the subspace that is spanned by the basis (r'j...., Fn Vi..... Vn, (1/,-.., V/.....)/m). To specify, we let Det (X) (W2o(W1x1) the Sigmoid function, Wi Raxd W2 E Rdxd and by E Rd Compared to the linear model, the generalized linear model needs additional computations for go (T())) and go (T(Y), which are at the order of (2 +d). It is because we need to include the complexity for matrix multiplication, e.g., W12 that costs (dd)), for Sigmoid function that costs Odd), and for adding bias by that costs (d). Therefore, the total computational complexity of the function f is OAAmd + 6md + d) while the number of parameters is 2mm + d² 1 dd).
Wp where gd. (X) (x',....,'", )mm) and De, (Y)
ba,-... (Wi + bo), where (-)
We finally propose another amortized model where we instead consider some mapping on the function wi - T(X)xs + +(YY)) Juz in the linear amortized model as increase the approximation power of the function fo.
Definition 5 Given X,Y Y. E Rdmm and the one-one "reshape" mapping T Rdd _ Rdxmm the non-linear amortized model is defined as:
1ap (wo + T(X) + w(Y)
fo(X,Y) = 11/aaa (wo + +T(X)w; w + T(Y) )
(11)
where W1, wy E R", wo E Rd, ii2 E 22 hwg IRd + Rdd and v (wo, W1, W2, 422).
In Definition the assumption that the optimal projecting direction lies on the image of the function heg (c) that maps from the subspace spanned by {x1,..., Vl...., Vm wo) where X = (I),..., Tmm and = (yi,...,,,) Mm). The computational complexity for (x) = (4o(W3z)) by when x E Rdd W3 Rdd W4 E Raxdd and by Rdd at the order of (2 + d)). Therefore, the total computational complexity of the function fy is O(2dd + 2d + 3d) while the number of parameters is 2(m + d² + d).
Using amortized models in Definitions 3 leads three amortized sliced Wasserstein losses, which are linear amortized sliced Wasserstein loss (C.A-SW), generalized linear amortized sliced Wasserstein loss (GW--W), and non-linear amortized sliced Wasserstein loss (A in turn.
The parametric forms in Definitions 3 are chosen as they are well-known choices for parametric functions. There are still several other ways of parumeterization that can be utilized in practice based on prior knowledge about data, e.g., 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.
Remark 1

---

Algorithm 3 Training generative models with amortized sliced Wasserstein loss
Input:
Data probability measure J., model learning rate 7h, amortized learning rate 722, maximum number of iterations T, number of Initialize 6, the model probability measure Vá-
& (is often set to 1).
Initialize 1), the amortized model fg.
while é, ve converge or reach do
V 0; a
Sample (X1, Yo.1), (Kx, Yo.k) ~ p om & 2m
for
1 to k do
V
V Φ +
(X,,, Yo.)) Ya.))) V WwwFo(X, Yo.)) (X,, (r.))
V W
V
V
end for
Φ = Ø - 7/1
e 3 1/2 V
end while
Return: 0,4g
3.2 Amortized Sliced Wasserstein Generative Models
Based on the amortized sliced Wasserstein losses, our objective function for training generative model v parametrized by ó E 4 now becomes:
min max E (X.Ya)=p < NEW
== min max C(u, Vss v) <E NEV
Since the above optimization forms a minimax problem, we can use an alternating stochastic gradient descent-ascent algorithm solve it. In particular, the stochastic gradients of & and v can be estimated (X&, Yo.k) ~ fm @ us
from mini-batches (X1,
@m as follows:
I
V VoL(u, Vo, +)
V W Y.. Yo.D..)
(12)
k
i=l
(oL(u, v)
V Wo Yo. Yx. No.))
(13)
For more details, we present the procedure in Algorithm
Computational complexity: From Algorithm 2 and Algorithm 31 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 mini-batch Max-SW is (2kT22 (22kT2(r (m log + dm)), which is composed by k mini-batches with T2 loops of the projection to one-dimension operator which costs ((2m)) and the computation of the sliced Wasserstein which costs (22m log m). For the mini-batch sliced Wasserstein, the overall computational complexity is (2kLL(m (m log m + dm)) where L the number of projections. For LA-SS, the overall computation complexity is O(2k(m log + 3rd + d)) where the extra complexity O(2m 1)d) comes from the computation of (-) (see Section 3.I. Similarly, the computational complexities GAA-SW and N-4-SW are respectively (22k(m (m log m 4md (ond2+7md+d))) d)) and O(2k (m log m-+3md+ +2 +3)).
Projection Complexity: Compared to the sliced Wasserstein, Max-SW reduces the space for projecting directions from ((L) O(1). For L.A-SW, GA-SW, and AA-SS, the projection complexity is also (1). However, compared to d parameters of Max-SW, CAA-SW needs 2m + d parameters for creating the projecting directions while GA-SWW and NAA-SW respectively need (22(m + d2 + d)) parameters for producing the directions (see Section 3.1
Remark The computational complexities and the projection complexities of GA-SW and N N-SW are based on the specific parameterization that we choose in Section 3 We would like recall that these complexities can reduced by lighter parameterization as the remark at the end of
Section 3
7

---

Max
Pre0Pii Sore
11 Coore
AX.50
Min 507
Man
A.--W
Epecta
Epasha
Sperra
37.30
STL.3
Caloo-19
54
50
044 04
01-88
SYGA
Figure 1: FID scores and IS scores over epochs of different training losses on datasets. We observe that members of A-SW usually help the generative models converge faster.
Related Works
introduced by changing the push-forward function from linear (x) = f x to non-linear (T( [ (0, 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 [7] Augmenting the data to nigher-dimensional space for better linear separation results in augmented sliced Wasserstein (6]. Projected Robust Wasserstein (PRW) metries appeared in 44 that finds the best orthogonal linear projecting operator onto d' > 1 dimensional space. Riemannian optimization
Generalized sliced Wasserstein
24
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 [[6]. Statistical guarantees of training generative models with sliced Wasserstein were derived in 38
Amortized optimization was first introduced the form of amortized variational inference Several techniques were proposed to improve the usage of amortized variational inference such as using meta sets
58 using iterative amortized variational inference {3 using regularization Amortized inference was also applied into many applications such as probabilistic reason-
in [53
ing 15 probabilistic programming [H8]. 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 [2] for tutorial about the amortized optimization.
Experiments
the sliced Wasserstein generator [ and the max-sliced Wasserstein generator [TO] The parameterization of model distribution is based on the neural network architecture of SNGAN 35
In this section, we focus on comparing A-SW generative models with SNGAN K5
The detail of the training processes of all models is given Appendix C For datasets, we choose standard benchmarks such CIFARIO (32x32) 20). STL10 (9xx96) 71. CelebA (64x64), and CelebAHQ (128x 32 For quantitative comparison, we use the FID score [&] 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 D The detailed settings about architectures, hyperparameters, and evaluation of FID and IS are given in Appendix E used this section are in their mini-ban batch version.
We would like to recall that all losses that
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,

---

Table 1: Summary of FID and IS scores of methods on CIFARIO (32x32), CelebA (64x64), STL10 (9xx96), and CelebA-HQ (128x128). We observe that A-SW losses provide the best results among all the training losses.
CIFARIO (32x32)
CelebA (64x64)
STL10 (96x96)
CelebA-HQ (128x128)
Method
FID ())
IS (1)
FID (L)
IS (t)
FID (L)
IS (1)
FID ())
IS (1)
SNGAN
17.09
8.07
12.41
2.61
59.48
9.29
19.25
2.32
SW
14.25=#0.84
8.12:00.7
10.45
2.70
56.32
10.37
16.17
2.65
Max-SW
31.03++3.02
6.67+-0.7
11.28
2.60
77.40
9.46
29.50
2.36
C.A-SW (ours)
13,21+00.99
8.19+0.03
9.82
2.72
52.08
10.52
14.94
2.50
G.A-SW (ours)
13.64:0.11
8.22
9.21
2.78
53,80
10.40
18.97
2.34
NA-SW (ours)
14.22+00.51
8.29+0.08
8.91
2.82
53.90
10.14
15.17
2.72
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 the last training step on all datasets in Table For SW and Max-SW, we select the best setting of hyper- parameters for each score. In particular, we search for the best setting of the number of projec- tions L E {1, 100, 1000, 10000}. Also, we do grid search on two hyperparameters of SW, namely, the slice maximum number of iterations T2 {1, 10, 100} and the slice learning rate 122 E (0.001. 0.01, 0.13. The detailed FID scores and IS scores for all settings are reported in Table 3 in Appendix D For amortized models, we fix the slice learning rate 122- = 0.01. From Table the best amortized model provides lower FID scores and IS scores than SNGAN, SW, and Max-SW on all datasets of multiple image resolutions. We would like recall that, SNGAN is reported to be better than WGAN 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- is still better than SW with L 1 (see Table 3 in Appendix D 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 CIFARIO, CelebA, STL10, and CelebA-HQ Figure | According to the figures in Figure 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 CAA-SW, GA-SW, and N can improve the convergence of training generative models.
Generated images: We show generated images on CIFARIO, CelebA, STLIO from SNGAN, and CA-SW in Figure 2 as a qualitative comparison. The generated images on CelebAHQ and the generated images of Max-SW, GAA-SW, and A-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 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 LAASSW is comparable to Max-SW and SW (L =1) about the computational memory and the computational time. More importantly, LA-SSW is faster and consumes less memory than SW (L 2 100) and Max-SW (T, 10). Compared to SNGAN, SW variants increase the demand for memory and computation slightly. From L.4-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 B 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, LAASSW is the best option since costs the least additional memory and time while gives consistently good results. We refer to Section for discussion of the time and projection complexities of A-SW.

---

SNGAN (CIFAR)
SNGAN (CelebA)
SNGAN (STLIO)
LAA-WW (CIFAR)
L.A-SW (CelebA)
LA-SW (STL10)
Figure 2:
Random generated images of SNGAN and L.A-SW from CIFARIO, CelebA, and STL10.
Table 2:
Computational time memory methods (in iterations per second and megabytes (MB)).
CIFARIO (32x32)
CelebA (64x64)
ST (96x96)
CelebA-HQ (128x128)
Method
Iters/s (*)
Mom (L)
Iters/s (t)
Mem ())
Iters's (1)
Mcm ())
Iters/s (1)
Mcm (J)
SNGAN (baseline)
19.97
1740
6,31
6713
9.33
3866
10.41
3459
SW (L=1)
18.73
2078
6.17
8011
9.31
4597
10.25
4111
SW (L=100)
18.42
2093
6.15
8015
9.11
4609
10.17
4120
SW (L=1000)
14.96
2112
6.13
8047
9.03
4616
9.63
4143
SW
(L=10000)
5.84
2421
4.21
8353
6.50
4780
5.17
4428
Max-SW (f}=1)
18.61
2078
6.17
8011
9.23
4597
10.22
4111
Max-SW (7:=10)
18.16
2078
6.15
801
9.17
4597
10.16
4111
Max-SW (2²=100)
13.47
5,78
8011
4597
CA-SW (ours)
18.58
2086
6.17
8021
9.23
4600
10.19
4115
GA-SW (ours)
17.27
6.07
10083
5251
NAWSW (ours)
17.67
4134
6.13
10068
9.11
5249
10.15
6152
6
Conclusion
We propose using amortized optimization for speeding up the training generative models that are based 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.
10

---

References
[1] J. Altschuler, J. Niles-Weed, and Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. In Advances in Neural Information Processing Systems, pages 1964-1974, 2017.
[2] B. Amos. Tutorial on amortized optimization for learning to optimize over continuous domains. arXiv preprint arXiv:2202. 00665, 2022.
[3] M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pages 214-223, 2017.
[4] N. Bonneel, J. Rabin, G. Peyré, and H. Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision 1(51):22-45, 2015.
[5] K.-W. Chang, S. Upadhyay, G. Kundu, and D. Roth. Structural learning with amortized inference. In Twenty-Ninth AAAI Conference on Artificial Intelligence, 2015.
[6] X. Chen, Y. Yang, and Y. Li. Augmented sliced Wasserstein distances. International Conference on Learning Representations, 2022.
[7] A. Coates, A. Ng, and H. Lee. An analysis of single-layer networks in unsupervised feature the learning. Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 215-223. JMLR Workshop and Conference Proceedings, 2011.
[8] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, pages 2292-2300, 2013.
[9] B. Dai and U. Seljak. Sliced iterative normalizing flows. In International Conference on Machine Learning, pages 2352-2364. PMLR, 2021.
[10] I. Deshpande, Y.-T. Hu, R. Sun, Pyrros, N. Siddiqui, S. Koyejo, Z. Zhao, Forsyth, and A. G. Schwing. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 10648-10656, 2019.
[11] I. Deshpande, Z. Zhang, and A. G. Schwing. Generative modeling using the sliced Wasserstein distance. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 3483-3491, 2018.
[12] K. Fatras, Y. Zine, R. Flamary, R. Gribonval, and N. Courty. Learning with minibatch Wasser- stein: asymptotic and gradient properties. In AISTATS 2020-23nd International Conference on Artificial Intelligence and Statistics, volume 108, pages 1-20, 2020.
[13] K. Fatras, Y. Zine, S. Majewski, R. Flamary, R. Gribonval, and N. Courty. Minibatch optimal transport distances; analysis and applications. arXiv preprint arXiv:2101.01792. 2021.
[14] A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608-1617. PMLR, 2018.
[15] S. Gershman and N. Goodman. Amortized inference in probabilistic reasoning. In Proceedings of the Annual Meeting of the Cognitive Science Society, volume 36, 2014.
[16] Z. Goldfeld and Greenewald. Sliced mutual information: A scalable measure of statistical dependence. Advances in Neural Information Processing Systems, 34, 2021.
[17] I. Goodfellow, J. Abadet- M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, pages 2672-2680, 2014.
[18] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. GANs trained by two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems, pages 6626-6637, 2017.
[19] Ho, Jain, and Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-683 18844006851, 2020.
[20] M. Huang, S. Ma, and I. Lai. A Riemannian block coordinate descent method for computing the projection robust Wasserstein distance. In International Conference on Machine Learning, pages 4446-4455 PMLR, 2021.

---

[21] M. Huang, S. Ma, and L. Lai. A riemannian block coordinate descent method for computing the projection robust wasserstein distance. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 4446 4455. PMLR, 18-24 Jul 2021.
[22] D. P. Kingma and Ba. Adam: A method for stochastic optimization. arXiv preprint
WrXiv:1412.69888 2014.
[23] D. P. Kingma
and and M. Welling, Auto-encoding variational bayes. arXiv preprint 13:
2013.
[24] S. Kolouri, K. Nadjahi, U. Simsekli, R. Badcau, and G. Rohde. Generalized sliced Wasserstein distances. In Advances in Neural Information Processing Systems, pages 261-272, 2019.
[25] S. Kolouri, P. F., Pope, C. E. Martin, and K. Rohde. Sliced Wasserstein auto-encoders. In International Conference on Learning Representations, 2018.
[26] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
[27] J. Lezama, W. Chen, and Q. Qiu. Run-sort-rerun. Escaping batch size limitations in sliced Wasserstein generative models. In International Conference on Machine Learning, pages 6275-6285. PMLR, 2021.
[28] Lin, Fan, N. Ho, M. Cuturi, and M. Jordan. Projection robust Wasserstein distance and Riemannian optimization. Advances in Neural Information Processing Systems, 33:9383-9397, 2020.
[29] T. Lin, N. Ho, X. Chen, M. Cuturi, and M. I. Jordan. Fixed-support Wasserstein barycenters: Computational hardness and fast algorithm. In NeurIPS, pages 5368-55880 -3888, 2020.
[30] T. Lin, N. Ho, and M. Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International Conference on Machine Learning, pages 3982-3991, 2019.
[31] T. Lin, N. Ho, and M. I. Jordan. On the efficiency of the Sinkhorn and Greenkhorn algorithms and their acceleration for optimal transport. ArXiv Preprint: 1906.01437, 2019.
[32] Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
[33] J. Marino, Y. Yue, and S. Mandt. Iterative amortized inference. In International Conference on Machine Learning, pages
PMLR, 2018.
[34] G. Mena and J. Weed. Statistical bounds for entropic optimal transport: sample complexity and the central limit theorem. In Advances in Neural Information Processing Systems, 2019.
[35] T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint arXiv:
2018.
[36] K. Nadjahi, A. Durmus, L. Chizat, S. Kolouri, S. Shahrampour, and U. Simsekli. Statistical and topological properties sliced probability divergences. Advances in Neural Information Processing Systems, 33:20802-2088 12, 2020.
[37] K. Nadjahi, A. Durmus, P. E. Jacob, R. Badeau, and U. Simsekli. Fast approximation of the distance using concentration of random projections. Advances in Neural Information Processing Systems, 34, 2021.
[38] K. Nadjahi, A. Durmus, U. Simsekli, and R. Badeau. Asymptotic guarantees for learning generative models with the Siiced-Wassersteii distance. Advances in Neural Information Processing Systems, 32, 2019.
[39] K. Nguyen, N. Ho, T. Pham, and H. Bui. Distributional dliced-asseeree and applications to generative modeling. In International Conference on Learning Representations, 2021.
[40] K. Nguyen. D. Nguyen, Q. Nguyen, T. Pham, H. Bui, D. Phung, T. Le, and N. Ho. On transportation of mini-matches batches: A hierarchical approach. In Proceedings of the 39th International Conference Machine Learning. 2022.
[41] K. Nguyen, D. Nguyen, T. Pham, and N. Ho. Improving mini-batch optimal transport via partial transportation. Proceedings of the 39th International Conference on Machine Learning, 2022.
12

---

[42] K. Nguyen, S. Nguyen, N. Ho, T. Pham, and H. Bui. Improving relational regularized au- toencoders with spherical sliced fused Gromov-Wase Wasserstein. In International Conference on Learning Representations, 2021.
[43] S. Nowozin, B. Cseke, and R. Tomioka. f-gan: Training generative neural samplers using variational divergence minimization. Advances in Neural Information Processing Systems, 29, 2016.
[44] F.PP. Paty and M. Cuturi. Subspace robust Wasserstein distances. In International Conference on Machine Learning, pages 5072-5081, 2019.
[45] G. Peyré and M. Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends0 in Machine Learning, (56-)))355-607 2019.
[46] G. Peyré and M. Cuturi. Computational optimal transport, 2020.
[47] D. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference Machine Learning, pages 1278-1286. PMLR, 2014.
[48] D. Ritchie, Horsfall, and N. D. Goodman. Deep amortized inference for probabilistic programs. arXiv preprint arXiv:1610.0573:: 2016.
[49] M. Rowland, J. Hron, Y. Tang, K. Choromanski, T. Sarlos, and A. Weller. Orthogonal estimation of Wasserstein distances. The 22nd International Conference on Artificial Intelligence and Statistics, pages 186-195. PMLR, 2019.
[50] T. Salimans, I. Goodfellow, W. Zaremba, V. Cheung, A. Radford, and X. Chen. Improved techniques for training GANs. Advances in Neural Information Processing Systems, 29, 2016.
[51] T. Salimans, H. Zhang, A. Radford, and D. Metaxas. Improving GANs using optimal transport. In International Conference on Learning Representations, 2018.
[52] R.
Shu.
Amortized optimization
http://ruishu..oo201177/
amortizad-optiiizationn Personal Blog, 2017.
[53] R. Shu, H. H. Bui, S. Zhao, M. J. Kochenderfer, and S. Ermon. Amortized inference regulariza- tion. Advances in Neural Information Processing Systems, 31, 2018.
[54] Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[55] Stanczuk, C. Etmann, L. M. Kreusser, and C.-B. Schönlieb. Wasserstein GANs work because they fail (to approximate the Wasserstein distance). arXiv preprint 10X.v:2103.01678 2021.
[56] C. Villani. Optimal transport: Old and New. Springer, 2008.
[57] J. Wu, Z. Huang, D. Acharya, W. Li, J. Thoma, D. P. Paudel, and L. V. Gool. Sliced Wasserstein generative models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3713-3722, 2019.
[58] M. Wu, K. Choi, N. Goodman, and S. Ermon. Meta-amortized variational inference and learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 6404 6412, 2020.
Checklist
1. 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]
2. 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]
13

---

3. 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 URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (c.g., with respect 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, cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) curating/releasing 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 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 using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "ynS7fFW-l1s", "reviewer": "Reviewer_nJDk", "summary": "In the paper, the authors investigate the application of amortized optimization procedures to find the most distinguishing direction that can be used to calculate the sliced Wasserstein distance between two empirical latent distributions. This leads to the construction of three variants of the amortized sliced Wasserstein model: linear amortized ($\\mathcal{LA}$-SW), generalized linear amortized ($\\mathcal{GA}$-SW), and non-linear amortized ($\\mathcal{NA}$-SW). The authors conduct experiments on CIFAR10, CelebA(-HQ), and  STL10 datasets, which prove the performance of their methods compared to SNGAN and (max-)SW models. They also analyze the computational complexity of the proposed solutions.", "questions": "My final opinion may depend on the response to issues (1) and (2).\n\n(1)  I think the presentation (including the English of the paper) needs some improvements. For example, the authors quite often use simplified phrases such as (amortized, sliced, max) Wasserstein, to name the models, distances, or losses (without adding these words, what I find inappropriate), e.g.:\n\nl. 43: Sliced Wasserstein is defined $\\rightarrow$ Sliced Wasserstein distance is defined,\n\nl. 69: named amortized sliced Wasserstein $\\rightarrow$ named amortized sliced Wasserstein losses,\n\nl. 141: mini-batch max-sliced Wasserstein $\\rightarrow$ mini-batch max-sliced Wasserstein loss,\n\nl. 222-223: leads to three corresponding amortized sliced Wasserstein $\\rightarrow$ leads to three corresponding amortized sliced Wasserstein models.\n\n(2) Do the sets of vectors listed in l. 192, 202, and 217 indeed form bases of linear subspaces in $\\mathbb{R}^d$, or do only span these spaces (without independence), and why does $w_0$ appear only one time? Moreover, what do you exactly mean by writing \"one-one \"reshape\" mapping\"?\n\n(3) Minor issues:\n\nl. 27: $\\phi$ is a parameter of NN (not NN),\n\nl. 23 and 30: $\\mathcal{D}(\\cdot,\\cdot)$ denotes an arbitrary discrepancy or an OT distance?,\n\nl. 38: what is $d$?\n\nl. 38: needs precising: I hope $m$ means the size of the sum of supports of two mini-batch measures,\n\nl. 91-92:  a language issue in this sentence (\"as the product measure which has the support is $m$ random variables follows $\\mu$\"?),\n\nl. 127: huge/large,e.g., $\\rightarrow$ huge/large, e.g., (missing gap),\n\nl. 133: samples $\\rightarrow$ sample,\n\nl. 172 and 184: equation (5)/5 $\\rightarrow$ Equation (5),\n\nl. 174: measure $\\rightarrow$ measures,\n\nl. 213: the one-one mapping $\\rightarrow$ the one-one \"reshape\" mapping,\n\nl. 235: $\\ldots$ $\\rightarrow$ $, \\ldots,$.", "limitations": "I would recommend including a separate paragraph to describe limitations (as it was done for potential impact, see Appendix F).", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposed method seems to be novel, well-motivated, and with sufficient theoretical background. Also, the experimental results show significant improvement over the SOTA techniques. Hence, I think this paper may be interesting for the ML community. However, I found some presentation issues (see below for the details).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "fcOM4NJ9mQY", "reviewer": "Reviewer_3bjP", "summary": "This work proposes to use amortized optimization to predict the optimal projection direction which can be used for learning Wasserstein generative models. The authors derived a family of amortized models which can be used for the framework. Empirical results show that the proposed approach is able to achieve good performance on benchmark datasets.", "questions": "1. Proposition 1 mentions that $A-SW(\\mu, \\nu)=0$ does not imply $\\mu=\\nu$. Would the proposed method still guarantee convergence to the right distribution?\n2. As mentioned in Line 233, the optimization forms a min-max problem. Does it cause training instabilities?\n3. Does the batch size affect the performance? In other words, does the proposed approach require significantly larger batch sizes to achieve good performance?\n ", "limitations": "The authors have not adequately addressed the limitations and potential negative societal impact of their work. It would be good to add more discussions.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n1. The paper is relatively well-written.\n2. The experiments are detailed.\n3. The proposed approach could be interesting for the community. However, learning the projection direction also seems straightforward. \n\n**Weakness**\n1. The technical novelty could be limited. The proposed approach seems like a direct application of amortized inference to sliced Wasserstein generative models. There are also no extra theoretical insights/analyses provided.\n2. The amortized model considered in the paper could be restrictive. It would be good to perform theoretical analyses on how powerful the model family is.\n3. Empirical results are good but not impressive. The improvements in Table 1 are marginal. In Table 2, the memory cost for the proposed approach is also larger than the baselines. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "cspeNTOajs", "reviewer": "Reviewer_hTkP", "summary": "Modern generative models aim at fitting the probability distribution of the observed data. To this end, a common strategy consists in minimizing a chosen divergence between the observed empirical distribution and a parametric distribution, over the set of parameters. In recent years, models based on the minimization of optimal transport (OT) metrics have attracted significant interest. However, the standard OT metric, namely the Wasserstein distance, suffers from an expensive computational cost and sample complexity in large-scale settings (i.e. when the compared distributions are high-dimensional or supported on a large number of samples). For that reason, using practical alternatives to the traditional Wasserstein distance is more convenient in generative modeling. \n\nDue to its lower computational complexity and favorable theoretical properties, the sliced Wasserstein distance has been successfully integrated within the generative modeling framework. Nevertheless, the sliced Wasserstein distance (SW) is defined as an expectation which is intractable in general. In practice, this distance is thus traditionally approximated with a simple Monte Carlo algorithm: the expectation is replaced by a finite sum over L projection directions. Computing the sliced Wasserstein distance then scales linearly with L, meaning that improving the accuracy of the Monte Carlo estimator (by increasing L) leads to higher computational requirements. This is a limitation of SW-based generative algorithms, since they need to compute SW at each training iteration. One solution is to pick \"more informative\" samples in the Monte Carlo strategy, which motivated the formulation of maximum sliced Wasserstein (max-SW): the expectation is replaced by a maximum operator.\n\nIn this paper, the authors argue that the mini-batch formulation of max-SW (mini-batch max-SW, equation (5)) can be computationally expensive when used in generative modeling, and they develop alternative pseudo-metrics inspired by amortized optimization to overcome this problem. Their main contributions are summarized below.\n\n1) This paper introduces a new family of pseudo-metrics, coined \"amortized sliced Wasserstein\" (A-SW), by reformulating the mini-batch max-SW using amortized optimization (Definition 2). In other words, the optimal projection directions found in mini-batch max-SW are approximated with a family of parametric functions (the amortized model).\n\n2) The authors prove that amortized SW are positive, symmetric (Proposition 1) and lower than mini-batch max-SW (Proposition 1).\n\n3) Three instances of amortized SW are proposed by specifying the amortized model (Definitions 3, 4 and 5) and an analysis of their computational complexity as compared to mini-batch-max-SW is provided. \n\n4) Finally, GANs based on A-SW are proposed and compared against SNGAN and (max-)SW-based generators in image generation in terms of training speed, memory and quality of results, on 4 standard datasets (Section 5). According to the empirical results, A-SW is able to produce images of higher quality, at the price of a higher execution time and memory due to the optimization of the amortized model (depending on its parameterization).", "questions": "To address my concerns, I encourage the authors to:\n- Clarify the motivation and take-home message of the paper: A-SW does not seem to provide computational advantages over mini-batch max-SW but instead, provides more accurate results ; this might be due to the design of the amortized models, which constraints the shape of the optimal projections, thus regularizes the underlying optimization problem.\n- Add a more extensive discussion on the amortization gap to better motivate the proposed amortized models and the training procedure.\n- Further illustrate the empirical superior performance of A-SW by: 1) running the experiments multiple times and report the associated error bars (Table 1, Figure 1); 2) finding an experiment where the training of mini-batch-max-SW-generator does not get stuck in a bad local optimum, and compare its performance against A-SW. It would also be interesting to see how mini-batch max-SW performs when the number of mini batches k is equal to 1. \n\nTypos:\nMain document: l.92, l.263 (unclear sentences); l.102 (definition of $W_p$: the norm should be Euclidean so indexed by 2); equation (2) ($\\nu_\\phi$ instead of $\\nu$); l.131 ($Y_{\\phi,i}$ instead of $Y_{\\phi_i}$); l.174 (\"two probability measureS\")\n\nAppendix: derivations below l.523 (parenthesis issue, equation repeated twice); l.553; legend of Table 7.", "limitations": "The limitations of A-SW are discussed in Section 5 and seem to be the computational complexity (which, again, contradicts the original motivation). The potential negative societal impact is discussed briefly in Appendix F.", "rating": 4, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The problem addressed in this paper is adequate for NeurIPS: generative models based on optimal transport metrics, in particular variants of sliced Wasserstein distance, have gained popularity in the machine learning community, and it is important to develop strategies which further reduce their computational requirements.\n- The main idea behind the paper is quite original since, to the best of my knowledge, amortized optimization (which has been previously deployed in other areas such as variational inference and auto-encoders, for example) had not been applied in optimal transport. \n- Overall, I found the paper clearly written and easy to follow. Some typos are reported in \"Questions\".\n\nWeaknesses:\n\nUnclear conclusion: In the end, the message of this paper is confusing to me: in the introduction, the authors motivate the development of A-SW as a strategy to reduce the computational complexity of mini-batch max-SW, but this aspect is not supported by the computational complexity analysis (l.237-247) nor the experimental section (l.324-326: A-SW is \"comparable to Max-SW ... about the computational memory and the computational time\", or even slower and more memory-intensive when the amortized model is not linear (Table 2)). According to Table 1, the main advantage in using A-SW instead of mini-batch-max-SW in generative models is not the speeding-up of the training, but the improved quality of the results (Table 1). However, it is not clear why this happens, which leads me to the next point (\"Lack of theoretical justification\").\n\nLack of theoretical justification: the proposed methodology is not theoretically-grounded in my opinion and the empirical analysis is not comprehensive enough to compensate this lack of insights, which makes me question the performance of the proposed generative algorithm. More specifically,\n1) Proof of Proposition 2 (Appendix A.2): The conditions guaranteeing the existence of $\\phi^*$ should be specified.\n2) The parameterization of the amortized models is not clearly motivated. Are the proposed parameterizations standard in amortized optimization (if so, please add some references)? In particular, the amortization gap seems to be an important aspect in amortized optimization and is usually taken into account in the final methodology [1,2], but the authors chose not to study it (l.162-163). Are there any guarantees on whether the optimal projection directions lie in the space parameterized by Definitions 3, 4 or 5? \n3) A-SW outperforms mini-batch max-SW in terms of FID/IS scores in the image generation problem, and the authors conjecture that mini-batch max-SW performs poorly because it is stuck in a local optimum. What prevents A-SW from suffering from the same issue? This aspect could have also been explored empirically, which relates to my last point (\"No uncertainty measure in the experiments\").\n\nNo uncertainty measure in the experiments: In Section 5, if I am not mistaken, generative models have been trained only once on each dataset. In the checklist, the authors specified that they reported error bars, but I cannot find these anywhere. Having multiple runs is essential and can support the conjecture that A-SW-based generative models do not get stuck in a local optimum or at least less often than mini-batch max-SW generators (I am not fully convinced by this aspect for now).\n\n[1] \"Amortized Inference Regularization\", Shu et al. (NeurIPS 2018)\n\n[2] \"Iterative Amortized Policy Optimization\", Marino et al. (NeurIPS 2021)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "GGLMUIQ9N96", "reviewer": "Reviewer_mdhx", "summary": "The presented draft proposes a method for estimating the Max-SW by learning the projection leading to the largest possible SW metric.", "questions": "Could the method be extended to “top-k SW” instead of max-SW? \nOrthogonality could be enforced by a orthogonality loss via a Lagrangian.\n", "limitations": "The method requires additional memory and model evaluations but this is extensively discussed in the presented draft.\n\nIn practice with large models, i.e., in the demonstrated experiments, the computational cost of the Wasserstein loss is often small, which means that the utility of the method is limited in some applications. However, due to the favorable asymptotical complexity of the method, it allows learning in new settings which were previously impossible or hardly feasible. Further, even in the demonstrated cases the method achieve superior performance.\n\nWhat I would have liked to see would be using SW and Max-SW with even larger L / T. It would be interesting to see at which point the methods break even, even if the baselines take a very long time. \nIt would strengthen the paper even further, and with such an evaluation, I would consider raising my score.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths \nGenerative modeling is an important subject in machine learning.\n\nThe paper is exceptionally well written and simple to understand.\nThe proposed method is simple yet very powerful as it simplifies computation for Wasserstein generative models while improving model performance.\n\nThe background is very well explained. The method section is easy to understand.\n\nThe empirical evaluation is extensive, covering four data sets of different scales, and the results are impressive.\n\nI greatly appreciate the submission of the code, which is well structured.\n\n\n### Weaknesses\n\n#### Miscellaneous Issues\n* The presented paper uses a different font / font size or line distance. As the font / line distance is increased, I will not consider it for the review score, but it has be fixed for the final version.\n\n#### Typos\n* 35: there have been remained certain problems\n* 85: we make some conclusion. Suggestions: we provide a conclusion.\n* 86: preposition: defer … in -> defer … to \n* 174: measure -> measures\n* 223: leads to three corresponding amortized sliced Wasserstein [a word seems to be missing here], …\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept