Learning disentangled representations with the Wasserstein Autoencoder
Authors: Benoit Gaujac
Ilya Feige
David Barber
Venue:  (2021)
Keywords: generative modeling
disentangle learning
wasserstein autoencoder

Abstract: Disentangled representation learning has undoubtedly benefited from objective function surgery. However, a delicate balancing act of tuning is still required in order to trade off reconstruction fidelity versus disentanglement. Building on previous successes of penalizing the total correlation in the latent variables, we propose TCWAE (Total Correlation Wasserstein Autoencoder). Working in the WAE paradigm naturally enables the separation of the total-correlation term, thus providing disentanglement control over the learned representation, while offering more flexibility in the choice of reconstruction cost. We propose two variants using different KL estimators and perform extensive quantitative comparisons on data sets with known generative factors, showing competitive results relative to state-of-the-art techniques. We further study the trade off between disentanglement and reconstruction on more-difficult data sets with unknown generative factors, where we expect improved reconstructions due to the flexibility of the WAE paradigm.

Under review as a conference paper at ICLR 2021
LEARNING DISENTANGLED REPRESENTATIONS WITH THE WASSERSTEIN AUTOENCODER
Anonymous authors
Paper under double-bling review
ABSTRACT
Disentangled representation learning has undoubtedly benefited from objective function surgery. However, a delicate balancing act of tuning still required in order to trade off reconstruction fidelity versus disentanglement. Building on previous successes of penalizing the total correlation in the latent variables, we propose TCWAE (Total Correlation Wasserstein Autoencoder). Working in the WAE paradigm naturally enables the separation the total-correlation term, thus providing disentanglement control over the learned representation, while offering more flexibility in the choice reconstruction cost. We propose two variants using different KL estimators and perform extensive quantitative comparisons data sets with known generative factors, showing competitive results relative to state-of-the-art techniques. We further study the trade off between disentanglement and reconstruction on more-difficult data sets with unknown generative factors, where the flexibility of the WAE paradigm in the reconstruction term improves reconstructions.
INTRODUCTION
Learning representations of data is at the heart of deep learning; the ability to interpret those representations empowers practitioners improve the performance and robustness of their models (Bengio et al., 20133 van Steenkiste et al. 2019). In the case where the data is underpinned by independent latent generative factors, a good representation should encode information about the data in semantically meaningful manner with statistically independent latent variables encoding for each factor. [Bengio al. (2013) define a disentangled representation as having the property that a change in one dimension corresponds to change in one factor of variation, while being relatively invariant to changes in other factors. While many attempts to formalize this concept have been proposed (Higgins) ct al. 2018: Eastwood & Williams, 2018 Do & Tran 2019). finding a principled and reproducible approach to assess disentanglement is still an open problem (Locatello et al.) 2019).
Recent successful unsupervised learning methods have shown how simply modifying the ELBO objective, either re- weighting the latent regularization terms or directly regularizing the statistical dependencies in the latent, (2017) and Burgess et al. (2018) control the information bottleneck capacity of Variational Autoen- coders (VAEs, (Kingma & Welling 2014, Rezende et al. (2014) by heavily penalizing the latent regularization term. Chen et al. (2018) perform ELBO surgery to isolate the terms at the origin of disentanglement in 3-VAE, improving the eiconstrccioo-ddientaaggeeeee
be effective in learning disentangled representation. (Higgins al.
trade off. [Esmaelli et al. (2018) further improve the reconstruction capacity of 3-TCVAE by introducing structural depen- dencies both between groups of variables and between variables within each group. Alternatively, directly regularizing the aggregated posterior to the prior with density-free divergences (Zhao et al.) 2019) or moments matching (Kumar et al. 2018), or simply penalizing high Total Correlation (TC, (Watanabe] 1960)) in the latent (Kim & Mnih 2018) has shown good disentanglement performances.
In fact, information theory has been a fertile ground to tackle representation learning Achille & Soatto (2018) re-interpret VAEs from an Information Bottleneck view (Tishby et al.) 1999), re-phrasing as a trade off between sufficiency and minimality of the representation, regularizing pseudo TC between the aggregated posterior and the true conditional posterior. Similarly, Gao et al. (2019) use the principle of total Correlation Explanation (CorEX) (Ver Steeg & Galstyan 2014) and maximize the mutual information between the observation and subset of anchor latent points. Maximizing the

---

Under review as a conference paper at ICLR 2021
mutual information (MI) between the observation and the latent has been broadly used (van den Oord et al. 2018 Hjelm et al. 2019 Bachman et al. [019 I'schannen et al. 2020), showing encouraging
results in representation learning. However, Tschannen et al. (2020) argued that MI maximization
alone cannot explain the disentanglement performances of these methods.
Building on the Optimal Transport (OT) problem (Villani 2008 Tolstikhin al. (2018) introduced the Wasserstein Autoencoder (WAE), an alternative to VAE for learning generative models. Similarly VAE, WAE maps the data into a (lowddimensiona)) latent space while regularizing the averaged encoding distribution. This is in contrast with VAEs where the posterior is regularized at each data point, and allows the encoding distribution to capture significant information about the data while still matching the prior when averaged over the whole data set. Interestingly, by directly regularizing the aggregated posterior, WAE hints at more explicit control on the way the information encoded, and thus better disentanglement. The reconstruction term of the WAE allows for any cost function on the observation space, opening the door to better suited reconstruction terms, for example when working with continuous RGB data sets where the Euclidean distance any metric the observation space can result in more accurate reconstructions of the data.
In this work. following the success of regularizing the TC in disentanglement, we propose to use the Hollback-Leibblee (KL) divergence as the latent regularization function in the WAE. We introduce the Total Correlation WAE (TCWAE) with explicit dependency on the TC of the aggregated posterior. Using two different estimators for the KI. terms, we perform extensive comparison with succesful methods number of data sets. Our results show that TCWAEs achieve competitive disentanglement performances while improving modelling performance by allowing flexibility in the choice of reconstruction cost.
2
IMPORTANCE OF TOTAL CORRELATION IN DISENTANGLEMENT
2.1 TOTAL CORRELATION
The TC a random vector E Z under defined by
d2 Hp (Za) Hp(Z)
TC(Z) &
where Pa(za) is the marginal density over only and Hp(Z) -Ep p(Z) is the Shannon differential entropy, which encodes the information contained in Z under P. Since
dz
Hp. H2(2)
with equality when the marginals Za are mutually independent, the TC can be interpreted as the loss of information when assuming mutual independence of the Zœ: namely, it measures the mutual dependence of the marginals. Thus, in the context of disentanglement learning, we seek low TC of the aggregated posterior, p(z) = [x p(2(x) p(z) dr, which forces the model to encode the data into statistically independent latent codes. High MI between the data and the latent is then obtained when the posterior, p(=(z)) manages to capture relevant information from the data.
2.2 TOTAL CORRELATION IN ELBO
We consider latent generative models (=()) = fe Po(x) 2) p(=) da with prior p(2) and decoder net- work, (=(2)2), parametrized by VAEs approximate the intractable posterior p(2[2) by introducing an encoding distribution (the encoder), (=(x))) and learning simultaneously 0 and @ when optimizing the variational lower bound, or ELBO, deñined in 3
LELBO ()
E
E
p((2)) s E plog (X|2) KL KL (Z (L(x(Z(X)
log pe (X)
Pa (x) 40 (2(2)x)
(a)
(3)
2

---

Under review as a conference paper at ICLR 2021
Following Hoffman & Johnson (2016), decompose the KL term in Eq. 3 as:
N
KL (z)))))) p(()) = KL (9(2.)) (l + KL a(2) (1
ll
(4)
Nbatch
① index-code MI
marginal KI.
where p(n) = t. ((z])) = g(zlzn), q(2,n) ((=]n)((n) and ((2) = = g(=(n) p(n).
refers to the index-code mutual information and represents the MI between the data and the latent under the join distribution g(2,1) and ni to the marginal KL matching the aggregated posterior to the prior. While discussion on the impact of a high index-code MI on disentanglement learning is still open, the marginal KI. term plays important role in disentanglement. Indeed, it pushes the encoder network to match the prior when averaged, as opposed to matching the prior for each data point. Combined with factorized prior p(z) -I Pa(za), as is often the case, the aggregated posterior is forced to factorize and align with the axis of the prior. More specifically, the marginal KL term in Eq. can be decomposed the as sum of TC term and a dimensionwise-KK KL term:
d3
KL KL(g ((((() 11 ((())) TC TC ((((())
Σ KL (L((((Za) pa((a))
(5)
d=1
Thus maximizing the BLOO implicitly minimizes the TC of the aggregated posterior, enforcing the
aggregated posterior disentangle as Higgins af. (2017) 7) and Burgess et al. (2018) observed when
strongly penalizing the KL term in Eq Chen al. (2018) leverage the KL decomposition in Eq. 5
by refining the heavy latent penalization to the TC only. However, the index-code MI term in Eq.
seems to have little to role in disentanglement (see ablation study of Chen et al. (0018), potentialy arming the reconstruction performances (Hoffman & Johnson 2016).
WAE NATURALLY GOOD AT DISENTANGLING?
In this section we introduce the OT problem and the WAE objective, and discuss the compelling properties of WAEs for representation learning. Mirroring 3-TCVAE decomposition, we derive the TCWAE objective.
3.1 WAE
The Kantorovich formulation of the OT between the trae-but-unknown data distribution Pp and the model distribution Po, for given cost function c, is defined by:
rEP(Po.Po) / c(x, ((x, .) drdī
OTo Po)
inf Sonfer///
(6)
where P(Pp, Pa) is the space of all couplings Pp and Pa; namely, the space of joint distributions I on X' x X whose densities y have marginals pD and pa- Tolstikhin et al. (2018) derive the WAE objective by restraining this space and relaxing the hard constraint on the marginal using a soft constraint with Lagrange multiplier (see Appendix A more details):
WD. Φ) E
E
IE c(x, + +AD g(Z) ll)(
po(e)gats(z)po((Bb)
where D is any divergence function and À a relaxation parameter. The decoder, p=(2))), and the encoder, qa (2(x), are optimized simultaneously by dropping the closed-form minimization over the encoder network, with standard stochastic gradient descent methods.
Similarly the ELBO, objective 7 consists of reconstruction cost term and a latent regularization term, preventing the latent codes to drift away from the prior. However, WAE explicitly penalizes the aggregate posterior. This motivates, following Section 2.2 the use of WAE in disentanglement learning. [Rubenstein et al. (2018) have shown promising disentanglement performances without modifying the objective Another important difference lies in the functional form of the reconstruc- tion cost in the reconstruction term. Indeed, WAE allows for more flexibility in the reconstruction term with any cost function allowed, and particular, it allows for cost functions better suited to the data at hand and for the use of deterministic decoder networks (Tolstikhin et al.) 2018 Frogner jet al. 2015). This can potentially result in an improved reconstruction eeconsruction-diisaaanggggo disentanglenennt trade off as empirically find in Sections 4.2 and 4.1
3

---

Under review as a conference paper at ICLR 2021
3.2 TCWAE
In this section, for notation simplicity, we drop the explicit dependency of the distributions to their respective parameters.
Following Section 2.2 and Eq. S we chose the divergence function, D, in Eq. 7 to be the KL diver- gence and assume factorized prior (e.g. p(z) [Odzz Idz)) obtaining the same decomposition than in Eq. 5) Re-weighting each term in Eq. 5 with hyper-parameters 3 and 7, and plugging into Eq. we obtain our TCWAE objective:
dz
Wro - E plan (m) Lp( 22)
E
3KL ((Z) (
II
aa (aa) +
Σ KL gd (o(Zo)
|1 Pal Zu)
Given the positivity of the KL divergence, the TCWAE in Eq. is an upper-bound of the WAE objective of Eq. with min(3,
Eq. can be directly related to the 3-TCVAE objective of Chen et al. 20188:
dæ
de
- CC--TC
E E plzn (en]o(=
- log logg(op 12) + SKL (())
II d=1
qa (Z)) + 2 KL (a((a)) || pa((aa)
+alg o(Z (1(2)));()))(N))
(9)
As already mentioned, the main differences are the absence of index-code MI and a different reconstruction cost function. Setting breaks the inequality in Eq. 3] Matching the two reconstruction terms would be possible if we could find a ground cost function such that 12(,.
in E makes the two latent regularizations match but
log 1ogg(xn1. 12).
3.3 ESTIMATORS
While being grounded and motivated by information theory and earlier works on disentanglement, using the KL the latent divergence function, as opposed to other sampled-based divergences (Tolstikhin et al. 2018 Patrini et al.] (2018), presents its own challenges. Indeed, the KL terms are intractable, and especially, we need estimators to approximate the entropy terms. We propose to use two estimators, one based on importance weight-sampling Chen et al, (2018), the other on adversarial estimation using the denisty-ratio trick (Kim & Mnih 2018).
TCWAE-MW3
Chen al. (2018) propose to estimate the intractable terms Ey q and 18 log gu(Z) in the KL terms of Eq. 8 with Minibatch- Weighted Sampling (MWS). Considering batch of observation {#1..... .. Nuaa), they sample the latent codes 2, g(2)x;)) and compute:
Nizzch
Nusa
log (2) R3 (x)
log
(((((1z))
(10)
Nhntch
N X Nbatch
This estimator, while being easily computed from samples, is biased estimator of Ey log 9( ())
Chen et al. (2018) also proposed an unbiased version, the Minibatch-Stratfiied Sampling (MSS).
However, they found that did not result in improved performances, and thus, Chen al. 2018
we chose to use the simpler MWS estimator. We call the resulting algorithm the M'CSAF-MWS. Other sampled-basee
estimators of the entropy the KL divergence have been proposed (Rubenstein et al. 2019 [Emaaili et al. 2018). However, we choose the solution of Chen et al. (2018) for 1) its simplicity and 2) the similarities between the TCWAE and 3-TCVAE objectives.
TCWAE-GAN
A different approach, similar in spirit to the WAE-GAN originally proposed by Tolstikhin al. (2018), is based on adversarial training.. While Tolstikhin et al. (2018) use the adversarial training to approximate the JS divergence, Kim & Mnih (2088) use the density-ratio trick and adversarial
4

---

Under review as a conference paper at ICLR 2021
(a) Rec.
(b) Latent reg.
(c) Rec. vs latent reg.
(d) Rec. vs MMD
Figure I: Reconstruction and latent regularization terms as functions of 3 for the NoisydSprites data set. (a): reconstruction error. (b): latent regularization term (MMD for WAE, KL for TCWAE), (c): reconstruction error against latent regularization. (d): reconstruction error against MMD. Shaded regions show + one standard deviation.
training to estimate the intractable terms in Eq. 8 The the density-ratio trick (ggyyen et al. 2008 Sugiyama al. 2011) estimates the KL divergence as:
de II (alzza) d=1
((2)
KL (L(((2)
((2) II
~
log ((=)
E
(11)
1 - D(2)
where D plays the same role than the discriminator in GANs and ouputs estimate of the probability that z sampled from q(2) and not from I[2=, ga(za). Given that we can easily sample from g(2). we can use Monte-Carlo sampling to estimate the expectation in Eq. 11 The discriminator D is
adversarially trained alongside the decoder and encoder networks. We call this adversarial version
the TCWAE-GAN.
4 EXPERIMENTS
We perform series of quantitative and qualitative experiments, starting with ablation study the impact of using different latent regularization functions in WAEs followed by a quantitative comparison of the disentanglement performances of our methods with existing ones toy data sets before moving to qualitative assessment of our method on more challenging data sets. Details of the data sets, the experimental setup well the networks architectures are given in Appendix [B] In all the experiments we fix the ground-cost function of the WAE-based methods to be the square Euclidean distance: c(I, 3) = ||æ y v
4.1 QUANTITATIVE ANALYSIS: DISENTANGLEMENT ON TOY DATA SETS
Ablation study of the latent divergence function We compare the impact of the different latent regularization functions in WAE-MMD (Tolstikhin et al. 2018), TCWAE-MWS and TCWAE-GAN,
We take B y in the TCWAE objectives isolating the impact of the different latent divergence functions used in the TCWAE and the original WAE. We train the methods with 3 E {1, 2, 4, 6, 8, 10), in the case of the NoisydSprites data set (Locatello al. 2019). As
and report the results Figure I expected, the higher the penalization on the latent regularization (high 3), the poorer the recon- structions. We can see that the trade off between reconstruction and latent regularization more sensible for TCWAE-GAN. where relatively modest improvement in latent regularization results in an important deterioration reconstruction performances while TCWAE-MWS less sensible. This is better illustrated in Figure Ic with a much higher slope for TCWAE-GAN than for TCWAE-MWS. WAE seems to be relatively little impacted by the latent penalization weight. We note Figure the bias of the MWS estimator (Chen et al. 2018). Finally, we plot the reconstruction versus the MMD between the aggregated posterior and the prior for all the models in Figure (Id). Interestingly, TCWAEs actually achieved lower MMD (left part of the plot) even if they not being trained
with that regularization function. However, as expected given that the TCWAE do not optimized the reconstruction-Min MMD trade off, the WAE achieved better reconstruction (bottom part of the plot).
5

---

Under review as a conference paper at ICLR 2021
Table I: Reconstruction and disentanglement scores (= one standard deviation).
Method
MSE
MIG
lasorVAE
SAP
TOWAE MWS (3 = 6)
34.95 + 0.90
0.323 +: 0.04
0.77 + 0.01
0.072 + 0,004
TCWAE GAN
11.39
1: 0.28
0.181 0.01
0.76 4: 0.03
0.074 1: 0.003
0-ICVAE
14.30
0.235 0.03
0.81 0.03
0.070 0.006
VAlo (7 10)
8.17 1 0.86
0.24 0.08
0.78 + 0.03
0.077 + 0.011
(a) dSprites
Method
MSE
MIG
factorVAE
SAP
WAB
982.01 .20
0.019
0.40 + .U9
0.011 + .005
TCWAE MWS (3 2)
9918. 17 + 3.82
0.118 + .08
0.57 + .07
0.011 + .005
TCWAE
+
0.055
0.58 +
(:TCVAE
998.67 3.71
0.101
.11
0.015
PacooVVAE 25)
988.10 .81
0.006
.03
0.52 + .07
0.019 1 .008
(b) NoisydSprites
Method
MSE
MIG
VActorVA
SAP
WAE
24.40
0.014 + .01
0.41
0.010 + .004
ICWAE MWS (d 2)
39.53 + 24
0.322 + .00
0.73 + .01
0.067 + .001
TOWAR
33.67
+ .009
(-ICVAR
43.64
.11
0.67
0.053 .220
FactoVAFF 26)
33.23
+ .53
0.256 + .07
0.69 .09
0.066 + .013
(c) ScreamdSprites
Method
MSE
fectorVAF
SAP
WAE
3.85
.0.03
0,010 +
0.38 +
0.008
TCWAF MWS 2)
11.48 + 26
0.029 + .003
0.44 + .03
0.017 + .002
TCWAF
.10
0.030
0.015 .0001
(3
10.34
0.038
.02
FacteVVAF 100)
8.60 2: .15
0.038 + .00
0.47 + .02
0.015 : 000
(d) smalINORB
Disentanglement performances We compare our methods with 3-TCVAE (Chen et al. 2018)
FactorVAE (Kim & Mnih (2018) and the original WAE-MMD (Tolstikhin et al.) 2018 on the dSprites (Matthey et al. 2017). NoisydSprites (Locatello et al. 2019). ScreamdSprites (Locatello et al. 2019) and smalINORB (LeCun al. 2004) data sets whose ground-trutt generative-factors are known and given in Table Appendix B.T We use three different disentanglement metrics to assess the disentanglement performances: the Mutual Information Gap (MIG, Chen et al. (2018)) the factorVAB metric (Kim & Mnih 2018) and the Separated Attribute Predictability score (SAP, Kumar let al. )018)). We follow Locatello et al. (2019) for the implementation of these metrics. We use the Mean Square Error (MSE) of the reconstructions to assess the reconstruction performances of the methods. For each model, we use 6 different values for each parameter, resulting in thirty-six different models for TCWAEs, and six for the remaining methods (see Appendix [.] for more details).
Mirroring the benchmark methods, we first tune y in the TCWAEs, regularizing the dimensionwise- KL, subsequently focusing on the role of the TC term in the disentanglement performances. The heat maps of the different scores for each method and data set are given Figures 5
6 7 and 8 in
Appendix C As expected, while 3 controls the trade off between reconstruction and disentangleme affects the range achievable when tuning B. Especially, for > 1, we can see Figure that better disentanglement is obtained without much deterioration in reconstruction.
and
Table I reports the results, averaged over 5 random runs, for the four different data sets. For each method, we report the best taken to be the one achieving an overall best ranking on the four different metrics (see Appendix ?? for details). Note that the performances of WAE on the dSprites data set, both in term of reconstruction and disentanglement where significantly worse and meaningless, thus, in order to avoid unfair extra tuning of the parameters, we chose not include them. TCWAEs achieve competitive performances across all the data sets, with top scores in several metrics. Especially, the square Euclidean distance seems to improve the trade off and perform better than the cross-entropy with color images (NoisydSprites, ScreandSprites) but less so with black and white images (dSprites). See Appendix C for more results on the different data sets.

---

Under review as a conference paper at ICLR 2021
As sanity check, plot Figure the latent traversals of the different methods the smalINORB data set. More specifically, we encode one observation and traverse the latent dimensions one at the time (rows) and reconstruct the resulting latent traversals (columns). Visually, all methods, with the exception of WAF, learn to disentangle, capturing four different factors in line with the ground-truth generative factors. Models reconstructions and samples for the different data sets given in Appendix C
8
WAE
ICWAE-MWS TCWAE-GAN
3-TCVAE
FactorVAE
Figure Latent traversals for each model on smallNORB. The parameters are the same than the ones reported in Tables and 7 Each row corresponds to the traversal of the latent 2; while the columns correspond to step in the that traversal. The rows are order by increasing KL 11/Mestt Etestset (<1/22) p(z)) and the traversal range is [-
Finally, visualise the reconsruutioon-disettngggee tanglement metrics against the MSE in Figure [3] As expected, when the TC regularization weight increased, the reconstruction deteriorates while the disentanglement improves up certain point. Then, when too much penalization is put on the TC term, the poor quality of the recon- structions prevents any disentanglement in the generative factors. Reflecting the results of Table [ TCWAE-MWS seems to perform better (top-left corner represents better reconstruction and dis- entanglement). TCWAE-GAN presents better reconstruction but slightly lower disentanglement performances (bottom left corner).
trade off by plotting the different disen-
McG
mstric=factorrV V AE
metric=SAP
0.75
10WAR-105 10WARKGAN
E06
E07
025
500
600
0.0
50.5
3.0
50.0
500
200
20.0
10.0
1000
200
000
900
1800
Figure 3: Disentanglement versus reconstruction on the ScreamdSprites data set. Annotations each point are values of B. Points with low reconstruction error and high scores (top-left corner) represent better models.
4.2 QUALITATIVE ANALYSIS: DISENTANGLEMEN ON REAL-WORLD DATA SETS
We train our methods on 3Dchairs Aubry et al. 2014) and CelebA (Liu et al. 2015) whose generative factors are not known and qualitatively find that TCWAEs achieve good disentanglement. Figure shows the latent traversals of four different factors learned by the TCWAEs, while Figures [6]and 18 in Appendix D show the models reconstructions and samples. Visually, TCWAES manage to capture different generative factors while retaining good reconstructions and samples. This confirms our intuition that the flexibility offered in the construction of the reconstruction term, mainly the possibility to chose the reconstruction cost function and use deterministic decoders, improves the wiwttructtoonddisttanggleie
trade off. In order to assess the quality of the reconstructions, we compute the MSE of the reconstructions and the FID scores (Heusel et al. (017) of the reconstructions and samples. Results are reported in Table TCWAEs indeed beat their VAËS counterparts both data sets. It is worth noting that, while the performances of Factor in Table 2 seem good, the
7

---

Under review as a conference paper at ICLR 2021
Legs type
Size
Orientation
Back rest size
TCWAF- MWS
*
x
4
A
A
M
M
f
#
A R
R
A
4
1
A
SCWAE-GAN
13
17
P
Glasses/Beard
Smile
Gender
Hue
TCWAB-GAN -AWAEGGAN TCWAE MWSS
Figure 4: Latent traversals for TCWAE-MWS and TCWAE-GAN, Each line corresponds to one input data point while each subplot corresponds to one latent factor. We vary evenly the encoded latent codes in the interval [44 4].
Table 2: MSE and FID scores for the different data sets. Details of the methodology is given in Appendix B
3D chairs
ColebA
Mathod
MSF
Roe.
Samples
MSF
Rec.
Semples
TUWAE-WWS
45.8 +: 4.72
1.227
1.821
147.5 do 33.
1.204
1.264
IAW
29.5 + 3.46
0.518
0.362
129.8
+ 34.45
1.003
0.975
3 3-TCVAE
43.0 1: 4.85
1.346
1.845
180.8
51.1
1.360
1.411
VAE
42.1
7.68
0.895
0.684
201.4 + 01.84
1.017
0.982
inspection of the reconstructions and samples in Appendix shows that FactorVAB in fact struggle to generalize and to learn smooth latent manifold.
5 CONCLUSION
Leveraging the surgery of the KL regularization term of the ELBO objective, we design new disentanglement method based on the WAE objective whose latent divergence function taken to be the KL divergence between the aggregated posterior and the prior. The WAE framework naturally enables the latent regularization to depend explicitly on the TC of the aggregated posterior, quantity previously associated with disentanglement. Using two different estimators of the KI terms, we show that our methods achieve competitive disentanglement on toy data sets. Moreover, the flexibility in the choice of the reconstruction cost function offered by the WAE framework makes method more compelling when working with more challenging data sets.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
A. Achille and S. Soatto. Information dropout: Learning optimal representations through noisy computation. In IEEE Transactions on Pattern Analysis and Machine Intelligence, 2018.
M. Aubry, D. Maturana, A. Efros, B. Russell, and J. Sivic. Seeing 3D chairs: exemplar part-based based 2D-3D alignment using large dataset of CAD models. In CVPR, 2014.
Bachman, D. Hjelm, and W. Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Processing Systems, 2019.
Y. Bengio, A. Courville, and P. Vincent. Representation learning: review and new perspectives. In IEEE Transactions on Pattern Analysis and Machine Intelligence, 2013.
C. P. Burgess, I. Higgins, A. Pal, L. Matthey, N. Watters, G. Desjardins, and A. Lerchner. Under- standing disentangling in 3- V-V.E. arXiv:804.03599, 2018.
R. T. K. Chen, X. Li, R. Grosse, and D. Duvenaud. Isolating sources of disentanglement in VAEs. In Advances in Neural Information Processing Systems, 2018.
K. Do and T. Tran. Theory and evaluation metrics for learning disentangled representations. arXiv: 1908. 09961, 2019.
C. Eastwood and C. I. Williams. framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations, 2018.
B. Esmaeili, H. B. Wu, S. Jain, A. Bozkurt, N. Siddharth, B. Paige, D. H. Brooks, J. Dy, and J.-W. van de Meent. Structured disentangled representations, In AISTATS, 2018.
C. Frogner, C. Zhang, H. Mobahi, M. Araya, and T. A. Poggio. Learning with Wasserstein loss. In Advances in Neural Information Processing Systems, 2015.
S. Gao, R. Brekelmans, G. Ver Steeg, and A. Galstyan. Auto-encoding total correlation explanation. In International Conference on Artificial Intelligence and Statistics, 2019.
M. Heusel, H. Ramsauer, Unterthiner, B. Nessler, and S. Hochreiter. GANs trained by a two time-scale update rule converge to local Nash equilibrium. In Advances in Neural Information Processing Systems, 2017.
I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. M. Botvinick, S. Mohamed, and A. Ler- chner. beta-VAE: Learning basic visual concepts with constrained variational framework. In International Conference Learning Representations, 2017.
I. Higgins, D. Amos, D. Pfau, S. Racanière, L. Matthey, D. J. Rezende, and A. Lerchner. Towards definition of disentangled representations. arXiv:1812..022300, 2018
R. D. Hjelm, A. Fedorov, S. ManiiiMMarchiloon. K. Grewal, P. Bachman, A. Trischler, and Y. Bengio. Learning deep representations by mutual information estimation and maximization. In International Conference Learning Representations, 2019.
M. D. Hoffman and M. Johnson. ELBO surgery: yet another way carve up the variational evidence lower bound. In NIPS Workshop Advances in Approximate Bayesian Inference, 2016.
H. Kim and A. Mnih. Disentangling by factorising. International Conference on Machine Learning, 2018.
D. P. Kingma and J. Ba. Adam: a method for stochastic optimization. In International Conference Learning Representations, 2015.
D. P. Kingma and M. Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
A. Kumar, P Sattigeri, and Balakrishnan. Variational inference of disentangled latent concepts from unlabeled observations. In International Conference Learning Representations, 2018.

---

Under review as a conference paper at ICLR 2021
Y. LeCun, J. Huang, and L. Bottou. Learning methods for generic object recognition with invariance to pose and lighting. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2004.
Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In International Conference on Computer Vision, 2015.
F. Locatello, S. Bauer, M. Lucic, G Raetsch, S. Gelly, B. Schölkopf, and O. Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. International Conference Machine Learning, 2019.
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dSprites: Disentanglement testing Sprites dataset. http://hww...ooeeeemmm/////eeeed
2017.
X. Nguyen, M. J. Wainwright, and I. J. Michael. Estimating divergence functionals and the likelihood ratio by penalized convex risk minimization. In Advances Neural Information Processing Systems 20, 2008.
G. Patrini, M. Carioni, Forré, S. Bhargav, M. Welling, R. Van Den Berg, T. Genewein, and F. Nielsen. Sinkhorn autoencoders. arXiv:1810.01 2018.
D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference Machine Learning. 2014.
P. Rubenstein, O. Bousquet, J. Djolonga, C. Riquelme, and I. Tolstikhin. Practical and consistent estimation of f-divergences. In Advances Neural Information Processing Systems, 2019.
P. K. Rubenstein, B. Schoelkopf. and I. Tolstikhin. Learning disentangled representations with Wasserstein Auto-Encoders. In ICLR Workshop, 2018.
M. Sugiyama, T. Suzuki, and T. Kanamori. Density ratio matching under the Bregman divergence: A unified framework of density ratio estimation. In Annals of the Institute of Statistical Mathematics, 2011.
N. Tishby, F. C. Pereira, and W. Bialek. The information bottleneck method. In Annual Allerton Conference on Communication, Control and Computing. 1999.
I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf. Wasserstein Auto-Encoder, In International Conference on Learning Representations, 2018
M. Tschannen, Djolonga, P. K. Rubenstein, S. Gelly, and M. Lucic. On mutual information maximization for representation learning. In International Conference on Learning Representations, 2020.
A. van den Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. 18X.v:1807.03748, 2018
S. van Steenkiste, F. Locatello, J. Schmidhuber, and O. Bachem. Are disentangled representations helpful for abstract visual reasoning? In Advances in Neural Information Processing Systems, 2019.
G. Ver Steeg and A, Galstyan. Discovering structure in high-dimensional data through correlation explanation. In Advances Neural Information Processing Systems, 2014.
C. Villani. Optimal Transport: Old and New. Springer Berlin Heidelberg, 2008
S. Watanabe. Information theoretical analysis of multivariate correlation. In IBM Journal of Research and Development, 1960.
S. Zhao, J. Song, and S. Ermon. InfoVAE: Balancing learning and inference in variational autoen- coders. In AAAI Conference on Artificial Intelligence, 2019.
10

---

Under review as a conference paper at ICLR 2021
A WAE DERIVATION
We recall the Kantorovich formulation the OT between the true-but-unknown data distribution PD and the model distribution Pa. with given cost function
(Pp Po) =
inf TeP(Pn,Po) xxxx c(z,0) ((x,)
[1 x) dr di
(12)
where P(PD. Po) is the space of all couplings of PD and Po:
(T] f ((x 8) di = pD(a), 7(0 () -oo(t)
P(PD, Po)
(13)
Tolstikhin et al. 2018) first restrain the space of couplings to the joint distributions of the form:
(x, 8) = po (+12) (=2[2) (x) dz
(14)
where {<][z), for X, plays the same role the variational distribution in variational inference.
While the marginal constraint on æ (first constraint in Eq. in Eq. is satisfied by construction, the second marginal constraint (that over x giving po in Eq. 13 not guaranteed. A sufficient condition to have for all 2 2:
[ ((=|=) do p(z)
(15)
Secondly. Toistikhin et al. (2018) relax the constraint in Eq. 15 using a soft constraint with a Lagrange multiplier:
[ c(2,5) (x, de dā AD(1(() v(2))
W.(PD, inf <(7xx) xxx
(16)
where D is any divergence function, a relaxation parameter, y is defined Eq. 14 and () the aggregated posterior as define in Section [ Finally, they drop the closed-form minimization over the variational distribution (2(xx), obtain the WAE objective, as defined Section 3.1
(0,,(0,0)) -
E
E
E
c(x,5) AD (((Z) ( 1(())
po(X)aaaalappaa1 po
E
E
c(a, Rx) AD (() p(())
(17)
B IMPLEMENTATION DETAILS
B. EXPERIMENTAL SETUP
We train and compare our methods on four different data sets, two with known ground-truth generative factors (see Table By: dSprites (Matthey al., 2017) with 737,280 binary, 64 x 64 images and smalINORB (LeCun et al. (204) with 48,600 greyscale, 64 x 64 images; and two with unknown ground-truth generative factors: 3Dchairs (Aubry al. 2014) with 86,366 RGB, 64 x 64 images and CelebA (Liu et al. 2015) with 202,599 RGB 64 x 64 images.
Table 3: Ground-truth generative-factors the dSprites and smalINORB data sets.
data set
Generative factors (number of different values)
dSprites and variations smallNORB
Shape (3), Orientation (40), Position X (32), Position Y (32)
categories (5), lightings (6), elevations (9), azimuths (18)
We use batch size of 64 in Section 4.2 while in the main experiments of Section 4.1 we take a batch size of 100. In the ablation study of Section 4.] we use a bigger batch size of 256 in order reduce the impact of the bias of the MWS estimator (Chen et al. (2018) however show that
11

---

Under review as a conference paper at ICLR 2021
Table 4: Hyper parameters values ranges used in the different Sections.
Method
Section 4.2
Section 4.1
TCWAE-MWS
[1,2, 8, 10]2
{1,2,5, 10, 15, 20
TCWAE-GAN
4,2,4,6,8,100
11,2,5100, 10, 20)
3-TCVAE
(1,2,4,688,10) [1,10, 25, 75, 100}
[1, 15, 20
FactorVAE
(1,2,5, 10,
there is very little impact the performance of the MWS when using smaller batch size). For all experiments, we use the Adam optimizer (Kingma & Ba 2015) with learning rate of 0.0005, betal of 0.9, beta2 of 0.999 and epsilon of 0.0008 and train for 300,000 iterations. For all the data sets of Section 4.1 we take the latent dimension dæ 10. while we use dæ = 16 for 3Dchairs and dæ 32 for CelebA. We use Gaussian encoders with diagonal covariance matrix in all the models and deterministic decoder networks when possible (WAB-based methods), We follow Locatello et al. the experiments expect for CelebA where we follow [Tolstikhin' al. 2018) 8) (details of the networks architectures given Section B.2]. We use (positive) mixture of Inverse MultiQuadraaic (IMQ) kernels and the associated reproductive Hilbert space to compute the MMD when it is needed (WAE and ablation study of Section 4.1
(2019) for the architectures
The different parameter values used for each experiment are given Table 4 is chosen such that the resulting method achieves the best score 8, when averaging over all the 3 values, where the score is defined as the sum of the ranking on each individual metric: 8 = TMSE + Emetric Tmetric where TMSE designed the ranking of the MSE (lower is better) and Tmetrics for metric (MIG, FactorVAE, SAP), is the ranking of the disentanglement performances as measured by the given metric (higher better). / then chosen such that the resulting method, with the previously found y, achieves the best overall score s defined above. In Section 4.1 we use a validation run to select the parameters values and report the MSE and FID scores on test run. MSE are computed on test set of size 10,000 with batch size of 1,000, while we follow Heusel et al. (2017) for the FID implementation: we first compute the activation statisties of the features maps on the full test set for both the reconstruction, respectively samples, and the true observations. We then compute the Frechet distance between two Gaussian with the computed statistics.
B.2 MODELS ARCHITECTURES
The Gaussian encoder networks, go (2(r)) and decoder network, p=(2))), are parametrized by neural networks as follow:
po(x|2) (o)(==) if WAE based method, (2(Ho(2) otherwise.
(0(2|>) = (((o))))(((((())))
where fe. #40 o-- Me and 2 are the outputs of convolutional neural networks. All the experiments use the architectures of Locatello et al. (2019) except for CelebA where we use the architecture inspired by Tolstikhin et al. (2018). The details for the architectures are given Table
All the discriminator networks, D, are fully connected networks and share the same architecture given T The optimisation setup for the discriminator is given Table
12

---

Under review as a conference paper at ICLR 2021
Table 5: Networks architectures
Encoder
Decoder
Discriminator
Input: 64 x 64x CONV. 4 32 stride ReLU
Input: dz
Input: dz
FC 256 ReLU
FC 1000 ReLU
CONV. X 4 x 32 stride ReLU
FC x 4 x 64 ReLU
FC 1000 ReLU
CONV. 4 64 stride ReLU
CONV, 4
64 stride 2 ReLU
FC 1000 ReLU
CONV. x 4 x 64 stride ReLU
CONV. 4 X 4 x 32 stride 2 ReLU
FC 1000 ReLU
FC 256 Relu
CONV, 4
32 stride 2 ReLU
FC 1000 ReLU
FC x dz
CONV. x 4 x e stride 2
FC 1000 ReLU
FC
Locatello (2019) architectures
Encoder
Decoder
Discriminator
Input: 64 64 c CONV. X 4 x 32 stride BN ReLU
Input: dz
Input: dz FC 1000 ReLU
FC 8 x 8 x 256 BN ReLU ReLU
CONV. 4 x stride BN ReLU
CONV. 4 X x 128 stride
FC 1000 ReLU
CONV. 4 X 4 x 128 stride 2 BN ReLU
CONV, 4 X 4 x 64 stride 2 BN ReLU
FC 1000 ReLU
CONV. x 4 x 256 stride 2 BN ReLU
CONV. 4 X x 32stride 2 BN Relu
FC 1000 ReLU
2 X dz
CONV. 4 X x
FC 1000 ReLU
FC 1000 ReLU
FC 2
(b) CelebA networks architectures
Table 6: FactorVAF discriminator setup
Parameter
Value
Learning rate
-4 (Section 4.1
(Section 4.2
beta 1
0.5
beta 2
0.9
epsilon le-08
13

---

Under review as a conference paper at ICLR 2021
C
QUANTITATIVE EXPERIMENTS
HYPER PARAMETER TUNING
Figure 5: Heat maps for the different dSprites.
Figure 6: Heat maps for the different scores on NoisydSprites.
Figure 7: Heat maps for the different scores on ScreamdSprites.
Figure 8: Heat maps for the different scores on SmalINORB.
14

---

Under review as a conference paper at ICLR 2021
Table 7: y values for methods for each data
Method
dSprites
NoisydSprites
ScreamdSprites
smalINORB
TCWAE MWS
2
1
1
TCWAE GAN
1
1
10
2
DISENTANGLEMENT SCORES vs B
For each method, we plot the distribution (over five random runs) of the different metrics for different values.
y.:
Figure 9: Violin plots of the different scores versus on dSprites.
15

---

Under review as a conference paper at ICLR 2021
180
M
Figure 10: Violin plots of the different scores versus on NoisydSprites.
TUNADGAR
CTEVE
Figure 11: Violin plots of the different scores versus ScreamdSprites.
16

---

Under review as a conference paper at ICLR 2021
TOAAR MATres)
ICH4 400 yed)
221
1
Figure 12: Violin plots of the different scores smallNORB.
RECONSTRUCTIONS AND SAMPLES
17

---

Under review as a conference paper at ICLR 2021
Reconstructions
(b) Samples
Figure 13: Samples and reconstructions for each model the NoisydSprites. (a): Reconstructions.
Top-row: input data, from second-to-top to bottom row: WAE, TCWAE-MWS, TCWAE-GAN 3-TCVAE, FactorVAE. (b) Samples. From top bottom row: WAE, TCWAE-MWS, TCWAE-GAN, 3-TCVAF, FactorVAE. Parameters are the ones reported in Tables
and 7
18

---

Under review as a conference paper at ICLR 2021
(a) Reconstructions
(b) Samples
Figure 14: Same than Figure 13 but for ScreamdSprites.
19

---

Under review as a conference paper at ICLR 2021
(a) Reconstructions
(b) Samples
Figure 15: Same than Figure 13 but for NORR.
20

---

Under review as a conference paper at ICLR 2021
D QUALITATIVE EXPERIMENTS
3DCHAIRS
Reconstructions
Samples
TCWAB-MWS
TCWAE-GAX
Figure 16: Reconstructions (left quadrants) and samples (right quadrants) for TCWAE-MWS (top quadrants) and TCWAE-GAN (bottom quadrants).
21

---

Under review as a conference paper at ICLR 2021
Reconstructions
Samples
h
4
*
hn
h
th
H
#
h
Factorvan VA (1 100
Figure 17: Reconstructions (left quadrants) and samples (right quadrants) for 3-TCVAE (top quad- rants) and FactorVAE (bottom quadrants).
22

---

Under review as a conference paper at ICLR 2021
CELEBA
Reconstructions
Samples
TCWAE-MWS
TCWAE.COM
Figure 18: Same as Figure 16 for the CelebA data set.
23

---

Under review as a conference paper at ICLR 2021
Reconstructions
Samples
B-TCVAE (3 = 10)
Factor VAE (Y = 20)
Figure 19: Same as Figure 18 for 3-TCVAB (top quadrants) and FactorVAB (bottom quadrants).
24


Reviews:
[{"review_id": "u00b_ORE7U", "reviewer": "AnonReviewer5", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper extends the Wasserstein Autoencoder (WAE) work by splitting the divergence on the variational marginal into 2 terms, akin to what was done in TC-VAE. This enables directly controlling the explicit contribution of the total correlation term, which is likely to contribute to disentanglement more directly. They explore 2 variations of their model, based on different estimators of the TC term (TCWAE-MWS, using minibatch-weighted sampling; TCWAE-GAN, using a density ratio trick).\n\nOverall, I found this work to be a nicely complete exploration of a simple extension of an existing framework. They mostly rederive existing methods in the WAE framework, but results are promising and the paper addresses several datasets, compares to baselines well and seems well executed. It seems to lack comparison and discussion to a paper which seems directly related [Xiao et al 2019]. But I feel this is still a worthy piece of research to showcase at ICLR.\n\nQuestions/comments:\n1. A cursory search indicated the following paper which also addresses disentanglement with the Wassertein Total Correlation: [Xiao et al 2019]. They use another estimator of the TC, instead opting for the Kantorovich-Rubinstein formulation.\n   1. Can you comment on how their work relates to this current paper?\n   2. A direct comparison would be rather interesting, but might be out of scope for a rebuttal.\n2. Reconstructions for the TCWAE-MWS appear rather bad (Figures 13-19 in the Appendix ), but Figure 1c doesn’t seem to reflect that, which is slightly surprising.\n   1. Could you comment on this discrepancy?\n3. Relatedly the TCWAE-GAN disentanglement doesn’t seem particularly exciting (metric-wise), would you still recommend using it instead of TCWAE-MWS? \n   1. It is still a clear improvement over vanilla WAE, so there’s value to the work in this current state; but I’d wonder when one would prefer choosing this versus TCVAE?\n4. It might be appropriate to discuss 2-Stage VAE [Dai et al 2019] and associated family of models, which currently obtain really good results on more complex datasets.\n\n\nReferences:\n*\t[Xiao et al 2019] https://arxiv.org/abs/1912.12818\n*\t[Dai et al 2019] https://arxiv.org/abs/1903.05789\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "WYHfwVnLTbc", "reviewer": "AnonReviewer4", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper addresses disentanglement in the latent space of autoencoders. To this end, it combines ideas from four existing papers, namely the reconstruction loss of the Wasserstein autoencoder, the regularization term decomposition from the total correlation autoencoder, and entropy estimation using minibatch-weighted sampling or the density-ratio trick. This combination certainly makes sense, as it brings together methods that have previously been shown to work well in isolation.\n\nThe main part of the paper is devoted to an empirical evaluation of the new autoencoder training procedure. The new method is compared against various baselines in terms of L2 reconstruction error and three disentanglement scores on four toy datasets. In addition, latent space traversals on 3Dchairs and CelebA are shown to qualitatively demonstrate the disentanglement capabilities of the proposed methods.\n\nUnfortunately, the description of the experiments is not very precise. \n* The role of the hyperparameter gamma remains unclear. In the ablation study, the authors simply set gamma=beta without further explanation, and in the comparison, they just state \"we first tune gamma\" and \"for gamma >1, better disentanglement is obtained\", again without further explanation.\n* In the comparison experiment, they report results for the values of beta that achieve \"an overall best ranking on the four different metrics\" without explaining what an \"overall best ranking\" is. Choices like this must not be taken lightly, as the analysis in \"Why rankings of biomedical image analysis competitions should be interpreted with care\" (Nature Communications 9: 5217, 2018) impressively demonstrates.\n* The experiment in figure 2 seems to have three degrees of freedom (the data instance x, the latent index i, and the size of the modification in direction z_i). However, only two degrees of freedom are shown, and it remains unclear from the caption and associated main text, which ones. Moreover, I cannot deduce justification for the statement \"all methods .. learn to disentangle, capturing four different factors\" from the figure -- I do not see any obvious disentanglement.\n\nThe bigger problem with the paper, however, is the question: What have we learned from these experiments? The rankings in table 1 are pretty inconsistent between different metrics, and the corresponding figure 3 appears to be cherry picked, as the ScreamdSprites is the dataset where the proposed methods perform best.\n\nI also do not agree with the claim that \"TCWAEs achieve good disentanglement\" on real-world datasets. Figure 4 shows severe entanglement between unrelated factors. For example, the size feature for the chairs also changes the type of chair. All features in the CelebA examples have a tendency also to change the background appearance. The gender feature dramatically influences person identity in the MWS results, whereas it does not change the gender at all in the GAN variant. Substantial variations in person identity are also visible in most other examples.\n\nIn summary, while the paper provides numbers, it lacks new insight. In light of mathematical proofs indicating that the true generative factors are generally unidentifiable in non-linear unsupervised settings (cf. the work of Aapo Hyvärinen and others), I am skeptical that heuristic trial-and-error investigations of disentanglement like the present one will yield interesting results. In a sense, this is also acknowledged by the authors, who merely state in the conclusion that \"our methods achieve competitive disentanglement on toy data sets\" -- that's not much, given the effort that went into the experiments.", "extra_scores": {}, "extra_text": {}}, {"review_id": "TH2oUW9qVbA", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This submission proposes to add a KL term to the Wasserstein auto-encoder objective in order to improve its disentanglement capabilities.\nThis combines the idea of Hoffman & Jonhson (2016) of using a marginal KL term, with the Wasserstein auto-encoder framework.\nChallenges regarding the estimation of the KL term are also addressed with two previous works. \nThis results in a two regularization parameter objective, whose superiority to existing approaches (using a single parameter) is not clear.\n\nStrengths: WAE with a disentanglement term was as far as I now not attempted before, the authors offer two well justified techniques to do it.\n\nWeaknesses: (1) The work is very iterative, existing approaches are only combined, (2) Superiority to WAE without this term is not surprising, and I failed to see a clear superiority to competing unsupervised disentanglement approaches. (3) Given the emphasis on the Wasserstein distance of the original approach, it is also a bit disappointing to resort to a KL term for disentanglement. (4) Most importantly, comparison to simpler alternative KL (non-marginal) losses is absent as far as I can tell. That was for me the most interesting appeal of the paper.\n\nOverall, I tend to think the paper would require a more exhaustive investigation of disentanglement approaches, contextualized to the Wasserstein distance and issues raised regarding marginal versus non-marginal divergences. I recommend rejection.\n\nOn this last point: It remains unclear to me whether the original hypothesis of the paper (page 3), that the index-code MI term of the KL divergence may be detrimental to disentanglement, is supported by the current study, and thus whether the extra technicalities required to eliminate it are worth the effort. Perhaps the authors could elaborate on that with an alternative objective close to the classical KL term, and thus easier to optimize?\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "7RZ46SAoRDl", "reviewer": "AnonReviewer3", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe paper is motivated by the need for a better trade-off between the reconstruction and disentanglement performance of an autoencoder. The proposed solution is to use KL as a latent regularizer in the framework of Wassestain autoencoders, which allows for a natural interpretation of total correlation.\n\n\nThe paper reads well, all related work and relevant background concepts are nicely integrated throughout the text. The experiments are exhaustive and the results show competitive performance wrt disentanglement while improving reconstruction/modeling of the AEs.\n\nIf a dataset is of dynamical nature, how difficult would it be to extend the current version of TCWAE to dynamical systems? Do the authors have any intuition/hint on what should change to make their method applicable to dynamical setups? Significantly changing the probabilistic model or modifying only the and encoder/decoder architecture could suffice? \n\nMinor:\n- Consider changing the naming of the baselines either in tables or figures to make them consistent  Chen et al (2018) -> TCVAE Kim & Mnih (2018) -> factorVAE.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: There were both positive and negative assessments of this paper by the reviewers: It was deemed a well written paper that explores cleanly rederiving the TC-VAE in the Wasserstein Autoencoder Framework and that has experiments comparing to competing approaches. However, there are two strong concerns with this paper: First, novelty appears to be strongly limited as it appears a rederivation using known approaches. Second, two reviewers were not convinced by the experimental results and do not agree with the claim that the proposed approach is better than competing methods in providing disentangled representations. I agree with this concern, in particular as assessing unsupervised disentanglement models is known to be very hard and easily leads to non-informative results (see e.g. the paper cited by the authors  from Locatello et al., 2019). Overall, I recommend rejecting this paper.
Decision: Reject