Learning Disentangled Representation by Exploiting Pretrained Generative Models:  A Contrastive Learning View

1 INTRODUCTION . Disentangled representation learning aims to identify and decompose the underlying explanatory factors hidden in the observed data , which is believed by many to be the only way to understand the world for AI fundamentally ( Bengio & LeCun , 2007 ) . To achieve the goal , as shown in Figure 1 ( a ) , we need an encoder and a generator . The encoder to extract representations from images with each dimension corresponds to one factor individually . The generator ( decoder ) decodes the changing of each factor into different kinds of image variations . With supervision , we can constrain each dimension of the representation only sensitive to one kind of image variation caused by changing one factor respectively . However , this kind of exhaustive supervision is often not available in real-world data . The typical unsupervised methods are based on a generative model to build the above encoder and generator framework , e.g. , VAE ( Kingma & Welling , 2014 ) provides encoder and generator , and GAN ( Goodfellow et al. , 2014 ; Miyato et al. , 2018 ; Karras et al. , 2019 ) provides generator . During the training process of the encoder and generator , to achieve disentangled representation , the typical methods rely on an additional disentanglement regularization term , e.g. , the total correlation for VAE-based methods ( Higgins et al. , 2017 ; Burgess et al. , 2018 ; Kumar et al. , 2017 ; Kim & Mnih , 2018 ; Chen et al. , 2018 ) or mutual information for InfoGAN-based methods ( Chen et al. , 2016 ; Lin et al. , 2020 ) . However , the extra terms usually result in a trade-off between disentanglement and generation quality ( Burgess et al. , 2018 ; Khrulkov et al. , 2021 ) . Furthermore , those unsupervised methods have been proved to have an infinite number of entangled solutions without introducing inductive bias ( Locatello et al. , 2019 ) . Recent works ( Shen & Zhou , 2021 ; Khrulkov et al. , 2021 ; Karras et al. , 2019 ; Härkönen et al. , 2020 ; Voynov & Babenko , 2020 ) show that , for GANs purely trained for image generation , traversing along different directions in the latent space causes different variations of the generated image . This phenomenon indicates that there is some disentanglement property embedded in the latent space of the pretrained GAN . The above observations indicate that training the encoder and generator simultaneous may not be the best choice . We provide an alternative route to learn disentangled representation : fix the pretrained generator , jointly discover the factors in the latent space of the generator and train the encoder to extract disentangled representation , as shown in Figure 1 ( b ) . From the intuitive notion of disentangled representation , similar image variations should be caused by changing the same factor , and different image variations should be caused by changing different factors . This provide a novel contrastive learning view for disentangled representation learning and inspires us to propose a framework : Disentanglement via Contrast ( DisCo ) for disentangled representation learning . In DisCo , changing a factor is implemented by traversing one discovered direction in the latent space . For discovering the factors , DisCo adopts a typical network module , Navigator , to provides candidate traversal directions in the latent space ( Voynov & Babenko , 2020 ; Jahanian et al. , 2020 ; Shen et al. , 2020 ) . For disentangled representation learning , to model the various image variations , we propose a novel ∆-Contrastor to build a Variation Space where we apply the contrastive loss . In addition to the above architecture innovations , we propose two key techniques for DisCo : ( i ) an entropy-based domination loss to encourage the encoded representations to be more disentangled , ( ii ) a hard negatives flipping strategy for better optimization of Contrastive Loss . We evaluate DisCo on three major generative models ( GAN , VAE , and Flow ) on three popular disentanglement datasets . DisCo achieves the state-of-the-art ( SOTA ) disentanglement performance compared to all the previous discovering-based methods and typical ( VAE/InfoGAN-based ) methods . Furthermore , we evaluate DisCo on the real-world dataset FFHQ ( Karras et al. , 2019 ) to demonstrate that it can discover SOTA disentangled directions in the latent space of pretrained generative models . Our main contributions can be summarized as : ( i ) To our best knowledge , DisCo is the first unified framework for jointly learning disentangled representation and discovering the latent space of pretrained generative models by contrasting the image variations . ( ii ) We propose a novel ∆-Contrastor to model image variations based on the disentangled representations for utilizing Contrastive Learning . ( iii ) DisCo is an unsupervised and model-agnostic method that endows non-disentangled VAE , GAN , or Flow models with the SOTA disentangled representation learning and latent space discovering . ( iv ) We propose two key techniques for DisCo : an entropy-based domination loss and a hard negatives flipping strategy . 2 RELATED WORK . Typical unsupervised disentanglement . There have been a lot of studies on unsupervised disentangled representation learning based on VAE ( Higgins et al. , 2017 ; Burgess et al. , 2018 ; Kumar et al. , 2017 ; Kim & Mnih , 2018 ; Chen et al. , 2018 ) or InfoGAN ( Chen et al. , 2016 ; Lin et al. , 2020 ) . These methods achieve disentanglement via an extra regularization , which often sacrifices the generation quality ( Burgess et al. , 2018 ; Khrulkov et al. , 2021 ) . VAE-based methods disentangle the variations by factorizing aggregated posterior , and InfoGAN-based methods maximize the mutual information between latent factors and related observations . VAE-based methods achieve relatively good disentanglement performance but have low-quality generation . InfoGAN-based methods have a relatively high quality of generation but poor disentanglement performance . Our method supplements generative models pretrained without disentanglement regularization term with contrastive learning in the Variation Space to achieve both high-fidelity image generation and SOTA disentanglement . Interpretable directions in the latent space . Recently , researchers have been interested in discovering the interpretable directions in the latent space of generative models without supervision , especially for GAN ( Goodfellow et al. , 2014 ; Miyato et al. , 2018 ; Karras et al. , 2020 ) . Based on the fact that the GAN latent space often possesses semantically meaningful directions ( Radford et al. , 2015 ; Shen et al. , 2020 ; Jahanian et al. , 2020 ) , Voynov & Babenko ( 2020 ) propose a regression-based method to explore interpretable directions in the latent space of a pretrained GAN . The subsequent works focus on extracting the directions from a specific layer of GANs . Härkönen et al . ( 2020 ) search for important and meaningful directions by performing PCA in the style space of StyleGAN ( Karras et al. , 2019 ; 2020 ) . Shen & Zhou ( 2021 ) propose to use the singular vectors of the first layer of a generator as the interpretable directions , and Khrulkov et al . ( 2021 ) extend this method to the intermediate layers by Jacobian matrix . All the above methods only discover the interpretable directions in the latent space , except for Khrulkov et al . ( 2021 ) which also learns disentangled representation of generated images by training an extra encoder in an extra stage . However , all these methods can not outperform the typical disentanglement methods . Our method is the first to jointly learn the disentangled representation and discover the directions in the latent spaces . Contrastive Learning . Contrastive Learning gains popularity due to its effectiveness in representation learning ( He et al. , 2020 ; Grill et al. , 2020 ; van den Oord et al. , 2018 ; Hénaff , 2020 ; Li et al. , 2020 ; Chen et al. , 2020 ) . Typically , contrastive approaches bring representations of different views of the same image ( positive pairs ) closer , and push representations of views from different images ( negative pairs ) apart using instance-level classification with Contrastive Loss . Recently , Contrastive Learning is extended to various tasks , such as image translation ( Liu et al. , 2021 ; Park et al. , 2020 ) and controllable generation ( Deng et al. , 2020 ) . In this work , we focus on the variations of representations and achieve SOTA disentanglement with Contrastive Learning in the Variation Space . Contrastive Learning is suitable for disentanglement due to : ( i ) the actual number of disentangled directions is usually unknown , which is similar to Contrastive Learning for retrieval ( Le-Khac et al. , 2020 ) , ( ii ) it works in the representation space directly without any extra layers for classification or regression . 3 DISENTANGLEMENT VIA CONTRAST . 3.1 OVERVIEW OF DISCO From the contrastive view of the intuitive notion of disentangled representation learning , we propose a DisCo to leverage pretrained generative models to jointly discover the factors embedded as directions in the latent space of the generative models and learn to extract disentangled representation . The benefits of leveraging a pretrained generative model are two-fold : ( i ) the pretrained models with high-quality image generation are readily available , which is important for reflecting detailed image variations and downstream tasks like controllable generation ; ( ii ) the factors are embedded in the pretrained model , severing as an inductive bias for unsupervised disentangled representation learning . DisCo consists of a Navigator to provides candidate traversal directions in the latent space and a ∆-Contrastor to extract the representation of image variations and build a Variation Space based on the target disentangled representations . More specifically , ∆-Contrastor is composed of two sharedweight Disentangling Encoders . The variation between two images is modeled as the difference of their corresponding encoded representations extracted by the Disentangling Encoders . In the Variation Space , by pulling together the variation samples resulted from traversing the same direction and pushing away the ones resulted from traversing different directions , the Navigator learns to discover disentangled directions as factors , and Disentangling Encoder learns to extract disentangled representations from images . Thus , traversing along the discovered directions causes distinct image variations , which causes separated dimensions of disentangled representations respond . Different from VAE-based or InfoGAN-based methods , our disentangled representations and factors are in two separate spaces , which actually does not affect the applications . Similar to the typical methods , the Disentangling Encoder can extract disentangled representations from images , and the pretrained generative model with discovered factors can be applied to controllable generation . Moreover , DisCo can be applied to different types of generative models . Here we provide a detailed workflow of DisCo . As Figure 2 shows , given a pretrained generative model G : Z → I , where Z ∈ RL denotes the latent space , and I denotes the image space , the workflow is : 1 ) A Navigator A provides a total of D candidate traversal directions in the latent space Z , e.g. , in the linear case , A ∈ RL×D is a learnable matrix , and each column is regarded as a candidate direction . 2 ) Image pairs G ( z ) , G ( z′ ) are generated . z is sampled from Z and z′ = z + A ( d , ε ) , where d ∈ { 1 , ... , D } and ε ∈ R , and A ( d , ε ) denotes the shift along the dth direction with ε scalar . 3 ) The ∆-Contrastor , composed of two shared-weight Disentangling Encoders E , encodes the image pair to a sample v ∈ V as v ( z , d , ε ) = |E ( G ( z + A ( d , ε ) ) ) −E ( G ( z ) ) | , ( 1 ) where V ∈ RJ+ denotes the Variation Space . Then we apply Contrastive Learning in V to optimize the Disentangling Encoder E to extract disentangled representations and simultaneously enable Navigator A to find the disentangled directions in the latent space Z . 3.2 DESIGN OF DISCO We present the design details of DisCo , which include : ( i ) the collection of query set Q = { qi } Bi=1 , positive key set K+ = { k+i } Ni=1 and negative key set K− = { k − i } Mi=1 , which are three subsets of the Variation Space V , ( ii ) the formulation of the Contrastive Loss . According to our goal of contrasting the variations , the samples from Q and K+ share the same traversal direction and should be pulled together , while the samples from Q and K− have different directions and should be pushed away . Recall that each sample v in V is determined as v ( z , d , ε ) . To achieve the contrastive learning process , we construct the query sample qi = v ( zi , di , εi ) , the key sample k+i = v ( z + i , d + i , ε + i ) and the negative sample k − i = v ( z − i , d − i , ε − i ) . Specifically , we randomly sample a direction index d̂ from a discrete uniform distribution U { 1 , D } for { di } Bi=1 and { d+i } Ni=1 to guarantee they are the same . We randomly sample { d − i } Mi=1 from the set of the rest of the directions U { 1 , D } \ { d̂ } individually and independently to cover the rest of directions in Navigator A . Note that the discovered direction should be independent with the starting point and the scale of variation , which is in line with the disentangled factors . Therefore , { zi } Bi=1 , { z + i } Ni=1 , { z − i } Mi=1 are all sampled from latent space Z , and { εi } Bi=1 , { ε + i } Ni=1 , { ε − i } Mi=1 are all sampled from a shared continuous uniform distribution U [ − , ] individually and independently . We normalize each sample in Q , K+ , and K− to a unit vector to eliminate the impact caused by different shift scalars . For the design of Contrastive Loss , a well-known form of Contrastive Loss is InfoNCE ( van den Oord et al. , 2018 ) : LNCE = − 1 |B| B∑ i=1 N∑ j=1 log exp ( qi · k+j /τ ) ∑N+M s=1 exp ( qi · ks/τ ) , ( 2 ) where τ is a temperature hyper-parameter and { ki } N+Mi=1 = { k + i } Ni=1 ⋃ { k−i } Mi=1 . The InfoNCE is originate from BCELoss ( Gutmann & Hyvärinen , 2010 ) . BCELoss has been used to achieve contrastive learning ( Wu et al. , 2018 ; Le-Khac et al. , 2020 ; Mnih & Kavukcuoglu , 2013 ; Mnih & Teh , 2012 ) . We choose to follow them to use BCELoss Llogits for reducing computational cost : Llogits = − 1 |B| B∑ i=1 ( l−i + l + i ) , ( 3 ) l+i = N∑ j=1 log σ ( qi · k+j /τ ) , l − i = M∑ m=1 log ( 1− σ ( qi · k−m/τ ) ) , ( 4 ) where σ denotes the sigmoid function , l+i denotes the part for positive samples , and l − i denotes the part for the negative ones . Note that we use a shared positive set for B different queries to reduce the computational cost . 3.3 KEY TECHNIQUES FOR DISCO Entropy-based domination loss . By optimizing the Contrastive Loss , Navigator A is optimized to find the disentangled directions in the latent space , and Disentangling Encoder E is optimized to extract disentangled representations from images . To further make the encoded representations more disentangled , i.e. , when traversing along one disentangled direction , only one dimension of the encoded representation should respond , we thus propose an entropy-based domination loss to encourage the corresponding samples in the Variation Space to be one-hot . To implement the entropy-based domination loss , we first get the mean c of Q and K+ as c = 1 |B +N | ( B∑ i=1 qi + N∑ i=1 k+i ) . ( 5 ) We then compute the probability as pi = exp c ( i ) / ∑J j=1 exp c ( j ) , where c ( i ) is the i-th element of c and J is the number of dimensions of c. The entropy-based domination loss Led is calculated as Led = − 1 J J∑ j=1 pj log ( pj ) . ( 6 ) Hard negatives flipping . Since the latent space of the generative models is a high-dimension complex manifold , many different directions carry the same semantic meaning . These directions with the same semantic meaning result in hard negatives during the optimization of Contrastive Loss . The hard negatives here are different from the hard negatives in the works of self-supervised representation learning ( He et al. , 2020 ; Coskun et al. , 2018 ) , where they have reliable annotations of the samples . Here , our hard negatives are more likely to be “ false ” negatives , and we choose to flip these hard negatives into positives . Specifically , we use a threshold T to identify the hard negative samples , and use their similarity to the queries as the pseudo-labels for them : l̂−i = ∑ αij < T log ( 1− σ ( αij ) ) + ∑ αij≥T αij log ( σ ( αij ) ) , ( 7 ) where l̂−i denotes the modified l − i , and αij = qi · k − j /τ . Therefore , the modified final BCELoss is : Llogits−f = 1 |B| B∑ i=1 ( l+i + l̂ − i ) . ( 8 ) Full objective . With the above two techniques , the full objective is : L = Llogits−f + λLed , ( 9 ) where λ is the weighting hyper-parameter for entropy-based domination loss Led .

Summary: This paper presents a framework to model disentangled directions for pretrained models. Such an approach mitigates the problems with poor generation quality arising while training models with additional regularization terms to force disentanglement. The underlying idea is contrastive-based: similar image variations are caused by changing the same factors in contrast to the remaining image variations. The proposed framework is model-agnostic: it can be applied to GANs, VAEs and flow models.