Hybrid Generative-Contrastive Representation Learning

Given a batch of $N$ unlabeled samples $\{\mathbf{x}_{i}\}_{i=1}^{N}$ and the corresponding set of augmentation pairs $\{(\mathbf{x}_{i}^{(1)},\mathbf{x}_{i}^{(2)})\}_{i=1}^{N}$ where $\mathbf{x}_{i}^{(1)}$ and $\mathbf{x}_{i}^{(2)}$ are two random augmentations of $\mathbf{x}_{i}$, our hybrid loss function is defined as

where $\alpha$ and $\beta$ are the loss weights, and the generative loss ${\mathcal{L}}_{g}$ and the contrastive loss ${\mathcal{L}}_{c}$ are respectively formulated as

Here, $D$ is the data dimension, $\mathbf{z}_{i}^{(1)}$ and $\mathbf{z}_{i}^{(2)}$ are projected and normalized representations extracted from $\mathbf{x}_{i}^{(1)}$ and $\mathbf{x}_{i}^{(2)}$, respectively, and $\mathbf{x}_{i,k}^{(1)}$ is the $k$th element of $\mathbf{x}_{i}^{(1)}$. The generative loss $\mathcal{L}_{g}$ is the likelihood of the model autoregressively factorized in a typical raster scan order [42], and the contrastive loss $\mathcal{L}_{c}$ is the symmetric normalized temperature cross-entropy loss as in SimCLR [6].

In order to optimize the loss function, we propose a transformer-based architecture composed of encoder and decoder, motivated by empirical observations presented in iGPT [4], where early transformer blocks behave similarly to an encoder in typical auto-encoders and the remaining blocks behave like a decoder. Therefore, without altering the spatial resolution, we explicitly split a set of transformer blocks into an encoder and a decoder as shown in Figure 1. For the contrastive learning, the encoder takes two versions of an input, one is weakly augmented and the other is strongly augmented. Then, the two representations obtained from the encoder are used for computing the contrastive loss $\mathcal{L}_{c}$. The decoder processes only the weakly augmented image to compute the generative loss $\mathcal{L}_{g}$. We observe that training of the generative loss with strong augmentations deteriorates the robustness, since it focuses on producing unrealistic samples from a severely distorted distribution.

We use the similar transformer block as iGPT for the encoder and decoder, where each block is constructed as follows. Let $\mathbf{h}_{l-1}$ be the output from the $(l-1)$th block. The $l$th transformer block is composed of the following operations,

where LN refers to layer normalization [1], causal multi-head self-attention module computes interactions across a sequence only in a raster scan order, and MLP block of the same structure as iGPT computes nonlinear embeddings of the tokens in a sequence independently. To accelerate a self-attention module, we reduce the size of the sequence by color quantization [4] and adopt the sparse attention block used in AxialTransformer [23] with slight modification. For color quantization, we compress RGB color channels into one of 512 codes as in iGPT, where the codebook is simply constructed by $k$-means clustering. We use the codebook provided from the official implementation of iGPT for fair comparison. Given a colored image of 32x32 resolution, this simple processing effectively reduces the sequence length from 3,072 to 1,024.

AxialTransformer [23] accelerates a self-attention module by using a two-dimensional structure of an image, where row-wise and column-wise attentions are combined to access a full image in a raster scan order. In GCRL, unlike the original implementation of AxialTransformer, we sequentially apply row-wise attention, causal column-wise attention, and causal row-wise attention, to access a full image in every transformer block, as shown in Table 1. We empirically observe that our implementation of axial attention block reduces working memory and training time significantly compared to dense attention, without loosing performance. For detailed description including empirical results on the type of attention blocks, please refer to our PyTorch-like pseudo-code and ablation study in Appendix.

To fairly compare GCRL to iGPT and SimCLR, we use the same architecture for all objectives. We perform experiments with several models by varying the number of transformer blocks (6, 12, and 24 blocks). For all models, two-dimensional positional embeddings are applied after the first linear embedding, and the number of heads for axial attention was set to 8. Table 1 compares the sizes of the models tested for iGPT, SimCLR, and GCRL. Note that since GCRL is composed of an encoder and a decoder, we only need a half of the total number of blocks for evaluating the representations.

We use three datasets: CIFAR10, CIFAR100 [28], and ImageNet32 [10], where CIFAR contains 50K samples for training and ImageNet32 provides 1.2M training samples of resolution 32x32. We use Adam optimizer with decoupled weight decay (AdamW) [39] with gradient clipping of norm 1.0 and weight decay factor $10^{-4}$, and do not decay trainable parameters in layer normalization and token embeddings. We apply dropout with rate 0.1 before every residual connection in transformer blocks. We linearly increase the learning rate from zero to a specified value over 5 epochs and apply the cosine decay after it [38]. We use PyTorch version 1.6.0 and conduct all experiments with V100s with 32GB memories. For CIFAR experiments, we report mean and standard deviations from repetitions with random seed 0, 1, and 2. For ImageNet experiments, we report the result with random seed 0.

We observe that it is practically useful for training GCRL only with the generative loss ${\mathcal{L}}_{g}$ for the first half of the total training epochs and then train with the full objective for the remaining half. Also, we set both of the loss weights $\alpha$ and $\beta$ for the hybrid objective to $1.0$. Below, we summarize the data augmentation policies we used.

iGPT. We use the standard random crop with padding followed by horizontal flip with probability 0.5 for all experiments, where the pad size is $4$ and the padding is done with reflections of images.

SimCLR. We use the similar augmentation policy as in [6]. We use the same two-layer MLP projection head as in SimCLR, where batch normalization is replaced by layer normalization, the hidden dimension is set to 1,024, and the final embedding dimension is set to 64 for CIFAR and 128 for ImageNet experiments.

GCRL. This requires both strong and weak augmentations where the weakly augmented images are processed for the generative loss. We choose the weak augmentation as the one used for iGPT and the strong augmentation as the one used for SimCLR. Table 2 shows a pipeline of data augmentations for SimCLR and GCRL. We think that this augmentation policy is fair enough to see the different behaviors of SimCLR and our approach, though it is possible to obtain more discriminative representations by employing a sophisticated augmentation strategy. For the projection head, we use the same network as in SimCLR.

It is not clear what representations should be used for downstream tasks in iGPT. Hence, we test with two versions of representations for iGPT. The first one is pooled from the middle transformer block, and the second one is pooled from the last transformer block. For SimCLR and GCRL, the output before the projection head is selected as the final representation.

We use the linear evaluation protocol on representations to measure the discriminative performance of visual features, and report bits-per-dimension (bpd) on the color-quantized space ²²2 Note that the bpd values reported in this paper are not directly comparable to the ones in literature [8, 26], because we quantize the color space to reduce the sequence length. for assessing the generative performance. For linear evaluation, we follow the setup in the previous work [4, 6], where we train a single-layer linear classifier with AdamW optimizer for 100 epochs with weight decay $10^{-4}$ and batch size 512. We use the learning rate of $0.01$ for CIFAR experiments, and $0.1$ for ImageNet experiments. No data augmentation is applied for linear evaluation.

Table 4 compares GCRL with baselines under various model capacity and datasets, where the hyper-parameters used for this experiment are reported in Table 3. For iGPT-76M, SimCLR-76M, and GCRL-(38+38)M with CIFAR10 and CIFAR100, we fine-tune our pre-trained ImageNet32 models by 15 epochs. Table 4 shows that GCRL is able to learn representations having both discriminative and generative features. When the model capacity is high enough, we observe that GCRL consistently improves the linear evaluation performance, compared to iGPT (our implementation), iGPT-S [4], and SimCLR, while maintaining the generative performance. Specifically, in case of ImageNet32, GCRL achieves 43.6% linear evaluation accuracy, which is superior to iGPT-S (41.9%). In addition, when fine-tuning the model of 76M and (38+38)M trained on ImageNet to CIFARs, GCRL achieves lower bpd than the one from iGPT (2.6784 vs. 2.6969 in CIFAR10), which indicates the generative performance of GCRL is superior to iGPT.

Figure 2 compares the perceptual quality of generated images from iGPT and GCRL on ImageNet32, where both approaches are able to generate realistic images. In addition, we measure FID-10K [21] for iGPT and GCRL of 10M parameters trained on CIFAR10, where the score of iGPT is 45.02 $\pm$ 0.99 and the score of GCRL is 44.43 $\pm$ 0.36. Considering these observations with the small difference between bpds in Table 4, we can argue that GCRL maintains the generative performance of iGPT.

We expect that our hybrid scheme learns more robust features for out-of-distribution (OOD) samples than both generative and discriminative learning since it can combine the merits of the two methods. Specifically, the generative loss is helpful for OOD discrimination because it learns the data distribution directly, and contrastive loss also helps OOD discrimination to some extent because it pulls the representations of similar in-distribution samples together to form a cluster-like structure. To validate our hypothesis, we conduct two types of OOD detection tasks using the representations learned from each objective. First, we consider a supervised OOD detection setting with the OOD detection scores that are computed from class-conditional Gaussian densities [31, 51]: $s(\mathbf{x})_{\textrm{sup}}=\max_{c}\left\{{\mathcal{N}}\left({\tilde{\mathbf{z}}}(\mathbf{x})|\boldsymbol{\mu}_{c},\boldsymbol{\Sigma}_{c}\right)\right\}_{c=1}^{|C|}$, where ${\tilde{\mathbf{z}}}(\mathbf{x})$ is the representation from $\mathbf{x}$ without projection, $|C|$ is the number of classes, and the parameters $(\boldsymbol{\mu}_{c},\boldsymbol{\Sigma}_{c})$ are estimated from training samples of class $c$. Second, we consider an unsupervised OOD detection setting where OOD scores are computed by area of high probability region [16] approximated with the magnitude of the score functions, $s(\mathbf{x})_{\textrm{unsup}}=-\left\|\frac{\partial\log p(\mathbf{x})}{\partial\mathbf{x}}\right\|_{2}$. Since iGPT and GCRL quantize the input space, we cannot exactly compute this gradient with respect to the input image. ³³3Using a stochastic quantization [24] makes it possible to derive the input gradient.. Instead, we compute the gradient with respect the token embedding. Note that SimCLR is not applicable to the unsupervised OOD detection because it does not provide a density.

Table 5 presents the supervised OOD detection results where GCRL consistently outperforms the baselines in case of SVHN. Since STL-10 is a subset of ImageNet, it may be considered as an in-distribution set. However, we want to report how our approach behaves when a subset of in-distribution sets is given as an OOD dataset. The performance of OOD detection is measured by area under the ROC curve (AUROC), and area under the precision-recall curve (AUPRC) results are placed in Appendix. Figure 3 shows several distributions of $s(\mathbf{x})_{\textrm{sup}}$ of in/out-distribution test samples, supporting the results of Table 5. Table 6 shows the results of the unsupervised OOD detection where GCRL outperforms iGPT. To empirically analyze the reason why GCRL works better than baselines, we visualize the representations learned from each method in a two-dimensional space by $t$-SNE [49], as shown in Figure 4. We can observe that iGPT is able to discriminate in-distribution and out-of-distribution samples, but the in-distribution samples are not well clustered. On the other hand, SimCLR pulls visually-similar samples together in the embedding space, but it is not able to discriminate in-distribution and out-of-distribution samples. GCRL achieves both in/out discrimination and cluster-like structure.

This section includes several additional experiments to support our claims; (a) ablation study on the motivation of network design and low-shot transfer learning, and (b) the regularization effects of a generative loss in GCRL.

Table 7 shows an ablation study to support our network design choice, showing that when the number of blocks in encoder is half of the total number of blocks, the final performance is robust to the hyperparameters in the GCRL objective ($\alpha$ and $\beta$). It is noted that noticeable trade-offs between the two objectives can be observed when the hybrid loss is imposed directly to the same target representation from the final block, whereas the proposed separation of target representations can significantly reduce the trade-offs.

Table 8 shows the low-shot transfer classification accuracy from ImageNet32 to CIFAR10. We observe that GCRL improves the classification accuracy as well as expected calibration error (ECE) [18] by a large margin. We believe that the generative loss greatly improves the calibration of predictions and generalization ability because it prevents overfitting to some extent, especially when labeled samples are scarce.

Table 9 shows the regularization effect of the generative loss, where we deliberately use an easy pretext task for constrastive learning to cause overfitting. We adjust the random crop ratio from 0.2 to 0.8, meaning that at least 80% of an original image will be included in the cropped image. To understand the effect of such weak augmentations for SimCLR, Table 9 shows 5-NN classification losses of SimCLR and GCRL on test data during training. We observe that the loss of SimCLR starts increasing after 400 epochs, meaning that the model is overfitted to the pretext task, and thus hurting generalization. In addition, we observe that ECE of SimCLR at 800 epoch is higher than the one at 400 epoch, which also indicates the overfitting. Yet, GCRL does not suffer from such an overfitting and still performs comparable to the ones with strong augmentations thanks to the generative loss.

To show that GCRL prevents a classification model from overfitting, we conduct a supervised learning for classification with two baselines: ResNet-18 [20] and iGPT by replacing the generative loss into the supervised cross-entropy loss. In case of GCRL, we replace the contrastive loss into the supervised loss. For training ResNet, we employ a default augmentation policy [53], where a 32x32 crop is randomly selected from a padded image with reflection of 4 pixels on each side, followed by its horizontal flip of probability 0.5. For GCRL as well as iGPT, we use the same augmentation policy of SimCLR in section 4.2. Table 10 shows classification accuracy and ECE of compared methods on CIFAR10 and CIFAR100. Compared to iGPT, GCRL clearly achieves the better accuracy, proving that the regularization by a generative loss is effective. In addition, GCRL achieves very competitive performances compared to ResNet, even though ours has minimal inductive biases in network structure. Figure 5 presents reliability diagrams of iGPT-10M and GCRL-(5+5)M models, showing that our approach provides well-calibrated predictions than the baseline.