Self-Supervised Variational Auto-Encoders

Density estimation , compression , and data generation are crucial tasks in artificial intelligence . Variational Auto-Encoders ( VAEs ) constitute a single framework to achieve these goals . Here , we present a novel class of generative models , called self-supervised Variational Auto-Encoder ( selfVAE ) , that utilizes deterministic and discrete transformations of data . This class of models allows performing both conditional and unconditional sampling while simplifying the objective function . First , we use a single self-supervised transformation as a latent variable , where a transformation is either downscaling or edge detection . Next , we consider a hierarchical architecture , i.e. , multiple transformations , and we show its benefits compared to the VAE . The flexibility of selfVAE in data reconstruction finds a particularly interesting use case in data compression tasks , where we can trade-off memory for better data quality , and vice-versa . We present the performance of our approach on three benchmark image data ( Cifar10 , Imagenette64 , and CelebA ) . 1 INTRODUCTION . The framework of variational autoencoders ( VAEs ) provides a principled approach for learning latentvariable models . As it utilizes a meaningful low-dimensional latent space with density estimation capabilities , it forms an attractive solution for generative modelling tasks . However , its performance in terms of the test log-likelihood and quality of generated samples is often disappointing , thus , many modifications were proposed . In general , one can obtain a tighter lower bound , and , thus , a more powerful and flexible model , by advancing over the following three components : the encoder ( Rezende et al. , 2014 ; van den Berg et al. , 2018 ; Hoogeboom et al. , 2020 ; Maaløe et al. , 2016 ) , the prior ( or marginal over latents ) ( Chen et al. , 2016 ; Habibian et al. , 2019 ; Lavda et al. , 2020 ; Lin & Clark , 2020 ; Tomczak & Welling , 2017 ) and the decoder ( Gulrajani et al. , 2016 ) . Recent studies have shown that by employing deep hierarchical architectures and by carefully designing building blocks of the neural networks , VAEs can successfully model high-dimensional data and reach state-of-the-art test likelihoods ( Zhao et al. , 2017 ; Maaløe et al. , 2019 ; Vahdat & Kautz , 2020 ) . In this work , we present a novel class of VAEs , called self-supervised Variational Auto-Encoders , where we introduce additional variables to VAEs that result from discrete and deterministic transformations of observed images . Since the transformations are deterministic , and they provide a specific aspect of images ( e.g. , contextual information through detecting edges or downscaling ) , we refer to them as self-supervised representations . The introduction of the discrete and deterministic variables allows to train deep hierarchical models efficiently by decomposing the task of learning a highly complex distribution into training smaller and conditional distributions . In this way , the model allows to integrate the prior knowledge about the data , but still enables to synthesize unconditional samples . Furthermore , the discrete and deterministic variables could be used to conditionally reconstruct data , which could be of great use in data compression and super-resolution tasks . We make the following contributions : i ) We propose an extension of the VAE framework by incorporating self-supervised representations of the data . ii ) We analyze the impact of modelling natural images with different data transformations as self-supervised representations . iii ) This new type of generative model ( self-supervised Variational Auto-Encoders ) , which is able to perform both conditional and unconditional sampling , demonstrate improved quantitative performance in terms of density estimation and generative capabilities on image benchmarks . 2 BACKGROUND . 2.1 VARIATIONAL AUTO-ENCODERS . Let x 2 XD be a vector of observable variables , where X ✓ R or X ✓ Z , and z 2 RM denote a vector of latent variables . Since calculating p # ( x ) = R p # ( x , z ) dz is computationally intractable for non-linear stochastic dependencies , a variational family of distributions could be used for approximate inference . Then , the following objective function could be derived , namely , the evidence lower bound ( ELBO ) ( Jordan et al. , 1999 ) : ln p # ( x ) Eq ( z|x ) [ ln p✓ ( x|z ) + ln p ( z ) ln q ( z|x ) ] , ( 1 ) where q ( z|x ) is the variational posterior ( or the encoder ) , p✓ ( x|z ) is the conditional likelihood function ( or the decoder ) and p ( z ) is the prior ( or marginal ) , , ✓ and denote parameters . The expectation is approximated by Monte Carlo sampling while exploiting the reparameterization trick in order to obtain unbiased gradient estimators . The models are parameterized by neural networks . This generative framework is known as Variational Auto-Encoder ( VAE ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . 2.2 VAES WITH BIJECTIVE PRIORS . Even though the lower-bound suggests that the prior plays a crucial role in improving the variational bounds , usually a fixed distribution is used , e.g. , a standard multivariate Gaussian . While being relatively simple and computationally cheap , the fixed prior is known to result in over-regularized models that tend to ignore most of the latent dimensions ( Burda et al. , 2015 ; Hoffman & Johnson , 2016 ; Tomczak & Welling , 2017 ) . Moreover , even with powerful encoders , VAEs may still fail to match the variational posterior to a unit Gaussian prior ( Rosca et al. , 2018 ) . However , it is possible to obtain a rich , multi-modal prior distribution p ( z ) by using a bijective ( or flow-based ) model ( Dinh et al. , 2016 ) . Formally , given a latent code z , a base distribution pV ( v ) over latent variables v 2 RM , and f : RM ! RM consisting of a sequence of L diffeomorphic transformations1 , where fi ( vi 1 ) = vi , v0 = v and vL = z , the change of variable can be used sequentially to express the distribution of z as a function of v as follows : log p ( z ) = log pV ( v ) LX i=1 log @ fi ( vi 1 ) @ vi 1 , ( 2 ) where @ fi ( vi 1 ) @ vi 1 is the Jacobian-determinant of the ith transformation . Thus , using the bijective prior yields the following lower-bound : ln p ( x ) Eq ( z|x ) h log p✓ ( x|z ) log q ( z|x ) + log pV ( v0 ) + LX i=1 log @ f 1i ( vi ) @ vi i . ( 3 ) In this work , we utilize RealNVP ( Dinh et al. , 2016 ) as the prior , however , any other flow-based model could be used ( Kingma & Dhariwal , 2018 ; Hoogeboom et al. , 2020 ) . For the experiments and ablation study that shows the impact of the bijective prior on VAEs , we refer to the appendix A.1 . 3 METHOD . 3.1 MOTIVATION . The idea of self-supervised learning is about utilizing original unlabeled data to create additional context information . It could be achieved in multiple manners , e.g. , by adding noise to data ( Vincent et al. , 2008 ) or masking data during training ( Zhang et al. , 2017 ) . Self-supervised learning could also be seen as turning an unsupervised model into a supervised by , e.g. , treating predicting next pixels as a classification task ( Hénaff et al. , 2019 ; Oord et al. , 2018 ) . These are only a few examples of a quickly growing research line ( Liu et al. , 2020 ) . 1That is , invertible and differentiable transformations . i ) Stochastic dependencies of self-supervised VAE ii ) Hierarchical ssVAE self-supervised VAE Here , we propose to use non-trainable transformations to obtain information about image data . Our main hypothesis is that since working with highly-quality images is challenging , we could alleviate this problem by additionally considering partial information about them . Fitting a model to images of lower quality , and then enhancing them to match the target distribution seems to be overall an easier task ( Chang et al. , 2004 ; Gatopoulos et al. , 2020 ) . By incorporating compressed transformations ( i.e. , the self-supervised representations ) that still contain global information , with the premise that it would be easier to approximate , the process of modelling a high-dimensional complex density breaks down into simpler tasks . In this way , the expressivity of the model will grow and gradually result into richer , better generations . A positive effect of the proposed framework is that the model allows us to integrate prior knowledge through the image transformations , without losing its uncon- ditional generative functionality . Overall , we end up with a two-level VAE with three latent variables , where one is a data transformation that can be obtained in a self-supervised fashion . In Figure 1 a schematic representation of the proposed approach with downscaling is presented . A number of exemplary image transformations are presented in Figure 2 . We notice that with these transformations , even though they discard a lot of information , the global structure is preserved . As a result , in practice the model should have the ability to extract a general concept of the data , and add local information afterwards . In this work , we focus on downscaling ( Figure 2.b , c & d ) and edge detection or sketching ( Fig . 2.i ) . 3.2 MODEL FORMULATION . In our model , we consider representations that result from deterministic and discrete transformations of an image . Formally , we introduce a transformation d : XD ! XC that takes x and returns an image representation y , e.g. , a downscaled image . Since we lose information about the original image , z could be seen as a variable that compensates lost details in x . Further we propose to introduce an additional latent variable , u 2 RN to model y and z . We can define the joint distribution of x and y as follows : p ( x , y ) = p ( y|x ) p ( x ) , where p ( y|x ) = ( y d ( x ) ) due to the deterministic transformation d ( · ) , where ( · ) is the Kronecker delta . Thus , the empirical distribution is ( y d ( x ) ) pdata ( x ) . However , since we are interested in decomposing the problem of modeling a complex distribution p ( x ) , we propose to model p ( x|y ) p ( y ) instead , and utilize the variational inference of the form Q ( u , z|x , y ) = q ( u|y ) q ( z|x ) that yields : ln p ( x , y ) EQ ⇥ ln p✓ ( x|y , z ) + ln p ( z|u , y ) + ln p ( y|u ) + ln p ( u ) ln q ( z|x ) ln q ( u|y ) ⇤ . ( 4 ) Intuitively , the premise for selfVAE is that the latents u will capture the global structure of the input data and the latents z will encode the missing information between y and x , guiding the model to discover the distribution of the target observations . In order to highlight the self-supervised part in our model , we refer to it as the self-supervised Variational Auto-Encoder ( or selfVAE for short ) . Further , we propose to choose the following distributions : p ( v ) = N ( v|0,1 ) p ( u ) = p ( v ) FY i=1 det @ fi ( vi 1 ) @ vi 1 1 p✓1 ( y|u ) = IX i=1 ⇡ ( u ) i Dlogistic ⇣ µ ( u ) i , s ( u ) i ⌘ q 1 ( u|y ) = N ( u|µ 1 ( y ) , diag ( 1 ( y ) ) ) q 2 ( z|x ) = N ( z|µ 2 ( x ) , diag ( 2 ( x ) ) ) . p✓2 ( z|y , u ) = N ( z|µ✓2 ( y , u ) , diag ( ✓2 ( y , u ) ) ) p✓3 ( x|z , y ) = IX i=1 ⇡ ( z , y ) i Dlogistic ⇣ µ ( z , y ) i , s ( z , y ) i ⌘ where Dlogistic is defined as the discretized logistic distribution ( Salimans et al. , 2017 ) , and we utilize a flow-based model for p ( u ) . Notice that we use the discretized logistic distribution , because images are represented by values between 0 and 255 . For integer-valued random variables , other distributions like Gaussian are inappropriate .

Summary: $\bullet$ VAEs can ignore some dimensions of the latent code. Enforcing the posterior distributions to consider desired factors of variations in the input can be fulfilled by either making it more structured (i.e., quantization as in VQ-VAE-2) or introducing additional constraints. This paper tackles this problem by applying the latter, two self-supervised tasks: edge maps and downscaled versions of inputs.