Online Prototypes and Class-Wise Hypergradients for Online Continual Learning with Pre-Trained Models

In recent years, pre-trained models (PTMs) have been widely utilized in offCL (McDonnell et al., 2024; Lin et al., 2023; Zhou et al., 2024b; Smith et al., 2023; Wang et al., 2022b). However, their application in onCL remains largely unexplored, partly because most existing methods heavily depend on task boundaries, i.e., explicit knowledge of when the task changes. This is often assumed in pair with clear boundaries, meaning that all classes from the previous task are suddenly unavailable while all new classes encountered belong to the new task. With real-world streaming data, such a situation is equally unlikely.

In onCL, incoming data can be seen only once, analogous to a continuous data stream (He et al., 2020). Therefore, clear boundaries are unlikely to be available and several studies suggest working in boundary-free scenarios (Buzzega et al., 2020) where task change is unknown. However, if the change is clear, it can easily be inferred. In that sense, blurry boundary setting have been proposed (Moon et al., 2023; Koh et al., 2023; Bang et al., 2022; Michel et al., 2024) in previous work. In particular, we are interested in the Si-Blurry setting (Moon et al., 2023) where not only task change is blurry, but some classes can appear or disappear during multiple tasks, which brings the experimental setup one step closer to real-world scenarios while being more challenging. While numerous studies rely on prototypes for Continual Learning (McDonnell et al., 2024; Lin et al., 2023; Zhou et al., 2024b), such representation-based methods must generally be combined with task boundary knowledge as prototypes are updated at the end of each task. In onCL, prototypes are harder to capitalize on when training a model from scratch due to the shift of representations hindering prototype computation (Caccia et al., 2022). However, when working with PTM, such a shift is drastically reduced as representations are already of high quality, making the usage of prototypes more efficient.

Hypergradient (Baydin et al., 2018; Almeida et al., 1999) addresses the problem of finding the optimal learning rate in conventional training scenarios. In that sense, the authors proposed to derive a gradient descent algorithm to learn the LR. Notably, they demonstrate that computing the dot product between gradients from previous steps $\nabla\mathcal{L}(\theta_{t})\cdot\nabla\mathcal{L}(\theta_{t-1})$ is sufficient to complete one step of the learning rate update rule, with $t$ the index of the current step, $\theta$ the parameters, and $\mathcal{L}$ the loss function. However, such techniques have been, to the best of our knowledge, developed solely for offline scenarios at a global level. In Continual Learning, gradient re-weighting strategies have been designed for replay-based CL methods. Notably, previous work proposed to re-weight the gradient at the loss level to mitigate its accumulation during training in CL context, also called gradient imbalance (Guo et al., 2023). Recently, to compensate for the class imbalance, class-wise manually defined weights in the last Fully Connected (FC) layer have been leveraged (He, 2024). Our work on Class-Wise Hypergradients lies at the cross-road between Hypergradients and Gradient Re-Weighting.

Generally, the problem of CL is defined as training a model $f_{\theta}(\cdot)$ parameterized by $\theta$ on a sequence of $T$ tasks where each task of index $k\in\{1,\cdots,K\}$ is defined by its corresponding dataset $\mathcal{D}_{k}$, each potentially being drawn from a different distribution. In the case Class Incremental Learning (Hsu et al., 2018), we have $\mathcal{D}_{k}=(\mathcal{X}_{k},\mathcal{Y}_{k})$, the data-label pairs. In offCL, the model is trained sequentially on each task while other task data is unavailable. For onCL, while the model is similarly trained sequentially on each task, only the data of the current batch is available and can be seen only once during training. The final objective in offCL and onCL is the same: to maximize performance across all tasks. Equivalently, to minimize the average loss ${\mathcal{L}}$ across all tasks:

As discussed in Section 3, LR is critical in onCL as it not only has a direct impact on the plasticity-stability trade-off, it is near impossible to use the optimal value. While previous studies highlight the need for different LR at the layer level, the case of different learning rates at the class level is still in its infancy. Following the analysis described in Section 3.3, we make the hypothesis that a class-wise LR strategy focusing on the last FC layer is crucial for onCL. The intuition is that different classes require different learning rates so that the model can adapt its stability-plasticity tradeoff not only over time, but also over classes. In that sense, we reckon hypergradients (Baydin et al., 2018) to be adequate for improving over non-optimal initial LR values. However, hypergradients are not designed for the onCL scenario and a naive implementation leads to a severe performance drop. Indeed, hypergradient computation is identical regardless of the network layer or the classes. Inspired by the work of He et al. (He, 2024), we propose learning how to rescale the gradient coefficient of the last FC layer, class-wise, using hypergradient learning theory.

Additionally, to adapt PTM-based methods to the onCL scenario further, we leverage the stability of the output representation of PTM by computing Online Prototypes.

Let us consider a model $f_{\theta}$ parameterized by $\theta$ such that for an input $x\in\mathbb{R}^{d}$, with $d$ the dimension of the input space, we have $f_{\theta}(x)=h_{w}(x)^{T}\cdot W$, with $W\in\mathbb{R}^{l,c}$, $c$ the number of classes, $l$ the dimension of the output of $h_{w}$ and $\theta=\{w,W\}$. In this context, $h_{w}$ would typically be a PTM and $W$ the weight of the final FC layer (including the bias). Looking at the final FC layer $W$, for a learning rate $\eta$, we can write the gradient-based weight update:

with $t$ the iteration index. Here, we omit the input data for simplicity. Then, for any class index $j\in\{1,c\}$, we can write the update rule of the weights corresponding to $j$:

with $W^{j}$ the $j^{th}$ column of $W$. Building upon previous studies (He, 2024), we introduce step-dependent class-wise weighting coefficients, leading to the following update rule:

with $\alpha_{t}^{j}\in\mathbb{R}^{+\star}$, class dependent gradient weighting coefficient at index $t$. While those coefficients were traditionally introduced to compensate for class interference, and computed with hand-crafted rules, we propose to learn them through an online adaptive rule based on hypergradients theory. In particular, we want to construct a higher level update for $\{\alpha_{t}^{j}\}_{j}$ such that, in the case:

with $\beta\in\mathbb{R}^{+}$ the hypergradient learning rate. To compute the partial derivative we apply the chain rule and make use of the fact that $W^{j}_{t}=W^{j}_{t-1}-\alpha^{j}\eta\nabla\mathcal{L}(W_{t-1}^{j})$, such that:

The resulting Class-Wise Hypergradient update becomes:

With $\gamma=\beta\eta$. Naturally, this introduces an undesired extra-hyperparameter. We discuss this limitation in Section 6. For clarity, the relation presented in Eq. (8) relies on an SGD update. In practice, we favor a momentum-based update, notably Adam (Kingma & Ba, 2014). We follow the guidelines of Baydin et al. (Baydin et al., 2018) regarding its implementation and give more detail in Appendix A.

In order to reduce forgetting in the last layer, we compute Online Prototypes (OP) $\mathcal{P}=\{p_{k_{1}}^{1},p_{k_{2}}^{2},\cdots,p_{k_{c}}^{c}\}$ of each class during training. For a given class $j$, the class prototype $p_{k_{j}}^{j}$ computed over $k_{j}$ samples is updated when encountering a new sample $x_{k_{j}+1}^{j}$. For simplicity, we omit the $j$ index in $k_{j}$ going forward. Therefore, the prototype update rule is:

with $x_{k+1}^{j}$ the ${k+1}^{th}$ encountered sample of class $j$. For all classes, prototypes are initialized such that $p_{0}^{j}=0$. Prototypes are then used to recalibrate the final FC layer, analogous to replaying the average of past data representation during training. In that sense, we define the prototype-based loss term as:

with $\mathcal{C}_{old}=\{j\in\{1,c\}\ |\ p^{j}_{k_{j}}\neq 0\}$. $\mathcal{L}_{OP}$ is the cross-entropy with regard to prototypes of encountered classes. OP acts as simple and low-budget memory data.

The approach relies on two components that are orthogonal to most existing methods, as long as the model optimizes a final FC layer for classification. In that sense, we propose to integrate it into various state-of-the-art baselines relying on PTM, notably prompt-based approaches. To do so, if we consider $\mathcal{L}_{base}$ to be the loss of the baseline method, to which we attach both components OP and CWH, we have the overall loss to minimize:

A Pytorch-like (Paszke et al., 2019) pseudo code is given in Algorithm 1. Additionally, we leverage batch-wise masking to consider the logits of classes that are only present in the current batch. More details can be found in the Appendix A. Similarly, note that the bias has been omitted for simplicity.

We combined our method with four offline approaches and one online state-of-the-art approach, all using PTM. As mentioned in Section 4.4, our method must be applied on the classification head, therefore prompt-based approaches are especially suited as they all leverage a final FC layer on top of the PTM representation for classification.

Despite its apparent simplicity, OP gives the largest performance gain. This is partially due to the fact that only the FC layer and the input prompts are trained, so the output prototypes are stable over time. In Table 4 we show that unfreezing intermediate weight not only drastically decreases overall performances but similarly negates the effect of OP significantly. Additionally, in Figure 2, we show the average Euclidean distance of Online Prototypes between two training step iterations. As we can see, the learned prototypes tend to stabilize rapidly over time, even more so when the intermediate weights remain frozen. This stability of the prototype indicates that even when computed online, they can be used as a reasonable proxy for class average representation.

The main drawback of leveraging CWH is the addition of an extra hyper-parameter $\gamma$, whereas our objective is to reduce the dependency on hyper-parameters for onCL. However, we argue that selecting $\gamma$ remains particularly simple. To demonstrate this, we experiment with $\gamma\in\{0,1\times 10^{-5},1\times 10^{-3},1\}$, for an initial LR $\eta\in\{5\times 10^{-5},5\times 10^{-3}\}$. In that case, $\gamma=0$ is equivalent to disabling CWH. Such results are presented in Table 5. When the initial LR is lower, including CWH leads to an improvement in performance, especially when $\gamma$ is large. Experimentally, this can be seen as a consequence of $\{\alpha^{j}\}_{j}$ values increasing during training in all cases, as discussed in Section 6.3. Therefore, it is natural that lower initial LR values would benefit more from such a strategy. When the initial LR is large, as discussed in Section 5.3, a marginal drop in performance can be observed in some cases. However, the final performances remain stable for all values of $\gamma$, minimizing the need for hyperparameter search. Additionally, the lower the value of $\gamma$, the closer it is to the original method.

In conclusion, for higher LR we recommend lower values of $\gamma$, as their effect on the training is reduced. For lower LR, any value of $\gamma$ leads to an improvement, however, higher values are recommended for optimal accuracy gain. Default values such as $0.001$ should lead to an increase in performances in most cases or a slight decrease in worst cases. Overall, selecting $\gamma$ remains a simpler and more realistic task than doing a grid search in onCL.