text:
\section{Introduction}
\blfootnote{
    %

    %

    %

    %
     \hspace{-0.65cm}
     This work is licensed under a Creative Commons
     Attribution 4.0 International Licence.
     Licence details:
     \url{http://creativecommons.org/licenses/by/4.0/}.
    %

    %

}
Pre-trained language models have received great interest in the natural language processing (NLP) community in the last recent years \citep{dai_semi-supervised_2015,radford_improving_2018,howard_universal_2018,baevski_cloze-driven_2019,dong_unified_2019}.
These models are trained in a semi-supervised fashion to learn a general language model, for example, by predicting the next word of a sentence \citep{radford_improving_2018}.
Then, transfer learning \citep{pan_survey_2010,ruder-etal-2019-transfer,moeed-etal-2020-evaluation}
can be used to leverage the learned knowledge for a down-stream task, such as text-classification \citep{NIPS2005_2843,aggarwal_survey_2012,Reimers:2019:ACL,10.1007/978-3-030-32381-3_16}.

\citet{devlin_bert:_2019} introduced the  ``Bidirectional Encoder Representations from Transformers'' (BERT), a pre-trained language model based on the Transformer architecture \citep{vaswani_attention_2017}. BERT is a deeply bidirectional model that was pre-trained using a huge amount of text with a masked language model objective where the goal is to predict randomly masked words from their context \citep{taylor_cloze_1953}. The fact is, BERT has achieved state of the art results on the ``General Language Understanding Evaluation'' (GLUE) benchmark \citep{wang_glue:_2018} by only training a single, task-specific layer at the output and fine-tuning the base model for each task.
Furthermore, BERT demonstrated its applicability to many other natural language tasks since then including but not limited to sentiment analysis \citep{sun-etal-2019-utilizing,xu-etal-2019-bert,li-etal-2019-exploiting}, relation extraction \citep{baldini-soares-etal-2019-matching,10.1007/978-3-030-32236-6_65} and word sense disambiguation \citep{huang-etal-2019-glossbert,hadiwinoto-etal-2019-improved,huang-etal-2019-glossbert}, as well as its adaptability to languages other than English \citep{martin-etal-2020-camembert,antoun-etal-2020-arabert,agerri-etal-2020-give}.
However, the fine-tuning data set often contains thousands of labeled data points.
This plethora of training data is often not available in real world scenarios \citep{tan_empirical_2008,wan_using_2008,salameh_sentiment_2015,fang_model_2017}.

In this paper, we focus on the low-resource setting with less than 1,000 training data points.
Our research attempts to answer the question if pool-based active learning can be used to increase the performance of a text classifier based on a Transformer architecture such as BERT.
That leads to the next question: How can layer freezing techniques \citep{yosinski_how_2014,howard_universal_2018,peters_tune_2019}, i.e. reducing the parameter space, impact model training convergence with fewer data points?

To answer these questions, we explore the use of recently introduced Bayesian approximations of model uncertainty \cite{gal_dropout_2015} for data selection that potentially leads to faster convergence during fine-tuning by only introducing new data points that maximize the knowledge gain of the model. To the best of our knowledge, the work presented in this paper is the first demonstration of combining modern transfer learning using pre-trained Transformer-based language model such as the BERT model with active learning to improve performance in low-resource scenarios.
Furthermore, we explore the effect of trainable parameters reduction on model performance and training stability by analyzing the layer-wise change of model parameters to reason about the selection of layers excluded from training.

The main findings of our work are summarized as follows: a) we found that the model's classification uncertainty on unseen data can be approximated by using Bayesian approximations and therefore, used to efficiently select data for manual labeling in an active learning setting; b) by analyzing layer-wise change of model parameters, we found that the active learning strategy specifically selects data points that train the first and thus more general natural language understanding layers of the BERT model rather than the later and thus more task-specific layers.

\section{Methods}

\subsection{Base Model}
In \citep{devlin_bert:_2019} a simple classification architecture subsequent to the output of the Transformer is used to calculate the cross-entropy of the classifier for a text classification task with $C$ classes.
Specifically, first, a dropout operation \citep{srivastava_dropout:_2014} is applied to the Transformer's last layer hidden state of the special \texttt{[CLS]} token that is inserted at the beginning of each document. The regularized output is then fed into a single fully-connected layer with $C$ output neurons and a softmax activation function to scale the logits of its output to probabilities of class association.

\begin{figure}[ht]
      \centering
      \includegraphics[width=\textwidth]{figures/architecture.pdf}
      \caption{High-level overview of the architecture used in our experiments including a depiction of the simple FFNN decoder used by \citet{devlin_bert:_2019} highlighting where our model differs from the original experiment setup.}
      \label{fig:arch}
\end{figure}

In contrast, we use a more complex classification architecture based on a convolutional neural network (CNN) following \citet{kim_convolutional_2014}\footnote{We found that using a CNN is slightly better than using a simple feed forward neural network in our experiments.}.
All hidden states in the last layer of the BERT model are arranged in a 2-dimensional matrix.
Then, convolutional filters of height ($3, 4, 5$) and length corresponding to the hidden state size ($786$) are shifted over the input to calculate $64$ 1-dimensional feature maps per filter size. These feature maps are then batch-normalized \citep{ioffe_batch_2015}, dropout regularized and global max-pooled before they are concatenated and fed into 2 fully-connected layers, each of which applies another dropout operation on its input \footnote{All implementation details can be found in the source code provided in
\ifcolingfinal
\raggedright\url{https://gitlab.mi.hdm-stuttgart.de/griesshaber/bayesian-active-learning}
\else
\raggedright\url{https://anonym}
\fi}.

\subsection{How to Select Data}
When labeled training data is sparse, but unlabeled in-domain data is readily available, manual labeling of all data is often not feasible due to cost.
In this scenario, it is advantageous to let the current model select $Q$ data points that it is most confused about from the pool of potential but yet unlabeled training elements ($U$) to be labeled by an human annotator. Then, they can be included in the training set $T_{new} = T_{old}\cup U_{{x\in U}_{arg max}(a(x, \mathcal{M}))\left<1,...,Q\right> }$ where $a(x, \mathcal{M})$ is a function predicting the knowledge gain of the model $\mathcal{M}$ when trained on the datum $x$.
Thus, the model is specifically trained on data that it can not yet confidently classify to maximize the knowledge gain in each training step, while keeping the cost of labeling new data constant.

We propose to estimate model prediction uncertainty by using Bayesian approximations as presented in \citet{gal_dropout_2015} for data selection process. The main idea is to leverage stochastic regularization layers (e.g. Dropout or Gaussian noises) that can be used to approximate model uncertainty ($\phi(x)$) on any datum ($x$) by performing multiple stochastic forward passes for each element in $U$.
This is implemented by applying the layer operation not only during training but also during the interference stage of the model.
Multiple forward passes of the model $\mathcal{M}$ with the same parameters ($\theta$) and inputs ($x$) thus yield different model outputs ($y$), as each pass samples a concrete model from a approximate distribution $q^*_\theta(\omega)$ that minimizes the Kullback-Leibler divergence to the true model posterior $p(\omega | T)$.
\citet{gal_dropout_2015} call this Monte-Carlo-Dropout (MC-Dropout) as the repeated, non-deterministic forward passes of the model can be interpreted as a Monte-Carlo process.

To decide which elements in $U$ are chosen, different acquisition functions ($a(x, \mathcal{M})$) could be used. In this work, we focus on the ``Bayesian Active Learning by Disagreement (BALD)'' \citep{houlsby_bayesian_2011} acquisition strategy because it demonstrated very good performance in comparison to other strategies in the experiments by \citet{gal_deep_2017} as well as in our own preliminary experiments.
BALD calculates the information gain for the model's parameters that can be achieved with the new data points, that is, the mutual information between predictions and parameters $\mathbb{I}\left[ y, \omega | x, T\right]$.
Hence, this acquisition function has maximum values for inputs that produce disagreeing predictions with high uncertainty.
This is equivalent to high entropy in the logits (the unscaled output of the network before the final softmax normalization) in a classifier model, as multiple (stochastic) forward passes of the model with the same input yield different classification results.
We use the same approximation of BALD that \citet{gal_deep_2017} used in their work (equation \ref{eqn:baldapprox}) where $\widehat{{p}}^{s}_{c}$ is the probability of class association (i.e. the softmax scaled logits) for input $x$ and class $c$ for one of $S$ samples $\widehat{\omega}_s \sim q^*_\theta(\omega)$ from the model's approximated posterior distribution, i.e. $\widehat{{p}}^{s} = \mathrm{softmax}(\mathcal{M}(y | x, \theta, \widehat{\omega}_s))$.

\begin{equation}
\label{eqn:baldapprox}
\begin{split}
&a_{BALD}(x, \mathcal{M}) = \\*
&-\sum_{c\in C}{\left(\frac{1}{S} \sum_{s=1}^{S}{ \widehat{{p}}^{s}_{c} }\right)\;
                log\left(\frac{1}{S} \sum_{s=1}^{S}{ \widehat{{p}}^{s}_{c} }\right)}
+ \frac{1}{S} \sum_{c\in C,s=1}^{s=S}{ \widehat{{p}}^{s}_{c}\;
                                    log\,\widehat{{p}}^{s}_{c}  }
\approx \mathbb{I}\left[ y, \omega | x, T\right]
\end{split}
\end{equation}

As a baseline, we compare the BALD strategy with random sampling of elements in $U$. That is, $a_{Rand}(x, \mathcal{M}) = unif[0,1)$, where $unif[a,b)$ is a sample from the uniform distribution in the half-closed interval $[a,b)$.

When using the BALD acquisition function, we sample for $S=50$ forward passes and add $Q=100$ data points with the highest model uncertainty according to the calculated BALD scores. When randomly acquiring new data points for the baseline, no forward passes are needed ($S=0$) while the number of acquisitions $Q$ stays constant. The MC-Dropout layers
that apply the dropout operation also during a stochastic forward pass of the model, are only placed in the classification architecture, thus the base model is used unaltered with regular dropout layers only active during the training phase.
In our CNN architecture, the dropout is only applied in the penultimate layer of the network. We use the same dropout rate of $0.1$ as \citet{devlin_bert:_2019} in the decoder. This means that the stochastic forward passes only affect the output of the layers after the first MC-Dropout layer in the model. Thus, only a single and relativley expensive pass through the transformer is needed in each learning step while the required multiple passes though the subsequent classifiers are comparatively cheap.

\subsection{How to Fine-tune Models}
\noindent \textbf{Reducing the Number of Trainable Parameters}
Freezing of parameters can be useful when fine-tuning a complex model, as it effectively lowers the number of parameters that need to be tuned during training \citep{howard_universal_2018}.
This is especially relevant in the low-resource setting with a small amount of training data, since the pre-trained BERT\textsubscript{BASE} model with $\sim\!110M$ parameters is over-parametrized for such small data sets.
However, since freezing parameters also lowers the adaptability of a model \citep{peters_tune_2019}, it is crucial to determine which parameters are frozen and which can be fine-tuned during training to not negatively affect the model's performance.

To our best knowledge no prior work has considered the training set size as a dependent parameter.
This parameter is especially important in low-resource settings.
Therefore, we will conduct experiments to visualize the model's performance in dependence of training data size with different sets of layers that are frozen during training.
We denote the number of frozen layers with $F$. For positive values, the layers in the half-closed integer interval $\left[ 0\;..\;F\right)$ are frozen, while negative values represent the interval $\left[ 12\!+\!F\;..\;12\right)$, i.e. the last $-F$ layers of the BERT model.
Layers of the classification architecture are always trainable during training because they are initialized randomly and thus need to be tuned.
We also visualize the change of model parameters in each layer during fine-tuning to reason about the choice of frozen layers.
Results of these experiments are presented in section \ref{sec:freezeresults}.

\noindent \textbf{Analyzing Changes in Layer-Parameters}
To analyze the changes in model parameters during the fine-tuning of the BERT model, we capture the initial state of all model parameters before starting the training.
This includes all pre-trained model weights, as well as the randomly initialized parameters of the CNN denoted by $\theta_0$.
We can then snapshot the same parameters after training for 3 epochs denoted by $\theta_3$.
The change in the model parameters can thus be described by $\Delta\theta = \theta_3 - \theta_0$.
Since the number of parameters is very high ($7,087,872$ for every layer of the BERT model while the exact number of parameters of the CNN depends on the output configuration of the model), we aggregate the change of parameters per layer by calculating the mean absolute difference (MAD) of the parameters.
The mean absolute difference in layer $x$ can be described by equation \ref{eqn:mad} where ${{\theta_e}_x}_n$ is the value of the $n$-th parameter in layer $x$ after training epoch $e$.
\begin{equation}
\label{eqn:mad}
\begin{split}
\overline{\left|\Delta\theta_x\right|}
&= \frac{1}{N}\sum_{i=0}^{N}{\left|\Delta{{\theta}_x}_n\right|}  \\*
&= \frac{1}{N}\sum_{i=0}^{N}{\left|{{\theta_0}_x}_n-{{\theta_3}_x}_n\right|}
\end{split}
\end{equation}

\section{Experiment Setup}
\noindent \textbf{Base Model}
The BERT model, as introduced by \citet{devlin_bert:_2019}, is used to create contextualized embeddings in all experiments.
Specifically, the pre-trained language model \textbf{BERT\textsubscript{BASE}}\footnote{Available under \raggedright\url{https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip}} with 12 layers, 768 hidden states and 12 self-attention heads made available online by the authors is used.

\noindent \textbf{Data Sets}
Similar to the experiments presented by \citet{devlin_bert:_2019}, we use the ``GLUE Multi-Task Benchmark'' \citep{wang_glue:_2018} consisting of multiple NLP classification tasks including sentiment analysis and textual entailment to evaluate the performance of the different model configurations. Contrary to \citeauthor{devlin_bert:_2019}, we report the average accuracy of $N=3$ runs instead of the maximum achieved value for each setting. By doing so, we are able to analyze training stability by comparing the distribution of the results over different runs.

\noindent \textbf{Low-Resource Scenarios}
To simulate the low-resource setting in the GLUE tasks, we repeatedly evaluate the model's performance after training for $3$ epochs on a subset of the available data.
Since data points in the training set provided by the GLUE Benchmark are already shuffled, the subset $S_x$ simply contains the first $x$ data points ($S_x=\{s_1, s_2, \dots, s_x\}$) to ensure the same data selection between experiments.

\section{Results}
\label{sec:results}
\begin{figure}[ht]
\begin{minipage}{.45\textwidth}

    \centering
    \includegraphics[width=\columnwidth]{figures/dataset_sizes.pdf}
    \caption{Accuracies of the BERT model when adapted on $100\%$ and $10\%$ (using random aquisition) of the available training data respectively. }
    \label{fig:res-size}

\end{minipage}
\hspace{0.09\textwidth}
\begin{minipage}{.45\textwidth}

    \centering
    \includegraphics[width=\columnwidth]{figures/active_learn_cnn.pdf}
    \caption{Model accuracies when adapted with data selected with the BALD acquisition strategy vs. random data selection.}
    \label{fig:res-activelearn}

\end{minipage}
\end{figure}

\subsection{Effect of Training Data Size}

As expected, all models benefit from an increased number of training elements until reaching comparable performance to the results reported by \citeauthor{devlin_bert:_2019} when trained on the whole data set.
However, Figure \ref{fig:res-size} shows that, while the accuracy is generally increasing for all experiments when introducing more training data, the model has a much larger increase in performance at the beginning, while adding another partition of the data set to an already large training set can only marginally improve the model's accuracy.
Adding more data, however, generally reduces the variance in model performance between runs when randomly initializing the classification architecture's parameters.
Due to these observations and for comparability of the results between data sets, the experiments in the following subsections are all performed on an initial training set of $S_{10}$ and reevaluated repeatably for $9$ iterations after adding another partition of $100$ elements until the final subset $S_{910}$ is reached. While at this point the performance has not yet converged to the final accuracies for any dataset, the performance changes drastically in these early training stages, which is ideal for our experimental comparison of training speed increases.

\subsection{Active Learning}
The active learning setting requires a pool of unlabeled data to pick the next training elements to acquire. Ideally, the knowledge gain $a(x, \mathcal{M})$ of all elements in the dataset would be approximated, however due to the need of multiple forward passes for each element in the pool to approximate $a$ only a subset $U = S_{20,000} \setminus S_{10} $ was chosen for our experiments to seed up training.
Thus, the size of $U$ can be seen as a hyperparameter to speed up training with the trade off of the risk to exclude highly relevant data points from the available pool data.
The size of $U$ was chosen by increasing it for multiple runs until the performance increase leveled off and acquisition times were still acceptable.
Since training with the random strategy does not have the same trade-off, all available data points in the pool can be used for the acquisition without lengthening the training.

The results in Figure \ref{fig:res-activelearn} show a better mean accuracy for all models and generally across all training set sizes when the BALD strategy is used, compared to picking random training elements from $U$.
Another noteworthy observation is the lower spread of model accuracy between experiments when training data is selected based on the active learning strategy.
These results show a strong indication that the Monte-Carlo approximation of model uncertainty works for state of the art Transformer architectures like the BERT model and can improve training performance when an active learning scenario in a low-resource setting is feasible.

\newcolumntype{Y}{>{\centering\arraybackslash}X}

\subsection{Layer Freezing to Reduce Number of Parameters}
\label{sec:freezeresults}

\begin{figure}[ht]

    \begin{subfigure}{\textwidth}
      \centering
      \includegraphics[width=\columnwidth]{figures/weightdelta_comparison_modelweights_normal.pdf}
      \caption{random sampling of new training elements from $U$}
      \label{fig:sfig1}
    \end{subfigure}
    \newline%
    \begin{subfigure}{\textwidth}
      \centering
      \includegraphics[width=\columnwidth]{figures/weightdelta_comparison_cnn.pdf}
      \caption{selective acquisition of elements with maximum BALD score from $U$}
      \label{fig:sfig2}
    \end{subfigure}

    \caption{Visualization of the mean absolute difference of parameters when training a classifier model on the QNLI data set. The blue bars indicate the variance over all weight changes. The red line indicates the border between the BERT model and the decoder CNN. Layers 12, 13 and 14 are the parallel convolutional layers, 15 and 16 the fully-connected hidden and output layers.
    \label{fig:res-weightdelta}

\end{figure}

Figure \ref{fig:res-weightdelta} visualizes the mean absolute difference of model parameters on a per-layer basis when fine-tuning on the QNLI data set.
Note that the frozen layers have a mean absolute difference of $0$ as the training does not alter any parameters in these layers.
When randomly sampling from $U$ is used to choose new data elements, layers closer to the output of the model generally have a higher change in their parameters during training which is in line with the findings by  \citet{yosinski_how_2014}, \citet{howard_universal_2018} and \citet{hao_visualizing_2019} which show that, while fine-tuning a Transformer-based language model, the front layers have a more general language understanding while the later layers capture more task specific concepts and thus need to be trained more.
However, if the active learning strategy is used, layers closer to the output of the BERT model have a smaller value change while the relative mean absolute difference in the layers of the CNN is comparable between the two strategies.
This indicates that the active learning strategy specifically selects data points that train the first and thus more general layers of the BERT model.

Freezing of layers with a smaller overall change in parameters during training may be beneficial in low-resource scenarios, since this reduces the number of parameters in the model equally while potentially enabling more freedom in the layers that need to be tuned more.
Based on this hypothesis, we compare the model's performance when freezing different numbers of layers starting on both, the input and output of the BERT model.

\begin{table}[hb]

\begin{minipage}{.45\textwidth}
    \newcommand{\specialcell}[2][c]{%
    \begin{tabular}[#1]{@{}c@{}}#2\end{tabular}
    }

\begin{tabularx}{\columnwidth}{|c|Y|Y|Y|Y|}
	\hline
    $F$ &                MNLI                &                QNLI                &                SST2                &                SNLI                 \\
	\hline
	 0  &  $0.53 \pm 0.021$  &  $0.76 \pm 0.010$  &  $0.78 \pm 0.059$  &  $0.67 \pm 0.015$  \\ \hline
	 3  &  $0.51 \pm 0.021$  &  $0.78 \pm 0.003$  &  $0.80 \pm 0.045$  &  $0.69 \pm 0.002$  \\ \hline
	-3  &  $0.52 \pm 0.010$  &  $0.78 \pm 0.002$  &  $0.84 \pm 0.013$  &  $0.69 \pm 0.008$  \\ \hline
	 6  &  $0.51 \pm 0.024$  &  $0.75 \pm 0.014$  &  $0.81 \pm 0.006$  &  $0.63 \pm 0.014$  \\ \hline
	-6  &  $0.47 \pm 0.020$  &  $0.77 \pm 0.010$  &  $0.64 \pm 0.094$  &  $0.64 \pm 0.067$  \\ \hline
\end{tabularx}
    \caption{Mean accuracies of the model after training on $910$ data points using the BALD acquisition strategy.}
    \label{tab:res-freezeaccuracies}
    \end{minipage}
    \hspace{0.09\textwidth}
\begin{minipage}{.45\textwidth}

  \centering
    \begin{tabularx}{\columnwidth}{|c|Y|Y|Y|Y|}
	\hline
    $F$ &  MNLI   &  QNLI   &  SST2   &  SNLI    \\
	\hline
	 0  & $0.054$ & $0.049$ & $0.108$ & $0.061$ \\
	 3  & $0.043$ & $0.032$ & $0.078$ & $0.032$ \\
	-3  & $0.038$ & $0.024$ & $0.047$ & $0.028$ \\
	 6  & $0.050$ & $0.033$ & $0.071$ & $0.082$ \\
	-6  & $0.050$ & $0.050$ & $0.113$ & $0.161$ \\ \hline
\end{tabularx}
    \caption{Mean width of the confidence intervals, i.e. the difference between upper and lower limit, over different runs using the BALD acquisition strategy.}
    \label{tab:res-decoder}

\end{minipage}
\end{table}
Table \ref{tab:res-freezeaccuracies} shows a general increase in model performance when freezing 25\% of the BERT's layers, indicating that a reduction of parameters in the low-resource setting is indeed beneficial.
However, with 50\% of the layers fixed during training, the average performance of the model decreases again, in many cases even below the baseline where all parameters are trainable.
This is an indicator that the BERT model needs fine-tuning and freezing a large ratio of layers may result in a model that can not adapt adequately to the task and is in line with the results of \citet{peters_tune_2019}.

Figure \ref{fig:res-freeze} and Table \ref{tab:res-decoder} also show the mean accuracy and bounds of these experiments over a range of different numbers of training elements.
The average accuracy of the models after multiple training trials do not show any conclusive advantage of freezing the later BERT layers over the the ones in the front.
However, for all data sets the models are more stable during the training over different runs, indicated by a lower average width of the bounds when using $F=-3$.

\begin{figure}[ht]
    \centering
    \includegraphics[width=\columnwidth]{figures/layer_freezing.pdf}
    \caption{The effect of layer-freezing on the stability of the training during active learning using the BALD strategy. The semi-transparent area in the chart is the confidence interval over multiple runs of the configuration, while the solid line denotes the average accuracy.}
    \label{fig:res-freeze}
\end{figure}

\section{Qualitative Analysis}
\noindent \textbf{Overall Observation}
Figure \ref{fig:res-alaqu} visualizes the number of training elements in $T$ grouped by their label ${}_c$ after each acquisition iteration of picking $100$ elements from $U$.
As the labels in the data sets and thus the subset $U$ are equally distributed, we expect the same equal distribution when sampling randomly from the pool elements, which is apparent in the lower row of Figure \ref{fig:res-alaqu} where the training data was actually randomly chosen from the pool of data.
In contrast, we observed that when using the BALD acquisition with non-deterministic forward passes, the distribution shows a stronger bias to a particular class.
This bias increases when sampling more pool data, whereas the difference between the biggest and smallest class in $T$ stays constant during random acquisition (see table \ref{tab:res-acqdiff}).
Furthermore, when randomly sampling, the class with the most or least training elements (${\mathrm{arg\,min}}_{c}\left|T_c\right|$, ${\mathrm{arg\,max}}_{c}\left|T_c\right|$) is changing in many iterations, whereas in the case where active learning is used, this is constant for the most part and changing, if any, only in the first iterations.

\noindent \textbf{Which Class?}
As the goal of active learning is to maximize the knowledge gain of the model with minimum cost, the choice of data the acquisition strategy selects from $U$ may give further insight into the models understanding of the input data.
One characteristic when applying active learning on the MNLI data set, that is apparent in Figure \ref{fig:res-alaqu}, is that samples from the \texttt\small\texttt{neutral} class are queried the most. It might indicate that the model is most confused about the samples from this class. This observation may be justified by the fact that the hypothesis stated in that datum may be unrelated to the associated premise, confusing the model.

Another observation is that the model queries data points from the \texttt\small\texttt{contradiction} class more often than data points from the \texttt\small\texttt{entailment} class at the beginning of the training, while at the end number of queries from the different classes is almost equal.
This indicates that, at the start of training, the model is more confused about contradictions, which may be explained by the fact that a contradiction can differ from an entailment only by a single negation at the correct place in the input, making it harder to differentiate from an entailment.
An example of negated statements in a contradiction can be seen in the last row of Table \ref{tbl:queryexamples} which also shows that the model indeed gave this example a high BALD score.

Table \ref{tbl:queryexamples} shows some of the examples of the MNLI dataset with their calculated BALD scores in the middle of the training after 5 acquisition iterations.
These examples show another noteworthy behaviour of the active learner.
Samples from the entailment class where many words in the two input sentences match or and have a similar wording in general, get a low BALD score, indicating that the model is already confident that it is able to correctly classify those examples.
The same is true for input pairs that differ in wording and contradict themselves.
However, entailing pairs where the wording is mostly different between the two inputs, as well as contradicting pairs with similar phrasing get a much higher BALD score and thus are more likely to be sampled during the acquisition phase.
The model at this point in training thus seems to already have learned to perform its task confidently on the \textit{simpler} examples and can thus concentrate more on the \textit{non-trivial} data-points.

\begin{table}[h!t]
\begin{minipage}{.45\textwidth}
    \begin{tabularx}{\columnwidth}{cc|c|c|c|c|}
    \cline{3-6}
    \multicolumn{1}{l}{}                                & \multicolumn{1}{l|}{} & \multicolumn{4}{c|}{Data set}                            \\ \cline{3-6}
    \multicolumn{1}{l}{}                                &                       & \multicolumn{2}{c|}{BALD} & \multicolumn{2}{c|}{Random} \\ \cline{3-6}
                                                        &                       & MNLI        & QNLI        & MNLI          & QNLI        \\ \hline
    \multicolumn{1}{|c|}{\multirow{9}{*}{\rotatebox{90}{$\left|T\right|$}}} & 110                     & 13          & 30          & 3             & 4           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 210                     & 19          & 24          & 3             & 2           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 310                     & 20          & 16          & 1             & 2           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 410                     & 29          & 18          & 2             & 2           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 510                     & 30          & 36          & 9             & 4           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 610                     & 28          & 40          & 6             & 6           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 710                     & 37          & 44          & 3             & 2           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 810                     & 34          & 44          & 4             & 4           \\ \cline{2-6}
    \multicolumn{1}{|c|}{}                              & 910                     & 25          & 36          & 9             & 4           \\ \hline
    \end{tabularx}
    \caption{Differences in number of elements in the largest and smallest group in the trainset: $\Delta\left|T\right| = \mathrm{max}(\left|T_c\right|) - \mathrm{min}(\left|T_c\right|)$}
    \label{tab:res-acqdiff}
\end{minipage}
\hspace{.09\columnwidth}
\begin{minipage}{.45\textwidth}

\centering
\includegraphics[width=\columnwidth]{figures/al_acquisitions.pdf}
\captionof{figure}{Distribution of training elements in $T$ for different acquisition strategies with the BERT model. $\left|T_c\right|$ is the number of elements in $T$ labeled with ${}_c$}
\label{fig:res-alaqu}
\end{minipage}
\end{table}

\begin{table}[h!t]
    \begin{tabularx}{\textwidth}{|c|r|X|X|l|}
    \hline
queried & $a_{BALD}(x, \mathcal{M})$ & $x_1$                                                                                & $x_2$                                                    & $y$           \\ \hline
no      & $-8.2\times 10^8$          & The view of the illuminated mont at night is spectacular from the other side of the bay.                                                                                & The view of the lit up mont at night is wonderful.       & entailment    \\ \hline
no      & $-5.9\times 10^8$          & At the end of Rue des Francs-Bourgeois is what many consider to be the city's most handsome residential square, the Place des Vosges, with its stone and red brick facades. & Place des Vosges is constructed entirely of gray marble. & contradiction \\ \hline
yes     & $6.2\times 10^8$           & The grounds, like those of any other private house after nightfall, seemed untenanted.                                                                                & The grounds were empty.                                  & entailment    \\ \hline
yes     & $8.6\times 10^8$           & And Forrest saying: "We don't give medals, Sergeant.                                                                                & Forrest said they gave medals ever week.                 & contradiction \\ \hline
yes     & $5.2\times 10^8$           & I did not wonder at John objecting to his beard.                                                                                & John had no objections to his beard.                     & contradiction \\ \hline
\end{tabularx}
    \caption{Examples from the MNLI Dataset with the two input sentences and the associated class label. The first column shows whether or not this data point was queried by the active learner during the complete training, $a_{BALD}(x, \mathcal{M})$ is the calculated BALD score of the input after 5 acquisition iterations ($\left|T\right|=510$).}
    \label{tbl:queryexamples}
\end{table}

\section{Related Work}
\label{sec:related}
\noindent \textbf{Low-resource NLP} Previous work in low-resource NLP tasks includes feature-engineering \citep{tan_empirical_2008} which requires a recurring effort when adapting to a new data set.
Another approach is to transfer knowledge across domains to increase the amount of data that is available for training \citep{zoph_transfer_2016,nguyen_transfer_2017,kocmi_trivial_2018,yang_transfer_2017,gupta_semi-supervised_2018}.
One of these approaches relied on adversarial training \citep{goodfellow_generative_2014} to learn a domain adaptive classifier \citep{ganin_domain-adversarial_2016} in another domain or language where training data was plentiful while ensuring that the model generalizes to the low-resource domain \citep{chen_adversarial_2017}.
However, these approaches have not used a pre-trained generic language model, but perform pre-training for each task individually.

\noindent \textbf{Adapting pre-trained models} The effectiveness of transfer learning in low-resource settings was previously demonstrated for machine translation \citep{zoph_transfer_2016,nguyen_transfer_2017,kocmi_trivial_2018}, sequence tagging \citep{yang_transfer_2017} and sentiment classification \citep{gupta_semi-supervised_2018}. However, all this prior work does not use general LMs but transfers knowledge from a pre-trained model to a close to target corpus.
Some previous work has analyzed the performance behavior of the BERT model in different scenarios.
\citet{hao_visualizing_2019} showed that a classifier that fine-tunes a pre-trained BERT model generally has wider optima on the training loss curves in comparison to models trained for the same task from scratch, indicating a more general classifier \citep{chaudhari_entropy-sgd:_2017,li_visualizing_2018,izmailov_averaging_2018}.
\citet{peters_tune_2019} examine the adaption phase of LM based classifiers by comparing fine-tuning and feature extraction where the LMs parameters are fixed.
In contrast, we focus on the low-resource setting where less than 1,000 data points are available for fine-tuning.

\noindent \textbf{Layer freezing in deep Transformers} Experiments by \citet{yosinski_how_2014} indicated that the first layers of a LM capture a more general language understanding, while later layers capture more task-specific knowledge.
With this motivation, \citet{howard_universal_2018} introduced \textit{gradual unfreezing} of the Transformer layers during each epoch, beginning with the last layer.
\citet{hao_visualizing_2019} analyzed the \textit{loss surfaces} in dependency of the model parameters before and after training and came to the same conclusion that lower layers contain more transferable features. However, none of the work has considered the training set size as a dependent parameter as our experiments presented in this paper.

\noindent \textbf{Active learning in NLP}
There is some prior work regarding active learning for NLP tasks using deep neural networks.
\citet{zhang_active_2017} explored pool-based active learning for text classification using a model similar to our setting. However, they used word-level embeddings \citep{mikolov_efficient_2013} and focus on representation learning, querying pool points that are expected to maximize the gradient changes in the embedding layer.
\citet{shen_deep_2017} used active learning for named entity recognition tasks. They proposed a acquisition strategy named \textit{Maximum Normalized Log-Probability} which is a normalized form of the \textit{Constrained Forward-Backward} confidence estimation \citep{culotta_confidence_2004,culotta_reducing_2005}.
Using this strategy, they achieved on-par performance in comparison to a model using the BALD acquisition function and MC Dropout without needing multiple forward passes.
However, this approach is not suitable for any arbitrary model architecture but requires conditional random fields (CRFs) for the approximation of model uncertainty.

\section{Conclusion}
In this paper, we evaluated the performance of a pre-trained Transformer model - BERT - in an active learning scenario for text classification in low-resource settings.
We showed that using Monte-Carlo Dropout in the classification architecture is an effective way to approximate model uncertainty on unlabeled training elements.
This technique enables us to select data for annotation that maximize the knowledge gain for the model fine-tuning process. Experimental results on GLUE data set show that it improves both model performance and training stability.
Finally, in order to improve the efficiency of the fine-tuning process with a small amount of data, we explored the reduction of trainable model parameters by freezing layers of the BERT model up to a certain level of depth. Comparing the exclusion of layers in the front or the back of the BERT model from training, we found it to be advantageous for training stability when freezing the layers closest to the output.

\ifcolingfinal
\section*{Acknowledgments}
This research and development project is funded within the "Future of Work" Program by the German Federal Ministry of Education and Research (BMBF) and the European Social Fund in Germany. It is implemented by the Project Management Agency Karlsruhe (PTKA). The authors are responsible for the content of this publication.
\fi

\bibliographystyle{coling}

redpajama_set_name: RedPajamaArXiv


text:
\section{Experiment Details}
\label{sec:details}

In Table~\ref{tbl:hyperparameter}, we show the complete set of hyperparameters in our semi-supervised learning experiments.

\input{table/hyperparameter}

The strong augmentation $\mathrm{Aug_s}$ on ImageNet unlabeled data uses color distortion in addition to the standard crop-and-flip.
A pseudo-code for the color distortion in PyTorch is as follows:

\noindent\texttt{\small
from torchvision import transforms as T\newline
color\_jitter = T.ColorJitter(0.4,0.4,0.4,0.1)\newline
transforms.Compose([}

\texttt{\small
	T.RandomApply([color\_jitter], p=0.8)
}

\texttt{\small
	T.RandomGrayscale(p=0.2)])
}

\section{MB and MQ in CoMatch}
Figure~\ref{fig:mb} illustrates how the EMA model is utilized in CoMatch to construct the memory bank (MB) and the momentum queue (MQ).
The memory bank contains the class probability and the low-dimensional embeddings for both weakly-augmented labeled samples and weakly-augmented unlabeled samples.
The momentum queue contains the pseudo-labels for the unlabeled samples and their strongly-augmented embeddings.

\input{table/fig_mb_mq}

\section{Pseudo-code of CoMatch}
Algorithm~\ref{alg} presents the pseudo-code of CoMatch.
\input{table/algorithm}

\end{appendices}

\section{Conclusion}
\vspace{-0.5ex}
To conclude,
CoMatch is a state-of-the-art semi-supervised learning method that achieves substantial improvement.
The key of its success can be attributed to three novel contributions: (1) co-training of two representations, (2) memory-smoothed pseudo-labeling, and (3) graph-based contrastive learning.
We believe that CoMatch will help enable machine learning to be deployed in practical domains where labels are expensive to acquire.

\section{Experiment}
\label{sec:experiment}
\vspace{-0.5ex}
We evaluate CoMatch on several standard SSL benchmarks including CIFAR-10~\cite{cifar}, STL-10~\cite{stl10} and ImageNet~\cite{imagenet}.
We focus on the most challenging label-scarce scenario where fewer labels are available.

\vspace{-0.5ex}
\subsection{CIFAR-10 and STL-10}
\vspace{-0.5ex}
First,
we conduct experiments on CIFAR-10 and STL-10 datasets.
CIFAR-10 contains 50,000 images of size $32\times32$ from 10 classes.
We vary the amount of labeled data and experiment with fewer labels than previously considered.
We evaluate on 5 runs with different random seeds.
STL-10 contains 5,000 labeled images of size $96\times96$ from 10 classes and 100,000 unlabeled images including ood samples.
We evaluate on the 5 pre-defined folds.
\input{table/cifar_stl}

\input{table/imagenet}
\noindent\textbf{Baseline methods.}
We improve the current state-of-the-art method FixMatch~\cite{fixmatch} with distribution alignment~\cite{remixmatch} to build a stronger baseline.
We also compare with the original FixMatch and MixMatch~\cite{mixmatch}.
We omit previous methods such as $\Pi$-model~\cite{pi_model}, Pseudo-Labeling~\cite{pseudo}, and Mean Teacher~\cite{mean_teacher} due to their poorer performance as reported in~\cite{fixmatch}.
Following~\cite{realistic}, we reimplemented the baselines and performed all experiments using the same model architecture,
the same codebase (PyTorch~\cite{pytorch}), and the same random seeds.

\noindent\textbf{Self-supervised pre-training}
can provide a good model initialization for semi-supervised learning.
Therefore, we also experiment with models pre-trained using SimCLR~\cite{simclr} for 100 epochs,
following the implementation in~\cite{simclr_pytorch}.

\noindent\textbf{Implementation details.}
We use a Wide ResNet-28-2~\cite{wrn} with 1.5M parameters for CIFAR-10,
and a ResNet-18~\cite{resnet} with 11.5M parameters for STL-10.
The projection head is a 2-layer MLP which outputs 64-dimensional embeddings.
The models are trained using SGD with a momentum of 0.9 and a weight decay of 0.0005.
We train for 200 epochs,
using an learning rate of 0.03 with a cosine decay schedule.
All baselines follow the same training protocol, except for MixMatch~\cite{mixmatch} which is trained for 1024 epochs.
For the hyperparameters in CoMatch that also exist in~\cite{fixmatch},
we follow~\cite{fixmatch} and set $\lambda_{cls}=1$, $\tau=0.95$, $\mu=7$, $B=64$.
For the additional hyperparameters,
we set $\alpha=0.9$, $K=2560$, $t=0.2$, and choose $\lambda_{ctr}\in\{1,5\}$, $T\in\{0.7,0.8\}$.
A complete list of hyperparameters is reported in appendix~\ref{sec:details}.

\noindent\textbf{Augmentations.}
CoMatch uses ``weak'' and ``strong'' augmentations.
The weak augmentation for all experiments is the standard crop-and-flip strategy.
For strong augmentations,
CIFAR-10 uses RandAugment~\cite{randaug} which randomly selects from a set of transformations (e.g., color inversion, translation, contrast adjustment) for each sample.
STL-10 uses the augmentation strategy in SimCLR~\cite{simclr} which applies random color jittering and grayscale conversion.

\noindent\textbf{Results.}
Table~\ref{tbl:cifar_stl} shows the results.
CoMatch substantially outperforms the best baseline (FixMatch w. DA) across all settings.
For example, CoMatch achieves an average accuracy of 81.19\% on CIFAR-10 with only 2 labels per class,
whereas FixMatch (w. DA) has a lower accuracy of 59.42\% and a larger variance.
On STL-10,
CoMatch also improves FixMatch (w. DA) by 19.26\%.
Self-supervised pre-training with SimCLR leads to improved performance for both CoMatch and FixMatch.
However, CoMatch without SimCLR still outperforms FixMatch (w. DA) with SimCLR.
\vspace{-0.5ex}
\subsection{ImageNet}

\input{table/fig_training}
\input{table/ablation}
\vspace{-0.5ex}
We evaluate CoMatch on ImageNet ILSVRC-2012 to verify its efficacy on large-scale datasets.
Following~\cite{S4L,simclr},
we randomly sample 1\% or 10\% of images with labels in a class-balanced way (13 or 128 samples per-class, respectively),
while the rest of images are unlabeled.
Our results are not sensitive to different random seeds hence we use a fixed random seed.

\noindent\textbf{Baseline methods.}
The baselines include (1) semi-supervised learning methods and (2) self-supervised pre-training followed by fine-tuning.
Furthermore, we construct a state-of-the-art baseline which combines FixMatch (w. DA) with self-supervised pre-training using MoCov2~\cite{mocov2} (pre-trained for 800 epochs).
Self-supervised methods require additional model parameters during training due to the projection network.
We count the number of training parameters as those that require gradient update.

\noindent\textbf{Implementation details.}
We use a ResNet-50~\cite{resnet} model as the encoder.
Following~\cite{mocov2,simclr}, the projection head is a 2-layer MLP which outputs 128-dimensional embeddings.
We train the model using SGD with a momentum of 0.9 and a weight decay of 0.0001.
The learning rate is 0.1, which follows a cosine decay schedule for 400 epochs.
For models that are initialized with MoCov2,
we use a smaller learning rate of 0.03.
The momentum parameter is set as $m=0.996$.
Other hyperparameters are shown in appendix~\ref{sec:details}.

\noindent\textbf{Results.}
Table~\ref{tbl:imagenet} shows the result, where CoMatch achieves state-of-the-art performance.
CoMatch obtains a top-1 accuracy of 66.0\% on 1\% of labels.
Compared to the the best baseline (MoCov2 followed by FixMatch w. DA),
CoMatch achieves 6.1\% improvement with $3\times$ less training time.
With the help of MoCov2 pre-training,
the performance of CoMatch can further improve to 67.1\% on 1\% of labels,
and 73.7\% on 10\% of labels.
In Figure~\ref{fig:training},
we further show that CoMatch produces pseudo-labels that are more confident and accurate.
Pre-training with MoCov2 helps speed up the convergence rate.

\input{table/voc}
\input{table/coco}
\subsection{Ablation Study.}
We perform extensive ablation study to examine the effect of different components in CoMatch.
We use ImageNet with 1\% labels as the main experiment.
Due to the number of experiments in our ablation study,
we report the top-1 accuracy after training for 100 epochs,
where the default setting of CoMatch achieves 57.1\%.

\noindent\textbf{Graph connection threshold.}
The threshold $T$ in eqn.(\ref{eqn:graph_q}) controls the sparsity of edges in the pseudo-label graph.
Figure~\ref{fig:ablation}(a) presents the effect of $T$.
As $T$ increases, samples whose pseudo-labels have lower similarity are disconnected.
Hence their embeddings are pushed apart by our contrastive loss.
When $T=1$, the proposed graph-based contrastive loss downgrades to the self-supervised loss in eqn.(\ref{eqn:contrastive}) where the only connections are the self-loops.
Using the self-supervised contrastive loss decreases the performance by 2.8\%.

\noindent\textbf{Contrastive loss weight.}
We vary the weight $\lambda_{ctr}$ for the contrastive loss $\mathcal{L}_u^{ctr}$ and report the result in Figure~\ref{fig:ablation}(b),
where $\lambda_{ctr}=10$ gives the best performance.
With 10\% of ImageNet labels, $\lambda_{ctr}=2$ yields better performance.
We find that in general,
fewer labeled samples require a larger $\lambda_{ctr}$ to strengthen the graph regularization.

\noindent\textbf{Pseudo-label smoothness.}
Our memory-smoothed pseudo-labeling uses $\alpha$ to control the balance between the model's prediction and smoothness constraint.
Figure~\ref{fig:ablation}(c) shows its effect,
where $\alpha=0.9$ results in the best performance.
When $\alpha=1$, the pseudo-labels completely rely on the model's prediction,
which decreases the accuracy by 2.1\% due to confirmation bias.
When $\alpha<0.9$, the pseudo-labels are over-smoothed.
A potential improvement is to apply sharpening~\cite{mixmatch} to pseudo-labels with smaller $\alpha$, but is not studied here due to the need for an extra sharpening hyperparameter.

\noindent\textbf{Size of memory bank and momentum queue.}
$K$ controls both the size of the memory bank for pseudo-labeling
and the size of the momentum queue for contrastive learning.
A larger $K$ considers more samples to enforce a structural constraint on the label space and the embedding space.
As shown in Figure~\ref{fig:ablation}(d),
the performance increases as $K$ increases from 10k to 30k,
but plateaus afterwards.

\subsection{Transfer of Learned Representations}

We further evaluate the quality of the representations learned by CoMatch by transferring it to other tasks.
Following~\cite{benchmark, pcl},
We first perform linear classification on two datasets:  PASCAL VOC2007~\cite{voc} for object classification and Places205~\cite{places} for scene recognition.
We train linear SVMs using fixed representations from ImageNet pre-trained models.
We preprocess all images by resizing them to 256 pixels along the shorter side and taking a 224$\times$224 center crop.
The SVMs are trained on the global average pooling features of ResNet-50.
To study the transferability of the representations in few-shot scenarios,
we vary the number of samples per-class ($k$) in the downstream datasets.

Table~\ref{tbl:voc} shows the results.
We compare CoMatch with standard supervised learning on labeled ImageNet and self-supervised learning (MoCov2~\cite{mocov2} and SwAV~\cite{swav}) on unlabeled ImageNet.
CoMatch with 10\% labels achieves state-of-the-art performance on both datasets,
except for Places with $k=256$.
It is interesting to observe that self-supervised learning methods do not perform well in few-shot transfer,
and only catch up with supervised learning when $k$ increases.

In Table~\ref{tbl:coco},
we also show that compared to supervised learning,
CoMatch provides a better backbone for object detection and instance segmentation on COCO~\cite{coco}.
We follow the exact same setting as~\cite{moco} to fine-tune a Mask-RCNN model~\cite{mask_rcnn} for $1\times$ or $2\times$ schedule.
\section{Introduction}
\label{sec:introduction}
\vspace{-0.5ex}
Semi-supervised learning (SSL)~\textendash~learning from few labeled data and a large amount of unlabeled data~\textendash~has been a long-standing problem in computer vision and machine learning.
Recent state-of-the-art methods mostly follow two trends:
(1) using the model's class prediction to produce a pseudo-label for each unlabeled sample as the label to train against~\cite{pseudo,mixmatch, remixmatch, fixmatch};
(2) unsupervised or self-supervised pre-training, followed by supervised fine-tuning~\cite{simclr,moco,byol,swav} and pseudo-labeling~\cite{bigsemi}.

However,
existing methods have several limitations.
Pseudo-labeling (also called self-training) methods heavily rely on the quality of the model's class prediction,
thus suffering from confirmation bias where the prediction mistakes would accumulate.
Self-supervised learning methods are task-agnostic,
and the widely adopted contrastive learning~\cite{simclr,moco} may learn representations that are suboptimal for the specific classification task.
Another branch of methods explore graph-based semi-supervised learning~\cite{teacher_graph,deep_label_prop},
but have yet shown competitive performance especially on larger datasets such as ImageNet~\cite{imagenet}.

We propose CoMatch, a new semi-supervised learning method that addresses the existing limitations.
A conceptual illustration is shown in Figure~\ref{fig:overview}.
In CoMatch,
each image has two compact representations:
a class probability produced by the classification head and a low-dimensional embedding produced by the projection head.
The two representations interact with each other and jointly evolve in a co-training framework.
Specifically,
the classification head is trained using memory-smoothed pseudo-labels,
where pseudo-labels are smoothed by aggregating information from nearby samples in the embedding space.
The projection head is trained using contrastive learning on a pseudo-label graph,
where samples with similar pseudo-labels are encouraged to have similar embeddings.
CoMatch unifies dominant ideas including consistency regularization,
entropy minimization,
contrastive learning, and graph-based SSL.
Furthermore,
we extend CoMatch with an exponential-moving-average (EMA) model~\cite{moco,mean_teacher} for learning from large-scale data.

\input{table/fig_overview.tex}

We perform experiments on multiple datasets and compare with state-of-the-art semi-supervised and self-supervised methods.
We create strong baselines which perform self-supervised pre-training followed by semi-supervised learning.
CoMatch substantially outperforms all baselines across all benchmarks,
especially in label-scarce scenarios.
On CIFAR-10 with 20 labeled samples, CoMatch outperforms FixMatch~\cite{fixmatch} by 21.77\% in accuracy.
On STL-10, CoMatch outperforms FixMatch by 19.26\%.
On ImageNet with only 1\% of labels, CoMatch achieves a top-1 accuracy of 66.0\% (67.1\% with self-supervised pre-training),
whereas the best baseline (MoCov2~\cite{mocov2} followed by FixMatch~\cite{fixmatch}) has an accuracy of 59.9\%.
Furthermore,
we demonstrate that CoMatch achieves better representation learning performance on down-stream image classification and object detection tasks,
outperforming both supervised learning and self-supervised learning.

\section{Background}
\label{sec:literature}

To set the stage for CoMatch,
we first introduce existing SSL methods,
mainly focusing on current state-of-the-art methods that are relevant.
More comprehensive reviews can be found in~\cite{survey_2005,survey_2020}.
In the following,
we refer to a deep encoder network (a convolutional neural network) as $f(\cdot)$,
which produces a high-dimensional feature $f(x)$ given an input image $x$.
A classification head (a fully-connected layer followed by softmax) is defined as $h(\cdot)$,
which outputs a distribution over classes $p(y|x)=h(f(x))$.
We also define a non-linear projection head (a MLP) $g(\cdot)$,
which transforms a feature $f(x)$ into a normalized low-dimensional embedding $z(x)=g(f(x))$.

\vspace{0.5ex}
\noindent\textbf{Consistency regularization} is a crucial piece for many state-of-the-art SSL methods.
It utilizes the assumption that a classifier should output the same class probability for an unlabeled sample even after it is augmented.
In the simplest form,
prior works~\cite{consistency,temporal} add the following consistency regularization loss on unlabeled samples:
\begin{equation}
\label{eqn:consistency}
\left\lVert {p(y|\mathrm{Aug}(x))-p(y|\mathrm{Aug}(x))} \right\rVert^2_2,
\end{equation}
where $\mathrm{Aug}(\cdot)$ is a stochastic transformation that does not alter the label of the image.
Mean Teacher~\cite{mean_teacher} replaces one of the terms in eq.(\ref{eqn:consistency}) with the output of an EMA model.
VAT~\cite{VAT} uses an adversarial transformation in place of $\mathrm{Aug}$.
MixMatch~\cite{mixmatch} averages predictions across multiple augmentations to produce $p(y)$.
UDA~\cite{UDA}, ReMixMatch~\cite{remixmatch}, and FixMatch~\cite{fixmatch} use a cross-entropy loss in place of the squared error, and apply stronger augmentation.

\vspace{0.5ex}
\noindent\textbf{Entropy minimization} is a common method in many SSL algorithms,
which encourages the classifier's decision boundary to pass through low-density regions of the data distribution.
It is either achieved explicitly by minimizing the entropy of $p(y|x)$ on unlabeled samples~\cite{EntMin},
or implicitly by constructing low-entropy \textbf{pseudo-labels} on unlabeled samples and using them as training targets in a cross-entropy loss~\cite{pseudo,mixmatch,remixmatch,fixmatch}.
Some methods~\cite{UDA,mixmatch,remixmatch} post-process the ``soft'' pseudo-labels with a sharpening function to reduce entropy,
whereas FixMatch~\cite{fixmatch} produces ``hard'' pseudo-labels for samples whose largest class probability fall above a predefined threshold.
Most methods~\cite{fixmatch,remixmatch,UDA} use weakly-augmented samples to produce pseudo-labels and train the model on strongly-augmented samples.
However,
since the pseudo-labels purely rely on the classifier,
such self-training strategy suffers from the confirmation bias problem,
where the error in the pseudo-labels would accumulate and harms learning.

\vspace{0.5ex}
\noindent\textbf{Self-supervised contrastive learning}
has attracted much attention,
due to its ability to leverage unlabeled data for model pre-training.
The widely adopted contrastive learning~\cite{instance,CPC,simclr,bigsemi,moco} optimizes for the task of instance discrimination,
and formulates the loss using the normalized low-dimensional embeddings $z$:
 \begin{equation}
 \label{eqn:contrastive}
 -\log\frac{\exp(z(\mathrm{Aug}(x_i))\cdot z(\mathrm{Aug}(x_i))/t)}{\sum_{j=1}^N\exp(z(\mathrm{Aug}(x_i))\cdot z(\mathrm{Aug}(x_j))/t)}
 \end{equation}
where $\mathrm{Aug}(\cdot)$ is a stochastic transformation similar as in eq.(\ref{eqn:consistency}),
and $x_j$ include $x_i$ and $N-1$ other images (\ie~negative samples).
Self-supervised contrastive learning can be interpreted as a form of class-agnostic consistency regularization,
which enforces the same image with different augmentations to have similar embeddings,
while different images have different embeddings.
Among recent methods,
SimCLR~\cite{simclr} uses images from the same batch to calculate pairwise similarity,
whereas MoCo~\cite{moco} maintains a queue of embeddings from an EMA model.

Self-supervised pre-training followed by supervised fine-tuning has shown strong performance on semi-supervised learning tasks~\cite{simclr,moco,byol,pcl,swav}.
SimCLR v2~\cite{bigsemi} further utilizes larger models for distillation.
However,
since self-supervised learning is a task-agnostic process,
the contrastive loss in eq.(\ref{eqn:contrastive}) optimizes for an objective that partially contradicts with task-specific learning.
It enforces images from the same class to have different representations,
which is undesirable for classification tasks.

\vspace{0.5ex}
\noindent\textbf{Graph-based semi-supervised learning} defines the similarity of data samples with a graph and encourages smooth predictions with respect to the graph structure~\cite{label_prop,Zhu_ICML_2013}.
Recent works use deep networks to generate graph representations.
\cite{deep_label_prop,dense_label_prop} perform iterative label propagation and network training.
\cite{teacher_graph,rem} connect data samples that have the same pseudo-labels,
and perform metric learning to enforce connected samples to have similar representations.
However,
these methods define representations as the high-dimensional feature $f(x)$,
which leads to several limitations: (1) since the features are highly-correlated with the class predictions, the same types of errors are likely to exist in both the feature space and the label space;
(2) due to the curse of dimensionality, Euclidean distance becomes less meaningful;
(3) computation cost is high which harms the scalability of the methods.
Furthermore, the loss functions in~\cite{teacher_graph,rem} consider the absolute distance between pairs,
whereas CoMatch optimizes for relative distance.

\section{Method}
\label{sec:method}

\subsection{Overview}

In this section, we introduce our proposed semi-supervised learning method.
Different from most existing semi-supervised and self-supervised learning methods,
CoMatch jointly learns the encoder $f(\cdot)$,
the classification head $h(\cdot)$,
and the projection head $g(\cdot)$.
Given a batch of $B$ labeled samples $\mathcal{X}=\{(x_b,y_b)\}_{b=1}^B$ where $y_b$ are one-hot labels,
and a batch of unlabeled samples $\mathcal{U}=\{u_b\}_{b=1}^{\mu B}$ where $\mu$ determines the relative size of $\mathcal{X}$ and $\mathcal{U}$,
CoMatch jointly optimizes three losses:
(1) a supervised classification loss on labeled data $\mathcal{L}_x$,
(2) an unsupervised classification loss on unlabeled data $\mathcal{L}_u^{cls}$,
and (3) a graph-based contrastive loss on unlabeled data $\mathcal{L}_u^{ctr}$.
Specifically, $\mathcal{L}_x$ is defined as the cross-entropy between the ground-truth labels and the model's predictions:
\begin{equation}
	\mathcal{L}_x = \frac{1}{B} \sum_{b=1}^{B} \mathrm{H}(y_b,p(y|\mathrm{Aug_w}(x_b))),
\end{equation}
where $\mathrm{H}(y,p)$ denotes the cross-entropy between two distributions $y$ and $p$, and $\mathrm{Aug_w}$ refers to weak augmentations.

The unsupervised classification loss $\mathcal{L}_u^{cls}$ is defined as the cross-entropy between the pseudo-labels $q_b$ and the model's predictions:
\begin{equation}
\label{eqn:u_cls}
\mathcal{L}_u^{cls} = \frac{1}{\mu B} \sum_{b=1}^{\mu B} \mathbbm{1} (\max q_b\geq \tau) \mathrm{H}(q_b,p(y|\mathrm{Aug_s}(u_b))),
\end{equation}
where $\mathrm{Aug_s}$ refers to strong augmentations.
Following FixMatch~\cite{fixmatch},
we retain pseudo-labels whose largest class probability are above a threshold $\tau$.
Different from FixMatch, our soft pseudo-labels $q_b$ are not converted to hard labels for entropy minimization.
Instead, we achieve entropy minimization by optimizing the contrastive loss $\mathcal{L}_u^{ctr}$.
Section~\ref{sec:comatch} explains the details of pseudo-labelling and contrastive learning.

Our overall training objective is:
\begin{equation}
\mathcal{L}=\mathcal{L}_x + \lambda_{cls} \mathcal{L}_u^{cls} + \lambda_{ctr} \mathcal{L}_u^{ctr},
\end{equation}
where $\lambda_{cls}$ and $\lambda_{ctr}$ are scalar hyperparameters to control the weight of the unsupervised losses.

\input{table/fig_framework}

\subsection{CoMatch}
\label{sec:comatch}
In CoMatch,
the high-dimensional feature of each sample is transformed to two compact representations:
its class probability $p$
and its normalized low-dimensional embedding $z$,
which reside in the label space and the embedding space, respectively.
Given a batch of unlabeled samples $\mathcal{U}$,
we first perform memory-smoothed pseudo-labeling on weak augmentations $\mathrm{Aug_w}(\mathcal{U)}$ to produce pseudo-labels.
Then, we construct a pseudo-label graph $W^q$ which defines the similarity of samples in the label space.
We use $W^q$ as the target to train an embedding graph $W^z$,
which measures the similarity of strongly-augmented samples $\mathrm{Aug_s}(\mathcal{U)}$ in the embedding space.
An illustration of CoMatch is shown in Fig~\ref{fig:framework},
and a pseudo-code is given in the appendix.
Next, we first introduce the pseudo-labeling process,
then we describe the graph-based contrastive learning algorithm.

\vspace{0.5ex}
\noindent\textbf{Memory-smoothed pseudo-labeling.}
\vspace{0.5ex}

Given each sample in $\mathcal{X}$ and $\mathcal{U}$,
we first obtain its class probability.
For a labeled sample, it is defined as the ground-truth label:
$p^w=y$.
For an unlabeled sample, it is defined as the model's prediction on its weak-augmentation:
$p^w=h\circ f(\mathrm{Aug_w}(u))$.
Following~\cite{remixmatch},
we perform distribution alignment (DA) on unlabeled samples:
$p^w=\mathrm{DA}(p^w)$.
DA prevents the model's prediction from collapsing to certain classes.
Specifically,
we maintain a moving-average $\tilde{p}^w$ of $p^w$ during training,
and adjust the current $p^w$ with $p^w = \mathrm{Normalize}(p^w/\tilde{p}^w)$,
where $\mathrm{Normalize}(p)_i=p_i / \sum_j p_j$ renormalizes the scaled result to a valid probability distribution.

For each sample in $\mathcal{X}$ and $\mathcal{U}$, we also obtain its embedding $z^w$ by forwarding the weakly-augmented sample through $f$ and $g$.
Then, we create a memory bank to store class probabilities and embeddings of the past $K$ weakly-augmented samples:
$\mathrm{MB}=\{(p_k^w,z_k^w)\}_{k=1}^K$.
The memory bank contains both labeled samples and unlabeled samples and is updated with first-in-first-out strategy.

For each unlabeled sample $u_b$ in the current batch with $p^w_b$ and $z^w_b$,
we generate a pseudo-label $q_b$ by aggregating class probabilities from neighboring samples in the memory bank.
Specifically,
we optimize the following objective:
\begin{equation}
	J(q_b) = (1-\alpha) \sum_{k=1}^K a_k \left\lVert q_b-p^w_k \right\rVert_2^2 + \alpha \left\lVert q_b-p^w_b \right\rVert_2^2 \big.
\end{equation}
The first term is a smoothness constraint which encourages $q_b$ to take a similar value as its nearby samples' class probabilities,
whereas the second term attempts to maintain its original class prediction.
$a_k$ measures the affinity between the current sample and the $k$-th sample in the memory,
and is computed using similarity in the embedding space:
\begin{equation}
a_k = \frac{\exp({z^w_b}\cdot{z^w_k}/t)}{\sum_{k=1}^{K}\exp({z^w_b}\cdot{z^w_k}/t)},
\end{equation}
where $t$ is a scalar temperature parameter.

Since $a_k$ is normalized (\ie~$a_k$ sums to one), the minimizer for $J(q_b)$ can be derived as:
\begin{equation}
\label{eqn:pseudo_label}
q_b  = \alpha  p^w_b + (1-\alpha) \sum_{k=1}^K a_k p^w_k.
\end{equation}

\noindent\textbf{Graph construction and contrastive learning.}
\vspace{0.5ex}

Given the pseudo-labels $\{q_b\}_{b=1}^{\mu B}$ for the batch of unlabeled samples,
we build the pseudo-label graph by constructing a similarity matrix $W^q$ of size $\mu B \times \mu B$:
\begin{equation}
\label{eqn:graph_q}
W^q_{bj} =
\begin{cases}
1 & \text{if}~b=j\\
q_b \cdot q_j & \text{if}~b\neq j \text{~and~} q_b \cdot q_j\geq T\\
0 & \text{otherwise}
\end{cases}
\end{equation}
Samples with similarity lower than a threshold $T$ are not connected,
and each sample is connected to itself with the strongest edge of value 1 (\ie self-loop).

The pseudo-label graph serves as the target to train an embedding graph.
To construct the embedding graph,
we first perform two strong augmentations on each unlabeled sample $u_b\in \mathcal{U}$,
and obtain their embeddings $z_b=g\circ f(\mathrm{Aug_s}(u_b))$, $z'_b=g\circ f(\mathrm{Aug'_s}(u_b))$.
Then we build the embedding graph $W^z$ as:
\begin{equation}
W^z_{bj} =
\begin{cases}
\exp(z_b \cdot z'_b/t)& \text{if}~b=j\\
\exp(z_b \cdot z_j/t)& \text{if}~b\neq j \\
\end{cases}
\end{equation}

We aim to train the encoder $f$ and the projection head $g$ such that the embedding graph has the same structure as the pseudo-label graph.
To this end,
we first normalize $W^q$ and $W^z$ with $\hat{W}_{bj}=W_{bj} / \sum_j W_{bj}$, so that each row of the similarity matrix sums to 1.
Then we minimize the cross-entropy between the two normalized graphs.
The contrastive loss is defined as:

\begin{equation}
	\mathcal{L}_u^{ctr} = \frac{1}{\mu B} \sum_{b=1}^{\mu B} \mathrm{H}(\hat{W^q_b},\hat{W^z_b})
\end{equation}
$\mathrm{H}(\hat{W^q_b},\hat{W^z_b})$ can be decomposed into two terms:
\begin{equation}
  \resizebox{1\hsize}{!}{$
- \hat{W}^q_{bb}  \log (\frac{\exp(z_b \cdot z'_b/t)}{\sum_{j=1}^{\mu B}\hat{W}^z_{bj}} ) - \sum\limits_{j=1,j\neq b}^{\mu B}\hat{W}^q_{bj}  \log (\frac{\exp(z_b \cdot z_j/t)}{\sum_{j=1}^{\mu B}\hat{W}^z_{bj}} )
	$}
\end{equation}
The first term is a self-supervised contrastive loss that comes from the self-loops in the pseudo-label graph.
It encourages the model to produce similar embeddings for different augmentations of the same image,
which is a form of consistency regularization.
The second term encourages samples with similar pseudo-labels to have similar embeddings.
It gathers samples from the same class into clusters,
which achieves entropy minimization.

During training,
a natural curriculum would occur from CoMatch.
The model would start with producing low-confidence pseudo-labels,
which leads to a sparse pseudo-label graph.
As training progresses,
samples are gradually clustered,
which in turns leads to more confident pseudo-labels and more connections in the pseudo-label graph.

Another advantage of CoMatch appears in open-set semi-supervised learning,
where the unlabeled data contains out-of-distribution (ood) samples.
Due to the smoothness constraint,
ood samples would have low-confidence pseudo-labels.
Therefore, they are less connected to in-distribution samples,
and will be pushed further away from in-distribution samples by the proposed contrastive loss.

\vspace{-0.5ex}
\subsection{Scalable learning with an EMA model}
\vspace{-0.5ex}
In order to build a meaningful pseudo-label graph,
the unlabeled batch of data should contain a sufficient number of samples from each class.
While this requirement can be easily satisfied for datasets with a small number of classes (\eg~CIFAR-10),
it becomes difficult for large datasets with more classes (\eg~ImageNet) because a large unlabeled batch would exceed the memory capacity of 8 commodity GPUs (\eg NVIDIA~V100).
Therefore,
we improve CoMatch for SSL on large-scale datasets.

Inspired by MoCo~\cite{moco} and Mean Teacher~\cite{mean_teacher},
we introduce an EMA model $\{\bar{f},\bar{g},\bar{h}\}$ whose parameters $\bar{\theta}$ are the moving-average of the original model's parameters $\theta$:
\vspace{-0.5ex}
\begin{equation}
\bar{\theta} \leftarrow m \bar{\theta} + (1-m) \theta.
\end{equation}
\vspace{-0.5ex}
The advantage of the EMA model is that it can evolve smoothly as controlled by the momentum parameter $m$.

We also introduce a momentum queue
which stores the pseudo-labels and the strongly-augmented embeddings for the past $K$ unlabeled samples:
$\mathrm{MQ}=\{(\bar{q}_k,\bar{z}_k=\bar{g}\circ \bar{f}(\mathrm{Aug_s}'(u_k)))\}_{k=1}^K$,
where $\bar{q}_k$ and $\bar{z}_k$ are produced using the EMA model.
Different from the memory bank,
the momentum queue only contains unlabeled samples.

We modify the pseudo-label graph $W^q$ to have a size of $\mu B \times K$.
It defines the similarity between each sample in the current batch and each sample in the momentum queue (which also contains the current batch).
Different from eqn.(\ref{eqn:graph_q}), the similarity is now calculated as $\bar{q}_b \cdot \bar{q}_j$,
where $b=\{1,...,\mu B\}$ and $j=\{1,...,K\}$.

The embedding graph $W^z$ is also modified to have a size of $\mu B \times K$,
where the similarity is calculated using the model's output embedding $z_b$ and the momentum embedding $\bar{z}_j$:
$W_{bj}^z=\exp(z_b \cdot \bar{z}_j / t)$.
Since gradient only flows back through $z_b$,
we can use a large $K$ without much increase in GPU memory usage.

Besides the contrastive loss,
we also leverage the EMA model for memory-smoothed pseudo-labeling,
by forwarding the weakly-augmented samples through the EMA model instead of the original model.
A graphical illustration of the memory bank and the momentum queue is given in the appendix.