On the Transfer of Disentangled Representations in Realistic Settings
Authors: Andrea Dittadi
Frederik Träuble
Francesco Locatello
Manuel Wuthrich
Vaibhav Agrawal
Ole Winther
Stefan Bauer
Bernhard Schölkopf
Venue: ICLR 2021 Poster (2021)
Keywords: representation learning
disentanglement
real-world

Abstract: Learning meaningful representations that disentangle the underlying structure of the data generating process is considered to be of key importance in machine learning. While disentangled representations were found to be useful for diverse tasks such as abstract reasoning and fair classification, their scalability and real-world impact remain questionable.
We introduce a new high-resolution dataset with 1M simulated images and over 1,800 annotated real-world images of the same setup. In contrast to previous work, this new dataset exhibits correlations, a complex underlying structure, and allows to evaluate transfer to unseen simulated and real-world settings where the encoder i) remains in distribution or ii) is out of distribution.
We propose new architectures in order to scale disentangled representation learning to realistic high-resolution settings and conduct a large-scale empirical study of disentangled representations on this dataset. We observe that disentanglement is a good predictor for out-of-distribution (OOD) task performance.

Published as a conference paper at ICLR 2021
ON THE TRANSFER OF DISENTANGLED REPRESENTA- TIONS IN REALISTIC SETTINGS
Andrea Dittadi,1! Frederik Träuble," Francesco Locatello, Manuel Wüthrich, Vaibhav Agrawal, Ole Winther, 14.5 Stefan Bauer, Bernhard Schölkopf
'Technical University Denmark, Max Planck Institute for Intelligent Systems, 'ETH Zurich, Department for Computer Science, *Copenhagen University Hospital, University of Copenhagen, "CIFAR Azrieli Global Scholar
ABSTRACT
Learning meaningful representations that disentangle the underlying structure of the data generating process is considered to be of key importance in machine learn- ing. While disentangled representations were found to be useful for diverse tasks such as abstract reasoning and fair classification, their scalability and real-world impact remain questionable. We introduce a new high-resolution dataset with 1M simulated images and over 1,800 annotated real-world images of the same setup. In contrast to previous work, this new dataset exhibits correlations, complex underlying structure, and allows to evaluate transfer unseen simulated and real- world settings where the encoder i) remains in distribution or ii) is out of distribu- tion. We propose new architectures in order to scale disentangled representation learning to realistic high-resolution settings and conduct a large-scale empirical study of disentangled representations this dataset. We observe that disentan- glement is a good predictor for duttrf-distributtin (OOD) task performance.
INTRODUCTION
Disentangled representations hold the promise of generalization unseen scenarios (Higgins let al. al. 2018 Higgins al., 2018) and faster
2017b), increased interpretability (Adel
learning on downstream tasks (van Steenkiste al. 2019; Locatello al. 2019a). How-
ever, most of the focus in learning disentan- gled representations has been on small syn- thetic datasets whose ground truth factors ex- hibit perfect independence by design. More re- alistic settings remain largely unexplored. We hypothesize that this is because real-world sce- narios present several challenges that have not been extensively studied to date. Important challenges are scaling (much higher resolution in observations and factors), occlusions, and correlation between factors. Consider, for instance, robotic moving cube: Here, the robot arm can occlude parts of the cube, and its end-effector position exhibits correlations with the cube s position and orientation, which might be problematic for common disentanglement learners (Träuble et al., 2020). Another difficulty is that we typically have only limited access to ground truth labels in the real world, which requires robust frameworks for model selection when no or only weak labels are available.
Figure I: Images from the simulated dataset (left) and from the real-world setup (right).
"Equal contribution. Correspondence to: Cadit@dtu.dk: <frederik.traeee @ nuebingen.mpg.de. Work done during internship the Max Planck Institute for Intelligent Systems.

---

Published as a conference paper at ICLR 2021
The goal of this work is to provide path towards disentangled representation learning in realistic settings. First, we argue that this requires a new dataset that captures the challenges mentioned above. We propose dataset consisting of simulated observations from scene where a robotic arm interacts with a cube in a stage (see Fig. D This setting exhibits correlations and occlusions that are typical in real-world roboties. Second, methods to perform well on this dataset. Third, we extensively analyze the usefulness of disen- tangled representations in terms of out-of-distribtiion downstream generalization, both in terms of held-out factors of variation and sim2real transfer. In fact, our dataset is based on the TriFinger robot from Wüthrich al. (2020), which can be built to test the deployment of models in the real world. While the analysis in this paper focuses on the transfer and generalization of predictive mod- els, we hope that our dataset may serve as a benchmark to explore the usefulness of disentangled representations in real-world control tasks.
show how to scale the architecture of disentanglement
The contributions of this paper can be summarized as follows:
We propose a new dataset for disentangled representation learning, containing 1M simu- lated high-resolution images from a robotic setup, with seven partly correlated factors of variation. Additionally, we provide a dataset of over 1,800 annotated images from the cor- responding real- woaldworld setup that can be used for challenging sim2real transfer tasks. These datasets are made publicly available
We propose new neural architecture successfully scale VAE-based disentanglement learning approaches to complex datasets.
We conduct a large-scale empirical study on generalization to various transfer scenarios on this challenging dataset. We train 1,080 models using state-of-the-ar disentangle- ment methods and discover that disentanglement is a good predictor for out-of outoof-disttibution (OOD) performance of downstream tasks.
RELATED WORK
Disentanglement methods. Most sfate-of-the-arr disentangled representation learning approaches are based on the framework of variational autoencoders (VAEs) (Kingma & Welling 2014 Rezende ctt al. 2). A (high-dimensional) observation æ is assumed to be generated according to the latent variable model po (((z==))(()))))))))))))))))))( where the latent variables z have a fixed prior p(z). The generative model po [z) x) and the approximate posterior distribution (0(2(a)) are typically parameterized by neural networks, which optimized by maximizing the evidence lower bound (ELBO):
LV gga(=l [log Po a =)] (KL (=[=))<<(<)) < log)(æ)
(I)
As the above objective does not enforce any structure on the latent space except for some similar- ity to p(z), different regularization strategies have been proposed, along with evaluation metrics
to gauge the disentanglement of the learned representations (Higgins et al.) 2017a] Kim & Mnih,
2018 Burgess et al. 2018 Kumar et al., 2018 Chen et al. 2018) [astwood & Williams] (0018). Recently, Locatello et al. 2011b) Theorem 1) showed that the purely unsupervised learning of disentangled representations is impossible. This limitation can be overcome without the need for explicitly labeled data by introducing weak labels (Locatello et al.) 2020 Shu et al. 2019). Ideas related to disentangling the factors of variation date back to the non- non-linear ICA literature (Comon)
1994 Hyvärinen & Pajunen| 1999 Bach & Jordan 2002 Juutten & Karhunen 2003 Hyvarinen & Morioka 2016 Hyvarinen et al. 2019) Gresele et al. 2019)) Recent work combines non-linear ICA
with disentanglement (Khemakhem al. 2020 Sorrenson et al.
2020 Klindt et al. 2020).
Evaluating disentangled representations. The BetaVAE (Higgins et al., 2017a) and Factor- VAE (Kim & Mnih, 2018) scores measure disentanglement by performing an intervention on the factors of variation and predicting which factor was intervened on. The Mutual Information Gap (MIG) (Chen et al., 2018). Modularity (Ridgeway & Mozer 2018), DCI Disentanglement (East- wood & Williams| 2018) and SAP scores (Kumar al. (2018) based on matrices relating factors of variation and codes (e.g. pairwise mutual information, feature importance and predictability).
httpp:/ ngen. http Finger_datasets.o 68 GB)
2

---

Published as a conference paper at ICLR 2021
Datasets for disentanglement learning. dSprites (Higgins et al. 2017a), which consists of binary low-resolution 2D images of basic shapes, is one of the most commonly used synthetic datasets for disentanglement learning. Color-dSprites. Noisy-dSprites, and Screan-dSprites are slightly more challenging variants of dSprites. The SmallNORB dataset contains toy images rendered under dif- ferent lighting conditions, elevations and azimuths (LeCun et al. 2004). Cars31) (Reed et al.) 2015) exhibits different car models from Fidler et al. (2012) under different camera viewpoints. 3dshapes popular dataset of simple shapes in 3D scene (Kim & Mnih 2018). Finally, [Gondal et al. (2019) proposed MPI3D, containing images of physical 3D objects with seven factors of variation, such as object color, shape, size and position available in a simulated, simulated and highly realistic rendered simulated variant. Except MPI3D which has over IM images, the size of the other datasets limited with only 17, 568 to 737, 280 images. All of the above datasets exhibit perfect indepen- dence of all factors, the number of possible states is on the order of 1M or less, and due to their static setting they do not allow for dynamic downstream tasks such as reinforcement learning. In addition, except for SmallNORB, the image resolution is limited to 64x64 and there are occlusions.
Other related work. Locatello et al. 20200) probed the out-of-distributtin generalization of down- stream tasks trained on disentangled representations. However, these representations are trained on the entire dataset. Generalization and transfer performance especially for representation learning has likewise been studied in Dayan [199)); Muandet et al. (2013) Heinze-Deml & Meinshausen) (2017);
Rojas-Carulla et al. (2018) Suter et al. 2019);
al. (2018); Arjovsky et (2019); Krueger al.
(2020): Gowal et al. (2020) J. For the role of disentang lement in causal representation learning we
refer to the recent overview by Schölkopf et al. (2021). Träuble et al. (2020) systematically investi-
gated the effects of correlations between factors of variation on disentangled representation learners.
Transfer of learned disentangled representations from simulation to the real world has been recently investigated by Gondal et al. (2019) on the MPI3D dataset, and previously by HHigiins et al. (201
in the context of reinforcement learning. Sim2real transfer is of major interest in the robotic learning
community, because of limited data and supervision in the real world (Tobin et al.
2017 Rusu et al.
2017
Peng et al. 2018: James et al. 2019 Yan et al. 2020; Andrychowicz et al.
2020
3 SCALING DISENTANGLED REPRESENTATIONS TO COMPLEX SCENARIOS
A new challenging dataset. Simulated im- ages in our dataset are derived from the trifinger robot platform introduced by Wüüthrich et al.) (2020). The motivation for choosing this set- ting is that (1) it is challenging due occlu- sions, correlations, and other difficulties en- countered in robotic settings, (2) it requires modeling of fine details such as tip links at high resolutions, and (3) it corresponds to a robotic setup, so that learned representations can be used for control and reinforcement learning in simulation and in the real world. The scene comprises a robot finger with three joints that be controlled to manipulate a cube in bowl-shaped stage. Fig. [ shows examples of scenes from our dataset. The data is gener- ated from 7 different factors of variation (FoV) listed in Table I Unlike in previous datasets, not all FoVs are independent: The end-effector (the tip of the finger) can collide with the floor or the cube, resulting in infeasible combinations of the factors (see Appendix I. We argue that such correlations are a key feature in real-world data that is not present in existing datasets. The high FoV resolution results in approximately 1.52 bil- lion feasible states, but the dataset itself only contains one million of them (approximately 0.065% of all possible FoV combinations), realistically rendered into 128 x 128 images. Additionally, we recorded an annotated dataset under the same conditions in the real-world setup: we acquired 1,809 camera images from the same viewpoint and recorded the labels of the 7 underlying factors of vari- evaluations, shot learning, and testing
Fo
Values
Upper joint
30 values in
[0.65, +0.65]
Middle joint
30 values in
+0.5
Lower joint
30 values in
+0.8
Cube position x
30 values in
+0.11
Cube position y
30 values in -0.11,
Cube rotation
10 values in
08.81%]
Cube color hue
12 values in
[0 3330
Table 1: Factors of variation in the proposed dataset. Values are linearly spaced in the spec- ified intervals. Joint angles are in radians, cube positions in meters.
ation. This dataset can be used for other sim 2real aspects.
3

---

Published as a conference paper at ICLR 2021
Latent traversals of trained model that perfectly disentangles the dataset's Fo In each column, all latent variables but one are fixed.
Figure 2:
Model architecture. When scaling disentangled representation learning to more complex datasets, such as the one proposed here, one of the main bottlenecks in current VAE-based ap- proaches is the flexibility of the encoder and decoder networks. In particular, using the architecture from Locatelo et al. (2019)), none of the models we trained correctly captured all factors of varia- tion or yielded high-quality reconstructions. While the increased image resolution already presents a challenge, the main practical issue in our dataset is the level of detail that needs be modeled. In particular, we identified the cube rotation and the lower joint position to be the factors of variation that were the hardest capture. This is likely because these factors only produce relatively small changes in the image and hence the reconstruction error.
To overcome these issues, we propose a deeper and wider neural architecture than those commonly used in the disentangled representation learning literature, where the encoder and decoder typically have 4 convolutional and 2 fully-connected layers. Our encoder consists of a convolutional layer, 10 residual blocks, and 2 fully-connected layers. Some residual blocks are followed by 1x1 con- volutions that change the number of channels, or by average pooling that downsamples the tensors by factor of along the spatial dimensions. Each residual block consists of two 3x3 convolutions with leaky ReLU nonlinearity, and a learnable scalar gating mechanism (Bachlechner al., 2020). Overall, the encoder has 23 convolutional layers and 2 fully connected layers. The decoder mirrors this architecture, with average pooling replaced by bilinear interpolation for upsampling. The total number of parameters is approximately 16.3M. See Appendix A for further implementation details.
Experimental setup. duced above by training 1,080 3-VAE models For further experimental details we refer the reader to Appendix The hyperparameter sweep is defined follows:
We perform large-scale empirical study on the simulated dataset intro- 2
We train the models using either unsupervised learning or weakly supervised learning Lo- catello al. 2020). In the weakly supervised case, model trained with pairs of images that differ in k factors of variation. Here we fix A = 1 as it was shown to lead to higher disentanglement by Locatello et al. (2020). The dataset therefore consists of 500k pairs of images that differ in only one FoV.
We vary the parameter 3 in (1, 4), and use linear deterministic warm-up Bowman et al. 2015 Sønderby al. 20166) over the first (0, 10000, 50000) training steps.
The latent space dimensionality is in {10, 20350.
Half of the models trained with additive noise in the input image. This choice is mo- tivated by the fact that adding noise to the input of neural networks has been shown to be beneficial for out-of distribbinn generalization (Sietsma & Dow 1991 Bishop 1995).
Each of the 108 resulting configurations trained with 10 random seeds.
Can we scale up disentanglement learning? Most of the trained VAEs in our empirical study fully capture all the elements of a scene, correctly model heavy occlusions, and generate detailed,
Training these models requires approximately 2.8 GPU years NVIDIA Tesla V100 PCIe.
4

---

Published as a conference paper at ICLR 2021
unsupervised
weakly supervised
ELBO
-52
47
5
-22
Recons. loss
69
66
2
35
Disentanglement score
GBT10000
97
67
-17
0.6
ELBO
90
71
22
50
0.4
loss
-88
-60
-25
-39
GBT10000
94
65
24
0.2
DCI
M1G
Modularity
SAP
DCI MIG
Modularity
SAP
Figure 3: Left: Disentanglement metrics aggregating all hyperparameters except for supervision type. Right: Rank correlations (Spearman) of ELBO, reconstruction loss, and the test error a GBT classifier trained 10,000 labelled data points with disentanglement metries. The upper rank correlations correspond to the unsupervised models and the lower to the weakly supervised models.
high-quality samples and reconstructions (see Appendix B.2). From visual inspections such as the latent traversals in Fig. 2 we observe that many trained models fully disentangle the ground-truth factors of variation. This, however, appears to only be possible in the weakly supervised scenario. The fact that models trained without supervísion learn entangled representations in line with the impossibility result for the unsupervised learning of disentangled representations from Locatello ct (2019))) Latent traversals from selection models with different degrees of disentanglemen presented in Appendix B.3] Interestingly, the high-disentane disentanglement models seem correct for correlations and interpolate infeasible states, i.e. the fingertip traverses through the cube or the floor.
(2019b).
Summary: The proposed architecture scale disentanglement learning to more realistic settings, but a form of weak supervision is necessary to achieve high disentanglement.
How useful are common disentanglement metrics realistic scenarios? The violin plot in Fig. 3l (left) shows that DCI and MIG measure high disentanglement under weak supervision and lower disentanglement in the unsupervised setting. This is consistent with our qualitative conclu- sion from visual inspection of the models (Appendix B.3) and with the aforementioned impossibility result. Many of the models trained with weak supervision exhibit very high DCI score (29% of them have >99% DCI, some of them up to 99.89%) SAP and Modularity appear to ineffective at capturing disentanglement in this setting, as also observed by Locatello et al. (2019b). Finally, note that the Beta VAE and Factor metrics are not straightforward be evaluated on datasets that do not contain all possible combinations of factor values. According to Fig. 3 (right), DCI and MIG strongly correlate with test accuracy of GBT classifiers predicting the FoVs. In the weakly supervised setting, these metrics strongly correlated with the EI ELBO (positively) and with the construction loss (negatively). We illustrate these relationships in more detail Appendix [B.4 Such correlations were also observed by Locatello al. (2020) on significantly less complex datasets, and can be exploited for unsupervised model selection: these unsupervised metrics can be used as prox- ies for disentangiement metrics, which would require fully labeled data.
Summary: DCI and MIG appear to be useful disentanglement metrics in realistic scenarios, whereas other metrics seem fall short of capturing disentanglement or can be difficult com- pute. When using weak supervision, we can select disentangled models with unsupervised metrics.
FRAMEWORK FOR THE EVALUATION OF OOD GENERALIZATION
Previous work has focused on evaluating the usefulness of disentangled representations for various downstream tasks, such predicting ground truth factors of variation. fair classification, and ab- stract reasoning. Here we propose a new framework for evaluating the out-of--iittrbutiion (OOD) generalization properties of representations. More specifically, consider downstream task
5

---

Published as a conference paper at ICLR 2021
case, regression of ground truth factors trained a learned representation of the data, and evaluate the performance on a held-out test set. While the test set typically follows the same dis- tribution the training set fin-distribution generalization), we also consider test sets that follow different distribution (out-of-distribttion generalization). Our goal is to investigate to what extent, if all, downstream tasks trained on disentangled representations exhibit higher degree of OOI) generalization than those trained on entangled representations.
Let D denote the training for disentangled representation learning. To investigate OOD gen- eralization, we train downstream regression models on a subset Dj C D to predict ground truth factor values from the learned representation computed by the encoder. We independently train one predictor per factor. We then test the regression models on set D2 that differs distributionally from the training set D1, as either contains images corresponding to held-out values of chosen FoV (e.g. unseen object colors), or it consists of real-world images. We now differentiate between two scenarios: (1) D2 C D, i.e the OOD test set subset of the dataset for representation learning; (2) D and D2 are disjoint and distributionally different. These two scenarios will be denoted by OODI and OOD2, respectively. For example, consider the case in which distributional shifts are based one FoV: the color of the object. Then, we could define these datasets such that images in D always contain a red or blue object, and those in D2 C D always contain a red object. In the OOD1 scenario, images in D2 would always contain blue object, whereas in the OOD2 case they would always contain an object that is neither red nor blue.
The regression models considered here Gradient Boosted Trees (GBT), random forests, and MLPs with {1, 3 hidden layers. Since random forests exhibit a similar behavior to GBTs, and all MLPs yield similar results each other, we choose GBTs and the 2-layer MLP representative models and only report results for those. To quantify prediction quality, we normalize the ground truth factor values to the range [0, 1), and compute the mean absolute error (MAE). Since the values are normalized, we can define our transfer metric as the average of the MAE over all factors (except for the FoV that is OOD).
BENEFITS AND TRANSFER OF STRUCTURED REPRESENTATIONS
Experimental setup. We evaluate the transfer metric introduced in Section across all 1,080 trained models. To compute this metric, we train regression models to predict the ground truth factors of variation, and test them under distributional shift. We consider distributional shifts in terms of cube color or sim2real, and we do not evaluate downstream prediction of cube color. We report scores for two different regression models: Gradient Boosted Tree (GBT) and an MLP with 2 hidden layers of size 256. In Appendix [ we provide details on the datasets used in this section.
In the OOD1 setting, we have D2 C D, hence the encoder is in-distribution: we are testing the predictor on representations images that were the training set of the representation learning algorithm. Therefore, we expect the representations to be meaningful. We consider three scenarios:
00D1-A:
The regression models are trained 1 cube color (red) and evaluated the remaining 7 colors.
0OD1-B:
The regression models are trained on 4 cube colors with high hue in the HSV space, and evaluated on 4 cube colors with low hue (extrapolation).
OODI-C: The regression models are again trained and evaluated on 4 cube colors, but the training and evaluation colors are alternating along the hue dimension (interpolation).
In the more challenging setting where even the encoder goes out-of-distribttion (OOD2, with D2 n D = Ø), we train the regression models on subset the training set D that includes all cube colors, and we consider the two following scenarios:
00D2-A: The regression models evaluated simulated data, on 4 cube colors that out of the encoder's training distribution.
00D2-B: The regression models are evaluated on real-world images of the robotic setup, without any adaptation or fine-tuning.

---

Published as a conference paper at ICLR 2021
0.30
score
MLP 00D1-C
-24
-18
-20
-30
<0.44
0.4 < 0.99
GBT 00D1-C
-91
-66
8
-58
Transfer score
0.20
MLP O0D1-B
-3
47
-15
44
0.15
GBT 00D1-B
-82
-84
-66
0.10
MLP 00D1-A
-17
-10
-13
-10
0.05
GBT 00D1-A
-87
-86
-60
MLP 00D1-C
GBT 0001-C
MLP 0001-8
GBT 0001-8
MUP OODEA
GBT OODLA
DCI
MIG
Modularity
SAP
Figure 4: Higher disentanglement corresponds to better generalization across all OOD1 scenarios, as seen from the transfer scores (left). The transfer score is computed as the mean absolute prediction еπоо of ground truth factor values (lower is better). This correlation is particularly evident in the GBT case, whereas MLPs appear exhibit better OODI transfer with very high disentanglement only. These results are mirrored in the Spearman rank correlations between transfer scores and disentanglement metries (right).
Is disentanglement correlated with OOD1 generalization? In Fig. 4 we consistently observe a negative correlation between disentanglement and transfer error across all OOD1 settings. The correlation is mild when using MLPs, strong when using GBTs. This difference is expected, as GBTs have an axis-alignment bias whereas MIPs can given enough data and capacity - disentan- gle entangled representation more easily. Our results therefore suggest that highly disentangled representations are useful for generalizing aotoo-ddistrbutiion as long as the encoder remains in-distribution. This is in line with the correlation found by Locatello et al. (2019b) between disen- tanglement and the GBT10000 metric. There, however, GBTs are tested on the same distribution as the training distribution, while here test them under distributional shift. Given that the compu- tation of disentanglement scores requires labels, this is of little benefit in the unsupervised setting. However, it can be exploited in the weakly supervised setting, where disentanglement was shown to correlate with ELBO and reconstruction loss (Section 3). Therefore, model selection for repre- sentations that transfer well in these scenarios is feasible based the ELBO or reconstruction loss, when weak supervision is available. Note that, in absolute terms, the OOD generalization error with encoder in-distribution (OOD1) is very low in the wigh-disentanglemmet case (the only exception being the MLP in the OOD1-C case, with the 1-7 color split, which seems to overfit). This suggests that disentangled representations can be useful downstream tasks even when transferring out of the training distribution.
Summary: Disentanglement seems to be positively correlated with OOD generalization of down- stream tasks, provided that the encoder remains in-distribution (OODI). Since in the weakly super- vised case disentanglement correlates with the ELBO and the reconstruction loss, model selection can be performed using these metrics as proxies for disentanglement. These metries have the advan- tage that they can be computed without labels, unlike disentanglement metrics.
Is disentanglement correlated with OOD2 generalization? As seen Fig, 5 the negative cor- relation between disentanglement and GBT transfer error is weaker when the encoder is out of distribution (OOD2). Nonetheless, we observe a non-negligible correlation for GBTs in the OOD2- A case, where investigate out-of-distribution generalization along one FoV, with observations in Do still generated from the same simulator. In the OOD2-B setting, where the observations are taken from cameras in the corresponding real-world setting, the correlation between disentanglement and transfer performance appears to be minor at best. This scenario can be considered a variant of zero-shot sim2real generalization.
Summary: Disentanglement has a minor effect on out-of-distribution generalization outside of the training distribution of the encoder (OOD2).
7

---

Published as a conference paper at ICLR 2021
0.30
DCI score
0.4
0.25
x
Transfer score
MLP 00D2-A
4
0.20
GBT 00D2-A
-50
-50
13
-5
0.15
MLP OOD2-B
-13
-11
-11
88
0.10
GBT QOD2-B
-20
-20
-5
-13
0.05
DCI
MIG
Modulanity
SAP
0.00
MLP O0D2-A
GBT 00D2-A
MLP OOD2-B
GBT 00D2-8
Figure 5: Disentanglement affects generalization across the OOD2 scenarios only minimally seen
from transfer scores (left) and corresponding rank correlations with disentanglement metrics (right).
-2
0.30
Noise on input
0.30
Noise inpul
0.25
False
0.25
False
True
Transfer score
0.20
Transfer score
0.20
0.10
0:10
0.00
MLP 0001-C
0ST 00D1-G
MLP 0001-8
GBT 00018
MLP 0001A
MLP 0002A
GBT 0002-A
MLP 0002-8
GBT 0002-9
GBT 00014
Figure 6: Noise improves generalization across the O0D2 scenarios and less for the OOD1
scenarios as seen from the transfer scores. Top row: Spearman rank correlation coefficients between transfer metries and presence of noise in the input.
What else matters for 00D2 generalization? Results in Fig. 6 suggest that adding Gaussian noise the input during training described in Section 3 significantly better OOD2 gener- alization, and has no effect on OODI generalization. Adding noise to the input of neural networks is known lead to better generalization (Sietsma Dowj 1991 Bishop 1995). This in agreement with our results, since 0OD1 generalization does not require generalization of the encoder, while OOD2 does. Interestingly, closer inspection reveals that the contribution different factors of vari- ation to the generalization error can vary widely. See Appendix B.5 further details. In particular, with noisy input, the position of the cube predicted accurately even in real-world images (<5% mean absolute error each axis). This is promising for robotics applications, where the true state of the joints is observable but inference the cube position relies on object tracking methods. Fig shows example of real-world inputs and reconstructions of their simulated equivalents.
Summary: Adding input noise during training appears to be significantly beneficial for O0D2 generalization, while having effect when the encoder is kept in its training distribution (OOD1).
8

---

Published as a conference paper at ICLR 2021
6 CONCLUSION
Despite the growing importance of the field and the potential societal impact in the medical domain
Chartsias et al. 2018) and fair decision making (Locatello et al. (019a), state-of-the-art approaches
for learning disentangled representations have so far only been systematically evaluated on synthetic toy datasets. Here introduced new high-resolution dataset with 1M simulated images and over 1,800 annotated real-world images of the same setup. This dataset exhibits number of challenges and features which are not present in previous datasets: it contains correlations between factors, occlusions, complex underlying structure, and it allows for evaluation of transfer to unseen simulated and real-world settings. We proposed a new VAE architecture scale disentangled representation learning this realis- tic setting and conducted a large-scale empirical study of disentangled representations on this dataset. We discovered that disentanglement is good predictor of OOD generalization of downstream tasks and showed that, in the context of weak supervision, model selection for good OOD performance can be based the ELBO or the reconstruc- tion loss, which are accessible without explicit labels. Our setting al- lows for studying wide variety of interesting downstream tasks in the future, such as reinforcement learning or learning a dynamics model of the environment. Finally, believe that in the future it will be impor- tant to take further steps in the direction of this paper by considering settings with even more complex structures and stronger correlations between factors.
Figure Zero-shot trans- fer to real-world observa- tions of our models trained in simulation. Left: input; right: reconstruction.
ACKNOWLEDGEMENTS
The authors thank Shruti Joshi and Felix Widmaier for their useful comments on the simulated setup, Anirudh Goyal for helpful discussions and comments, and CIFAR for the support. We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Frederik Träuble.
REFERENCES
Tameem Adel, Zoubin Ghahramani, and Adrian Weller. Discovering interpretable representations for both deep generative and discriminative models. In International Conference Machine Learning, pp. 50-59, 2018.
OpenAl: Marcin Andrychowicz, Bowen Baker, Maciek Chociej. Rafal Jozefowiez, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-20, 2020.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: (el:::9070000000 2019.
Francis Bach and Michael Jordan. Kernel independent component analysis. Journal of Machine Learning Research, 3(7):1-48, 2002.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. arXiv preprint 2083::0003.04887. 2020.
Chris M Bishop. Training with noise is equivalent to tikhonov regularization. Neural computation, 7(1):108-116, 1995.
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Ben- gio. Generating sentences from a continuous space. arXiv preprint arXiv:1511.06344 2015.
Christopher P Burgess, Irina Higgins, Arka Pal, Loie Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in beta- VAE. arXiv preprint arXiv: 1804 03599, 2018.

---

Published as a conference paper at ICLR 2021
Agisilaos Chartsias, Thomas Joyce, Giorgos Papanastasiou, Scott Semple, Michelle Williams, David Newby, Rohan Dharmakumar, and Sotirios A Tsaftaris. Factorised spatial representation learn- ing: Application in semi-supervised myocardial segmentation. In International Conference Medical Image Computing and Computer-Assisel Intervention, pp- 490-498. Springer, 2018.
Tian Qi Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentangle- ment in variational autoencoders. In Advances in Neural Information Processing Systems, 2018.
Pierre Comon. Independent component analysis, a new concept? Signal Processing, 36(3):287-314, 1994.
Peter Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural Computation, 5(4):613-624, 1993.
Cian Eastwood and Christopher KI Williams. A framework for the quantitative evaluation of disen- tangled representations. In International Conference on Learning Representations, 2018.
Sanja Fidler, Sven Dickinson, and Raquel Urtasun. 3d object detection and viewpoint estimation with a deformable 3d cuboid model. In Advances neural information processing systems, pp- 611-619, 2012.
Muhammad Waleed Gondal, Manuel Wüthrich, Djordje Miladinović, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the transfer of inductive bias from simulation the real world: new disentanglement dataset. In Advances Neural Information Processing Systems, 2019.
Sven Gowal, Chongli Qin, Po-Sen Huang, Taylan Cemgil, Krishnamurthy Dvijotham, Timothy Mann, and Pushmeet Kohli. Achieving robustness in the wild via adversarial mixing with disen- tangled representations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1211-1220, 2020.
Luigi Gresele, Paul K. Rubenstein, Arash Mehrjou, Francesco Locatello, and Bernhard Schölkopf. The incomplete rosetta stone problem: Identifiability results for multi-view nonlinear ica. In Conference on Uncertainty in Artificial Intelligence (UAI), 2019.
Christina Heinze-Deml and Nicolai Meinshausen. Conditional variance penalties and domain shift robustness. arXiv preprinl 1710. 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAB: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations, 2017a.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot trans- fer in reinforcement learning. In International Conference on Machine Learning, 2017b.
Irina Higgins, Nicolas Sonnerat, Loic Matthey, Arka Pal, Christopher Burgess, Matko Bošnjak, Murray Shanahan, Matthew Botvinick, Demis Hassabis, and Alexander Lerchner. Scan: Learning hierarchical compositional visual concepts. International Conference on Learning Represen- fations, 2018
Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. In Advances in Neural Information Processing Systems, 2016.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural Networks, 1999.
Aapo Hyvarinen, Hiroaki Sasaki, and Richard E Turner. Nonlinear ica using auxiliary variables and generalized contrastive learning. In International Conference on Artificial Intelligncce and Statisties, 2019.
10

---

Published as a conference paper at ICLR 2021
Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to-sim: Data- efficient robotic grasping via -o-aniizddto- adaptation networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 12627-12637, 2019.
Christian Jutten and Juha Karhumen. Advances in nonlinear blind source separation. In International Symposium on Independent Component Analysis and Blind Signal Separation, pp. 245-256, 2003.
Hyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoen- coders and nonlinear ica: A unifying framework. International Conference on Artificial Intel- ligence and Statistics, pp. 2207-2217, 2020.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning, 2018.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.69801 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
David Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge, and Dylan Paiton. Towards nonlinear disentanglement natural data with temporal sparse coding. arXiv preprint wrXiv2207.110930 2020.
David Krueger, Ethan Caballero. Joern- Heen-Herrik- Jacobsen, Amy Zhang, Jonathan Binas, Remi Le Priol, and Aaron Courville. Out-of--istribuut generalization via risk extrapolation (rex). arXiv preprint arXiv:2003.0066 2020.
Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan. Variational inference of disen- tangled latent concepts from unlabeled observations. In International Conference on Learning Representations, 2018
Yann LeCun, Fu Jie Huang, and Leon Bottou. Learning methods for generic object recognition with invariance to pose and lighting. In IEEE Conference on Computer Vision and Pattern Recognition, 2004.
Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 624-639, 2018.
Francesco Locatello, Gabriele Abbati, Thomas Rainforth, Stefan Bauer, Bernhard Schölkopf, and Olivier Bachem. On the fairness of disentangled representations, In Advances in Neural Informa- fion Processing Systems, pp. 14611-14 1-14224, 2019a.
Francesco Locatello, Stefan Bauer, Mario Lucic, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions the unsupervised learning of disentangled repre- sentations. In International Conference on Machine Learning, 2019b.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly disentanglement without compromises. arXiv preprint arXiv:2002.02886 02880.00...8............. 2020.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning, pp. 10-18, 2013.
Xue Bin Peng, Marcin Andrychowicz. Wojeiech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 1-8. IEEE, 2018.
Scott Reed, Yi Zhang, Yuting Zhang, and Honglak Lee. Deep visual amalogy-making In Advances in Neural Information Processing Systems, 2015.

---

Published as a conference paper at ICLR 2021
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. arXiv preprint arXiv:1401.4083: 2014.
Karl Ridgeway and Michael C Mozer. Learning deep disentangled embeddings with the f-statistic loss. In Advances Neural Information Processing Systems, 2018
Mateo
Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. The Journal of Machine Learning Research, 1331)133--134, 2018.
Andrei A Rusu, Matej Večerík, Thomas Rothörl, Nicolas Heess, Razvan Pascanu, and Raia Hadsell. Sim-to-real robot learning from pixels with progressive nets. In Conference on Robot Learning, pp. 262-270, 2017.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Towards causal representation learning. arXiv preprint arXiv:2102.11 100.111071111 2021.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disen- tanglement with guarantees. arXiv preprint arXiv: 191.0099, 2019.
Jocelyn Sietsma and Robert JF Dow. Creating artificial neural networks that generalize. Neural networks, 4(1):67-79, 1991.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. In Advances in neural information processing systems, pp. 3738-3746, 2016.
Peter Sorrenson, Carsten Rother, and Ullrich Köthe. Disentanglement by nonlinear ica with general incompressible-fot networks (gin). arXiv preprint arXiv:2001. 044872, 2020.
Raphael Suter. Djordje Miladinovic, Bernhard Schölkopf, and Stefan Bauer. Robustly disentangled causal mechanisms: Validating deep representations for interventional robustness. In Interna- tional Conference on Machine Learning, pp. 6056-6065. PMLR, 2019.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba. and Pieter Abbeel. Do- main randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 23-30. IEEE, 2017.
Frederik Träuble, Elliot Creager, Niki Kilbertus, Francesco Locatello, Andrea Dittadi, Anirudh Goyal, Bernhard Schölkopf. and Stefan Bauer. On disentangled representations learned from correlated data, arXiv preprint arXiv:2006. 200::20006078866, 2020.
Sjoerd van Steenkiste, Francesco Locatello, Jürgen Schmidhuber, and Olivier Bachem. Are disen- tangled representations helpful for abstract visual reasoning? arXiv preprint arXiv: /2506, 2019.
Manuel Wüthrich, Felix Widmaier, Felix Grimminger, Joel Akpo, Shruti Joshi, Vaibhav Agrawal, Bilal Hammoud, Majid Khadiv, Miroslav Bogdanovic, Vincent Berenz, et al. Trifinger: An open- source robot for learning dexterity. arXiv preprint arXiv:2008.033500 2020.
Mengyuan Yan, Qingyun Sun, Iuri Frosio, Stephen Tyree, and Jan Kautz. How to close sim-real gap? transfer with segmentation! arXiv preprint arXiv:2005.07695, 2020.
12

---

Published as a conference paper at ICLR 2021
A IMPLEMENTATION DETAILS
Training. We train the B-VAEs by maximizing the following objective function:
c Blgg - (0ol)))) = log p(z)
with 3 > 0 Qusing the Adam optimizer (Kingma & Ba 2014) with default parameters. We use a batch size of 64 and train for 400k steps. The learning rate is initialized to le-4 and halved at 150k and 300k training steps. We clip the global gradient norm to 1.0 before each weight update. Following Locatello et al. (2019b), we use Gaussian encoder with an isotropic Gaussian prior for the latent variable, and a Bernoulli decoder. Our implementation of weakly supervised learning is based on Ada-GVAE (Locatello al. 2020), but uses a symmetrized KI. divergence:
(KL(P q) = (kI()lla) + 2 (KL(l))
to infer which latent dimensions should be aggregated.
The noise added to the encoder's input consists of two independent components, both iid Gaussian with zero mean: one independent for each subpixel (RGB) and has standard deviation 0.03, the other is 8 x 8 pixel-wise (greyscale) noise with standard deviation 0.15, bilinearly upsampled by factor of 16. The latter has been designed (by visual inspection) to roughly mimic observation noise in the real images due to complex lighting conditions.
Neural architecture. Architectural details are provided in Tables 2 and S and Fig. g provides a high-level overview. In preliminary experiments, we observed that batch normalization, layer nor- malization, and dropout did not significantly affect performance in terms of ELBO, model samples, and disentanglement scores, both in the unsupervised and weakly supervised settings. On the other hand, layer normalization before the posterior parametcrization (last layer of the encoder) appeared to be beneficial for stability in early training. While using an architecture based on residual blocks leads to fast convergence, in practice we observed that it may be challenging to keep the gradients in check at the beginning of training In order to solve this issue, we resorted simple scalar gating mechanism in the residual blocks (Bachlechner et al. 2020) such that each residual block is initialized to the identity.
Datasets and OOD evaluation. Because we evaluate OOD generalization in terms of cube color hue (except in the sim2real case), we first sampled 8 color hues at random from the 12 specified in Table 1 The chosen hues are: [0 120", 150 180%, 210", ,700, ,300 3300). Then, the dataset D used for training VABS is generated by randomly sampling values for the factors of variation from Table 0 with the color hue restricted to the above-mentioned values. This makes O0D2 evaluation possible, specifically OOD2-A where the learned predictors are tested on representations extracted from images with held-out values of the cube hue.
For evaluation of out-of duttof-distribution generalization, we train the downstream predictors on a subset D1 C D of the representation training set. The downstream training set D2 sampled random from D but only contains a (not necessarily proper) subset of the 8 cube colors. This subset contains 1 color in the OOD1-A case, colors in OOD1-B and OOD1-C, all colors in OOD2 (in this case Dj is simply random subset of D). Then we test the downstream predictors on set Dr distributionally different from D2 in terms of cube color (all OODI scenarios as well as OOD2-A) or sim2real (OOD2-B), B). In the OODI case, D2 is also a subset of D and is generated the same way. In each OODI case, the test set D2 is paired with its corresponding Dj that was used to train the downstream predictors. D2 contains all colors in D minus those in Dj. In the O0D2-A case, D2 is separate dataset containing 5k simulated images like those in D. except that these only contain the 4 colors that were left out from the VAE training set D (hue in [30%, 60 90% 240%).. In the OOD2-B case, the set D2 is the dataset of real images. Following previous work (e.g. the GBT10000 metric in Locatello et al. (20196)) the training set D2 and test D2 for downstream tasks contain 10k and 5k images, respectively, except in the OOD2-B case, where the size is limited by the size of the real dataset.
'This instability may also be exacerbated in probabilistic models by the sampling step in latent space, where a large log variance causes the decnder input to take very large values. Intuitively, this might he reason why layer normalization hefore latent space appears to he heneficial for training stability.
13

---

Published as a conference paper at ICLR 2021
Encoder
Decoder
Operation
Output Shape
Operation
Output Shape
Input
128x 128x 64x64x64
Input
d
Conv 5x5, stride 2, 64 ch. Le.0))eL.U(0.02)
FC(512)
512
LeakyReLU(0.02)
2x ResidualBlock(64)
FC(4096)
4096
Conv 1x1, 128 channels AvcragePool(2)
64x64x 128
Reshape
4x4x256
x2 x 128
2x ResiduulBBock((56)
Residual/Block(128)
Bllinealnneeppoaatio(())
8x8x250
AveragePool((2)
16 16 128
2x ResidualBlock(256)
2x ResidnalBlock((133)
Conv 1x1, 128 channels
8x8x 128
Conv 1x1, 256 channels
16 x x
Bilinearlnterpolation(()
16 16 128
AveragePool(()
5x8x222
2x ResidualBlock
2x Residual
Bilinearlnterpolation((]
x2 128
AveragePool((2)
4x4x256
2x ResidualBlock(128)
2x ResidualBlook((55)
Conv 1x1, 64 channels
x2
Flatten
4096
Bilinearlnterpolationn()
64 64
LeakyReLU(0.002)
2x ResidualBlock((4) Bilinearinterpolatinn2)) LeakyReLU Conv 5x5, K channels
FC(512)
512
128 x 128 64
Leak)ReLU(0.022)
LayerNorm
128 x 128 x K
FC(d)
2d
Table 2: Encoder (left) and decoder (right) architectures. The latent space dimensionality is denoted by d, and = 3 indicates the number of image channels. Last line in the encoder architecture: the fully connected layer parameterizing the log variance of the approximate posterior distributions of the latent variables has custom initialization. The weights are initialized with 1/10 standard deviation than the default value, and the biases initialized -1 instead of O. Empirically, this together with (learnable) LayerNorm was beneficial for training stability at the beginning of training.
Residual Block
Input: shape H x xW X LcakyReLU(0.02) Conv 3x3, C' channels LeakyReLU(O.02)
Conv 3x3, C channels
Scalar gate
Sum with input
Table 3: Architecture of one residual block. The scalar gate is implemented by multiplying the tensor by a learnable scalar parameter before adding it to the block input. Initializing the residual block to the identity by setting this parameter has been originally proposed by Bachlechner et al. (2020). The tensor shape constant throughout the residual block.
14

---

Published as a conference paper at ICLR 2021
pool
pool
1x1
pool
sd
1x1
Bacteen
42
276
356
Ioput
ceslaaee
756
b14
125
125
64
101
up comm
Figure 8: Schemes of the encoder (top) and decoder (bottom) architectures. In both schemes, in- formation flows left to right. Blue blocks represent convolutional layers: those labeled "conv" have 5x5 kernels and stride 2, while those labeled "1x1" have 1x1 kernels. Each orange block represents a pair of residual blocks (implementation details of residual block are provided in Table Green blocks in the encoder represent average pooling with stride 2, and those in the decoder denote bi- linear upsampling by a factor of 2. Red blocks represent fully-connected layers. The block labeled "norm" indicates layer normalization. Dashed lines denote tensor reshaping.
15

---

Published as a conference paper at ICLR 2021
B ADDITIONAL RESULTS
B.1 DATASET CORRELATIONS
Feasible states when 1st DoF
0.4
angle of 2nd DoF
0.2
0.0
-0.2
-0.4
0.00
angle 3rd DoF
Figure 9: Feasible states of the 2nd and 3rd DoF when the angle of the 1st DoF is 0. Angles are in radians.
Distribution of feasible states
0.4
angle of 2nd DoF
0.2
0.0
-0.2
-0.4
-0.8
-0.6
-0.4
-0.2
0.0
0.2
0.4
0.6
0.8
angle of 3rd DoF
Figure 10: Density of feasible states of 2nd and 3rd DoF the whole training dataset. Darker shades of blue indicate regions of higher density. Angles are in radians.
16

---

Published as a conference paper at ICLR 2021
B.2 SAMPLES AND RECONSTRUCTIONS
Figure 11: Samples generated by trained model. This model was selected based on the ELBO.
Figure 12: Input reconstructions by a trained model. This model was selected based on ELBO. Image inputs are on odd columns, reconstructions on even columns.
17

---

Published as a conference paper at ICLR 2021
B.3 LATENT TRAVERSALS
(a) Low disentanglement
(b) Medium disentanglement
(c) High disentanglement
Figure 13: Latent traversals for model with low DCI score (0.15) in (a), medium DCI score (0.5) in (b), and high DCI score (1.0) in (c).
18

---

Published as a conference paper at ICLR 2021
B.4
UNSUPERVISED METRICS AND DISENTANGLEMENT
Supervision
Supervision
unsupervised
unsupervised
0.8
weakly supervised
0.8
weakly supervised
0.6
MG
0.4
02
0.0
0.0
-26440
-26430
-26420
-26410
-26400
-26300
26360
26370
26380
26390
25400
26410
26420
Reconstruction
1.0
0.8
0.8
0.6
0.6
g
0.4
02
Supervision
0.2
Supervision
unsupervised
unsupervised weakly supervised
weakly supervised
-26400
-26430
-6420
-26410
-26400
-26390
26360
26370
26380
26390
26400
26410
26420
ELBO
Reconstruction
Figure 14: Scatter plots of unsupervised metrics (left: ELBO; right: reconstruction loss) disentan- glement (top: MIG; bottom: DCI) for 1,080 trained models, color-coded according to supervision. Each point represents trained model.
19

---

Published as a conference paper at ICLR 2021
B.5 OUT-OF-DISTRIBUTION TRANSFER
0.40
Noise on input
0.35
0.30
Transfer scare
0.20
0.10
0.05
0.00
mean
GBT OOD2-A moan
MLP OOD2-4 Tox DoF
GBT 00D2-4 Tet
MLP CO02-A 2nd DoF
GBT OOD2-A 2nd 2nd DoF
GBT OOD2-A 3rd DoF
MLP OOD2-A cube
cubex
MLP 0002 A cube y
GBT OOD2-A cubo
GBT OOD2-A cube rot
O0D2-A
OOD2-A
Noise on input
0.35
False
True
Transfer scote
0.25
0.20
0.15
0.10
0.05
0.00
MLP B mean
GBT 0002-8 22-- mean
MLP O0D2-8 Tet
GBT 0002-8 DoF
MLP 0002-B 2nd DoF
GBT 0OD2-B 2nd DoF
MLP 00D2-8 DoF
GBT 0002-8 3rd DoF
MLP 00D2-B cube
x
MLP 00D2-8 8 cubo y
GBT 00D2-8 cube
GBT 00D2-8 cube rot
GBT
MUP
Figure 15: Transfer metric in OOD2-A (top) and OOD2-B (bottom) settings, decomposed according to the factor of variation and presence of input noise. When noise is added to the input during training, the inferred cube position error is relatively low (the scores the mean absolute error, and they are normalized to (0, 1)). This is particularly useful in the OOD2-B setting (real world) where the joint state is anyway considered known, while object position has be inferred with tracking methods.
20

---

Published as a conference paper at ICLR 2021
B.6 OUT-OF-DISTRIIUTIIO RECONSTRUCTIONS
(a) Low disentanglement
(b) Medium disentanglement
(c) High disentanglement
Figure 16: Reconstructions of real-world images (OOD2-B) for a model with low DCI score (0.15)
in (a), medium DCI score (0.5) (b), and high DCI score (1.0) in (c).
21

---

Published as a conference paper at ICLR 2021
Low disentanglement
(b) Medium disentanglement
(c) High disentanglement
Figure 17: Reconstructions of simulated images with encoder out-of-distribun colors (OOD2-A) for a model with low DCI score (0.15) in (a), medium DCI score (0.5) in (b), and high DCI score (1.0) (c).
22


Reviews:
[{"review_id": "ci2ulrcgR-O", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "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": "The work has sufficient quality, and is sufficiently clear and original. For detail, please see the below pros and cons. \n\nOverall Pros:\n- Dataset contains dependencies in FoV, not artificially induced but present due to the attempt at realism in the simulation, a realistic aspect which hasn't been addressed in previous work. \n-The complete set of factors (e.g. product of number of possible values for each factor), totals to 2.92B, thus the dataset contribution allows for addressing the generalization problem which hasn't been addressed in previous work. \n-Disentanglement metrics notably cannot be evaluated on acquired camera images, so the recording of an annotated dataset under the same conditions in a real-world setup enables researchers to study this problem in a more controlled manner. \n-The observation that deeper and wider autoencoders can arguably scale to datasets with fine-grained factors is an informative observation for the community. \n-Similar to [1], the results collected in the extensive experimental study provide information to the community which can be analyzed and built upon in future work. \n-The OOD evaluation is novel, interesting, and informative. Though the downstream task is quite connected with the disentanglement evaluation itself, the consideration of this setting is a contribution to the community. \n-The negative observation in sim-to-real transfer is a net positive for this work, as it is informative to the community in specifying a problem which can be worked on in future research. \n-The observation that adding input noise in training, a common tool in classic sim-to-real robotics work, is useful for transfer in representations is practical and meaningful. \n\nOverall Cons:\n- For the evaluation prior to OOD evaluation, it appears disentanglement is evaluated on the full collected dataset. If that is the case, though the dataset itself provides the ability to evaluate generalization, the evaluation prior to OOD evaluation does not evaluate generalization. It is not clear then if overfitting played a role in the observed results. \n- Fixing k=1 for the weakly supervised method implies control over the data generating process often unavailable in practical applications. At the same time, said experiment can be extended to test how performance degrades with a mismatch in assumptions. Considering k != 1, as well as PCL [2] and SlowVAE [3] (which assume laplacian transitions) could provide interesting comparisons. \n- The authors state that \"many\" trained models fully disentangle the factors, but this is \"only\" possible in the weakly supervised scenario. Clarifying details would be beneficial here, is it true that none of the unsupervised models fully disentangled the factors? If some did, was there an underlying correlation between successful models that can be discussed, or was it simply random? Is using latent traversals alone as a basis sufficient for claiming \"full disentanglement\"? More specific statements in the description would be desirable. \n- The authors appear to pre-assume that since latent traversals agree with the theory of [1], quantitative metrics which do not score weakly supervised models over unsupervised models are \"ineffective at capturing disentanglement in this setting\". This reads as confirmation bias and should clarified, either by justifying why metrics which do not agree with latent traversals are \"ineffective\", or presenting the information without implying judgement on whether the metrics are \"ineffective\" or not. \n- \"We stress that the OOD2 scenario, which is typically not studied extensively, is only possible because the representations are not trained on the entire dataset.\" Does this mean the representations learned for OOD evaluation are trained on less data than what was considered for disentanglement evaluation? Clarifying details would be beneficial here, as this standalone statement simply yields unanswered questions. \n- How was D1 selected from D, by what sampling process, and how was the number of samples decided upon? For in-distribution generalization, how was the split between training and held-out set in D1 conducted, what was the percentage for the split, and how was it selected? Clarifying detail would be beneficial. \n- \"Since the values are normalized, we can take the average of the MAE over all factors (except for the FoV which is OOD).\" Why can you not take the average of the MAE over the FoV which are OOD? This statement could be interpreted as some factors being OOD while some not, when a split within the factor set was not mentioned previously. This sentence requires clarification. \n- Many questions brought up by Section 4 (such as the above), are addressed in part in the experimental setup section for Section 5. I'd suggest providing said information earlier so the confusion yielded in Section 4 is mitigated. \n- Notably, the color hues simply mentioned in Table 1, without any explanation on why they are not included within Table 1, appear to be what is ablated upon for OOD1. It would have been beneficial to have earlier clarification on this point. Furthermore, distribution shift in the color hue but all underlying factors being the same is specific type of distribution shift not clarified to the reader. Consideration of discrepancies in renderer would be interesting, but more importantly, discrepancies in the factor value set. The distribution shift considered here is shift in nuisance factors, a useful test but limited in its scope. \n- \"Our results therefore suggest that highly disentangled representations are useful for generalising out-of-distribution as long as the encoder remains in-distribution\" This statement can be seen as a bit misleading given the data is only out-of-distribution in terms of nuisance factors, not the factors the model was evaluated on disentanglement with respect to. It's intuitive that a highly disentangled representation will capture the ground truth factors well, and thus be more robust to shift in nuisance factors, and it is useful to show this, but the fact that we are looking at shift in nuisance factors should be clear to the reader.\n- The justification for training half of the models with gaussian noise is provided in the final page of the paper. As with the previous comments, it would be beneficial to the reader to make this clear when introducing the experimental design decisions.  \n\nConclusion: This work provides significant contribution in experimental analysis on the performance of disentanglement with the inclusion of realistic complexities. Both the positive and negative observations are clear, interesting, and informative. While I suggested many ablations which could provide the reader with further information, I don't see the lack of inclusion of said ablations as an issue in the work, the experimental study is already extensive. The main issue I see with the work as it is merely comes from a writing standpoint, ensuring the statements made are justified, and ensuring the detail is there in the descriptions for the reader to understand, and most importantly, be able to reproduce, said results. \n\nOverall, I see this is a valuable contribution, but would like to see improvement in the writing given the points brought up in the revision. \n\nQuestions:\n- Can you describe how given a particular state for the factors of variation, how said factors are processed to render the observation? A description of how each factor individually influences the observation would be beneficial in the main paper, if only briefly, such that the reader can see why it is (1) challenging and (2) requires modeling of fine details. \n- \"The training set for the VAEs contains 8 randomly chosen color hues.\" How do these color hues differ from the cube color hues? How do they specifically affect the rendering? Clarifying detail for this statement is needed. \n- Were any ablation studies performed for the model architecture? Many model details are presented without any experimental testing discussed, an ablation study would be informative to the reader. \n- Did you test the behavior of models with latent_dim < factor_dim, e.g. latent space dimensionality of 5, as well as latent_dim = factor_dim, e.g. latent space dimensionality of 7? Such results could be interesting to see if, and how much, performance degrades not only as latent_dim progressively increases from factor_dim, but also when latent_dim progressively decreases from factor_dim, considering the equality case as well. \n- \" Finally, note that the BetaVAE and FactorVAE metrics are not straightforward to be evaluated on datasets\nthat do not contain all possible combinations of factor values.\" Could you detail why they are not straightforward to evaluate? Clarifying detail would be helpful. \n- Did the authors consider other unsupervised model selection techniques, such as UDR, to see if this under or outperformed the weakly supervised loss selection method?\n- Why the choice of MAE? Was this choice ablated on?\n\n[1]: Francesco Locatello, Stefan Bauer, Mario Lucic, Sylvain Gelly, Bernhard Scholkopf, and Olivier ¨\nBachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In International Conference on Machine Learning, 2019.\n\n[2]: Aapo Hyvärinen and Hiroshi Morioka. Nonlinear ica of temporally dependent stationary sources. In Proceedings\nof Machine Learning Research, 2017. \n\n[3]: David Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaninov, Wieland Brendel, Matthias Bethge,\nand Dylan Paiton. Towards nonlinear disentanglement in natural data with temporal sparse coding.\narXiv preprint arXiv:2007.10930, 2020.", "extra_scores": {}, "extra_text": {}}, {"review_id": "RdIANJDfnSG", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 9, "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\n\nThe paper presents a new, more complex, dataset for the use of disentangled representation learning. The dataset is based on real and simulated images of the trifinger robot platform. There are 7 factors of variation with high-resolution measurements of these factors. The dataset contains over 1 million simulated images and another ~1000 annotated images of a real trifinger robot arm.\n\nThe authors also present a new neural architecture to scale disentanglement on more complex datasets and present a large empirical study on the performance of various techniques on out-of-distribution downstream task performance.\n\n\n## Quality & Clarity\nThe paper itself is well written, with a structure that effectively guides the reader through the work and results.\n\nThe dataset, significance and experiments are clearly outlined.\n\n## Originally & Significance\n\nThe novelty of the dataset is clear. Providing a complex disentanglement dataset where the underlying factors of variation are inherently correlated. The in- and out-of-distribution experiments are possible because of the presence of both synthetic and real-world data.\n\nThe experiments run are repeated multiple times and the results are convincing. They use both unsupervised and weakly supervised approaches and the results are both intuitive and supported by the literature. \n\nThe experiments on out-of-distribution representation transfer are interesting and show that disentangled representations can lead to better transfer to out-of-distribution tasks.\n\n\n## Outcome Rationale\n\nThis dataset is likely to be extremely useful to the community going forward and work disentangled representation learning is likely to benefit from it. The experimental setups are sensible and the largescale benchmarks support the use of disentangled representations when transferring from simulated to real-world scenarios.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "DFD-INSEhnN", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 5, "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": "The authors proposed a unique learning scheme for representation disentanglement. However, unlike infoGAN or ACGAN which explicitly learn disjoint feature representations for describing the attributes of interest (via unsupervised and supervised settings, respectively), the authors chose to address this task in a questionable \"weakly supervised setting\". More specifically, the authors chose to train AE-like model using pairwise images, which the difference between each pair of the inputs is only associated with one attribute of interest (e.g., angle, position, etc.). \n\nFor some reasons, the authors expected such a training scheme would result in learning feature representation in which only \"one\" feature dimension would reflect such attribute differences. This is a very strong assumption, since it is very likely that more than one feature dimensions would correspond to such changes. \n\nMoreover, assuming that precisely one feature dimension would be associated with the attribute of interest by feeding in a pair of images with exactly this attribute change would not be practical either. Most real-world images would be complex and contain multiple attributes. Making this assumption would imply that the training images are not realistic. \n\nAs for the evaluation, there is no comparison to any baseline or SOTA representation disentanglement methods, I found the quantitative metrics selected by the authors not sufficiently informative or supportive either. Most importantly, the authors claimed that the features trained by VAE allowed improved performances (e.g., Figs 3~5). Since VAE are simply trained in a unsupervised way (even the authors called their setting a weakly supervised setting), I see no evidence why the resulting features would be any different from those derived from standard VAEs, and why improved disentanglement results could be achieved.\n\nBased on the above observations and remarks, I feel that the authors would not be able to deliver a work which is technically strong with sufficiently complete evaluation. Therefore, I do not think this paper is above the ICLR standard for acceptance.", "extra_scores": {}, "extra_text": {}}, {"review_id": "CxilaU-B9I0", "reviewer": "AnonReviewer1", "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": "Summary: \nThis paper identifies that traditional datasets used for learning disentangled representation have several shortcomings such as no correlation between variables and simple structure.\nIt proposes a new dataset that has 1M higher-resolution simulated images along with 1K annotated real-world images of the same setup and gives analysis on disentangled representations on the dataset. \nIts results suggest that disentangled representations can result in better out of distribution task performances. \n\n==================================\n\nStrength:\n- Identified weaknesses of the previous datasets and proposed a new dataset that exhibits correlations between different variables. This is an important aspect of real-world scenarios. \n- Provided thorough experiments on disentangled representations and their metrics on the proposed dataset. \n\nWeakeness:\n- Experimental results for showing more disentangled representation results in better OOD task performances is somewhat expected. Unlike the title, it is not clear if the paper shows sufficient transfer of disentangled representations in realistic settings. It would be great to propose a way to make the generalization better for the settings the models have trouble with (OOD2 generalization). \n- This paper tried one approach by adding noise during training, but it leads to my second concern that the real world observation is very similar to the simulated data. With some gaussian noise, they would look similar. Therefore, it might not be sufficient to show that we can use this approach for sim-to-real transfer. \n- The paper claims that the proposed dataset (which is interesting) is challenging and highly complex, but the rendered images look easy enough for a simple neural network model to learn to reconstruct. Datasets like CLEVR, although they might not be used for measuring disentanglement, seems more complex and exhibits occlusions.  \n\n==================================\n\nWhile the proposed dataset is interesting, I am not confident this paper is showing evidence for the usefulness of disentangled representations for transfer learning in realistic settings. It is also not clear if the dataset is more useful than others because of the reasons stated above. \n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new dataset for evaluating disentanglement and its impact on out of distribution generalization based on the trifinger robotics platform. Using this dataset, the authors rigorously investigate the performance of beta-VAEs in this setting under a number of conditions, finding that weak supervision is necessary to induce disentangled representations, and that, perhaps surprisingly, disentanglement does not help for sim2real settings despite the similarity between the simulator and the real data. Reviewers were divided on the work, but had a number of concerns related to the claims of novel architecture, comparisons to baselines, and issues with the clarity of the paper, some of which were addressed in the authors' response. I agree with some of these concerns, particularly with respect to the claims of novel architectures since the modifications could simply be viewed as tweaking hyperparameters and are not rigorously compared to baselines. However, I think the novelty of the dataset and the rigorous evaluation of OOD generalization settings is likely to be valuable enough to the community to merit acceptance. I'd encourage the authors, however, to tone down some of the claims regarding the architecture (or provide sufficient baseline comparisons), and instead focus on the dataset and the OOD results. I recommend acceptance. 
Decision: Accept (Poster)