🤖 AI

Attention Retention for Continual Learning with Vision Transformers

This paper proposes a novel attention-retaining framework for continual learning in Vision Transformers that mitigates catastrophic forgetting by identifying attention drift and applying instance-adaptive gradient masking to preserve previously learned visual concepts, achieving state-of-the-art performance across diverse scenarios.

Original authors: Yue Lu, Xiangyu Zhou, Shizhou Zhang, Yinghui Xing, Guoqiang Liang, Wencong Zhang

Published 2026-02-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yue Lu, Xiangyu Zhou, Shizhou Zhang, Yinghui Xing, Guoqiang Liang, Wencong Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a very smart, eager student (an AI) how to recognize different animals. First, you show them pictures of cats. They learn to look specifically at the pointy ears and whiskers to know it's a cat. Then, you show them pictures of dogs. They learn to look at floppy ears and wet noses.

The problem with traditional AI is a phenomenon called "Catastrophic Forgetting." When the student learns about dogs, their brain gets so excited about the new information that they completely forget what a cat looks like. They might start looking at the dog's floppy ears and thinking, "Oh, that's a cat!" or they might stop looking at the whiskers entirely because the new lesson overwrote the old one.

This paper introduces a new way to teach this student, called ARCL-ViT, which acts like a "Memory Guardian" to stop them from forgetting.

The Core Problem: "Attention Drift"

The authors discovered that in modern AI models (specifically Vision Transformers, or ViTs), the issue isn't just that the student forgets; it's that their focus shifts.

Think of the AI's "attention" like a flashlight.

When learning about cats, the flashlight shines brightly on the cat's whiskers.
When the AI learns about dogs, the flashlight starts to drift. It moves away from the whiskers and starts shining on the dog's nose.
If the flashlight moves too much, the AI loses its ability to recognize the cat later because it's no longer looking at the right spot.

The Solution: The "Do Not Touch" Mask

The authors propose a clever trick to keep the flashlight steady. Instead of trying to memorize every old picture (which takes up too much space), they use a two-step process:

Step 1: Mapping the "Golden Zones"

After the student finishes learning about cats, the system takes a snapshot of exactly where the flashlight was shining. It creates a map (a mask) that highlights the "Golden Zones"—the specific parts of the image that were crucial for recognizing the cat (like the whiskers).

The Analogy: Imagine the teacher draws a red circle around the cat's whiskers on a piece of paper and says, "This is the most important part. Do not change your focus here."

Step 2: The "Stop-Sign" for New Learning

When the student starts learning about dogs, the system looks at that red circle. If the new lesson tries to change how the AI looks at the whiskers (because the math says it should), the system hits the brakes.

The Mechanism: In the language of AI, this is called Gradient Masking. When the AI calculates how to update its brain, the system puts a "Stop Sign" on the parts of the brain that control the whisker-focus. It tells the AI: "You can learn about the dog's nose, but you are strictly forbidden from changing how you look at the whiskers."

To make sure this doesn't slow the AI down or confuse its learning engine (the optimizer), the system also adjusts the speed of the learning, ensuring the AI learns the new dog facts smoothly without accidentally erasing the cat facts.

Why This is Special

Most other methods try to solve this by:

Replaying old data: Showing the student old cat pictures every time they learn a new dog picture. (This is hard because you need to store all those old pictures).
Building new rooms: Adding new parts of the brain for every new animal. (This makes the brain huge and messy).

The authors' method is different because it locks the focus in place. It doesn't need to store old pictures, and it doesn't need to build new rooms. It just ensures the flashlight stays steady on the important features.

The Results

The paper tested this method on various difficult challenges, like recognizing animals in different artistic styles or from different domains.

The Outcome: The AI using this method remembered the old concepts (cats) much better than the standard AI, while still learning the new ones (dogs) very well.
The Proof: They showed pictures of the "flashlight" (attention maps). The standard AI's flashlight had drifted all over the place, while the new method's flashlight stayed perfectly focused on the original features, just like a human would.

In short, this paper teaches AI how to learn new things without losing its focus on the things it already knows, mimicking how human attention naturally stabilizes important concepts.

Technical Summary: Attention Retention for Continual Learning with Vision Transformers

Problem Statement

Continual Learning (CL) aims to enable AI systems to acquire knowledge progressively from non-stationary data streams. However, catastrophic forgetting remains a critical barrier, where learning new tasks degrades performance on previously learned ones. While various strategies exist (replay, expansion, regularization), this paper identifies a specific mechanism of forgetting in Vision Transformers (ViTs): attention drift.

The authors observe that as ViTs learn new tasks, their attention mechanisms shift significantly away from the discriminative features of previously learned visual concepts. Unlike human visual systems, which maintain selective attention on key features even as new concepts are acquired, ViTs lose this stability, leading to the erosion of learned knowledge. Existing methods often fail to address this specific drift in the attention mechanism of ViTs, particularly in replay-free settings where storing previous task data is prohibited.

Methodology: ARCL-ViT

The proposed framework, ARCL-ViT (Attention-Retaining Continual Learning with Vision Transformers), addresses attention drift by explicitly constraining gradient updates during the learning of new tasks. The method operates without storing previous task data and consists of two primary stages:

1. Adaptive Mask Generation

After completing a task $T_{t-1}$ , the model generates a binary mask to identify critical attention regions.

Layer-wise Rollout: Instead of using raw attention maps from a single layer, the method employs a layer-wise rollout mechanism. It aggregates activated attention matrices across all layers ( $L$ ) to compute a comprehensive attention map $\hat{S}^l_{t-1}$ . This involves multiplying transformed attention matrices ( $\tilde{S} = I + S$ ) to capture the flow of information from the input to the class token.
Instance-Adaptive Thresholding: To handle varying sizes of discriminative regions across images, a fixed threshold is avoided. The method sorts attention values and identifies an inflection point (via the second-order derivative) where values transition from background to salient regions. This determines an adaptive threshold $\tau$ , generating a binary mask $M$ where high-attention regions are marked as 0 (to be preserved) and low-attention regions as 1.

2. Gradient Masking and Optimizer Adaptation

During the training of the current task $T_t$ , the generated mask is used to suppress updates that would alter the attention patterns of previous tasks.

Gradient Zeroing: The mask is applied element-wise to the gradients of the Query ( $Q$ ), Key ( $K$ ), and Value ( $V$ ) projection weights ( $W_q, W_k, W_v$ ). Specifically, gradients corresponding to the high-attention regions of the previous task are zeroed out. This prevents the optimization process from shifting the model's focus away from previously learned concepts.
Optimizer Compatibility: Simply zeroing gradients can disrupt modern optimizers (e.g., Adam) that rely on momentum and variance statistics, potentially causing unstable updates. To mitigate this, the authors propose a scaling mechanism. The parameter update $\Delta W'$ is scaled proportionally to the ratio of the masked gradient to the original gradient. This ensures that the relative magnitude of updates is maintained, allowing the optimizer to function correctly while still respecting the attention constraints.

Key Contributions

Identification of Attention Drift: The paper establishes attention drift in ViTs as a primary source of catastrophic forgetting in continual learning, highlighting the necessity of preserving selective attention to prior visual concepts.
Novel Framework (ARCL-ViT): A gradient-based framework is introduced that constrains attention drift without requiring data replay. It utilizes a layer-wise rollout for robust attention map extraction and an adaptive thresholding strategy for precise mask generation.
Optimizer-Aware Gradient Masking: The method introduces a scaling technique to ensure gradient masking is compatible with adaptive optimizers, preventing optimization instability while effectively preserving learned attention patterns.
State-of-the-Art Performance: The approach achieves superior results across diverse benchmarks and pre-training schemes, demonstrating robust generalizability.

Experimental Results

The method was evaluated on four class-incremental learning benchmarks: 10-split and 20-split ImageNet-R, 10-split CIFAR-100, and 10-split DomainNet.

Performance: ARCL-ViT achieved state-of-the-art (SOTA) performance, outperforming recent methods (e.g., VPT-CPG, SD-LoRA, BiLoRA) by an average of 1.8% to 3.7% in accuracy across benchmarks.
Forgetting Reduction: The method significantly reduced catastrophic forgetting. For instance, on 10-split ImageNet-R, it achieved 81.17% accuracy with 4.84% forgetting, compared to the sequential fine-tuning (Seq-FT) baseline which suffered from 44.31% forgetting and only 49.43% accuracy.
Ablation Studies:
- Mask Construction: The combination of layer-wise rollout and adaptive thresholding yielded the best results, outperforming raw attention or fixed thresholds.
- Masked Regions: Masking the specific attention regions of previous tasks was proven essential; masking random regions or non-attention regions resulted in significantly higher forgetting.
Generalization:
- Pre-training Weights: The method generalized well across different pre-trained weights (DINO-1k and iBOT-1k), consistently outperforming baselines by ~37% in accuracy gain over Seq-FT.
- Long-Sequence CL: In challenging long-sequence settings (50 and 100 splits), ARCL-ViT maintained superior accuracy and achieved a better trade-off between plasticity and stability compared to existing SOTA methods.
Visualization: Quantitative analysis of attention drift showed ARCL-ViT maintained the lowest drift (6.0%) compared to Seq-FT (25.9%), visually confirming the preservation of attention on key visual concepts.

Significance and Claims

The paper claims that by aligning the ViT attention mechanism with biological principles of selective attention, it is possible to mitigate catastrophic forgetting effectively without the storage overhead of replay-based methods. The authors position their work as a significant step toward building continual learning systems that resemble biological intelligence, specifically by stabilizing the "neural resources" (attention) dedicated to previously learned concepts. The method is presented as a robust, data-free solution that is compatible with modern optimization techniques and generalizes across diverse pre-training strategies and long-term learning scenarios.