i -M IX : A S TRATEGY FOR R EGULARIZING C ONTRASTIVE R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Contrastive representation learning has shown to be an effective way of learning representations from unlabeled data.
However, much progress has been made in vision domains relying on data augmentations carefully designed using domain knowledge.
In this work, we propose i -Mix, a simple yet effective regularization strategy for improving contrastive representation learning in both vision and nonvision domains.
We cast contrastive learning as training a non-parametric classiﬁer by assigning a unique virtual class to each data in a batch.
Then, data instances are mixed in both the input and virtual label spaces, providing more augmented data during training.
In experiments, we demonstrate that i -Mix consistently improves the quality of self-supervised representations across domains, resulting in signiﬁcant performance gains on downstream tasks. Furthermore, we conﬁrm its regularization effect via extensive ablation studies across model and dataset sizes.
The code will be released.
I NTRODUCTION
Representation learning (Bengio et al., 2013) is a fundamental task in machine learning since the success of machine learning relies on the quality of representation.
Self-supervised representation learning (SSL) has been successfully applied in several domains, including image recognition (He et al., 2020;
Chen et al., 2020a), natural language processing (Mikolov et al., 2013; Devlin et al., 2018), robotics (Sermanet et al., 2018; Lee et al., 2019), speech recognition (Ravanelli et al., 2020), and video understanding (Korbar et al., 2018; Owens & Efros, 2018).
Since no label is available in the unsupervised setting, pretext tasks are proposed to provide self-supervision: for example, context prediction (Doersch et al., 2015), inpainting (Pathak et al., 2016), and contrastive learning (Wu et al., 2018b; Hjelm et al., 2019;
He et al., 2020; Chen et al., 2020a).
SSL has also been used as an auxiliary task to improve the performance on the main task, such as generative model learning (Chen et al., 2019), semi-supervised learning (Zhai et al., 2019), and improving robustness and uncertainty (Hendrycks et al., 2019).
Recently, contrastive representation learning has gained increasing attention by showing state-ofthe-art performance in SSL for large-scale image recognition (He et al., 2020;
Chen et al., 2020a), which outperforms its supervised pre-training counterpart (He et al., 2016) on downstream tasks.
However, while the concept of contrastive learning is applicable to any domains, the quality of learned representations rely on the domain-speciﬁc inductive bias: as anchors and positive samples are obtained from the same data instance, data augmentation introduces semantically meaningful variance for better generalization.
To achieve a strong, yet semantically meaningful data augmentation, domain knowledge is required, e.g., color jittering in 2D images or structural information in video understanding.
Hence, contrastive representation learning in different domains requires an effort to develop effective data augmentations.
Furthermore, while recent works have focused on largescale settings where millions of unlabeled data is available, it would not be practical in real-world applications. For example, in lithography, acquiring data is very expensive in terms of both time and cost due to the complexity of manufacturing process (Lin et al., 2018; Sim et al., 2019).
Meanwhile, MixUp (Zhang et al., 2018) has shown to be a successful data augmentation for supervised learning in various domains and tasks, including image classiﬁcation (Zhang et al., 2018), generative model learning (Lucas et al., 2018), and natural language processing (Guo et al., 2019; Guo, 2020).
In this paper, we explore the following natural, yet important question: is the idea of MixUp useful for unsupervised, self-supervised, or contrastive representation learning across different domains?
To this end, we propose instance Mix ( i -Mix) , a data-driven augmentation for contrastive representation learning effective across different domains.
The key idea of i -Mix is to introduce virtual labels in a batch and mix data instances and their corresponding virtual labels in the input and label spaces, respectively.
We ﬁrst introduce the general formulation of i -Mix, and then we show the applicability of i -Mix to state-of-the-art contrastive representation learning methods, SimCLR (Chen et al., 2020a) and MoCo (He et al., 2020), and a method without negative pairs, BYOL (Grill et al., 2020).
Through the experiments, we demonstrate the efﬁcacy of i -Mix in a variety of settings. First, we show the effectiveness of i -Mix by evaluating the discriminative performance of learned representations in multiple domains.
Speciﬁcally, we adapt i -Mix to the contrastive representation learning methods, advancing state-of-the-art performance across different domains, including image (Krizhevsky & Hinton, 2009; Deng et al., 2009), speech (Warden, 2018), and tabular (Asuncion & Newman, 2007) datasets.
Then, we study i -Mix in various conditions, including when 1) the model and training dataset is small or large, 2) domain knowledge is limited, and 3) transfer learning.
Contribution.
In summary, our contribution is three-fold:
• We propose i -Mix, a method for regularizing contrastive representation learning, motivated by MixUp (Zhang et al., 2018).
We show how to apply i -Mix to state-of-the-art contrastive representation learning methods (Chen et al., 2020a;
He et al., 2020; Grill et al., 2020).
• We show that i -Mix consistently improves contrastive representation learning in both vision and non-vision domains.
In particular, the discriminative performance of representations learned with i -Mix is on par with fully supervised learning on CIFAR-10/100 (Krizhevsky & Hinton, 2009) and Speech Commands (Warden, 2018).
• We verify the regularization effect of i -Mix in a variety of settings.
We empirically observed that i -Mix signiﬁcantly improves contrastive representation learning when 1) the training dataset size is small, or 2) the domain knowledge for data augmentations is not enough.
R ELATED W ORK
Self-supervised representation learning (SSL) aims at learning representations from unlabeled data by solving a pretext task that is derived from self-supervision.
Early works on SSL proposed pretext tasks based on data reconstruction by autoencoding (Bengio et al., 2007), such as context prediction (Doersch et al., 2015) and inpainting (Pathak et al., 2016).
Decoder-free SSL has made a huge progress in recent years.
Exemplar CNN (Dosovitskiy et al., 2014) learns by classifying individual instances with data augmentations.
SSL of visual representation, including colorization (Zhang et al., 2016), solving jigsaw puzzles (Noroozi & Favaro, 2016), counting the number of objects (Noroozi et al., 2017), rotation prediction (Gidaris et al., 2018), next pixel prediction (Oord et al., 2018; H´enaff et al., 2019), and combinations of them (Doersch & Zisserman, 2017; Kim et al., 2018; Noroozi et al., 2018) often leverages image-speciﬁc properties to design pretext tasks.
Meanwhile, alhough deep clustering (Caron et al., 2018; 2019; Asano et al., 2020) is often distinguished from SSL, it also leverages clustering assignments as self-supervision for representation learning.
Contrastive representation learning has gained lots of attention for SSL (He et al., 2020;
Chen et al., 2020a).
As opposed to early works on exemplar CNN (Dosovitskiy et al., 2014; 2015), contrastive learning maximizes similarities of positive pairs while minimizes similarities of negative pairs instead of training an instance classiﬁer.
As the choice of negative pairs is crucial for the quality of learned representations, recent works have carefully designed them.
Memory-based approaches (Wu et al., 2018b; Hjelm et al., 2019; Bachman et al., 2019; Misra & van der Maaten, 2020; Tian et al., 2020a) maintain a memory bank of embedding vectors of instances to keep negative samples, where the memory is updated with embedding vectors extracted from previous batches.
In addition, MoCo (He et al., 2020) showed that differentiating the model for anchors and positive/negative samples is effective, where the model for positive/negative samples is updated by the exponential moving average of the model for anchors.
On the other hand, recent works (Ye et al., 2019;
Misra & van der Maaten, 2020; Chen et al., 2020a; Tian et al., 2020a) showed that learning invariance to different views is important in contrastive representation learning.
The views can be generated through data augmentations carefully designed using domain knowledge (Chen et al., 2020a), splitting input channels (Tian et al., 2020a), or borrowing the idea of other pretext tasks, such as creating jigsaw puzzles or rotating inputs (Misra & van der Maaten, 2020).
In particular, SimCLR (Chen et al., 2020a) showed that a simple memory-free approach with a large batch size and strong data augmentations has a comparable performance to memory-based approaches.
InfoMin (Tian et al., 2020b) further studied a way to generate good views for contrastive representation learning and achieved state-of-the-art performance by combining prior works.
Different from other contrastive representation learning methods, BYOL (Grill et al., 2020) does not require negative pairs, where the proposed pretext task aims at predicting latent representations of one view from another.
While prior works have focused on SSL on large-scale visual recognition tasks, our work focuses on contrastive representation learning in both small- and large-scale settings in different domains.
Data augmentation is a technique to increase the diversity of data, especially when training data are not enough for generalization.
Since the augmented data must be understood as the original data, data augmentations are carefully designed using the domain knowledge about images (DeVries & Taylor, 2017b; Cubuk et al., 2019a;b; Zhong et al., 2020), speech (Amodei et al., 2016; Park et al., 2019), or natural languages (Zhang et al., 2015; Wei & Zou, 2019).
Some works have studied data augmentation with less domain knowledge: DeVries & Taylor (2017a) proposed a domain-agnostic augmentation strategy by ﬁrst encoding the dataset and then applying augmentations in the feature space.
MixUp (Zhang et al., 2018) is an effective data augmentation strategy in supervised learning, which performs vicinal risk minimization instead of empirical risk minimization, by linearly interpolating input data and their labels on the data and label spaces, respectively.
On the other hand, MixUp has also shown its effectiveness in other tasks and non-vision domains, including generative adversarial networks (Lucas et al., 2018), improved robustness and uncertainty (Hendrycks et al., 2020), and sentence classiﬁcation in natural language processing (Guo, 2020; Guo et al., 2019).
Other variations have also been investigated by interpolating in the feature space (Verma et al., 2019) or leveraging domain knowledge (Yun et al., 2019).
MixUp would not be directly applicable to some domains, such as point clouds, but its adaptation can be effective (Harris et al., 2020).
i -Mix is a kind of data augmentation for better generalization in contrastive representation learning, resulting in better performances on downstream tasks.
A PPROACH
In this section, we review MixUp (Zhang et al., 2018) in supervised learning and present i -Mix in contrastive learning (He et al., 2020; Chen et al., 2020a; Grill et al., 2020).
Throughout this section, let X be a data space, R D be a D -dimensional embedding space, and a model f : X → R D be a mapping between them.
For conciseness, f i = f
( x i ) and ˜ f i = f
(˜ x i ) for x i , ˜ x i ∈ X , and model parameters are omitted in loss functions.
3. M IX U P IN S UPERVISED L EARNING
Suppose an one-hot label y i ∈ { 0 , 1 } C is assigned to a data x i , where C is the number of classes.
Let a linear classiﬁer predicting the labels consists of weight vectors { w 1 , . .
. , w C } , where w c ∈ R D .
Then, the cross-entropy loss for supervised learning is deﬁned as:
[Equation]
While the cross-entropy loss is widely used for supervised training of deep neural networks, there are several challenges of training with the cross-entropy loss, such as preventing overﬁtting or networks being overconﬁdent.
Several regularization techniques have been proposed to alleviate these issues, including label smoothing (Szegedy et al., 2016), adversarial training (Miyato et al., 2018), and conﬁdence calibration (Lee et al., 2018).
MixUp (Zhang et al., 2018) is an effective regularization with negligible computational overhead.
It conducts a linear interpolation of two data instances in both input and label spaces and trains a model by minimizing the cross-entropy loss deﬁned on the interpolated data and labels.
Speciﬁcally, for two labeled data ( x i , y i ) , ( x j , y j ) , the MixUp loss is deﬁned as follows:
[Equation]
where λ ∼ Beta ( α, α ) is a mixing coefﬁcient from a beta distribution.
MixUp is a vicinal risk minimization method (Chapelle et al., 2001) that augments data and their labels in a data-driven manner.
Not only improving the generalization on the supervised task, it also improves adversarial robustness (Zhang et al., 2018; Pang et al., 2019) and conﬁdence calibration (Thulasidasan et al., 2019).
3. i -M IX IN C ONTRASTIVE L EARNING
We introduce instance mix ( i -Mix) , a data-driven augmentation strategy for contrastive representation learning to improve the generalization of learned representations.
Intuitively, instead of mixing class labels, i -Mix interpolates their virtual labels, which indicates their identity in a batch.
Let B = { ( x i , ˜ x i ) } Ni =1 be a batch of data pairs, where N is the batch size, x i , ˜ x i ∈
X are two views of the same data, which are usually generated by different augmentations.
For each anchor x i , we call ˜ x
i and ˜ x j (cid:54) =
i positive and negative samples, respectively.
Then, the model f learns to maximize similarities of positive pairs (instances from the same data) while minimize similarities of negative pairs (instances from different data) in the embedding space.
The output of f is L2-normalized, which has shown to be effective (Wu et al., 2018a;
He et al., 2020; Chen et al., 2020a).
Let v i ∈ { 0 , 1 } N be the virtual label of x i and ˜ x i in a batch B , where v i,i = 1 and v i,j (cid:54)
=
i = 0 .
For a general sample-wise contrastive loss with virtual labels (cid:96) ( x i , v i ) , the i -Mix loss is deﬁned as follows:
[Equation]
where λ ∼ Beta ( α, α ) is a mixing coefﬁcient and Mix is a mixing operator, which can be adapted depending on target domains: for example, MixUp ( x i , x j ; λ )
=
λx
i + (1 − λ ) x j (Zhang et al., 2018) when data values are continuous, and CutMix ( x i , x j ; λ ) =
M λ (cid:12) x i + (1 − M λ ) (cid:12) x j
(Yun et al., 2019) when data values have a spatial correlation, where M λ is a binary mask ﬁltering out a region whose relative area is (1 − λ ) , and (cid:12) is an element-wise multiplication.
In the following, we show how to apply i -Mix to contrastive representation learning methods.
We use the MixUp operator for i -Mix formulations and experiments, unless otherwise speciﬁed.
SimCLR (Chen et al., 2020a) is a simple contrastive representation learning method without a memory bank, where each anchor has one positive sample and ( 2 N − 2 ) negative samples.
Let x N + i =
˜ x i for conciseness.
Then, the ( 2 N − 1 )-way discrimination loss is written as follows:
[Equation]
where τ is a temperature scaling parameter and s ( f, ˜ f )
= ( f (cid:62) ˜ f ) / (cid:107) f
(cid:107)(cid:107) ˜ f (cid:107) is the inner product of two
L2-normalized vectors.
In this formulation, i -Mix is not directly applicable because virtual labels are deﬁned differently for each anchor.
3 To resolve this issue, we simplify the formulation of SimCLR by excluding anchors from negative samples.
Then, with virtual labels, the N -way discrimination loss is written as follows:
[Equation]
where we call it the N-pair contrastive loss, as the formulation is similar to the N-pair loss in the context of metric learning (Sohn, 2016).
For two data instances ( x i , v i ) , ( x j , v j ) and a batch of data pairs B = { ( x i , ˜ x i ) } Ni =1 , the i -Mix loss is deﬁned as follows:
[Equation] [Equation]
We present the application of i -Mix to the original SimCLR formulation in Appendix A.
InfoNCE (Oord et al., 2018) is a similar loss inspired by the a, b = aug(x), aug(x) # two different views of input x lam = Beta(alpha, alpha).sample() # mixing coefficient randidx = randperm(len(x)) a = lam * a + (1-lam) * a[randidx] logits = matmul(normalize(model(a)), normalize(model(b)).T) / t loss =
lam * CrossEntropyLoss(logits, arange(len(x))) + \ (1-lam) * CrossEntropyLoss(logits, randidx) Algorithm 1 provides the pseudocode of i -Mix on N-pair contrastive learning for one iteration.
Pair relations in contrastive loss.
To use contrastive loss for representation learning, one needs to properly deﬁne a pair relation { ( x i , ˜ x i ) } Ni =1 .
For contrastive representation learning, where semantic class labels are not provided, the pair relation would be deﬁned in that 1) a positive pair, x i and ˜ x i , are different views of the same data and 2) a negative pair, x i and ˜ x j (cid:54)
=
i , are different data instances.
For supervised representation learning, x i and ˜ x i are two data instances from the same class, while x i and ˜ x j (cid:54)
=
i are from different classes.
Note that two augmented versions of the same data also belong to the same class, so they can also be considered as a positive pair.
i -Mix is not limited to self-supervised contrastive representation learning, but it can also be used as a regularization method for supervised contrastive representation learning (Khosla et al., 2020) or deep metric learning (Sohn, 2016; Movshovitz-Attias et al., 2017).
MoCo (He et al., 2020).
In contrastive representation learning, the number of negative samples affects the quality of learned representations (Arora et al., 2019).
Because SimCLR mines negative samples in the current batch, having a large batch size is crucial, which often requires a lot of computational resources (Chen et al., 2020a).
For efﬁcient training, recent works have maintained a memory bank M = { µ k } Kk =1 , which is a queue of previously extracted embedding vectors, where K is the size of the memory bank (Wu et al., 2018b;
He et al., 2020; Tian et al., 2020a;b).
In addition,
MoCo introduces an exponential moving average (EMA) model to extract positive and negative embedding vectors, whose parameters are updated as θ f EMA ← mθ f EMA + (1 − m ) θ f , where m ∈ [0 , 1) is a momentum coefﬁcient and θ is model parameters.
The loss is written as follows:
[Equation]
The memory bank M is then updated with { ˜ f EMA i } in the ﬁrst-in ﬁrst-out order.
In this ( K +1 )-way discrimination loss, data pairs are independent to each other, such that i -Mix is not directly applicable because virtual labels are deﬁned differently for each anchor.
To overcome this issue, we include the positive samples of other anchors as negative samples, similar to the N-pair contrastive loss in
Eq.
Let ˜ v i ∈ { 0 , 1 } N + K be a virtual label indicating the positive sample of each anchor, where ˜ v i,i = 1 and ˜ v i,j (cid:54)
=
i = 0 .
Then, the ( N + K )-way discrimination loss is written as follows:
[Equation]
As virtual labels are bounded in the same set in this formulation, i -Mix is directly applicable: for two data instances ( x i , ˜ v i ) , ( x j , ˜ v j ) , a batch of data pairs B = { ( x i , ˜ x i ) } Ni =1 , and the memory bank M , the i -Mix loss is deﬁned as follows:
[Equation]
BYOL.
Different from other contrastive representation learning methods, BYOL is a self-supervised representation learning method without contrasting negative pairs. For two views of the same data λ ) j , i + λ )
j , j ) , i. ., optimization to the mixed label is equivalent to joint optimization to original labels. The proof for losses in contrastive learning methods is provided in Appendix B. x i , ˜ x i ∈ X , the model f learns to predict a view embedded with the EMA model ˜
f EMA i from its embedding f i .
Speciﬁcally, an additional prediction layer g is introduced, such that the difference between g ( f i ) and ˜
f EMA i is learned to be minimized.
The BYOL loss is written as follows:
[Equation]
This formulation can be represented in the form of the general contrastive loss in Eq.
(3), as the second view ˜ x i can be accessed from the batch B with its virtual label v i .
To derive i -Mix in BYOL, let ˜ F
=
[ ˜ f 1 / (cid:107) ˜
f
1 (cid:107) , . . ., ˜ f N / (cid:107) ˜
f N (cid:107) ]
∈ R D
× N be the collection of L2-normalized embedding vectors of the second views, such that ˜ f
i / (cid:107) ˜ f
i (cid:107) = ˜
Fv i .
Then, the BYOL loss is written as follows:
[Equation]
For two data instances ( x i , v i ) , ( x j , v j ) and a batch of data pairs B = { ( x i , ˜ x i ) } Ni =1 , the i -Mix loss is deﬁned as follows:
[Equation]
3. I NPUT M IX
The contribution of data augmentations to the quality of learned representations is crucial in contrastive representation learning.
For the case when the domain knowledge about efﬁcient data augmentations is limited, we propose to apply InputMix together with i -Mix, which mixes input data but not their labels.
This method can be viewed as introducing structured noises driven by auxiliary data to the principal data with the largest mixing coefﬁcient λ , and the label of the principal data is assigned to the mixed data. We applied InputMix and i -Mix together on image datasets in Table 3.
E XPERIMENTS
In this section, we demonstrate the effectiveness of i -Mix.
In all experiments, we conduct contrastive representation learning on a pretext dataset and evaluate the quality of representations via supervised classiﬁcation on a downstream dataset. We report the accuracy averaged over up to ﬁve runs.
In the ﬁrst stage, a convolutional neural network (CNN) or multilayer perceptron (MLP) followed by the two-layer MLP projection head is trained on an unlabeled dataset.
Then, we replace the projection head with a linear classiﬁer and train only the linear classiﬁer on a labeled dataset for downstream task.
Except for transfer learning, datasets for the pretext and downstream tasks are the same.
For i -Mix, we sample a mixing coefﬁcient λ ∼ Beta ( α, α ) for each data, where α = 1 for image and speech datasets and α = 2 for tabular datasets unless otherwise stated.
6 Additional details for the experimental settings and more experiments can be found in Appendix C.
4. E XPERIMENTAL S ETUP
Baselines and datasets.
We consider 1) N-pair contrastive learning as a memory-free contrastive learning method, 7 2) MoCo v2 (He et al., 2020; Chen et al., 2020b) 8 as a memory-based contrastive learning method, and 3) BYOL (Grill et al., 2020), which is a self-supervised learning method without negative pairs.
We apply i -Mix to these methods and compare their performances.
To show the effectiveness of i -Mix across domains, we evaluate the methods on datasets from multiple domains, including image, speech, and tabular datasets.
CIFAR-10/100 (Krizhevsky & Hinton, 2009) consist of 50k training and 10k test images, and ImageNet (Deng et al., 2009) has 1.3M training and 50k validation images, where we use them for evaluation. For some experiments, we use ImageNet-100, a subset with randomly chosenclasses out of 1k classes.
We apply a set of data augmentations randomly in sequence including random resized cropping, horizontal ﬂipping, color jittering, gray scaling, and Gaussian blurring for Domain Dataset [Table]
Image
[Table] Speech Commands [Table]
Tabular
[Table]
[Figure] ImageNet, which has shown to be effective (Chen et al., 2020a;b). We use ResNet-50 (He et al., 2016) as a backbone network. Models are trained with a batch size of 256 (i.e., 512 including augmented data) for up to 4000 epochs on CIFAR-10 and 100, and with a batch size of 512 for 800 epochs on
ImageNet. For ImageNet experiments, we use the CutMix (Yun et al., 2019) version of i -Mix.
The Speech Commands dataset (Warden, 2018) contains 51k training, 7k validation, and 7k test data in 12 classes.
We apply a set of data augmentations randomly in sequence: changing amplitude, speed, and pitch in time domain, stretching, time shifting, and adding background noise in frequency domain.
Augmented data are then transformed to a 32 × 32 mel spectogram.
We use the same architecture with image experiments.
Models are trained with a batch size of 256 for 500 epochs.
For tabular dataset experiments, we consider Forest Cover Type (CovType) and Higgs Boson (Higgs) from UCI repository (Asuncion & Newman, 2007).
CovType contains 15k training and validation, and 566k test data in 7 classes, and Higgs contains 10.5M training and 0.5M test data for binary classiﬁcation. For Higgs, we also use a subset of 100k training data to experiment at a different scale.
Since the domain knowledge for data augmentations on tabular data is limited, only a masking noise with the probability 0.2 is considered as a data augmentation.
We use a 5-layer MLP with batch normalization (Ioffe & Szegedy, 2015) as a backbone network.
Models are trained with a batch size of 512 for 1000 epochs.
4. M AIN R ESULTS
Table 1 shows the wide applicability of i -Mix to state-of-the-art contrastive representation learning methods for multiple domains.
i -Mix results in consistent improvements on the classiﬁcation accuracy, e.g., up to 6.1% when i -Mix is applied to MoCo v2 on CIFAR-100.
Interestingly, we observe that linear classiﬁers on top of representations learned with i -Mix without ﬁne-tuning the pre-trained part often yield a classiﬁcation accuracy on par with simple end-to-end supervised learning from random initialization, 9 e.g., i -Mix vs. end-to-end supervised learning performance is 96.1% vs. 95.5% on
CIFAR-10, 78.9% for both on CIFAR-100, and 98.2% vs. 98.0% on Speech Commands.
4. R EGULARIZATION E FFECT OF i -M IX
A better regularization method often beneﬁts from longer training of deeper models, which is more critical when training on a small dataset. To investigate the regularization effect of i -Mix, we ﬁrst
[Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Figure]
[Table]
make a comparison between N-pair contrastive learning and i -Mix by training with different model sizes and number of training epochs on the pretext task.
We train ResNet-18, 50, 101, and 152 models with varying number of training epochs from 500 to 4000.
Figure 2 shows the performance of N-pair contrastive learning (solid box) and i -Mix (dashed box).
When models are trained for 500 epochs, the performance of i -Mix is just on par with the baseline.
However, i -Mix improves the performance by signiﬁcant margins when trained longer.
This is because i -Mix introduces more variance to the pretext dataset, such that it requires more update to learn them.
In addition, deeper models beneﬁt from strong regularization of i -Mix, achieving 96.1% on CIFAR-10 and 75.6% on CIFAR-100 using
ResNet-152 after 4000 epochs of training.
On the other hand, models trained without i -Mix start to show decrease in performance, possibly due to overﬁtting to the pretext task when trained longer.
The trend clearly shows that i -Mix results in better representations via improved regularization.
Next, we study the effect of i -Mix with varying dataset sizes for the pretext tasks.
Table 2 shows the effect of i -Mix on large-scale datasets 10 from image and tabular domains.
We observe that i -Mix is particularly effective when the amount of training data is reduced, e.g., ImageNetconsists of images from 100 classes, thus has only 10% of training data compared to ImageNet-1k.
However, the performance gain is reduced when the amount of training data is large.
we further study representations learned with different pretext dataset sizes from 1% to 100% of the ImageNet training data in Figure 3. Here, different from ImageNet-100, we reduce the amount of data for each class, but maintain the number of classes the same.
We observe that the performance gain by i -Mix is more signiﬁcant when the size of the pretext dataset is small. Our study suggests that i -Mix is effective for regularizing self-supervised representations when training from a limited amount of data. We believe that this is aligned with ﬁndings in Zhang et al.
(2018) for MixUp in supervised learning. Finally, when large-scale unlabeled data is available, we expect i -Mix would still be useful in obtaining better representations when trained longer with deeper and larger models.
4. C ONTRASTIVE L EARNING WITHOUT D OMAIN -S PECIFIC D ATA A UGMENTATION
Data augmentations play a key role in contrastive representation learning, and therefore it raises a question when applying them to domains with a limited or no knowledge of such augmentations.
In this section, we study the effectiveness of i -Mix as a data-driven augmentation strategy for contrastive representation learning, which can be adapted to different domains.
Table 3 shows the performance of N-pair contrastive learning and i -Mix with and without data augmentations.
We observe signiﬁcant performance gains with i -Mix when other data augmentations are not applied.
For example, compared to the accuracy of 92.5% on the downstream task when other data augmentations are applied, contrastive learning achieves only 17.0% when trained without any data augmentations.
This suggests that data augmentation is an essential part for the success of contrastive representation [Table] learning (Chen et al., 2020a).
However, i -Mix is able to learn meaningful representations without other data augmentations and achieves close to the accuracy of 80% on CIFAR-10. Moreover, we verify its effectiveness on other domains beyond the image domain.
For example, the performance improves from 62.4% to 72.7% on the Speech Commands dataset when we assume no other data augmentations are available. We also observe consistent improvements in accuracy for tabular datasets, even when the training dataset size is large. Although the domain knowledge for data augmentations is important to achieve state-of-the-art results, our demonstration shows the potential of i -Mix to be used for a wide range of application domains where domain knowledge is particularly limited. 4. T RANSFERABILITY OF i -M IX
In this section, we show the improved transferability of the representations learned with i -Mix.
The results are provided in Table 4.
First, we train linear classiﬁers with downstream datasets different from the pretext dataset used to train backbone networks and evaluate their performance, e.g.,
CIFAR-10 as pretext and CIFAR-100 as downstream datasets or vice versa.
We observe consistent performance gains when learned representations from one dataset are evaluated on classiﬁcation tasks of another dataset.
Next, we transfer representations trained on ImageNet to the PASCAL VOC object detection task (Everingham et al., 2010).
We follow the settings in prior works (He et al., 2020;
Chen et al., 2020b): the parameters of the pre-trained ResNet-50 are transferred to a Faster R-CNN detector with the ResNet50-C4 backbone (Ren et al., 2015), and ﬁne-tuned end-to-end on the {+VOC 
+}
+12 trainval dataset and evaluated on the VOC 07 test dataset.
We report the average precision (AP) averaged over IoU thresholds between 50% to 95% at a step of 5%, and AP 50 and AP 75 , which are AP values when IoU threshold is 50% and 75%, respectively.
Similar to Table 2, we observe small but consistent performance gains in all metrics.
Those results conﬁrm that i -Mix improves the quality of learned representations, such that performances on downstream tasks are improved.
C ONCLUSION
We propose i -Mix, a data-driven augmentation strategy applicable to a class of self-supervised learning.
The key idea of i -Mix is to introduce a virtual label to each data instance, and mix both inputs and the corresponding virtual labels.
We show that i -Mix is applicable to state-of-the-art self-supervised representation learning methods including SimCLR, MoCo, and BYOL, which consistently improves the performance in a variety of settings and domains.
Our experimental results indicate that i -Mix is particularly effective when the training dataset size is small or data augmentation is not available, each of which are prevalent in practice.
A M ORE A PPLICATIONS OF i -M IX
In this section, we introduce more variations of i -Mix. For conciseness, we use v i to denote virtual labels for different methods. We make the deﬁnition of v i for each application clear. A. i -M IX FOR S IM CLR
For each anchor, SimCLR takes other anchors as negative samples such that the virtual labels must be extended.
Let x N + i =
˜ x i for conciseness, and v i ∈ { 0 , 1 } 2 N be the virtual label indicating the positive sample of each anchor, where v i,N + i = 1 and v i,j (cid:54)
=
N + i = 0 .
Note that v i,i = 0 because the anchor itself is not counted as a positive sample.
Then, Eq.
(4) can be represented in the form of the cross-entropy loss:
[Equation]
The application of i -Mix to SimCLR is straightforward: for two data instances ( x i , v i ) , ( x j , v j ) and a batch of data B = { x i } 2 Ni =1 , the i -Mix loss is deﬁned as follows:
[Equation]
Note that only the input data of Eq.
(A.2) is mixed, such that f i in Eq.
(A.1) is an embedding vector of the mixed data while the other f n ’s are the ones of clean data.
Because both clean and mixed data need to be fed to the network f , i -Mix for SimCLR requires twice more memory and training time compared to SimCLR when the same batch size is used.
A. i -M IX FOR S UPERVISED C ONTRASTIVE L EARNING
Supervised contrastive learning has recently shown to be effective for supervised representation learning and it often outperforms the standard end-to-end supervised classiﬁer learning (Khosla et al., 2020).
Suppose an one-hot label y i ∈ { 0 , 1 } C is assigned to a data x i , where C is the number of classes.
Let x N + i =
˜ x i and y N + i = y i for conciseness.
For a batch of data pairs and their labels
{+B = { ( x i , y i ) } 2 Ni =1 , let v i ∈ { 0 , 1 } 2 N be the virtual label indicating the positive samples of each anchor, where v i,j = 
+}
if y i =
y j (cid:54) =
i , and otherwise v i,j = 0 .
Intuitively, (cid:80) 2 Nj =1 v i,j = 2 N y i − 1 where N y i is the number of data with the label
y i .
Then, the supervised version of the SimCLR (SupCLR) loss function is written as follows:
[Equation]
The application of i -Mix to SupCLR is straightforward: for two data instances ( x i , v i ) , ( x j , v j ) and a batch of data B = { x i } 2 Ni =1 , the i -Mix loss is deﬁned as follows:
[Equation]
i -MixSupCLR (cid:0) ( x i
[Equation]
v i ) , ( x j , v j ); B
[Equation]
λ (cid:1) =
(cid:96) SupCLR ( λx i + (1 − λ )
x j
[Equation]
λv i + (1 − λ ) v j
[Equation]
B ) .
[Equation]
A.4)
A. i -M IX FOR N-P AIR S UPERVISED C ONTRASTIVE L EARNING
Note that i -Mix in Eq.
(A.4) is not as efﬁcient as SupCLR in Eq.
(A.3) due to the same reason in the case of SimCLR.
To overcome this, we reformulate SupCLR in the form of the N-pair loss (Sohn, 2016).
Suppose an one-hot label y i ∈ { 0 , 1 } C is assigned to a data x i , where C is the number of classes.
{+For a batch of data pairs and their labels B = { ( x i , ˜ x i , y i ) } Ni =1 , let v i ∈ { 0 , 1 } N be the virtual label indicating the positive samples of each anchor, where v i,j = 
+}
if y i =
y j (cid:54) =
i , and otherwise v i,j = 0 .
Then, the supervised version of the N-pair (Sup-N-pair) contrastive loss function is written as follows:
[Equation]
Then, the i -Mix loss for Sup-N-pair is deﬁned as follows:
[Equation]
B P ROOF OF THE LINEARITY OF LOSSES WITH RESPECT TO VIRTUAL LABELS
Cross-entropy loss.
The loss used in contrastive representation learning works, which is often referred to as InfoNCE (Oord et al., 2018), can be represented in the form of the cross-entropy loss as we showed for N-pair contrastive learning, SimCLR (Chen et al., 2020a)
[Equation]
and MoCo (He et al., 2020).
Here we provide an example in the case of N-pair contrastive learning.
Let f λij =
f
( λx i + (1 − λ ) x j ) for conciseness.
[Equation]
L2 loss between L2-normalized feature vectors.
The BYOL (Grill et al., 2020) loss is in this type.
Let ˜ F =
[ ˜ f 1 / (cid:107) ˜
f
1 (cid:107) , . . ., ˜ f N / (cid:107) ˜
f N (cid:107) ]
∈ R D × N such that ˜
f
i / (cid:107) ˜ f
i (cid:107) = ˜
Fv i , and ¯ g = g ( f ( λx i + (1 − λ ) x j )) / (cid:107) g
( f ( λx i + (1 − λ ) x j ))
(cid:107) for conciseness.
[Equation]
MixBYOL
[Equation]
=
[Equation]
=
[Equation]
const =
[Equation]
+ const .
(B.2)
Because ˜ F is not backpropagated, it can be considered as a constant.
C M ORE E XPERIMENTS C. S ETUP
In this section, we describe details of the experimental settings.
Note that the learning rate is scaled by the batch size (Goyal et al., 2017):
ScaledLearningRate = LearningRate
× BatchSize / 256 .
Image.
The experiments on CIFAR-10 and 100 (Krizhevsky & Hinton, 2009) and ImageNet (Deng et al., 2009) are conducted in two stages: following Chen et al.
(2020a), the convolutional neural network (CNN) part of ResNet-50 (He et al., 2016)
12 followed by the two-layer multilayer perceptron (MLP) projection head (output dimensions are 2048 and 128, respectively) is trained on the unlabeled pretext dataset with a batch size of 256 (i.e., 512 augmented data) with the stochastic gradient descent (SGD) optimizer with a momentum of 0.9 over up to 4000 epochs.
BYOL has an additional prediction head (output dimensions are the same with the projection head), which follows the projection head, only for the model updated by gradient.
10 epochs of warmup with a linear schedule to an initial learning rate of 0.125, followed by the cosine learning rate schedule (Loshchilov & Hutter, 2017) is used.
We use the weight decay of 0.0001 at the ﬁrst stage.
Then, the head of the CNN is replaced with a linear classiﬁer, and only the linear classiﬁer is trained with the labeled downstream dataset.
For N-pair contrastive learning and SimCLR, we use a batch size of 512 with the SGD optimizer with a momentum of 0.9 and an initial learning rate of 1 over epochs, where the learning rate is decayed by 0.2 after 80, 90, 95 epochs.
For MoCo and BYOL, we use the same learning schedule but an initial learning rate chosen among { 10, 30, 50, 70 } .
No weight decay is used at the second stage.
The quality of representation is evaluated by the top-1 accuracy on the downstream task.
We sample a single mixing coefﬁcient λ ∼ Beta (1 , 1) for each training batch.
The temperature is set to be τ = 0 .
5 for N-pair contrastive learning and SimCLR, and 0 .
2 for MoCo.
Note that the optimal distribution of λ and the optimal value of τ varies over different architectures, methods, and datasets, but the choices above result in a reasonably good performance.
The memory bank size is 4096 for MoCo, and the momentum for the exponential moving average (EMA) update is 0.999 for MoCo and BYOL.
We do not symmetrize the BYOL loss, as it does not change the performance while increasing computational complexity.
For data augmentation, we follow Chen et al.
(2020a):
We apply a set of data augmentations randomly in sequence including resized cropping (Szegedy et al., 2015), horizontal ﬂipping with a probability of 0.5, color jittering, 13 and gray scaling with a probability of 0.2.
A Gaussian blurring with σ ∈ [0 . 1 , 2] and kernel size of 10% of the image height/width is applied for ImageNet.
For evaluation on downstream tasks, we apply padded cropping with the pad size of 4 and horizontal ﬂipping for
CIFAR-10 and 100, and resized cropping and horizontal ﬂipping for ImageNet.
To improve the performance when domain-speciﬁc data augmentation is not available, InputMix is applied together with i -Mix on image datasets in Table 3.
For each principal data, we mixed two auxiliary data.
More speciﬁcally, ( 0 . 5 λ 1 + 0 .
5 λ 2 , 0 .
5 λ 3 ) are used as mixing coefﬁcients, where λ 1 , λ 2 , λ 3 ∼ Dir (1 , 1 , 1) where Dir is the Dirichlet distribution.
Speech.
In the experiments on Speech Commands (Warden, 2018), the network is the same with the image domain experiments, except that the number of input channels is one instead of three.
We use the temperature is set to be τ = 0 .
10% of silence data (all zero) are added when training.
At the ﬁrst stage, the model is trained with the SGD optimizer with a momentum of 0.9 and an initial learning rate of 0.0001 over 500 epochs, where the learning rate decays by 0.1 after 300 andepochs and the weight decay of 0.0001.
The other settings are the same with the experiments in the image domain.
At the second stage, the MLP head is replaced with a linear classiﬁer.
For N-pair contrastive learning, we use a batch size of 512 with the SGD optimizer with a momentum of 0.9 and an initial learning rate of 1 over 100 epochs, where the learning rate is decayed by 0.2 after 80, 90,epochs.
For MoCo and BYOL, we use the same learning schedule but an initial learning rate of 10.
No weight decay is used at the second stage.
For data augmentation, 14 we apply a set of data augmentations randomly in sequence including changing amplitude, speed, and pitch in time domain, stretching, time shifting, and adding background noise in frequency domain.
Each data augmentation is applied with a probability of 0.5.
Augmented data are then transformed to the mel spectogram in the size of 32 × 32.
Tabular.
In the experiments on CovType and Higgs (Asuncion & Newman, 2007), we take a ﬁvelayer MLP with batch normalization as a backbone network.
The output dimensions of layers are (2048-2048-4096-4096-8192), where all layers have batch normalization followed by ReLU except for the last layer.
The last layer activation is maxout (Goodfellow et al., 2013) with 4 sets, such that the output dimension is 2048.
On top of this ﬁve-layer MLP, we attach two-layer MLP (2048-128) as a projection head.
We sample a single mixing coefﬁcient λ ∼ Beta (2 , 2) for each training batch.
We use the temperature τ = 0 .
For N-pair contrastive learning, at the ﬁrst stage, the model is trained with a batch size of 512 with the Adam optimizer with beta values of 0.9 and 0.999, an initial learning rate of 5 × 10 − 5 and 2 .
5 × 10 − 5 over 1000 epochs for CovType and Higgs, respectively, and the weight decay of 0.0001.
For MoCo and BYOL with CovType, we use the same learning schedule but an initial learning rate of 0.0002 and 5 × 10 − 5 , respectively.
The other settings are the same with the experiments in the image domain.
At the second stage, the MLP head is replaced with a linear classiﬁer.
The classiﬁer is computed by linear regression from the feature matrix obtained without data augmentation to the label matrix using the pseudoinverse.
Since the prior knowledge on tabular data is very limited, only the masking noise with a probability of 0.2 is considered as a data augmentation.
[Table]
[Table]
C. V ARIATIONS OF i -M IX
In this section we additionally compare the MixUp (Zhang et al., 2018) and CutMix (Yun et al., 2019) variation of i -Mix on N-pair contrastive learning and SimCLR.
To distinguish them, we call them i -MixUp
and i -CutMix, respectively.
To be fair with the memory usage in the pretext task stage, we reduce the batch size of i -MixUp
and i -CutMix by half (256 to 128) for SimCLR.
Following the learning rate adjustment strategy in Goyal et al.
(2017), we also decrease the learning rate by half (0.125 to 0.0625) when the batch size is reduced.
We note that
i -MixUp
and i -CutMix on SimCLR take approximately 2.5 times more training time to achieve the same number of training epochs.
The results are provided in Table C.1.
We ﬁrst verify that the N-pair formulation results in no worse performance than that of SimCLR.
This justiﬁes to conduct experiments using the N-pair formulation instead of that of SimCLR, which is simpler and more efﬁcient, especially when applying i -Mix, while not losing the performance.
When pretext and downstream tasks share the training dataset, i -CutMix often outperforms i -MixUp, though the margin is small.
However, i -CutMix shows a worse performance in transfer learning.
Table C.2 compares the performance of SimCLR, N-pair contrastive learning, and i -Mix on N-pair contrastive learning when the pretext task is self-supervised and supervised contrastive learning.
We conﬁrm that the N-pair formulation results in no worse performance than that of SimCLR in supervised contrastive learning as well.
i -Mix improves the performance of supervised contrastive learning from 95.7% to 97.0% on CIFAR-10, similarly to improvement achieved by MixUp for supervised learning where it improves the performance of supervised classiﬁer learning from 95.5% to 96.6%.
On the other hand, when the pretext dataset is CIFAR-100, the performance of supervised contrastive learning is not better than that of supervised learning: MixUp improves the performance of supervised classiﬁer learning from 78.9% to 82.2%, and i -Mix improves the performance of supervised contrastive learning from 74.6% to 78.4%.
[Figure]
Pretext
Downstream
FED ( × 10 − 4 ) ( ↓ )
Training Acc (%) ( ↑ )
Test Acc (%) ( ↑ )
N-pair
+ i -Mix
N-pair
+ i -Mix
N-pair
+ i -Mix
CIFARCIFAR- 30.
CIFAR13.
CIFARCIFAR- 15.
CIFAR30.
learned representation is more generalized.
While supervised i -Mix improves the classiﬁcation accuracy on CIFAR-10 when trained on CIFAR10, the representation does not transfer well to CIFAR-100, possibly due to overﬁtting to 10 class classiﬁcation.
When pretext dataset is CIFAR-100, supervised contrastive learning shows a better performance than self-supervised contrastive learning regardless of the distribution shift, as it learns sufﬁciently general representation for linear classiﬁer to work well on CIFAR-10 as well.
C. Q UALITATIVE E MBEDDING A NALYSIS
Figure C.1 visualizes embedding spaces learned by N-pair contrastive learning and i -Mix on CIFAR10 and 100.
When the downstream dataset is the same with the pretext task, both contrastive learning and i -Mix cluster classes well, as shown in Figure C.1(a) and C.1(b).
However, when the downstream task is transferred to CIFAR-100, i -Mix in Figure C.1(d) clusters classes better than contrastive learning in Figure C.1(c).
Speciﬁcally, clusters of “apple,” “chair,” and “dolphin,” can be found in Figure C.1(d) while they spread out in Figure C.1(c).
Also, “rose” and “squirrel” are more separated in Figure C.1(d) than C.1(c).
This shows that the representation learned with i -Mix is more generalizable than vanilla contrastive learning.
C. Q UANTITATIVE E MBEDDING A NALYSIS
To estimate the quality of representation by the similarity between training and test data distribution, we measure the Fr´echet embedding distance (FED): similarly to the Fr´echet inception distance (FID) introduced in Heusel et al.
(2017), FED is the Fr´echet distance (Fr´echet, 1957; Vaserstein, 1969) between the set of training and test embedding vectors under the Gaussian distribution assumption.
For conciseness, let ¯ f i = f
( x i ) / (cid:107) f
( x i ) (cid:107) be an (cid:96) 2 normalized embedding vector; we normalize embedding vectors as we do when we measure the cosine similarity.
{+Then, with the estimated mean m = 1 N (cid:80) Ni =1 ¯ f i and the estimated covariance S = 1 N (cid:80) Ni =
+}
( ¯ f i − m )( ¯ f i − m ) (cid:62) , the FED can be deﬁned as
[Equation]
As shown in Table C.3, i -Mix improves FED over contrastive learning, regardless of the distribution shift.
Note that the distance is large when the training dataset of the downstream task is the same with that of the pretext task.
This is because the model is overﬁt to the training dataset, such that the distance from the test dataset, which is unseen during training, has to be large.
On the other hand, Table C.3 shows that i -Mix reduces the gap between the training and test accuracy.
This implies that i -Mix is an effective regularization method for pretext tasks, such that the learned representation is more generalizable in downstream tasks.