Contrastive Learning for Unpaired Image-to-Image Translation
Taesung Park, Alexei A. Efros, Richard Zhang, Jun-Yan Zhu
Introduction
Consider the image-to-image translation problem in Figure 1. We wish for the output to take on the appearance of the target domain (a zebra), while retaining the structure, or content, of the specific input horse. This is, fundamentally, a disentanglement problem: separating the content, which needs to be preserved across domains, from appearance, which must change. Typically, target appearance is enforced using an adversarial loss , while content is preserved using cycle-consistency . However, cycle-consistency assumes that the relationship between the two domains is a bijection, which is often too restrictive. In this paper, we propose an alternative, rather straightforward way of maintaining correspondence in content but not appearance – by maximizing the mutual information between corresponding input and output patches.
In a successful result, given a specific patch on the output, for example, the generated zebra forehead highlighted in blue, one should have a good idea that it came from the horse forehead, and not the other parts of the horse or the background vegetation. We achieve this by using a type of contrastive loss function, InfoNCE loss , which aims to learn an embedding or an encoder that associates corresponding patches to each other, while disassociating them from others. To do so, the encoder learns to pay attention to the commonalities between the two domains, such as object parts and shapes, while being invariant to the differences, such as the textures of the animals. The two networks, the generator and encoder, conspire together to generate an image such that patches can be easily traceable to the input.
Contrastive learning has been an effective tool in unsupervised visual representation learning . In this work, we demonstrate its effectiveness in a conditional image synthesis setting and systematically study several key factors to make it successful. We find it pertinent to use it on a multilayer, patchwise fashion. In addition, we find that drawing negatives internally from within the input image, rather than externally from other images in the dataset, forces the patches to better preserve the content of the input. Our method requires neither memory bank nor specialized architectures .
Extensive experiments show that our faster, lighter model outperforms both prior one-sided translation methods and state-of-the-art models that rely on several auxiliary networks and multiple loss functions. Furthermore, since our contrastive representation is formulated within the same image, our method can even be trained on single images. Our code and models are available at GitHub.
Related Work
Image translation and cycle-consistency. Paired image-to-image translation maps an image from input to output domain using an adversarial loss , in conjunction with a reconstruction loss between the result and target. In unpaired translation settings, corresponding examples from domains are not available. In such cases, cycle-consistency has become the de facto method for enforcing correspondence , which learns an inverse mapping from the output domain back to the input and checks if the input can be reconstructed. Alternatively, UNIT and MUNIT propose to learn a shared intermediate “content” latent space. Recent works further enable multiple domains and multi-modal synthesis and improve the quality of results . In all of the above examples, cycle-consistency is used, often in multiple aspects, between (a) two image domains (b) image to latent , or (c) latent to image . While effective, the underlying bijective assumption behind cycle-consistency is sometimes too restrictive. Perfect reconstruction is difficult to achieve, especially when images from one domain have additional information compared to the other domain.
Relationship preservation. An interesting alternative approach is to encourage relationships present in the input be analogously reflected in the output. For example, perceptually similar patches within an input image should be similar in the output , output and input images share similar content regarding a pre-defined distance , vector arithmetic between input images is preserved using a margin-based triplet loss , distances between input images should be consistent in output images , the network should be equivariant to geometric transformations . Among them, TraVeLGAN , DistanceGAN and GcGAN enable one-way translation and bypass cycle-consistency. However, they rely on relationship between entire images, or often with predefined distance functions. Here we seek to replace cycle-consistency by instead learning a cross-domain similarity function between input and output patches through information maximization, without relying on a pre-specified distance.
Contrastive representation learning. Traditional unsupervised learning has sought to learn a compressed code which can effectively reconstruct the input . Data imputation – holding one subset of raw data to predict from another – has emerged as a more effective family of pretext tasks, including denoising , context prediction , colorization , cross-channel encoding , frame prediction , and multi-sensory prediction . However, such methods suffer from the same issue as before — the need for a pre-specified, hand-designed loss function to measure predictive performance.
Recently, a family of methods based on maximizing mutual information has emerged to bypass the above issue . These methods make use of noise contrastive estimation , learning an embedding where associated signals are brought together, in contrast to other samples in the dataset (note that similar ideas go back to classic work on metric learning with Siamese nets ). Associated signals can be an image with itself , an image with its downstream representation , neighboring patches within an image , or multiple views of the input image , and most successfully, an image with a set of transformed versions of itself . The design choices of the InfoNCE loss, such as the number of negatives and how to sample them, hyperparameter settings, and data augmentations all play a critical role and need to be carefully studied. We are the first to use InfoNCE loss for the conditional image synthesis tasks. As such, we draw on these important insights, and find additional pertinent factors, unique to image synthesis.
Methods
We wish to translate images from input domain to appear like an image from the output domain . We are given a dataset of unpaired instances . Our method can operate even when and only contain a single image each.
Our method only requires learning the mapping in one direction and avoids using inverse auxiliary generators and discriminators. This can largely simplify the training procedure and reduce training time. We break up our generator function into two components, an encoder followed by a decoder , which are applied sequentially to produce output image .
Adversarial loss. We use an adversarial loss , to encourage the output to be visually similar to images from the target domain, as follows:
Mutual information maximization. We use a noise contrastive estimation framework to maximize mutual information between input and output. The idea of contrastive learning is to associate two signals, a “query” and its “positive” example, in contrast to other points within the dataset, referred to as “negatives”. The query, positive, and negatives are mapped to -dimensional vectors and , respectively. denotes the n-th negative. We normalize vectors onto a unit sphere to prevent the space from collapsing or expanding. An –way classification problem is set up, where the distances between the query and other examples are scaled by a temperature and passed as logits . The cross-entropy loss is calculated, representing the probability of the positive example being selected over negatives.
Our goal is to associate the input and output data. In our context, query refers to an output. positive and negatives are corresponding and noncorresponding input. Below, we explore several important design choices, including how to map the images into vectors and how to sample the negatives.
Multilayer, patchwise contrastive learning. In the unsupervised learning setting, contrastive learning has been used both on an image and patch level . For our application, we note that not only should the whole images share content, but also corresponding patches between the input and output images. For example, given a patch showing the legs of an output zebra, one should be able to more strongly associate it to the corresponding legs of the input horse, more so than the other patches of the horse image. Even at the pixel level, the colors of a zebra body (black and white) can be more strongly associated to the color of a horse body than to the background shades of grass. Thus, we employ a multilayer, patch-based learning objective.
Since the encoder is computed to produce the image translation, its feature stack is readily available, and we take advantage. Each layer and spatial location within this feature stack represents a patch of the input image, with deeper layers corresponding to bigger patches. We select layers of interest and pass the feature maps through a small two-layer MLP network , as used in SimCLR , producing a stack of features , where represents the output of the -th chosen layer. We index into layers and denote , where is the number of spatial locations in each layer. We refer to the corresponding feature as and the other features as , where is the number of channels at each layer. Similarly, we encode the output image into .
We aim to match corresponding input-output patches at a specific location. We can leverage the other patches within the input as negatives. For example, a zebra leg should be more closely associated with an input horse leg than the other patches of the same input, such as other horse parts or the background sky and vegetation. We name it as the PatchNCE loss, as illustrated in Figure 2. Appendix C.3 provides pseudocode.
In Section 4.1, we show that our encoder learns to capture domain-invariant concepts, such as animal body, grass, and sky for horse zebra, while our decoder learns to synthesize domain-specific features such as zebra stripes. Interestingly, through systematic evaluations, we find that using internal patches only outperforms using external patches. We hypothesize that by using internal statistics, our encoder does not need to model large intra-class variation such as white horse vs. brown horse, which is not necessary for generating output zebras. Single image internal statistics has been proven effective in many vision tasks such as segmentation , super-resolution, and denoising .
Final objective. Our final objective is as follows. The generated image should be realistic, while patches in the input and output images should share correspondence. Figure 1 illustrates our minimax learning objective. Additionally, we may utilize PatchNCE loss on images from domain to prevent the generator from making unnecessary changes. This loss is essentially a learnable, domain-specific version of the identity loss, commonly used by previous unpaired translation methods .
We choose when we jointly train with the identity loss , and choose a larger value without the identity loss () to compensate for the absence of the regularizer. We find that the former configuration, named Contrastive Unpaired Translation (CUT) hereafter, achieves superior performance to existing methods, whereas the latter, named FastCUT, can be thought as a faster and lighter version of CycleGAN. Our model is relatively simple compared to recent methods that often use 5-10 losses and hyper-parameters.
Discussion. Li et al. has shown that cycle-consistency loss is the upper bound of conditional entropy (and ). Therefore, minimizing cycle-consistency loss encourages the output to be more dependent on input . This is related to our objective of maximizing the mutual information , as . As entropy is a constant and independent of the generator , maximizing mutual information is equivalent to minimizing the conditional entropy. Notably, using contrastive learning, we can achieve a similar goal without introducing inverse mapping networks and additional discriminators. In the unconditional modeling scenario, InfoGAN shows that simple losses (e.g., L2 or cross-entropy) can serve as a lower bound for maximizing mutual information between an image and a low-dimensional code. In our setting, we maximize the mutual information between two high-dimensional image spaces, where simple losses are no longer effective. Liang et al. proposes an adversarial loss based on Siamese networks that encourages the output to be closer to the target domain than to its source domain. The above method still builds on cycle-consistency and two-way translations. Different from the above work, we use contrastive learning to enforce content consistency, rather than to improve the adversarial loss itself. To measure the similarity between two distributions, the Contextual Loss used softmax over cosine disntances of features extracted from pre-trained networks. In contrast, we learn the encoder with the NCE loss to associate the input and output patches at the same location.
Experiments
We test across several datasets. We first show that our method improves upon baselines in unpaired image translation. We then show that our method can extend to single-image training. Full results are available at our website.
Datasets We conduct experiments on the following datasets.
CatDog contains 5,000 training and 500 val images from AFHQ Dataset .
HorseZebra contains 2,403 training and 260 zebra images from ImageNet and was introduced in CycleGAN .
Cityscapes contains street scenes from German cities, with 2,975 training and 500 validation images. We train models at 256256 resolution. Unlike previous datasets listed, this does have corresponding labels. We can leverage this to measure how well our unpaired algorithm discovers correspondences.
Evaluation protocol. We adopt the evaluation protocols from , aimed at assessing visual quality and discovered correspondence. For the first, we utilize the widely-used Fréchet Inception Distance (FID) metric, which empirically estimates the distribution of real and generated images in a deep network space and computes the divergence between them. Intuitively, if the generated images are realistic, they should have similar summary statistics as real images, in any feature space. For Cityscapes specifically, we have ground truth of paired label maps. If accurate correspondences are discovered, the algorithm should generate images that are recognizable as the correct class. Using an off-the-shelf network to test “semantic interpretability” of image translation results has been commonly used . We use the pretrained semantic segmentation network DRN. We train the DRN at 256x128 resolution, and compute mean average precision (mAP), pixel-wise accuracy (pixAcc), and average class accuracy (classAcc). See Appendix C.2 for more evaluation details.
Comparison to baselines. In Table 1, we show quantitative measures of our and Figure 3, we compare our method to baselines. We present two settings of our method in Eqn. 5: CUT with the identity loss (), and FastCUT without it (). On image quality metrics across datasets, our methods outperform baselines. We show qualitative results in Figure 3 and additional results in Appendix A. In addition, our Cityscapes semantic segmentation scores are higher, suggesting that our method is able to find correspondences between output and input.
Speed and memory. Since our model is one-sided, our method is memory-efficient and fast. For example, our method with the identity loss was 40% faster and 31% more memory-efficient than CycleGAN at training time, using the same architectures as CycleGAN (Table 1). Furthermore, our faster variant FastCUT is 63% faster and 53% lighter, while achieving superior metrics to CycleGAN. Table 1 contains the speed and memory usage of each method measured on NVIDIA GTX 1080Ti, and shows that FastCUT achieves competitive FIDs and segmentation scores with a lower time and memory requirement. Therefore, our method can serves as a practical, lighter alternative in scenarios, when an image translation model is jointly trained with other components .
2 Ablation study and analysis
We find that in the image synthesis setting, similarly to the unsupervised learning setting , implementation choices for contrastive loss are important. Here, try various settings and ablations of our method, summarized in Figure 4. By default, we use the ResNet-based generator used in CycleGAN , with patchNCE using (a) negatives sampled from the input image, (b) multiple layers of the encoder, and (c) a PatchNCE loss on domain . In Figure 4, we show results using several variants and ablations, taken after training for 400 epochs. We show qualitative examples in Figure 5.
Internal negatives are more effective than external. By default, we sample negatives from within the same image (internal negatives). We also try adding negatives from other images, using a momentum encoder . However, the external negatives, either as addition (int and ext) or replacement of internal negatives (ext only), hurts performance. In Figure 5, we see a loss of quality, such as repeated texture in the Cityscapes dataset, indicating that sampling negatives from the same image serves as a stronger signal for preserving content.
regularizer stabilizes training. Given an image from the output domain , this regularizer encourages the generator to leave the image unchanged with our patch-based contrastive loss. We also experiment with a variant without this regularizer, no id. As shown in Figure 4, removing the regularizer improves results for the horsezebra task, but decreases performance on Cityscapes. We further investigate by showing the training curves in Figure 6, across 400 epochs. In the Cityscapes results, the training can collapse without the regularizer (although it can recover). We observe that although the final FID is sometimes better without, the training is more stable with the regularizer.
Visualizing learned similarity by encoder To further understand why our encoder network has learned to perform horse zebra task, we study the output space of the 1st residual block for both horse and zebra features. As shown in Figure 7. Given an input and output image, we compute the distance between a query patch’s feature vector (highlighted as red or blue dot) to feature vectors of all the patches in the input using (Eqn. 2). Additionally, we perform a PCA dimension reduction on feature vectors from both horse and zebra patches. In (d) and (e), we show the top three principal components, which looks similar before and after translation. This indicates that our encoder is able to bring the corresponding patches from two domains into a similar location in the feature embedding space.
Additional applications. Figure 8 shows additional results: Parisian street Burano’s brightly painted houses and Russian Blue cat Grumpy cat.
3 High-resolution single image translation
Finally, we conduct experiments in the single image setting, where both the source and target domain only have one image each. Here, we transfer a Claude Monet’s painting to a natural photograph. Recent methods have explored training unconditional models on a single image. Bearing the additional challenge of respecting the structure of the input image, conditional image synthesis using only one image has not been explored by previous image-to-image translation methods. Our painting photo task is also different from neural style transfer (photo painting) and photo style transfer (photo photo).
Since the whole image (at HD resolution) cannot fit on a commercial GPU, at each iteration we train on 16 random crops of size 128128. We also randomly scale the image to prevent overfitting. Furthermore, we observe that limiting the receptive field of the discriminator is important for preserving the structure of the input image, as otherwise the GAN loss will force the output image to be identical to the target image. Therefore, the crops are further split into 6464 patches before passed to the discriminator. Lastly, we find that using gradient penalty stabilizes optimization. We call this variant SinCUT.
Conclusion
Acknowledgments. We thank Allan Jabri and Phillip Isola for helpful discussion and feedback. Taesung Park is supported by a Samsung Scholarship and an Adobe Research Fellowship, and some of this work was done as an Adobe Research intern. This work was partially supported by NSF grant IIS-1633310, grant from SAP, and gifts from Berkeley DeepDrive and Adobe.
References
Appendix A Additional Image-to-Image Results
We first show additional, randomly selected results on datasets used in our main paper. We then show results on additional datasets.
In Figure 10, we show additional, randomly selected results for HorseZebra and CatDog. This is an extension of Figure 3 in the main paper. We compare to baseline methods CycleGAN , MUNIT , DRIT , Self-Distance and DistanceGAN , and GcGAN .
A.2 Additional datasets
In Figure 11 and Figure 12, we show additional datasets, compared against baseline method CycleGAN . Our method provides better or comparable results, demonstrating its flexibility across a variety of datasets.
AppleOrange contains 996 apple and 1,020 orange images from ImageNet and was introduced in CycleGAN .
Yosemite SummerWinter contains 1,273 summer and 854 winter images of Yosemite scraped using the FlickAPI was introduced in CycleGAN .
GTACityscapes GTA contains 24,966 images and Cityscapes contains 19,998 images of street scenes from German cities. The task was originally used in CyCADA .
Appendix B Additional Single Image Translation Results
We show additional results in Figure 13 and Figure 14, and describe training details below.
Training details. At each iteration, the input image is randomly scaled to a width between 384 to 1024, and we randomly sample 16 crops of size . To avoid overfitting, we divide crops into tiles before passing them to the discriminator. At test time, since the generator network is fully convolutional, it takes the input image at full size.
We found that adopting the architecture of StyleGAN2 instead of CycleGAN slightly improves the output quality, although the difference is marginal. Our StyleGAN2-based generator consists of one downsampling block of StyleGAN2 discriminator, 6 StyleGAN2 residual blocks, and one StyleGAN2 upsampling block. Our discriminator has the same architecture as StyleGAN2. Following StyleGAN2, we use non-saturating GAN loss with R1 gradient penalty . Since we do not use style code, the style modulation layer of StyleGAN2 was removed.
In Figure 13 and Figure 14, we show additional comparison results for our method, Gatys et al. , STROTSS , WCT2 , and CycleGAN baseline . Note that the CycleGAN baseline adopts the same augmentation techniques as well as the same generator/discriminator architectures as our method. The image resolution is at 1-2 Megapixels. Please zoom in to see more visual details.
Both figures demonstrate that our results look more photorealistic compared to CycleGAN baseline, Gatys et al , and WCT2. The quality of our results is on par with results from STROTSS . Note that STROTSS compares to and outperforms recent style transfer methods (e.g., ).
Appendix C Unpaired Translation Details and Analysis
To show the effect of the proposed patch-based contrastive loss, we intentionally match the architecture and hyperparameter settings of CycleGAN, except the loss function. This includes the ResNet-based generator with 9 residual blocks, PatchGAN discriminator , Least Square GAN loss , batch size of 1, and Adam optimizer with learning rate 0.002.
Our full model CUT is trained up to 400 epochs, while the fast variant FastCUT is trained up to 200 epochs, following CycleGAN. Moreover, inspired by GcGAN , FastCUT is trained with flip-equivariance augmentation, where the input image to the generator is horizontally flipped, and the output features are flipped back before computing the PatchNCE loss. Our encoder is the first half of the CycleGAN generator . In order to calculate our multi-layer, patch-based contrastive loss, we extract features from 5 layers, which are RGB pixels, the first and second downsampling convolution, and the first and the fifth residual block. The layers we use correspond to receptive fields of sizes 11, 99, 1515, 3535, and 9999. For each layer’s features, we sample 256 random locations, and apply 2-layer MLP to acquire 256-dim final features. For our baseline model that uses MoCo-style memory bank , we follow the setting of MoCo, and used momentum value 0.999 with temperature 0.07. The size of the memory bank is 16384 per layer, and we enqueue 256 patches per image per iteration.
C.2 Evaluation details
We list the details of our evaluation protocol.
Fréchet Inception Distance (FID ) throughout this paper is computed by resizing the images to 299-by-299 using bilinear sampling of PyTorch framework, and then taking the activations of the last average pooling layer of a pretrained Inception V3 using the weights provided by the TensorFlow framework. We use the default setting of https://github.com/mseitzer/pytorch-fid. All test set images are used for evaluation, unless noted otherwise.
Semantic segmentation metrics on the Cityscapes dataset are computed as follows. First, we trained a semantic segmentation network using the DRN-D-22 architecture. We used the recommended setting from https://github.com/fyu/drn, with batch size 32 and learning rate 0.01, for 250 epochs at 256x128 resolution. The output images of the 500 validation labels are resized to 256x128 using bicubic downsampling, passed to the trained DRN network, and compared against the ground truth labels downsampled to the same size using nearest-neighbor sampling.
C.3 Pseudocode
Here we provide the pseudo-code of PatchNCE loss in the PyTorch style. Our code and models are available at our GitHub repo.
C.4 Distribution matching
In Figure 15, we show an interesting phenomenon of our method, caused by the training set imbalance of the horsezebra set. We use an off-the-shelf DeepLab model trained on COCO-Stuff , to measure the percentage of pixels that belong to horses and zebrasPretrained model from https://github.com/kazuto1011/deeplab-pytorch. The training set exhibits dataset bias . On average, zebras appear in more close-up pictures than horses and take up about twice the number of pixels ( vs ). To perfectly satisfy the discriminator, a translation model should attempt to match the statistics of the training set. Our method allows the flexibility for the horses to change the size, and the percentage of output zebra pixels () better matches the training distribution () than the CycleGAN baseline (). On the other hand, our fast variant FastCUT uses a larger weight () on the Patch NCE loss and flip-equivariance augmentation, and hence behaves more conservatively and more similar to CycleGAN. The strong distribution matching capacity has pros and cons. For certain applications, it can create introduce undesired changes (e.g., zebra patterns on the background for horsezebra). On the other hand, it can enable dramatic geometric changes for applications such as CatDog.
C.5 Additional Ablation studies
In the paper, we mainly discussed the impact of loss functions and the number of patches on the final performance. Here we present additional ablation studies on more subtle design choices. We run all the variants on horse2zebra datasets . The FID of our original model is 46.6. We compare it to the following two variants of our model:
Ours without weight sharing for the encoder and MLP projection network : for this variant, when computing features , we use two separate encoders and MLP networks for embedding input images (e.g., horse) and the generated images (e.g., zebras) to feature space. They do not share any weights. The FID of this variant is 50.5, worse than our method. This shows that weight sharing helps stabilize training while reducing the number of parameters in our model.
Ours without updating the decoder using PatchNCE loss: in this variant, we exclude the gradient propagation of the decoder regarding PatchNCE loss . In other words, the decoder only gets updated through the adversarial loss . The FID of this variant is 444.2, and the results contain severe artifacts. This shows that our not only helps learn the encoder , as done in previous unsupervised feature learning methods , but also learns a better decoder together with the GAN loss. Intuitively, if the generated result has many artifacts and is far from realistic, it would be difficult for the encoder to find correspondences between the input and output, producing a large PatchNCE loss.
Appendix D Changelog
v2 and v3 (1) Fix typos in Eqn. 3 and Eqn. 4. (2) Add additional related work.