\section{Introduction}
\label{sec:intro}
Data augmentation is widely used in supervised learning for image recognition, achieving excellent results on popular datasets like MNIST, CIFAR-10/100, and ImageNet \cite{ciregan2012multi,sato2015apac,simard2003best,wan2013regularization,krizhevsky2017imagenet}.
AutoAugment~\cite{cubuk2019autoaugment} and RandAugment~\cite{cubuk2020randaugment} use strong augmentations composed of a combination of different transformations. It has been shown that these manually designed data augmentations can effectively increase the amount of training data and enhance the model's ability to capture the invariance of the training data, thereby boosting supervised training performance.

Data augmentation (DA) is also a key component in recent semi-supervised learning techniques based on contrastive learning \cite{chen2020simple}.   A base encoder that learns good latent representations of the input data is trained with a contrastive loss. The contrastive loss is reflective of the following:  in the latent space a pair of two views of a given data example transformed by two different data augmentation functions are correlated (similar) while transformed views of different input examples are dissimilar. The quality of the encoder trained using unlabeled data as such is essential to the overall contrastive learning performance, and is contingent upon the employed data augmentations.

 ~\cite{wang2020understanding} interprets the minimization of the contrastive loss  is to achieve alignment and uniformity of the learned representations. To learn powerful and transferable representations, many studies have aimed at improving contrastive learning by choosing proper data augmentations  \cite{chen2020simple,tian2020makes,he2020momentum} have shown that combining different augmentations with a proper magnitude may improve performance. In these works, data augmentations are restricted to be a random composition of four or five types of augmentations with a very limited range of magnitude. To further improve contrastive learning in downstream tasks, PIRL~\cite{misra2020self} adopts two additional data augmentations, SwAV~\cite{caron2020unsupervised} introduced multiple random resized-crops to provide the encoder with more views of data and CLSA~\cite{wang2021contrastive} adopts combinations of stronger augmentations.

While exploiting a diverse set of data augmentations is appealing, thus far, there does not exist a systematic approach  for incorporating a large number of augmentations for contrastive learning. Importantly, the intricacies of data augmentation and representation learning may result in performance degradation if data augmentation functions are not properly chosen and used. To address the above challenges, this paper proposes \textbf{Co}ntrastive Learning with \textbf{Co}nsistent \textbf{R}epresentations (CoCor). First, we define a set of  \textit{composite augmentations}, denoted by $\bm{A}$ and composed by combining multiple basic augmentations.  We represent each composite augmentation   using a \emph{composition vector} that encodes the types and numbers of basic augmentations used and their augmentation magnitude.
A large number of diverse augmentations can be drawn from this composite set  $\bm{A}$ with widely varying transformation types and overall strength.

Critically, we propose a notion of \emph{consistency} for data augmentations, called \emph{DA consistency}, which imposes a desiring property on the latent representations of views produced by different composite data augmentations. For a given input example, a stronger data augmentation shall push its latent-space view further away from the latent representation of the original example. That is, the latent-space similarity between the transformed view and the original input decreases with the strength of the data augmentation, a very desiring characteristic, as illustrated in \cref{fig:overview}(b). In other words, a base encoder that does not satisfy the above property is deemed \emph{inconsistent} from a data augmentation point of view.
We consequently propose to add a new \textit{consistency loss} to impose  \emph{DA consistency} on  the base encoder. Different from the contrastive loss that is agnostic to the type and overall strength of employed data augmentations ~\cite{chen2020simple,tian2020makes,he2020momentum}, our \emph{consistency loss} places additional finer-grained constraints on the encoder learning such that the resulted latent-space views are consistent with the strength of data augmentations applied.
Enforcing \emph{DA consistency} is key to circumventing inconsistencies among diverse data augmentations, which  may result in performance degradation.

\begin{figure*}
    \centering
    \includegraphics[width=0.98\textwidth]{./figures/overview.png}
    \caption{(a): Contrastive learning aims at learning well-clustered representations on the latent hypersphere~\cite{wang2020understanding}. (b)-Left: Encoder trained with standard contrastive loss is inconsistent since different views of an instance are encouraged to be represented samely in the feature space, regardless of the actual difference between them. (b)-Right: For a consistent encoder, the latent view of stronger augmented data should be further away from the latent view of the raw data.}
    \label{fig:overview}
\end{figure*}

The proposed consistency loss is defined upon the optimal similarity between the latent-space representations of an  input example and a transformed view, which is a function of the strength of the composite data augmentation. However, this optimal similarity is not known \emph{a prior}. To cope with this difficulty, we take a data-driven approach to train a black-box neural network that maps from the composition vector of each data augmentation to the desired similarity using bi-level optimization. Recognizing that the similarity monotonically decreases with the strength of augmentation, we define and impose a partial monotonicity constraint to the black-box neural network such that strictly stronger composition data augmentations correspond to a strictly smaller valued similarity.

The main contributions to contrastive learning from this paper are as follows:

\begin{itemize}
  \item Systematically explore  a large set of diverse composite data augmentations for which we define a partial monotonicity concept.
  \item Propose a new notion of consistency  called  \emph{DA consistency} for data augmentation, quantifying the monotonic dependency of the latent-space similarity between an input example and its transformed view as a function of the strength of the augmentation.
  \item Introduce a new \emph{consistency loss} for training the base encoder towards satisfying DA consistency and for making use of diverse augmentations without inconsistency-induced performance loss.
  \item Propose a data-driven approach to learn the desired latent-space similarities with properly defined partial monotonicity  based on bi-level optimization.

\end{itemize}

 The proposed techniques give rise to a semi-supervised learning framework based on bi-level optimization, achieving new state-of-the-art results for image recognition.

\section{Related works}
\subsection{Contrastive Learning}
Recently, a lot of contrastive learning methods have shown impressive results in visual representation learning. SimCLR~\cite{chen2020simple} takes advantage of manually designed data augmentation to produce a pool of negative and positive data pairs in minibatches. A momentum updated encoder and a large representations memory bank were introduced to further improve the performance in MOCO~\cite{he2020momentum}. Alignment and uniformity~\cite{wang2020understanding} were introduced to describe the quality of representations. It was shown that by minimizing contrastive loss alignment and uniformity of features can be asymptotically achieved.

\subsection{Data Augmentation}
In supervised deep learning, data augmentation has already been applied to natural image datasets while training, which helps improve the generalization and robustness of learnt models. Works like AutoAugment~\cite{cubuk2019autoaugment} and RandAugment~\cite{cubuk2020randaugment} adopted combinations of various transformations to produce strong data augmentation.
In self-supervised learning~\cite{chen2020simple,tian2020makes,he2020momentum},  data augmentations are carefully designed by sampling from a highly restricted set of augmentations to maintain identity of instances and avoid noise. SwAV~\cite{caron2020unsupervised} used multiple random resized-crops to produce more views of data. CLSA~\cite{wang2021contrastive} utilized combinations of augmentations as stronger augmentation, in which the encoder is trained to match the distribution of the representations of weakly and strongly augmented instances.

\subsection{Bi-level Optimization}
In deep learning,  bi-level optimization is used to adaptively adjust the learning rate during training~\cite{baydin2017online}. DARTS~\cite{liu2018darts} proposed an approach of architecture search based on bi-level optimization. In semi-supervised and self-supervised learning, Meta Pseudo Labels~\cite{pham2021meta} adopted bi-level optimization  to simultaneously update a student model and a teacher model. JOAO~\cite{you2021graph} used bi-level optimization in adaptively adjusting data augmentation applied to graph-structured data.
\section{Preliminaries}
\label{sec:formatting}
\subsection{Contrastive Learning with Memory}

The goal of contrastive representation learning is to learn a powerful encoder parameterized by $\bm{\theta_e}$ that maps an input data $\bm{x} \in \mathbb{R}^n$ to an $\ell_2$ normalized feature vector $\bm{z}$ of dimension $m$, i.e., $f_{\bm{\theta_e}}(\cdot):\mathbb{R}^n \to \mathcal{S}^{m-1}$. The encoder is trained by minimizing a contrastive loss. Specifically, it defines two different augmented versions of the same input example  as a positive pair, which are expected to have similar representations in the feature space. Meanwhile,  the encoder shall be trained to discriminate any two instances augmented from different input examples, i.e., a negative pair, in the latent feature space. Empirically, a queue $\mathcal{Q}$ of negative embeddings is maintained for forming negative pairs with augmented data from each minibatch $\mathcal{B}$.
In this paper, we consider the wide adopted  loss that has been shown to be effective in representation learning \cite{chen2020simple,he2020momentum,tian2020contrastive}:
\begin{equation}
\begin{split}
    \mathcal{L}_{\text{contrast}}=
    \mathbb{E}_{i\in \mathcal{B}}[- \log \frac{e^{\bm{z}_i^T\bm{z}_i^+/\tau}}{e^{\bm{z}_i^T\bm{z}_i^+/\tau} + \sum_{z_j \in \mathcal{Q}} e^{\bm{z}_i^T\bm{z}_j/\tau}}]
\end{split} \label{eq:contrastive_loss}
\end{equation}
where $\bm{z}_i$ and $\bm{z}_i^+$ denote the representations of a positive pair, $\bm{z}_i$ and $\bm{z}_j$ are representations of a negative pair.  $\bm{z}^{T}\bm{z}$ denotes the cosine similarity of feature vectors, which is scaled by temperature $\tau$.

\subsection{Inconsistency in Contrastive Learning}

Data augmentation plays a pivotal role in contrastive learning and needs to be carefully designed to maintain the latent structure of the original data. Despite that a  variety of augmentations have been successfully applied to supervised settings, contrastive learning benefits from diverse augmentations to a lesser degree,  and only a subset of "weak" augmentations have been empirically shown to be effective. For instance, in the popular contrastive learning framework of \cite{chen2020simple,he2020momentum,chen2020improved,tian2020makes}, the adopted data augmentations are limited to a combination of three to five  types of basic transformations, which are chosen empirically according to the experience.   Adding to or removing from this pre-selected basic transformations can  deteriorate the performance, a phenomenon that is not well understood.  Moreover, although a few works have  attempted to introduce ``stronger"

augmentations\cite{caron2020unsupervised,wang2021contrastive,misra2020self}, there still lacks a systematic view of how augmentations with a widely varying strength affect the learning process and how to leverage a larger set of diverse augmentations to improve, as opposed to, degrade the performance .

Our key observation is that the vanilla contrastive learning framework\cite{chen2020simple,he2020momentum} does not imply \emph{data augmentation (DA) consistency}, a key property proposed by this work, which should be maintained during contrastive learning.
Without the consistency in place, the optimized encoder simply pulls positive pairs to the same position in the latent space $\mathcal{S}^{m-1}$, regardless of the actual difference between them. However, a weakly augmented view and a strongly augmented view of an input example do not necessarily share the same latent structure, and it may be desirable to encode them to different regions in the feature space. An illustrated example is shown in \cref{fig:overview} (b). In this case, the encoder that tries to pull the positive pairs together  learns an incorrect latent representation. We refer to this phenomenon as data augmentation \emph{inconsistency} in contrastive learning.

To prevent this inconsistency, the encoder should be trained to not only cluster positive pairs together, but also map them  to their respective optimal location on $\mathcal{S}^{m-1}$, based on the strength of the applied augmentations.
We thereby define this property as \emph{data augmentation consistency}, shortly \emph{DA consistency}, an additional constraint for contrastive learning. The proposed DA consistency  takes into account the critical role played by the augmentation strength in the learning process, removes potential inconsistencies across different augmentations, and facilities effective use of large numbers of diverse augmentations.

\section{Method}
\subsection{Composite Augmentations}

\begin{definition}[Composite Augmentations] Given a set of $k$ basic augmentations $\bm{T}=\left\{T_{1}, T_{2}\cdots T_{k}\right\}$, a composite augmentation with a length of $l$, namely, $A^{l}$ is defined as the composition of $l$ randomly sampled augmentations from $\bm{T}$:
\end{definition}

\begin{equation}
\begin{aligned}
    &\quad \ A^{l} = T^{1}\circ T^{2}\circ \cdots T^{l}\\
    &\text{where}\ \ T^{i}\in \bm{T},\ i={1, 2\cdots l}
\end{aligned}
\end{equation}

We denote the set that contains all composite augmentations with a length of $l$ by $\bm{A}^{l}$, and denote  the set that contains all composite augmentations with different lengths by $\bm{A}=\bm{A}^{1}\cup \bm{A}^{2}\cdots$.
Empirical studies showed that the ordering of the basic augmentations in a composite augmentation does not have a significant impact on performance. To simplify the discussion, we do not consider the effect of the ordering of the basic transformations in each composite augmentation  $A\in \bm{A}$ and hence assume that $A$ is order invariant, e.g., $T^{1}\circ T^{2}=T^{2}\circ T^{1}$. Thereby, we represent a composite augmentation with an unique composition vector.

\begin{definition}[Composition Vector] The composition vector of a composite augmentation $A^{l}$ is defined as a $k$ dimensional vector $\bm{V}(A^{l})\in {\mathbb{R}^{k}}$, where the $i_{th}$ entry $V_{i}$ is the number of times that the basic augmentation $T_{i}$ is used in $A^{l}$. Specifically, $\sum_{i=1}^{k} V_{i} = l$.
\end{definition}

Given an  encoder $f$, the location in $\mathcal{S}^{m-1}$ to which  an encoded augmented view $A(\bm{x})$ is mapped can be characterized by the cosine similarity between $f(\bm{x})$ and $f(A(\bm{x}))$, which we define as  latent deviation  below.
\begin{definition}[Latent Deviation]
Given a raw input example $\bm{x}$, an encoder $f$, and a composite augmentation $A\in \bm{A}$, the latent deviation $\Omega(\cdot;f,A)$ of $A(\bm{x})$ is defined as the cosine similarity between $\bm{x}$ and $A(\bm{x})$ in the latent space $\mathcal{S}^{m-1}$:
\begin{equation}
    \Omega(\bm{x};f,A)=f(\bm{x})^{T}\cdot f(A(\bm{x})).
\end{equation}
\end{definition}
We assume that encoder $f$ maps its input to a normalized vector in the latent space.

\subsection{Data Augmentation Consistency}

Defining the aforementioned \emph{DA consistency} in contrastive learning requires quantifying the strength of different augmentations and learning a mapping from the strength to the optimal mapped location in $\mathcal{S}^{m-1}$ with respect to a given raw input example $\bm{x}$. We specify a mapped latent space location using the corresponding latent deviation. To be DA consistent, latent deviations of different augmentations shall be monotonically decreasing with the augmentation strength.

However, a straightforward definition of augmentation strength is challenging. This is because that  two composite augmentations may consist of different basic augmentation types and it is not immediately evident how to define and compare their strength levels. To tackle this problem, we take a black-box data-driven approach to directly learn a mapping $g^{*}$ that takes a  composite augmentation as input and map it to the optimal mapped latent space location of  the augmented view of each given  raw input  $\bm{x}$, relaxing the necessity of modeling the strength of different augmentations.
Since a key objective of contrastive learning is to find an optimal encoder $f^*$ that delivers the best performance, the optimal mapped location outputted by mapping $g^{*}$ is  specified as the corresponding latent deviation of an optimal encoder $f^*$.

Under this optimal  $f^*$, the mapped latent space location $f(A(\bm{x}))$ in  $\mathcal{S}^{m-1}$ corresponds to the optimal latent deviation.
The optimal encoder $f^*$ is defined to be \emph{fully data-augmentation consistent}, or \emph{fully DA consistent} in short, and its  latent derivation for any given composite augmentation is considered optimal.  For other encoders that are not fully  DA consistent, we further define DA consistency level, as the degree at which the learned representations are consistent in terms of data augmentation.
\begin{definition}[Consistency Level]
Given an encoder $f$, a set of composite augmentation $\bm{A}$, a raw input $\bm{x}$, and a fully DA consistent encoder $f^{*}$, the DA consistency level (DACL) is define as:
\begin{equation}
    \text{DACL}(\bm{x},f;f^{*},\bm{A}) = \mathbb{E}_{A\sim \bm{A}}[\lvert \Omega(\bm{x};f,A)-\Omega(\bm{x};f^{*},A) \lvert]
\end{equation}
\end{definition}

DACL measures the level of consistency of a given encoder $f$ by comparing the latent deviation of the augmented view of $\bm{x}$ with the corresponding optimal latent deviation $\Omega(\bm{x};f^{*},A^{l})$ over all possible composite augmentations.

With DACL, an encoder that is not fully DA consistent can be trained to be more DA consistent by minimizing the \emph{consistency loss} defined below:
\begin{equation}
\begin{aligned}
    &\quad \mathcal{L}_{\text{consistent}}(\bm{\theta_{e}})\\
    &=\mathbb{E}_{\bm{x}\sim \mathcal{B}}[\text{DACL}(\bm{x}, f;f^{*}, \bm{A})]\\
    &=\mathbb{E}_{\bm{x}\sim \mathcal{B}, A\sim \bm{A}}[\lvert \Omega(\bm{x};f,A)-\Omega(\bm{x};f^{*},A) \lvert]
\end{aligned} \label{eq:consist_loss}
\end{equation}
where $\bm{\theta_{e}}$ is the parameter of the encoder $f$.

\subsection{Learning Optimal Latent Deviations with a Partially Monotonic Neural Network}\label{sec43}

Improving the consistency of an encoder by minimizing the proposed consistency loss needs to determine the optimal latent deviations $\Omega^{*} = \Omega(\bm{x};f^{*},A) $ in \cref{eq:consist_loss}.  In practice, we adopt a data-driven method to approximately learn the optimal latent deviations. We introduce a neural network $g_{\theta_d}(\cdot)$, parameterized by $\bm{\theta_d}$, to predict the optimal latent deviation for a given composite augmentation specified by its composite vector $\bm{V}$, i.e., $g_{\theta_d}(\bm{V})\cong\Omega^{*}(\bm{V})$.

Although it is possible to approximate the optimal latent deviation $\Omega^{*}$  with a blackbox neural network with no prior information, doing so may violate the important  partially monotonic constraint discussed next. As the overall strength of a composite augmentation increases,  e.g., by incorporating additional basic augmentations, the similarity between the augmented data and the raw data is expected to decrease. For instance, even though it is difficult to directly compare the impact of \emph{rotation} with \emph{saturation} on the input data, applying \emph{saturation} twice would certainly distort the data more than applying it once. Thus, we enforce partial monotonicity of the network's output w.r.t. the input vector $\bm{V}$, i.e., the output is monotonically decreasing w.r.t. every entry in the input vector. We consequently call the mapping network as a partially monotonic neural network (PMNN).

We replace $\Omega^{*}$ in the original consistency loss (\cref{eq:consist_loss}) by the value estimated by PMNN and rewrite the consistency loss as:

\begin{equation}
\begin{aligned}
    &\quad \mathcal{L}_{\text{consistent}}(\bm{\theta_{e}},\bm{\theta_d})\\
    &=\mathbb{E}_{\bm{x}\sim \mathcal{B}, A\sim \bm{A}}[\lvert \Omega(\bm{x};f,A)-g_{\theta_d}(\bm{V}(A) \lvert]
\end{aligned} \label{eq:consist_loss2}
\end{equation}

The overall loss function to optimize the encoder can be written as:

\begin{equation}
\begin{split}
    \mathcal{L}_u(\bm{\theta_e}, \bm{\theta_d})=\mathcal{L}_{\text{contrast}}(\bm{\theta_e})+\mathcal{L}_{\text{consist}}(\bm{\theta_e},\bm{\theta_d})
\end{split} \label{eq:overall_loss}
\end{equation}

PMNN aims to learn the optimal latent deviation based on the composition of augmentation which has been applied. However, to obtain the accurate optimal latent deviation is challenging since in an unsupervised learning framework, we lack the information of latent deviations' optimality in the feature space. We hence introduce a small amount of supervision to help PMNN more accurately predict the optimal latent deviation.

Note that in \cref{eq:overall_loss}, the update of the encoder's parameters $\bm{\theta_e}$ always depends on $\bm{\theta_d}$ via the optimal latent deviation prediction in consistency loss. We can then express this dependency as $\bm{\theta_e}(\bm{\theta_d})$, i.e., viewing $\bm{\theta_e}$ as a ``function" of $\bm{\theta_d}$. Therefore, we could further optimize $\bm{\theta_d}$ on labeled data:

\begin{equation}
\begin{split}
 &\quad \quad\ \min_{\bm{\theta_d}}\quad \text{CE}(\bm{x_{l}}, y;\bm{\theta_e^{*}}(\bm{\theta_d}))\\
 &\text{s.t.}\quad  \bm{\theta_e^{*}}(\bm{\theta_d})=\argmin_{\bm{\theta_e}} \mathcal{L}_u(\bm{\theta_e}, \bm{\theta_d})
\end{split} \label{eq:bilevel_optim}
\end{equation}
where $\bm{x_l}$ is labeled data, $y$ is the corresponding label, and $\text{CE}$ denotes the cross entropy criterion. Additionally, in order to update $\theta_d$ by evaluating the encoder's performance on labeled data. We adopt a linear classifier that is trained on the fly with the PMNN on the small amount of labeled data. When calculating the cross entropy loss in \cref{eq:bilevel_optim}, the classifier is connected to the stop-gradient backbone of the encoder.

Intuitively, in \cref{eq:bilevel_optim}, by updating encoder's parameter with $\mathcal{L}_u$ on unlabeled data, the encoder can approach optimal DA consistency predicted by PMNN. And by adjusting PMNN's parameter $\bm{\theta_d}$ according to the encoder's current performance on labeled data, PMNN can learn a more accurate mapping from augmentations to latent deviation.

With a learning rate $\eta_e$ and PMNN's prediction of optimal latent deviation, the update of the encoder can be rewritten as:

\begin{equation}\label{eq:thetae}
    \bm{\theta_e}^\prime=\bm{\theta_e}-\eta_e\nabla_{\bm{\theta_e}}\mathcal{L}_u(\bm{\theta_e},\bm{\theta_d})
\end{equation}
And with the current encoder, PMNN's parameters $\theta_d$'s update can be derived by applying chain rule. We rewrite its update as:

\begin{equation}\label{eq:thetad}
    \bm{\theta_d}^{\prime}=\bm{\theta_d}- \eta_d\nabla_{\bm{\theta_d}}\text{CE}(\bm{x_{l}}, y;\bm{\theta_e}(\bm{\theta_d}))
\end{equation}

The detailed derivation of the dependency of $\bm{\theta_e}$ on $\bm{\theta_d}$ and the gradient of $\bm{\theta_d}$ are included in Appendix A.

\subsection{CoCor Framework Overview}
The proposed CoCor framework consists of two parameterized neural network to be trained, i.e., the encoder and the PMNN. The encoder training is based purely on unlabeled data, while small amount of labeled data is introduced for monotonic neural network training. We utilized alternating gradient descent algorithm that alternates between \cref{eq:thetae} and \cref{eq:thetad}.

We summarized the overall algorithm flow of CoCor in \cref{alg:bilevel}.

\begin{algorithm}
\caption{Algorithm flow of CoCor}\label{alg:bilevel}

\KwInput{initial parameters $\bm{\theta_e^{(0)}}$ and $\bm{\theta_d^{(0)}}$, number of training epochs $N$.}
\For{i=1 \KwTo N}
{
1.Call \cref{eq:thetae} to update the encoder's parameter $\bm{\theta_e^{(i)}}$ on unlabeled data with monotonic neural network parameters $\bm{\theta_d^{(i-1)}}$ fixed.\\
2.Call \cref{eq:thetad} to update the monotonic neural net $\bm{\theta_d^{(i)}}$ on labeled data with the encoder parameters $\bm{\theta_e^{(i)}}$ fixed.
}
\KwOutput{Trained encoder with parameters $\bm{\theta_e^{(N)}}$}
\end{algorithm}

\section{Experiments}
We evaluate the proposed CoCor framework with different sets of experiments and benchmarks. We perform pretraining of the encoder using this proposed method and evaluate the pretrained encoder with several different downstream tasks. Its performance is compared with other unsupervised learning methods.
\subsection{Pretraining on ImageNet}
A ResNet-50~\cite{he2016deep} is adopted as the backbone of the encoder~\cite{chen2020simple,he2020momentum} for pretraining on ImageNet. We adopt a 128-dim projection head, which contains two layers with a 2048-dim hidden layer and ReLU. The projected features are $\ell_2$ before their cosine similarity are calculated. An SGD~\cite{bottou2010large} optimizer with initial $\text{lr}=0.03$, $\text{weight\ decay}=10^{-4}$, and $\text{momentum}=0.9$ is used as the optimizer. The learning rate decay is completed by a cosine scheduler. The batch size is set to 256. When running on our cluster of 8 A100 GPUs, the actual size of the batch allocated to each GPU is 32. Following the setup of MOCO~\cite{he2020momentum}, we adopt Shuffling Batch Normalization to prevent the encoder from approaching trivial solutions.

For the sake of fair comparison with SwAV~\cite{caron2020unsupervised}, we adopt multi-crop as it has been shown effective in providing the encoder with more information. In unsupervised pretraining, we apply five different crops $(224\times 224, 192\times 192, 160 \times 160, 128\times 128, 96\times 96)$.

In contrastive loss, the data augmentation we utilize follows~\cite{chen2020improved}, which includes $224\times 224$ \textit{cropping}, \textit{Gaussian Blur}, \textit{Color Jittering}, \textit{Grayscale}, and \textit{HorizontalFlip}. While in consistency loss, the composite augmentation is formed by randomly sampling augmentations from a pool of size 14, which includes \textit{Autocontrast}, \textit{brightness}, \textit{Color}, \textit{Contrast}, \textit{Equalize}, \textit{Identity}, \textit{Posterize}, \textit{Rotate}, \textit{Sharpness}, \textit{ShearX}, \textit{ShearY}, \textit{TranslateX}, \textit{TranslateY}, and \textit{Solarize}. Our preliminary experiments show that composite augmentations of any intensity can help boost the performance. However, we observed that lower intensity composite augmentation has better effects on downstream tasks. Thus, we report the results by applying composite augmentation that contains 1 sample from the augmentation pool.

\begin{table*}
\begin{center}
\caption{Transfer learning results on cross-dataset classification and object detection tasks.}\label{tab:transfer}
\small
\setlength{\tabcolsep}{7mm}{
\begin{tabular}{lccccc}
\toprule[1pt]
 & &\multicolumn{1}{c} {Classification} & \multicolumn{3}{c}{Object Detection} \\
 \cmidrule(lr){3-3}
 \cmidrule(lr){4-6}
 &  &VOC07 & VOC07+12 & \multicolumn{2}{c}{COCO} \\
Method & Epoch  & mAP & $AP_{50}$ & $AP$ & $AP_{S}$ \\
\midrule[1pt]
Supervised & - & 87.5 & 81.3 & 40.8 & 20.1 \\ \hline
MoCo~\cite{he2020momentum} & 800  & 79.8 & 81.5 & - & - \\
PIRL~\cite{misra2020self} & 800 & 81.1 & 80.7 & - & - \\
PCL v2~\cite{li2020prototypical} & 200 & 85.4 & 78.5 & - & - \\
SimCLR~\cite{chen2020simple} & 800 & 86.4 & - & - & - \\
MoCo v2~\cite{chen2020improved} & 800  & 87.1 & 82.5 & 42.0 & 20.8 \\
SWAV~\cite{caron2020unsupervised} & 800  & 88.9 & 82.6 & {\bf 42.1} & 19.7 \\
\midrule
CoCor & {\bf 200} & {\bf 93.4} & {\bf 83.1} &41.9  & {\bf 24.6}  \\
\bottomrule[1pt]
\end{tabular}}
\end{center}
\end{table*}

\subsection{Linear Evaluation on Imagenet}
We pretrained the encoder for 200 epochs. Then a fully connected classifier is trained on top of the average pooling features of the frozen pretrained backbone. The linear classifier is trained for 100 epochs, with a starting learning rate of 10.0, a cosine learning rate decay scheduler, and a batch size of 256.

\cref{tab:linear_eval} reports the linear evaluation results on ImageNet of different methods with the ResNet-50 backbone and 200 epochs pretraining. Our method outperforms all of the other listed methods. The fact that our method outperforms SwAV, which adopts exactly the same multi-crop as ours, indicates that the way by which we introduce more diverse augmentations indeed provides the encoder with additional beneficial information. These information helps the encoder to better explore the representations space and achieve better performance in linear classification.

\begin{table}
  \centering
  \begin{tabular}{lcc}
    \toprule
    Method & Batch size & Top-1 acc \\
    \midrule
    InstDisc~\cite{wu2018unsupervised} & 256 & 58.5 \\
    MoCo~\cite{he2020momentum} & 256 & 60.8 \\
    SimCLR~\cite{chen2020simple} & 256 & 61.9\\
    CPC v2~\cite{henaff2020data} & 512 & 63.8\\
    PCL v2~\cite{li2020prototypical} & 256 & 67.6\\
    MoCo v2~\cite{chen2020improved} & 256 & 67.5\\
    MoCHi~\cite{kalantidis2020hard} & 512 & 68.0\\
    PIC~\cite{cao2020parametric} & 512 & 67.6 \\
    SwAV~\cite{caron2020unsupervised} & 256 & 72.7\\
    \midrule
    CoCor & 256 & \textbf{72.8}\\
    \bottomrule
  \end{tabular}
  \caption{Linear evaluation results of different methods with ResNet-50 backbone over 200 epochs pretraining.}
  \label{tab:linear_eval}
\end{table}

\subsection{Transferring Representations}
An essential characteristic of general and robust representations is their transferability. We hence evaluate the backbone pretrained on ImageNet by transferring to some other downstream tasks.

Firstly, we conduct the transfer learning experiments on VOC07~\cite{everingham2009pascal}. We freeze the pretrained ResNet-50 backbone, and fine-tune a linear classifier connected to it on VOC07trainval. The testing is completed on VOC07test.

Secondly, we fine-tune the pretrained backbone for object detection. We adopt detectron2, which has been utilized by other works\cite{he2020momentum, chen2020simple, hu2021adco}, for evaluating on VOC and COCO~\cite{lin2014microsoft} datasets. For object detection in VOC, we fine-tune the network on VOC07+12 trainval set and test it on VOC07 test set. For COCO, we fine-tune the network on train2017 set (118k images) and test it on val2017 datatset.

The results of transfer learning is reported in \cref{tab:transfer}. The performance of our method with 200 epochs ImageNet pretraining is compared with other methods with 800 epochs pretraining. CoCor achieves outstanding results compared with other models' state-of-the-art performance. CoCor's good performance in transfer learning implies that this proposed method learns more robust and general representations than other listed methods. With only 200 epochs pretraining, CoCor outperforms other SOTAs with 800 epochs a lot on VOC in classification and object detection. In COCO, the CoCor achieved a result of 24.6\% of $APs$, which stands for the accuracy of small object detection. The performance on the well known difficult challenge in COCO has been notably improved by 3.8\% from MoCo v2's performance of 20.8\%. The diversity in views produced by various augmentations provides the encoder with the power of detecting and distinguishing these challenging small objects.

\subsection{Benefits of the Monotonic Neural Network}
We further conduct ablation study on CIFAR-10 to show the benefits of introducing the monotonic neural network. We experiment with single-crop situation, and pick for consistency loss the composite augmentation which contains two transformations from the pool.

Firstly, we select the optimal latent deviation $\Omega^{*}$ in consistency loss by repeatedly running experiments minimizing the overall loss function in \cref{eq:overall_loss} with different value of $\Omega$. We then fix the pre-tuned optimal $\Omega^{*}$ in further training, which means the variance caused by the transformations' type difference is neglected. Only the number of transformation sampled to form composite augmentation is considered to be able to influence the value $\Omega$. The optimal result by fixing the constant $\Omega^{*}$ is reported in \cref{tab:ablation_monot_nn}. Then we introduce the monotonic neural network to adaptively learn $\Omega$ w.r.t. the composition of the composite augmentation with a length of 2 and trained with the bi-level optimization framework in \cref{alg:bilevel}.

\begin{table}
  \centering
  \begin{tabular}{ccc}
    \toprule
     & without PMNN & with PMNN \\
    \midrule
    Top1 acc & 50.20 & 52.04\\
    \bottomrule
  \end{tabular}
  \caption{Ablation study of the partially monotonic neural network on CIFAR-10 dataset.}
  \label{tab:ablation_monot_nn}
\end{table}

The encoders with same backbone structure, learning rate, optimizer, batch size, and learning rate scheduler is trained with/without the monotonic neural network. They are pretrained for 20 epochs on Cifar-10. Linear classifiers connected to the frozen pretrained backbone are then trained for 20 epochs to evaluate their performance. \cref{tab:ablation_monot_nn} reports the results of training with/without the monotonic neural network. The encoder trained with the monotonic neural network outperforms the one without it. With this neural network, the variance of transformations are taken into consideration. Meanwhile, the small amount of labeled data for updating the neural network can deliver information of label to the encoder under the bi-level optimization update.

{\small
\bibliographystyle{ieee_fullname}

\section{Introduction}
\label{sec:intro}
Data augmentation is widely used in supervised learning for image recognition, achieving excellent results on popular datasets like MNIST, CIFAR-10/100, and ImageNet \cite{ciregan2012multi,sato2015apac,simard2003best,wan2013regularization,krizhevsky2017imagenet}.
AutoAugment~\cite{cubuk2019autoaugment} and RandAugment~\cite{cubuk2020randaugment} use strong augmentations composed of a combination of different transformations. It has been shown that these manually designed data augmentations can effectively increase the amount of training data and enhance the model's ability to capture the invariance of the training data, thereby boosting supervised training performance.

Data augmentation (DA) is also a key component in recent semi-supervised learning techniques based on contrastive learning \cite{chen2020simple}.   A base encoder that learns good latent representations of the input data is trained with a contrastive loss. The contrastive loss is reflective of the following:  in the latent space a pair of two views of a given data example transformed by two different data augmentation functions are correlated (similar) while transformed views of different input examples are dissimilar. The quality of the encoder trained using unlabeled data as such is essential to the overall contrastive learning performance, and is contingent upon the employed data augmentations.

 ~\cite{wang2020understanding} interprets the minimization of the contrastive loss  is to achieve alignment and uniformity of the learned representations. To learn powerful and transferable representations, many studies have aimed at improving contrastive learning by choosing proper data augmentations  \cite{chen2020simple,tian2020makes,he2020momentum} have shown that combining different augmentations with a proper magnitude may improve performance. In these works, data augmentations are restricted to be a random composition of four or five types of augmentations with a very limited range of magnitude. To further improve contrastive learning in downstream tasks, PIRL~\cite{misra2020self} adopts two additional data augmentations, SwAV~\cite{caron2020unsupervised} introduced multiple random resized-crops to provide the encoder with more views of data and CLSA~\cite{wang2021contrastive} adopts combinations of stronger augmentations.

While exploiting a diverse set of data augmentations is appealing, thus far, there does not exist a systematic approach  for incorporating a large number of augmentations for contrastive learning. Importantly, the intricacies of data augmentation and representation learning may result in performance degradation if data augmentation functions are not properly chosen and used. To address the above challenges, this paper proposes \textbf{Co}ntrastive Learning with \textbf{Co}nsistent \textbf{R}epresentations (CoCor). First, we define a set of  \textit{composite augmentations}, denoted by $\bm{A}$ and composed by combining multiple basic augmentations.  We represent each composite augmentation   using a \emph{composition vector} that encodes the types and numbers of basic augmentations used and their augmentation magnitude.
A large number of diverse augmentations can be drawn from this composite set  $\bm{A}$ with widely varying transformation types and overall strength.

Critically, we propose a notion of \emph{consistency} for data augmentations, called \emph{DA consistency}, which imposes a desiring property on the latent representations of views produced by different composite data augmentations. For a given input example, a stronger data augmentation shall push its latent-space view further away from the latent representation of the original example. That is, the latent-space similarity between the transformed view and the original input decreases with the strength of the data augmentation, a very desiring characteristic, as illustrated in \cref{fig:overview}(b). In other words, a base encoder that does not satisfy the above property is deemed \emph{inconsistent} from a data augmentation point of view.
We consequently propose to add a new \textit{consistency loss} to impose  \emph{DA consistency} on  the base encoder. Different from the contrastive loss that is agnostic to the type and overall strength of employed data augmentations ~\cite{chen2020simple,tian2020makes,he2020momentum}, our \emph{consistency loss} places additional finer-grained constraints on the encoder learning such that the resulted latent-space views are consistent with the strength of data augmentations applied.
Enforcing \emph{DA consistency} is key to circumventing inconsistencies among diverse data augmentations, which  may result in performance degradation.

\begin{figure*}
    \centering
    \includegraphics[width=0.98\textwidth]{./figures/overview.png}
    \caption{(a): Contrastive learning aims at learning well-clustered representations on the latent hypersphere~\cite{wang2020understanding}. (b)-Left: Encoder trained with standard contrastive loss is inconsistent since different views of an instance are encouraged to be represented samely in the feature space, regardless of the actual difference between them. (b)-Right: For a consistent encoder, the latent view of stronger augmented data should be further away from the latent view of the raw data.}
    \label{fig:overview}
\end{figure*}

The proposed consistency loss is defined upon the optimal similarity between the latent-space representations of an  input example and a transformed view, which is a function of the strength of the composite data augmentation. However, this optimal similarity is not known \emph{a prior}. To cope with this difficulty, we take a data-driven approach to train a black-box neural network that maps from the composition vector of each data augmentation to the desired similarity using bi-level optimization. Recognizing that the similarity monotonically decreases with the strength of augmentation, we define and impose a partial monotonicity constraint to the black-box neural network such that strictly stronger composition data augmentations correspond to a strictly smaller valued similarity.

The main contributions to contrastive learning from this paper are as follows:

\begin{itemize}
  \item Systematically explore  a large set of diverse composite data augmentations for which we define a partial monotonicity concept.
  \item Propose a new notion of consistency  called  \emph{DA consistency} for data augmentation, quantifying the monotonic dependency of the latent-space similarity between an input example and its transformed view as a function of the strength of the augmentation.
  \item Introduce a new \emph{consistency loss} for training the base encoder towards satisfying DA consistency and for making use of diverse augmentations without inconsistency-induced performance loss.
  \item Propose a data-driven approach to learn the desired latent-space similarities with properly defined partial monotonicity  based on bi-level optimization.

\end{itemize}

 The proposed techniques give rise to a semi-supervised learning framework based on bi-level optimization, achieving new state-of-the-art results for image recognition.

\section{Related works}
\subsection{Contrastive Learning}
Recently, a lot of contrastive learning methods have shown impressive results in visual representation learning. SimCLR~\cite{chen2020simple} takes advantage of manually designed data augmentation to produce a pool of negative and positive data pairs in minibatches. A momentum updated encoder and a large representations memory bank were introduced to further improve the performance in MOCO~\cite{he2020momentum}. Alignment and uniformity~\cite{wang2020understanding} were introduced to describe the quality of representations. It was shown that by minimizing contrastive loss alignment and uniformity of features can be asymptotically achieved.

\subsection{Data Augmentation}
In supervised deep learning, data augmentation has already been applied to natural image datasets while training, which helps improve the generalization and robustness of learnt models. Works like AutoAugment~\cite{cubuk2019autoaugment} and RandAugment~\cite{cubuk2020randaugment} adopted combinations of various transformations to produce strong data augmentation.
In self-supervised learning~\cite{chen2020simple,tian2020makes,he2020momentum},  data augmentations are carefully designed by sampling from a highly restricted set of augmentations to maintain identity of instances and avoid noise. SwAV~\cite{caron2020unsupervised} used multiple random resized-crops to produce more views of data. CLSA~\cite{wang2021contrastive} utilized combinations of augmentations as stronger augmentation, in which the encoder is trained to match the distribution of the representations of weakly and strongly augmented instances.

\subsection{Bi-level Optimization}
In deep learning,  bi-level optimization is used to adaptively adjust the learning rate during training~\cite{baydin2017online}. DARTS~\cite{liu2018darts} proposed an approach of architecture search based on bi-level optimization. In semi-supervised and self-supervised learning, Meta Pseudo Labels~\cite{pham2021meta} adopted bi-level optimization  to simultaneously update a student model and a teacher model. JOAO~\cite{you2021graph} used bi-level optimization in adaptively adjusting data augmentation applied to graph-structured data.
\section{Preliminaries}
\label{sec:formatting}
\subsection{Contrastive Learning with Memory}

The goal of contrastive representation learning is to learn a powerful encoder parameterized by $\bm{\theta_e}$ that maps an input data $\bm{x} \in \mathbb{R}^n$ to an $\ell_2$ normalized feature vector $\bm{z}$ of dimension $m$, i.e., $f_{\bm{\theta_e}}(\cdot):\mathbb{R}^n \to \mathcal{S}^{m-1}$. The encoder is trained by minimizing a contrastive loss. Specifically, it defines two different augmented versions of the same input example  as a positive pair, which are expected to have similar representations in the feature space. Meanwhile,  the encoder shall be trained to discriminate any two instances augmented from different input examples, i.e., a negative pair, in the latent feature space. Empirically, a queue $\mathcal{Q}$ of negative embeddings is maintained for forming negative pairs with augmented data from each minibatch $\mathcal{B}$.
In this paper, we consider the wide adopted  loss that has been shown to be effective in representation learning \cite{chen2020simple,he2020momentum,tian2020contrastive}:
\begin{equation}
\begin{split}
    \mathcal{L}_{\text{contrast}}=
    \mathbb{E}_{i\in \mathcal{B}}[- \log \frac{e^{\bm{z}_i^T\bm{z}_i^+/\tau}}{e^{\bm{z}_i^T\bm{z}_i^+/\tau} + \sum_{z_j \in \mathcal{Q}} e^{\bm{z}_i^T\bm{z}_j/\tau}}]
\end{split} \label{eq:contrastive_loss}
\end{equation}
where $\bm{z}_i$ and $\bm{z}_i^+$ denote the representations of a positive pair, $\bm{z}_i$ and $\bm{z}_j$ are representations of a negative pair.  $\bm{z}^{T}\bm{z}$ denotes the cosine similarity of feature vectors, which is scaled by temperature $\tau$.

\subsection{Inconsistency in Contrastive Learning}

Data augmentation plays a pivotal role in contrastive learning and needs to be carefully designed to maintain the latent structure of the original data. Despite that a  variety of augmentations have been successfully applied to supervised settings, contrastive learning benefits from diverse augmentations to a lesser degree,  and only a subset of "weak" augmentations have been empirically shown to be effective. For instance, in the popular contrastive learning framework of \cite{chen2020simple,he2020momentum,chen2020improved,tian2020makes}, the adopted data augmentations are limited to a combination of three to five  types of basic transformations, which are chosen empirically according to the experience.   Adding to or removing from this pre-selected basic transformations can  deteriorate the performance, a phenomenon that is not well understood.  Moreover, although a few works have  attempted to introduce ``stronger"

augmentations\cite{caron2020unsupervised,wang2021contrastive,misra2020self}, there still lacks a systematic view of how augmentations with a widely varying strength affect the learning process and how to leverage a larger set of diverse augmentations to improve, as opposed to, degrade the performance .

Our key observation is that the vanilla contrastive learning framework\cite{chen2020simple,he2020momentum} does not imply \emph{data augmentation (DA) consistency}, a key property proposed by this work, which should be maintained during contrastive learning.
Without the consistency in place, the optimized encoder simply pulls positive pairs to the same position in the latent space $\mathcal{S}^{m-1}$, regardless of the actual difference between them. However, a weakly augmented view and a strongly augmented view of an input example do not necessarily share the same latent structure, and it may be desirable to encode them to different regions in the feature space. An illustrated example is shown in \cref{fig:overview} (b). In this case, the encoder that tries to pull the positive pairs together  learns an incorrect latent representation. We refer to this phenomenon as data augmentation \emph{inconsistency} in contrastive learning.

To prevent this inconsistency, the encoder should be trained to not only cluster positive pairs together, but also map them  to their respective optimal location on $\mathcal{S}^{m-1}$, based on the strength of the applied augmentations.
We thereby define this property as \emph{data augmentation consistency}, shortly \emph{DA consistency}, an additional constraint for contrastive learning. The proposed DA consistency  takes into account the critical role played by the augmentation strength in the learning process, removes potential inconsistencies across different augmentations, and facilities effective use of large numbers of diverse augmentations.

\section{Method}
\subsection{Composite Augmentations}

\begin{definition}[Composite Augmentations] Given a set of $k$ basic augmentations $\bm{T}=\left\{T_{1}, T_{2}\cdots T_{k}\right\}$, a composite augmentation with a length of $l$, namely, $A^{l}$ is defined as the composition of $l$ randomly sampled augmentations from $\bm{T}$:
\end{definition}

\begin{equation}
\begin{aligned}
    &\quad \ A^{l} = T^{1}\circ T^{2}\circ \cdots T^{l}\\
    &\text{where}\ \ T^{i}\in \bm{T},\ i={1, 2\cdots l}
\end{aligned}
\end{equation}

We denote the set that contains all composite augmentations with a length of $l$ by $\bm{A}^{l}$, and denote  the set that contains all composite augmentations with different lengths by $\bm{A}=\bm{A}^{1}\cup \bm{A}^{2}\cdots$.
Empirical studies showed that the ordering of the basic augmentations in a composite augmentation does not have a significant impact on performance. To simplify the discussion, we do not consider the effect of the ordering of the basic transformations in each composite augmentation  $A\in \bm{A}$ and hence assume that $A$ is order invariant, e.g., $T^{1}\circ T^{2}=T^{2}\circ T^{1}$. Thereby, we represent a composite augmentation with an unique composition vector.

\begin{definition}[Composition Vector] The composition vector of a composite augmentation $A^{l}$ is defined as a $k$ dimensional vector $\bm{V}(A^{l})\in {\mathbb{R}^{k}}$, where the $i_{th}$ entry $V_{i}$ is the number of times that the basic augmentation $T_{i}$ is used in $A^{l}$. Specifically, $\sum_{i=1}^{k} V_{i} = l$.
\end{definition}

Given an  encoder $f$, the location in $\mathcal{S}^{m-1}$ to which  an encoded augmented view $A(\bm{x})$ is mapped can be characterized by the cosine similarity between $f(\bm{x})$ and $f(A(\bm{x}))$, which we define as  latent deviation  below.
\begin{definition}[Latent Deviation]
Given a raw input example $\bm{x}$, an encoder $f$, and a composite augmentation $A\in \bm{A}$, the latent deviation $\Omega(\cdot;f,A)$ of $A(\bm{x})$ is defined as the cosine similarity between $\bm{x}$ and $A(\bm{x})$ in the latent space $\mathcal{S}^{m-1}$:
\begin{equation}
    \Omega(\bm{x};f,A)=f(\bm{x})^{T}\cdot f(A(\bm{x})).
\end{equation}
\end{definition}
We assume that encoder $f$ maps its input to a normalized vector in the latent space.

\subsection{Data Augmentation Consistency}

Defining the aforementioned \emph{DA consistency} in contrastive learning requires quantifying the strength of different augmentations and learning a mapping from the strength to the optimal mapped location in $\mathcal{S}^{m-1}$ with respect to a given raw input example $\bm{x}$. We specify a mapped latent space location using the corresponding latent deviation. To be DA consistent, latent deviations of different augmentations shall be monotonically decreasing with the augmentation strength.

However, a straightforward definition of augmentation strength is challenging. This is because that  two composite augmentations may consist of different basic augmentation types and it is not immediately evident how to define and compare their strength levels. To tackle this problem, we take a black-box data-driven approach to directly learn a mapping $g^{*}$ that takes a  composite augmentation as input and map it to the optimal mapped latent space location of  the augmented view of each given  raw input  $\bm{x}$, relaxing the necessity of modeling the strength of different augmentations.
Since a key objective of contrastive learning is to find an optimal encoder $f^*$ that delivers the best performance, the optimal mapped location outputted by mapping $g^{*}$ is  specified as the corresponding latent deviation of an optimal encoder $f^*$.

Under this optimal  $f^*$, the mapped latent space location $f(A(\bm{x}))$ in  $\mathcal{S}^{m-1}$ corresponds to the optimal latent deviation.
The optimal encoder $f^*$ is defined to be \emph{fully data-augmentation consistent}, or \emph{fully DA consistent} in short, and its  latent derivation for any given composite augmentation is considered optimal.  For other encoders that are not fully  DA consistent, we further define DA consistency level, as the degree at which the learned representations are consistent in terms of data augmentation.
\begin{definition}[Consistency Level]
Given an encoder $f$, a set of composite augmentation $\bm{A}$, a raw input $\bm{x}$, and a fully DA consistent encoder $f^{*}$, the DA consistency level (DACL) is define as:
\begin{equation}
    \text{DACL}(\bm{x},f;f^{*},\bm{A}) = \mathbb{E}_{A\sim \bm{A}}[\lvert \Omega(\bm{x};f,A)-\Omega(\bm{x};f^{*},A) \lvert]
\end{equation}
\end{definition}

DACL measures the level of consistency of a given encoder $f$ by comparing the latent deviation of the augmented view of $\bm{x}$ with the corresponding optimal latent deviation $\Omega(\bm{x};f^{*},A^{l})$ over all possible composite augmentations.

With DACL, an encoder that is not fully DA consistent can be trained to be more DA consistent by minimizing the \emph{consistency loss} defined below:
\begin{equation}
\begin{aligned}
    &\quad \mathcal{L}_{\text{consistent}}(\bm{\theta_{e}})\\
    &=\mathbb{E}_{\bm{x}\sim \mathcal{B}}[\text{DACL}(\bm{x}, f;f^{*}, \bm{A})]\\
    &=\mathbb{E}_{\bm{x}\sim \mathcal{B}, A\sim \bm{A}}[\lvert \Omega(\bm{x};f,A)-\Omega(\bm{x};f^{*},A) \lvert]
\end{aligned} \label{eq:consist_loss}
\end{equation}
where $\bm{\theta_{e}}$ is the parameter of the encoder $f$.

\subsection{Learning Optimal Latent Deviations with a Partially Monotonic Neural Network}\label{sec43}

Improving the consistency of an encoder by minimizing the proposed consistency loss needs to determine the optimal latent deviations $\Omega^{*} = \Omega(\bm{x};f^{*},A) $ in \cref{eq:consist_loss}.  In practice, we adopt a data-driven method to approximately learn the optimal latent deviations. We introduce a neural network $g_{\theta_d}(\cdot)$, parameterized by $\bm{\theta_d}$, to predict the optimal latent deviation for a given composite augmentation specified by its composite vector $\bm{V}$, i.e., $g_{\theta_d}(\bm{V})\cong\Omega^{*}(\bm{V})$.

Although it is possible to approximate the optimal latent deviation $\Omega^{*}$  with a blackbox neural network with no prior information, doing so may violate the important  partially monotonic constraint discussed next. As the overall strength of a composite augmentation increases,  e.g., by incorporating additional basic augmentations, the similarity between the augmented data and the raw data is expected to decrease. For instance, even though it is difficult to directly compare the impact of \emph{rotation} with \emph{saturation} on the input data, applying \emph{saturation} twice would certainly distort the data more than applying it once. Thus, we enforce partial monotonicity of the network's output w.r.t. the input vector $\bm{V}$, i.e., the output is monotonically decreasing w.r.t. every entry in the input vector. We consequently call the mapping network as a partially monotonic neural network (PMNN).

We replace $\Omega^{*}$ in the original consistency loss (\cref{eq:consist_loss}) by the value estimated by PMNN and rewrite the consistency loss as:

\begin{equation}
\begin{aligned}
    &\quad \mathcal{L}_{\text{consistent}}(\bm{\theta_{e}},\bm{\theta_d})\\
    &=\mathbb{E}_{\bm{x}\sim \mathcal{B}, A\sim \bm{A}}[\lvert \Omega(\bm{x};f,A)-g_{\theta_d}(\bm{V}(A) \lvert]
\end{aligned} \label{eq:consist_loss2}
\end{equation}

The overall loss function to optimize the encoder can be written as:

\begin{equation}
\begin{split}
    \mathcal{L}_u(\bm{\theta_e}, \bm{\theta_d})=\mathcal{L}_{\text{contrast}}(\bm{\theta_e})+\mathcal{L}_{\text{consist}}(\bm{\theta_e},\bm{\theta_d})
\end{split} \label{eq:overall_loss}
\end{equation}

PMNN aims to learn the optimal latent deviation based on the composition of augmentation which has been applied. However, to obtain the accurate optimal latent deviation is challenging since in an unsupervised learning framework, we lack the information of latent deviations' optimality in the feature space. We hence introduce a small amount of supervision to help PMNN more accurately predict the optimal latent deviation.

Note that in \cref{eq:overall_loss}, the update of the encoder's parameters $\bm{\theta_e}$ always depends on $\bm{\theta_d}$ via the optimal latent deviation prediction in consistency loss. We can then express this dependency as $\bm{\theta_e}(\bm{\theta_d})$, i.e., viewing $\bm{\theta_e}$ as a ``function" of $\bm{\theta_d}$. Therefore, we could further optimize $\bm{\theta_d}$ on labeled data:

\begin{equation}
\begin{split}
 &\quad \quad\ \min_{\bm{\theta_d}}\quad \text{CE}(\bm{x_{l}}, y;\bm{\theta_e^{*}}(\bm{\theta_d}))\\
 &\text{s.t.}\quad  \bm{\theta_e^{*}}(\bm{\theta_d})=\argmin_{\bm{\theta_e}} \mathcal{L}_u(\bm{\theta_e}, \bm{\theta_d})
\end{split} \label{eq:bilevel_optim}
\end{equation}
where $\bm{x_l}$ is labeled data, $y$ is the corresponding label, and $\text{CE}$ denotes the cross entropy criterion. Additionally, in order to update $\theta_d$ by evaluating the encoder's performance on labeled data. We adopt a linear classifier that is trained on the fly with the PMNN on the small amount of labeled data. When calculating the cross entropy loss in \cref{eq:bilevel_optim}, the classifier is connected to the stop-gradient backbone of the encoder.

Intuitively, in \cref{eq:bilevel_optim}, by updating encoder's parameter with $\mathcal{L}_u$ on unlabeled data, the encoder can approach optimal DA consistency predicted by PMNN. And by adjusting PMNN's parameter $\bm{\theta_d}$ according to the encoder's current performance on labeled data, PMNN can learn a more accurate mapping from augmentations to latent deviation.

With a learning rate $\eta_e$ and PMNN's prediction of optimal latent deviation, the update of the encoder can be rewritten as:

\begin{equation}\label{eq:thetae}
    \bm{\theta_e}^\prime=\bm{\theta_e}-\eta_e\nabla_{\bm{\theta_e}}\mathcal{L}_u(\bm{\theta_e},\bm{\theta_d})
\end{equation}
And with the current encoder, PMNN's parameters $\theta_d$'s update can be derived by applying chain rule. We rewrite its update as:

\begin{equation}\label{eq:thetad}
    \bm{\theta_d}^{\prime}=\bm{\theta_d}- \eta_d\nabla_{\bm{\theta_d}}\text{CE}(\bm{x_{l}}, y;\bm{\theta_e}(\bm{\theta_d}))
\end{equation}

The detailed derivation of the dependency of $\bm{\theta_e}$ on $\bm{\theta_d}$ and the gradient of $\bm{\theta_d}$ are included in Appendix A.

\subsection{CoCor Framework Overview}
The proposed CoCor framework consists of two parameterized neural network to be trained, i.e., the encoder and the PMNN. The encoder training is based purely on unlabeled data, while small amount of labeled data is introduced for monotonic neural network training. We utilized alternating gradient descent algorithm that alternates between \cref{eq:thetae} and \cref{eq:thetad}.

We summarized the overall algorithm flow of CoCor in \cref{alg:bilevel}.

\begin{algorithm}
\caption{Algorithm flow of CoCor}\label{alg:bilevel}

\KwInput{initial parameters $\bm{\theta_e^{(0)}}$ and $\bm{\theta_d^{(0)}}$, number of training epochs $N$.}
\For{i=1 \KwTo N}
{
1.Call \cref{eq:thetae} to update the encoder's parameter $\bm{\theta_e^{(i)}}$ on unlabeled data with monotonic neural network parameters $\bm{\theta_d^{(i-1)}}$ fixed.\\
2.Call \cref{eq:thetad} to update the monotonic neural net $\bm{\theta_d^{(i)}}$ on labeled data with the encoder parameters $\bm{\theta_e^{(i)}}$ fixed.
}
\KwOutput{Trained encoder with parameters $\bm{\theta_e^{(N)}}$}
\end{algorithm}

\section{Experiments}
We evaluate the proposed CoCor framework with different sets of experiments and benchmarks. We perform pretraining of the encoder using this proposed method and evaluate the pretrained encoder with several different downstream tasks. Its performance is compared with other unsupervised learning methods.
\subsection{Pretraining on ImageNet}
A ResNet-50~\cite{he2016deep} is adopted as the backbone of the encoder~\cite{chen2020simple,he2020momentum} for pretraining on ImageNet. We adopt a 128-dim projection head, which contains two layers with a 2048-dim hidden layer and ReLU. The projected features are $\ell_2$ before their cosine similarity are calculated. An SGD~\cite{bottou2010large} optimizer with initial $\text{lr}=0.03$, $\text{weight\ decay}=10^{-4}$, and $\text{momentum}=0.9$ is used as the optimizer. The learning rate decay is completed by a cosine scheduler. The batch size is set to 256. When running on our cluster of 8 A100 GPUs, the actual size of the batch allocated to each GPU is 32. Following the setup of MOCO~\cite{he2020momentum}, we adopt Shuffling Batch Normalization to prevent the encoder from approaching trivial solutions.

For the sake of fair comparison with SwAV~\cite{caron2020unsupervised}, we adopt multi-crop as it has been shown effective in providing the encoder with more information. In unsupervised pretraining, we apply five different crops $(224\times 224, 192\times 192, 160 \times 160, 128\times 128, 96\times 96)$.

In contrastive loss, the data augmentation we utilize follows~\cite{chen2020improved}, which includes $224\times 224$ \textit{cropping}, \textit{Gaussian Blur}, \textit{Color Jittering}, \textit{Grayscale}, and \textit{HorizontalFlip}. While in consistency loss, the composite augmentation is formed by randomly sampling augmentations from a pool of size 14, which includes \textit{Autocontrast}, \textit{brightness}, \textit{Color}, \textit{Contrast}, \textit{Equalize}, \textit{Identity}, \textit{Posterize}, \textit{Rotate}, \textit{Sharpness}, \textit{ShearX}, \textit{ShearY}, \textit{TranslateX}, \textit{TranslateY}, and \textit{Solarize}. Our preliminary experiments show that composite augmentations of any intensity can help boost the performance. However, we observed that lower intensity composite augmentation has better effects on downstream tasks. Thus, we report the results by applying composite augmentation that contains 1 sample from the augmentation pool.

\begin{table*}
\begin{center}
\caption{Transfer learning results on cross-dataset classification and object detection tasks.}\label{tab:transfer}
\small
\setlength{\tabcolsep}{7mm}{
\begin{tabular}{lccccc}
\toprule[1pt]
 & &\multicolumn{1}{c} {Classification} & \multicolumn{3}{c}{Object Detection} \\
 \cmidrule(lr){3-3}
 \cmidrule(lr){4-6}
 &  &VOC07 & VOC07+12 & \multicolumn{2}{c}{COCO} \\
Method & Epoch  & mAP & $AP_{50}$ & $AP$ & $AP_{S}$ \\
\midrule[1pt]
Supervised & - & 87.5 & 81.3 & 40.8 & 20.1 \\ \hline
MoCo~\cite{he2020momentum} & 800  & 79.8 & 81.5 & - & - \\
PIRL~\cite{misra2020self} & 800 & 81.1 & 80.7 & - & - \\
PCL v2~\cite{li2020prototypical} & 200 & 85.4 & 78.5 & - & - \\
SimCLR~\cite{chen2020simple} & 800 & 86.4 & - & - & - \\
MoCo v2~\cite{chen2020improved} & 800  & 87.1 & 82.5 & 42.0 & 20.8 \\
SWAV~\cite{caron2020unsupervised} & 800  & 88.9 & 82.6 & {\bf 42.1} & 19.7 \\
\midrule
CoCor & {\bf 200} & {\bf 93.4} & {\bf 83.1} &41.9  & {\bf 24.6}  \\
\bottomrule[1pt]
\end{tabular}}
\end{center}
\end{table*}

\subsection{Linear Evaluation on Imagenet}
We pretrained the encoder for 200 epochs. Then a fully connected classifier is trained on top of the average pooling features of the frozen pretrained backbone. The linear classifier is trained for 100 epochs, with a starting learning rate of 10.0, a cosine learning rate decay scheduler, and a batch size of 256.

\cref{tab:linear_eval} reports the linear evaluation results on ImageNet of different methods with the ResNet-50 backbone and 200 epochs pretraining. Our method outperforms all of the other listed methods. The fact that our method outperforms SwAV, which adopts exactly the same multi-crop as ours, indicates that the way by which we introduce more diverse augmentations indeed provides the encoder with additional beneficial information. These information helps the encoder to better explore the representations space and achieve better performance in linear classification.

\begin{table}
  \centering
  \begin{tabular}{lcc}
    \toprule
    Method & Batch size & Top-1 acc \\
    \midrule
    InstDisc~\cite{wu2018unsupervised} & 256 & 58.5 \\
    MoCo~\cite{he2020momentum} & 256 & 60.8 \\
    SimCLR~\cite{chen2020simple} & 256 & 61.9\\
    CPC v2~\cite{henaff2020data} & 512 & 63.8\\
    PCL v2~\cite{li2020prototypical} & 256 & 67.6\\
    MoCo v2~\cite{chen2020improved} & 256 & 67.5\\
    MoCHi~\cite{kalantidis2020hard} & 512 & 68.0\\
    PIC~\cite{cao2020parametric} & 512 & 67.6 \\
    SwAV~\cite{caron2020unsupervised} & 256 & 72.7\\
    \midrule
    CoCor & 256 & \textbf{72.8}\\
    \bottomrule
  \end{tabular}
  \caption{Linear evaluation results of different methods with ResNet-50 backbone over 200 epochs pretraining.}
  \label{tab:linear_eval}
\end{table}

\subsection{Transferring Representations}
An essential characteristic of general and robust representations is their transferability. We hence evaluate the backbone pretrained on ImageNet by transferring to some other downstream tasks.

Firstly, we conduct the transfer learning experiments on VOC07~\cite{everingham2009pascal}. We freeze the pretrained ResNet-50 backbone, and fine-tune a linear classifier connected to it on VOC07trainval. The testing is completed on VOC07test.

Secondly, we fine-tune the pretrained backbone for object detection. We adopt detectron2, which has been utilized by other works\cite{he2020momentum, chen2020simple, hu2021adco}, for evaluating on VOC and COCO~\cite{lin2014microsoft} datasets. For object detection in VOC, we fine-tune the network on VOC07+12 trainval set and test it on VOC07 test set. For COCO, we fine-tune the network on train2017 set (118k images) and test it on val2017 datatset.

The results of transfer learning is reported in \cref{tab:transfer}. The performance of our method with 200 epochs ImageNet pretraining is compared with other methods with 800 epochs pretraining. CoCor achieves outstanding results compared with other models' state-of-the-art performance. CoCor's good performance in transfer learning implies that this proposed method learns more robust and general representations than other listed methods. With only 200 epochs pretraining, CoCor outperforms other SOTAs with 800 epochs a lot on VOC in classification and object detection. In COCO, the CoCor achieved a result of 24.6\% of $APs$, which stands for the accuracy of small object detection. The performance on the well known difficult challenge in COCO has been notably improved by 3.8\% from MoCo v2's performance of 20.8\%. The diversity in views produced by various augmentations provides the encoder with the power of detecting and distinguishing these challenging small objects.

\subsection{Benefits of the Monotonic Neural Network}
We further conduct ablation study on CIFAR-10 to show the benefits of introducing the monotonic neural network. We experiment with single-crop situation, and pick for consistency loss the composite augmentation which contains two transformations from the pool.

Firstly, we select the optimal latent deviation $\Omega^{*}$ in consistency loss by repeatedly running experiments minimizing the overall loss function in \cref{eq:overall_loss} with different value of $\Omega$. We then fix the pre-tuned optimal $\Omega^{*}$ in further training, which means the variance caused by the transformations' type difference is neglected. Only the number of transformation sampled to form composite augmentation is considered to be able to influence the value $\Omega$. The optimal result by fixing the constant $\Omega^{*}$ is reported in \cref{tab:ablation_monot_nn}. Then we introduce the monotonic neural network to adaptively learn $\Omega$ w.r.t. the composition of the composite augmentation with a length of 2 and trained with the bi-level optimization framework in \cref{alg:bilevel}.

\begin{table}
  \centering
  \begin{tabular}{ccc}
    \toprule
     & without PMNN & with PMNN \\
    \midrule
    Top1 acc & 50.20 & 52.04\\
    \bottomrule
  \end{tabular}
  \caption{Ablation study of the partially monotonic neural network on CIFAR-10 dataset.}
  \label{tab:ablation_monot_nn}
\end{table}

The encoders with same backbone structure, learning rate, optimizer, batch size, and learning rate scheduler is trained with/without the monotonic neural network. They are pretrained for 20 epochs on Cifar-10. Linear classifiers connected to the frozen pretrained backbone are then trained for 20 epochs to evaluate their performance. \cref{tab:ablation_monot_nn} reports the results of training with/without the monotonic neural network. The encoder trained with the monotonic neural network outperforms the one without it. With this neural network, the variance of transformations are taken into consideration. Meanwhile, the small amount of labeled data for updating the neural network can deliver information of label to the encoder under the bi-level optimization update.

{\small
\bibliographystyle{ieee_fullname}