Overlooked Implications of the Reconstruction Loss for VAE Disentanglement

1 Introduction

A fundamental challenge in machine learning is discovering useful representations from high-dimensional data that can be used to solve subsequent tasks effectively. Recently, deep learning approaches have showcased the ability of neural networks to extract meaningful features from high-dimensional inputs for tasks such as classification [Krizhevsky et al., 2012] and reinforcement learning [Mnih et al., 2015]. However, these learned representations are often not semantically meaningful, which can negatively impact interpretability, fairness [Locatello et al., 2019a], and downstream task performance [Locatello et al., 2019b].

Prior work has therefore argued that it is desirable to learn a representation that is disentangled [Bengio et al., 2013]. While there is no consensus on what constitutes a disentangled representation, it is generally agreed that it should be factorised so that each latent variable corresponds to a single explanatory variable responsible for generating the data [Burgess et al.,

2017]. For example, a single image from a video game may be represented by continuous latent variables governing the x and y positions of the player or enemies, as well as categorical variables governing their clothing or appearance.

A common approach to discovering these representations is variational autoencoders (VAEs) [Kingma and Welling, 2014], which are trained on unlabelled data to learn a lowerdimensional representation capable of reconstructing the input. However, it has been shown that unsupervised methods cannot reliably learn representations without the introduction of supervision or inductive biases [Locatello et al., 2019b]. The recently introduced Ada-GVAE framework partially overcame this problem by using a weakly supervised signal to discover underlying factors [Locatello et al., 2020], but there remains room for improvement.

Interestingly, VAEs do not have an explicit mechanism that encourages the learning of disentangled representations, but it is theorised that this behaviour is related to the regularisation term and the information bottleneck principle [Burgess et al., 2017; Mathieu et al., 2019; Rolinek et al., 2019]. However, despite this hypothesis, there is still no explicit reason for why the representations learnt by these frameworks should align with generative factors in the data. Nonetheless, these frameworks have been shown to produce disentangled representations when trained on synthetically generated data, as measured by appropriate metrics [Eastwood and Williams, 2018; Chen et al., 2018; Zaidi et al., 2020].

In this paper, we aim to understand why VAEs implicitly learn disentangled representations by investigating the interaction between the reconstruction loss of the VAE and the input data. We provide compelling evidence that disentanglement occurs not because of special algorithmic choices or the regularisation term, but because of how VAEs perceive distances between observations in the datasets themselves according to the reconstruction loss, and the fact that these distances accidentally correlate to the chosen ground-truth factors generating the data. In particular, we find that standardised benchmarks are constructed in such a way that they unintentionally encourage models to learn what appear to be disentangled representations.

The main, summarised contributions of this paper 1 are:

1 Our extended paper with factor importance results and full experiment details is available at: https://arxiv.org/abs/2202.13341

• (i) We introduce the concept of perceived distance, in terms of the VAE reconstruction loss, to measure overlap or similarity between dataset pairs. We demonstrate that perceived distances in existing datasets unintentionally correspond to the distances between ground-truth factors, and that VAEs learn these distances, explaining why learnt representations may appear disentangled.
• (ii) We provide a technique to visualise the correlation between perceived distances in the data and ground-truth factors generating the data. We use this understanding to provide a theory for what constitutes an adversarial dataset under a given reconstruction loss.
• (iii) We reveal the ineffectiveness of state-of-the-art models by using our theory to design a simple, example adversarial dataset with constant perceived distance between elements, over which VAE-based frameworks fail to learn disentangled representations.
• (iv) We provide an example solution to the adversarial dataset that modifies the reconstruction loss, and thus perceived distances across the dataset, so that VAE frameworks are again able to capture the ground-truth factors.
• (v) We contribute Disent, a general PyTorch [Paszke et al., 2017] disentanglement framework, with common models, metrics, and datasets. 2

2 Background

Assume a dataset X = {x(0) , ..., x(n)} is a set of independent and identically distributed (i.i.d) observations x ∈ RN, generated by some random process involving an unobserved random variable z ∈ RD of lower dimensionality D ≪ N. Additionally, the true prior distribution z ∼ p∗(z) and true conditional distribution x ∼ p∗(x|z) are unknown. Variational autoencoders (VAEs) aim to learn this generative process. Unlike autoencoders (AEs), which consist of an encoder fϕ(x) = z and decoder gθ(z) = ˆx with weights ϕ and θ, VAEs instead construct a probabilistic encoder by using the output from the encoder or inference model to parameterise approximate posterior distributions z ∼ qϕ(z|x). The approximate posterior is then sampled during training to obtain representations z, which are then decoded using the generative model to obtain reconstructions ˆx ∼ pθ(x|z).

A factorised Gaussian encoder [Kingma and Welling, 2014] is commonly used. The posterior is modelled using a multivariate Gaussian distribution with diagonal covariance z ∼ N(µϕ(x), σϕ(x)), and the prior is given by the multivariate normal distribution pθ(z) = N(0, I), with a mean of 0 and diagonal covariance I. To enable backpropagation, the reparameterisation trick in Equation (1) is used to sample from the posterior distribution by offsetting the distribution means by scaled noise values. 3

VAEs maximise the evidence lower bound (ELBO) by minimising the loss given by Equation (4). VAE-based approaches often make slight modifications to this loss [Higgins et al., 2016; Zhao et al., 2017; Hou et al., 2017; Kumar et al., 2018; Chen et al., 2018; Kim and Mnih, 2018; Locatello et al., 2020], but the terms of these modified loss functions can usually still be grouped into reconstruction and regularisation components, given by Equations 2 and 3 respectively. The regularisation term affects the representations learnt by the encoder, while the reconstruction term improves the outputs from the decoder. These terms usually contradict in practice, with strong regularisation leading to worse reconstructions but often better disentanglement [Higgins et al., 2016; Burgess et al., 2017].

2 Disent framework repository: https://github.com/nmichlo/disent. Code is provided under the MIT license.

3 The notation ⊙ represents the element-wise product.

2.1 Random Sampling Reorganises VAE Embeddings

Disentanglement in VAEs is generally attributed to the regularisation term in Equation (3); however, we highlight that regularisation only enables the underlying disentanglement mechanism. Disentanglement arises rather as a result of VAEs reorganising the latent space to minimise reconstruction mistakes due to random sampling from the probabilistic encoder during training. Through this mechanism, a VAE will place similar observations according to the reconstruction loss in Equation (2) close together in the latent space [Burgess et al., 2017; Mathieu et al., 2019; Zietlow et al., 2021], as this action minimises sampling errors.

The regularisation term enables this interaction by controlling the overlap between latent distributions corresponding to different inputs. If these distributions overlap sufficiently, the decoder will often attribute a random sample to an incorrect input, see Figure 1. Thus, a mistake will be made during the decoding process, which encourages reorganisation to minimise the reconstruction error.

3 RelatedWork

The following works are the most applicable to our research, falling under three general categories: (i) explanations for disentanglement, (ii) the role of the reconstruction loss in disentanglement, and (iii) problems with disentanglement.

Firstly, Burgess et al.[2017] relate VAEs to the information bottleneck principle. Which explains that random sampling leads to a local minimisation of the reconstruction loss which reorganises the latent space so that points close in pixel space are close in the latent space. Mathieu et al.[2019] argue that VAEs do not explicitly encourage disentanglement through their design. Rather, they provide the explanation that the diagonal prior typically used in VAEs when combined with random sampling produces a similar effect to PCA. Our work takes inspiration from these ideas to develop the theory of perceived overlap in VAEs, which we use to analyse datasets and improve or hinder disentanglement.

Secondly, inspired by Burgess et al.[2017] most modern frameworks offer some way to balance the regularisation and reconstruction components of the loss, the ControlVAE automates this process [Shao et al., 2020]. Hou et al.[2017] instead swap out the reconstruction loss of VAEs for a perceptual loss function, which can improve the representations learnt by the model. Zietlow et al.[2021] extend the analysis of Mathieu et al.[2019]; However, emphasis is placed on constructing adversarial datasets that hinder disentanglement performance using a mild transformation, obtained from trained models which achieve poor disentanglement scores. Our work provides intuition by constructing an example adversarial dataset that targets a specific reconstruction loss, and then remedies this problem by adjusting the loss.

Finally, Locatello et al.[2019b] show that useful representations cannot be reliably learnt with unsupervised methods, unless inductive biases are introduced, and Gondal et al.[2019] show that representations learnt on synthetic data often do not transfer well to real-world data. Our work investigates the interplay between the reconstruction loss and data as the main bias in VAEs, standard choices accidentally disentangle synthetic data.

4 Existing Disentanglement Datasets

Consider the 3D Shapes dataset [Burgess and Kim, 2018] in Figure 2a, which contains observations of shapes fixed in the centre of the image with progressively changing attributes or factors such as size and colour. If, as humans, we are given unordered observations from a traversal along the size factor of 3D Shapes, it would be easy to order these observations using a perceived increase or decrease in the size of the shape. We might even say that the shapes in the images overlap by different amounts, considering shapes that are closer in size to possess more overlap, and thus also considering them as closer together in terms of distance. This idea naturally extends to VAEs ordering pairs of observations, and so we seek to investigate the correspondence between how these frameworks perceive distances over data points to reorganise the latent space and the ground-truth factors themselves.

4.1 Dataset Ground-truth Distance

Synthetic datasets [Burgess and Kim, 2018; LeCun et al., 2004; Matthey et al., 2017; Reed et al., 2015; Gondal et al., 2019] used for benchmarking disentanglement frameworks are all generated from F ∈ N+ ground-truth factors of variation. Each factor i ∈ [F] represents some property about the data that can be varied, 4 and has a dimensionality or size of fi > 0

4 The bracket notation [F] gives the natural numbers set {1, ..., F}.

where fi ∈ N+. The set of all factors used for generating the dataset is written as Y = [f1] × ... × [fF]. The full dataset is generated from this set of factors using some ground-truth generative process X = {g∗(y) | y ∈ Y}. Examples of this generative process are given in Figure 2.

With this construction of synthetic datasets, it is fitting to describe the ground-truth distances between observations x(a) , x(b) ∈ X using the Manhattan or ℓ1 distance between their corresponding ground-truth factors y(a) , y(b) ∈ Y, 5 as in Equation (5). It is important to note that this choice may not be optimal for single factors; rather, ℓ1 distance naturally aligns with how the datasets are constructed.

4.2 VAE Perceived Distance

With the idea of ground-truth distances between observations, we need a distance measure between observations as perceived by VAE frameworks. We derive the perceived distance between dataset elements from the noisy sampling procedure and the chosen reconstruction loss in a VAE framework.

Let z(b) ∼ qϕ(z|x(a) ) be a (possibly incorrect) sample from the posterior distribution corresponding to some input element x(a) ∈ X . Since the regularisation term encourages latent distributions to overlap, this sample z(b) may be incorrectly attributed by the decoder to a distribution corresponding to some other element from the dataset x(b) ∈ X , with reconstruction ˆx(b) ≈ x(b) . As the VAE objective consisting of the regularisation and reconstruction losses is jointly optimised, the decoder becomes better at reconstructing the inputs. In an ideal scenario, the inputs map to outputs (ˆx → x), and reconstructions are samples from our dataset: ˆx ∈ X . While this is not the case in practice due to the regularisation term, we derive the perceived distance in Equation (7) from this assumption that ˆx → x. This allows us to directly compare the elements x(a) , x(b) ∈ X within a dataset using the reconstruction loss as a distance function:

The perceived distance depends on the choice of reconstruction loss, which in literature is usually the pixel-wise Mean Squared Error (MSE) for data that is assumed to be normally distributed. We assume MSE is used throughout the rest of this work, unless specified. Note that analyses are similar for other pixel-wise losses, such as Binary Cross-Entropy (BCE).

4.3 Perceived Distances Correspond to Ground-Truth

In Section 4.1, all the ground-truth factors of a dataset are defined as the set Y = [f1] × . . . × [fF]. In Equation (9), we now define a factor traversal Y (a,i) ⊂ Y as the ordered set of all the coordinates along a factor i ∈ [F] such that the set passes through a point y(a) ∈ Y. The number of elements in the traversal is equal to the size of the chosen factor |Y (a,i) | = fi, and each element in the traversal generates the same traversal ∀y(b) ∈ Y (a,i) , Y (a,i) = Y (b,i) . Figure 2 gives examples of traversals.

5 When indexing y(a) , z(a) , x(a) , for convenience a may be an integer or a ground-truth factor a = y(a)

We compute the distance matrix ˜D(a,i) ∈ Rfi×fi , for some distance function d, between pairwise elements along a factor traversal Y (a,i) , written in Equation (10) using matrix notation.

To examine the ground-truth factors within our datasets, we compute the average distance matrix D(i) = Ea∈Y[ ˜D(a,i) ] for each factor i ∈ [F]. We plot these results in Figure 3 for both the ground-truth distance dgt and perceived distance dpcv. It is immediately obvious from these plots that the groundtruth distances and the distances perceived by a VAE may accidentally correspond.

Finally, we relate our work to Burgess et al.[2017] by outlining a direct approach in our extended paper in Footnote 1 for computing relative factor importance using perceived distances. A factor is considered more important if a VAE prefers to learn it before another factor.

4.4 VAEs Learn Perceived Distances

We compute distance matrices over a trained β-VAE at various levels of the network, including the representation layer and reconstructions. At each level of the VAE, the learnt distances all correspond to the original perceived distances already present within the dataset, see Figure 4. Since VAEs reorganise the embedding space according to the perceived distances, and noting results from Section 4.3, VAEs may discover structures that are similar to the underlying ground-truth factors.

Our results in Figures 3 and 4 provide empirical evidence that VAEs mimic the distances already present in the dataset according to the reconstruction loss. To appear disentangled if the goal is factored representations, individual latent units will need to encode portions of the distances that correspond to factors within the data. However, it is known that VAEs with diagonal priors are rotationally invariant [Mathieu et al., 2019], thus the same distances between the means µ of latent distributions can be learnt for any arbitrary rotation of the latent space. This suggests that VAEs disentangle by accident, since ground-truth factors naturally correspond with distances in the dataset. If these perceived distances were to change such that they do not correspond to the ground-truth distances, VAEs might not be able to learn meaningful representations. This is highlighted by the fact that VAEs are already known to perform poorly on real-world data [Gondal et al., 2019].

5 Adversarial Datasets

In the previous section, we highlighted the striking similarity between the ground-truth distances and the perceived distances between observations in synthetic ground-truth datasets. This suggests that disentanglement occurs because latent distances accidentally correspond to ground-truth distances, when the latent space is reorganised to minimise reconstruction errors and perceived distances from the data space are captured.

Consider the example of a single chess piece moving across a chess board; there are no smooth transitions between grid points, since the piece is only valid when placed in the middle of squares. We describe such a dataset as having constant perceived distance. This property is adversarial in nature as it is impossible for a VAE to order these observations using pixel-wise perceived distance. It is tempting to think that a harder case is if the perceived distances do not correspond to ground-truth distances; however, an (incorrect) ordering can then still be found. Existing datasets such as Cars3D already satisfy this incorrect ordering, which may explain the generally worse disentanglement performance compared to other datasets, see Figure 3.

Formally, we say that a dataset has constant overlap when the pairwise distances over factor traversals are all equal. Let i ∈ [F] be a factor and y(a) ∈ Y be ground-truth coordinate vector. Then, for all elements over the factor traversal ∀y(b) ∈ Y (a,i) / {y(a)}, the corresponding perceived distance is constant such that dpcv(x(a) , x(b) ) = Cf with Cf ∈ R and Cf > 0. Along factor traversals in such a dataset, no distinct ordering of elements can be found when a VAE tries to minimise the sampling error over the reconstruction loss. Going forward, we only consider the case where ∀f ∈ [F], Cf = C for some C > 0.

5.1 Example XYSquares Adversarial Dataset

Taking inspiration from the chess piece example, we design a synthetic adversarial dataset called XYSquares (See Figure 5)

that specifically targets VAEs that use a pixel-wise reconstruction loss such as MSE, resulting in constant perceived distances. The dataset consists of three 8 × 8 pixel squares in a world of size of 64 × 64. This leaves 8 grid positions along each axis without any pixel-wise overlap. The three squares are each assigned a colour according to R (1, 0, 0), G (0, 1, 0) and B (0, 0, 1) to avoid any channel-wise overlap. With 6 ground-truth factors (three squares moving along two axes), each with 8 possible values, this gives a total dataset size of 86 = 262144 observations. In the rightmost column of Figure 8, we validate that this leads to constant perceived distances between observation pairs in factor traversals.

5.2 Experimental Setup

We now investigate the performance of VAEs on our new dataset. In particular, we use the unsupervised β-VAE [Higgins et al., 2016] and the state-of-the-art weakly supervised Ada-GVAE [Locatello et al., 2020]. The β-VAE scales the VAE regularisation term with a coefficient β > 0, while the Ada-GVAE breaks symmetry and encourages shared latent variables between pairs of observations. This is achieved by averaging together latent distributions between observation pairs that are estimated to remain unchanged when the KL divergence is below some threshold. We note that if the weakly supervised Ada-GVAE performs poorly, then it is highly likely that another unsupervised method will also perform poorly.

We use the same Adam [Kingma and Ba, 2015] optimiser and convolutional neural architecture as Burgess et al.[2017]. To evaluate disentangled representations, we use the MIG [Chen et al., 2018] (Mutual Information Gap) and DCI Disentanglement [Eastwood and Williams, 2018] scores. MIG measures the mutual information between the highest and second highest latent units for each factor, and DCI Disentanglement measures how much each latent unit captures a ground-truth factor using a predictive model.

Finally, we perform an extensive hyper-parameter grid search for existing frameworks and datasets before running our own experiments. Hyperparameters include the learning rate, size of the latent dimension, training steps, batch size and β values. See our extended paper in Footnote 1 for further details on all experiments conducted throughout the remainder of the paper.

5.3 Example Adversarial Dataset Results

Figure 6 shows that the disentanglement performance over XYSquares is extremely poor compared to existing datasets, even with the state-of-the-art Ada-GVAE. We are concerned only with the maximum score obtained for each model and dataset, as the graph is plotted over the hyper-parameter sweeps. This validates our adversarial dataset hypothesis in Section 5. Not only is the disentanglement performance poor, but much smaller values for β are needed when tuning the regularisation loss. Example latent traversals from a VAE trained over the adversarial dataset are given in Figure 7, results are far from disentangled and do not correspond in any way to the ground-truth factors in Figure 5.

5.4 Example of Varying Levels of Overlap

We have examined the effect of training on existing datasets with significant amounts of overlap, as well as our own adversarial dataset with constant perceived distances according to pixel-wise losses. However, we have not investigated increasing levels of overlap in datasets, or rather reducing perceived between observations that are also close in ground-truth factor space. To do so, we modify XYSquares by decreasing the spacing between grid points while keeping the number of grid points constant along each factor, ensuring the dataset size remains fixed at 86 = 262144 observations.

The original adversarial dataset, with a spacing of 8, has a constant distance value of dpcv(x(a) , x(b) ) = C8. As the spacing s decreases from 8 → 1 over the datasets, the probability increases that any two observations re-sampled along a single factor traversal overlap p(dpcv(x(a) , x(b) ) < C8) and should thus be placed closer together in the latent space. More overlap leads to more unique distance values which in turn allows for easier ordering of data points. We visualise this concept using ground-truth and perceived distance matrices in Figure 8.

We verify our statements through the experimental results in Figure 9, where the β-VAE and Ada-GVAE are trained on these datasets. As the spacing decreases and overlap is introduced, the disentanglement performance improves, since it is easier for a VAE to introduce an ordering over representations. Even for the XYSquares dataset with 1 pixel of overlap between grid points, an ordering of elements along factor traversals can be induced. However, the probability of a VAE encountering these scenarios in the latent space due to random sampling is low, and thus it is still not always easy for the model to learn disentangled representations over such a dataset.

6 Example of Introducing Overlap

The previous section focused on increasing overlap by changing the underlying dataset; however, this still does not solve the case for the original XYSquares dataset with constant pixel-wise perceived distance. Throughout this paper, we have provided evidence that VAEs disentangle based on their reconstruction loss, which happens to align with ground-truth factors of variation in common benchmark datasets. This correspondence is not optimal for all tasks and we propose that this leads to the poor disentanglement performance in these settings. Our solution is to choose a loss function that modifies perceived distances such that they also correspond to ground-truth distances.

The new loss function we choose cannot be a pixel-wise approach, as this does not capture the distances due to the spatial nature of the XYSquares dataset. For the sake of simplicity in this example, we convert the existing pixel-wise loss function into a spatially aware loss function by introducing a differentiable augmentation to its inputs. An appropriate augmentation for our dataset is a channel-wise box blur. The problem, however, is that the decoder needs to be able to reconstruct the data, and so purely replacing the pixel-wise loss with the augmented loss may not succeed. Rather, in Equation (11), we append the augmented term to the existing loss and scale it by a constant α > 0.

6.1 Example Augmented Loss Experiments

We choose a channel-wise box blur for our loss function with a radius of 31, or a total kernel size of 63 × 63. We efficiently implement large filters using the Fast Fourier Transform. The size of the filter ensures that if two observations have active pixels on opposite sides of the images, then overlap will still be introduced between them. We set α = 632 , while this appears large, a box blur kernel is normalised so that the sum of all its values is 1. We accordingly update our perceived distance measure and evaluate the new distances over the XYSquares dataset for each factor in Figure 10 after training and tuning β-VAEs.

Finally, in Figure 11, we compare the performance of the spatially-aware loss function to the original pixel-wise loss. Our new loss significantly improves the disentanglement performance over the adversarial dataset. This is because it allows our models to capture perceived distances between observations that align with the ground-truth factors.

While our choice of loss may not be optimal for disentanglement of these specific x and y factors from our adversarial dataset, disentanglement results are impressive. This is important because it provides the intuition that changing the loss function changes perceived distances and affects the ability of VAE frameworks to learn disentangled representations. We leave learning or identification of optimal reconstruction losses for different datasets, to improve disentanglement, as future work.

7 Considerations for Disentanglement Research

We highlight the similarity between introducing overlap in Section 6 through the reconstruction loss function and varying levels of overlap in Section 5.4 through modifications to the construction of the dataset itself. Both methods aim to improve disentanglement by changing perceived distances to better correspond to the ground-truth factors, while keeping groundtruth factors fixed.

The problem is that ground-truth factors can indeed change, and this choice, while at the discretion of the researcher, is largely ignored in literature. For example, a researcher may choose RGB, HSV or categorical representations for colours, they may choose binary or continuous encodings for positions, or they may split or merge various factors together.

As our work shows, disentanglement is largely dependent on the chosen reconstruction loss and not special algorithmic choices. Obtaining improved disentanglement results under current VAE disentanglement frameworks will ultimately require supervision from the researcher to adjust perceived distances of the model to the task at hand. This contradicts the current notion that unsupervised and weakly supervised disentanglement methods can automatically uncover these human interpretable ground-truth factors [Higgins et al., 2016].

Ultimately, benchmarking against synthetic datasets with already subjective ground-truth factors will thus always remain problematic. There are infinitely many datasets with infinitely many choices as to what constitutes their ground-truth factors. Accurate disentanglement through future methods may need general world knowledge so that the methods can adapt to the task at hand.

8 Conclusion

In this paper, we demonstrated that there are fundamental characteristics of existing datasets that encourage VAEs to learn disentangled representations. Our work provides a theory for how VAEs perceive distances between pairs of observations in datasets. We used this theory to provide intuition by constructing an adversarial dataset for pixel-wise losses over-which state-of-the-art VAEs fail to learn disentangled representations. Finally, we re-enabled disentanglement over the example adversarial dataset by again adjusting perceived distances, instead through a change of the VAE reconstruction loss to capture the ground-truth factors of the dataset.

Our results highlight issues in current representation learning approaches. We find that the focus on regularisation for disentanglement is misplaced, rather, disentanglement is largely accidental, and careful choice of the reconstruction loss or data is needed to capture the ultimately subjective ground-truth factors. This is impractical in the real world, since perceived distances cannot be a prerequisite for true disentanglement. More advanced methods are therefore required that can uncover true meaning within the data.

Acknowledgements

Computations were performed using the High Performance Computing infrastructure provided by the Mathematical Sciences Support unit at the University of the Witwatersrand.