Lost in Latent Space: Examining failures of disentangled models at combinatorial generalisation
Abstract
Recent research has shown that generative models with highly disentangled representations fail to generalise to unseen combination of generative factor values.
These findings contradict earlier research which showed improved performancein out-of-training distribution settings when compared to entangled representations.
Additionally, it is not clear if the reported failures are due to (a) encodersfailing to map novel combinations to the proper regions of the latent space, or (b)novel combinations being mapped correctly but the decoder is unable to render thecorrect output for the unseen combinations.
We investigate these alternatives bytesting several models on a range of datasets and training settings.
We find that (i)when models fail, their encoders also fail to map unseen combinations to correctregions of the latent space and (ii) when models succeed, it is either because thetest conditions do not exclude enough examples, or because excluded cases involvecombinations of object properties with it’s shape.
We argue that to generaliseproperly, models not only need to capture factors of variation, but also understandhow to invert the process that causes the visual stimulus.
Introduction
Disentangled representations extract factors of variations from data, and learning them has beenthe focus of much research in recent years [10].
Several approaches have been proposed to inducedisentanglement, including latent space penalization [10, 3, 17], different training regimes [23, 22],architectures [36], data-driven inductive biases [21] and model selection methods [5].
These haveproduced more interpretable representations [12] that improve sample efficiency and learning fordownstream models [13, 33, 5].
Importantly, researchers have also hypothesized that disentangled representations could provide away of improving generalisation performance by enabling the discovery of causal variables in data[29] or by capturing its compositional structure [5].
This claim is especially interesting as it allowsmachine learning systems to emulate a key property of human intelligence – the ability to generaliseto unseen combinations of known elements.
For example, if a model has learned to generate redtriangles and blue squares, then the model should also be able to correctly generate blue trianglesand red squares.
This property, which we refer to as combinatorial generalisation , gives humans theability to make “infinite use of finite means” [34, 4, 31, 25] and has been termed “a top priority for AIto achieve human-like abilities” [1].
Indeed, several authors have reported that unsupervised modelsthat are better at disentangling generative factors are also better at some forms of combinatorialgeneralisation [10, 13, 36]
However, two recent studies have reported evidence that contradicts this hypothesis.
Montero et al.[26] considered datasets where inputs varied along several dimensions (generative factors) and dividedgeneralisation conditions into different types.
They found that VAEs with highly disentangled latentrepresentations succeeded at the easiest generalisation conditions where only one combination of allgenerative factors was excluded (a condition they termed recombination-to-element ), but failed atmore challenging conditions, where all combinations of a subset of generative factors were excluded(termed recombination-to-range by Montero et al.
In another second study, Schott et al.
[30]tested 17 learning approaches, expanding the analysis to other paradigms, architectures and taskstructures.
They observed that models showed moderate success at generalising on some artificialdatasets, such as 3DShapes.
However, their ability to generalise dropped significantly on morerealistic datasets and in more challenging generalisation conditions.
These contradictory results could be down to a number of differences between studies that reportsuccesses and those that report failures.
For example, there were qualitative differences in the datasetsthat were tested and how the test conditions were set up.
This is indeed what Montero et al.
[26]conclude, writing that “it is not clear what exactly was excluded while training [previous models]”,suggesting that previously observed successes may have been on the simplest generalisation conditionwhere only very few combinations were left out.
But this still leaves open the question of why modelsthat achieve high degree of disentanglement nevertheless fail at combinatorial generalisation undermore challenging conditions.
We see two possible reasons why this could happen.
One possibility is that models correctly infer thevalues of latent variables (Fig. 1, left) for novel combinations of generative factors, but the decoderfails to map these unseen latent values to the (output) image space – decoder error (Fig. 1, right).
Some evidence supporting this hypothesis was observed by Watters et al.
[36], who looked at notonly the output image reconstructions, but also at latent representations in a simple reconstructiontask.
Watters et al.
[36] observed that, in their task, models that have more disentangled latentrepresentations do indeed map unseen combinations of generative factors to the correct values inthe latent space.
But they only tested a very simple dataset, involving only two generative factors.
Montero et al.
[26] showed that models can sometimes succeed to generalise in simple settings wherethey can solve the combinatorial generalisation problem through interpolation, but struggle at morechallenging settings where larger number of combinations are excluded from the training set.
Another possibility is that current models fail at harder forms of combinatorial generalisation (suchas recombination-to-range) due to an encoding failure – encoder error .
That is, the encoder fails tomap these harder unseen combinations of generative factors to the correct values of variables in thelatent space (Fig. 1, middle).
If this were the case, then it reflects a more fundamental limitation of current models, showing that even though models can map observed combinations to correct valuesof latent representations, they do not understand how generative factors are combined.
So whennovel combinations are presented, models cannot infer the values of generative factors that led to theobserved data.
In this work we explore these issues by examining the latent space of generative models.
Extendingthe findings of Montero et al.
[26] and Watters et al.
[36], we show that failure to generalise in theoutput space is accompanied by failure to generalise in latent space across several datasets.
We showthat these results hold for a broad range of architectures (the standard deconvolution decoder as wellas the spatial broadcast decoder ), loss functions ( β -VAE and WAE ) and task settings.
Finally, bylooking at the latent space, we discover that, in addition to the difficulty of the generalisation task, acrucial condition for failure of combinatorial generalisation is the way in which generative factorsare combined.
These results not only challenge research that argues that disentanglement leads tobetter generalisation [13, 36], but also shows that the failure to generalise is a lot more entrenched incurrent models than suspected by Montero et al.
[26] and Schott et al.
Testing generalisation in latent representations
In this section we first describe a carefully designed semi-supervised task that enables us to achievelatent representations that are highly disentangled.
This is crucial because standard disentanglementlearning objectives, which rely on minimizing total correlation [10, 3, 17], struggle to reliably achievea high degree of disentanglement.
This is problematic because it is difficult to test whether entangledmodels are projecting test data to the expected position in the latent space (see Figure 1).
Secondly,we describe our choice of corresponding model architectures as well as loss functions.
For our resultsto be broadly applicable to existing research, we chose standard encoder and decoder architecturesas well as a more sophisticated approach designed to induce a high degree of disentanglement [36].
Lastly, we describe some challenging test conditions for combinatorial generalisation on well knowndatasets.
2.1 Experimental Setup
2.1.1 Task
One of the main barriers to achieving disentanglement in standard unsupervised training is thenon-identifiability of the models when using iid data [15, 21].
In short, this is because there areinfinitely many linear combinations of the underlying generative factors that produce a valid basis onwhich they can be represented.
To get around this problem we used the composition task developed by Montero et al.
[26] (see Appendix A.1 for further discussion). This task takes two images and a query vector as inputs (see Figure 2(a)).
The goal of the task is to output an image that combines the two images based on anaction in the query vector: input =
[Equation]
q output = x out where x og and x trans are the two input images, x out is the output image, and q is the query vector.
Following Montero et al.
, we used actions that involve replacing one of the properties (generativefactors) of x og with a property of x trans , based on the value of q .
This vector was a one-hot encodingof the generative factor that was required to be changed.
By combining a semi-supervised component (the query vector plus image reconstruction) with sparsetransitions (only one generative factor needed to be modified at a time), the composition task provideda strong inductive bias towards high levels of disentanglement, as has been shown in previous workthat exploits these approaches [23, 22, 21].
2.1.2 Models
The model architecture used to solve the composition task is shown in Figure 2(b).
It involves twoencoders that map each image to a latent space.
The latent vectors are then “composed” (see below),based on the query (action) vector.
This composed vector is mapped from the latent space into the(output) image space using a decoder.
Encoders and decoders for generative models match the onesused in past work [10, 17] for both dSprites and 3DShapes datasets.
For MPI3D we added another2 more layers, which we found was necessary for models to learn the task.
For the Circles dataset weused the encoder defined in Watters et al, along with the corresponding Spatial Broadcast Decoderarchitecture.
Latent representations for the generative models are standard diagonal Gaussians withinput dependent means and variances, to which we applied both the Variational and Wasserstein loss[19, 32].
We can define the composition operation that combines both latent representations as:
[Equation]
where z out , z og and z trans are the latent representations corresponding to the output, original,and transform images and ⊙ is the element-wise product between vectors.
The variable c is theinterpolation vector of coefficients.
We defined two ways of computing c : (i) a learned interpolationfunction computed as σ ( W · [ z og ; z trans ; q ] + b ) with parameters W and b where [ . ; .
representsthe concatenation of vectors; and (ii) a fixed interpolation function c
= padded ( q ) which just pads qwith zeros to match the vector length of the latent representation (see Appendix A.2 for more details).
To facilitate our goal of testing combinatorial generalisation, we augmented the target output byrequiring the model to reconstruct the input images as well.
Thus the trained encoder and decoderform a valid standard autoencoder. We can then probe the image reconstructions and the latentrepresentations for unseen combinations to check if the models generalise.
An illustration of thismodel can be found in Figure 2.
2.1.3 Datasets
We tested our models on three standard datasets: dSprites (a sythesised dataset of 2D shapesprocedurally generated from 5 independent generative factors [24]), 3DShapes (another synthesiseddataset consisting of 3D shapes procedurally generated from 6 ground truth generative factors [2])and MPI3D (a collection of four different datasets consisting of synthetic as well as real-world imagesprocedurally generated from 7 generative factors [8]).
These are the most popular datasets used inthe disentanglement literature that also provide ground truth generative factor values.
We testedcombinatorial generalisation by systematically excluding combinations of values of generative factorsfrom each dataset.
Here we focus on the conditions described as recombination-to-range by Monteroet al..
These are the most interesting conditions where, according to Montero et al., one would expecta model that learns disentangled representations to succeed at combinatorial generalisation, but testedmodels typically fail.
Consider a dataset with, say, four generative factors [ g 1 , g 2 , g 3 , g 4 ] whereall g i ∈ [0 , 1] .
The recombination-to-range condition creates a training/test split where all exampleswith combinations of a subset of generative factors are excluded from the training set and added tothe test set.
Thus, an example of a dataset that tests recombination-to-range may consist of a trainingset where all combinations where [ g 1 > 0 .
5 , g 2 > 0 .
5 ] have been excluded from the training set andadded to the test set.
Note that the model trained on such a datasets would come across a number of train test train test train test o r
i g i n
a l r e c o n s examples where [ g 1 > 0 . 5 ] and also examples where [ g 2 > 0 . 5 ] , but never be trained on an examplewhere both these conditions are true simultaneously.
This method was used to create training / testsets for each of the datsets in the following manner:
• dSprites : All images such that [shape=square, posX > 0 . 5 ] were excluded from the training set.
Squares never appear on the right side of the image, but do appear on the left and other shapes (hearts, ellipsis) are observed on the right as well.
• 3DShapes : All images such that [shape=pill, object-hue= > 0 . 5 ] were excluded from the training set.
Thus, pills colored as any of the colors in the second half of the HSV spectrum did not appear in the training set.
These colors (shades of blue, purple, etc) were observed on the other
shapes, and the pill was observed with other colors such as red and orange.
• MPI3D : All images such that [shape=cylinder, vertical axis > 0 . 5 ] were excluded from the
training set.
We also excluded all images where [shape=cone] or [shape=hexagonal] as theseshapes are very
similar to the pyramid and cylinder, respectively and make it hard to accessreconstruction accuracy.
2.1.4 Measuring Disentanglement
We used the DCI metric from Eastwood & Williams [6] to measure the degree of disentanglement(see Appendix A.5).
This metric provides a set of regression weights between ground truth factorsand latent variables.
We selected the latent variable corresponding to each ground truth factorusing a matching algorithm based on these weights.
Crucially, for highly disentangled models anycombination of generative factors will have a corresponding combination of latent variables of thesame cardinality.
This allows us to create visualisations of the latent representations analogous to theones in Figure 1.
2.2 Results
We observed that most models trained on the composition task managed to achieve a reasonablyhigh degree of disentanglement, and successfully reconstructed images in the training data (see
Appendix B.1 and B.2 for quantitative results).
Here we present results of typical models that achieveda very high degree of disentanglement (with alignment scores typically > 0 .
95 , see Appendix B.1 fordetails).
Figure 3 shows reconstructions as well latent space representations of three typical modelstrained on the dSprites , 3DShapes and MPI3D datasets, respectively.
Replicating results of Montero et al.
[26], we observed that models showed poor generalisationto unseen combinations for all three datasets (see reconstruction of test images in Figure 3).
For
[Figure]
dSprites , the models typically produced images with correct location, but made an error in thetarget shape (e.g., replacing the square with an ellipse).
Similarly, with 3DShapes and MPI3D , modelstypically reconstructed shapes with the correct colour hue ( 3DShapes ) or vertical location ( MPI3D ),but made an error in the shape of the target object (e.g., replacing the pill ( 3DShapes ) or cylinder( MPI3D ) with a sphere( 3DShapes ) or pyramid ( MPI3D ).
Next, we looked at the latent representations of each of these cases by selecting a latent variable thatshowed highest correlation with each generative factor as described in Section 2.1.4.
Figure 3 plotsthe distributions of the value of each latent variable for every combination of values of the relevantgenerative factors.
Visualising these probability distributions for the combination of generative factorsseen during training (in black) confirmed that the models learned highly disentangled representations(note the low variance of the probability distributions for trained combinations, especially fordSprites and 3DShapes ).
Crucially, we observed that the encoder failed to map the generative factors for the test combinations to the correct location in the latent space (red distributions).
This was particularly true for purple]] ) are shifted to overlap the latent representations seen during training.
The shifts in the mean of distributions were less acute for MPI3D , but we nevertheless observed an increase in the variability of test distributions and a consistent shift in means for values of generative factors thatwere far from the trained values.
In summary, we made three key observations in these experiments.
First, models replicated thenegative results reported previously: i.e., a degradation in reconstruction performance in the challenging combinatorial generalisation conditions.
Second, by visualising the latent space, we observedthat failures of generalisation coincided with poor latent representations, showing that failures ofgeneralisation are not entirely due to decoder errors.
Third, the failures in latent representations aswell as output reconstructions showed that different generative factors have important qualitativedifferences, with models making larger errors in generating the correct shape for unseen combinationsthan in reproducing the position, scale or color (see Appendix B.2).
2.3 Exploring different decoder architectures
In the experiments above, we used the deconvolution network as the decoder, which is the standarddecoder used in VAEs [10].
However, some studies have shown that replacing the deconvolutionnetwork with a different architecture helps the model learn representations that are more disentangled.
One such architecture is the Spatial Broadcast Decoder (SBD from hereon) developed by Watterset al.
[36], who showed that such a decoder not only helped in learning representations that werehighly disentangled, but also helped in solving the problem of combinatorial generalisation.
correspond conditions the
In our next set of experiments, we tested whether using the SBD succeeds at combinatorial generalisation in more challenging settings, akin to the ones we tested above.
We first replicated the resultsfor both conditions tested by Watters et al.
[36], using our composition task on the Circles datasetand replacing the deconvolution network with the SBD.
Figure 4 (panel A) shows the results for thefirst condition, where the training set excluded images with circles presented in the middle of thecanvas (see Appendix B.2 for results on the second condition).
As we can see from this figure, thelatent space learned by the model was indeed highly structured and, crucially, the model mapped thelatent values for the left out combinations (red crosses) to the correct region of the latent space.
Next, we wanted to check how these results scaled when we made the generalisation test morechallenging, excluding a combination of values for a range of variables (the recombination-to-rangecondition developed by Montero et al.
We can do this with the Circles dataset by adding a thirdgenerative factor and the simplest way of doing this is by letting the shape take one of several values,instead of being a circle for all examples.
We tested this by constructing a new dataset called Simple.
Like the Circles dataset, this datasetconsisted of images that had a shape located at various (x and y) positions on the canvas.
Unlike the
Circles dataset, this shape could be either a triangle or a square (we chose these shapes as they aredifferent enough for models to not confuse them by accident).
We then tested the two conditionssimilar to the conditions tested by Watters et al.
[36], excluding all combinations where one of theshapes was presented in the middle (first condition) or the top-right (second condition) of the canvas.
The results for the first condition are shown in panel B of Figure 4 (see Appendix for results ofthe second condition).
Like the results for the first set of experiments above, we have plotted theprobability distribution of combinations of values of the the latent variables (here, shape and posX)for all values of the third latent variable (here, posY).
Similar to our earlier results, we observed thateven the model using the SBD failed to map the unseen (test) combinations to the correct positionin the latent space (compare the mean and variability of the red distributions compared to the blackones).
Correspondingly, we also observed that models failed to correctly reconstruct output imagesfor the test conditions, even when they had no problem reconstructing these images for combinationof generative factors seen during training (See Appendix B.2).
Replacing decoder with ground-truth latent representations
We have shown that errors in generalisation occur in latent space as well in the reconstructions forgenerative models.
To conclude this analysis we show that the problem does not lie in the generativenature of the task by testing the encoders on latent prediction task, where the model has to predictcorrect latent representations based on an input image.
We used the same datasets and generalisationconditions as before.
3.1 Experimental setup
We trained models on a prediction task where the model was given a single image as the input andthe target was to output the ground truth generative factor values.
This ensures that the resultinglatent representation will be completely disentangled, as the model must output each of the factor values separately.
We used the mean squared error of the output and target vectors as a learningsignal.
All other parameters remain same as the composition task.
We used the same architectures forthe feed-forward models as the encoders in the corresponding datasets above.
We evaluated the models using the R 2 metric as used in Schott et al.
[30] to check whether themodels had solved the task.
We also use the same method to visualise the results as described above,plotting the joint probability distributions for the combination of generative factors along which themodel was tested for generalisation.
In this case, it is straightforward to do so as the relevant outputdimensions shared the same index as the corresponding generative factor.
3.2 Results
The results for the three dataset are shown in Figure 5.
We observed a very similar pattern for thistask as the semi-supervised task above.
For all three datasets, the probability distributions of latentvalues for the left out combinations (in red) showed a much higher variance and means that wereshifted towards combinations of latent values that had been experienced during training.
Thus, evenwhen the encoders were trained to recognise perfectly disentangled generative factors, they failed togeneralise to combinations that were not experienced during training.
[Figure]
Explaining contradictory findings
How can we reconcile the failures of generalisation observed above with with past studies thatshowed successful combinatorial generalisation on some datasets and conditions?
One response isthat the generalisation conditions we test are more challenging than the ones tested in previous studiesbecause they exclude more data from the training set.
However, it is also possible that the success ofcombinatorial generalisation is not only determined by how many combinations are excluded fromthe test condition, but also by which combination of generative factors are tested.
All failures observed above were for cases where the generative factor ‘shape’ is combined withanother generative factor (position / color).
It is possible that this was because shape combines withother factors in a qualitatively different manner than, say, the floor hue combines with the wall hue in3DShapes .
In order to completely disentangle shape from position / color / orientation, the modelmust come up with a representation of shape that is invariant to a change in these generative factors.
And it must do so by learning from training examples where shape and orientation interact – that is,the observed image is obtained via a composition of the two generative factors affecting the samepart of the image.
But not all combinatorial generalisation problems face this problem of working out how generativefactors combine to determine the same part of the observed image.
In some cases, an image may bedetermined by a combination of generative factors, but each factor governs a different part of theimage, such as the generative factors [floor-hue, wall-hue] in 3DShapes .
In such cases, learning a representation of one factor that remains invariant to the other factor is trivial and the model maysucceed at combinatorial generalisation, even though the same amount of data is excluded from thetraining set.
To test this hypothesis, we carried out a set of experiments where the model had to solvethe problem of combinatorial generalisation, but the combinations that were excluded did not interactin the training examples.
We can do this for both the 3DShapes dataset as well as the MPI3D dataset,as both datasets involve images where the canvas contains several disjoint elements.
So, we repeatedthe semi-supervised learning experiments above, where models had to learn the composition task,but instead of excluding combinations where the generative factors interacted with each other, weexcluded the following combinations from the training set:
• 3DShapes : Exclude all combinations such that [floor-hue < 0 .
25 , wall hue > 0 .
75 ] : none of
the training images had the combination of floors with a “warm” hue and walls with a “cold” hue.
• MPI3D : Exclude all combinations such that [shape={cylinder,sphere}, backgroundcolor=salmon] .
Note, even though shape combinations are excluded, the combination ofgenerative factors excluded do not interact (i.e. determine different parts of the image).
4.1 Results
Unsurprisingly, we again observed that models learned highly disentangled representations forthis task.
More interestingly, we also observed that models now succeeded at the combinatorialgeneralisation task.
Figure 6 shows some typical reconstructions for training as well as test images.
These reconstructions successfully reproduce the unseen combination of floor hues and wall hues forthe 3DShapes dataset and the combination of shape and background arc color for the MPI3D dataset.
This figure also shows the joint probability distribution for trained as well as novel combination oflatent values.
In contrast to Figure 3, we now see that the means of test latent distributions (in red)fall in the expected location and a show a much smaller variance.
Discussion
Our world is inherently compositional – it can be decomposed into simpler parts and relationshipsbetween these parts.
The idea behind learning disentangled representations is to recover this underlying compositional structure of the world from perceptual inputs.
Unsupervised learning models,such as VAEs, aim to do this by separating out the factors that remain invariant under transformationsof other factors [11].
It is therefore tempting to conclude that models that manage to show a highdegree of disentanglement on the training set also capture the compositional structure of the world.
Ifthey did, they should be able to generalise to settings that present novel combination of the factors ofvariation.
However, in our experiments, we observed that models that learned highly disentangledrepresentations nevertheless failed at combinatorial generalisation, not only in the reconstructionspace, but also in the latent space.
These failures reproduced over several different datasets, differenttypes of decoders, different loss functions and under a variety of different task settings.
Our interpretation of these results is that models manage to achieve a high degree of disentanglementby discovering factors that remain invariant over training examples and simply associating perceptualinputs with these factors.
However, to capture the compositional structure of the world, models mustadditionally understand how factors interact to cause the perceptual input – that is, develop a goodcausal model of the world.
We think the failure of models to form such a causal model is the reasonbehind why models succeeded at problems of combinatorial generalisation that do not involve aninteraction between the left-out factors (Figure 6), but consistently failed at problems where thesefactors interacted (Figure 3).
When generative factors do not interact, the model does not need tolearn how the factors combine to determine the same part of the output space.
Instead, it can simplymap the value of each generative factor to a different location.
It appears to solve combinatorialgeneralisation, but it does not understand how generative factors combine.
In summary, our work shows that the problem of combinatorial generalisation remains unsolved inboth latent space and reconstruction space, even for highly disentangled models.
By highlighting thislimitation, we hope to inspire more work exploring new approaches that emulate this key capacity ofhuman cognition, thus endowing models with a better understanding of the compositional structure Checklist For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] Instead, we report probability distributions of relevant
output variables.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[No] We provide links to the Github
repositories containing the datasets and their respective licenses.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]