Soft-TransFormers for Continual Learning

Problem Statement. Continual Learning (CL) involves training deep neural networks (DNN) on time-variant data represented as a sequence of tasks, $\mathcal{D}=\{\mathcal{D}_{1},\cdots,\mathcal{D}_{\mathcal{T}}\}$. Each $t$-th task, $\mathcal{D}_{t}=\{(\bm{x}_{i}^{t},y_{i}^{t})^{n_{t}}_{i=1}\}$ consists of $n_{t}$ tuples where $\bm{x}_{i}^{t}\in\mathcal{X}_{t}$ is an input sample and $y_{i}^{t}\in\mathcal{Y}_{t}$ is the corresponding label. When a task $\mathcal{X}_{t}$ arrives, a model $f_{\theta}$ is trained for the current task, while data from previous tasks is inaccessible. This work focuses primarily on class incremental learning (CIL), in which the task-ID is not given during inference.

Soft & Subnetworks have been explored in continual learning through two notable approaches. One approach, known as supermasks (Wortsman et al., 2020), produces outputs by $\bm{p}=f(\bm{x},\bm{w}\odot\bm{m})$, where $\odot$ denotes elementwise multiplication. In this method, the weights $\bm{w}$ remain fixed at their initialization, with bias terms set to $\bm{0}$ and other parameters initialized to $\pm c$ with equal probability where the constant $c$ is the standard deviation of the corresponding Kaiming normal distribution (He et al., 2015). Another line of work includes WSN (Kang et al., 2022) and SoftNet (Kang et al., 2023), which jointly learn the model weights $\bm{w}$ and task-adaptive subnetworks $\bm{m}$. The parameter-efficient reusable subnetworks are obtained by iteratively selecting the top-$c\%$ of the weights based on an importance score $\bm{s}$ at each layer. WSN has primarily demonstrated its effectiveness in Convolutional Neural Networks (CNNs). However, its pruning mechanism for pre-trained Transformers, such as ViT, remains unexplored. To discover the competitive sparseness in Transformers, we detail the WSN-style task-adaptive fine-tuning and the learnable soft-networks $\bm{m}$ of Transformers, presenting these adaptations for the first time with empirical observations. The soft-networks originate from learned parameters distributed with $\mu\approx 1.0$ & various variances, as stated in Figure 6.

A simple yet effective prompt-based (prompt-tuning) CIL model: Learning to Prompt (L2P) (Wang et al., 2022c) is first proposed. In this model, a prompt $p$, a tiny set of trainable tokens combined with image features, is fed into the Vision Transformer (ViT) to help the model resist forgetting. To select suitable prompts for task-specific training, L2P utilizes a prompt pool $P$ containing numerous prompt-key pairs, $\{p_{t},k_{t}\}_{t=1}^{\mathcal{T}}$, where ${p}_{t}\in\mathbb{R}^{1\times D}$ represents the $t$-th task prompt, $k_{t}$ represents the $t$-th coresponding task key, and $\mathcal{T}$ is the total number of prompt-key pairs.

Building on L2P, DualPrompt (Wang et al., 2022b) divided the prompts into expert (E-) prompts and general (G-) prompts for distinct features learning. DualPrompt also replaced prompt-tuning with prefix-tuning, which was successfully proven in NLP. DyTox (Douillard et al., 2022) designed a novel task attention block that utilized task tokens to infer task identifiers. Coda-Prompt (Smith et al., 2023b) replaced the prompt pool with a decomposed prompt, represented by a weighted sum of learnable prompt components, which optimized itself in an end-to-end fashion, providing high plasticity. LGCL (Khan et al., 2023) introduced text information into the learning of prompt pool, improving performance without any additional learnable parameters.

Recently, Qiao et al. (2024) introduced Prompt Gradient Projection (PGP), which applies an orthogonal condition on the prompt gradient to reduce forgetting via the self-attention mechanism in ViT effectively. Although various prompt-based continual learners have demonstrated state-of-the-art performance, they do not explicitly model task-specific fine-tuning and forgetting within the continual learning framework. In this work, we address task-specific fine-tuning and gradient-based task identification in CIL and TIL scenarios by leveraging prompt-tuning and learnable sparse networks.

To address the task-specific fine-tuning of the pre-trained model, such as ViT, this work proposes a new Soft-TransFormer (Soft-TF), as illustrated in Figure 1. The proposed Soft-TF consists of a conventional neural network, like a multilayer transformer with multihead attention and forward networks. Using well-trained transformer parameters, we could discover task-specific soft-networks, as depicted in Figure 3. The Soft-TF incrementally learns model weights and task-adaptive soft-masks with well-pre-trained and soft-network parameters $\bm{m}$.

Given a pre-trained parameter $\bm{\theta}$ and learnable soft-parameters $\bm{m}$, Soft-ViT is represented as $f_{\bm{\theta}\odot{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}}}$, consisting of $N$ consecutive soft-MSA layers. We extend the notation by denoting the input embedding feature of the $l_{\ast}$-th learnable soft-MSA layer as $\bm{h}^{(l_{\ast})}$, where $l_{\ast}=1,2,\dots,N$, and $l_{\ast}$ can refer to either the G-Prompt layer $l_{g}$ or the E-Prompt layer $l_{e}$. Note that while the pre-trained parameters $\bm{\theta}$ remain fixed, the soft-parameters $\bm{m}$ are updated to provide task-specific solutions.

G-prompt. $\bm{g}\in\mathbb{R}^{L_{g}\times D}$ with sequence length $L_{g}$ and embedding dimension $D$, is a shared parameter for all tasks. G-Prompt is attached to the $l_{g}$-th MSA layer to transform $\bm{h}^{(l_{g})}$ via a prompting function as follows:

where $f^{prompt}_{\bm{\theta}}$ defines the approach for attaching the prompt to the hidden embeddings.

E-prompt & Soft-networks. $\bm{e}=\{\bm{e}_{t}\}^{\mathcal{T}}_{t=1}$ is a set of task-dependent parameters, where $\bm{e}_{t}\in\mathbb{R}^{L_{e}\times D}$ has as sequence length of $L_{e}$ and the same embedding dimension $D$ as the G-prompt, and $\mathcal{T}$ is the total number of tasks. Unlike the shared G-prompt, each $\bm{e}_{t}$ is associated with a task-specific key $\bm{k}_{t}\in\mathbb{R}^{D}$, which is also a learnable parameter aimed at capturing representative features of a task. For an input example from the $t$-th task, to attach E-prompt to the $l_{e}$-th soft-MSA layer, we apply the prompting function in a similar way:

G- and E-prompts, along with learnable soft-networks, encode specific types of instructions during training with the backbone and work together to guide the model’s predictions during inference. We have demonstrated the method for attaching prompts and learnable soft-networks to a single soft-MSA layer. Similarly to the approach taken in DualPrompt (Wang et al., 2022b), we also investigate layers of E-prompts with learnable soft-networks $\bm{m}$, while utilizing the layers designated for G-prompts.

Layers of G- and E-Prompts. We use the multilayered extension of both types of prompts: $\bm{g}=\{\bm{g}^{(l_{g})}\}_{l_{g}=start_{g}}^{end_{g}}$, where $\bm{g}^{(l_{g})}\in\mathbb{R}^{L_{g}\times D}$ represents the G-prompt attached to the $l_{g}$-th MSA layer. Similarly, we define $\bm{e}_{t}=\{\bm{e}_{t}^{(l_{e})}\}_{l_{e}=start_{e}}^{end_{e}}$ for the $l_{e}$-th conventional MSA layer. In this configuration, the G-prompt $\bm{g}^{(l_{g})}$ is attached from the $start_{g}$-th to the $end_{g}$-th conventional MSA layers, and the E-prompt $\bm{e}_{t}^{(l_{e})}$ is attached to the $\left[start_{e},end_{e}\right]$-th soft-MSA layers, ensuring that there is no overlap between them. In our experiments, we follow the $l_{g}\notin\left[start_{e},end_{e}\right]$ ($l_{g}=[1,2]$) settings used in DualPrompt and empirically search for the optimal $\left[start_{e},end_{e}\right]$ layers for the learnable subnetworks through ablation studies.

Learnable Soft-networks. The prompting function $f^{prompt}_{\bm{\theta\odot{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}}}}$ determines how prompts ($\bm{p}$) are combined with fine-tuned soft ($\bm{\theta}\odot{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}}$) embedding features. From another perspective, $f^{prompt}_{\bm{\theta\odot{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}}}}$ directly influences the interaction between high-level instructions in the prompts and low-level representations. Therefore, we believe that a well-designed prompting function, along with task-specific parameters, is crucial for optimizing overall continual learning performance.

Specifically, applying a prompting and fine-tuning function $f^{prompt}_{\bm{\theta\odot{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}}}}$ can be seen as modifying the inputs to the soft-MSA layers. Let the input to the soft-MSA layer be $\bm{h}\in\mathbb{R}^{L\times D}$, and denote the input query, key, and values for the soft-MSA layer as $\bm{h}_{Q}$, $\bm{h}_{K}$, and $\bm{h}_{V}$, respectively. A soft-MSA layer is defined by the following equation:

where $\bm{w}_{i}^{O},\bm{w}_{i}^{Q},\bm{w}_{i}^{K}$, and $\bm{w}_{i}^{V}$ are fixed projection matrices while ${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}^{O}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}^{Q}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}^{K}}$, and ${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\bm{m}^{V}}$ are learnable parameters. $s$ is the number of heads. In ViT, $\bm{h}_{Q}=\bm{h}_{K}=\bm{h}_{V}$. Here, we define a unified prompt parameter with a sequence length of $L_{p}$, such as $\bm{p}\in\mathbb{R}^{L_{p}\times D}$ for a single-layered G- or E-prompt.

In this framework, we concatenate the prompts $\bm{p}_{t}$ and the embedding sequence $\bm{x}_{t}$, i.e., inputs from $t$-th task, along the embedding dimension: $\bm{z}_{t}=[\bm{p}_{t};\bm{x}_{t}]$. With the weights of $\bm{w}^{Q}\odot\bm{m}^{Q},\bm{w}^{K}\odot\bm{m}^{K},\bm{w}^{V}\odot\bm{m}^{V}$, the soft-transformer takes query ($\bm{q}_{t}=(\bm{w}^{Q}\odot\bm{m}^{Q})\bm{z}_{t}$) and key ($\bm{k}_{t}=(\bm{w}^{K}\odot\bm{m}^{K})\bm{z}_{t}$) as input of the soft-MSA layer. The soft-attention matrix is then given by:

where we focus on $\bm{q}_{t}\bm{k}_{t}^{T}=(\bm{w}^{Q}\odot\bm{m}^{Q})\bm{z}_{t}$ and $\bm{z}_{t}^{T}(\bm{w}^{K}\odot\bm{m}^{K})^{T}$. First, the trainable prompt parameters can be denoted as:

Second, the trainable soft-attention layer’s parameters with $\bm{m}^{Q}$ and $\bm{m}^{K}$ are as follows:

where $\bm{w}^{Q}$ and $\bm{w}^{K}$ are frozen and unchanged during training and test.

The overall process of the Soft-TransFormers (Soft-TF) during training and testing is described as Algorithm 1 and Algorithm 2. We denote the architecture with attached prompts as $f_{\bm{g},\bm{e}_{t},{\bm{m}_{t}}}$. The input $\bm{x}$ of the $t$-th task is transformed using $f_{\bm{g},\bm{e}_{t},{\bm{m}_{t}}}$ and then passed to the classification head $f\phi$, parameterized by $\phi$, for prediction. Finally, we train the two prompts, the task keys, the soft-attention parameters, and the newly-initialized classification head in an end-to-end manner:

Here, $\mathcal{L}$ represents the cross-entropy loss, and $\mathcal{L}_{\text{match}}(\bm{x},\bm{k}_{t})=\gamma(\bm{q}(\bm{x}),\bm{k}_{t})$ denotes the matching loss, where $\bm{q}(\bm{x})=f(\bm{x})[0]$ corresponds to the feature vector associated with the [class] token (Dosovitskiy et al., 2020; Wang et al., 2022b), and $\gamma$ is the cosine similarity. The scalar $\lambda$ serves as a balancing factor between the losses; here, we follow the same DualPrompt setting as a baseline.

Analysis of Soft-TF for Convex-Lipschitz Functions. To analyze the convergence rate of the Soft-Transformer, we focus on the case of convex-Lipschitz functions. Let $\bm{w}^{\ast}=\{\bm{g}^{\ast},\bm{e}_{t}^{\ast},\bm{m}_{t}^{\ast}\}$ be any vector, and let $B$ be an upper bound on $||\bm{w}^{\ast}||$ when $\bm{w}^{(1)}=\bm{0}$, or $\bm{w}^{(1)}$ is an initial state. It is helpful to consider $\bm{w}^{\ast}$ as the minimizer of $f(\bm{w})$, although the following analysis applies to any $\bm{w}^{\ast}$.

We derive an upper bound on the sub-optimality of our solution relative to $\bm{w}^{\ast}$, specifically $f(\bar{\bm{w}})-f(\bm{w}^{\ast})$, where $\bar{\bm{w}}=\frac{1}{T}\sum_{t=1}^{T}\bm{w}^{(t)}$. By the definition of $\bar{\bm{w}}$ and applying Jensen’s inequality, we obtain the following (see Section A.1 stated in detail):

For every $t$, because of the convexity of $f$, we have that

Combining the preceding, we obtain

To bound the right-hand side of the above formula, we rely on the following lemma:

Datasets. We evaluate our method mainly on 1) 10/20-Split-CIFAR100 (Krizhevsky et al., 2009), constructed by splitting the 100 classes into 10 tasks/20 tasks. 2) 10-Split-TinyImageNet (Abai & Rajmalwar, 2019), constructed by splitting the 200 classes into 10 tasks. 3) 10-Split-ImageNet-R (Hendrycks et al., 2021), constructed by splitting the 200 classes into 10 tasks. To show our effectiveness, we additionally compare our method with the baselines on 5-Split-CUB200 and 10-Split-TinyImageNet. The detailed experimental settings are depicted in the Supplementary.

Implementation. For fair comparisons, we set L2P (Wang et al., 2022c), DualPrompt (Wang et al., 2022b), CLIP (Radford et al., 2021), and PGP (Qiao et al., 2024) as our baselines. We follow experimental settings Qiao et al. (2024) entirely.

Baselines. To validate the powerfulness of our method, we compare our results with various CIL baselines including ICaRL (Rebuffi et al., 2017), BiC (Wu et al., 2019), DER++ (Buzzega et al., 2020), LWF (Li & Hoiem, 2017), EWC (Kirkpatrick et al., 2017b), DER+MCG (Cai et al., 2023), and DualPrompt-PGP (Qiao et al., 2024). In addition, we investigate subnetwork solutions such as WSN (Kang et al., 2022) (obtained by selecting top-$c\%$ of weight scores while fixing pre-trained parameters) and SoftNet (Kang et al., 2023) (acquired by selecting top-$c\%$ of weight scores as major tickets while setting $100.0-\text{top-}c\%$ as minor tickets) in Vision Transformers (ViT) using prompt tuning methods. We adopt average accuracy (ACC) and forgetting (FOR) as our validation metrics (Wang et al., 2022b; Qiao et al., 2024).

Task Inference. At the inference time, we infer task identity for arbitrary pieces of task samples $\bm{x}$ for finding the proper task nuances and demonstrating full fine-tuning results. We summarize the following two methods:

• •

  Prompt ID: For a test example $\bm{x}$, we simply choose the best matched task index via $\text{argmin}_{t}\gamma(q(\bm{x}),\bm{k}_{t})$.

• •

  Gradient ID: To infer the task identity, we follow SupSup’s one-shot task inference (Wortsman et al., 2020). In short, we assign each learned subnetwork $\bm{m}_{t}$ a weight $\alpha_{t}$ such that $\sum_{t}\alpha_{t}=1$ and $\alpha_{t}=1/\mathcal{T}>0$ when evaluating all seen tasks. Given an example data point of batch $\bm{x}\in\bm{b}$ to classify, we can compute the loss as $\mathcal{L}=\mathcal{H}(f^{prompt}_{\bm{\theta}\odot(\sum_{t}\alpha_{t}\bm{m}_{t})}(\bm{x}))$ where$f^{prompt}_{\bm{\theta}}(\bm{x})$ is the pre-trained model which outputs logits and $\mathcal{H}$ is the entropy function. From here our inferred task is simply $\hat{t}=\mathrm{argmin}_{t}\frac{\partial\mathcal{H}}{\partial\alpha_{t}}$.

Performances of Soft-TF on CIL. We compare our Soft-TransFormers (Soft-TF) with state-of-the-art CIL baselines, as shown in Table 1. Our Soft-TF significantly outperforms all baselines and upper-bounds of Soft-TF, including L2P and DualPrompt, in both accuracy and forgetting measurements. The performance gain of Soft-TF is especially notable in DualPrompt-based learning compared to L2P, suggesting the importance of global prompt-tuning and multi-head attention prompt-tuning in DualPrompt. Additionally, task-identity inference using Gradient-ID is crucial for achieving full fine-tuning results in CIL. To demonstrate the effectiveness of Soft-TF, Figure 2(a) presents a radar chart comparing DualPrompt-PGP and Soft-TransFormers across four benchmark datasets.

Well-initialized LTH (WLTH) on CIL. To demonstrate the efficacy of our proposed method on Well-initialized Lottery Ticket Hypothesis (WLTH) backbones, we evaluate our Soft-TransFormers (Soft-TF) by extending two distinct pre-trained models, ViT-DINO and ViT-SAM (Caron et al., 2021; Chen et al., 2021). As shown in Table 2, we tested our method on the 10-Split-CIFAR100 and 5-Split-CUB200 datasets using three pre-trained ViTs: ImageNet-21K, DINO, and SAM, further validating the effectiveness of our approach on non-ImageNet datasets (Krizhevsky et al., 2009; Wah et al., 2011). Surprisingly, when initialized with ImageNet-21K and SAM, DualPrompt-Soft-TF with ImageNet-21K achieved the same performance levels. Moreover, DualPrompt-Soft-TF outperformed all baselines, i.e., DualPrompt-PGP, on both benchmark datasets, indicating that well-initialized weights provide better generalization in continual learning scenarios.

CLIP on CIL and TIL. We conduct our experiments on the 10-Split-CIFAR100 dataset under both Class Incremental Learning (CIL) and Task Incremental Learning (TIL) settings, as shown in Table 3. The results demonstrate that CLIP-Prompt-Soft-TF-L[1-12] significantly improves performance in both settings, indicating that our Soft-TF with Gradient ID is also effective in vision-language models, thereby broadening its applicability.

Layer-wise Inspections. We analyze the layer-wise performance of Soft-Transformer with respect to L2P and DualPrompt on the 10-Split-CIFAR100 dataset to identify the optimal configurations, as shown in Figure 4. Our observations reveal that the global prompt in DualPrompt influences Soft-Transformer’s performance differently in L2P and DualPrompt settings. In L2P-PGP, the best performance was achieved with Soft-TransFormers applied to the lower layers ((a) L2P-Soft-TF-L[1,2]-PGP), whereas in DualPrompt, the higher layers ((b) DualPrompt-Soft-TF-L[10,11,12]) yielded the best results. Notably, DualPrompt-Soft-TF-L[10,11,12] without PGP demonstrated impressive performance, achieving almost zero forgetting (0.21). These findings suggest that our approach could significantly enhance the effectiveness of large-scale Transformer models in continual learning scenarios.