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 presents a hierarchical latent variable model and accompanying sampling procedure for learning disentangled representations. Rather than a latent feature map/vector, the latent variables are used to condition affine transforms in the decoder. The authors combine this with a ‘hybrid’ sampling strategy, effectively sampling from the aggregate approximate posterior. Comparing with several baselines, the authors compare their model in terms of FID and various disentanglement scores. The authors also investigate generalization.