The Variational InfoMax AutoEncoder

1 INTRODUCTION . A common assumption in machine learning is that any visible data x ∈ X is completely described by some generative factor o , living in a smaller hidden space O , i.e . x = g ( o ) with g a ( possibly stochastic ) generative function . The aim of unsupervised representation learning research is to find a representation z of the generative factor o living in a known space Z describing , as well as o , the visible data x . This is particularly relevant because the learnt small representation z is task agnostic and , in principle , can be used as input for networks performing different tasks , leading to faster and more robust learning ( generalisation property ) , ( Rifai et al. , 2011 ) . Many models fφ : X → Z trying to learn such representations have been proposed ( Dinh et al. , 2016 ; Hinton et al. , 2006 ; Maddison et al. , 2017 ; Radford et al. , 2015 ) , but recently in order to solve this problem it was proposed to consider a dual problem : define a priori z and find a generator map gθ , such that for any z , gθ ( z ) is an element of X . In particular , two families of probabilistic generative models have become dominant : Variational AutoEncoder ( VAE ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) and Generative Adversarial Network ( GAN ) ( Goodfellow et al. , 2014 ) . The common idea of the two approaches is that a good generator pθ ( x|z ) is the one able to generate the data that is as close as possible to the visible one , i.e . that with respect a certain metric D , the distance between the marginal pθ ( x ) = Ep ( z ) [ pθ ( x|z ) ] and the visible distribution pD ( x ) is minimal . In this manuscript we restrict our attention to the VAE model , since by its architecture , it is the only one where the learnt representation , possibly from different datasets , can be used as input for networks performing different tasks , ( Achille et al. , 2018 ; Ramapuram et al. , 2017 ) . Although VAE , by its training robustness and general good generative performance is the most popular model for representation learning , in particular cases it suffers from the uninformative representation issue : the representation is entangled and the generative model tends to be independent of z , i.e . pθ ( x|z ) ≈ pθ ( x ) . As highlighted in the next section such behaviour is intrinsic in the variational loss , the Evidence Lower BOund ( ELBO ) , encouraging a less informative representation . Following the direction suggested in ( Alemi et al. , 2017 ; Zhao et al. , 2017 ) we describe the VAE from an information theoretic perspective . Such description lead us to two observations : the WAE ( Tolstikhin et al. , 2017 ) and InfoVAE ( Zhao et al. , 2017 ) models are actually maximising a lower bound of the mutual information associated to a generator pθ ( x|z ) belonging in a certain family P , the Variational InfoMax ( VIM ) ; and given that the capacity of the network is a function of the entropy of the prior p ( z ) , the VIM is the variational expression of the Capacity-Constrained InfoMax ( CCIM ) . Following the analogies between the CCIM and the Information Bottleneck ( IB ) , the theoretical principle associated with the ELBO , we deduce that the representation quality is associated to the capacity term . The theoretical arguments are confirmed by the performed experiments where we observe that , differently from what was argued in previous works ( Higgins et al. , 2017 ; Burgess et al. , 2018 ) , it is possible to train a model that is able to learn good ( able to generalise ) representations while maintaining optimal generative performance . The main contributions of the paper are summarised in the following points : • derivation of a variational lower bound for the maximal mutual information of a generaive model belonging in a certain family , see equation 7 ; • definition and bounds estimation for the network capacity for a variational autoencoder , see equation 9 ; • association of the two main properties of VAE , generation quality and good representation , to two different information concepts , respectively Mutual Information and network capacity ; • proposal of a new learning principle for unsupervised models : the Capacity-Constrained InfoMax , see equation 10 , that allows both to learn a good representation while maintaining optimal generative performance . The work is divided as follows : in the second section we describe briefly the VAE and its variants ; in the third and fourth sections we describe the variational infomax method and related work . We conclude the paper with the experimental results and the final observations . 2 BACKGROUND . The aim of this section is to describe VAE , understand principal issues of the ELBO objective and describe the three most relevant approaches to overcome such issues . 2.1 NOTATION AND PRELIMINARY DEFINITIONS . We use calligraphic letters ( i.e . X ) for sets , capital letters ( i.e . X ) for random variables , and lower case letters ( i.e . x ) for their samples . With abuse of notation we denote both the probability and the corresponding density with the lower case letters ( i.e . p ( x ) ) . KL divergence Given two random distributions p ( x ) and q ( x ) , the Kullback-Leibler ( KL ) divergence DKL ( p ( x ) ||q ( x ) ) = ∫ log ( p ( y ) q ( y ) ) p ( y ) dy ( 1 ) is an ( intuitive ) measure of the distance between the distributions p and q . Mutual Information and Capacity Given a channel Z → X with X and Z random variables , jointly distributed according to p ( x , z ) and with marginals p ( x ) and p ( z ) . The mutual information I ( X , Z ) = DKL ( p ( x , z ) ||p ( x ) p ( z ) ) , is a measure of the reduction of uncertainty in X due to the knowledge of Z , and the capacity C ( X , Z ) = sup p ( z ) ∈P I ( X , Z ) is the maximal information that can be shared for a fixed generator p ( x|z ) . 2.2 VARIATIONAL AUTOENCODER . From now on let us assume that the unknown distribution of the data p ( x ) coincides with the empirical one pD ( x ) , and that the distribution of the latent representation p ( z ) is known . In this context the VAE is a model solving the following optimisation problem : find the generative model pθ ( x , z ) ∈ Pθ , specified by the parameters θ of the associated neural network , maximising the ELBO objective ELBOθ , φ = Ep ( x ) [ −DKL ( qφ ( z|x ) ||p ( z ) ) + Eq ( z|x ) [ log p ( x|z ) ] ] , ( 2 ) a lower bound of the unfeasible-to-compute marginal likelihood Ep ( x ) [ log pθ ( x ) ] . The ELBO objective is optimised by a regularised autoencoder , with encoder and decoder parametetrising , respectively , the inference and generative distributions , qφ ( z|x ) and pθ ( x|z ) , with φ ∈ Φ , θ ∈ Θ and regulariser defined by the rate term DKL ( qφ ( z|x ) ||p ( z ) ) , measuring the excess number of bits required to encode samples from the encoder using the optimal code designed for p ( z ) . 2.3 UNINFORMATIVE REPRESENTATION ISSUE . As underlined in the introduction , the main issue of VAE is that the representations are not really informative of the input data and in the worst case , it is learned a Z-independent generative model pθ ( x|z ) = pθ ( x ) . Such issues are intrinsic in the ELBO objective , equation 2 , that can reach the optimum when DKL ( qφ ( z|x ) ||p ( z ) ) = 0 ( Zhao et al. , 2017 ) . The latter case means that the representation is completely uninformative , indeed the rate term , which can be rewritten as DKL ( qφ ( z|x ) ||p ( z ) ) = Iq ( X , Z ) +DKL ( qφ ( z ) ||p ( z ) ) , is a penalty on the encoding information , and is zero when Iq ( X , Z ) = 0 , with qφ ( z|x ) = qφ ( z ) = p ( z ) , i.e . when qφ does not encode any information about the input x . We now describe the three most relevant models that try to overcome the uninformative representation issue . InfoVAE In ( Zhao et al. , 2017 ) the InfoVAE family of models was proposed , a generalisation of the VAE model optimising the objective −αIq ( X , Z ) − λDKL ( qφ ( z ) ||p ( z ) ) + Ep ( x ) [ Eq ( z|x ) [ log p ( x|z ) ] ] , with α and λ two real positive hyper-parameters . The main advantage of this definition is that it is possible to consider separately the two components of the rate term . In particular , in ( Zhao et al. , 2017 ) it was observed that by eliminating the information penalty ( α = 0 ) , the generative performance of the model improves and the representation results are more informative . β-VAE In ( Higgins et al. , 2017 ) , starting from the observation that the optimal case is rare , but most of the learned features by VAE are not disentangled , it is proposed an opposite approach : put a high penalty to the rate term , in order to constrain the model to learn the most informative property of the data , and then have a disentangled representation of the data . The β-VAE family is a particular case of InfoVAE where α = λ 1 . This idea , that at first sight looks counter-intuitive , is based on the observation that by the additive property of the KL-divergence DKL ( qφ ( z|x ) ||p ( z ) ) = dim ( Z ) ∑ i=1 DKL ( qφ ( zi|x ) ||p ( zi ) ) ( 3 ) pushing the penalty associated with the rate is equivalent to penalising the informativeness of most features , leaving few features containing the relevant information . Starting from a bits-back coding argument , a similar conclusion was derived in ( Chen et al. , 2016 ) . Minimal rate bound In order to avoid the null rate issue , and maintain semantic informative representation , in ( Alemi et al. , 2017 ) was suggested to optimise the following variant of ELBO : Ep ( x ) [ Eq ( z|x ) [ log p ( x|z ) ] ] − β|R0 −DKL ( qφ ( z|x ) ||p ( z ) ) | , ( 4 ) where R0 > 0 , defines both the lower bound of the rate and the upper bound of the encoding information Iq ( X , Z ) . If , as we will see below , the objective in equation 4 is a generalisation of both InfoVAE ( R0 = maxθ Iθ ( X , Z ) ) and β-VAE ( R0 = 0 ) , the correct R0 is unknown and its optimisitation is not an easy task .

Summary: The paper develops an information-theoretic training scheme for Variational Auto-Encoders (VAEs). This scheme is tailored for addressing the well-known disentanglement problem of VAEs where an over-capacity encoder sometimes manages to both maximize data fit and shrink the KL divergence between the approximate posterior and prior to zero. Consequently, the latent representations of the observations become independent, making them unusable for any downstream task.