Continual Learners are Incremental Model Generalizers

Given a neural network $f(\cdot;\bm{w})$ parameterized by weights $\bm{w}$, recent works have addressed the broad machine learning problems described to $\mathcal{D}_{target}$ by optimizing learnable weights with respect to complex objective functions. Beyond statistical initialization of network weights (Glorot & Bengio, 2010; He et al., 2015), pre-training, where leveraging learned weights from scaled benchmark datasets (e.g., ImageNet (Deng et al., 2009)) as the initialization of $\bm{w}$, has been widely adopted to promote a rapid and stable convergence curve during training. Self-supervised learning (Chen et al., 2020a; He et al., 2020; Caron et al., 2021; Bardes et al., 2021; Xie et al., 2022a) has recently become prevalent for pre-training, demonstrating superior generalization performance compared to supervised counterparts by capturing task-agnostic input features. While multiple different frameworks are considered for self-supervised learning, we exemplify the encoder-decoder framework in this paragraph. Let $h$ and $g$ be an encoder and a decoder parameterized by $\bm{\theta}$ and $\bm{\phi}$, respectively, the objective function is to minimize self-supervised loss given input data $\bm{d}$ without supervision:

where $\circ$ indicates function composition. The loss function is often designed in several formulations based on similarity, identity correlation, and contrastive loss. After the pre-training phase, the encoder transfers learned features to backbone neural networks for fine-tuning, $\bm{w}\leftarrow\bm{\theta}^{*}$.

Supervised Continual Learning (SCL) (Mallya & Lazebnik, 2018; Riemer et al., 2019; Aljundi et al., 2019; Chaudhry et al., 2019, 2020; Chrysakis & Moens, 2020; Titsias et al., 2020; Shen et al., 2020; Douillard et al., 2022; Yoon et al., 2020, 2022) is about a sustainable adaptation to unlimited task sequences while maintaining proficiency on previous tasks. Let us consider an intuitive example with the image-based problem: suppose $\mathcal{T}=\{\mathcal{T}_{1},...,\mathcal{T}_{T}\}$ be a sequence of $T$ tasks, where the dataset $\mathcal{D}_{t}$ for the $t$-th task consists of $n_{t}$ training instances $\mathcal{X}_{t}\in\mathbb{R}^{n_{t}\times C\times H\times W}$ and corresponding labels $\mathcal{Y}_{t}\in\mathbb{R}^{n_{t}}$. That is, $C,H,\text{and~{}}W$ denotes a channel, height, and width of images, respectively. A continual learner $f_{\bm{w}}$, parameterized by a set of weights $\bm{w}$, aims to predict classes by minimizing the optimization problem: ${minimize}_{\bm{w}}~{}\sum^{T}_{t=1}L_{CE}\left(f\left(\mathcal{X}_{t};\bm{w}\right),\mathcal{Y}_{t}\right)$, where $L_{CE}$ is a cross-entropy loss. Yet, we assume that $f_{\bm{w}}$ can access each task in a specific timestep that loses the authorization to revisit previous tasks’ data instances when the next task arrives. That is, the model solves the following non-stationary problem at task $t$ throughout task sequential training:

Obtained models directly evaluate the performance of each task, categorizing several incremental learning setups according to the accessibility to task oracle during inference.

Unsupervised Continual Learning (UCL) is formulated in representation learning frameworks on a sequence of unlabeled tasks $\{\mathcal{X}_{t}\}_{t=1}^{T}$, often referred to as continual self-supervised learning. A learner $f_{\bm{w}}$ aims to find the best solution that learns the informative representation of multiple tasks sequentially. At each timestep $t$, the model resorts to the accessible dataset $\mathcal{X}_{t}$ without any human-annotated labels to solve the problem:

where $L$ is an arbitrary loss function for representation learning (e.g., self-supervised losses (Chen et al., 2020b; Grill et al., 2020; Zbontar et al., 2021)) and $\bm{w}_{ext}$ is an (optional) extra weights for additional structures not included in backbone weights $\bm{w}$, such as a decoder, a projection layer, and a predictor. Since a direct comparison of the quality of representation models is intractable, recent representation learning literature validates obtained representation models by probing generic transferability on multiple downstream tasks. In a similar vein, prior UCL works (Madaan et al., 2022; Fini et al., 2022) adopt supervised prediction tools like the KNN classifier and linear evaluation while keeping the learned backbone fixed. However, we argue that such evaluation paradigms cannot appropriately measure the transferability of representation on unseen tasks.

Prior continual learning literature has demonstrated that a model with a standard CL setting suffers from forgetting due to loss of local features and attention to past tasks. But, we argue that they have barely discussed the generalization of continually learned representation, which can be a great source for deploying an improvable foundation model by fine-tuning an unlimited number of tasks. We throw a question mark at this point:

”So, is the model generalization getting worse as it goes through training sequential tasks?”

Surprisingly, we found that this is not the case as the generalization (a.k.a., transfer quality) is consistently getting improved during CL as shown in Figure 1 Right. To explicate behaviors of multi-head attention using transformer backbones, similar to the experiment in Figure 1, we train supervised and unsupervised CL frameworks on the ImageNet-1K split dataset without any specific CL methods. After that, we investigate changes in regional inductive bias by measuring the average distance of attention heads (defined by Equation 8) in each layer, visualized as different points in Figure 2 (a) top and middle row. Supervised and contrastive self-supervised (SimSiam) continual learning models focus more on strong locality inductive bias at lower layers (decreased attention distance) and global attention at higher layers (increased attention distance) during continual learning as these frameworks are innately designed to cluster/classify input features in deeper layers. However, focusing on global attention to capture task-specific features is undesirable for transferring the representations to out-of-distribution tasks. Next, we visualize the entropy conditioned solely on the distribution of each attention head in Figure 2 (b) by computing $-\sum_{i}\bm{a}_{i}\log(\bm{a}_{i})$ for each attention head $\bm{a}$. Supervised and contrastive unsupervised frameworks broadly focus on most attention heads during continual learning at deeper layers. This indicates that they already substantially adapt to pre-trained tasks while losing a degree of freedom to transfer downstream tasks.

To build a new UCL framework for a better generalizable representation model across all layers, we survey Masked Image modeling (MIM) (Pathak et al., 2016; He et al., 2022) that self-trains input representation by minimizing regression loss to predict RGB pixel values in randomly zeroed patches in patchified input images. MIM enjoys locality inductive bias with diverse attention across layers, allowing better transfer ability to unseen tasks.

We formulate a representational learner $f_{\bm{w}}$, composed of a neural encoder $h_{\bm{\theta}}$ and a decoder $g_{\bm{\phi}}$. We build backbones using Vision Transformer variants (Dosovitskiy et al., 2020; Liu et al., 2021) due to their powerful generality and remarkable performance on high-resolution visual tasks. They are flexible to transfer the obtained representations to downstream tasks requiring various input image sizes in demands when existing UCL frameworks (Madaan et al., 2022; Rao et al., 2019) allow the fixed image size for representation learning and fine-tuning since the architectures are basically composed of multi-layer perceptrons and convolutional neural networks. At training $t$-th task with a training instance $\bm{x}_{t}\in\mathbb{R}^{C\times H\times W}\in\mathcal{X}_{t}$, a model segments $\bm{x}_{t}$ into smaller image patches where the width and height are $s<H,W$, and randomly zeros a fraction of image patches out with a fixed ratio $\tau$. An encoder tokenizes masked patches to the embedding space and fed into multiple self-attention blocks to capture latent representation features. A decoder reconstructs encoded features to approximate the input image. The objective is to minimize the following loss function for continual representation learning ($\|\cdot\|_{\mu}$ denotes any norm, often $\mu\in\{1,2\}$ is used). Let $K=\left\lfloor\frac{H}{s}\right\rfloor\cdot\left\lfloor\frac{W}{s}\right\rfloor$ be the number of tokens at each image, we formulate the loss function as follows:

With patch size $s$, $*_{s}$ denotes a patch-wise multiplication operation between a training instance $\bm{x}$ and a generated mask vector $\bm{m}$ drawn by the binary distribution $B$ with sparsity ratio $\rho$. $B(i,\rho)$ is a $i$ independent binary sampling with a ratio $\rho$ to pick $1$. The model updates a set of weights that predicts masked regions of input images, conditioning other available areas. We simply adopt $\ell_{1}$ regularization to minimize the distance between predicted patches and the targets, followed by earlier reconstruction-based works (Xie et al., 2022b; He et al., 2022), and after completing a sequential training, the obtained encoder $h_{\bm{\theta}}$ can be utilized for many different downstream tasks. And we find that our new framework outperforms supervised and contrastive benchmarks in model transferability during continual pre-training (Please see Figure 6). Interestingly, unsupervised continual learning with the masked autoencoder framework (SimMIM in Figure 2 bottom row) behaves very differently from the other two frameworks. Almost all layers have a diverse focus on locality, and this tendency becomes stronger as they continue to pre-train more tasks.

Our aim is to utilize the backbone of the fine-tuned model as a pre-training checkpoint for another problem in a supervised manner. Given the target task dataset $\mathcal{D}_{target}=\{\mathcal{X},\mathcal{Y}\}$ and a classifier $\delta(\cdot;{\bm{u}})$ parameterized by $\bm{u}$, we formulate the objective of continual pre-training as follows:

We suppose each fine-tuning step independently introduces its own classifier. The formulation is aligned with the continual learning problem described in Section 3.2, but this setting is about never-ending model generalization to achieve a consistently improved adaptation to the out-of-distribution task in the future. That is, obtained representation model for fine-tuning task $t$, $\widehat{\bm{w}}^{(t)}$, will be reused for future task training ($\bm{w}^{(t+1)}=\widehat{\bm{w}}^{(t)}$).

However, fine-tuning often reduces the general transferability to adapt to different tasks, demonstrating a suboptimal model generalization of supervised CL compared to our MIM-based framework (Please see Figure 6 Left). Motivated by our findings in Section 4.1 and Section 4.2, we propose a new method for continual pre-training, named GLobal Attention Discretization (GLAD). Our proposed method preserves diverse degrees of averaged distance at each attention head to preserve transferable backbone weights for future problems while capturing task-adaptive features guided by GLAD modules. As illustrated in Figure 3, a multi-head self-attention operation with adaptor weights $\bm{v}$. We transform the task-generic MSA features to input-dependent by propagating adaptor $\bm{v}$. Then, the model enables solving the current task problem while constraining the backbone weights to preserve abundant locality inductive bias. Let $\bm{a}^{l,i}$ be an averaged entropy of the attention passed over adaptor operation (dark dashed arrow) from $i$-th head at layer $l$, the objective function of our GLAD is:

where $E$ indicates the expectation, $\overline{\bm{a}}^{l}=\frac{1}{H^{l}}\sum^{H^{l}}_{i}\bm{a}^{l,i}$ at layer $l$ with the number of its attention head $H^{l}$, $\lambda$ is a scaling factor. $\epsilon$ is a small constant value. We jointly minimize the task loss with an additional regularizer that constrains the entropy variance of attention heads to sufficiently diverge at each layer as an average of their inverse standard deviation. We add to minimize a Frobenius norm for $\overline{\bm{a}}$ to promote abundant locality inductive bias for backbone attention weights. Note that our proposed method is robust to utilize any kind of multi-head self-attention modules, we demonstrate the efficacy in vanilla Vision Transformer (Dosovitskiy et al., 2020) and Swin Transformer (Liu et al., 2021). The learned backbone weights $\bm{w}$ excepting classifier and GLAD-adaptors can be reused for fine-tuning future tasks. We describe the overall continual pre-training procedure with GLAD in Algorithm 1.

We use ViT (Dosovitskiy et al., 2020) and Swin Transformer (Liu et al., 2021) as backbone architectures. We follow Siamese network structure by Madaan et al. (2022) and implement a MIM-based continual self-supervised learning framework under SimMIM (Xie et al., 2022b) and MAE (He et al., 2022) for UCL. CURL (Rao et al., 2019) is one of the pioneer works on unsupervised continual learning literature, but it is not scalable for high-resolution visual images by design. We utilize several CL methods: SI (Zenke et al., 2017), DER (Buzzega et al., 2020), and LUMP (Madaan et al., 2022). We further describe details, including hyperparameter setups in Appendix A.