Local Disentanglement in Variational Auto-Encoders Using Jacobian $L_1$ Regularization

1 Introduction . Unsupervised representation learning takes a collection of image data from the world and figures out how to organize and find patterns in the data without additional information about how the images were generated . The ideal representation learning algorithm would compress high-dimensional image data into a lower-dimensional latent representation that contains relevant information about the ground-truth factors of variation that generated the image . Inferring a good latent representation from a dataset is a difficult problem and is generally underspecified in algorithms . This underspecification is called the “ model identification ” problem , and one example is the fact that representation learning algorithms often struggle to precise the correct orientation of a latent space . That is , optimization criteria used to learn a representation function might be equally well satisfied by an equivalent representation that is just a rotation of the latent space by an arbitrary amount . As commonly implemented ( using axis-aligned Gaussian posterior distributions ) , Variational AutoEncoders ( VAEs ) [ 1 ] and their extension , the β-VAE [ 2 ] , solve the rotational part of the model identification issue by tending to ensure that the generation function ’ s Jacobian matrix has orthogonal columns ( i.e. , that the generative Jacobian matrix ’ s right singular values are aligned with the axes of the latent space ) [ 3 , 4 ] . Intuitively , this is because the VAE ’ s stochastic reconstruction cost prefers to budget higher precision ( lower embedding noise ) in the directions along which the generative Jacobian changes most rapidly . Kumar and Poole [ 4 ] draw the parallel between this preferred orientation and linear Principal Component Analysis ( PCA ) . However , we note that learning algorithms that resolve rotational identification issues through methods related to PCA will still suffer from an identifiability issue related to rotations that mix the directions for which the generative Jacobian matrix has equal 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . singular values . Along these directions , the posterior Gaussian would have approximately equal variance and be rotationally symmetric . We propose adding an L1 cost to the generation function ’ s Jacobian matrix as a way to resolve that rotational identifiability issue . Since L1 cost is not rotation invariant , L1 regularization creates a preferred latent-space orientation among directions whose singular values are equal . This use of the L1 norm to choose an orientation is inspired by similar use in linear models . For example , when using a Laplacian prior in Independent Component Analysis ( ICA ) [ 5 ] and applying it to already-whitened data , ICA rotates the data to minimize the L1 norm . As shown in Olshausen [ 6 ] , that ICA formulation is equivalent to sparse coding using an L1 cost [ 7 ] , with the same rotational effect . In Sparse PCA , the L1 norm encourages sparsity in the loadings/mixing matrix ( rather than the principal components/sources as in ICA ) [ 8 , 9 ] , similarly encouraging preferred orientations . These techniques for preferring certain orientations using the L1 norm are all linear model techniques , and we apply them to non-linear VAEs by regularizing the L1 norm of the generator Jacobian matrix , thereby encouraging a preferred orientation for our non-linear models . The motivation above suggests that an L1 norm on the generator Jacobian can address the rotational identifiability issue , and we think the L1 regularization will encourage useful orientations of the latent space for image data . This belief comes from results on the sparse linear coding of images . In particular , Olshausen and Field [ 7 ] showed how sparse linear coding on natural images generates local receptive fields similar to those discovered in the mammalian visual processing system . These types of localized receptive fields are shown in Figure 1 . Each image in that figure corresponds to a direction in the latent space , and shows how perturbing the latent value in that direction changes the generated image—white pixels means the image gets brighter there , black pixels means the image gets darker . For the ICA basis , we see that perturbations in different latent directions are associated with localized changes to the image , whereas for the PCA basis , the latent directions tend to affect the entire image . Adding an L1 penalty to our model should encourage sparsity within the generative Jacobian columns of our model , meaning that perturbations of latent values along individual latent directions should modify as few pixels as possible , leading to latent directions that affect the output image in localized regions . L1 generative Jacobian regularization should therefore disentangle representations of different objects in an image . We call the proposed model trained with this additional regularization a Jacobian L1 Regularized Variational Auto-Encoder ( JL1-VAE ) . The use of the generative Jacobian in the JL1-VAE implies a local linear approximation of the generative function , so in order for the JL1-VAE to be useful , the training image data should lie on a manifold [ 10 ] , and should be sufficiently well sampled . Additionally , JL1-VAE contains inductive bias , like every other unsupervised disentanglement algorithm [ 11 ] . As mentioned above , JL1-VAE ’ s regularization of the generative Jacobian encourages small changes in latent values to result in sparse ( impacting a small number of pixels ) changes to the resulting image , similar to local receptive fields . This inductive bias is well suited for disentangling motions of different objects in an image , but would presumably not be useful for whole-image changes , such as rotation of the entire image or brightness changes across the whole image . We apply our novel JL1-VAE framework to a variety of datasets , giving qualitative and quantitative results showing that our added L1 cost can encourage local alignment of the axes of the latent representation with individual factors of variation . 2 Background . We present a brief overview of VAEs and introduce the notation we will use throughout this paper . VAEs [ 1 ] train a model to generate a data distribution that approximately matches the distribution of unlabeled training data . For the VAEs we consider in this paper , generating a datapoint x̃ ∈ Rn from a trained VAE consists of sampling a latent variable z ∈ Rl from a standard multivariate Gaussian distribution N ( 0 , 1 ) and applying a generator function g : Rl → Rn to z to map the latent variable to a generated image g ( z ) . Around this generated image we assume a Bernoulli probability of similar images x̃ ∼ p ( x̃ ; g ( z ) ) with the generated image g ( z ) as its mean . To train the VAE we also define a multi-variate Gaussian embedding distribution q ( z|x ) for each training image x with mean h ( x ) ( using a learned embedding function h : Rn → Rl ) and diagonal covariance Σz|x ( x ) ( using a learned covariance function Σz|x : Rn → Rl that computes the diagonal elements ) . This embedding distribution q ( z|x ) is motivated by a desire to approximate the posterior distribution p ( z|x ) . The objective during training of the VAE is to maximize the Evidence Lower BOund ( ELBO ) , which is defined as ∑ x L ( x ) , where , for each data point x , L ( x ) = Ez∼q ( z|x ) [ p ( x ; g ( z ) ) ] − KL ( q ( z|x ) ‖N ( 0 , 1 ) ) ( 1 ) The ELBO is a lower bound on log ( p ( x ) ) = log ( ∫ z p ( x|z ) p ( z ) dz ) , which is the likelihood of the data point given our model . Thus , maximizing the ELBO is a proxy for maximum likelihood estimation . The β-VAE [ 2 ] is an extension to the VAE that multiplies the second term in the ELBO by an adjustable hyperparameter β . The first term in the ELBO is a stochastic reconstruction loss , so we sometimes refer to Σz|x as the “ embedding noise , ” since it adds noise to the embedding when estimating the stochastic reconstruction loss . Kumar and Poole [ 4 ] show that the stochastic reconstruction term of the ELBO , Ez∼q ( z|x ) [ log ( p ( x ; g ( z ) ) ) ] , can be approximated using a second-order Taylor expansion as : Ez∼q ( z|x ) [ log ( p ( x ; g ( z ) ) ] ≈ log p ( x|h ( x ) ) + 1 2 tr ( Jg ( h ( x ) ) > Hpx ( g ( h ( x ) ) ) Jg ( h ( x ) ) Σz|x ( x ) ) ( 2 ) Jg is the Jacobian of the generator function , and Hpx ( g ( h ( x ) ) ) is the Hessian with respect to g ( z ) of the log of the generative probability log ( p ( x ; g ( z ) ) ) evaluated at g ( h ( x ) ) . For standard VAE implementations using diagonal Gaussian posteriors and pixel-factorized generative probabilities , Σz|x ( x ) and Hpx ( g ( h ( x ) ) ) are both diagonal . Equation 2 shows that the stochastic reconstruction loss of the ELBO can be approximated by a deterministic reconstruction loss with a ( weighted ) L2 regularization cost on the Jacobian Jg ( h ( x ) ) . Kumar and Poole [ 4 ] use this weighted L2 regularization in the approximation above to show how the ELBO encourages local alignment of the right singular vectors of Jg to Σz|x ( x ) . That is , larger values of Σz|x ( x ) ( larger embedding noise ) in some directions leads to larger L2 regularization on the generator Jacobian in those directions . Equivalently , in directions with large changes in the generator Jacobian , the regularization encourages smaller embedding noise ( more precision ) in the posterior Σz|x ( x ) . For this work , we use the presence of an implicit L2 loss on the generator Jacobian as further motivation for our choice to add an explicit L1 regularization to the generator Jacobian . 3 Related Work . There are several areas of related work we wish to call to the reader ’ s attention . We note previous work in the analysis of the disentanglement properties of β-VAEs . We delve into previous uses of the term “ sparse VAE , ” as there are ( at least ) two other common and well studied meanings of “ sparse VAE , ” which we disambiguate from the type of sparsity we study in this work . We note previous VAE work inspired by ICA . We discuss architectural choices that have been shown to improve disentanglement . Finally , we discuss modifications to the VAE objective that have previously been studied to improve disentanglement . Disentanglement of β-VAEs Mathieu et al . [ 12 ] and Rolinek et al . [ 3 ] show that restricting the posterior covariance to diagonal ( called the mean-field assumption ) breaks the rotational symmetry of β-VAEs . Rolinek et al . [ 3 ] further show that this encourages the columns of the generator Jacobian to be orthogonal , relating local β-VAE latent directions with PCA decomposition . In our work , we take this a step further , explicitly regularizing the sparsity of the generator Jacobian , to break rotational symmetry between directions with equal singular values in of the generator Jacobian . Sparse VAEs : Sparsity in VAE codes Some previous work involving L1 regularization and VAEs uses the term “ sparse VAEs ” to refer to sparsity in the latent values taken on by the latent codes themselves . That is , these works attempt a minimization of something like ‖z‖1 . This meaning is studied in [ 13–15 ] . When we use sparsity in the present work , however , we are not concerned with the values of the latents z , but rather with how small axis-aligned changes in the latent values affect the output . That is , we are concerned with sparsity of Jg ( z ) , not z . Our associated cost is ‖Jg ( z ) ‖1 . Our use of “ sparsity ” has to do with local disentanglement , rather than sparsity in latent values , which is a type of global disentanglement . Sparse VAEs : Sparsity in Network Weights Likewise , other work involving L1 regularization and VAEs ( and neural networks more broadly ) uses “ sparsity ” to refer to the desire to make many network weights 0 , with the motivation of reducing the size of the stored neural network architecture . This meaning is seen in [ 16 , 17 ] . In this work , by contrast , we are interested in sparsity in the generator Jacobian , not in the network weights . We note the distinction between sparsity of individual network weights and sparsity in the Jacobian of the overall function . Multiplying sparse matrices does not necessarily result in sparse matrices , and non-linear activations can allow a sparse Jacobian even if the network weights themselves are not particularly sparse . Thus , the ‖Jg ( z ) ‖1 cost we study in this work is not what is referred to in studies of sparsity of neural networks , which consider regularizations like the L1 cost over network weights . ICA within the VAE Literature Independent Component Analysis ( ICA ) [ 5 ] is often mentioned in the VAE literature in relation to the role ICA has played in the theory of identifiability and disentanglement of representations [ 4 , 12 , 3 , 18 , 19 ] . Stuehmer et al . [ 20 ] propose using a structured , rotationally asymmetric prior to encourage disentanglement in the embedding . This , and other approaches that attempt to globally match the embedding distribution to a desired shape , are very different from the local , Jacobian-based approach we take in this paper . Khemakhem et al . [ 21 ] relate nonlinear ICA with VAEs in the case where the data has an additionally observed variable , and they give a proof that in that case their model is identifiable and correctly disentangles the ground-truth factors of variation . We focus on fully unsupervised training data and assume that we are not given access to any data labels . We are not aware of any prior work applying a sparsity cost to the generator Jacobian , which is the inspiration we take from ICA . Architectures shown to improve disentanglement Previous work has shown impressive results from modifying the network architecture in order to explicitly represent multiple objects by , for example , learning object masks [ 22 ] , or by modifying how the latent variable is read in to the generative model architecture [ 23 ] . Ainsworth et al . [ 24 ] and Khan et al . [ 25 ] both construct model architectures that explicitly encourage sparsity in the generative network weights . In this work , we focus on how we can regularize the objective function to improve disentanglement , rather than studying how different network architectures can improve disentanglement . Modification of VAE prior Previous work has also investigated modifications to the unit Gaussian prior commonly used in VAEs . Tomczak and Welling [ 26 ] use a learnable Gaussian mixture prior . Stuehmer et al . [ 20 ] use a generalized Gaussian distribution ( that is not rotationally invariant ) as the prior . Bauer and Mnih [ 27 ] use rejection sampling to form a more complicated , non-rotationallyinvariant prior . Davidson et al . [ 28 ] and Perez Rey et al . [ 29 ] even modify the prior to lie on non-Euclidean surfaces . Kim and Mnih [ 18 ] do not explicitly enforce a prior distribution , but rather use a regularization term to encourage the prior distribution q ( z ) to be a factorized distribution . While these approaches pressure the entire embedded distribution to have certain properties , we are instead focused on how to modify the learning objective to give local bias toward disentanglement , rather than using more global methods based on the overall distribution of the embedded dataset . Regularization of VAEs Several previous works explicitly or implicitly use L2 normalization of the generator Jacobian [ 4 , 30–32 ] . Chen et al . [ 33 ] regularize by ‖J > g ( z ) Jg ( z ) − c1‖2 for some constant c. We are not aware of previous investigations of ‖Jg ( z ) ‖1 regularization for VAEs .

Summary: Authors propose to regularize the Jacobian of the generator in a VAE setting so as to encourage local alignment of the latent axes, i.e. `local disantanglement'. They do so by finite differences, and empirically show on the MPI3D dataset the effectiveness of their approach.