Deep automodulators

1 INTRODUCTION . This paper introduces a new generative autoencoder for learning representations of image data sets , in a way that allows arbitrary combinations of latent codes to generate images ( see Fig . 1 ) . We achieve this with an architecture that uses adaptive instance normalization ( AdaIn , Dumoulin et al. , 2017b ; Huang & Belongie , 2017 ) , and training methods that let the model learn a highly disentangled latent space by utilizing progressively growing autoencoders ( Heljakka et al. , 2019 ) . In a typical autoencoder , input images are encoded into latent space , and the information of the latent variables is then passed through successive layers of decoding until a reconstruction of the input image has been formed . In our model , the latent vector independently modulates the statistics of each layer of the decoder—the output of layer n is no longer solely determined by the input from layer n− 1 . In image generation , the probability mass representing sensible images ( such as human faces ) lies concentrated on a low-dimensional manifold . Even if impressive results have been shown for image generation ( e.g. , by GANs , Goodfellow et al. , 2014 ; Karras et al. , 2019 ) , efficient reconstruction and manipulation remain open problems . Deep generative autoencoders provide a principled approach for feature extraction , editing , and image synthesis . We show that modulation of decoder layers with AdaIn further improves these capabilities . It also yields representations that are more disentangled , a property here broadly defined as something that allows for fine control of one semantic ( image ) sample attribute independently of others . Previous works on AdaIn are mostly based on GAN models ( Karras et al. , 2019 ; Chen et al. , 2019 ) with no encoder for new input images . Autoencoders are typically single-pass encoder–decoder structures , where a sample enters from one end and is reconstructed at the other . For robustness , the reconstructed samples could be reintroduced to the encoder , repeating the process , and requiring consistency between the passes . In comparison to stochastic models ( like VAEs , Kingma & Welling , 2014 ; Rezende et al. , 2014 ) , our deterministic model is better suited to take advantage of this . We can mix the latent codes of separate samples and measure the conservation of layer-specific information for each . This enforces disentanglement of layer-specific properties , because we can ensure that the latent code introduced from a certain layer onwards on the 1st pass should not affect the previous layers on the 2nd pass . Image autoencoders with convolutional architectures are often used for finding attributes that are disentangled from each other in the latent space . Often , they have relatively poor output image quality , and prior information must be fed in either via class-conditioning or by more complex loss functions for the latent variables . Fully implicit methods avoid this , but without a built-in encoding mechanism often cover only a relatively small fraction of the variation in the training data . For the first time , we show a full autoencoder architecture that has a state-of-the-art disentanglement performance even for unsupervised training setup in several data sets , sharp image output quality , good coverage of the variation of the training data , arbitrary mixing of latent representations and a principled approach for incorporating scale-specific prior information . Our contributions are as follows . ( i ) We introduce the automodulator , a new autoencoder-like model with powerful properties not found in regular autoencoders , including scale-specific style transfer ( Gatys et al. , 2016 ) . In contrast to architecturally similar ‘ style ’ -based GANs , we can now directly encode and manipulate new inputs . ( ii ) We present methodology that allows fully unsupervised training of such a model to learn rich latent representation in high resolutions , and specifically explain a novel multi-pass training principle that enables the model to learn a more disentangled representation of training data than regular autoencoders . ( iii ) We demonstrate qualitative and quantitative performance and applications on CELEBA-HQ and LSUN Bedrooms and Cars data sets , taking advantage of the highly disentangled model structure . 2 RELATED WORK . Our work builds upon several lines of previous work in unsupervised representation learning . The most relevant concepts are variational autoencoders ( VAEs , Kingma & Welling , 2014 ) and generative adversarial networks ( GANs , Goodfellow et al. , 2014 ) . In VAEs , an encoder maps data points to a lower dimensional latent space and a decoder maps the latent representations back to the data space . The model is learnt by minimizing the reconstruction error together with a regularization term that encourages the distribution of latents to match a predefined prior . Latent representations often provide useful features for applications ( e.g. , image analysis and manipulation ) and decoding random samples from the prior allows novel data generation . However , in case of images , the generated samples are often blurry and not fully photorealistic , with imperfect reconstructions . Current state-of-the-art in generative image modeling is represented by GAN models ( Brock et al. , 2019 ; Karras et al. , 2019 ) which achieve higher image quality than VAE based models . Nevertheless , these GANs lack an encoder for obtaining the latent representation for a given image , limiting their usefulness . In some cases , a given image can be semantically mapped to the latent space via generator inversion but this iterative process is prohibitively slow for many applications , and the result may depend on initialization ( Abdal et al. , 2019 ; Creswell & Bharath , 2019 ) . Bidirectional mapping has been targeted by VAE-GAN hybrids such as Makhzani et al . ( 2016 ) ; Mescheder et al . ( 2017 ) ; Makhzani ( 2018 ) , and adversarial models such as Donahue et al . ( 2017 ) ; Dumoulin et al . ( 2017a ) ; Donahue & Simonyan ( 2019 ) . These models learn mappings between the data space and latent space using combinations of encoders , generators and discriminators . IntroVAE ( Huang et al. , 2018 ) and AGE ( Ulyanov et al. , 2018 ) are more compact autoencoders , which both contain only two deep networks , encoder and decoder , learnt using a combination of reconstruction loss and adversarial loss , which encourages the distributions of real and generated images to match with the prior in latent space . A progressively grown version of AGE , PIONEER , was proposed by Heljakka et al . ( 2018 ; 2019 ) . There is also recent work on vector quantized variational autoencoders ( VQ-VAE , Razavi et al. , 2019 ) with a discrete latent space . However , while suitable for image compression , the discrete latent representation can not e.g . be interpolated and hence may not be optimal for semantic image manipulation and learning a disentangled latent space . The architecture of our generator and AdaIn utilization was inspired by the recent StyleGAN ( Karras et al. , 2019 ) , but has no discriminator ( nor the ‘ mapping network ’ f ) . Instead , our model simply contains the generator ( i.e. , decoder ) and an encoder , jointly trained applying a similar progressive growing strategy as Heljakka et al . ( 2018 ) , but with a modified loss formulation . Besides the reconstruction loss and adversarial loss from Ulyanov et al . ( 2018 ) and Heljakka et al . ( 2018 ) , we can also recirculate style-mixed reconstructions as ‘ second-pass ’ training samples in order to better encourage the independence and disentanglement of emerging styles and conservation of layerspecific information . The recirculation idea is biologically motivated and conceptually related to many works , such as that by Hinton & McClelland ( 1988 ) . Utilizing the outputs of the model as inputs for the next iteration is related to , e.g. , Zamir et al . ( 2017 ) , where feedback is shown to benefit image classification , and in RNN-based methods ( Rezende et al. , 2016 ; Gregor et al. , 2015 ; 2016 ) . 3 METHODS . We now describe the model architecture , starting from regular PIONEER and building up to the automodulator ( Sec . 3.1 ) . We then describe the unsupervised training via our extensions to the regular PIONEER loss function , to support flexible mixing of latents ( Sec . 3.2 ) . Finally , we introduce an additional weakly supervised training method to utilize known invariances ( Sec . 3.3 ) . 3.1 AUTOMODULATOR ARCHITECTURE . Our interest is in unsupervised training of an autoencoder wherein the inputs x are images fed through an encoder φ to form a low-dimensional latent space representation z ( we use z ∈ R512 , normalized to unity ) . This representation can then be decoded back into an image x̂ through a decoder θ . For stable training , we adopt the progressively growing architecture of Balanced PIONEER ( Heljakka et al. , 2018 ; 2019 ) . The convolutional layers of the symmetric encoder and decoder are faded in gradually during the training , in tandem with the resolutions of training images and generated images ( Fig . 2 ) . Adaptive Instance Normalization ( AdaIn ) A traditional decoder architecture would start from a small-resolution image and expand it layer by layer until the full image is formed , feeding the full information of the latent code via the decoder layers . In contrast , our decoder is composed of layer-wise functions θi ( ξi−1 , z ) that separately take a ‘ canvas ’ variable ξi−1 denoting the content input from the preceding decoder layer ( see Figs . 2 and 3a ) , and the actual ( shared ) latent code z . First , each deconvolutional layer # i computes its feature map output χi from ξi−1 as in traditional decoders . Second , we take the AdaIn normalization step that uses z to modulate the channel-wise mean µ and variance σ2 of the output . To do this , we need a map gi : z 7→ ( µi , σ2i ) , arriving at : ξi = AdaIn ( χi , gi ( z ) ) = σi ( χi − µ ( χi ) σ ( χi ) ) + µi . ( 1 ) We implement gi for layer # i as a fully connected linear layer , with output size 2Ci for Ci channels . Layer # 1 starts from a constant input ξ ( 0 ) ∈ R4×4 . Here , we focus on pyramidal decoders with monotonically increasing resolution and decreasing number of channels , but any deep decoder would be applicable . The AdaIn-based architecture allows the output of layer # i to be not solely determined by the input from layer # i− 1 , enabling finer control over the image information , image mixing schemes ( Sec . 3.2 ) , and utilization of scale-specific invariances in training data ( Sec . 3.3 ) . 3.2 STRONG CONSERVATION OF CYCLIC INFORMATION : LAYER-SPECIFIC LOSS METRICS . In training of regular AGE and PIONEER models , the encoder φ and the decoder θ are trained at separate steps , optimizing either the loss Lφ or Lθ , correspondingly : Lφ = DKL [ qφ ( z |x ) ‖N ( 0 , I ) ] −DKL [ qφ ( z | x̂ ) ‖N ( 0 , I ) ] + λX dX ( x , θ ( φ ( x ) ) ) , Lθ = DKL [ qφ ( z | x̂ ) ‖N ( 0 , I ) ] + λZ dcos ( z , φ ( θ ( z ) ) ) , ( 2 ) where x is sampled from the training set , x̂ ∼ qθ ( x | z ) , z ∼ N ( 0 , I ) , dX is L1 or L2 distance , and dcos is cosine distance . Since the model allows sampling from latent space ( similarly to VAEs and GANs ) , the KL divergence terms during training can be calculated from empirical distributions , despite the model itself being fully deterministic . Yet we retain , in principle , the full information about the image , at every stage of the processing . For an image encoded into a latent vector z , decoded back to image space as x̂ , and re-encoded as latent vector z′ , it is possible and desirable to require that z is as close to z′ as possible , yielding the latent reconstruction error dcos ( z , φ ( θ ( z ) ) ) . Layer-specific Loss We now show how to encourage the latent space to become hierarchically disentangled with respect to the levels of image detail , allowing one to separately retrieve ‘ coarse ’ and ‘ fine ’ aspects of a latent code . This allows for e.g . conditional sampling by fixing the latent code for specific layers of the decoder , or mixing the scale-specific features of two or more input images—impossible feats for a traditional autoencoder with mutually entangled decoder layers . The latent reconstruction error as such can be trivially generalized to any layer of θ . We simply pick the layer of measurement , and from there , pass the sample once through a full encoder-decoder cycle . But now , in the automodulator , latent codes can be introduced on a per-layer basis , enabling more powerful reconstruction measurements . ( Without loss of generality , here we only consider mixtures of 2 codes . ) We can present a decoder ( Fig . 3b ) with N layers , split after the jth one , as a composition of θj+1 : N ( θ1 : j ( ξ ( 0 ) , zA ) , zB ) . Crucially , we can choose zA 6= zB ( extending the method of Karras et al. , 2019 ) , such as zA = φ ( xA ) and zB = φ ( xB ) for ( image ) inputs xA 6= xB . Because the earlier layers # 1 : j operate on image content at lower resolutions , the output fusion image mixes the ‘ coarse ’ features of zA with ‘ fine ’ features of zB . Now , z holds feature information at different levels of detail , some of which are mutually independent . Hence , when re-encoding an image , we should keep the representation of those levels disentangled in z , even if they originate from separate source images . Hence , when we re-encode the fusion image into ẑAB , and decode once more , the output of θ1 : j ( ξ ( 0 ) , ẑAB ) should be unaffected by zB . This leads to the following conjecture . Assume that the described network reconstructs input samples perfectly , i.e . x = θ ( φ ( x ) ) . Any zA and zB can be mixed , decoded and re-encoded as ẑAB ∼ qφ ( z |x ) qθj+1 : N ( x | ξ ( j ) , zB ) qθ1 : j ( ξ ( j ) | ξ ( 0 ) , zA ) . ( 3 ) Now , between ξ ( j ) of the first and ξ ( j ) of the second pass ( see Fig . 3c ) , the mutual information is I [ qθ1 : j ( ξ ( j ) | ξ ( 0 ) , ẑAB ) ; qθ1 : j ( ξ ( j ) | ξ ( 0 ) , zA ) ] . ( 4 ) With zA , zB ∼ N ( 0 , I ) , for each j , we can now maximize ( 4 ) by minimizing Lj = d ( θ1 : j ( ξ ( 0 ) , ẑAB ) , θ1 : j ( ξ ( 0 ) , zA ) ) ( 5 ) for some distance function d ( here , L2 norm ) . In other words , the fusion image can be encoded into a new latent vector ẑAB in such a way that , at each layer , the decoder will treat the new code similarly to whichever of the original two separate latent codes was originally used there ( see Fig . 3b ) . For a perfect network , Lj can be viewed as a layer entanglement error . Randomizing j during the training , we can measure Lj for any layers of the decoder . zB zAB xA x̂AB Final Unsupervised Automodulator Training Loss Our results improved by using a robust image reconstruction loss dρ of Barron ( 2019 ) instead of L1/L2 . dρ generalizes various norms and exposes robustness as an explicit continuous parameter vector α . The complete loss functions are Lφ = max ( −Mgap , DKL [ qφ ( z |x ) ‖N ( 0 , I ) ] −DKL [ qφ ( z | x̂ ) ‖N ( 0 , I ) ] ) + λX dρ ( x , θ ( φ ( x ) ) ) Lθ = DKL [ qφ ( z | x̂ ) ‖N ( 0 , I ) ] + λZ dcos ( z , φ ( θ ( z ) ) ) + Lj , ( 6 ) where x̂1 : 34M ∼ qθ ( x | z ) with z ∼ N ( 0 , I ) , and x̂ 34M : M ∼ qθ ( x | ẑAB ) , with a set 3:4 ratio of regular and mixed samples of batch size M and j ∼ U { 1 , N } . Margin Mgap = 0.5 in low resolutions , then 0.2 for 1282 and 0.4 for 2562 ( see Heljakka et al. , 2019 ) . To avoid discontinuities in α , we utilize a progressively-growing variation of dρ , where we first learn theα in the lower resolutions ( e.g. , 4×4 ) . Each αi corresponds to one pixel . Then , when switching to the higher resolution stage , we take take each parameter αi that corresponds to pixels px , y in the lower resolution , to initialize the α1×4j that , in the higher resolution , corresponds to px , y , px+1 , y , px , y+1 and px+1 , y+1 , respectively .

Summary: Deep Automodulators introduces a generative autoencoder architecture that replaces the canonical encoder decoder autoencoder architecture with one inspired by StyleGAN. The encoder interacts with the decoder by modulating layer statistics via Adaptive Instance Normalization (AdaIN) conditioned on the latent. The paper trains this architecture with the loss framework of the Adversarial Generator–Encoder (AGE) and utilizes the progressive growing trick originally introduced in Progressive GAN which is also adapted by the Pioneer models, recent followups to AGE.