Contrastive Learning for Unpaired Image-to-Image Translation

Abstract

In image-to-image translation, each patch in the output should reflect the content of the corresponding patch in the input, independent of domain. We propose a straightforward method for doing so – maximizing mutual information between the two, using a framework based on contrastive learning. The method encourages two elements (corresponding patches) to map to a similar point in a learned feature space, relative to other elements (other patches) in the dataset, referred to as negatives. We explore several critical design choices for making contrastive learning effective in the image synthesis setting. Notably, we use a multilayer, patch-based approach, rather than operate on entire images. Furthermore, we draw negatives from within the input image itself, rather than from the rest of the dataset. We demonstrate that our framework enables one-sided translation in the unpaired image-to-image translation setting, while improving quality and reducing training time. In addition, our method can even be extended to the training setting where each “domain” is only a single image.

Appendices

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.

1.1 Additional Comparisons

In Fig. 10, we show additional, randomly selected results for Horse$\rightarrow$Zebra and Cat$\rightarrow$Dog. This is an extension of Fig. 3 in the main paper. We compare to baseline methods CycleGAN [89], MUNIT [30], DRIT [41], Self-Distance and DistanceGAN [4], and GcGAN [18].

1.2 Additional Datasets

In Fig. 11 and 12, we show additional datasets, compared against baseline method CycleGAN [89]. Our method provides better or comparable results, demonstrating its flexibility across a variety of datasets.

• Apple$\rightarrow$Orange contains 996 apple and 1,020 orange images from ImageNet and was introduced in CycleGAN [89].

• Yosemite Summer$\rightarrow$Winter contains 1,273 summer and 854 winter images of Yosemite scraped using the FlickAPI was introduced in CycleGAN [89].

• GTA$\rightarrow$Cityscapes GTA contains 24,966 images [63] and Cityscapes [13] contains 19,998 images of street scenes from German cities. The task was originally used in CyCADA [29].

Fig. 10.

Randomly selected Horse$\rightarrow$ Zebra and Cat$\rightarrow$Dog results. This is an extension of Fig. 3 in the main paper.

Fig. 11.

Apple$\rightarrow$Orange and Summer$\rightarrow$Winter Yosemite. CycleGAN models were downloaded from the authors’ public code repository. Apple$\rightarrow$Orange shows that CycleGAN may suffer from color flipping issue.

Fig. 12.

GTA$\rightarrow$Cityscapes results at $1024\times 512$ resolution. The model was trained on $512\times 512$ crops.

Additional Single Image Translation Results

We show additional results in Fig. 13 and Fig. 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 $128\,\times \,128$. To avoid overfitting, we divide crops into $64\,\times \,64$ 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 [36] 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 [61] with R1 gradient penalty [53]. Since we do not use style code, the style modulation layer of StyleGAN2 was removed.

Single image results.

In Fig. 13 and 14, we show additional comparison results for our method, Gatys et al. [19], STROTSS [39], WCT$^2$ [82], and CycleGAN baseline [89]. 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. [19], and WCT$^2$. The quality of our results is on par with results from STROTSS [39]. Note that STROTSS [39] compares to and outperforms recent style transfer methods (e.g., [22, 52]).

Fig. 13.

High-res painting to photo translation (I). We transfer Monet’s paintings to reference natural photos shown as insets at top-left corners. The training only requires a single image from each domain. We compare our results to recent style and photo transfer methods including Gatys et al. [19], WCT$^2$ [82], STROTSS [39], and our modified patch-based CycleGAN [89]. Our method can reproduce the texture of the reference photos while retaining structure of the input paintings. Our results are at 1k $\sim$ 1.5k resolution.

Fig. 14.

High-res painting to photo translation (II). We transfer Monet’s paintings to reference natural photos shown as insets at top-left corners. The training only requires a single image from each domain. We compare our results to recent style and photo transfer methods including Gatys et al. [19], WCT$^2$ [82], STROTSS [39], and our modified patch-based CycleGAN [89]. Our method can reproduce the texture of the reference photos while retaining structure of the input paintings. Our results are at 1k $\sim$ 1.5k resolution.

Unpaired Translation Details and Analysis

1.1 Training Details

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 [34] with 9 residual blocks, PatchGAN discriminator [31], Least Square GAN loss [50], batch size of 1, and Adam optimizer [38] 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 [18], 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 $G_{\text {enc}}$ is the first half of the CycleGAN generator [89]. 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 1 $\times$ 1, 9 $\times$ 9, 15 $\times$ 15, 35 $\times$ 35, and 99 $\times$ 99. 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 [24], 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.

1.2 Evaluation Details

We list the details of our evaluation protocol.

Fréchet Inception Distance (FID [26]) 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 [70] 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 [83] 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 256 $\times$ 128 resolution. The output images of the 500 validation labels are resized to 256 $\times$ 128 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.

1.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.

Fig. 15.

Distribution matching. We measure the percentage of pixels belonging to the horse/zebra bodies, using a pre-trained semantic segmentation model. We find a distribution mismatch between sizes of horses and zebras images – zebras usually appear larger (36.8% vs. 17.9%). Our full method CUT has the flexibility to enlarge the horses, as a means of better matching of the training statistics than CycleGAN [89]. Our faster variant FastCUT, trained with a higher PatchNCE loss ($\lambda _{X}=10$) and flip-equivariance augmentation, behaves more conservatively like CycleGAN.

1.4 Distribution Matching

In Fig. 15, we show an interesting phenomenon of our method, caused by the training set imbalance of the horse$\rightarrow$zebra set. We use an off-the-shelf DeepLab model [7] trained on COCO-Stuff [6], to measure the percentage of pixels that belong to horses and zebras¹. The training set exhibits dataset bias [74]. On average, zebras appear in more close-up pictures than horses and take up about twice the number of pixels ($37\%$ vs $18\%$). 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 ($31\%$) better matches the training distribution ($37\%$) than the CycleGAN baseline ($19\%$). On the other hand, our fast variant FastCUT uses a larger weight ($\lambda _{X} = 10$) 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 horse$\rightarrow$zebra). On the other hand, it can enable dramatic geometric changes for applications such as Cat$\rightarrow$Dog.

1.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  [89]. 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 $G_{\text {enc}}$ and MLP projection network $H$: for this variant, when computing features $\{\textit{\textbf{z}}_l\}_L=\{H_l(G_{\text {enc}}^l(\textit{\textbf{x}}))\}_L$, 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 $G_{\text {dec}}$ using PatchNCE loss: in this variant, we exclude the gradient propagation of the decoder $G_{\text {dec}}$ regarding PatchNCE loss $\mathcal {L}_\text {PatchNCE}$. In other words, the decoder $G_{\text {dec}}$ only gets updated through the adversarial loss $\mathcal {L}_\text {GAN}$. The FID of this variant is 444.2, and the results contain severe artifacts. This shows that our $\mathcal {L}_\text {PatchNCE}$ not only helps learn the encoder $G_{\text {enc}}$, as done in previous unsupervised feature learning methods [24], but also learns a better decoder $G_{\text {dec}}$ 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.