Structure by Architecture: Disentangled Representations without Regularization

1 INTRODUCTION . Deep learning has achieved strong results on a plethora of challenging tasks . However , performing well on a highly specific dataset is usually insufficient to satisfactorily solve real-world problems ( Tan et al. , 2018 ; Zhuang et al. , 2019 ) . This has lead to a particular interest in consistently learning more structured representations with useful properties to help with a variety of downstream tasks ( Bengio et al. , 2013 ; Tschannen et al. , 2018 ; Bengio et al. , 2017 ; Tschannen et al. , 2018 ; Radhakrishnan et al. , 2018 ) . Here , deep learning provides a flexible paradigm to train complex architectures based on autoencoders ( Ballard , 1987 ; Lange and Riedmiller , 2010 ) which are typically latent variable models , thus allowing us to embed powerful inductive biases into the models to further structure the representations . However , it is still largely an open question as to what kinds of structure in a representation are the most effective for generative modeling and how to learn such structures without supervision ( Locatello et al. , 2018 ; Khrulkov et al. , 2021 ; Shu et al. , 2019 ; Chen et al. , 2020 ; Nie et al. , 2020 ; Mathieu et al. , 2019 ; Vaswani et al. , 2017 ; Kwon and Ye , 2021 ) . One direction that may contribute to an answer is causal modeling , as it focuses on the underlying ( causal ) mechanisms that generate the observations , instead of relying on ( possibly spurious ) correlations ( Pearl , 2009 ; Peters et al. , 2017 ; Schölkopf , 2019 ; Louizos et al. , 2017 ; Mitrovic et al. , 2020 ; Shen et al. , 2020 ) . With the versatility of deep learning on one hand , and the conceptual insights of causality on the other , our contributions herein include : • We propose an architecture called the Structural Autoencoder ( SAE ) , where the structural decoder emulates a general acyclic structural causal model to learn a hierarchical representation that can separate and order the underlying factors of variation in the data . • We provide a sampling method that can be used for any autoencoder-based generative models which does not use an explicit regularization objective and instead relies on independence in the latent space . • We investigate how well the encoder and decoder are able to extrapolate upon seeing novel samples . We release our code at * anonymized * . 1.1 RELATED WORK . The most popular autoencoder based method is the Variational Autoencoder ( VAE ) ( Kingma and Welling , 2013a ) and the closely related βVAE ( Higgins et al. , 2017 ) . These methods focus on matching the distribution in the latent space to a known prior distribution by regularizing the reconstruction training objective ( Locatello et al. , 2020a ; Zhou et al. , 2020 ) . Although this structure is convenient for generative modeling and even tends to disentangle the latent space to some extent , it comes at the cost of somewhat blurry images due to posterior collapse and holes in the latent space ( Locatello et al. , 2018 ; Higgins et al. , 2017 ; Burgess et al. , 2018 ; Kim and Mnih , 2018 ; Stühmer et al. , 2020 ) . To mitigate the double-edged nature of VAEs , less aggressive regularization techniques have been proposed such as the Wasserstein Autoencoder ( WAE ) , which focuses on the aggregate posterior ( Tolstikhin et al. , 2018 ) . Unfortunately , WAEs generally fail to produce a particularly meaningful or disentangled latent space ( Rubenstein et al. , 2018 ) , unless weak supervision is available ( Han et al. , 2021 ) . A more structured alternative is the Variational Ladder autoencoder ( VLAE ) ( Zhao et al. , 2017 ) which separates the latent space into separate chunks each of which is processed at different levels of the encoder and decoder ( called ” rungs ” ) . Our proposed architecture takes inspiration from the VLAEs but crucially we do not use the variational regularization , and instead use a sampling method based on hybridization ( Besserve et al. , 2018 ) . We also infuse the information from the latent variables using Str-Tfm layers ( see section 2.2 ) which are based on the Ada-IN layers from Style-GANs ( Karras et al. , 2019 ) , which significantly improves on the ladder rungs . Lastly , unlike the VLAEs , the SAEs use a simple feed-forward CNN encoder ( like all the other unstructured models ) , which simplifies analysis without significantly impacting the ability to structure the representation ( see further discussion in section 4.2 ) . 2 METHODS . 2.1 CAUSAL REPRESENTATION LEARNING . Graphical causal modeling represents the relationship between random variables Si using a directed acyclic graph ( DAG ) whose edges indicate direct causation and structural equations of the form Si : = fi ( PAi , Ui ) , ( i = 1 , . . . , D ) , ( 1 ) encoding the dependence of variable Si on its parents PAi in the graph and on an unexplained noise variable Ui . The noises U1 , . . . , UD are assumed to be jointly independent . The DAG along with the mechanisms ( 1 ) is referred to as a Structural Causal Model ( SCM ) ( Pearl , 2009 ) . Any joint distribution of the Si can be expressed as an SCM using suitable fi and Ui . However , the SCM entails additional assumptions regarding how statistical dependencies between the Si are generated by mechanisms ( 1 ) , such that changes due to interventions can be modelled as well ( e.g. , by setting some Ui or Si to constants ) . Real-world observations are often not structured into meaningful causal variables and mechanisms to begin with . E.g. , images are high-dimensional , and it is hard to learn objects and their causal relationships from data ( Lopez-Paz et al. , 2017 ) . One may thus attempt to learn a representation consisting of causal variables or disentangled “ factors ” which are statistically independent ( Higgins et al. , 2017 ) . However , in an SCM it is not the Si that should be statistically independent , but the Ui . For this reason , our representations will comprise the Ui as latent variables , driving causal mechanisms via learned functions fi ( 1 ) . This embeds an SCM into a larger model whose inputs and outputs may be high-dimensional and unstructured ( e.g. , images ) ( Suter et al. , 2018a ; Schölkopf , 2019 ) . Given ( high-dimensional ) X = ( X1 , . . . , Xd ) , our goal is to learn a low-dimensional representation U = ( U1 , . . . , UD ) ( D d ) using an encoder Rd → RD , and model the generative process ( including the SCM ) that produced X from the inferred latent variables using a decoder RD → Rd . If the causal graph of the true generative process were known , the topology of the decoder could be fixed accordingly . However , in the fully unsupervised setting , the generator must be capable of modeling a general SCM , so our architecture specifies a causal ordering ( while learning what information to embed in the specified ordering ) , while the edges between parents and children in the SCM are learned implicitly in the computation layers . Specifically , the learned generative process , a.k.a . decoder , produces a reconstruction X̂ of X by feeding each of the Ui into subsequent computation layers . Here , the root node S1 in the DAG only depends on U1 , while later Si depend on their noise Ui and potentially their parents Sj ( j < i ) . Thus the depth in the network corresponds to a causal ordering ( see Supplement for further discussion ) . 2.2 STRUCTURAL DECODERS This model architecture is implemented by the structural decoder , using D StructuralTransform ( Str-Tfm ) layers placed evenly in between the convolution blocks . From the corresponding latent variable Ui , the ith Str-Tfm layer estimates a scale αi and bias βi which are then used to pixelwise transform the intermediate features of the lth layer vl ( as seen in Figure 1 ) much like in Ada-IN ( Karras et al. , 2019 ) except without the preceding normalization . Each Str-Tfm layer thus acts like an fi in ( 1 ) by integrating the information from a latent variable Ui to transform the features PAi from earlier layers . This variable decoder depth for each of the latent variables biases high-level non-linear information towards the earlier ( and thereby deeper ) latent variables , while the model capacity is reduced for the later variables so they can only capture relatively low-level linear features with respect to the data space . Meanwhile , the random initialization of the MLPs in the Str-Tfm layers produce more distinct activations with re- spect to each of the latent variables than if they are all transformed linearly by same dense layers . This architectural asymmetry between latent variables thereby encourages statistical independence and induces a relatively intuitive hierarchical structure of the latent space . In theory , without some form of supervision or side information , the learned latent variables are not guaranteed to disentangle the true factors of variation , much less that the learned SCM matches the true one ( Locatello et al. , 2020b ) . Instead , all we are guaranteed from training on observational data is that the model is optimized to reproduce the same observational distribution as the true generative process . However , for datasets with independent factors of variation , the independent latent variables may align with the true factors , thus disentangling the factors in the representation . Regardless of disentanglement , independence between latent variables also enables principled interventions on individual variables for generating new samples , using hybrid sampling , which is akin to intervening on the learned SCM for generative modeling . 2.3 HYBRID SAMPLING . For generative modeling , it is necessary to sample novel latent vectors that are transformed into ( synthetic ) observations using the decoder . Usually , this is done by regularizing the training objective so the posterior matches some simple prior ( e.g . the standard normal ) . However , in practice , regularization techniques can fail to match the prior perfectly and actually exacerbate the information bottleneck , leading to blurry samples from holes in the learned latent distribution and unused latent dimensions due to posterior collapse ( Dai and Wipf , 2019 ; Stühmer et al. , 2020 ; Lucas et al. , 2019 ; Hoffman and Johnson , 2016 ) . Instead of trying to match some prior distribution in the latent space , we suggest an alternative sampling method that eliminates the need for any regularization of the loss . Inspired by Besserve et al . ( 2018 ) , we refer to it as hybrid sampling : the model stores a finite set of N ( = 128 in our case ) latent vectors , selected uniformly at random from the training set . To generate a new latent vector Ũ , a value for each of theD latent variables Ũi is selected independently from theN stored latent vectors { U ( j ) } Nj=1 uniformly at random , such that : Ũi ← U ( j ) i where j ∼ { 1 , ... , N } This allows the model to generate ND distinct latent vectors , and since we assume D d , N can easily be increased up to the full training set size without significant memory costs . Note that hybrid sampling is directly applicable to any learned representation as it does not affect training at all , however the fidelity of generated samples will diminish if there are strong correlations between latent dimensions . Consequently , the goal is to achieve maximal independence between latent variables without compromising on the fidelity of the decoder ( i.e . reconstruction error ) . Hybrid sampling implicitly relies on the statistical independence of the latent variables since resampling the marginal of the aggregate posterior independently breaks any existing correlations . Not only does this align well with the objectives of unsupervised disentanglement methods , but it is also consistent with the causal perspective of the latent variables as independent noises Ui driving an SCM .

Summary: This paper structures how latents are used in an autoencoder to improve its performance. They are motivated by causal structure and independence of latent variables. They also sample using a sort of discrete mixup between latent codes from pushing forward data, and show various improvements there.