MetAug: Contrastive Learning via Meta Feature Augmentation

We recap the preliminaries of contrastive learning(Tian et al., 2019; Chen et al., 2020): the foundational idea behind contrastive learning is to learn an embedding that maximizes agreement between the views of the same sample and separates the views of different samples in latent space. Given a multi-view dataset $X$, we treat pairs of the views of the same sample $\{x_{i}^{j},x_{i}^{j^{\prime}}\}$, where $j,j^{\prime}\in\{1,...,M\}$, as positives, versus pairs of the views of the different samples $\{x_{i}^{j},x_{i^{\prime}}^{j^{\prime}}\}$, where $i\neq i^{\prime}$, as negatives. To impose contrastive learning, we feed the input $x_{i}^{j}$ into a view-specific encoder $f_{\theta_{j}}(\cdot)$ to learn a representation $h_{i}^{j}$, and $h_{i}^{j}$ is mapped into a feature $z_{i}^{j}$ by a projection head $g_{\vartheta_{j}}(\cdot)$, where $\theta_{j}$ and $\vartheta_{j}$ are the network parameters of $f_{\theta_{j}}(\cdot)$ and $g_{\vartheta_{j}}(\cdot)$, respectively. A discriminating function $d(\cdot)$ is adopted to measure the similarity of $\{z_{i}^{j},z_{i^{\prime}}^{j^{\prime}}\}$, where $i\neq i^{\prime}$. The encoder $f_{\theta_{j}}(\cdot)$ and projection head $g_{\vartheta_{j}}(\cdot)$ are trained by using a contrastive loss (Oord et al., 2018), which is formulated as follows:

where $X_{S}=\{\{z^{+}\},\{z^{-}\}_{1},\{z^{-}\}_{2},...,,\{z^{-}\}_{k}\}$ is a set of pairs randomly sampled from $X$, which includes a positive $\{z^{+}\}$ and $K$ negatives $\{z^{-}\}_{k}$, $k\in\{1,...,K\}$, because contrastive loss can only use one positive in an iteration. In test, the projection head $g_{\vartheta_{j}}(\cdot)$ is discarded, and the representation $h_{i}^{j}$ is directly used for downstream tasks.

Recent contrastive methods rely on complex data augmentations to increase the informativeness of views. Yet this lack of guidance approach leads to the demand for a large number of training data (e.g., large batch size and memory bank). We propose a meta feature augmentation method, which creates informative augmented features by updating parameters of its own network according to the performance (gradient) of the encoder (see Appendix A.3 for our rethinking of augmented features). A visualization of the overall MetAug architecture is shown in Figure 1.

To this end, we build a group of MAGs $a_{\omega}(\cdot)=\{a_{\omega_{1}}(\cdot),...,a_{\omega_{M}}(\cdot)\}$ for all $M$ views, where $\omega=\{\omega_{1},...,\omega_{M}\}$. To be simplified, we define $f_{\theta}$ and $g_{\vartheta}$ as the groups of view-specific encoders and projection heads, respectively, i.e., $\theta=\{\theta_{1},...,\theta_{M}\}$ and $\vartheta=\{\vartheta_{1},...,\vartheta_{M}\}$.

In training, the encoders $f_{\theta}(\cdot)$ and the projection heads $g_{\vartheta}(\cdot)$ are trained alongside the MAG $a_{\omega}(\cdot)$ (with the network parameters $\omega$). Following the protocol of meta learning (Finn et al., 2017; Liu et al., 2019), we firstly train $f_{\theta}(\cdot)$ and $g_{\vartheta}(\cdot)$ under the learning paradigm of self-supervised contrastive learning. Then, $a_{\omega}(\cdot)$ is updated by computing its gradients with respect to the performance of $f_{\theta}(\cdot)$ and $g_{\vartheta}(\cdot)$. Here, we measure the performance of $f_{\theta}(\cdot)$ and $g_{\vartheta}(\cdot)$ by the gradients of them when the corresponding contrastive loss is back-propagated. Concretely, all of $f_{\theta}(\cdot)$, $g_{\vartheta}(\cdot)$, and $a_{\omega}(\cdot)$ are iteratively trained until convergence.

Specifically, we first update network parameters $\theta$ and $\vartheta$ of the encoders and projection heads by adopting the conventional contrastive loss. Then, we train the MAG $a_{\omega}(\cdot)$ in a meta learning manner. We encourage the augmented features to be informative, and the encoders $f_{\theta}(\cdot)$ can better explore the discriminative information by jointly using the original and augmented features to contrast. Hence, the performance of the encoders would be promoted on the same training data. To update network parameters $\omega$ of the $a_{\omega}(\cdot)$, we formalize the meta updating objective as follows:

where $\widetilde{X}$ represents a minibatch sampled from the training dataset $X$, $\Big{\{}g_{\mathring{\vartheta}}\big{(}f_{\mathring{\theta}}\big{(}\widetilde{X}\big{)}\big{)},a_{\omega}\Big{(}g_{\mathring{\vartheta}}\big{(}f_{\mathring{\theta}}\big{(}\widetilde{X}\big{)}\big{)}\Big{)}\Big{\}}$ denotes a set including both original features and meta augmented features. $\mathring{\vartheta}$ and $\mathring{\theta}$ represent the parameter sets of the encoders and projection heads, respectively, which are computed by the updating of one gradient back-propagation:

where $\ell$ is the learning rate shared between $\theta$ and $\vartheta$. The idea behind the meta updating objective is that we perform the second-derivative technique (Finn et al., 2017; Zhang et al., 2018; Liu et al., 2019) to train $a_{\omega}(\cdot)$. Specifically, a derivative over the derivative (Hessian matrix) of the combination $\{\theta,\vartheta\}$ is used to update $\omega$, where $\{\theta,\vartheta\}$ is a parameter set conjoining $\theta$ and $\vartheta$. We compute the derivative with respect to $\omega$ by using a retained computational graph of $\{\theta,\vartheta\}$.

However, in practice, we find a critical issue: when the original features are not informative enough, large gradients are difficult to generate by contrasting the uninformative features, the MAGs $a_{\omega}(\cdot)$ are inclined to create collapsed augmented features, e.g., the augmented features and the original features are very similar. We consider the reason for the feature collapse is that small gradient changes of the encoders alongside projection heads $g_{\vartheta}(f_{\theta}(\cdot))$ lead to the update step-size of $a_{\omega}(\cdot)$ to become extensively small, which leaves the optimization of $a_{\omega}(\cdot)$ to fall into a local optimum. The augmented features are such that without any extra useful information. To tackle this issue, we further inject a margin to encourage $a_{\omega}(\cdot)$ to generate more complex and informative augmented features, which can be considered as a regularization term in the meta updating objective. See Figure 2(a) for the details of the augmented feature collapse issue, and we observe that, without margin-injected regularization, MAGs tend to generate collapsed features that are very similar with the original features. Formally, we formulate the approach to generate margins for $a_{\omega}(\cdot)$ by

where $\big{\{}d(\{z^{+}\}_{k^{+}})\big{\}}$ is a set of the outputs (similarities) of positives computed by the discriminating function $d(\cdot)$, and $k^{+}\in\{1,...,K^{+}\}$ where $K^{+}$ represents the number of positives in a minibatch. $\Big{\{}d(\{z^{-}\}_{k^{-}})\Big{\}}$ is a set of the discriminating outputs of negatives, and $k^{-}\in\{1,...,K^{-}\}$ where $K^{-}$ represents the number of negatives. Note that only original features are used in Equation 4. We call the formulated margin generation approach as "Large", and we also propose two more approaches, called "Medium" and "Small". In Appendix A.2, we conduct comparisons to evaluate the effects of the three margin generation approaches. We inject the margins between the augmented features and original features by adding a regularization term in the meta updating objective, and the regularization is defined as:

where $\{\hat{z}^{+}\}_{\hat{k}^{+}}$ denotes a positive that includes one original feature and one augmented feature, and $\hat{K}^{+}$ denotes the number of such positives. $\{\hat{z}^{-}\}_{\hat{k}^{-}}$ represents likewise one of $\hat{K}^{-}$ negatives, each of which includes one original feature and one augmented feature. $[\cdot]_{+}$ denotes the cut-off-at-zero function, which is defined as $[a]_{+}=\max(a,0)$. We then integrate such regularization to the updating of $\omega$ by

where $\ell^{\prime}$ represents the learning rate of $\omega$, and $\alpha$ is a hyperparameter balancing the impact of the loss of margin-injected regularization term. $\mathcal{R_{\sigma}}$ restricts MAGs to generate informative features that are more different with the original features (see Figure 2(b)). In practice, Figure 2(c) and (d) show that the features learned by our method (with margin-injected regularization) are more concentrated, e.g., the features of the same image are more similar and the gap between the features of the different images are enlarged, which proves informative augmented features can further lead the encoders to learn non-collapsed (scattered) features.

We propose to jointly contrast all features (including the original features and the meta augmented features) in one gradient back-propagation step. Motivated by (Schroff et al., 2015), we introduce the following optimization-driven unified loss function to replace the conventional contrastive loss as follows:

where $[\cdot]_{+}$ ensures that $\mathcal{L}\geq 0$ is always held. Note that all original features and augmented features are involved. $\lambda$ is a margin between the summarized instance similarities to enhance the capability of the similarity separation. However, we find that the difference between $\sum_{k^{-}=1}^{K^{-}}d(\{z^{-}\}_{k^{-}})$ and $\sum_{k^{+}=1}^{K^{+}}d(\{z^{+}\}_{k^{+}})$ is not the larger the better. Excessive increases of the difference may undermine the convergence in optimization. We thereby wish to adopt a margin $\lambda$ that leads to preferable convergence. We reform the loss in Equation 7 by adding a temperature coefficient $\beta$ as follows:

when $\beta\to+\infty$, Equation 8 is Equation 7. Inspired by (Sun et al., 2020), we use weighting factors $\Gamma^{-}$ and $\Gamma^{+}$ to modulate the impacts of $d(\{z^{-}\}_{k^{-}})$ and $d(\{z^{+}\}_{k^{+}})$. Such approach aims to give greater weight to the similarity that deviates from the optimum and smaller weight to the similarity that has the close proximity with the optimum. $\Gamma^{-}=[d(\{z^{-}\}_{k^{-}})-O^{-}]_{+}$ and $\Gamma^{+}=[O^{+}-d(\{z^{+}\}_{k^{+}})]_{+}$, where $O^{-}$ and $O^{+}$ represents the expected optimums of $d(\{z^{-}\}_{k^{-}})$ and $d(\{z^{+}\}_{k^{+}})$. Note that, we further propose a variation of $\Gamma$ including $\Gamma^{-}$ and $\Gamma^{+}$, and the comparisons of them are demonstrated in Section 4.4. $\gamma^{+}$ and $\gamma^{-}$ is used to replace $\lambda$ and add $\Gamma^{-}$ and $\Gamma^{+}$ in Equation 8:

We limit $d(\{z^{-}\}_{k^{-}})$ and $d(\{z^{+}\}_{k^{+}})$ in the range of $[0,1]$ by normalizing the features in $\{z^{-}\}_{k^{-}}$ and $\{z^{-}\}_{k^{-}}$, such that theoretically, the optimum of $d(\{z^{-}\}_{k^{-}})$ is $0$, the optimum of $d(\{z^{+}\}_{k^{+}})$ is $1$. The positive of $d(\{z^{-}\}_{k^{-}})-O^{-}$ and $O^{+}-d(\{z^{+}\}_{k^{+}})$ can easily be guaranteed. To cut the number of hyperparameters, we reform Equation 9 into

which is derived by setting $O^{+}=1+\gamma$, $O^{-}=-\gamma$, $\gamma^{+}=1-\gamma$, and $\gamma^{-}=\gamma$.

Concretely, we adopt margin-injected meta feature augmentation in the contrastive learning paradigm to achieve desired discriminative multi-view representations, and the proposed $\mathcal{L}_{MetAug}$ is incorporated to replace the conventional contrastive loss $\mathcal{L}$. The final model objective is defined as:

where $\mathcal{L}_{OUCL}^{ori}$ represents the loss NOT including the meta augmented features, $\mathcal{L}_{OUCL}^{aug}$ represents the loss including such features, and $\delta$ is a coefficient that controls the balance between them (we perform parameter comparisons in Appendix A.4). It is worthy to note that the margin-injected regularization $\mathcal{R_{\sigma}}$ is only used in meta training the MAGs, i.e., $a_{\omega}(\cdot)$, while in regular training of encoders and projection heads, $\mathcal{R_{\sigma}}$ is discarded. $\mathcal{R_{\sigma}}$ restricts the augmented features to be informative so that such features can lead the encoder to efficiently and effectively learn discriminative representations. The training process is detailed by Algorithm 1.

Implementations. To efficiently perform CL within a restricted amount of the inputs in training, we uniformly set the batch size as 64 (see Appendix A.1 for the comparisons under different setting of batch size). For the experiments with conv and fc as the backbone networks, we adopt a network with the 5 convolutional layers in AlexNet (Krizhevsky et al., 2012) as conv and a network with further 2 fully connected layers as fc. Inspired by the backbone splitting setting of SplitBrain (Zhang et al., 2017), we evenly split the AlexNet into sub-networks across the channel dimension, and each sub-network is the view-specific encoder (see Appendix B for the detailed implementation). For the experiments with ResNet-50, we directly change the encoder network to ResNet-50. All backbone encoders are not pre-trained. MetAug (only OUCL) is the ablation variant without margin-injected meta feature augmentation.

Given an RGB image, we convert it to the Lab image color space and split it into L and ab channels. During contrastive learning, RGB, L, and ab are used as three views of the image. Before feeding the views into our model, we simply adopt the same data augmentations in CMC (Tian et al., 2019). Especially, the major contribution of DACL is the proposed data augmentation (i.e., mixup-noise) so that we particularly add mixup data augmentation for DACL. In training, a memory bank (Wu et al., 2018) is adopted to facilitate calculations. We retrieve 4096 past features from the memory bank to derive negatives. The learning rates and weight decay rates are uniform over comparisons.

Comparison on downstream tasks. We collect the results of 20 trials for comparisons. The average result of the last 20 epochs is used as the final result of each trial, and the 95% confidence intervals are also reported, while the results without 95% confidence intervals are quoted from the published papers. We compare MetAug against a fully-supervised method (similar to AlexNet (Krizhevsky et al., 2012)) and the state-of-the-art unsupervised methods. Table 1 shows the comparisons on four benchmark datasets. The last two rows of tables represent the results of our methods. As demonstrated in tables, MetAug beats the best prior methods on all datasets. Even compared with the fully-supervised method trained end-to-end (without fine-tuning) for the architecture presented, the proposed method has a significant improvement on most downstream tasks, which demonstrates that MetAug can better model discriminative information when supervision is insufficient (e.g., the training data is limited). The ablation model (i.e., MetAug (only OUCL)) outperforms most unsupervised methods but falls short of the performance of MetAug. Thus, the ablation study proves the effectiveness of our proposed margin-injected meta feature augmentation and optimization-driven unified contrast.

DACL and LooC propose to enhance contrastive learning from the perspective of data augmentation, while MetAug improves contrastive learning from the perspective of feature augmentation. The idea behind our method is simple but effective, since contrastive learning works directly on features, and the augmented images need one step of encoding to become features. The experimental results support that MetAug achieves better performance on benchmarks.

Performing MetAug on ResNet. We perform classification comparisons on the CIFAR10 and STL-10 by using ResNet-50. Table 2 shows that MetAug and the ablation variant outperform the compared methods, which indicates that MetAug has strong adaptability to different encoders.

Implementation. To comprehensively understand the performance of our proposed MetAug, we conduct comparisons on ImageNet and make fair comparisons with benchmark methods. The backbone encoder is conv or ResNet-50, and the results are demonstrated in Table 3. MetAug is a decoupled approach so that we can introduce MetAug in the learning paradigm of state-of-the-art to improve the performance, e.g, for the experiments using conv or ResNet-50, we perform MetAug in CMC or NNCLR, respectively.

Results. As shown in Table 3, we find that MetAug can effectively promote the performance of benchmark methods in the comparisons using both conv and ResNet-50. The results support that our proposed meta feature augmentation can enable different encoders to model discriminative information even in the large-scale dataset.

To illustrate the impacts of different data augmentations, we conducted multiple comparisons on CIFAR10 shown in Table 4. Note that horizontal flip and rotate are similar, and we use them together in the 1-th comparison. In the 5-th comparison, we take the same data augmentations as the setting of comparisons in Section 4.1. The data augmentations adopted in the 6-th comparison are as same as the setting of LooC (Xiao et al., 2021). Additionally, mixup is proposed by DACL (Verma et al., 2021).

We observe from Table 4 that MetAug outperforms the compared methods in all comparisons. It is worth noting that even using weak data augmentation degenerates the performance of our method as well as benchmark methods, but the performance degeneration of our method is minimal compared to others, e.g., from 8-th and 1-th comparison, we find that the gap of MetAug is 1.60%, while that of LooC is 2.44%. The results support that MetAug is robust for various data augmentations.

In practice, we find that the introduction of the weighting factor $\Gamma$ cannot directly improve our proposed method. Our conjuncture lies in that $\Gamma$ may cause the loss to converge excessively fast, which leaves the network parameters at a local minimum. Therefore, we propose a variant to replace $\Gamma$ in Equation 9, i.e., $\bar{\Gamma}=\dfrac{\Gamma}{\phi_{dec}}$ where $\phi_{dec}$ is a linear attenuation coefficient to linearly attenuate the impact of $\Gamma$ so that the difference between the current value and the optimum becomes smaller.

We use MetAug (only OUCL) to demonstrate the effectiveness of the proposed variant. The results are shown in Figure 3. We observe that the performance of our method get peak value when $\phi_{dec}$ is 6, which manifests that introducing a certain linear attenuation to $\Gamma$ can promote MetAug.