The role of Disentanglement in Generalisation

1 INTRODUCTION . Generalisation to unseen data has been a key challenge for neural networks since the early days of connectionism , with considerable debate about whether these models can emulate the kinds of behaviours that are present in humans ( McClelland et al. , 1986 ; Fodor & Pylyshyn , 1988 ; Smolensky , 1987 ; 1988 ; Fodor & McLaughlin , 1990 ) . While the modern successes of Deep Learning do indeed point to impressive gains in this regard , human level generalisation still remains elusive ( Lake & Baroni , 2018 ; Marcus , 2018 ) . One explanation for this is that humans encode stimuli in a compositional manner , with a small set of independent and more primitive features ( e.g. , separate representations of size , position , line orientation , etc . ) being used to build more complex representation ( e.g. , a square of a given size and position ) . The meaning of the more complex representation comes from the meaning of it ’ s parts . Critically , compositional representations afford the ability to recombine primitives in novel ways : if a person has learnt to recognize squares and circles in context where all squares are blue and all circles are red , they can nevertheless also recognise red squares , even though they have never seen these in the training data . This ability to perform combinatorial generalisation based on compositional representations is thought to be a hallmark of human level intelligence ( Fodor & Pylyshyn , 1988 ) ( See McClelland et al . ( 1986 ) for a diverging opinion ) . Recently it has been proposed that generalisation in neural networks can be improved by extracting disentangled representations ( Higgins et al. , 2017 ) from data using ( variational ) generative models ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . In this view , disentangled representations capture the compositional structure of the world ( Higgins et al. , 2018a ; Duan et al. , 2020 ) , separating the generative factors present in the stimuli into separate components of the internal representation ( Higgins et al. , 2017 ; Burgess et al. , 2018 ) . It has been argued that these representations allow downstream models to perform better due to the structured nature of the representations ( Higgins et al. , 2017 ; 2018b ) and to share information across related tasks ( Bengio et al. , 2014 ) . Here we are interested in the question of whether networks can support combinatorial generalisation and extrapolation by exploiting these disentangled representations . In this study we systematically tested whether and how disentangled representations support three forms of generalisation : two forms of combinatorial generalisation that varied in difficulty as well as extrapolation , as detailed below . We explored this issue by assessing how well models could render images when we varied ( 1 ) the image datasets ( dSprites and 3DShape ) , ( 2 ) the models used to reconstruct these images ( β-VAEs and FactorVAEs with different disentanglement pressures , and decoder models in which we dropped the encoders and directly input perfectly disentangled latents ) , and ( 3 ) the tasks that varied in their combinatorial requirements ( image reconstruction vs. image transformation ) . Across all conditions we found that models only supported the simplest versions of combinatorial generalisation and the degree of disentanglement had no impact on the degree of generalisation . These findings suggest that models with entangled and disentangled representations are both generalising on the basis of overall similarity of the trained and test images ( interpolation ) , and that combinatorial generalisation requires more than learning disentangled representations . 1.1 PREVIOUS WORK . Recent work on learning disentangled representations in unsupervised generative models has indeed shown some promise in improving the performance of downstream tasks ( Higgins et al. , 2018b ; van Steenkiste et al. , 2019 ) but this benefit is mainly related to sample efficiency rather than generalisation . Indeed , we are only aware of two studies that have considered the importance of learned disentanglement for combinatorial generalisation and they have used different network architectures and have reached opposite conclusions . Bowers et al . ( 2016 ) showed that a recurrent model of shortterm memory tested on lists of words that required some degree of combinatorial generalisation ( recalling a sequence of words when one or more of words at test were novel ) only succeeded when it had learned highly selective ( disentangled ) representations ( `` grandmother cell '' units for letters ) . By contrast , Chaabouni et al . ( 2020 ) found that models with disentangled representations do not confer significant improvements in generalisation over entangled ones in a language modeling setting , with both entangled and disentangled representations supporting combinatorial generalisation as long as the training set was rich enough . At the same time , they found that languages generated through compositional representations were easier to learn , suggesting this as a pressure to learn disentangled representations . A number of recent papers have reported that VAEs can support some degree of combinatorial generalisation , but there is no clear understanding of whether and how disentangled representations played any role in supporting this performance . Esmaeili et al . ( 2019 ) showed that a model trained on the MNIST dataset could reconstruct images even when some particular combination of factors were removed during training , such as a thick number 7 or a narrow 0 . The authors also showed that the model had learned disentangled representations and concluded that the disentangled representations played a role in the successful performance . However , the authors did not vary the degree of disentanglement in their models and , accordingly , it is possible that a VAE that learned entangled representations would do just as well . Similarly , Higgins et al . ( 2018c ) have highlighted how VAEs that learn disentangled representations can support some forms of combinatorial generalisation when generating images from text . For example , their model could render a room with white walls , pink floor and blue ceiling even though it was never shown that combination in the training set . This is an impressive form of combinatorial generalisation but , as we show below , truly compositional representations should be able to support several other forms of combinatorial generalisations that were not tested in this study . Moreover , it is not clear what role disentanglement played in this successful instance of generalisation . Finally , Zhao et al . ( 2018 ) assessed VAE performance on a range of combinatorial generalisation tasks that varied in difficulty , and found that the model performed well in the simplest settings but struggled in more difficult ones . But again , they did not consider whether learning disentangled representations was relevant to generalisation performance . Another work that has significant relation to ours is Locatello et al . ( 2019 ) , who examine how hard it is to learn disentangled representations and their relation to sampling efficiency for downstream tasks . We are interested in a related , but different question : even if a model learns a disentangled representation in an intermediate layer , does this enable models to achieve combinatorial generali- sation ? So while Locatello et al . ( 2019 ) train their models on complete datasets to investigate the degree of disentanglement and sampling efficiency , we systematically exclude generative factors from training in order to test for combinatorial generalisation ( see Methods and Results ) . 2 METHODS AND RESULTS . We assessed combinatorial generalisation on two different datasets . The dSprites image dataset ( Matthey et al. , 2017 ) contains 2D images in black and white that vary along five generative factors : shape , scale , orientation , position-x and position-y and focuses on manipulations of single objects . The 3D Shapes dataset ( Burgess & Kim , 2018 ) contains 3D images in colour that vary along six generative factors : floor-hue , wall-hue , object-hue , object-shape , object-scale , object-orientation . In contrast to dSprites , the images are more realistic , which has been shown to aid reconstruction performance ( Locatello et al. , 2019 ) . To test combinatorial generalisation , we systematically excluded some combinations of these generative factors from the training data and tested reconstruction on these unseen values . Test cases can be divided into three broad categories based on the number of combinations excluded from training . • Recombination-to-Element ( red squares in Figure 1 ) : The model has never been trained on one combination of all of the generative factors . In dSprites , an example of this case would be excluding the combination : [ shape=ellipse , scale=1 , orientation < 120◦ , position-x > 0.5 , positiony > 0.5 ] from the training set – i.e . the model has never seen a large ellipse at < 120◦ in the bottom-right corner , though it has seen all other combinations . • Recombination-to-Range ( green squares in Figure 1 ) : The model has never been trained on all combinations of some of the factors ( i.e . a subset of generative factors ) . For example , in the 3D Shapes dataset , all combinations with [ object-hue=1 , shape=sphere ] have been left out of the training set – i.e . none of the training images contain a blue sphere . This condition is more complex than Recombination-to-Element as an entire range of combinations [ floor-hue=0 . . . 1 , wall-hue=0 . . . 1 , ob ject-hue=1 , shape=sphere , scale=0 . . . 1 , orientation=0 . . . 1 ] have been left out ( here bold text indicates the range of values excluded ) . When the number of generative factors is larger than three , “ Recombination-to-Range ” is , in fact , a set of conditions that vary in difficulty , depending upon how many generative factors have been excluded . Another example would be excluding all combinations where [ floor-hue=1 , wall-hue=1 , object-hue=1 , shape=1 , scale=1 ] . Here a smaller range of combinations [ floor-hue=1 , wall-hue=1 , object-hue=1 , shape=1 , scale=1 , orientation=0 . . . 1 ] have been excluded . • Extrapolation ( blue squares in Figure 1 ) : This is the most challenging form of generalisation where models are tested on values of generative factors that are beyond the range of values observed in the training dataset . For example , in the dSprites dataset , all combinations where [ posi tion-x > 0.5 ] have never been seen . Each of these conditions is interesting for different reasons . A model that learns compositional representations should be able to combine observed values of shape ( ellipses ) , translation ( bottom-right ) and rotation ( 0◦ to 120◦ ) to generalise to all unseen combination of factors . The simplest case is the recombination-to-element condition in which all combinations but one have been trained , but a model that learns entangled representations might also succeed based on its training on highly similar patterns ( generalisation by interpolation ) . A more challenging case is recombination-to-range condition given that more combinations have been excluded , making generalisation by similarity ( interpolation ) more difficult . The final condition is not a form of combinatorial generalisation as the model can not combine observed values of generative factors to render images . Indeed compositional representations may be inadequate for this form of generalisation . 2.1 IMAGE RECONSTRUCTION WITH DSPRITES DATASET . In the dSprites dataset , for testing the Recombination-to-element case , we split each range of values of a generative factor into three bins , so that we had 3 × 3 × 3 × 3 × 3 such combinations of bins for all five generative factors . We then remove one of these 243 combinations during training , namely those that satisfied [ shape=ellipsis , position-x > = 0.6 , , position-y > = 0.6 , 120◦ < = rotation < = 240◦ , scale < 0.6 ] . In other words ellipses in the bottom right corner with those given rotations , which is a relatively small number of combinations that are all very similar to each other . For the Recombination-to-range case , we tested three different variants . First , we excluded all combinations where [ shape=square , position-x > 0.5 ] . The model sees other shapes at those positions during training and it sees squares on the left-hand side of the screen . Thus the models experiences both generative factor values independently and has to recombine them to produce a novel image at test time . In the second case , we excluded all combinations where [ shape=square , scale > 0.5 ] . In the third case , we excluded all combinations where [ shape=square , rotation > 90◦ ] . We observed very similar results for all three cases and below we report the results for the first variant . Finally , for the Extrapolation case , we excluded all combinations of generative factors where [ po sition-x > x ] . We chose a set of different values for x : x ∈ 0.16 , 0.25 , 0.50 , 0.75 , where x is normalised in the range [ 0 , 1 ] ( results shown in Figure 2 for x = 0.50 ) . At test time the model needed to reconstruct images where translation along the x-axis , x , was greater than the cutoff value . We tested three classes of models on all three types of generalisation : standard Variational Autoencoder ( VAEs , Kingma & Welling ( 2013 ) ; Rezende et al . ( 2014 ) ) , β-VAE ( Higgins et al. , 2017 ; Burgess et al. , 2018 ) with β = 8 and β = 12 , FactorVAE ( Kim & Mnih , 2019 ) with γ = 20 , γ = 50 and γ = 100 . The architectures are the ones found in Higgins et al . ( 2017 ) , Burgess et al . ( 2018 ) and Kim & Mnih ( 2019 ) ( Details in the Appendix ) . We used a batch size of 64 and a learning rate of 5e− 4 for the Adam optimizer ( Kingma & Ba , 2017 ) . In each case , we simulated three seeds and we report results for runs where we obtained largest disentanglement . As shown by Locatello et al . ( 2019 ) , none of the models trained end-to-end in an unsupervised manner produce perfectly disentangled representations . Since we were interested in studying the effect of disentanglement on generalisation , we compared our results with a model where we removed the encoder and directly gave disentangled latents as inputs to the decoder . We call this model the ground-truth decoder ( GT Decoder from here on ) . This decoder uses the same MLP architecture as the one used in Higgins et al . ( 2017 ) . We tested deeper decoders with convolutions and batch norm as well , but found no benefit or a decrease in performance . We measured the level of disentanglement using the framework introduced in Eastwood & Williams ( 2018 ) . The procedure consists of using the latent representations generated for each image to predict the true generative factors using a regression model ( in our case , Lasso regression ; see Ap- to-Element condition where the models did not see [ shape = ellipse , scale = 1 , orientation < 120◦ , po sition-x > 0.5 , position-y > 0.5 ] , Middle ) Recombination-to-Range condition where models did not see [ shape = square , position-x > 0.5 ] , Right ) Extrapolation condition where models did not see [ posi tion-x > 0.5 ] ( b ) Visualisation of disentanglement . In each panel , columns show latent variables and rows show the generative factors . The size of the square represents the relative importance of the latent variable for predicting the generative factor . Sparse matrices indicate higher disentanglement ( Eastwood & Williams , 2018 ) . Each disentanglement matrix corresponds to the model on that row in ( a ) in the Reconstruction-to-Range condition . The visualisation of the entire set of models and all conditions is shown in Appendix B pendix A ) . The level of disentanglement is quantified by their ‘ Overall disentanglement metric ’ , which we call D-score here . Figure 2 shows examples of model reconstructions for each of the conditions which help assess the reconstruction success qualitatively ( more examples are shown in Appendix C ) . A more quantitative assessment of the models can be made by examining the negative-log-likelihood of reconstructions for different conditions , plotted in Figure 3 . The amount of disentanglement achieved by the models trained end-to-end varied over a broad range and was a function of model architecture and the hyperparameter ( β and γ ) values . In general , reconstruction accuracy was better for smaller values of β both during training and testing . This has been observed before and is a known issue encountered when increasing the value of β parameter ( Hoffman & Johnson , 2016 ) . We found that models were able to perform the Recombination-to-Element generalisation but failed in the Recombination-toRange and Extrapolation cases . In these cases , models either showed really poor reconstruction of the critical element or substituted one of the excluded combination with a combination that had been observed during training ( see reconstructions for test cases in Figure 2 ( a ) ) . Moreover , the amount of generalisation did not depend on the degree of disentanglement . Indeed , the GT Decoder using perfectly disentangled representations was no better than the end-to-end models . Even though this model achieved a lower NLL score , examining the image reconstructions showed that it failed to reconstruct the essential combination excluded from the training data ( see Appendix B ) . The Recombination-to-Range condition shows another interesting qualitative difference between the entangled and disentangled models . All models failed to generalise , but in different ways . Entangled models tended to put a blob in the correct location , which allows them to minimise loss in pixel space over a large set of test examples . In contrast , the models with higher level of disentanglement fell back to the most similar shape ( in pixel space ) that they had seen at that location . Finally , the Recombination-to-Element condition was solved by all the models , regardless of disentanglement score . In fact , the entangled models tended to achieve better reconstructions as evidenced by the disentangled models with β=12 which had a hard time reconstructing ellipses at small scales and tended to just produce a circle instead . The second panel in Figure 2 shows the coefficients computed by the disentanglement metric for the Reconstruction-to-Range condition . The size of each square denotes the relative importance of a latent ( column ) in predicting the corresponding generative factor ( row ) . The higher the disentanglement , the sparser the matrices . An examination of these matrices revealed that different models achieved a large range of disentanglement though none of the end-to-end models achieved perfect disentanglement .

Summary: Learning disentangled representation is often considered an important step to achieve human-like generalization. This paper studies how the degree of disentanglement affects various forms of generalization.  Variational autoencoders (VAEs) is trained with different levels of disentanglement on an unsupervised task by excluding combinations of generative factors during training. At test time the models are used to reconstruct the missing combinations in order to measure generalization performance. The paper shows that the models support only weak combinatorial generalization. The paper also tests the models in a more complex task which explicitly required independent generative factors to be controlled. The paper concludes that learning disentanglement representation is not sufficient for supporting more difficult forms of generalization.