Lost in Latent Space: Disentangled Models and the Challenge of Combinatorial Generalisation

TL;DR

The paper investigates why highly disentangled representations often fail at combinatorial generalisation, proposing that simply isolating factors is insufficient without inverting the data-generating process. Using a semi-supervised composition task and datasets such as $dSprites$, $3DShapes$, and $MPI3D$, it shows that high disentanglement does not guarantee correct generalisation to unseen factor combinations, with encoder mappings to latent space drifting for test cases. Through latent-space visualisations and experiments with alternative decoders and encoder-only tasks, the authors demonstrate that encoder failures frequently accompany output-space generalisation errors, and that even advanced architectures (e.g., Spatial Broadcast Decoder, CascadeVAE, LieGroupVAE) do not fully overcome interactive-factor generalisation challenges. The work concludes that future models must learn how generative factors combine and invert the generative process, moving beyond disentanglement alone toward causal, invertible representations for robust 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 performance in out-of-training distribution settings when compared to entangled representations. Additionally, it is not clear if the reported failures are due to (a) encoders failing to map novel combinations to the proper regions of the latent space or (b) novel combinations being mapped correctly but the decoder/downstream process is unable to render the correct output for the unseen combinations. We investigate these alternatives by testing 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 correct regions of the latent space and (ii) when models succeed, it is either because the test conditions do not exclude enough examples, or because excluded generative factors determine independent parts of the output image. Based on these results, we argue that to generalise properly, models not only need to capture factors of variation, but also understand how to invert the generative process that was used to generate the data.

Figures (38)

• Figure 1: Possible problems for combinatorial generalisation. Left: visualisation of how a combinatorial generalisation condition can be defined for a dataset containing two generative factors. Black and red circles denote training and test examples, respectively. In the Middle and Right panels, the outline of each circle shows the target position of the representation in latent space, while a shaded circle shows the position at which the latent representation is projected by the encoder. Middle: first type of combinatorial generalisation error, encoder error. The encoder projects unseen inputs to different parts of the latent space than what was expected based on their generative factor values (shown as shaded circles falling outside their target outlines) Right: second type of combinatorial generalisation error, decoder error. Observed representations (shaded circles), are mapped to the correct position in latent space (circle outlines), but a decoder/downstream process mixes-up the black and red representations (blue dashed lines).
• Figure 2: (a) Image composition task. An example of the composition task. In this case, the [shape] of the output must match the transform image and the rest of the values ([position, orientation]) must match the original image. (b) Model architecture. The model uses the same encoder to represent both images. Then a transform takes these latent representation and combines them to produce a transformed representation, which is used to reconstruct the output image. Reproduced from Montero et al. with permission from the authors.
• Figure 3: Generalisation in latent spaceTop: Typical examples of reconstructions for training and test images for models that learn highly disentangled representations. Bottom: Visualisation of the latent space of models trained on three datasets -- dSprites (left), 3DShapes (middle) and MPI3D (right). Each panel shows contours from the joint distribution of two latent variables that best predict the corresponding generative factors. In all cases, the red distributions indicate test data and the black ones indicate training data. Note that the latent values in this figure are not necessarily the same as the values of generative factors excluded. This is because every model ends up with a different internal representation based on its initialisation and sequence of training trials. However, in each case, we observed that these internal representations were highly structured when the model showed a high degree of disentanglement.
• Figure 4: Latent space induced by SBD Visualisation of the latent space for models that use a Spatial Broadcast Decoder as described in Watters et al. Both visualisations correspond to a condition where combinations in the middle of the image are removed. (a) The result when this is done for the Circles dataset which uses only one shape and where the training set excludes all combinations [shape=circle, $0.35<$ posX $<0.65$, $0.35<$ posY $<0.65$] (b) The result for the a similar condition for the Simple dataset which uses two shapes and excludes all combinations [shape=square, $0.35<$ posX $<0.65$, $0.35<$ posY $<0.65$].
• Figure 5: Replacing decoder with ground-truth Visualisation of the latent space for models trained to predict the correct generative factor values for each image. The visualizations correspond to the same test conditions as the ones presented in Figure \ref{['fig:results-comp']}: (a) dSprites, where all combinations where [shape=square, posX $>0.5$] were excluded from the training set (b) 3DShapes, where all combinations where [shape=pill, object-hue $>0.5$] were excluded, and (c) MPI3D, where all combinations where [shape=cylinder, vertical-axis $>0.5$] were excluded (note that, in this case, [vertical-axis $>0.5$] corresponds to the corresponding latent variable value $<0$).