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 paper proposes a new two-step approach for learning disentangled representation. In the first step, a generator is learned to map latent embeddings to reconstruct images. Similar to LORD, the embedding corresponding to each image is jointly optimized with the generator weights. The second step uses contrastive learning together with regression loss to learn an encoder that maps images to the latent space, which is the main contribution of the paper. Results show that the proposed approach achieves state-of-the-art performance on several datasets for disentangled representation learning.