Does RoBERTa Perform Better than BERT in Continual Learning: An Attention Sink Perspective

We characterize the presence of attention sinks (Xiao et al., 2024) in LMs using data from NLI datasets SST and MNLI (Socher et al., 2013; Williams et al., 2018a; Wang et al., 2019).

Attention on sink tokens. Fig. 1 illustrates the presence of attention sinks using attention maps, which show high attention scores are allocated to specific input tokens (i.e., sink tokens). The figure also shows sink tokens might receive similar (high) attention scores from all tokens. To empirically quantify these observations, we devise the measurements below.

Let ${{\bf{A}}}\in\mathbb{R}^{n\times n}$ denote an attention matrix over $n$ tokens for a single attention head. An element $a_{ij}\in\bf{A}$ denotes the attention on the $j$-th key token for the $i$-th query token. For the $i$-th (query) token, we have the following measurements :

We average $d_{i}$ and $\Delta_{i}$ across all attention-heads in each layer.t $d_{i}$ is the averaged attention score allocated to the $i$-th token, and $\Delta_{i}$ is the per-degree attention score deviation to the $i$-th token’s average outer degree. We study sink tokens with the largest average outer degrees, and calculate their attention deviations as sink attention deviations.

Fig.2 shows the layer-wise average outer degrees and sink attention deviations for BERT and RoBERTa models. In Fig.2 (a), tokens with top-3 largest outer degrees obtain 60% attention scores from input tokens, while the top-1 tokens obtain over 20% attention. This shows that a small number of tokens obtain major attention in self-attention layers. In Fig.2 (b), we observe that the sink attention deviations are mostly small, except for the first two layers. This shows that all tokens pay similar attention to the sink tokens in the top layers.

Sink tokens are usually common tokens. We also find that sink tokens in pre-trained LMs turn out to be common tokens, ones that appear in many language tasks. These common tokens include special tokens such as (‘[SEP]’), punctuation (‘.’), or the initial tokens in inputs, which are not semantically significant. The right side of Fig. 3 shows the ratio of sink tokens with the top-1 largest outer degrees that are also common tokens at each layer. Almost all sink tokens are common tokens in the first several layers of the pre-trained model. Even after fine-tuning, the models still have high attention on them in the middle layers.

We show the impact of attention sinks on single tasks by connecting them to an over-smoothing phenomenon, where the attention deviations on sink tokens are a key factor.

Over-Smoothing in transformers. Previous works have identified an over-smoothing phenomenon in transformer-based models: token representations become identical after several self-attention layers. Over-smoothing is closely related to models’ task learning ability and their overfitting on a task (Shi et al., 2022). There are several factors that can lead to over-smoothing, and we focus on the effect of self-attention matrices here.

For a self-attention layer with an attention matrix ${\bf{A}}\in\mathbb{R}^{n\times n}$, let ${\bf{H}}\in\mathbb{R}^{n\times d}$ be its input token representations and ${\bf{{A}H}}$ its output. The over-smoothing problem is described as:

The distance $d_{\mathcal{M}}({\bf{H}}):=||({\bf{I}}-{\bf{ee}}^{T}){\bf{H}}||_{F}$ measures the closeness between each token representation in ${\bf{H}}$ and the averaged token representation ${\bf{ee}}^{T}{\bf{H}}$, where ${\bf{e}}=n^{-\frac{1}{2}}[1,1,...,1]\in\mathbb{R}^{n}$. $\lambda_{\max}$ is the largest eigenvalue of ${\bf{A}}^{T}({\bf{I}}-{\bf{ee}}^{T}){\bf{A}}$. When $\lambda_{\max}$ is small, representations after the attention layer will be closer to the averaged representations, which causes over-smoothing.

Connection to attention sinks. Over-smoothing has been identified in many transformer-based models. We analyze the eigenvalue $\lambda_{\max}$ to see the property of the attention matrix $\bf{A}$ under the over-smoothing circumstances (i.e., $\lambda_{\max}$ is small). With each attention score $a_{ij}$ and average outer degree $d_{i}$ defined in Section 3.1,it $\lambda_{\max}$ is lower bounded as:

The details are in Appendix A. When $\lambda_{\max}$ is small, the RHS of Eq. (2) has to be small. In particular, the $i$-th token which has the largest outer degree must have its deviation $\sum\nolimits_{k=1}^{n}(a_{ki}-d_{i})^{2}$ to be small. When $d_{i}$ is large (as shown in Fig. 2), to make the deviation small each attention $a_{ki}$ needs to be close to $d_{i}$. Therefore, all tokens have similar (high) attention on the token with the largest outer degree, which is an attention sink phenomenon.

We empirically show the connection between attention deviations $\Delta$ and over-smoothing in Fig. 3. The over-smoothing degree is reflected by the average cosine similarity between token representations (Shi et al., 2022). Going from lower to higher layers, we observe that the attention deviation decreases while the representation similarity increases. This validates the connection between attention sinks and the over-smoothing phenomenon.

Impact of over-smoothing with attention sinks. When over-smoothing occurs with attention sinks above, the sink token representations may dominate the averaged token representation, and make other token representations (including [CLS]) close to them. To learn a task, models may push sink token representations close to the task data representation (Fig. 4(a)). Since sink tokens may be irrelevant to tasks (Fig. 3 right), this may distort pre-trained features and make models less generalizable to OOD data (Kumar et al., 2022).

Comparing BERT and RoBERTa, we observe that pre-trained RoBERTa suffers more from over-smoothing (i.e., high representation similarity), corresponding with low attention deviations on sink tokens at the second and last several layers (Fig. 2(b)). Therefore, we hypothesize that RoBERTa may be more vulnerable to feature distortion in task learning, which is reflected by its distorted attention patterns (Fig. 1 and Fig. 3 right) after fine-tuning. This may also influence RoBERTa’s CL capacity.

We study the following CL problem: A model continually learns a sequence of tasks, with no previous tasks’ data accessible when learning new tasks, and no storage of previous tasks’ model parameters. The model has different predictors for different tasks, while the encoder is shared. Each task $i$ consists of data $\mathcal{D}_{i}=({\bf{X}}_{i},y_{i})$ where ${\bf{X}}_{i}$ is the input feature for task $i$, and $y_{i}\in\mathbb{R}$ is its target output.

When learning task $j$ after task $i$, one way to quantify the cross-task interference on the shared parameter $\theta$ is through the dot product between its (vectorized) gradients on the two tasks’ losses (Riemer et al., 2019):

where $\cdot$ is the dot product operator on the flattened gradients. The interpretation is that interference that leads to forgetting happens when $I(\theta;i,j)<0$, while the positive knowledge transfer may happen if $I(\theta;i,j)>0$. For tasks that do not have knowledge transfer, the interference is expected to be 0.

We use a case study to show that attention sinks can propagate unexpected interference between irrelevant tasks. The study is based on the attention sink phenomenon characterized in Section 3, which showed: (1). models allocate high attention to sink tokens with small deviations; (2). sink tokens are usually common tokens shared across different tasks.

Data. We consider two irrelevant NLP tasks in a CL setting. For task $i$, we have data instance (${\bf{X}}_{i},y_{i}$) where ${\bf{X}}_{i}$ consists of embeddings of input tokens. We make the following assumptions about tasks and data:

1. 1.

   There is no knowledge transfer between the tasks and there should be no interference (positive or destructive) when learning one task after the other.

2. 2.

   Assume there are $k$ common tokens (e.g., special tokens) in two tasks’ data instances. For all other tokens in a task, we assume they are irrelevant to non-common tokens in the other task, with corresponding embeddings being orthogonal.

Model. We use a model consisting of two single-head self-attention layers. For each task $i$, the input ${\bf{X}}_{i}\in\mathbb{R}^{n_{i}\times d}$ consists of $d$-dimensional embeddings for $n_{i}$ tokens. Considering a regression problem (generalizable to classification), the prediction $\hat{y}_{i}$ is calculated as:

where ${\bf{A}}_{i}\in\mathbb{R}^{n_{i}\times n_{i}}$ is the attention matrix in the first attention layer, ${\bf{b}}_{i}\in\mathbb{R}^{n_{i}}$ is the attention vector in the second attention layer that integrates all hidden representations for the target task prediction. Both ${\bf{A}}_{i}$ and ${\bf{b}}_{i}$ are obtained under the self-attention mechanism (Vaswani et al., 2017). ${\bf{W}}\in\mathbb{R}^{d\times d}$ is a transformation matrix. ${\bf{v}}_{i}\in\mathbb{R}^{d\times 1}$ is the predictor that maps the representation to an output value $\hat{y}_{i}$ for task $i$. The loss function is: $L(\hat{y}_{i},y_{i})=\mathbb{E}[\frac{1}{2}(\hat{y}_{i}-{y}_{i})^{2}].$

For simplicity, we sort ${\bf{X}}_{i}$ and ${\bf{A}}_{i}$ to make the $k$ common tokens have indices $\{1,...,k\}$ and others have indices $\{(k+1),...,n_{i}\}$. We assume the common tokens are sink tokens in ${\bf{A}}_{i}$.

Since attention sinks on common tokens can propagate unnecessary interference, should we exclude sink tokens that are common tokens (‘common sink tokens’) when calculating attention in CL? To answer this question, we first train models with and without attention on common sink tokens, and then compare their in-task and transfer learning capacities. Results in Fig. 4(b) show that when discarding the attention on common sink tokens in training, models have a significant performance drop in the in-task evaluation.

In addition, common sink tokens may benefit tasks with positive knowledge transfer. In Fig. 4(b), we use models trained on MNLI data for zero-shot evaluation on SNLI data, which is for the same sentence entailment task but with a different data distribution. Results show that the models’ transfer ability significantly drops after discarding the attention on common sink tokens. We hypothesize that this could be because sink tokens that are common across tasks can easily transfer knowledge on them.

The analysis above motivates us to balance the transfer and interference caused by attention sinks in task learning. Specifically, when allowing models to preserve relatively high attention on sink tokens, we in turn encourage them to pay attention to non-sink tokens that have relatively low attention on sink tokens. This may increase sink attention deviations, which help reduce interference and oversmoothing. We develop a pre-scaling model to achieve this goal, detailed in Section 5.

Datasets. We evaluate four sequences of CL tasks: (1) Yahoo Split: a split of Yahoo dataset for news question-answer categorization (Zhang et al., 2015) with 5 disjoint tasks containing 2 classes each; (3) DB: a split of DBPedia data for Wikipedia article classification (Zhang et al., 2015) with 7 disjoint tasks containing 2 classes each; (4) News Series: a sequence of tasks on news-related data, including AG_news (news classification, 4 classes), MRPC (paraphrase detection, 2 classes) (Dolan & Brockett, 2005), RTE (text entailment, 2 classes) (Williams et al., 2018b) and SST (sentiment analysis, 2 classes) (Socher et al., 2013). For the above sequences, we randomly sample 1245 samples per class, which is the least number of class samples in our datasets.

Baselines. We consider two categories of baselines:

One category performs vanilla sequential learning for CL but has different training strategies on each single task, including (1) FT: a model where all parameters are sequentially updated; (2) PT+FT (Kumar et al., 2022): a model first trains the classifier in the probing stage and then fine-tunes the whole model; (3) Prescale (ours): a model first trains the classifier and a scaling layer in the probing stage and then fine-tunes the whole model (with scaling).

Another category is designed with specific CL techniques like experience replay, including (1) ER: a FT model storing all seen examples and performs sparse (1%) experience replay; (2) A-GEM (Chaudhry et al., 2019a): a FT model constraining on gradients to prevent degrading performance of previous tasks; (3) MBPA++ (d’Autume et al., 2019): a FT model that stores and retrieves samples to locally adapt the model at inference time (Sprechmann et al., 2018). (4). IDBR (Huang et al., 2021b): a FT model with information-disentanglement-based regularization and replay. We also compare to IDBR without replay, denoted as IDBR(-R); (5) CTR (Ke et al., 2021): an adapter-based task-incremental model with capsules and task transfer routing; (6) L2P (Wang et al., 2022): a prompt-based model that learns to dynamically prompt for different data and tasks. We also compare models under multi-task learning (MTL) and separate learning for each task (Separate) as non-CL baselines to show the performance gap from CL models to them. Detailed settings are in Appendix C.

We compare BERT-base and RoBERTa-base models in sequential learning, and use BERT-base for CL-specific models. For BERT models, we use learning rate 2e-5 to train 3 epochs for each task; and for RoBERTa models, we use learning rate 1e-5 to train 5 epochs per task. For probing and pre-scaling, we use the learning rate 5e-4 to train the classifier.

Metrics. We train models in a task-incremental setting where task identifiers are given (task-aware). We evaluate models’ CL performance by evaluating their average accuracy (ACC) and forgetting (FGT) on the sequence (Chaudhry et al., 2019b), with or without task identifiers (task-agnostic). For analysis, we evaluate models’ over-smoothing and attention deviations on sink tokens using metrics in Eq. (1).

RoBERTa does not always outperform BERT in CL. Table 1 compares BERT and RoBERTa’s CL performance under sequential learning. With fine-tuning, RoBERTa does not always outperform BERT in CL despite its high pre-trained capacity. Specifically, RoBERTa has a lower accuracy than BERT on Yahoo Split and a higher forgetting on News Series. This may relate to its higher over-smoothing and lower sink attention deviations than BERT, which make it more vulnerable to feature distortion and easier to propagate interference.

Prescaling improves RoBERTa’s CL capacity. With PT+FT, RoBERTa consistently outperforms BERT on CL tasks. We believe that is because PT+FT first learns good class vectors $\bf{v}$, which reduces feature distortion in each single task and then benefits CL. After applying our prescaling method, BERT and RoBERTa achieve further improvements in CL tasks.

Prescaling increases attention deviations. In Fig. 6, we compare models’ representational similarity and attention deviations after CL. After prescaling, RoBERTa’s representation similarity decreases while the attention deviations increase. This suggests that our pre-scaling can encourage diverse attention, which reduces over-smoothing and benefits CL.

Prescaling model outperforms CL models with replay. In Fig. 5, we compare our pre-scaling model to CL-specific models to evaluate its overall CL capacity. Without experience replay or progressively storing model parameters, Prescale achieves overall best accuracies on CL tasks. Even for the News Series sequence which has knowledge transfer between tasks, Prescale outperforms replay-based models that are effective in this scenario. This validates the effectiveness of our prescaling mechanism in CL.

Scaling strategies. In Table 6.2, we compare our prescaling strategy (Full) to two other scaling strategies: one uniformly distributing attention to all tokens (Uniform); another learning to scale attention only on common sink tokens including special tokens, the punctuation ‘.’ and the second token in the sentence (Sink). Results show that only scaling attention on common sink tokens does not yield improvements to PT+FT, and uniform scaling does not give as much improvement as full scaling. These suggest that the effectiveness of our prescaling strategy does not come only from having distributed attention on tokens.

Scaling visualization. Fig. 7 shows a heapmap of the scaled attention on the SST data (with task classes {positive, negative}) after training the model on News Series. For the corresponding positive/negative classes, we observe that attention is also distributed on task-related tokens (e.g., charming, affecting) besides common tokens.

Task-agnostic evaluation. In Table 6.2, we evaluate models in the task-agnostic setting after task-aware training to show the separation of data representations across tasks. Both PT+FT and Prescale perform better than FT. This may relate to their larger attention deviations and less over-smoothing (Fig. 6), which make data representations contain more information from non-sink tokens with different distributions across tasks. On BERT, PT+FT outperforms Prescale. We hypothesize that this is because BERT is pre-trained with the next sentence prediction, and thus [CLS] may contain more general sentence-level information across tasks than the learned scaling. On the other hand, for RoBERTa which does not have sentence-level pre-training, Prescale performs better than PT+FT.