IB-GAN: Disentangled Representation Learning with Information Bottleneck GAN

1 INTRODUCTION . Learning good representations for data is one of the essential topics in machine learning community . Although any strict definition for it may not exist , the consensus about the useful properties of good representations has been discussed throughout many studies ( Bengio et al. , 2013 ; Lake et al. , 2017 ; Higgins et al. , 2018 ) . A disentanglement , one of those useful properties of representation , is often described as a statistical independence or factorization ; each independent factor is expected to be semantically well aligned with the human intuition on the data generative factor ( e.g . a chair-type from azimuth on Chairs dataset ( Aubry et al. , 2014 ) , or age from azimuth on CelebA dataset ( Liu et al. , 2015 ) ) . The learned representation distilling each important factors of data into a single independent direction is hard to be done but highly valuable for many other downstream tasks ( Ridgeway , 2016 ; Higgins et al. , 2017b ; 2018 ) . Many models have been proposed for disentangled representation learning ( Hinton et al. , 2011 ; Kingma et al. , 2014 ; Reed et al. , 2014 ; Narayanaswamy et al. , 2017 ; Denton et al. , 2017 ) . Despite their impressive results , they either require knowledge of ground-truth generative factors or weak-supervision ( e.g . domain knowledge or partial labels ) . In contrast , among many unsupervised approaches ( Desjardins et al. , 2012 ; Kingma & Welling , 2013 ; Rezende et al. , 2014 ; Springenberg , 2015 ; Dumoulin et al. , 2017 ) , yet the two most successful approaches for the independent factor learning are β-VAE ( Higgins et al. , 2017a ) and InfoGAN ( Chen et al. , 2016 ) . Higgins et al . ( 2017a ) demonstrate that encouraging the KL-divergence term of Variational autoencoder ( VAE ) objective ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) by multiplying a constant β > 1 induces a high-quality disentanglement of latent factors . As follow-up research , Burgess et al . ( 2018 ) provide a theoretical justification of the disentangling effect of β-VAE in the context of Information Bottleneck theory ( Tishby et al. , 1999 ; Tishby & Zaslavsky , 2015 ) Chen et al . ( 2016 ) propose another fully unsupervised approach based on Generative Adversarial Network ( GAN ) ( Goodfellow et al. , 2014 ) . He achieves the goal by enforcing the generator to learn disentangled representations through increasing the mutual information ( MI ) between the generated samples and the latent representations . Although InfoGAN can learn to disentangle representations for relatively simple datasets ( e.g . MNIST , 3D Chairs ) , it struggles to do so on more complicated datasets such as CelebA . Moreover , the disentangling performance of the learned representations from InfoGAN is known as not good as the performance of the β-VAE and its variant models ( Higgins et al. , 2017a ; Kim & Mnih , 2018 ; Chen et al. , 2018 ) . Stimulated by the success of β-VAE models ( Burgess et al. , 2018 ; Higgins et al. , 2017a ; Kim & Mnih , 2018 ; Chen et al. , 2018 ; Esmaeili et al. , 2018 ) with the Information Bottleneck theory ( Alemi et al. , 2017 ; 2018 ; Achille & Soatto , 2017 ) in disentangled representations learning task , we hypothesize that the weakness of InfoGAN in the representation learning may originate from that it can only maximize the mutual information but lacks any constraining mechanisms . In other words , InfoGAN misses the term upper-bounding the mutual information from the perspective of IB theory . We present a novel unsupervised model named IB-GAN ( Information Bottleneck GAN ) for learning disentangled representations based on IB theory . We propose a new architecture of GANs from IB theory so that the training objective involves an information capacity constraint that InfoGAN lacks but β-VAE has . We also derive a new variational approximation algorithm to optimize IBGAN objective in practice . Thanks to the information regularizer , the generator can use the latent representations in a manner that is both more interpretable and disentangled than InfoGAN The contributions of this work are summarized as follows : 1 . IB-GAN is a new GAN-based model for fully unsupervised learning of disentangled representations . To the best of our knowledge , there is no other unsupervised GAN-based model for this sake except the InfoGAN ’ s variants ( Higgins et al. , 2017a ; Kim & Mnih , 2018 ) . 2 . Our work is the first attempt to utilize the IB theory into the GAN-based deep generative model . IB-GAN can be seen as an extension to the InfoGAN , supplementing an information constraining regularizer that InfoGAN misses . 3 . IB-GAN surpasses state-of-the-art disentanglement scores of ( Higgins et al. , 2017a ; Eastwood & Williams , 2018 ) on dSprites dataset ( Matthey et al. , 2017 ) . The quality of generated samples by IB-GAN on 3D Chairs ( Aubry et al. , 2014 ) and CelebA ( Liu et al. , 2015 ) is also much realistic compared to that of the existing β-VAE variants of the same task . 2 PRELIMINARIES . We remind some backgrounds : IB principle in section 2.1 and the connection between β-VAE and IB theory ( Burgess et al. , 2018 ) in section 2.2 . Lastly , InfoGAN ( Chen et al. , 2016 ) is briefly reviewed in section 2.3 . 2.1 INFORMATION BOTTLENECK PRINCIPLE . Let the input variable X and the target variable Y distributed according to some joint data distribution p ( x , y ) . The goal of the IB ( Tishby et al. , 1999 ; Alemi et al. , 2017 ; 2018 ) is to obtain a compressive representation Z from the input variable X , while maintaining the predictive information about the target variable Y as much as possible . The objective for the IB is max qφ ( z|x ) LIB = I ( Z , Y ) − βI ( Z , X ) ( 1 ) where I ( · , · ) denotes MI and β ≥ 0 is a Lagrange multiplier . The goal is to obtain the optimal representation encoder qφ ( z|x ) that balances the trade-off between the maximization and minimization of both MI terms . Hence , the IB objective in Eq . ( 1 ) provides a natural means for good representations by enforcing the representation Z to ignore irrelevant information from the input and simultaneously to be predictive about the target , which can act as a minimal sufficient statistic of X for predicting Y ( Tishby et al. , 1999 ; Alemi et al. , 2017 ; 2018 ) . A growing body of studies ( Alemi et al. , 2017 ; Achille & Soatto , 2018 ; 2017 ) supports that the learned representations adapting the IB objective tend to be highly efficient and distilled in terms of its code length ( Alemi et al. , 2018 ; Shannon et al. , 1951 ) . As a consequence , the learned representation is more generalizable and robust to adversarial attack ( Alemi et al. , 2017 ) , disentangled ( Burgess et al. , 2018 ) and invariant to nuance factors ( Achille & Soatto , 2017 ) . Moreover , the IB framework prevents weight over-fitting ( Alemi et al. , 2018 ; Achille & Soatto , 2018 ; Vera et al. , 2017 ) , and can be used to visualize high dimensional embedding in a low dimensional latent space ( Rabinowitz et al. , 2018 ) . 2.2 β-VAE FROM THE IB THEORY β-VAE ( Higgins et al. , 2017a ) is one of the state-of-the-art unsupervised disentangled representation learning models . The key idea of β-VAE is to multiply a constant β ≥ 1 to the KL-divergence term of the original VAE ’ s objective ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) : max pθ , qφ Lβ-VAE = Ep ( x ) [ Eqφ ( z|x ) [ log pθ ( x|z ) ] − βKL ( qφ ( z|x ) ||p ( z ) ) ] , ( 2 ) where the encoder qφ ( z|x ) is generally known as the variational approximation to the intractable p ( z|x ) , p ( z ) is a prior for the latent representation and pθ ( x|z ) is the decoder in the VAE context . Recently , a notable connection between β-VAE and the IB theory has been discovered in ( Alemi et al. , 2017 ; 2018 ) . Eq . ( 2 ) can be derived from the variational approximation to the IB objective Eq. ( 1 ) . To clarify this connection , see the variational upper and lower bound of the MI : Ep ( x ) [ Eqφ ( z|x ) [ log pθ ( x|z ) ] +H ( x ) ≤ Iq ( Z , X ) ≤ Ep ( x ) [ KL ( qφ ( z|x ) ||p ( z ) ) ] . ( 3 ) The MI in Eq . ( 3 ) subscribed with q : Iq ( Z , X ) = Eqφ ( z|x ) p ( x ) [ qφ ( z|x ) p ( x ) /qφ ( z ) p ( x ) ] , is called as the representational MI1 . Given that computing marginal qφ ( z ) is intractable , we can use any prior p ( z ) to substitute for qφ ( z ) , forming the variational upper-bound in Eq. ( 3 ) 2 . Likewise , we can use any decoder model pθ ( x|z ) to approximate qφ ( x|z ) = qφ ( z|x ) p ( x ) /qφ ( z ) of the MI , forming the variational lower-bound in Eq. ( 3 ) . If the target variable Y in Eq . ( 1 ) is replaced with X , the task is to reconstruct ( auto-encode ) data from the representation Z . The variational lower-bound of Eq . ( 1 ) obtained by leveraging the upper and lower bound of MI in Eq . ( 3 ) corresponds to Eq. ( 2 ) 3 . The disentanglement-promoting behavior of the β-VAE based on IB theory is discussed in ( Burgess et al. , 2018 ) . Constraining the MI ( or minimizing KL-divergence in practice ) forces the encoder to learn representation containing only strongly relevant information to the data reconstruction , while ignoring other unnecessary ( or less-necessary ) features . The encoder becomes reluctant to use more channels ( or dimensions ) of the latent vector to lower the MI constraining cost . Hence , the most distinctive and principle features of data are grouped and aligned along with each independent dimension of the representation space . 2.3 INFOGAN : INFORMATION MAXIMIZING GAN . Generative Adversarial Networks ( GAN ) ( Goodfellow et al. , 2014 ) establish a min-max adversarial game between two neural networks , a generator G and a discriminator D. The discriminator D aims to distinguish well between real sample x ∼ p ( x ) and synthetic sample created by the G ( z ) with a random noise z ∼ p ( z ) , while the generator G is trained to produce a realistic sample that is indistinguishable from the true sample . The adversarial game is formulated as follow : min G max D LGAN ( D , G ) = Ep ( x ) [ log ( D ( x ) ) ] + Ep ( z ) [ log ( 1−D ( G ( z ) ) ] . ( 4 ) Under an optimal discriminator D∗ , Eq . ( 4 ) theoretically involves with the Jensen-Shannon divergence between the synthetic and the true sample distribution : JS ( G ( z ) ||p ( x ) ) . However , Eq . ( 4 ) does not have any specific guidance on how G utilizes a mapping from z to x . That is , the variation of z in any independent dimension often yields entangled effects on a generated sample x . On the other hand , InfoGAN ( Chen et al. , 2016 ) is capable of learning disentangled representations . InfoGAN introduces an additional latent code c and encourages it to describe the semantic features of the data . To do so , the training objective of InfoGAN accommodates a mutual information maximization term between the latent code c and the generated sample x = G ( z , c ) : max G min D LInfoGAN ( D , G ) = −LGAN ( D , G ) + λI ( c , G ( z , c ) ) , ( 5 ) where I ( · , · ) denote MI and λ is a weight coefficient . To optimize Eq . ( 5 ) , the variational lower bound of MI is also exploited similar to that of the IM algorithm ( Barber & Agakov , 2003 ) . 1We distinguish it from the generative in the next section . 2The variational inference relies on the positivity of the KL divergence : Ep ( · ) [ log p ( · ) ] ≥ Ep ( · ) [ log q ( · ) ] for any variational ( or approximating ) distribution q ( · ) ( Jordan et al. , 1999 ; Wainwright & Jordan , 2008 ; Alemi & Fischer , 2018 ) . 3A constant data entropy term H ( X ) = −Ep ( x ) [ log p ( x ) ] is ignored for brevity .

Summary: This work addresses the problem of unsupervised disentangled representation learning, and leverages insights and intuitions about utilizing an information bottleneck (IB) approach to encourage disentangling. In particular, building upon insights of how beta-VAE can be understood (and improved upon) by understanding it in terms of IB, the authors propose to modify GANs to include an IB, so as to leverage similar disentangling benefits. The promise is that this approach could utilise the strengths of the GAN framework over VAEs, such as the often sharper reconstructions and the ease of including discrete latents in addition to continuous ones.