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: The paper proposes a novel representation learning technique to disentangle the latent space of pre-trained generative models, by discovering semantically meaningful directions in them.  The method trains a navigator and a delta-contrastor network, which consists of 2 encoders sharing weights. First, random samples are perturbed along the directions obtained from the navigator. The perturbed vectors are then decoded with the pre-trained generator, then encoded and the difference between 2 samples are taken. The output is in the variation space, where a contrastive learning technique clusters together the samples that were perturbed with the same direction.