Overcoming the Disentanglement vs Reconstruction Trade-off via Jacobian Supervision

1 INTRODUCTION . A desired characteristic of deep generative models is the ability to output realistic images while controlling one or more of the factors of variation underlying the image formation . Moreover , when each unit in the model ’ s internal image representation is sensitive to each of these factors , the model is said to obtain disentangled representations . Learning such models has been approached in the past by training autoencoders where the latent variables ( or a subset of them ) are constrained to correspond to given factors of variation , which can be specified ( supervised ) or learned from the data ( unsupervised ) ( Bengio et al. , 2013 ; Mathieu et al. , 2016 ; Hu et al. , 2017 ; Szabo et al. , 2018 ; Kim & Mnih , 2018 ) . The remaining latent variables are typically considered nuisance variables and are used by the autoencoder to complete the reconstruction of the image . There exists one fundamental problem when learning disentangled representations using autoencoders , sometimes referred to as the “ shortcut problem ” ( Hu et al. , 2017 ; Szabo et al. , 2018 ) . If the dimension of the latent code is too large , the decoder ignores the latent variables associated to the specified factors of variation , and achieves the reconstruction by using the capacity available in the nuisance variables . On the other hand , if the dimension of the latent code is small , the decoder is encouraged to use the specified variables , but is also limited in the amount of information it can use for reconstruction , so the reconstructed image is more distorted with respect to the autoencoder ’ s input . Szabo et al . ( 2018 ) showed that this trade-off between reconstruction and disentangling can indeed be traversed by varying the dimension of the latent code . However , no principled method exists to choose the optimal latent code dimension . The shortcut problem was also addressed by using additional mechanisms to make sure the decoder output is a function of the specified factors in the latent code . One approach , for example , consists in swapping the specified part of the latent code between different samples , and using adversarial 1Source code available at https : //github.com/jlezama/disentangling-jacobian . training to make sure the output distribution is indeed conditioned to the specified factors ( Mathieu et al. , 2016 ; Lample et al. , 2017 ; Szabó et al. , 2017 ; Szabo et al. , 2018 ) . However , adversarial training remains a difficult and unstable optimization problem in practice . Based on these observations , we propose a method for avoiding the shortcut problem that requires no adversarial training and achieves good disentanglement and reconstruction at the same time . Our method consists in first training an autoencoder model , the teacher , where the dimension of the latent code is small , so that the autoencoder is able to effectively disentangle the factors of variation and condition its output on them . These factors can be specified in a supervised manner or learned from the data in an unsupervised way , as we shall demonstrate . After the teacher model is trained , we construct a student model that has a larger latent code dimension for the nuisance variables . For the student , we optimize the reconstruction loss as well as an additional loss function that constrains the variation of the output with respect to the specified latent variables to be the same as the teacher ’ s . In what follows , we consider autoencoder models ( E , D ) , that receive an image x as input and produce a reconstruction x̂ : D ( E ( x ) ) = x̂ . We consider that the latent code is always split into a specified factors part y ∈ Rk and a nuisance variables part z ∈ Rd : E ( x ) = ( y , z ) , D ( y , z ) = x̂ . Consider a teacher autoencoder ( ET , DT ) , with nuisance variables dimension dT , and a student autoencoder ( ES , DS ) with nuisance variables dimension dS ; dS > dT . Because the dimension of the nuisance variables of the student is larger than in the teacher model , we expect a better reconstruction from it ( i.e . ||x− x̂S || < ||x− x̂T || , for some norm ) . At the same time , we want the student model to maintain the same disentangling ability as the teacher as well as the conditioning of the output on the specified factors . A first order approximation of this desired goal can be expressed as ∂x̂Sj ∂yi ≈ ∂x̂Tj ∂yi , ( 1 ) where j ∈ { 1 ... H ·W · C } , H , W and C are the dimensions of the output image , and i ∈ { 1 ... k } indexes over the specified factors of variation . In this paper we propose a method to impose the first-order constraint in ( 1 ) , which we term Jacobian supervision . We show two applications of this method . First , we propose an unsupervised algorithm that progressively disentangles the principal factors of variation in a dataset of images . Second , we use the Jacobian supervision to train an autoencoder model for images of faces , in which the factors of variation to be controlled are facial attributes . Our resulting model outperforms the state-of-theart in terms of both reconstruction quality and facial attribute manipulation ability . 2 RELATED WORK . Autoencoders ( Hinton & Salakhutdinov , 2006 ; Bengio et al. , 2013 ; Kingma & Welling , 2014 ) are trained to reconstruct an input image while learning an internal low-dimensional representation of the input . Ideally , this representation should be disentangled , in the sense that each hidden unit in the latent code should encode one factor of variation in the formation of the input images , and should control this factor in the output images . There exist extensive literature on learning disentangled representations ( Rifai et al. , 2012 ; Bengio , 2013 ; Cheung et al. , 2014 ; Kingma et al. , 2014 ; Cogswell et al. , 2015 ; Chen et al. , 2016 ; Mathieu et al. , 2016 ; Szabó et al. , 2017 ; Perarnau et al. , 2016 ; Hu et al. , 2017 ; Kim & Mnih , 2018 ; Burgess et al. , 2018 ) . Disentangled representations have two important applications . One is their use as rich features for downstream tasks such as classification ( Rifai et al. , 2012 ; Tran et al. , 2017 ) or semi-supervised learning ( Kingma et al. , 2014 ) . In the face recognition community , for example , disentanglement is often used to learn viewpoint- or pose-invariant features ( Yang et al. , 2015 ; Peng et al. , 2017 ; Tran et al. , 2017 ) . A second important application is in a generative setting , where a disentangled representation can be used to control the factors of variation in the generated image ( Yan et al. , 2016 ; Perarnau et al. , 2016 ; Higgins et al. , 2016 ; Mathieu et al. , 2016 ; Szabó et al. , 2017 ; Lample et al. , 2017 ) . In this work we concentrate on the second one . In recent years , with the advent of Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , a broad family of methods uses adversarial training to learn disentangled representations ( Mathieu et al. , 2016 ; Szabó et al. , 2017 ; Lample et al. , 2017 ; Perarnau et al. , 2016 ; Chen et al. , 2016 ) . In a generative setting , the adversarial discriminator can be used to assess the quality of a reconstructed image for which the conditioning factors do not exist in the training set ( Mathieu et al. , 2016 ; Szabó et al. , 2017 ; Chen et al. , 2016 ) . Another alternative , proposed in Fader Networks ( Lample et al. , 2017 ) , is to apply the adversarial discriminator on the latent code itself , to prevent it from containing any information pertaining to the specified factors of variation . Then , the known factors of variation or attributes are appended to the latent code . This allows to specify directly the amount of variation for each factor , generating visually pleasing attribute manipulations . Despite being trained on binary attribute labels , Fader Networks generalize remarkably well to real-valued attribute conditioning . However , despite recent advances ( Arjovsky et al. , 2017 ; Gulrajani et al. , 2017 ) , adversarial training remains a non-trivial min-max optimization problem , that in this work we wish to avoid . Other remarkable disentangling methods that require no adversarial training are : Cheung et al . ( 2014 ) , where the cross-covariance between parts of the latent representation is minimized , so that the hidden factors of variation can be learned unsupervised and Higgins et al . ( 2016 ) ; Kim & Mnih ( 2018 ) ; Burgess et al . ( 2018 ) where a factorized latent representation is learned using the Variational Autoencoder ( VAE ) framework . In particular , the authors of Burgess et al . ( 2018 ) , propose to overcome the disentangling versus reconstruction trade-off by progressively allowing a larger divergence between the factorized prior distribution and the latent posterior in a VAE . Related to the task of varying the factors of image generation is that of domain-transfer ( Reed et al. , 2015 ; Zhu et al. , 2017 ; Isola et al. , 2017 ; Choi et al. , 2017 ; Donahue et al. , 2016 ; Liu & Tuzel , 2016 ) . Here the challenge is to “ translate ” an image into a domain for which examples of the original image are unknown and not available during training . For example , in the face generation task , the target domain can represent a change of facial attribute such as wearing eyeglasses or not , gender , age , etc . ( Liu & Tuzel , 2016 ; Perarnau et al. , 2016 ; Yan et al. , 2016 ; Choi et al. , 2017 ) . 3 UNSUPERVISED PROGRESSIVE LEARNING OF DISENTANGLED . REPRESENTATIONS In this section we detail how the Jacobian supervision motivated in Section 1 can be applied , by ways of a practical example . We will use the Jacobian supervision to learn a disentangled image representation , where the main factors of variation are progressively discovered and learned unsupervised . We start with a simple autoencoder model , the teacher T , identified by its encoder and decoder parts ( ET , DT ) . The output of the encoder ( the latent code ) is split into two parts . One part corresponds to the factors of variation y ∈ Rk and the other part corresponds to the nuisance variables , z ∈ Rd . We begin by using k = 2 and d = 0 , meaning that the latent code of the teacher is only 2- dimensional . We consider the information encoded in these two variables as the two principal factors of variation in the dataset . This choice was done merely for visualization purposes ( Figure 1 ) . For this example , we trained a 3-layer multi-layer perceptron ( MLP ) on MNIST digits , using only the L2 reconstruction loss . We used BatchNorm at the end of the encoder , so that the distribution of y is normalized inside a mini-batch . In Figure 1 ( a ) we show the result of sampling this twodimensional variable and feeding the samples to the decoder DT . The resulting digits are blurry , but the hidden variables learned to encode the digit class . Next , we create a student autoencoder model ( ES , DS ) , similar to the teacher , but with a larger latent code . Namely , k = 2 and d = 1 instead of d = 0 , so that the latent code has now an extra dimension and the reconstruction can be improved . In order to try to maintain the conditioning of the digit class by the 2D hidden variable y , we will impose that the Jacobian of the student with respect to y be the same as that of the teacher , as in ( 1 ) . How to achieve this is described next . We take two random samples from the training set x1 and x2 , and feed them to the student autoencoder , producing two sets of latent codes : ( yS1 , z S 1 ) and ( y S 2 , z S 2 ) , and two reconstructions x̂ S 1 and x̂S2 , respectively . We then swap the parts of the latent code to form ( y S 2 , z S 1 ) and ( y S 1 , z S 2 ) and feed them to the student decoder to to obtain their respective reconstructions x̂S21 and x̂ S 12 . We also feed the same pair of images to the teacher autoencoder to obtain yT1 , y T 2 , x̂ T 1 , x̂ T 2 . Note that the teacher encoder in this case does not produce a z . We observe , by a first-order Taylor expansion , that DT ( yT2 ) = D T ( yT1 ) + JT ( y T 1 ) ( y T 2 − yT1 ) + oT ( ||yT2 − yT1 || ) , ( 2 ) and DS ( yS2 , z S 1 ) = D S ( yS1 , z S 1 ) + JS ( y S 1 , z S 1 ) [ ( yS2 , z S 1 ) − ( yS1 , zS1 ) ] + oS ( || ( yS2 , zS1 ) − ( yS1 , zS1 ) || ) ( 3 ) = DS ( yS1 , z S 1 ) + JS ( y S 1 , z S 1 ) ( y S 2 − yS1 ,0 ) + oS ( ||yS2 − yS1 || ) , ( 4 ) where JT and JS are the Jacobian of the teacher and student decoders respectively . Suppose yS1 = y T 1 and y S 2 = y T 2 , ( 5 ) and DT ( yT2 ) −DT ( yT1 ) = DS ( yS2 , zS1 ) −DS ( yS1 , zS1 ) ( 6 ) then , by simple arithmetic , JS ( y1 , z1 ) [ ( y2 − y1,0 ) ] ≈ JT ( y1 ) ( y2 − y1 ) , ( 7 ) where , since we assume ( 5 ) holds , we dropped the superscripts for clarity . What ( 7 ) expresses is that the partial derivative of the output with respect to the latent variables y in the direction of ( y2 − y1 ) is approximately the same for the student model and the teacher model . To achieve this , the proposed method consists essentially in enforcing the assumptions in ( 5 ) and ( 6 ) by simple reconstruction losses used during training of the student . Note that one could exhaustively explore partial derivatives in all the canonical directions of the space . In our case however , by visiting random pairs during training , we impose the constraint in ( 7 ) for random directions sampled from the data itself . This allows for more efficient training than exhaustive exploration . Putting everything together , the loss function for training the student autoencoder with Jacobian supervision is composed of a reconstruction part Lrec and a Jacobian part Ljac : Lrec ( x , ES , DS ) : = ||x−DS ( ES ( x ) ) ||22 = ||x−DS ( yS , zS ) ||22 = ||x− x̂S ||22 ( 8 ) Ljac ( x , ES , DS ) : = λy||yS − yT ||22 + λdiff || ( DT ( yTj ) −DT ( yT ) ) − ( DS ( ySj , z S ) −DS ( yS , zS ) ) ||22 ( 9 ) where the subscript j indicates a paired random sample . For the experiments in Figure 1 we used λy = 0.25 , λdiff = 0.1 . Table 4 in the appendix presents ablation studies on these hyperparameters . In practice , we found it also helps to add a term computing the cross-covariance between y and z , to obtain further decorrelation between disentangled features ( Cheung et al. , 2014 ) : Lxcov ( y , z ) : = ∑ ij [ 1 M M∑ m=1 ( zmi − z̄i ) ( ymj − ȳj ) ] 2 , ( 10 ) where M is the number of samples in the data batch , m is an index over samples and i , j index feature dimensions , and z̄i and ȳj denote means over samples . In our experiments we weigh this loss with λxcov = 1e−3 . Once the student model is trained , it generates a better reconstructed image than the teacher model , thanks to the expanded latent code , while maintaining the conditioning of the output that the teacher had . The extra variable in the student latent code will be exploited by the autoencoder to learn the next important factor of variation in the dataset . Examples of factors of variations progressively learned in this way are shown in Figure 2 . To progressively obtain an unsupervised disentangled representation we do the following procedure . After training of the student with k = 2 , d = 1 is finished , we consider this model as a new teacher ( equivalent to k = 3 ) , and we create a new student model with one more hidden unit ( equivalent to k = 3 , d = 1 ) . We then repeat the same procedure . Results of repeating this procedure 14 times , using 100 epochs for each stage are shown in Figure 1 . In Figure 1 ( b ) , we show how the resulting final model can maintain the conditioning of the digit class , while obtaining a much better reconstruction . A model trained progressively until reaching the same latent code dimension but without Jacobian supervision , and only the cross-covariance loss for disentangling ( Cheung et al. , 2014 ) , is shown in Figure 1 ( c ) . This model also obtains good reconstruction but loses the conditioning . For this model we also found λxcov = 1e−3 to give the best result . To quantitatively evaluate the disentangling performance of each model , we look at how the first two latent units ( k = 2 ) control the digit class in each model . We take two images of different digits from the test set , feed them to the encoder , swap their corresponding y subvector and feed the fabricated latent codes to the decoder . We then run a pre-trained MNIST classifier in the generated image to see if the class was correctly swapped . The quantitative results are shown in Table 1 . We observe that the reconstruction-disentanglement trade-off is indeed more advantageous for the student with Jacobian supervision . To complement this section , we present results of the unsupervised progressive learning of disentangled representations for the SVHN dataset ( Netzer et al. , 2011 ) in Section A.5 in the Appendix .

Summary: The paper proposes a method to tackle the disentanglement-reconstruction tradeoff problem in many disentangling approaches. This is achieved by first training the teacher autoencoder (unsupervised or supervised) that learns to disentangle the factors of variation at the cost of poor reconstruction, and then distills these learned representations into a student model with extra latent dimensions, where these extra latents can be used to improve the reconstructions of the student autoencoder compared to the teacher autoencoder. The distillation of the learned representation is encouraged via a novel Jacobian loss term that encourages the change in reconstructions of the teacher and student to be similar when the latent representation changes. There is one experiment for progressive unsupervised disentangling (disentangling factor by factor) on MNIST data, and one experiment for semi-supervised disentangling on CelebA-HQ.