Improving Representational Continuity with Supervised Continued Pretraining

We consider the setting of continual representation learning [15, 14]. As a motivating example, suppose we are developing a foundation model [2] for satellite remote sensing [9]. A resource rich organization pretrains a model $M$ on lots of satellite data from North America. Over time, the organization needs to upgrade the capabilities of its model by incorporating more geographical regions—retraining the model is expensive so it takes $M$ and continues the pretraining process on data from Africa to get a model $M^{\prime}$.

The organization releases the pretrained model, and resource constrained academic labs across the world use the updated model $M^{\prime}$ for important applications—a common way to use the model $M^{\prime}$ is to fine-tune it on a small amount of task-specific labeled data. We want the updated model $M^{\prime}$ to be useful for both new data (from Africa) and old data (from North America).

More formally, we sequentially pretrain a model on a sequence of tasks $\mathcal{T}_{1},\ldots,\mathcal{T}_{T}$ to get a model $f_{\theta_{T}}$—that is, we pretrain on data from $\mathcal{T}_{1}$, then data from $\mathcal{T}_{2}$, etc. Naively pretraining in this way (which we call SGD) leads to forgetting—the model performs poorly on data from older tasks such as $\mathcal{T}_{1}$. Prior works in continual learning propose a variety of methods to reduce forgetting: popular methods include Synaptic intelligence (SI), Dark Experience Replay (DER), and Unsupervised Continual Learning (UCL). To examine if $f_{\theta_{T}}$ has retained knowledge about every task $\mathcal{T}_{i}$, we adapt $f_{\theta_{T}}$ using a small amount of labeled data from task $\mathcal{T}_{i}$, and evaluate its accuracy on task $\mathcal{T}_{i}$.

Practitioners will typically use the adaptation method that gets the highest accuracy for their use case—in this setting, we find that continual learning methods such as UCL, DER, and SI, can actually do worse than naive training (SGD). For each task, we consider adapting the pretrained model with three adaptation methods (training a kNN or linear classifier on frozen representations, or full fine-tuning of the model) on a small amount of task specific data. Prior work uses the kNN evaluation protocol—partly because this is cheap to evaluate. Under this scheme, we indeed see that continual learning methods improve over naive training. However, when considering the best adaptation method, we find that strong continual learning baselines (SI, DER, and UCL) can perform even worse than naive training—see Table 2. For example, under a kNN evaluation protocol, SI gets 86% which is 8% better than naive SGD (78%), but under the best adaptation method SI gets 91% accuracy which is 2% lower than naive SGD (93%).

However, we find that a method from the transfer learning community (LP-FT) [12] outperforms naive SGD. For each new task $\mathcal{T}_{i}$, naive SGD updates the entire model via gradient descent—this can distort representations learned from previous tasks. Instead,[12] proposes first training the linear ‘head’ layer on the new task $\mathcal{T}_{i}$ (to find the best way to use existing representations for the new task) and then fine-tune the entire model to incorporate information from $\mathcal{T}_{i}$ into its representations. LP-FT consistently improves over naive SGD under all adaptation methods, and gets the highest overall accuracy. Even under the exact kNN evaluation protocol used in [14], LP-FT matches or outperforms continual learning methods on three popular datasets: Sequential-CIFAR10 [11], Sequential-CIFAR100 [11], and Sequential-TinyImageNet [6].

Finally, we find that LP-FT performs well in two other domains as well. 1. We consider continual representation learning in a real world satellite remote sensing task (Functional Map of the World [5])—LP-FT (56%) outperforms naive SGD (53%). 2. LP-FT gets state-of-the-art accuracies on an NLP continual learning setting [10] where they update BERT for a sequence of sentiment analysis tasks. [10] propose a new continual learning method that outperforms 12 strong baselines. We show that a variant of LP-FT gives a further 2% accuracy boost over their strongest method.

Over the years, many continual learning approaches have been proposed in three main categories. Regularization-based methods [21, 13, 4] minimize the drift of model representations during sequential learning to prevent the forgetting of learned knowledge. Architectural-methods [16, 19, 20] expand the network structure to learn the task-independent and task-shared knowledge across the sequence of tasks. Replay-based methods [17, 1, 3] select and revisit a representative subset of the past-task examples during the training of future tasks to alleviate forgetting. While most of these methods were restricted to supervised settings, recent works [14, 7] have extended them to continual learning with the unlabelled data stream.

In continual learning, there is a continuum of $T$ tasks $\mathcal{T}_{1},\ldots,\mathcal{T}_{T}$. For each task $\mathcal{T}_{t}$, we have a training dataset $\mathcal{D}_{tr}^{(t)}$, few-shot dataset $\mathcal{D}_{ft}^{(t)}$, evaluation dataset $\mathcal{D}_{te}^{(t)}$ sampled from a distribution $P_{t}$ over $\mathcal{X}_{t}\times\mathcal{Y}_{t}$. We parameterize the models by a base feature extractor $\theta\in\mathcal{B}$ and task-specific head $\phi_{i}\in\mathcal{V}$ to predict $h_{\phi_{t}}(f_{\theta}(x))$ on task $t$, where $f_{\theta}(x)\in R^{k}$ maps inputs into a lower dimensional feature space, and $h_{\phi}$ represents a linear head.

We include formal definitions for the three different instantiations of the probing classifier used for our evaluation.

1. 1.

   Linear probe

   We train a linear classifier on frozen representations produced by the model $\theta_{T}$:

   $$\begin{split}\operatorname*{\text{probe}}(\theta_{T},D_{ft}^{(i)}):&=\operatorname*{arg\,min}_{\xi}L_{i}(\xi),\theta_{T}\mbox{, where }\\ L_{i}(\xi)&=\sum_{x,y\in D_{ft}^{(i)}}l(h_{\xi}(f_{\theta_{T}}(x)),y)\end{split}$$

   (4.1)

2. 2.

   KNN probe $\xi$ follows [18] by building a nearest-neighbor classifier based on ${\theta_{T}(x):x\in D_{ft}^{(i)}}.$

3. 3.

   LPFT probe We define a loss function:

   $$L_{i}(\xi,\theta)=\sum_{x,y\in D_{ft}^{(i)}}l(h_{\xi}(f_{\theta}(x)),y).$$

   For each task $i$, we first train the head $\xi$ and then fine-tune the entire model $\xi,\theta_{T}$.

   	$\displaystyle\widehat{\xi_{i}}_{lp}$	$\displaystyle=\operatorname*{arg\,min}_{\xi}L_{i}(\xi,\theta_{T})$		(4.2)
   	$\displaystyle\operatorname*{\text{probe}}(\theta_{T},D_{ft}^{(i)}):$	$\displaystyle=\operatorname*{arg\,min}_{\xi,\theta}L_{i}(\xi,\theta)$		(4.3)

   initialized at $\xi=\widehat{\xi_{i}}_{lp},\theta=\theta_{T}.$

We show quantitative results on popular benchmarks Sequential-CIFAR-10 $(T=5)$, CIFAR-100 $(T=20$), Tiny-ImageNet $(T=10)$, as well as a benchmark of 19 Aspect Sentiment Classification Datasets [10] $(T=19)$ and an initial result on real-world satellite data Functional Map of the World $(T=6)$. We abbreviate these as C10, C100, Tiny, ASC, and FMOW. We show that LPFT significantly improves over fine-tuning on C10, C100, and Tiny, and is almost as good as the best unsupervised continual learning algorithms and SOTA supervised continual learning algorithms. In order to better evaluate representations, we introduce an evaluation protocol where we are given a small percentage of data to tune on at test time and show that the benefit from popular regularization and replay continual learning algorithms is largely negated when we follow this evaluation protocol.

This works aims to learn better representations via continued pretraining and find better protocols to evaluate them. We introduce a probe evaluation framework that changes the ranking of continual learning methods. We introduce a simple yet effective technique that boosts performance across all datasets, can be applied in different domains and settings, couples with existing continual learning methods, and does not incur additional costs.