Representation Disentanglement in Generative Models with Contrastive Learning

1 INTRODUCTION . In recent years , disentanglement of factors in images has attracted many researchers ’ attention , which mainly includes two folds : adversarial and non-adversarial methods . Adversarial methods ( Mathieu et al. , 2016 ; Denton & Birodkar , 2017 ; Hadad et al. , 2018 ; Razavi et al. , 2019 ; Gabbay & Hoshen , 2021 ) often apply a min-max optimization framework ( Goodfellow et al. , 2014 ) for disentanglement of images , which costs much time on hyper-parameters tuning . In terms of non-adversarial models , several variational autoencoders ( Higgins et al. , 2017 ; Kim & Mnih , 2018 ) variants have been proposed to disentangle the generative factors in an unsupervised manner without inductive biases , which did not achieve satisfactory results as proven in an empirical study ( Locatello et al. , 2019 ) . With the extra class supervision , semi-supervised methods achieve promising performance in disentanglement . Typically , comprehensive experiments in ( Locatello et al. , 2020 ) validate the effectiveness of a limited amount of supervision in state-of-the-art unsupervised disentanglement models . LORD ( Gabbay & Hoshen , 2020 ) applies a latent optimization framework with a noise regularizer on content embeddings to achieve superior performance over amortized inference . Based on LORD , OverLORD ( Gabbay & Hoshen , 2021 ) is proposed to disentangle class , correlated and uncorrelated attributes for image translation . A more recent work ( Gabbay et al. , 2021 ) adopts a pre-trained CLIP ( Radford et al. , 2021 ) model to generate partial annotations for image manipulation . However , there exist two main drawbacks among these methods : 1 ) using different separate encoders for different factors is resource-wasteful for real-world applications and requires expensive human design . 2 ) just learning the content embeddings inside each sample is not sufficient to learn the diversity existing in the dataset . Driven by the shortcomings discussed above , we propose a simple yet effective method named ContraLORD , where we incorporate contrastive learning into latent optimization for representation disentanglement . Recent works ( Deng et al. , 2020 ; Ojha et al. , 2020 ) apply the contrastive learning on the discriminator of the GAN ( Goodfellow et al. , 2014 ) for disentangling representations . Typically , the 3D imitative-contrastive learning in ( Deng et al. , 2020 ) is used for controllable face image generation by comparing pairs of generated images . However , in this work , we focus on applying contrastive learning on the encoders to learn the discriminative and generative embeddings with disentangled information . Specifically , we first apply a generator to learn discriminative and generative embeddings via latent optimization . Then we apply an encoder and a momentum encoder to dynamically learn disentangled information across a large number of samples with content-level and residual-level contrastive loss . In the meanwhile , we use the learned discriminative and generative embeddings to tune the encoder in an amortized manner . We evaluate our ContraLORD on two main tasks : linear classification and disentanglement . Extensive experiments show the effectiveness of the learned discriminative embeddings on linear classification and generative embeddings on the disentanglement of factors . We conduct comprehensive studies on three benchmarks on linear classification and seven benchmarks on disentanglement to investigate if contrastive self-supervised models can learn disentangled features . In the meantime , we achieve superior performance on linear classification compared to baselines . Our ContraLORD also achieves promising results over state-of-the-art methods in terms of disentanglement . The main contributions of this work can be summarized as follows : • We present a simple yet effective method called ContraLORD by incorporating contrastive learning into latent optimization for representation disentanglement and linear classification . • We formally explore the disentangled features across a large number of samples with content-level and residual-level contrastive losses . • Extensive experiments on ten benchmarks show the effectiveness of our approach on learning disentangled representations . 2 RELATED WORK . Discriminative representations learning . Discriminative representations learning has addressed researchers ’ attention for a long time since discriminative representations are significant for image classification . Most of the previous works adopt supervised ( Lezama et al. , 2018 ) and unsupervised learning ( Ji et al. , 2017 ; Zhang et al. , 2018 ; Zhou et al. , 2018 ; Zhang et al. , 2019 ) to learn embeddings that most discriminate between classes in the dataset . Typically , the principle of maximal coding rate reduction ( Yu et al. , 2020 ) is applied to maximize the coding rate difference between the whole dataset and the sum of each separate class . However , there exists little work of contrastive learning to explore the discriminative representations for the pre-training stage . In this work , we mainly focus on learning discriminative embeddings for linear classification by incorporating contrastive learning into latent optimization to improve the performance of baselines . Disentangled Representations Learning . Disentangled representation learning aims at learning generative factors existing in the dataset , that is , disentanglement learning . A bunch of previous work focuses on unsupervised learning with variational autoencoders , such as β-VAE ( Higgins et al. , 2017 ) , Factor-VAE ( Kim & Mnih , 2018 ) . Following those work , DCI disentanglement ( Eastwood & Williams , 2018 ) , SAP score ( Kumar et al. , 2018 ) , and Mutual Information Gap ( MIG ) ( Chen et al. , 2018 ) are often utilized as quantitative metrics to measure the quality of disentangled representations . In recent years , semi-supervised models have been used by many researchers in the literature . Adding a limited amount of supervision to unsupervised models is proven in ( Locatello et al. , 2020 ) to be effective in learning disentangled representations for real-world scenarios . LORD ( Gabbay & Hoshen , 2020 ) leverages the latent optimization framework with a noise regularizer on content embeddings for class and content disentanglement . More recently , a simple framework for disentangling labeled and unlabeled attributes is utilized in OverLORD ( Gabbay & Hoshen , 2021 ) for high-fidelity images synthesis . A study ( Gabbay et al. , 2021 ) uses a CLIP ( Radford et al. , 2021 ) pre-trained model to annotate a set of attributes for disentangled image manipulation . In this work , we intend to learn disentangled embeddings via combining latent optimization and contrastive selfsupervised learning . Contrastive Learning . Recently , contrastive self-supervised learning ( Tian et al. , 2020 ; Chen et al. , 2020a ; b ; Grill et al. , 2020 ; He et al. , 2020 ; Chen et al. , 2020c ; Chen & He , 2021 ) has been explored a lot by many effective methods . SimCLR ( Chen et al. , 2020a ) , an end-to-end structure , is proposed to pull away from the features of each instance from those of all other instances in the training set . In the self-supervised setting , low-level image transformations such as cropping , scaling , and color jittering are utilized for encoding the in-variances from samples . The InfoNCE loss , that is , the normalized temperature-scaled cross-entropy loss , are often optimized to maximize the similarity between positive samples and minimize the similarity between negative samples . Large batch size is always used in this kind of end-to-end structure ( Chen et al. , 2020a ; b ) to accumulate a large bunch of negative samples in the contrastive loss . PIRL ( Misra & Maaten , 2020 ) without a large batch size applies a memory bank to store negative samples and update representations at a specified stage . MoCo ( He et al. , 2020 ) and MoCov2 ( Chen et al. , 2020c ) replace the memory bank with a memory encoder to queue new batch samples and to dequeue the oldest batch . In this work , we leverage content-level and residual-level momentum encoders to store a queue of negative samples with disentangled information for learning generative embeddings , where content-level and residuallevel contrastive losses are applied to capture content and residual representations . 3 METHODOLOGY . 3.1 PROBLEM SETUP . In this part , we first begin with the problem setup , and formally define the notations for easy reading . In terms of the problem , our goal is to demystify the disentangled and discriminative features learned by contrastive self-supervised learning . To address this problem , we propose a simple yet effective method by combining contrastive learning and latent optimization together for representation disentanglement . To explain it better , we define the notations below in a unified manner . Notations . Given a set of training examples X = { x1 , x2 , · · · , xn } . For each image xi , i ∈ { 1 , 2 , · · · , n } , we need to learn one discriminative embeddings d and a generative embeddings gi from a pre-defined set of embeddings { g1 , g2 , · · · , gm } , where m denotes the total number of generative factors in the training data . That is , di ∈ R1×d , gi ∈ R1×g , where d , g denote the dimensionality of discriminative and generative embedding , separately . In our setting , we split the generative embedding gi into two folds : content embeddings gci and residual embeddings g r i . The content embeddings contain the information that is related to the discriminative embedding , while the residual embedding includes the uncorrelated information . Overall , the objective of our work is to learn di , gci , g r i for each image xi from a training dataset . In the next part , we present the technical details of our method . To learn di , gci , g r i from a set of training examples , we propose a simple yet effective approach called ContraLORD , as shown in Figure 1 . Our ContraLORD mainly includes two parts : 1 ) embedding optimization : we first use a generator G ( · ) to learn discriminative and disentangled embeddings via latent optimization . 2 ) encoder pretraining : we apply a encoder f ( · ) and a momentum encoder fm ( · ) to dynamically learn disentangled information across large amount of samples with content- and residual-level contrastive loss . 3.2 EMBEDDING OPTIMIZATION . In order to learn the discriminative and disentangled embeddings , we are motivated by LORD ( Gabbay & Hoshen , 2020 ) to introduce the latent optimization in the first stage . Specifically , we apply a generator G ( · ) to reconstruct the original image xi by using the discriminative and disentangled embeddings of each sample . Instead of using the KL-divergence in variation auto-encoders ( Kingma & Welling , 2014 ) , we equally add a regularizer with Gaussian noise of a fixed variance to the disentangled embeddings gci and gri . Thus , the objective of embeddings optimization is defined as Lopt = s n∑ i=1 ` ( G ( di , gci + zi , g r i + zi ) ) , xi ) + λ · ( ||gci ||2 + ||gri ||2 ) ( 1 ) where ` ( · ) denotes ` 2 loss for synthetic data and VGG perceptual loss for real images . λ is the penalty weight of the capacity of the disentangled embeddings . zi ∼ N ( 0 , σ2I ) . In this way , we can learn the disentangled embeddings di , gci , gri without any adversarial learning involved , that is , d̃i , g̃ c i , g̃ r i = argmin Lopt . For training sets with annotations , d̃i is given .

Summary: The authors describe a method for learning disentangled representations (categorical and continuous). It is based on the generative latent optimization scheme for learning disentangled latent codes for a given dataset, followed by training an encoder network to predict these latent codes given an arbitrary image.  The main contribution of the authors is a proposition to include contrastive losses into the encoder training to prevent it from converging to a suboptimal state where it could predict similar embedding to dissimilar input images. They validate the resulting model using standard image generation and disentanglement benchmarks across multiple datasets.