Gromov-Wasserstein Autoencoders
Authors: Nao Nakagawa
Ren Togo
Takahiro Ogawa
Miki Haseyama
Venue: ICLR 2023 poster (2023)
Keywords: representation learning
deep generative models
variational autoencoders
optimal transport
implicit distributions
meta-prior
disentanglement
clustering

Abstract: Variational Autoencoder (VAE)-based generative models offer flexible representation learning by incorporating meta-priors, general premises considered beneficial for downstream tasks. However, the incorporated meta-priors often involve ad-hoc model deviations from the original likelihood architecture, causing undesirable changes in their training. In this paper, we propose a novel representation learning method, Gromov-Wasserstein Autoencoders (GWAE), which directly matches the latent and data distributions using the variational autoencoding scheme. Instead of likelihood-based objectives, GWAE models minimize the Gromov-Wasserstein (GW) metric between the trainable prior and given data distributions. The GW metric measures the distance structure-oriented discrepancy between distributions even with different dimensionalities, which provides a direct measure between the latent and data spaces. By restricting the prior family, we can introduce meta-priors into the latent space without changing their objective. The empirical comparisons with VAE-based models show that GWAE models work in two prominent meta-priors, disentanglement and clustering, with their GW objective unchanged.

Published as a conference paper at ICLR 2023
GROMOV- WAANO--WASSESTEEN AUTOENCODERS
Nao Nakagawa', Ren Togo?, Takahiro Ogawa', & Miki Haseyama" Graduate School of Information Science and Technology, Hokkaido University, Japan 2 Faculty of Information Science and Technology, Hokkaido University, Japan Iakagaaa, togo, ogawa, mhaseyama ama @ @1ld. ist hww.hoi......................... acc.
ABSTRACT
Variational Autoencoder (VAE)-based generative models offer flexible representa- tion learning by incorporating meta-priors, general premises considered beneficial for downstream tasks. However, the incorporated meta-priors often involve ad-hoc model deviations from the original likelihood architecture, causing undesirable changes in their training. In this paper, we propose novel representation learning method, Gromov- Wasserstein Autoencoders (GWAE), which directly matches the latent and data distributions using the variational autoencoding scheme. Instead of likelihood-basee objectives, GWAE models minimize the (GW) metric between the trainable prior and given data distributions. The GW met- ric measures the distance structure-oriented discrepancy between distributions even with different dimensionalities, which provides a direct measure between the latent meta-priors and data spaces. By restricting the prior family, we can introduce meta- priors into the latent space without changing their objective. The empirical comparisons with VAE-based models show that GWAE models work in two prominent meta-priors, disentanglement and clustering, with their GW objective unchanged.
INTRODUCTION
One fundamental challenge in unsupervised learning capturing the underlying low-dimensional structure of high-dimensional data because natural data (e.g., images) lie in low-dimensional man- ifolds (Carlsson et al., 2008; Bengio et al., 2013). Since deep neural networks have shown their potential for non-linear mapping, representation learning has recently made substantial progress in its applications
-igh--imensional and complex data (Kingma & Welling, 2014; Rezende et al., 2014; Hsu et al., 2017; Hu et al., 2017). Learning low-dimensional representations is in mounting demand because the inference of concise representations extracts the essence of data to facilitate various downstream tasks (Thomas et al., 2017; Higgins et al., 2017b; Creager et al., 2019; Locatello et al., 2019a). For obtaining such general-purpose representations, several meta-priors have been proposed (Bengio et al., 2013; Tschannen et al., 2018). Meta Meta-priors priors are general premises about the world, such as disentanglement (Higgins et al., 2017a; Chen et al., 2018; Kim & Mnih, 2018; Ding al., 2020), hierarchical factors (Vahdat & Kautz, 2020; Zhao et al., 2017; Sønderby al., 2016), and clustering (Zhao et al., 2018; Zong et al., 2018; Asano et al., 2020).
A prominent approach to representation learning is deep generative model based on the variational autoencoder (VAE) (Kingma & Welling, 2014). VAE-based models adopt the variational autoencoding scheme, which introduces inference model in addition to a generative model and thereby offers bidirectionally tractable processes between observed variables (data) and latent variables. In this scheme, the reparameterization trick (Kingma & Welling, 2014) yields representation learning capability since reparameterized latent codes are tractable for gradient computation. The introduction of additional losses and constraints provides further regularization for the training process based on vetarsriors. However, controlling representation learning remains a challenging task in VAE-based models owing to the deviation from the original optimization. Whereas the existing VAE-based approaches modify the latent space based the meta-prior (Kim & Mnih, 2018; Zhao al., 2017; Zong et al., 2018), 8), their training objectives still partly rely on the evidence lower bound (ELBO). Since the ELBO objective is grounded on variational inference, ad-hoc model modifications cause implicit and undesirable changes, e.g., posterior collapse (Dai et al., 2020) and implicit prior change (Hoffman al., 2017) in 3-V (Higgins et al., 2017a). Under such modifications, it is also unclear whether

---

Published as a conference paper at ICLR 2023
latent representation retains the underlying data structure because VAE models implicitly interpolate data points to form a latent space using noises injected into latent codes by the reparameterization trick (Rezende & Viola, 2018a;b; Aneja et al., 2021).
As another paradigm of variational modeling, the ELBO objective has been reinterpreted from the optimal transport (OT) viewpoint (Tolstikhin al., 201 20.8). Tolstikhin al. (2018) 8) have derived family of generative models called the Wasserstein autoencoder (WAE) by applying the variational autoencoding model nigh-dimensional OT problems as the couplings (Appendix A.4 for more details). Despite the OT-based model derivation, the WAE objective is equivalent to that of Info- VAE (Zhao al., 2019), whose objective consists of the ELBO and the mutual information term. The WAE formulation is derived from the estimation and minimization of the OT cost (Tolstikhin et al., 2018; Arjovsky et al., 2017) between the data distribution and the generative model, i.e., the generative modeling by applying the Wasserstein metric. It furnishes a wide class of models, even when the prior support does not cover the entire variational posterior support. The OT paradigm also applies to existing representation learning approaches originally derived from re- re-weighting the Kullback-Leible (KL) divergence term (Gaujac et al., 2021).
Another technique for optimizing the VAE-based ELBO objective called implicit variational infer- ence (IVD) (Huszár, 2017) has been actively researched. While the VAE model has an analytically tractable prior for variational inference, IVI aims at variational inference using implicit distributions, in which one can use its sampler instead of its probability density function. A notable approach to IVI is the density ratio estimation (Sugiyama al., 2012), which replaces the f f-divergence term in the variational objective with an adversarial discriminator that distinguishes the origin of the samples. For distribution matching, this algorithm shares theoretical grounds with generative models based on the generative adversarial networks (GANs) (Goodfellow et al., 2014; Sønderby et al., 2017), which induces the application of IVI toward the distribution matching in complex and high-dimensional variables, such as images. See Appendix A.6 for more discussions.
In this paper, we propose a novel representation learning methodology. Gromov- Wasserstein Au- toencoder (GWAE) based on the in (GW) metric (Mémoli, 2011) 1), an OT-based metric between distributions applicable even with different dimensionality (Mémoli, 2011; Xu et al., 2020; Nguyen et al., 2021). Instead of the ELBO objective, we apply the GW metric objective in the variational autoencoding scheme to directly match the latent marginal (prior) and the data distribution. The GWAE models obtain latent representation retaining the distance structure of the data space to hold the underlying data information. The GW objective also induces the variational autoencoding to perform the distribution matching of the generative and inference models, despite the OT--based derivation. Under the OTT-baeed variational autoencoding, one can adopt a prior of a GWAE model from rich class of trainable priors depending on the assumed meta-prior even though the KL divergence from the prior to the encoder is infinite. Our contributions are listed below.
We propose a novel probabilistic model family GWAE, which matches the latent space to the given unlabeled data via the variational autoencoding scheme. The GWAE models estimate and minimize the GW metric between the latent and data spaces to directly match the latent representation closer to the data in terms of distance structure.
We propose several families of priors in the form of implicit distributions, adaptively learned from the given dataset using stochastic gradient descent (SGD). The choice of the prior family corresponds to the meta-prior, thereby providing a more flexible modeling scheme for representation learning.
We conduct empirical evaluations on the capability of GWAF in prominent meta-priors: disentanglement and clustering. Several experiments on image datasets CelebA (Liu et al.. 2015), MNIST (LeCun et al., 1998), and 3D Shapes (Burgess & Kim, 2018), show that GWAE models outperform the VAE-based representation learning methods whereas their GW objective is not changed over different meta- priorss
2
RELATED WORK
VAE-based Representation Learning. VAE (Kingma & Welling, 2014) is prominent deep generative model for representation learning. Following its theoretical consistency and explicit handling of latent variables, many state-of-the--rt representation learning methods are proposed
2

---

Published as a conference paper at ICLR 2023
based on VAE with modification (Higgins et al., 2017a; Chen et al., 2018; Kim & Mnih, 2018; Achille & Soatto, 2018; Kumar et al., 2018; Zong et al., 2018; Zhao et al., 2017; Sønderby et al., 2016; Zhao et al., 2019; Hou et al., 2019; Detlefsen & Hauberg, 2019; Ding et al., 2020). The standard VAF. learns an encoder and a decoder with parameters & and 0, respectively, to learn a low-dimensional representation in its latent variables z using a bottleneck layer of the autoencoder. Using data x E Pdata (x) supported on the data space X, the VAE objective is the ELBO formulated by the following optimization problem:
maximize 0.0
E Bxsaa(x) (B)g(()x)) pe pe
DKL (KL(
(I)
where the encoder ( (z x) and decoder po (x(z) are parameterized by neural networks, and the prior r2)) is postulated before training. The first and second terms (called the reconstruction term and the KL term, respectively) in Eq. (1) are in a trade-off relationship (Tschannen et al., 2018). This implies that learning is guided to autoencoding by the reconstruction term while matching the distribution of latent variables to the pre-defined prior using the KL term.
Implicit Variational Inference. IVI solves the variational inference problem using implicit distribu- tions (Huszár, 2017). A major approach to IVI is density ratio estimation (Sugiyama et al., 2012), in which the ratio between probability distribution functions is estimated using a discriminator instead of their closed-form expression. Since IVI-based and GAN-based models share density ratio estimation mechanisms in distribution matching (Sønderby et al., 2017), the combination of VAEs and GANs has been actively studied, especially from the aspect of the matching of implicit distributions. The successful results achieved by GAN-based models in high-dimensional data, such as natural images, have propelled an active application and research of IVI in unsupervised learning (Larsen et al., 2016; Makhzani, 2018).
Optimal Transport. The OT cost is used as a measure of the difference between distributions supported on high-dimensional space using SGD (Arjovsky et al., 2017; Tolstikhin et al., 2018; Gaujac et al., 2021). This provides the Wasserstein metric for the discrepancy between distributions.
For a constant v 2 1, the ę- E-Wasserstein metric between distributions r and s is defined as
We(r, - (Y inf (TEP(r(0))(())) Exxxx)) [d(xx
1/E
where x denotes the random variable in which the distributions r and are defined, and P(r(x), s(x')) denotes the set consisting of all couplings whose x-marginal is r r(x) and whose x' x'-marginal s(x'). Owing to the difficulty of computing the exact infimum in Eq. (2) for high-dimensional, large-scale data, several approaches try to minimize the estimated Waaassesttein metric using neural networks and SGD (Tolstikhin al., 2018; 8; Arjovsky et al., 2017). The form in Eq. (2) is the primal form of the Wasserstein metric, particularly compared with its dual form for the case of & = 1 (Arjovsky et al., 2017). The two prominent approaches for the OT in high-dimensional, complex large-scale data (i) minimizing the primal form using a probabilistic autoencoder (Tolstikhin et al., 2018), and (ii) adversarially optimizing the dual form using a generator-critic pair (Arjovsky et al., 2017).
Wasserstein Autoencoder (WAE). WAE (Tolstikhin et al., 2018) is a family of generative models whose autoencoder estimates and minimizes the primal form of the Wasserstein metric between the generative model pa (x) and the data distribution Paata (x) using SGD in the variational autoencoding settings, i.e., the VAE model architecture (Kingma & Welling, 2014). This primal-basse formulation induces a representation learning methodology from the OT viewpoint because the WAE objective is equivalent to that of VAE (Zhao et al., 2019). which learns the variational autoencoding model by retaining the mutual information of the probabilistic encoder.
Duality. The Wasserstein GAN models (Arjovsky et al., 2017) adopt an objective based on the I-Wasserstein metric between the generative model px(x) and data distribu- tion Pdata(x)) (x). This objective estimated using the Mantorovvic-RRuiistttr duality (Villiani, 2009; Arjovsky et al., 2017), which holds for the I-Wasserstein as
W1 (r,
sup Ex(x) [f(x) - Ex(x)) [f(x)]
f:1-Lipschitz
To estimate this function of using SGD, a I-Lipschitz neural network called a critic is introduced, as with discriminator in the GAN-based models. The training process using is
3

---

Published as a conference paper at ICLR 2023
adversarially conducted, f.e., by repeating updates of the critic parameters and the generative param- eters alternatively. During this process, the critic maximizes the objective in Eq. (3) to approach the supremum, whereas the generative model minimizes the objective for the distribution match- ing po x) xs (x).
3 PROPOSED METHOD
Our GWAE models minimize the OT cost between the data and latent spaces, based on generative modeling in the variational autoencoding. GWAE models learn representations by matching the distance structure between the latent and data spaces, instead of likelihood maximization.
3.1 OPTIMAL. TRANSPORT BETWEEN SPACES
Although the OT problem induces metric between probability distributions, its application is limited to distributions sharing one sample space. The GW metric (Mémoli, 2011) measures the discrepancy between metric measure spaces using the OT of distance distributions. A metric measure space consists of a sample space, metric, and probability measure. Given a pair of different metric spaces, ie., sample spaces and metrics, the GW metric measures the discrepancy between probability distributions supported on the spaces. terms of the GW metric, two distributions are considered to be equal if there is an isometric mapping between their supports (Sturm, 2012; Sejourne et al., 2021). For a constant o 2 1, the formulation of the p-GW metric between probability distributions r(x) supported metric space (X, dxx) and s(z) supported on (Z, d2) ) is given by
(x [xdx(x,x] x') d2 (z, (2zz) 1/p
inf GW (r, = ((E(((()).()))
(4)
where Pxr(x), sz)) denotes the set all couplings with r(x) as x-marginal and sz) as z-marginal. The metries dx and dz are the metrics in the spaces X' and Z, respectively.
3.2 APPLICATION TO REPRESENTATIOE LEARNING: GROMOV- WASSERSTEIN AUTOENCODER
In this work, we propose novel GWAF modeling methodology based the GW metric for distance structure modeling in the variational autoencoding formulation. The objectives of generative models typically aim for distribution matching the data space, e.g., the likelihood (Kingma Welling, 2014) and the Jensen-Shanee Shannon divergence (Goodfellow et al., 2014). The GWAE objective differs from these approaches and aims to directly match the latent and data distributions based on their distance structure.
3.2.1 MODEL SETTINGS: VARIATIONAL AUTOENCODING
Given N-sized set of data points (x, N supported data space X', representation learning aims to build latent space Z and obtain mappings between both the spaces. For numerical computation, we postulate that the spaces X' and Z respectively have tractable metrics dx and dæ such as the Euclidean distance (see Appendix B.1 1 for details), and let M, E N (0). X S RMM and Z C LR We mention the bottleneck case M > L similarly to the existing representation learning methods (Kingma& & Welling, 2014; Higgins et al., 2017a; Kim & Mnih, 2018) because the data space X is typically an L-dimensional manifold (Carlsson et al., 2008; Bengio al., 2013).
We construct model with trainable latent prior To (2) to approach the data distribution Pdata in terms of distance structure. Following the standard VAE (Kingma Welling, 2014), we consider generative model (x z) with parameters 0 and inference model go (x. with parameters Φ. The generation process consists of the prior To (z) and decoder po Po (z) parameterized with neural networks. Since the inverted generation process (z)x) To (z) (x /))po(x) is intractable in this scheme, an encoder Is (z|x) zo Po (z(x))) is instead established using neural networks for parameterization. Thus, the generative po(x, z) and inference g z) models defined
Po x, 2) = To (x)
g 2) = Paata (x gos (z|x).
(5)
The empirical Paata (x) 1/N EN N 5(x x;) is used for the estimation of Pdata (x). A Dirac decoder and a diagonal Gaussian encoder are used to alleviate deviations from the data manifold as in Tolstikhin et al. (2018) 8) (see Appendix B.1 for these details and formulations).
4

---

Published as a conference paper at ICLR 2023
3.2.2 OPTIMAL TRANSPORT OBJECTIVE
Here, we focus on the latent space Z to transfer the underlying data structure to the latent space. This highlights the main difference between the GWAE and the existing generative approaches. The training objective of GWAE is the GW metric between the metric measure spaces (A², dx Pdata (x)) and (2,dz, (2)2)) as
minimize GW, (Pdata(x), mo(2)))
(6)
where p 2 a constant, and we adopt p= to alleviate the effect of outlier samples distant from the isometry for training stability. Computing the exact GW value is difficult owing to the high dimensionality of both x and z. Hence, we estimate and minimize the GW metric using the variational autoencoding scheme, which captures the latent factors of complex data in stable manner. We recast the GW objective into a main GW estimator Law with three regularizations: a reconstruction loss Lw, joint dual loss Lp, and entropy regularization RH-
Estimated GW metric Low- We use the generative model Po(x, z) as the coupling of Eq. (6) similarly to the WAE (Tolstikhin et al., 2018) methodology. The main loss Low estimates the GW metric as:
minimize Law == (p x)) [xxx(x,x) 22)(z1
subject Pdats (x)= (x();;
where a trainable scale constant to cancel out the scale degree of freedom, and pe (x) denotes the marginal po (x) -fz (x,
WAE-based X-mmarginal condition Lw. To obtain a numerical solution with stable training, Tolstikhin et al. (2018) relax the X' X-matching condition of Eq. (8) into E-Wasserstein minimiza- tion (E 2 1) using the variational autoencoding coupling. The WAE methodology (Tolstikhin et al., 2018) uses the inference model (x(x, z) to formulate the &- Wassersteiin minimization as the reconstruction loss LW with a 2-matching condition as:
minimize 0.6
Lw (xp (d)(xxx')]
(9)
subject to (z) Tz()).
(10)
where dx is distance function based on the Le metric. We adopt the settings { = 2 to retain the conventional Gaussian reconstruction loss.
Merged sufficient condition Lp- We merge the marginal coupling conditions of Eq. (8) and Eq. (10)
into the joint X' x Z-matching sufficient condition p z) = 9x(xx z) to attain bidirectional inferences while preserving the stability of autoencoding. Since such joint distribution matching can also be relaxed into the minimization of W1 (gs(x, [x, 2x, 2)), this condition is satisfied by minimizing the Practically, a I-Lipschitz neural network (critic) fe estimates the supremum of Eq. (3), and the main model minimizes this estimated supremum as:
duality introduced by Arjovsky al. (2017) in Eq. (3).
minimize maximize 0.4
LD := (xx( [f (x, 2)] (x. [f (x, 2)]
(11)
where 2b is the critic parameters. To satisfy the I-Lipschitz constraint, the critic fos is implemented with techniques such as spectral normalization (Miyato al., 2018) and gradient penalty (Gulrajani et al., 2017) (see Appendix B.3 for the details of the gradient penalty loss).
Entropy regularization Ru- We further introduce the entropy regularization Ru using the inference entropy to avoid degenerate solutions in which the encoder (z|x) becomes Dirac and deterministic for all data points. In such degenerate solutions, the latent representation simply becomes a look-up table because such point-to-point encoder maps the set of data points into set of latent code points with measure zero (Hoffman et al., 2017: 17; Dai et al., 2018), causing overfitting into the empirical data distribution. An effective way to avoid it is regularization with the inference entropy H. of the latent variables z conditioned on data x as
RH == 1l = Ens (x.2) [- log 9 (z|x)]
(12)
Since the conditioned entropy (z|z))) diverges to negative infinity in the degenerate solutions, the regularization term -R facilitates the probabilistic learning of GWAE models.
5

---

Published as a conference paper at ICLR 2023
Stochastic Training with Single Estimated Objective. Applying the Lagrange multiplier method to the aforementioned constraints, we recast the GW metric of Eq. (6) into single objective with multipliers Xw, Ap, and Xx
minimize maximize 0.4
C: Low + Xw Lw + ADCD RH
(13)
One efficient solution to optimize this objective is using the mini-batch gradient descent in alter- native steps (Goodfellow et al., 2014; Arjovsky et al., 2017), which we can conduct in automatic differentiation packages, such as PyTorch (Paszke al., 2019). One step of mini-bate descent the minimization of the total objective L. in Eq. (13), and the other step is the maximization of the critic objective LD Eq. (11). By alternatively repeating these steps, the critic estimates the Wasserstein metric using the expected potential difference LD (Arjovsky et al., 2017). Although the objective in Eq. (13) involves three auxiliary regularizations including an adversarial term, the GWAE model can be efficiently optimized because the adversarial mechanism and the variational autoencoding scheme share the goal distribution matching pe (x, z) x q (x, z) (see Appendix C.5 for more details).
3.2.3 PRIOR SAMPLING
GWAE models apply to the cases in which the prior Te (z) takes the form of an implicit distribution with a sampler. An implicit distribution Te (z) provides its sampler z ~~ (T(2) while a closed-fort expression of the probability density function is not available. The adversarial algorithm of GWAE handles such cases and enables a wide class of priors to provide meta-prior- based inductive biases for unsupervised representation learning, e.g-, for disentanglement (Locatello al., 2019b; 2020). Note that the GW objective in Eq. (6) becomes a constant function in non-trainable prior cases.
Neural Prior (NP). A straightforward way build differentiable sampler of trainable prior is using a neural network to convert noises. The prior of the latent variables z is defined via sampling using a neural network ge _² R with parameters (see Appendix B.2 for its formulation). Notably, the neural network need not be invertible unlike Normalizing Flow (Rezende & Mohamed, 2015) since the prior is defined an implicit distribution not requiring a push-forward measure.
Factorized Neural Prior (FNP). For disentanglement, we can constitute factorized prior using
an element-wise independent neural network ğo { ) (see Appendix B.2 for its formu- lation). Such factorized priors can be easily implemented utilizing the I-dimensional grouped convolution (Krizhevsky et al., 2012).
Gaussian Mixture Prior (GMP). For clustering structure, we construct a class of Gaussian mixture priors. Given that the prior contains K components, the k-tth component is parameterized using the weights Wh. means m E RR,, and square-root covariances M, E L. as
K
To (z) =
(14)
1=1
where the weights {w are normalized as Σk= wk = 1. To sample from a prior of this class, randomly chooses component i from the K K-way categorical distribution with probabili- ties (w,, Wy..... wp) and draws a sample z as follows:
m MAC.
E C~,LO,
(15)
where O and I, denote the zero vector and the n-sized identity matrix, respectively. In this class of 1 priors, the set of trainable parameters consists of {(wk, my. M Note that this parameterization can be easily implemented in differentiable programming frameworks because M&M is positive semidefinite for any Mg E LxL
EXPERIMENTS
We investigated the wide capability of the GWAE models for learning representations based on meta- meta-priors. We evaluated GWAES in two principal meta-priors: disentanglement and clustering.
In the tables of the quantitative evaluations, 1 and indicate scores in which higher and lower values are hetter, respectively.

---

Published as a conference paper at ICLR 2023
Ensinited Engliial
3820
220
2520
2050
140
11
120
-0
23
Number Spods Traned
(a) The estimation of the GW metric in each epoch.
(b) The isometry in GWAE.
The estimation and minimization of the GW metric. This trial of training is conduced in GWAE (NP, Ap=1. Aw=l, xx=l) using the MNIST (LeCun et al., 1998) dataset. (a) The curves show the GW values estimated by the loss term Law (solid, blue) and the empirical GW computed by the POT package (Flamary et al., 2021) (dashed, orange). The values are computed using the validation set. (b) The axes Ar = (x x') (vertical) and A (z, 2') (horizontal) respectively denote the difference the data and latent spaces between generated samples (x, z), (x', zz) Po z). The histogram contains 10,000 generated sample pairs.
Figure I:
To validate the effectiveness of GWAE on different tasks for each meta-prior, we conducted each experiment in corresponding experimental settings. We further studied their autoencoding and generation for the inspection of general capability.
4.1 EXPERIMENTAL SETTINGS
We compared the GWAE models with existing representation learning methods (see Appendix A for the details of the compared methods). For the experimental results in this section, we used four visual datasets: CelebA (Liu et al., 2015), MNIST (LeCun et al., 1998), 3D Shapes (Burgess & Kim, 2018), and Omniglot (Lake et al., 2015) (see Appendix C.1 for dataset details). For quantitative evaluations, we selected hyperparameters from Aw E [10", 10'], Ap E [10", 10 and Ax E [10 4, 10 using their performance on the validation set. For fair comparisons,
trained the networks with consistent architecture from scratch in all the methods (see Appendix C.2 for architecture details).
4.2 GROMOV- WASSERSTEIN ESTIMATION AND MINIMIZATION
We validated the estimation and minimization of the GW metric in Fig. 1. First, to validate the estimation of the GW metric, we compared the GW metric estimated in GWAE and the empirical GW value computed in the conventional method in Fig. 1a. Against the GWAE models estimating the GW metric as in Eq. (7), the empirical GW values are computed by the standard OT framework POT (Flamary et al., 2021). Although the estimated Low is slightly higher than the empirical values, the curves behave in very similar manner during the entire training process. This result supports that the GWAE model successfully estimated the GW values and yielded their gradients to proceed with the distribution matching between the data and latent spaces. Second, to validate the minimization of the GW metric, we show the histogram of the differences of generated samples in the data and latent space in Fig. 1b. The isometry of generated samples is attained if the generative coupling Pe (x, z) attains the infimum in Eq. (4). This histogram result shows that the generative model pe (x, z) acquired nearly-isometric latent embedding, and suggests that the GW metric was successfully minimized although the objective of Eq. (13) contains three regularization loss terms (refer to Appendix C.8 for ablation studies, and Appendix C.4 for comparisons). These two experimental results support that the GWAE models successfully estimated and optimized the GW objective.
4.3 LEARNING REPRESENTATIONS BASED ON META-PRION
Disentanglement. We investigated the disentanglement of representations obtained using GWAE models and compared them with conventional VAE-based disentanglement methods. Since the element-wiss wise independence in the latent space is postulated as a meta-prior for disentangled represen- tation learning, we used the FNP class for the prior To (z).. Considering practical applications with
7

---

Published as a conference paper at ICLR 2023
1.0
1.0
10
20
30
0.
0.8
08
1.0
05
005
06
69
04
04
04
02
02
-9
02
-2.
-20
-1.0
0.0
1.0
2.0
055
4.0
655
00
05
10
15
-20
-10
10
20
(a) VAB (Kingma & Welling, (b) FactorVAF (Kim & Mnih, (c) GWAF (FNP, Aw=l, \p=10, 2014).
(y=10).
0x=0.3).
Figure 2: Comparison of the learned latent spaces in 3D Shapes (Burgess & Kim, 2018) and L = 16. The vertical and horizontal axes in the scatter plots respectively represent two of the 16 (= L) latent variables with the highest and the second-higs highe- informativeness (Do & Tran, 2020) w.r.t. the object hue factor. Note that a single factor value varies along only one axis in disentangled representation.
Table I: Quantitative comparison of disentanglement. The reported scores were calculated in 3D Shapes (Burgess Kim, 2018), and the latent size 16. Since the latent size L is larger than the number of the ground truth factors, the hyperparameter tuning was based on the validation set DCI-C (Eastwood Williams, 2018) values. To deal with the probabílistic scores (Zaidi et al., 2021). we reported the ranges for five measurements. The details of the scores are provided in Appendix C.3.
Model
DCI-C
DCI-D f
DCI-I t
VAE (Kingma & Welling, 2014)
0.7734 0.0004
0.6831
0.0002
0.9914 + 0.0003
B VAE (Higgins al., 2017a)
0.8245 + 0.0002
0.7328 0.0002
0,9796 = 0.0002
WAE (Tolstikhin et al., 2018)
0.8288 + 0.0004
0.7544 = 0.0004
0.9959 = 0.0001
3-TCVAF (Chen et al., 2018)
0.8347 + 0.0003
0.7085 + 0.0002
0.9880 + 0.0002
FactorVAR (Kim & Mnih, 2018)
0.7963 + 0.0004
0.7390 + 0.0004
0.9961 + 0.0002
DIP-VAF-I (Kumar al., 2018)
0.8609 + 0.0003
0.6984
0.0003
0.9961 + 0.0001
DIPIVAE-II (Kumar al., 2018)
0.8236 1 0.0001
0.7498 0.0003
0.9957 1 0.0002
GWAE (FNP)
0.9080 + 0.0002
0.7024 0.0002
0,9966 1 0.0002
* The ranges are denoted by (mean) (standard error of mean).
unknown ground-truth factor, we set relatively large latent size I. to avoid the shortage of dimension- ality. The qualitative and quantitative results are shown Fig. 2 and Table 1, respectively. These results support the ability to learn disentangled representation in complex data. The scatter plots in Fig. 2 suggest that the GWAE model successfully extracted one underlying factor of variation (object hue) precisely along axis, whereas the standard VAE (Kingma & Welling, 2014) formed several clusters for each value, and FactorVAL (Kim & Mnih, 2018) 8) obtained the factor in quadrants.
Clustering Structure. We empirically evaluated the capabilities of capturing clusters using MNIST (LeCun et al., 1998). We compared the GWAE model using GMP with other VAE-based methods considering the out-of-distribution (OoD) detection performance in Fig. 3. We used MNIST images as in-distribution (ID) samples for training and Omniglot (Lake et al., 2015) images as unseen OoD samples. Quantitative results show that the GWAE model successfully extracted the clustering structure, empirically implying the applicability of multimodal priors.
4.4 AUTOENCODING MODEL
We additionally studied the autoencoding and generation performance of GWAB models in Table (see Appendix C.7 for qualitative evaluations). Although the distribution matching px(x) ss Paita (x) a collateral condition of Eq. (7), quantitative results show that the GWAE model also favorably compares with existing autoencoding models in terms of generative capacity. This result suggests the substantial capture of the underlying oiw-dinesiional distribution in GWAE models, which can lead to the applications to other types of meta-priors.
8

---

Published as a conference paper at ICLR 2023
Tou positee rate
Fabe positive rats
Fable posiee
Falle post rota
(a) VAB (Kingma & Welling, 2014), AUC=0.9957.
(b) DAGMM (Zong et al., 2018), (c) GWAE (GMP), AUC=1.0000. AUC=0.9654.
Figure 3: The ROC curves of the OoD detection in MNIST (LeCun et al., 1998) against Om- niglot (Lake al., 2015). We trained these models using MNIST ID samples and used Omniglot as OoD samples. We upsampled Omniglot to 10,000 samples for data balancing. For the anomaly detection using the latent codes z, we applied the negative log- likelihood energy loog ग (z) for VAE and DAGMM, and used the estimated Kantorovich potential (x)()) [f (x, z)] for GWAE (see Appendix C.10 for more latent space details).
Table 2: Quantitative comparisons of generation and reconstruction. The FID scores (Heusel et al., 2017) evaluate random sample set from the generative model pe (without using dataset images) against the entire test set, and both consist of an equal number of 19,962 samples. The PSNR scores measure reconstruction Tsl (z)po(s (ppp(xpp x|z) using test images (see Appendix C.3 for details). All reported values were computed in CelebA (Liu et al., 2015) with a latent size of L, = 64. For all the methods, we applied early stopping (patience=100 10) and hyperparameter tuning using the validation set. The bold and underlined values respectively denote the best and the second-best performance in each score.
Model
FID
PSNR [dB]
Baseline
VAE (Kingma & Welling, 2014)
130.9
19.96
3- VAE (Higgins et al., 2017a)
92.6
22.71
KL re-weighting
GECO (Rezende & Viola, 2018a)
162.1
21.19
a- VAE (Rybkin et al., 2021)
53.13*
20.03
LAd (Sonderby al., 2016)
255.6
12.35
Hicrarchical factors
LLaderAAE (Zhao et al., 2017)
147.1
19.76
WAE (Tolstikhin et al., 2018)
55"
22.70
WVI (Ambrogioni et al., 2018)
295.0
14.45
OT-haaed models
SWAE (Kolouri al., 2019)
102.2
21.85
RAF (Xu et al., 2020)
52.20'
21.34
VampPrior (Tomezak & Welling, 2018)
243.8
16.23
Trainable priors
2-Stage VAE (Dai & Wipf, 2019)
34'
16.15
VAE-GAN (Larsen al., 2016)
111.8
19.51
IVI-based models
AVB (Mescheder al. 2017)
93.0
22.60
ALI (Dumoulin al., 2017)
171.8
12.26
Ours GWAE (NP)
45.3
22.82
The values are cited from the original papers annotated after the model names.
CONCLUSION
In this work, we have introduced novel representation learning method that performs the distance distribution matching between the given unlabeled data and the latent space. Our GWAE model family transfers distance structure from the data space into the latent space in the OT viewpoint, replacing the ELBO objective of variational inference with the GW metric. The GW objective provides direct measure between the latent and data distribution. Qualitative and quantitative evaluations empirically show the performance of GWAE models in terms of representation learning. In future work. further applications also remain open to various types of meta-priors, such as spherical representations and non- non-Euce -on-Euclidean embedding spaces.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
We describe the implementation details in Section 4, Appendix B, and Appendix C. The dataset details provided in Appendix C.I. To ensure reproducibility, code is available online at httttts:/ /githubb wttphhuwwcocooooooooooooo..e.eoe wttpgaaaiooigoooooooooooo. gwaae gwae and is provided as the supplementary material.
ACKNOWLEDGMENTS
This work was partly supported by AMED Grant Number JP212f0127004 and JSPS KAKENHI Grant Number JP21H0
REFERENCES
Alessandro Achille and Stefano Soatto. Information dropout: Learning optimal representations through noisy computation. IEEE Transactions on Pattern Analysis & Machine Intelligence, 40 (12):2897-2905,, 2018. doi: 10.1 109/TPAM..001.2788//8
3, 16
Alexander A. Alemi, lan Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-19, 2018. URL wtt.t:pps//mm wwwpeeview.ee.cc/////p/.. wet.eiiwwweeeeeeeeeeeeeeeeeeeeee net/ forum Himmmidd--yxiicooo
16
Luca Ambrogioni, Umut Güçlü, Yağmur Güęlütürk, Max Hinne, Marcel A. J. van Gerven, and Eric Maris. Wasserstein variational inference. In Proceedings of Neural Information Processing Systems (NIPS), pp. 2473-2482, 2018. URL ht /wtppaaappppaeeaaa ni cc wtc.c//pp/ae paper/2015
69bf8
Bww.eff-AAbbtrrre Abstract html.
Jyoti Aneja, Alex Schwing, Jan Kautz, and Arash Vahdat. A contrastive learning approach for training variational autoencoder priors. In Proceedings of Neural Information Processing Sys- lems (NewrlPS), pp. 480-493, 2021. URL htw.s: /ttppcee///////eeepepepeppp/eeee ngs. nwwriesupp..er ccc sec/paper/ 2021 /ww.hhaaaaat 0499/04966664 d80 f66.66666666666 wtw.oeeebbeceeeeee0e.eee 20622226666 Abstract. hhml1. 2
Martín Arjovsky and Léon Bottou. Towards principled methods for training generative adversarial
networks. In Proceedings of the International Conference Learning Representations (ICLR), pp. 1-17, 2017. URL https: htltttts////ppeeevea/
net/ forum?id=Hlk _awddd-Hk4__wwxee_e___dd_ 20
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks.
In Proceedings of the International Conference on Machine Learning (ICML), pp. 214-223, 2017. URL https: http::://ppcceeiiie
mllr. www.:aaww..................m...e
a htmm. 2, 3, 5, 6, 20
Yuki M. Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultaneous clustering and representation learning. In Proceedings of the International Conference on Learning Rep- resentations (ICLR), 2020. URL https:/ https:://pennevaaeeee /openzeview. net wwtp//eet//ete. cefoorummidHHye. jyBFFr.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 1788 2013. doi: 10.1109/TPAMI.2013.55 1,4
Representation learning: A review and new
Leo Breiman. Random forests. Machine Learning, 2000 doi: 10.1023/A:1010093344422 25
Chris Burgess and Hyunjik Kim. 3D Shapes Dataset. https: https://github. wwt. http:/tttwb..oo////eemddmdd
3d- shapes/, 2018. Accessed May 13, 2022. 2, 7, 8, 22
Gunnar Carlsson, Tigran Ishkhanov, Vin de Silva, and Afra Zomorodian. On the local behavior of
spaces of natural images. International Journal of Computer Vision (IJCV), 76(1):1-2, 2008. doi: 10.10/7/s 226-007--0056- 1,4
Ricky T. Q. Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disen- tanglement in variational autoencoders. In Proceedings of Neural Information Processing Sys- tems (NIPS), pp. 2610-2624 2018. URL https: httpttttss//poccceeanggg/ooo
www.ine...c..ooo wwweeuups.ccpppppeeeeeeeeee 201 8 htt/haah/aae ee 2tt. atw.::22222e2222aaaaaaaaaa
- .hmll. 1,3, 21
10

---

Published as a conference paper at ICLR 2023
Casey Chu, Kentaro Minami, and Kenji Fukumizu. Smoothness and stability in GANs. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-15, 2020. URL https: http:///ppeeeeeeeeeeeee net
21
Elliot Creager, David Madras, Joern-Henrik Jacobsen, Marissa Weis, Kevin Swersky, Toniann Pitassi, and Richard Zemel. Flexibly fair representation learning by disentanglement. In Proceedings of the International Conference on Machine Learning (ICML), pp. 1436-1445, 2019. URL pettps://proeeet wtwpeeeeiiing.m...pe.......
Bin Dai and David Wipf. Diagnosing and enhancing vae models. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-12, 2019. URL https: httptttsss//ppeeeeeeeeeeee /ppnnreveee....
forum? 1111ddd111e1111 9, 18, 27, 28, 35
Bin Dai, Yu Wang. John Aston, Gang Hua, and David Wipf. Connections with robust PCA and the role of emergent sparsity in variational autoencoder models. Journal of Machine Learning Research, 19(41):1-44 42, 2018. URL http:: jmlr. org httpp/aaaaaaa//aa/aaeeaae
html.
Bin Dai, Ziyu Wang, and David Wipf. The usual suspects? collapse. In Proceedings of the International Conference on Machine Learning (ICML), pp. 2313-2322 2020. URL https: https::/ http///ooeeenngsse//ee/ee/e mlr. press peress/vil9 https/vll///////// /ail0//ai0....mm/./ html. I
Reassessing blame for VAE posterior
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale
hierarchical image database. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 248-252-5 2009. doi: 10.1 20.1099/CVRR.2205520648.
26
Nicki Skafte Detlefsen and Søren Hauberg. Explicit disentanglement of appearance and perspective in generative models. In Proceedings of Neural Information Processing Systems (NeurIPS), pp. 1018-1028, 2019. URL htttps://proceeeiiiie .ng.urrrpppi. neurips....... cccccc/p/ooo.. http/ppprr///00 9/ 9/hash/ 3433344feee03ccff3333333333
Abst 3tml.
Zheng Ding, Yifan Xu, Weijian Xu, Gaurav Parmar, Yang Yang. Max Welling, and Zhuowen Tu. Guided variational autoencoder for disentanglement learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 7920- 7929, 2020. URL https://openaccee
thectheeeeeee com/ wwwpecnntteenteteeooeoo httml/Ding_tt Gurg__uiiee htw.o//ariiiiiaaaaattttreooeeoo. @earning_cVPR.cc 2020 200__pppere_pp_e_____ html. 1,3
Kien Do and Truyen Tran. Theory and evaluation metrics for learning disentangled representations. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-30, 2020. URL https::///........... ///p//ppeeviieeeee/ee/eeeeeeeeee met/fooum? h
Jeff Donahue, Philipp Krähenbühl, and Trevor Darrell. Adversarial feature learning. Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-18, 2017. URL
https: http:///ppeeeeeeeeeee/eeeeeeeeee www.w.eet.. wttee/ooommiiidddttaaa..
20
Vincent Dumoulin, Ishmael Belghazi. Ben Poole, Olivier Mastropietro, Alex Lamb, Martin Arjovsky, and Aaron Courville. Adversarially learned inference. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-18, 2017. URL htthst/::/ hto//ppereeee revi ew. net Dutt/fose Eorum? iiwiiid111eee 31E1R4cgg. 9, 20, 28, 35
Cian Eastwood and Christopher K. Williams. A framework for the guantitative evaluation of disen- tangled representations. In Proceedings of the International Conference on Learning Representa- tions (ICLR), pp. 1-15, 2018. URL https: http:///ppeeeeeeeeeeeeeeeeeee.ea net forum? forum?id=By- Bd= Adz-AZ. 8, 24, 26
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021. URL http: / /mml= r. orgg haagg/apppe pappas/pppp/pspppppp x22 /0 451. html. 7,37
11

---

Published as a conference paper at ICLR 2023
Benoit Gaujac, Ilya Feige, and David Barber. Learning disentangled representations with the wasserstein autoencoder. In Proceedings of the European Conference on Machine Learning and Principles and Practice Knowledge Discovery in Databases (ECML PKDD), Part III, pp. 69-84, 2021. doi: 10. 10000/978-3-030-86553-3
2,
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Proceedings of Neural Infor- mation Processing Systems (NIPS), pp. 2672 2680, 2014. URL https: /tpppraappsssp//////eeeeeaaaesss nips. cc/ -bbtrffctttb.tttt..o.ooooaaaaa.. html.
5cale//caeeeee 9b paae:///l///astt/scaeeeeeee
4, 6, 19
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. Proceedings of Neural Information Processing Systems (NIPS), pp. 5769-5779, 2017. URL https: / 322c3blccdccd5222 ctteee2ccddffffff8oo Abstract - html. 5,
//tpprss//ppprrs../ papers.pppppppppp/pppp nips cc/ htpppppp//////// 2017/hhah/
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv: 1606.08415, 2016. 23, 24
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Proceedings of Neural Information Processing Systems (NIPS), 2017. URL https: //aperss//pprrsssee nips cc/ httpppppr//100/ 2ww-88811133944222--bbb2a-- Abstract. html. 24, 26, 29
pp.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 3-VAE: Learning basic visual concepts with constrained variational framework.
Proceedings of the International Conference on Learning Representattions (ICLR), pp. 1-22, 2017a. URL https: htttts:://pennevieew..oo net wtt.://wwww.eeee///eeiie. forumiid- Sy2fzU9g1. 1, 3, 4, 8, 9, 16, 21, 28, 32, 33,
Irina Higgins, Arka Pal, Andrei A. Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In Proceedings of the International Conference on Machine Learning (ICML), pp. 1480-1490, 201 7b. URL /ptp::e/rrcceeeiiiii ngs. ml p......eresv00 higgins a html.
Matt Hoffman, Carlos Riquelme, and Matthew Johnson. The beta VAE's implicit prior. In Proceedings of Neural Information Processing Systems (NIPS) Workshop on Bayesian Deep Learning, pp. 1-5, 2017. URL https: research goog 1 e pubs/ tubbs/pube ppbp433500 1,5
Xianxu Hou, Ke Sun, Linlin Shen, and Guoping Qiu. Improving variational autoencoder with deep feature consistent and generative adversarial training. Neurocomputing, 341:183-194 2019. doi: hww.://j.neccom.21990
3
Wei-Ning Hsu, Yu Zhang, and James Glass. Learning latent representations for speech generation and transformation. Proceedings of the Anmual Conference of the International Speech Communica- tion Association (INTERSPEECH) pp. 1273-1277, 2017. doi: 09.214457/nerrseehh22117333
Zhiting Hu, Zichao Yang, Xiaodan Liang, Ruslan Salakhutdinov, and Eric P. Xing. Toward controlled generation of text. Proceedings of the International Conference on Machine Learning (ICML), pp. 1587-1596, 2017. URL https: httttps:/pprcceeiingg/ wtw.aiggs................e.e.eee /vcoss//vv/o v70 /hwl7e html. 1
Ferenc Huszár. Variational inference using implicit distributions. arXiv: 1702.08235, 2017. 2,3,
Hyunjik Kim and Andriy Mnih. Disentangling by factorising, In Proceedings of the International Conference on Machine Learning (ICML), pp. 2649-2643 2018. URL MtLLtt:://pprcedgnng
mlr. www.spresssveeee v8/ / Kiml 8' html. 1, 4. 8, 17, 21,
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings
of the International Conference on Learning Representations (ICLR), pp. 1-13, 2015. URL htttp:://pperrevieee net f0
12

---

Published as a conference paper at ICLR 2023
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Proceedings of the International Conference on Learning Representations ( pp. 1-13, 2014. URL https: /ttpee/sseeee..eet/ooooaaaaaaaa-
// //oppeeevieww.oo nnt forummid=33X9Eda 9Py2d. 1, 3, 4, 8, 9, 16, 20, 21, 27, 29, 30, 38
Soheil Kolouri, Phillip E. Pope, Charles E. Martin, and Gustavo K. Rohde. Sliced wasserstein auto-encoders, In Proceedings of International Conference on Learning Representations (ICLR), 2019. URL https: /ectpss:/openzeeeet
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images, 2009. Master's thesis, Technical Report, University of Toronto. 22, 28, 31, 34
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In Proceedings of Neural Information Processing Systems (NIPS), pp. 1097-1105, 2012. URL /tttps://papeee nips cc/ wwtpppppp//////////e /hash/ c399862d3b9ddbb 0tw.r/e666666633222222666e6 -ateaeeeeeeeaeao----taattaatat Abstract.htm 6
Abhishek Kumar, prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-16, 2018. URL https: / http//ppwwweweeeeeeeeee nett/ forum? d111 KG7GZAN 3,
Brenden M. Lake, Ruslan Salakhutdinov, and Joshua B. Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 33066266) 1338, 2015. doi: 10.1 10.1126/ science.aab3050. 7, 8, 9, 22
Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Au- toencoding beyond pixels using a learned similarity metric. In Proceedings of the International Conference on Machine Learning (ICML), pp. 1558-1566, 2016. URL http: httttt::/proceee pproceediee..... mlr.press/ https/ss/////aaaaeeeeee larsenle 6, html. 3, 9, 20, 21, 28, 35
Yann LeCun, Léeon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. Proceedings of the IEEE, 86(1 (6(11):2278-233 1998. doi: 10.1109/5.72679 7,7,8, 8, 9, 22, 27, 28, 30, 33, 37
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), pp. 3730-3738, 2015. doi: 10. 10.110//CCC.2015..40 2, 7, 9, 22, 28, 29, 31, 32, 35,
Francesco Locatello, Gabriele Abbati, Thomas Rainforth, Stefan Bauer, Bernhard Schölkopf, and Olivier Bachem. On the fairness of disentangled representations. In Proceedings of Neural Information Processing Systems (NeurIPS), pp. 14584- 14597, 2019a. URL. 1b486d7a511 B 4
wttp:/ss.....p....pppeeeeeeee html.
9 //hash/
Francesco Locatello, Stefan Bauer, Mario Lucic, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled represen- tations. In Proceedings of the International Conference on Machine Learning (ICML), pp. 41 4114 4124, 2019b. URL http: http://////occceengg...... mlr. prwss/ / http://ccaaellllooooo.eo.e htmll. 6
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. A sober look at the unsupervised learning of disentangled representations and their evaluation. Journal of Machine Learning Research, 21(209): 2020. URL http: //jmlr r. org/ pwtpg//ppes/////ee wwllll1ll--1 976 html.
Andrew L. Maas, Awni Y. Hannun, and Andrew Y. Ng. Rectifier nonlinearities improve neural network acoustic models. Proceedings of the Workshop on Deep Learning for Audio, Speech, and Language Processing, ICML (WDLASL), pp. 3-9, 2013. URL ht robotics st stanforr ford.. http:/ppppssrrreee aell/-amaasspapers//eeee icm fww.liilll.ll. fdffl.pdf.. pa.f. 25
edu/-amaas
Alireza Makhzani. Implicit autoencoders. arXiv: 1805.09804, 2018.
13

---

Published as a conference paper at ICLR 2023
Facundo Mémoli. Gromov- wasserstein distances and the metric approach to object matching. Foun- dations of Computational Mathematics, 11(1):417-487 2011. doi: 00000//10208-011-999999
2,
Lars Mescheder, Sebastian Nowozin, and Andreas Geiger. Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks. In Proceedings of the International Conference on Machine Learning (ICML), pp. 2391-2400, 2017. URL http:/ http:::/pproceeeee /ttp///eecccee..../. / v www.//eeccceeeeeeeeeeeeeeeeeeeee 17a 1at.tm................ 9, 20, 32, 33, 34
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In Proceedings of International Conference on Learning Representations (ICLR), pp. 1-26, 2018. URL https:/ /oppnreviewww..//////p net ftrumrrrmnn ?id B10RgziT-, 21,
Khai Nguyen, Son Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Improving relational regularized autoencoders with spherical sliced fused gromov wasserstein. In Proceedings of the International Conference on Learning Representations (ICLR), pp. 1-11, 2021. URL https: openreview. met/Eorunt 2 dd= WDW.L2233.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Proceedings of Neural Information Processing Systems (NeurlPS), pp. 8024 8035, 2019. URL httpt:ttttppppeeeeeeeee nips cc/ htppp/aper////////// 221 9/ hash/ bdbca288fee www.feefffffbffbcecce baboo228Eee1112EEE03337772222211
Abstrrattbbbrttbaaaa cttpstrrcttttttl...... html. 6, 22
Danilo J. Rezende and Fabio Viola. Generalized elbo with constrained optimization, geco. In Proceed- ings of Neural Information Processing Systems (NIPS) Workshop on Bayesian Deep Learning, pp. 1-11, 2018a. URL http:/ / http//aaasssaeeeeeeeee seyyessandeeplearniii. org. Pttp///pappp333333 pdf. 2,
Danilo J. Rezende and Fabio Viola. Taming vaes. arXiv: 1810.00597, 2018b. 2, 29
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In
Proceedings of the International Conference on Machine Learning (ICMI), pp. 1530-1538, 2015.
URL http: http::://rrooeeeeee proceediiiiiooo ngs mll. pww.ppressvvee vv7 htt.///eeeeee1...5/ html.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings the International Conference on Machine Learning (ICML), pp. 1278-1286. 2014. URL https: /ttpe//eioo/eegsss////o..//...se mlr.
press/v32 rttp///eeeeeeeeeem html.
Oleh Rybkin, Kostas Daniilidis, and Sergey Levine. Simple and effective vae training with calibrated decoders. In Proceedings of the International Conference on Machine Learning (ICML), pp. 9179-9189, 2021. URL http: http///pppppegsseeeeeee www..rr......m www.ppeessssvvee https////1///l/kkieeeee ryw.innlkknna.aa html.
Thibault Sejourne, Xraveeis-Xavier Vialard, and Gabriel Peyré. The unbalanced gro- mov wasserstein distance: Conic formulation and relaxation. In Proceedings of Neurul Information Processing Systems (NeurIPS), volume 34, pp. 3766-8779, URL tttpss:/proceed:: ngs neurips wwtppppr//0211111111 /aah//hhhhh/// http:///e0ddddeeeeee555555553333
2021.
Abst html. 4
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Proceedings of Neural Information Processing Systems (NIPS), pp. 3753, 2016. URL hittps:/ppapes nips cc/ papppappp/////eeaa /0/0/hhhhh////
da -bt.teea-Abbttttttcccc .tt.trr...ttttl.......... 1, 18
Casper Kaae Sønderby, Jose Caballero, Lucas Theis, Wenzhe Shi, and Ferenc Huszár. Amortised map inference for image super-resolution. In Proceedings International Conference Learning Representations (ICLR), pp. 1-17, 2017. URL https:/ //ppnrevieeeeeeeee/////e net www.mmiiriiioooooo C1RP6G1.1 2,3
https:://penzeeieew
14

---

Published as a conference paper at ICLR 2023
Karl-Theodor Sturm. The space of spaces: curvature bounds and gradient flows the space of metric measure spaces. arXiv: 1208.0434, 2012. 4
Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density Ratio Estimation in Machine Learning, Cambridge University Press, 2012. doi: 10. 10.1017/CB09781 (3900099991111031111 139035613. 2, 3. 17, 19, 20
Valentin Thomas, Emmanuel Bengio. William Fedus, Jules Pondard, Philippe Beaudoin, Hugo Larochelle, Joelle Pineau, Doina Precup, and Yoshua Bengio. Disentangling the independently controllable factors of variation by interacting with the world. In Proceedings of Neural Information Processing Systems (NIPS) Workshop, Learning Disentangled Representations: from Perception to Control, pp. 1-9, 2017. URL https:/ aassweb. ucsd. wdw/-wfedus/ /pdf/ IIF _NPP__ 2017_workshop pdf.
Naftali Tishby, Fernando C. Pereira, and William Bialek. The information bottleneck method. In Proceedings the 37th Annual Allerton Conference Communication, Control, and Computing, pp. 368-377, 1999. URL https: / /wwww CS
http///aaaa/eaaannnggg.ga Paperss/
allerton. pdf. 16
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. In Proceedings of the International Conference on Learning Representations (ICL.R) pp. 1-16, 2018. URL https: wpwpreveeeeeeeeee////ee wwt//orrmn watoorummmddddilnnaao Ob. 2, 4, 5, 8, 9, 18, 19, 27, 29
Jakub Tomezak and Max Welling. Vae with vampprior. In Proceedings of the International
Conference Artificial Intelligence and Statistics (AISTATS), pp. 1214-1223, 2018. URL
walr.presse vv84 www.acck...aa.....mm..m.m
9, 18
Michael Tschannen, Olivier Bachem, and Mario Lucie. Recent advances in autoencoder based repre- sentation learning. In Proceedings of Neural Information Processing Systems (NIPS) Workshop on Bayesian Deep Learning, pp. 1-25, 2018. URL https::///// www. mins ee wete. ch/ bubb/
p/ 20tt 1, 1.3,88 28
Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Proceedings of Neural Information Processing Systems (NeurIPS), volume 33, pp. 19667- 19679, 2020. URL httt::::/prrcceeingg. neurips cctppaper//222//ileeee e #bb21266183cff227776666
6357 wwwpiinnnnn.neuurseeeeeeeeru
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of Ma- chine Learning Research, 9(86):2579-2088 2008. URL htLp:/ttt//mmml org/ ptgpppppss/v///////o
vandermaaten html. 30, 38
Cédric Villiani. Optimal Transport: Old and New. Springer Berlin, 2009. doi: 10.1007/
778-3-540-71050- 3
Hongteng Xu, Dixin Luo, Ricardo Henao, Svati Shah, and Lawrence Carin. Learning autoencoders with relational regularization. In Proceedings of the International Conference on Machine Learn- ing (ICML), pp. 1057660 2020. URL https: http::///occeeeeeeeeeeeeeeeeeeee mlr. wwls.prres//vll xu20e html. 2. 19
Julian Zaidi, Jonathan Boilard, Ghyslain Gagnon, and Marc-André Carbonneau. Measuring disentan- glement: A review of metrics. arXiv: 2012.09276, 2021. 8, 26
Junho Zhao, Yoon Kim, Kelly Zhang, Alexander Rush, and Yann LeCun. Adversarially regularized
autoencoders. In Proceedings of the International Conference Machine Learning (ICML), pp. 5902-5911, 2018. URL https: /ttppppecenngs....//e////o mmlr. presss / vv00 z/zzolbbbb 18b. .tml.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Learning hierarchical features from generative models. In Proceedings of the International Conference Machine Learning (ICML), pp. 4091- 4099, 2017. URL h / /pw.c/edoeeeeeeppoooooooocco.cee ngs m1 r. prepprrsssseee /ttp///0/////// 7c html. 1, 9, 18
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Info Balancing learning and inference in variational autoencoders. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 5885-5892, 2019. doi: 10. /3t.:/aaii..300033331588888 2, 17, 19
15

---

Published as a conference paper at ICLR 2023
Bo Zong, Qi Song, Martin Renquang Min, Wei Cheng, Cristian Lumezanu, Dacki Cho, and Haifeng Chen. Deep autoencoding gaussian mixture model for unsupervised anomaly detection. In Proceedings the International Conference Learning Representations (ICLR), pp. 2018. 2018. URL htttps://openreviiee . htwrrorum/f//e forum? d=BJJLHbbO 1,3,9,30, 38
A DETAILS OF RELATED WORK
For cenf-containment, we describe VAE-based representation learning methods. As with Section 3, x and denote data and latent variables, respectively, and the data x are M-dimeesional and the latent variables Z are T--dimensional. Unless otherwise noted, each VAE-based model consists of generative model po x, z) with parameters 0, an inference model dø (x, z) with parameters Φ, and pre-defined (non-trainable) prior (z) as in the standard VAE model architecture.
A.1 VAE-BASED MODELS WITH ELBO EXTENSION
Utilizing the latent variables of Several models with extended ELBO-based objectives aim overcome the shortcomings of the original VAE model, such as posterior collapse. VAE-based models are mainly grounded on the ELBO objective, where denote the ELBO for the data point
models is a prominent approach to representation Iearning.
(x 0, Φ) (xlx)) [lo (x)z) [2] - (o (()
(16)
which is mentioned as the expected objective of the original VAE (Kingma & Welling, 2014) in Eq. (1).
A.1.1 8-VAE
3-VAE (Higgins et al., 2017a) is VAE-based model for learning disentangled representations by re-weighting the KL term of the ELBO. Given a KKT multiplier > O, the 3-VAE objective is expressed as
maximize 6,
Epasea (xx) [Ego(zix) [x|2] 3D (2x
(17)
The KKT multiplier 3 works as the weight of the regularization to impose factorized prior (e.g., the standard Gaussian (O,, I 1) on the latent variables. This re-weighting induces the capability of disentanglement in the case the latent variables "forget" the information of the input data.
> 1; however, a large value causes posterior collapse, which
From the Information Bottleneck (IB) (Tishby et al., 1999) point of view, the 3-VAE objective is re-interpreted as the following optimization problem (Alemi et al., 2018; Achille & Soatto, 2018):
maximize 0.0
I (z:y)
(18)
subject
(z; Ie.
(19)
where I. is bottleneck capacity, y is a task to be estimated, and (=):)) denotes the mutual information on the inference model. Introducing the Lagrange multiplier B, the IB problem is given
as
maximize 0.
Lo(zyy) - [lo(z;
(20)
Alemi et al. (2018) have given the lower bound of this IB objective as (z (2; (ppaaa(y)go(el)) [log pe (y [z)] H(y) 3 E (xx) (2)
The lower hound of (x:
The upper bound of (x(
(21)
where the task entropy H(y) is independent of the parameters 0 and Φ. The autoencoding task y gives the objective equivalent to that of the original VAE. This (B-based formulation of the 3-VAE objective implies that the larger value of the multiplier guides the training process to minimize the mutual information L (z; x) to make the encoder forget the input data, i.e to cause posterior collapse.
16

---

Published as a conference paper at ICLR 2023
A.1.2 FACTORVAE
FactorVAE (Kim & Mnih, 2018) is state-of-thh-arr disentanglement method that minimizes the Total Correlation (TC) of the aggregated posterior 90(2) ( 4 (z|x)] in addition to the original ELBO objective. The TC expressed as the KL divergence between distribution and its factorized counterpart. In the FactorVAE case, the TC of the aggregated posterior is the KL divergence from the
factorized aggregated posterior (z) = IT l (2()) the aggregated posterior (z)... The training
objective of Factor VAE is the weighted sum of the ELBO and the TC term as
maximize 0.0
ELBO(x; 0, TTC(gp(z),
(22)
where (C(z) denotes the TC of the latent variables 2 defined
TC(z) DKL. (go (2) (1 (z))
(23)
Disc(4 (z)
(x)(() log 1 Disc
(24)
In Eq. (24), Disc(z) denotes discriminator estimate the TC term by density ratio estima- tion (Sugiyama et al., 2012) as
Disc(z) = arg max (x) [log f())] (x) [log(1 - (og(1) f02-1
(25)
Practically, the discriminator estimated using SGD in parallel using samples from To (z) by permuting the latent codes along the batch dimension independently in each latent variable.
A.1.3 INFOVAE
InfoVAE (Zhao et al., 2019) is an extension of VAE to prevent posterior collapse by the retention of data information in the latent variables. The IAE objective is the sum of the ELBO and the inference model mutual information I& in Eq. (19). To this end, the following maximization problem solved via SGD:
maximize 0.4
Epascal (x) [EIBO 0, Φ)] + (x;
(26)
E Epsnta (x) Es (x(x) [x(z]] DKL (20(())
(27)
The main difference between the VAE and InfoVAE objectives is using the regularization term DKL (x))l(( (z)] instead of the original VAE regularization (z)
The orig- inal KI. term becomes zero if all the data points are encoded into the standard Gaussian N( Iz) to cause posterior collapse. The Info KL term DKL (9o [z)][[[ (z))) alleviates this problem by adopting the aggregated posterior go for optimization instead of the encoder go (z x). The authors of VAE (Zhao al., 2019) further provide the model family in which the KL term is replaced with other divergences. They introduce an alternative divergence D(g (z), (z)) and its weight À to conduct representation learning by the following training objective:
maximize 0.4
E (x)sta()) [ELBO(x: 0, Φ)] + (x;
(28)
=B Bppaaa (x) (2|x) (xa(x(z))
XD( (2),
(29)
In the original VAE paper (Zhao et al., 2019), the authors reported that the Maximum-Mean Dis- crepancy (MMD) the best choice for the divergence D. The MMD divergence MMD(1s(z). m (z)) is defined as
(2)(go(Z))T(()) (xs (no [k(z, 2')] + (n) [k(z, 2'] BB Br(xe) [k(z, z')],
(30)
where k(-,-) is any universal kernel, such as the radial basis function kernel
k(zz 2') exp(- ||z z' o2
(31)
for constant 0.
17

---

Published as a conference paper at ICLR 2023
A.2 VAE-BASER METHODS BASED ON HIERARCHICAL FACTORS
Several VAE-based methods postulate the existence of hierarchical factors as its meta-prior to learn representations with the abstractness of different levels (Sønderby al., 2016; Zhao al., 2017). These methods involve the change in their network architecture to utilize the feature hierarchy often captured the hidden layers of deep neural networks.
A.2.1 LADDER VARIATIONAL AUTOENCODER (LADDERVAE)
Ladder Variational Autoencoder (LadderVAE) (Sønderby et al., 2016) introduces hierarchical latent variables to the VAE model. Whereas the objective is still the ELBO, the LadderVAE model structure has hierarchical latent variables. The generative process is modeled as the Markov chain of several latent variable groups. and the inference model consists of deterministic feature encoders and the decoders shared with generative models. In the original paper (Sønderby et al., 2016), the authors claim that the LadderVAF models provide tighter log-likelihood lower bounds than the standard VAE.
A.2.2 VARIATIONAL LADDER AUTOENCODER (VLADDERAE)
Variational Ladder Autoencoder (VLadderAE) (Zhao al., 2017) is a VAE-based model for hierar- chical factors. Instead of the hierarchical models based on Markov chains, the VLadderAE models introduce the hierarchical structure in the network architecture parameterizing the generative and the inference model. Since it constrains feature hierarchy by the process of feature extraction, VLad- derAE also performs disentanglement, e.g., the latent variables from different hidden convolutional layers capture textural or global features of visual data.
A.3 VAE-BASED METHODS INVOLVING PRIOR LEARNING
The standard VAE model has a pre-defined prior, which may cause the discrepancy between the underlying data structure and the postulated prior (Dai Wipf, 2019). Several methods overcome this problem by involving the prior itself in the training process.
A.3.1 VAMPPRIOR
VampPrior (Tomczak & Welling, 2018) is type of prior consisting of the mixture of the encoder distributions from several pseudo-input. The pseudo-inputs are introduced as trainable parameters, which are input into the encoder to build mixture prior. Thus, the VAE models with VampPriors have trainable priors while retaining the main training procedure using the reparameterization trick to apply SGD.
A.3.2 2-STAGE VAE
2-Stage VAE (Dai Wipf, 2019) a generative model with two probabilistic autoencoders. The process of 2- 2-Stage VAE consists of two steps: (i) training standard VAE using the given dataset as the input, and (ii) training another VAE using the latent variables of the previous VAE as the input. The 2-Stage VAE model attempts to overcome the discrepancy between the pre-defined prior and the learned latent representation by introducing the second VAE in stage (ii), which yields the prior training using the VAE in stage (i).
A.4 WASSERSTEIN AUTOENCODER (WAE)
WAE (Tolstikhin et al., 2018) is family of generative models whose autoencoder tries to estimate and minimize the primal form the Wasserstein metric between the generative model pe (x) and the data distribution Paata (x) using SGD with the following objective:
minimize 0.4b
(pa Epant (x (xe [d(x,x)] + XZ(go(z),
(32)
where À is Lagrange multiplier, the generative model is defined as latent variable model po (x, z)
т (z)po (x/z) postulating the prior the latent variables 7 (z), and a conditional distribution g [x)
probabilistic encoder optimize instead of all couplings supported X' X. The WAE
18

---

Published as a conference paper at ICLR 2023
objective is indeed equivalent to that of InfoVAE (Zhao al., 2019) in Eq. (27), which provides the OT-bd perspective on VAE-based models. Following the Info (Zhao et al., 2019), we adopt the MMD for the divergence D, which is denoted by "WAE-MMD" in the original WAE paper (Tolstikhin et al., 2018). Although the WAF-based approaches rewrite VAF-based objectives with the Wasserstein metric, these metrics between x-marginal distributions and do not directly include the latent space Z. To learn representations z, the Wasserstein-bast objective is further modified (Gaujac et al., 2021).
A.5 RELATIONAL REGULARIZED AUTOENCODER (RAE)
Relational Regularized Autoencoder (RAE) (Xu et al., 2020) is variational autoencoding generative model with regularization loss based on the fused Gromov- Wasserstein (FGW) metric. RAE introduces the FGW metric between the aggregated posterior and the latent prior as the regularization divergence to fortify the WAE constraint Te (z) (z(z) introduced by Tolstikhin et al. (2018) for generative modeling. The FGW regularization is introduced with weight hyperparameter 3 E [0,
and given as
minimize 0.0
Epaac (x) Bag x) potx x) [d(x, x') + DDFOw (z) (z));
(33)
where DPGW denotes the FGW metric being the upper bound of the weighted sum of the Wasserstein and Gromov- Wasserstein metrics. The FGW metric DFGW is given as
D DFGW (z), (z)
inf (o)
(11 - B [dz z')] +
22) d2 (2 x))1
=
(34)
2 (1 3)
inf TEP (()
Wasserstein comparison
+ 8 inf (EP(Oo(x).TO(x) Ba 23 (ldd)// dz 2))
(35)
Wassers-a torm for relational comparison
where P(ga(z) (z)) is a set of all couplings whose marginals are 4s (z), TZ( The discrepancy between the prior no (z(z) and the aggregated posterior Q (a(z) causes the degradation of generative performance since the processes of decoding Paata (xant(x)app (x)q (z) x x)po (x x(zz and generation wo (z)pel (x|z) are modeled in different regions of the latent space. This formulation enables learning a prior distribution g (z) with flexibly assuming the structures of data, where the prior To Te(z) is modeled as a Gaussian mixture model the original settings by Xu al. (2020). They aim at matching the distributions on the latent space 2, which can have an identical dimensionality but may differ in terms of distance structure.
A.6 IVT METHODS
Beyond the analytically tractable distributions, implicit distributions applied to variational in- ference. An implicit distribution only requires its sampling method, which extends the variety of modeling and applications in variational inference and VAE-based models.
A. DENSITY RATIO ESTIMATION BY ADVERSARIAL DISCRIMINATORS
The density ratio estimation technique (Sugiyama et al., 2012) is essential to the mechanism of GANs (Goodfellow al., 2014) and IVI methods (Huszár, 2017), which conducted via an optimal discriminator f between distributions r(x) and sxx as
(kL(r(x)[((xx)) Ex(x)) log r(x)
(x)
Ex(x) log 1- (x)
sxx)
Br(x) [log (x) log(1 (((11-x(x))) (x))] where (x) arg
(36)
max (x)))) [log f(x) + Es(s)) [log(1 f(x))]
(37)
/ (0.1)
19

---

Published as a conference paper at ICLR 2023
The discriminator is estimated via maximizing Eq. (37) with neural network f= fr. The training of discriminators often suffers from instability and mode collapse owing to its alternative parameter updates based Eq. (36) and Eq. (37) (Arjovsky & Bottou, 2017; Arjovsky al., 2017). One approach to tackle this problem is imposing the Lipschitz continuity on the discriminator based on the Senborovich-Rubinsser duality (Arjovsky et al., 2017).
A.6.2 ADVERSARIAL VARIATIONAL BAYES (AVB)
Adversarial Variational Bayes (AVB) (Mescheder et al., 2017) is an ELBO optimization method using the adversarial training process instead of the analytical KL term, Let us recall that the KL term in Eq. (1) is defined by the expected density ratio
(z(x)
(Z)xx
E (x(()(x)) T (z)
(38)
Adopting the density ratio trick (Sugiyama et al., 2012), the analytical KL term can be replaced with the optimal discriminator, which takes data point x and its encoder sample z ~ gæ (z(x))) to output the density ratio 4e (rxx)/(()) It enables implicit distributions in the prior while retaining the ELBO objective of variational inference.
A.6.3 ADVERSARIALLY LEARNED INFERENCE (ALI) BIDIRECTIONAL GENERATIVE ADVERSARIAL NETWORKS (BIGAN)
Adversarially Learned Inference (ALI) (Dumoulin et al., 2017) / Bidirectional Generative Adversarial Networks (BiGAN) (Donahue et al., 2017) are models introducing the distribution matching of generative model and the inference model as implicit distributions. These models have been proposed in different papers (Dumoulin et al., 2017: Donahue et al., 2017); however, they share an equivalent methodology. One draw samples from the generative model (z))pe(x(z by decoding prior samples and also from the inference model Paata (x) (z) x) by encoding data points. Here the ALI/BiGAN models introduce discriminator estimate the Jensen-Shanne divergence between the generative model (x(x, z) and the inference model 4s (x, z). The model matching between the encoder and the decoder also learns latent representations by the bidirectional mappings.
A.6.4 VAE-GAN
VAE-GAN (Larsen et al., 2016) is hybrid model based on VAE and GANs. The VAE-GAN models introduce a discriminator for the generative modeling w.r.t. the data and utilize the hidden layers of the discriminator to model the decoder likelihood pe(x z) along the manifolds supporting the data. It provides the outstanding performance of data generation to the VAE framework by measuring the similarity of data utilizing the GANs-like network architecture.
B DETAILS OF PROPOSED METHOD
B. MODELING DETAILS
The decoder (x(z) is modeled with neural network Do 2 ++ RRM and its parameters 0 as p = 5(x (z))))
(39)
Following the standard VAE settings (Kingma & Welling, 2014), the encoder 4ø (z|x) defined as diagonal Gaussian parameterized by neural networks pap R and o Z R with parameters as
qa (z(x)) (Zzluu((x) diag(o (x))).
(40)
For the distance functions dx and dz in Eq. (7) and Eq. (9), we used the La distance defined as
xx)xx')= 2 x x'
(41)
dz(2,z')
1z z'
(42)
/2
20

---

Published as a conference paper at ICLR 2023
As another choice, also utilized the adversarially learned metric (Larsen al., 2016) in Eq. (9). In the adversarially learned metric, the distance is measured in the feature space formed by the hidden outputs of the critic fa. Let hi (x) denote the critic hidden outputs in which the critic takes x as its input. We can then define a distance d' based on the adversarially learned metric as
d' (xxxx) -xxx)xxxxxxxx)))x
1
[x)x(x) hx((x')1
(43)
Since the critic network fu (x, z) has the Y-shaped architecture (see Appendix C.2) and concatenates x- x-based features and z based features one the hidden layers to take a pair (x, z) as the inputs, we use the x-side branch hy (x).
B.2 PRIOR DETAILS
Neural Prior (NP). Formally, the NP (z) with neural network go is defined as:
T / T(E) del Oe x(E) (e(
θgo(E)
de,
(44)
where
(45)
We can implement this class of prior with sampling noises as go(c), avoiding the calculation of the integral.
Factorized Neural Prior. For disentanglement in the variational autoencoding settings, element-wise independence is often imposed on latent variables z. Following the standard VAE settings (Kingma & Welling, 2014), postulate Z Rt. where the latent variables Z are expressed as
L-dimensional vector z = [21, .2..... 1 As with the NP, the FNP class of prior is defined as
(z) = II
To (z+),
(1)
(46)
i=1
-16) (x)) / ))) To
Oge (c())) Oe(i)
where
de
1,2,....
...
(47)
7() = -
1,2....;
(48)
This prior can be implemented with N disjoint neural networks, or 1-dimensional grouped con- volutions. The difference between the NP and the FNP is element-wise independence, in which the prior (z) is factorized into distributions for each latent variable. Factorized priors enable disentanglement by obtaining a representation comprising independent factors of variation (Higgins al., 20172 7a; Chen al., 2018; Kim & Mnih, 2018).
B.3 GRADIENT PENALTY
In the case of gradient penalty (Gulrajani al., 2017), the maximization in Eq. (11) is further modified as
Be (II (Z) 2)1l2 -1)
maximize +
LD +
(49)
where AGP > 0 a constant, and x ex+ and ez (1- are interpolated samples by the random uniform noise e. We adopt AGp 10 in all the experiments reported in this paper. Introducing the gradient penalty together with other techniques such as spectral normalization (Miyato et al., 201 8) is effective and essential for adversarial learning in general (Chu et al., 2020; Miyato et al., 2018).
C EXPERIMENTAL DETAILS
For the reported experimental results, we used a single GPU of NVIDIA GeForce@ RTX 2080 Ti, and single run of the entire GWAR training process until convergence takes about eight hours.
21

---

Published as a conference paper at ICLR 2023
C.I DATASET DETAILS
For the reported experiments in Section 4, we used the following datasets:
MNIST
(LeCun et al., 1998). The MNIST dataset contains 70,000 handwritten digit images of 10 classes, comprising 60,000 training images and 10,000 test images. We used the original test set and randomly split the original training set into 54,000 training images and 6,000 validation images. We used the class information its approximate factors of variation in the form of 10-dimensional dummy variables. This dataset is available online in its original format or via the forchiision package" in the PyTorch (Paszke et al., 2019) tensor format. The MNIST dataset is licensed under the terms of the Creative Commons Attribution-Shan Share Alike 3.0 license
CelebA (Liu al., 2015). The CelebA dataset contains 202,599 aligned face images with 40 binary attributes. We cropped 144 x 144 pixels in the center of the 178 x 218-sized aligned images in the original dataset omit excessive backgrounds. We used the train/ crain/validation/ees partitions that the original authors provided. We used the binary attributes as its approximate factors of variation in the form of 40-dimensional vectors. As in the website this dataset', the CelebA dataset is available for non-commercial research purposes only.
3D Shapes (Burgess & Kim, 2018). The 3D Shapes dataset contains 480,000 8,,0000 synthetic images with six ground truth factors of variation. The images in this dataset contain a single-colored 3D object, single-colored wall of rectangular room, single-colored floor. These images are procedurally generated from the independent factors of variation, floor colour, wall colour, object colour, scale, shape, and orientation (Burgess & Kim, 2018). We randomly split the entire dataset into 384,000/48,000 /tt///88800088,,.,1 8,000 images for the train/validation/tess set, respectively. Since the factor shape is categorical variable in four classes, we converted it into four dummy variables to obtain quantitative factors of variation in the form of 9- dimensional vectors. The repository of this dataset is licensed under Apache License 2.0"
Omniglot (Lake et al., 2015). The Omniglot dataset contains 1,623 images of hand-written charac- ters from 50 different alphabets written by 20 different people. The images are 105 x 105- sized, binary-valued. We used this dataset as OoD samples over MNIST in the evaluations on the OoD detection utilizing cluster structure. The repository of this dataset is licensed under the MIT License?.
CIFAR-10 (Krizhevsky & Hinton, 2009). The CIFARI0 dataset contains 60,000 images with 10 classes, comprising 50,000 training images and 10,000 test images. The images are 32x32 color images in 10 natural image classes, such as airplane and cat. This dataset is provided online without any specific license.
In all the datasets above, we used all the images the raster (bitmap) representation and resized them to 64 x 64 pixels with three channels, where each image is 3 x 64 x 64-sized tensor value and M 12, 288. For gray-scale (one-channeled) images such as in MNIST, we repeated these images along the channel dimension three times to uniform these sizes to 3 x 64 x 64 elements.
C.2 ARCHITECTURE DETAILS
The architecture of neural networks in GWAE and the compared methods are built with convolutions and deconvolution (transposed convolution) in the same settings as shown in Tables 3 and 4. In all the experiments on GWAE, we applied the gradient penalty and the spectral normalization in
hhttp://yann. htt.:/cwn...cm/eee////m//tttttte
html
22

---

Published as a conference paper at ICLR 2023
Table 3: Model architecture for the encoders in the GWAE models and the compared models. For the 64 x 64 RGB images used in the experiments, the input size is set to (Channels, Height, Width) (3, 64, 64). FC and Conv denote fully-connected (linear) layers and convolutional layers, respectively.
Layer
Input Shape
Output Shape
Options
Inverse Sigmoid o (x) = log I-s
Conv
3 x 64 x 64
32 x 32 x 32 kernel size=4, stride=2, padding-1
SiLU activation (Hendrycks & Gimpel, 2016)
Conv
32 x 32 x 32
64 x 16 x 16
kernel size=4, stride=2, padding=1
SiLU activation (Hendrycks & Gimpel, 2016)
Conv
64 x 16 x 16
128 x 8 x 8 kernel size=4, stride=2, padding=1
SiLU activation (Hendrycks & Gimpel, 2016)
Conv
128 x 8 x 8
256 x 4 x 4 kernel size=4, stride=2, padding=1
SiLU activation (Hendrycks & Gimpel, 2016)
FC
256 x 4 x 4
256 SiLU activation (Hendrycks & Gimpel, 2016)
bias=True
FC
256
L for μ. L for o²
bias=True
the critic networks to impose the I-Lipschitz continuity on the critic fe. as shown in Table 7. In the neural samplers of GWAE models, we used the fully-connected architecture in Table 5 for NP and the grouped-convolution unconstrained priors in NP, and I-dimensional grouped convolution layers (converting sequences with length I and L channels) for factorized priors in FNP. For the optimizers of GWAE, we used RMSProp" with learning rate of 10 for the main autoencoder network and used RMSProp with learning rate 5 x 10-6 for the critic network. For all the compared methods except for GWAE, we used the Adam (Kingma & Ba, 2015) optimizer with learning rate of 10 In the experiments, used an equal batch size of 64 for all evaluated models. The batch size is relatively small, since the computational cost of GWAE for each batch is quadratic to the batch size B and the GW estimation runs in time B) for each epoch using [N/B] B batches.
architecture in Table for FNP. We used fully connected layers for
In the case that a batch normalization layer is introduced in the encoder outputs (z(zx))) (x(xx), diag (6 (x))), the mean and variance are computed w.r.t. the aggregated posterior 4s (z) rather than the element-wise sample mean and variance of L-dimensional output values. The normal-
ized parameters (fu(x), (x)) against the original parameters (μ(()), or (x)) are given
m(x) Eg (x) [z)
(x(x)
(50)
V Vao [z]
(x(x)
xx(x)
(51)
(x) [z]
where the division is element-wise conducted, and V denotes the variance. The mean ( [>] and variance Vos (z) [z] are approximated using unbiased estimators consisting of mini-batch samples. Given a batch index set B C (1,2..... N}, the unbiased estimations are expressed using the law of total variance as
1
E [2)(2)11
R
(x;;) =: it,
(52)
#B
iEB
1
V
[2)
(x,) +
#B - Σ (m(x) - 22)2.
(53)
#
iEB
iEB
pdf
23

---

Published as a conference paper at ICLR 2023
Table Model architecture for the decoders in the GWAE models and the compared models. The image shape is set to the same Table 3. FC and DeConv denote fully-connected layers and deconvolutional layers, respectively.
Layer
Input Shape
Output Shape
Options
FC
256
bias=True
SiLU activation (Hendrycks & Gimpel, 2016)
FC
256
256 x 4
bias=True
SiLU activation (Hendrycks & Gimpel, 2016)
DeConv
256 x x 4
128 x 8 x 8
kernel size=4, stride=2, padding=1
SiLU activation (Hendrycks & Gimpel, 2016)
DeConv
128 x 8 x 8 64 x 16 x 16
kernel size=4, stride=2, padding=1
SiLU activation (Hendrycks & Gimpel, 2016)
DeConv
64 x 16 x 16 SiLU activation (Hendrycks & Gimpel, 2016)
32 x 32 32
kernel size=4, stride=2, padding=1
DeConv
32 x 32 32
3 64 x 64 kernel size=4, stride=2, padding=1
Sigmoid (x)
Ite-
Table 5: Model architecture for the samplers in the GWAE models with NP. FC denotes a fully- connected layer.
Layer
Input Shape
Output Shape
Options
FC
L
256
bias=True
SiLU activation (Hendrycks & Gimpel, 2016)
FC
256
256
bias=True
SiLU activation (Hendrycks & Gimpel, 2016)
FC
256
256
bias=Truee
SiLU activation (Hendrycks & Gimpel, 2016)
FC
256
bias=True
Batch Normalization with affine=False
Table 6: Model architecture for the samplers in the GWAE models with FNP. GroupConv denotes
1 1-dimensional grouped convolutional layers.
Layer
Input Shape
Output Shape
Options
GroupConv
256
bias=True, groups=L
SiLU activation (Hendrycks & Gimpel, 2016)
GroupConv
256
256
bias=True, groups=T.
SiLU activation (Hendrycks & Gimpel, 2016)
GroupConv
256
256 bias=True, groups=L
SiLU activation (Hendrycks & Gimpel, 2016)
GroupConv
L Batch Normalization with affine=False
256
bias=True, groups-L
C.3 QUANTITATIVE EVALUATION DETAILS
For quantitative evaluations, we used the DCI scores (Fastwood Williams, 2018) for disentan- glement, the FID score (Heusel et al., 2017) for image generation, and the PSNR score for image reconstruction.
C.3.1 DCI SCORES
The DCI scores (Eastwood & Williams, 2018) measure representation in terms of disentangled representation learning. In the DCI scores, disentanglement measured from three aspects: (i) each representation variable represents single factor of variation, (ii) each factor of variation is expressed
24

---

Published as a conference paper at ICLR 2023
Table 7: Model architecture for the critics in the GWAE models. We concatenated the outputs of the x-side and z-side branches and multiplied the concatenated outputs by 0.5 to input into the stem network for the sake of the gradient norm, resulting in a Y-shaped network. We applied spectral normalization (Miyato et al., 2018) to all the layers in the critic networks and used the LeakyReLU (Maas et al., 2013) activation for the critic to retain the 1-Lipschitz continuity. FC and Conv denote fully-connected layers and convolutional layers, respectively.
Layer
Input Shape
Output Shape
Options
x-side branch
Conv
3 x 64 x 64
8 x 32 x 32
kernel stride==,, padding=1
LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
Conv
8 x 32 x 32
16 x 16 x 16 kernel size=4, stride=2, padding=1
LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
Conv
16 x 16 x 16
32 8 x 8
kemnl size=4, stride=2, padding=1
LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
Conv
kernel size=4, stride=2, padding-1 LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
32 x 8 x 8
64 4 x 4
Conv
128 x 2 x 2 kemee size=4, stride=2, padding=1 LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
64 x 4 x 4
Conv
kernel size=4, stride=2, padding=1 LeakyReLUU activation (Maas et al., 2013) with negative slope 0.2
128 x 2
256 1
FC
256
64
bias=True
z-side branch
FC
bias=True LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
L
256
FC
256
256
bias=True
LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
FC
bias=True LeakyReL.U activation (Maas et al., 2013) with negative slope 0.2
256
64
Stem network
z-side branch
FC
64+64
256
bias=True
LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
FC
256
256
bias=True
LeakyReLU activation (Maas et al., 2013) with negative slope 0.2
FC
bias=True LeakyReL.U activation (Maas et al., 2013) with negative slope 0.2
256
by a single representation variable, and (iii) representation is informative w.r.t the original data. The correspondence of variables and factors is computed via estimating the ground truth factors from the representation using random forest (Breiman, 2001). DCI Disentanglement (DCI-D) measures (i) the factor singleness for each variable. DCI Completeness (DCI-C) measures (ii) the variable singleness for each factor. DCI Informativeness (DCI-I) measures (iii) whether the representation is informative for estimating the ground truth factors. These metrics are computed via the variable importances (e.g., the Gini impurity (Breiman, 2001)) of the random forest (Breiman, 2001), in which the random forest regressor estimates the ground truth factors using the representation variables. Using L-dimensional representation variables 2, Vddimensional factors y and their importance Rig of the i-th variable z;
25

---

Published as a conference paper at ICLR 2023
for the k-th factor yks the DCI-D and DCI-C scores for each variable and each factor are defined
DCI-D, = ++ Σ
Pik logv Pik:
(i = 1,2, ......,,, ,,
(54)
where Pik
Rich
Σ
Raj
(55)
j=1
L
DCI-C& = 1
Σ
Tik logv qikr
(k (k= 162,.......
(56)
V Σ
where qik
Raxe
RzK
(57)
The DCI-D score for the entire variable set is given by the weighted sum E Pj P;DCI-D; where the weight P: is weighted importance P: = =Ex=1 Rat/ E/ Rax). The DCI-C score for the entire factor set is given by the average score 1/V E DCI-Cx. The DCI-D and DCI-C metrics take values within the range [0, 1]. where higher values indicate better performance. For DCI-I, we used the normalized definition by Zaidi et al. (2021) because the normalized DCI-I values are within the range [0, 1] and the higher values mean better informativeness, while DCI-I score DCI-Doriginal is the estimation mean square error in the original definition. The DCI-I definition that we used is expressed
DCI- Following the original paper (Eastwood & Williams, 2018), we set the number random trees to 10 and decided the tree depth with noss-validation.
DCI-Ioriginal-
(58)
C.3.2 FRÉCHET INCEPTION DISTANCE (FID)
Fréchet Inception Distance (FID) (Heusel et al., 2017) is a score for evaluating the quality of the generated images by generative models. The FID score defined as the squared 2-Wasserstein metric between the features of the real images with mean (Mr, Σ,) and that of the generated images with mean (Hg.. Eg). Assuming that the features are normally distributed in the feature space, the FID score is expressed as
FID W3 E.).N Eg))
(59)
= |4,
+ -
(60)
Since the Wasserstein metric measures the discrepancy between distributions, lower values indicate better generation performance in the FID score. Following the original FIDD paper (Heusel al., 2017),
we used the features obtained from the final pooling layer outputs of the Inception-và v3 pre-trained in the ImageNet dataset (Deng et al., 2009).
C.3.3 PEAK SIGNAL-TO-NON NOISE RATIO (PSNR)
For measuring the image reconstruction, used the Peak Signal-to-Noise Ratio (PSNR) value. The PSNR value is defined as
PSNR = 20 log (MAX) 10 log (MSE), where MAX denotes the maximum value of the pixel values, and MSE indicates the mean square error (MSE). all the experiments conducted in Section 4, the value of MAX set to MAX = because the images input as dataset data are scaled within the range [0, 1].
(61)
C.4 ISOMETRY COMPARISON
Regarding the evaluations Section 4.2, we further conducted comparisons on isometry in Fig. The results show that the GWAE models provide more isometric autoencoders compared with other VAE-based representation learning methods. The existing VAE-base based methods did not yield as far GWAE, which supports that the GW metric works different objective class from the ELBO. This implies that the GW metric loss substantially affects the training procedure of learning representations.
26

---

Published as a conference paper at ICLR 2023
60
2000
1750
52
2000
1500
1200
1000
1000
10
(a) VAE (Kingma & Welling, 2014).
(b) FactorVAE (Kim & Mnih, 2018).
50
45.00
60
4000
3050
22
3
2530
30
15:0
20
1000
12
1000
(c) 2-Stage VAE (Dai Wipf, 2019).
(d) WAE (Tolstikhin al., 2018).
3000
130
2000
1500
1000
20
20 40
20
100
GWAE
Figure 4: Histograms of the differences in MNIST (LeCun et al., 1998). Each histogram consists of 10,000 samples of (Ax, Az), where Ar (vertical) and Az (horizontal) respectively denote the differences Ar dx x') and A= dz (z, z') of two generative samples (x. z), (x', z') ~ [x(x, z). In all reported results including Factor (Kim Mnih, 2018) (y=3), WAE (Tolstikhin al., 2018), and GWAE (NP, \ >w=l, Ax=1), the latent dimension L was set to L 16, and their priors were set to the standard Gaussian.
27

---

Published as a conference paper at ICLR 2023
Table The effect of prior familyy selection in the GWAE model. The same settings in Table 1 applied to all the reported models.
Model
DCI-C t
DCI-D 1
DCI-1
GWAE (NP)
0.3966
0.3113
0,9403
GWAE (FNP)
0,9080
0.7024
0.9966
GWAE (GMP)
0.4247
0.4373
0.9655
C.5 TRAINING PROCESS
We present the training process of GWAE in Fig. 5. Although the objective seems complex for its composition of four different losses, the training process successfully converged and the values of the terms Lcw. LW, and LD jointly descended in the most part of training. Although the term R increased, its values did not diverge to prevent the degenerate solutions. These results imply that the three different losses Law. Lw, and Ln did not conflict during the training process even for the complicated data, balancing these three terms against - RR in the trade-off of the reconstruction against the regularization in 3-VAB (Higgins al., 2017a; T'channen al., 2018).
C.6 PRIOR FAMILY SELECTION
We show the effect of prior family selection regarding a meta-prior, disentanglement, in Table While GWAE models with the NP and GMP retain the informativeness of the FNP, the other two priors than FNP did not comparably disentangle the latent factors. Although the NP covers a more general family of prior, these results suggest that choosing a prior family suitable to the postulated meta-prior greatly facilitates learning representations.
C.7 QUALITATIVE EVALUATIONS OF GENERATION AND RECONSTRUCTION
We show the reconstructed images by GWAE and state-of-the-art variational autoencoding methods in Fig. 6. The shown images are the first ten samples of the test split in the CelebA (Liu et al., 2015) dataset under the latent size L 64. Compared with the other methods, the reconstruction of the GWAE model tends to retain edges (see the bottom rows of Fig. 6), while VAE-based models generate smooth, blurry images due the noise injected in the latent space to perform probabilistic modeling and manifold learning. We also show the reconstruction results of MNIST (LeCun et al., 1998) in Fig. 7 and CIFAR-10 10 (Krizhevsky the GWAE models consistently perform autoencoding also in more simple dataset (MNIST). In a more complex dataset (CIFAR-1O), the GWAE model attained the best evaluation generation albeit its reconstruction. suggesting that the GWAE model successfully captured the abstract structure of data rather than reconstructed the given images. This difference highlights the difference in their objectives, i.e,, the GW objective aims at distribution matching in the latent space, while the 3-VAE (Higgins et al., 2017a) objective with B < 1 puts weight on reconstruction.
Hinton, 2009) in Fig. 8. These results support that
We further study the generated images by GWAE and state-of-the-arr VAE-based generative models in Fig. 9. These qualitative results show that the GWAE generation successfully obtains diverse of images compared with those of state-of-the-art autoencoding generative models. The ALI model (Dumoulin et al., 2017) (Fig. 9 (a)) also generates various images by the distribution matching of bidirectional models, but the generated images have wavy contours, failing composing images with a consistent appearance owing to the lack of an autoencoding process. Although the VAE- GAN model (Larsen al., 2016) (Fig. 9 (b)) adequately yields organized images with smooth textures, the azimuth of these images is less diverse, i.e., the great majority of the images are facing forward or looking slightly sideways. The images generated by 2-Stage VAE (Dai Wipf, 2019) (Fig. 9 (c)) have diverse azimuth, color, and background; however, these images tend to incline toward the majority attributes, e.g., not wearing eyeglasses or sunglasses. The GWAE model (Fig. 9 (d)) successfully generates facial images with various skin colors, diversified backgrounds, and assorted facial expressions (e.g., wearing a mustache). These results imply that the GWAE models also function as generative models while has been built representation learning method owing the
28

---

Published as a conference paper at ICLR 2023
Table 9: The ablation study on generation and reconstruction in CelebA (Liu al., 2015). The same
settings as Table 2 are applied in these experiments.
Model
FID (Heusel et al., 2017) 4
PSNR [dB] f
GWAE (NP)
45.3
22.82
GWAE (NP) w/o LW
233.7
9.80
GWAE (NP) w/o LD
403.8
18.63
GWAE (NP) w/ MMD LD
158.4
22.61
GWAE (NP) w/ Z-only critic
102.4
22.41
GWAE (NP) w/o RH
179.6
21.57
GWAE (NP, p = E
123.5
16.03
collateral condition po (x z) se qxl(x, 2) in Eq. (11) and the generative modeling po(x) R3 Pdaza its necessary condition.
C.8 ABLATION STUDY
We conducted the ablation study of the losses and regularizations introduced in Eq. (13). Table 9 shows the results the ablation study the three sub-constraints Lw, Ln, and Rx. The ablations yielded the performance degradation of GWAE, especially in Lw. These results suggest the necessity of each regularization term and reveal their roles in representation learning.
Ablation of Lw. The ablation of the term Lw brought low-quality reconstruction, which suggests that Lw works as the autoencoding constraint as can be seen from taking the reconstruction loss in Lw. It also reduced generation capability as well reconstruction, suggesting that the gen- erative modeling via autoencoding is inherited from the variational autoencoding architecture of VAEs (Kingma & Welling, 2014).
Ablation of Lp- Without the term LD. the GWAE models suffer from the lack of distribution matching in data generation, while it successfully conducted data reconstruction. These phenomena
could be caused by the discrepancy between the encoded latent distribution (œ(z) and the prior Te (z) Similar results are also obtained the ablation of the merged sufficient condition (see Eq. (11)) for the regularization LD, where LD is defined as the MMD loss between the prior Te(z) and the encoded latent 96 (z), as in the WAE-MMD model (Tolstikhin et al., 2018). This choice of Lp on the diw-dimensional space Z appears to be a replacement for the Kantorovich potential adversarially learned in the high-dimensional joint space X' 2; however, lacking the merged sufficient condition seems to have caused the crucial reduction of generation performance in the gross ablation of CD. These results imply that the term LD with adversarial learning regularizes the generative model pe xx to match the inference (x,
Ablation of Ry. Removing Rx slightly increased the reconstruction error but deteriorated the generation quality. To confirm this behavior, also show the samples generated by the GWAE model without the regularization Rx in Fig. 10 and its reconstruction in Fig. 11. These qualitative results that the decoder without Rx successfully reconstructs the images from the inference 4ø (z) but generates corrupted images from the prior (z).. It suggests the "hole" problem (Rezende & single latent point to
Viola, 201 in the degenerate solution, where each data point is mapped cover the zero-measure area of the latent space and the latent space is almost everywhere not covered by the inference g (z). Thus, the entropy regularization R. seems have worked for retaining the probabilistic mappings in the encoder 10 (z)x) to avoid this phenomenon.
In addition, for ablating p intuitively natural although causing unstable training process due to the outlier samples in Lew. The GWAE model with p suffered from performance degradation both in the generation and reconstruction, suggesting that our settings
we also experimented with the p & settings that appear to be
< & affect the learning process of the entire model.
C.9 THE META-PRION EFFECT ON GW MINIMIZATION AND ESTIMATION
For further inspection of Section 4.2, we also studied the GW minimization and estimation using FNP in Fig. 12. Compared with the NP case in Fig. GWAE with FNP presents less stable and
29

---

Published as a conference paper at ICLR 2023
more biased estimation and minimization. The learning curve of Lew in Fig. 12a is largely biased in the first 40 epochs and then seems to be converged at approximately 3.2, a higher value than that of Fig. la (lower than 2). The isometry histogram also suggests the degradation of GW minimization in FNP. In Fig. 12h, more samples fell in off-diagonal areas, showing that the isometry is less tight than that of Fig. 1b. These results presumably due the mismatch of disentanglement meta-prior in MNIST (LeCun et al., 1998) because one of the major generative factors of MNIST images is the kind of digits, categorical variable typically learned one-hot variables in contrast the factorization imposed by FNP.
C.10 PRIORS IN CLUSTERING STRUCTURE
For more detailed investigation of the capture of clustering structure studied in Fig. 3, we further study the latent spaces of VAE (Kingma & Welling, 2014), DAGMM (Zong al., 2018), and GWAE with GMP. The t-SNE visualization (van der Maaten & Hinton, 2008) of the latent spaces are shown in Fig. 13, which suggests that the GWAE model with GMP clearly captured the clustering structure in its latent space. The prior of VAE (Kingma & Welling, 2014) 4) is defined as the standard Gaussian (O, I.) which does not consist of multiple clusters. The learned prior of DAGMM contains multiple clusters; however, adjacent clusters were overlapping to some extent. From the learned prior in GWAE, can observe clear clusters densely concentrating themselves and separating each other. These results support the quantitative OoD results in Fig. 3, in which the GWAE model outperforms the other two models with and without explicit clustering modeling, respectively.
30

---

Published as a conference paper at ICLR 2023
Tra station
Trannig
400
300
phyy Mmm mw M
23
40
100
0
20
105
(a) Law, CelchA
(b) Lw. CelchA.
400
860
Traiking
Unlldo Caliooiion
(c) LD. CelebA.
(d) - RRH. CelebA.
800
Trunesg Val
500
700
500
400
(c) LGW, CIFAR- 10.
(f) Lw. CIFAR-10.
Thaei station
50
Trantis
M 10M0
700
22
10
200
100
20
62
20
102
(g) LD. CIFAR-10.
(h) RRx. CIFAR-10.
Figure 5: The training process of GWAE model. The model is trained using NP and Xy = Aw Au = The plot (a)-(d) are training curves during one trial of training using CelebA (Liu al., 2015), and (h) are using CIFAR-10 (Krizhevsky & Hinton, 2009). In each plot, the horizontal axis represents the number epochs elapsed, and the vertical axis expresses the loss value. The blue and orange curves represent the training and validation losses, respectively.
31

---

Published as a conference paper at ICLR 2023
(a) VAE.
29
(b) 9- (Higgins et al., 2017a) 1...
(c) AVB (Mescheder et al., 2017) (8=1).
2
(d) WAE (A=100).
25
(e) GWAE (NP, Ap=L, Aw=10, 0==0000
Figure 6: Reconstructed images in CelebA (Liu et al., 2015). The images denote original data samples (top rows), reconstructed images (middle rows), and zoomed reconstructions (bottom rows). Each column corresponds to one data instance in the test set.
32

---

Published as a conference paper at ICLR 2023
7
2
/
0
4
5
9
7
42/0414
9
5
9
(a) VAE. FID: 16.8, PSNR: 23.66 dB.
7
2
/ 6 4
1 4
6
9
7
4
(b) 3-VEE (Higgiin et al., 2017a) (8=0.1). FID: 15.5, PSNR: 25.45 dB.
7
4
4 5
9
7
72/0415
4
9
9
(c) AVB (Mescheder et al., 2017) (8=1). FID: 39.2, PSNR: 24.31
7
4
O 4
1 4
9
7
9
5 9
(d) WAE (A=100), FID: 16.9, PSNR: 25.28 dB.
7 2 /
O 4
1 4
9
7 7
9 9
9
(c) GWAE (NP, Ap=l, Aw=10, 0.000001.. FID: 14.4, PSNR: 26.11 dB.
Figure 7: Reconstructed images in MNIST (LeCun al., 1998). The images denote original data samples (top rows), reconstructed images (bottom rows). Each column corresponds to one data instance in the test set.
33

---

Published as a conference paper at ICLR 2023
(a) VAE, FID: 111.3, PSNR: 19.84 dB,
(b) 3-VA (Higgins et al. 2017a) (8=0.1). FID: 84.5, PSNR: 22.48 dB.
(c) AVB (Mescheder et al., 2017) (8=1). FID: 109.9, PSNR: 21.14 dB.
(d) WAE (A=100). FID: 87.3, PSNR: 22.45 dB.
(c) GWAE (NP, Ap=1, Aw=10, 0.0000.... FID: 59.9, PSNR: 17.64 dB.
Figure 8: Reconstructed images in CIFAR-10 (Krizhevsky & Hinton, 2009). The images denote original data samples (top rows), reconstructed images (bottom rows). Each column corresponds to one data instance in the test set.
34

---

Published as a conference paper at ICLR 2023
ALI (Dumoulin et al., 2017).
(b) VAE-GAN (Larsen et al., 2016) (o=1).
(c) 2-Stage VAE (Dai Wipf, 2019).
(d) GWAE (NP, Ap=1, A1= 0000000.
Figure 9: Generated images in CelebA (Liu et al., 2015). We show 100 images sampled from the generative model (x) without conducting therry-picking.
35

---

Published as a conference paper at ICLR 2023
Figure 10: Generated images in CelebA (Liu et al., 2015) using the GWAE model without the regularization term Ry.
7
Figure 11:
Reconstructed images in CelebA (Liu et al., 2015) using the GWAE model without the regularization term Ry. Each column corresponds to one test data instance. The rows denote original (top) and reconstructed (bottom) images.
36

---

Published as a conference paper at ICLR 2023
4.4
Estimated Cow
Empirical GW
Boomo-WWaseesten Value
M
4.0
3.8
3.4
3.2
20
40
60
80
100
Number of Epochs Trained
(a) The estimation the GW metric using FNP.
-60
50
2000
40
1500
30
1000
20
500
10
20
40
60
80
(b) The isometry in GWAE with FNP.
The estimation and minimization of the GW metric. This trial of training is conduced in GWAE (FNP, Ap=1, > )x=1) using the MNIST (LeCun et al., 1998) dataset, which the same settings as Fig. except for FNP. (a) The curves show the GW values estimated by the loss term Low (solid, blue) and the empirical GW computed by the POT package (Flamary et al., 2021) (dashed, orange). The values computed using the validation set. (b) The axes Ar dx (x, x') (vertical) and A= = dz (z, z') (horizontal) respectively denote the difference in the data and latent spaces between generated samples (x, z), (x', z') po (x z). The histogram contains 10,000 generated sample pairs.
Figure 12:
37

---

Published as a conference paper at ICLR 2023
30
t-SNE Dimensio
LINN
-20
8
-30
Dimension
18NE Dinarsnnn
(a) VAE (Kingma & Welling, 2014).
40
I-SNE Dimension 2
20
JONE Dereminn
-20
-30
-40
-20
20
t-SNE Dimension
54NE Dinwers
(b) DAGMM (Zong et al., 2018).
40
L-SNE Dimension 2
0
:80
-40
-20
t-SNE Dimension
1ENE Dinoreoo
(c) GWAE (GMP).
Figure 13: The t-SNE visualizations (van der Maaten & Hinton, 2008) for latent space samples
TØ(z) for the OoD detection in Fig. 3. The left plot presents the sampled points of the t-SNE embeddings, and the right one presents the kernel density estimation (KDE) of these embeddings. The sample size is equally 1 1,024 in each reported model.
38


Reviews:
[{"review_id": "r8ntBQ3HLnk", "reviewer": "Reviewer_JFYx", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "A representation learning scheme is proposed using the Gromov-Wasserstein (GW) distance, a variant on optimal transport distances which allows comparison between distributions in different metric spaces. An auto-encoder based on the GW distance is proposed and evaluated on common image datasets.  Empirical tests on disentangling and clustering are made.\n", "summary_of_the_review": "This work is a natural successor to many OT based generative modeling works existing in the literature. The empirical advantage of the proposed method appears limited with regards to generative quality measures, but does some improvements in disentanglement. I would like to see more discussion of cost of estimating GW. Nevertheless this work appears to me to be a worthwhile contribution and therefore I support acceptance.\n", "strength_and_weaknesses": "[+] Convincing empirical evidence that GW metric is indeed estimated (Figure 1a).\n\n[+] Good evidence that proposed method quantitatively improves disentanglement scores (Table 1).\n\n[+] Comprehensive review of literature.\n\n[-] Unclear empirical advantage of GWAE for sample generation task (Table 1). Prior work 2-stage VAE appears better.\n\n[-] No discussion of computational cost or scalability of GW estimation. Isn't L_GW (Eq. 8) quadratically expensive?\n\n[-] Only (relatively) small scale datasets are considered (CelebA, 3D Shapes, Omniglot). Is this related to the cost of estimating GW? More discussion is desired.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "kPrh0kWD82", "reviewer": "Reviewer_ZoWi", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper proposes an integrated VAE-based approach for representation learning by minimizing Gromov-Wasserstein distance between the metric-measure space of the data and latent vectors. Isometric encoding is obtained from the trained prior and their performances in meta-priors (e.g. disentanglement and clustering) are empirically demonstrated without additional penalty terms. In addition, the model has good generation performance although its goal is matching latent and data distributions.\n", "summary_of_the_review": "An autoencoder model that approximately minimizes the Gromov-Wasserstein distance seems novel, although its derivation needs further justification. The proposed model is extensively compared with other models. ", "strength_and_weaknesses": "## Strengths \n\nThis paper proposes a method that can perform disentanglement and clustering tasks within a single model. The objective for the proposed model is intuitively derived. The performance of the proposed method is extensively demonstrated with many supporting experiments. The effect of each regularization term is experimentally confirmed.\n\n## Weaknesses\n\n1. Merging Eqs. (8) and (10) into the constraint $p_{\\theta}(x,z) = q_{\\phi}(x,z)$ needs a justification. The latter implies Eqs. (8) and (10), but the vice versa. I wonder if you can explain further about the advantages of using $\\mathcal{L}_D$ in (11) compared to the maximum-mean discrepancy (MMD) between the aggregated posterior and the prior.\n\n2. That the generated CelebA images preserves the details of hair is encouraging, but this may be a consequence of adding high-frequency noise to the generated smooth image. In fact, other parts of the faces exhibit spotty noise.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "fbkGsRCO8iP", "reviewer": "Reviewer_PxfM", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The authors propose Gromov-Wasserstein Autoencoders (GWAE). The application of the GW metric allows to match data distributions with a given learnable prior, even when the distributions lie in spaces of different dimensions, thus allowing matching between latent and data space (in equation 4). Further, the GW objective allows rich meta-priors for flexible representation learning. Experiments on disentanglement and clustering using benchmark datasets (CelebA, MNIST and 3D Shapes) demonstrate the capabilities of the proposed approach.", "summary_of_the_review": "The authors offer an interesting and seemingly flexible approach to optimize variational autoencoders by leveraging the GW metric, which allows direct matching between observed and latent spaces, however, not without the help of complementary objective functions, specifically, marginal matching given the approximate posterior (WAE), critic-based joint distribution matching (KR), and entropy regularization to alleviate collapsing.", "strength_and_weaknesses": "The authors present a new perspective of using GW for the construction of flexible variational autoencoders. For the most part the proposed approach is well motivated and properly justified, though requiring three auxiliary and complementary (based on the ablation studies) loss functions. Though mostly presented in the supplementary material, the authors present a thorough overview of the variational autoencoding landscape.\n\nThe work has some key weaknesses:\n\nThe justification for \\rho=1 citing alleviating outlier effects is not necessarily convincing, more so when in (9) they adopt a \\zeta=2 reconstruction loss, so why are outliers a concern in (7) but not in (9)?\n\nThough the ablation study is a welcome addition to the experiments as they demonstrate the contribution of L_W, L_D and R_H to (13), it will be important to include results without L_GW.\n\nThe results in Figure 2 though impressive at first raise some questions: i) it seems that both VAE and GWAE separate (disentangle) object hues in 2 out of the 16 dimensions considered, and ii) even though it seems that the x-axis in Figure 2c captures the object hue, either it does not preserve the order (see color map) or there is no ordering in the hue, in which case one dimension captures the hue, but one wonders how consistently across runs and model hyperparameters.\n\nIn Table 2, it is not clear how biased are the results considering that hyperparameters were selected to optimize DCI-C that is used as performance metric and happens to show the largest difference with existing methods.\n\nFrom Table 2 it is not clear why the authors could not obtain results for WAE considering the closeness with the proposed approach. Further, the table seems to indicate that FID was taken from the original papers but not PSNR? \n\nThe results for clustering structure are interesting, however, are probably not the best way to showcase the clustering capabilities of the prior used in the proposed approach. In fact, as the authors show, a relatively simple VAE lacking a clustering prior achieves near perfect outlier detection.\n\nThe authors do not discuss why for Section 2 they use NP exclusively. It will be interesting to see results for GW estimation and minimization using FNP.\n\nMinor points:\n\nThe marginal p_\\theta(x) in (8) is not defined.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "__y092qXE9p", "reviewer": "Reviewer_y4yE", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper is an extension of variational auto-encoders that uses a different type of metric for performing training. The suggested metric directly matches the latent and data distributions using the variational autoencoding scheme; this is based on the Gromov-Wasserstein (GW) metric between the trainable prior and given data distributions.\n\n", "summary_of_the_review": "The method is interesting and sound, though marginally novel. The experiments show promise and depth, but we need more datasets for these to be convincing.", "strength_and_weaknesses": "Pros: \n1. The authors perform an extensive ablation study of the method which convinces me that the shown results are not statistical artifacts.\n2. The method is based on a solid theory which is motivated well in the manuscript. \n3. The derivations, although straightforward, are correct. \n\n\nCons:\n1. The idea is a delta upon the large corpus of VAEs; importantly, it heavily builds upon WAE.\n2. Even though the method shows promise in the considered datasets, I think that much more datasets are needed for the empirical evidence to be convincing.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster