Yue Lu, Xiangyu Zhou, Shizhou Zhang, Yinghui Xing, Guoqiang Liang, Wencong Zhang
9 min
Abstract
Continual learning (CL) empowers AI systems to progressively acquire knowledge from non-stationary data streams. However, catastrophic forgetting remains a critical challenge. In this work, we identify attention drift in Vision Transformers as a primary source of catastrophic forgetting, where the attention to previously learned visual concepts shifts significantly after learning new tasks. Inspired by neuroscientific insights into the selective attention in the human visual system, we propose a novel attention-retaining framework to mitigate forgetting in CL. Our method constrains attention drift by explicitly modifying gradients during backpropagation through a two-step process: 1) extracting attention maps of the previous task using a layer-wise rollout mechanism and generating instance-adaptive binary masks, and 2) when learning a new task, applying these masks to zero out gradients associated with previous attention regions, thereby preventing disruption of learned visual concepts. For compatibility with modern optimizers, the gradient masking process is further enhanced by scaling parameter updates proportionally to maintain their relative magnitudes. Experiments and visualizations demonstrate the effectiveness of our method in mitigating catastrophic forgetting and preserving visual concepts. It achieves state-of-the-art performance and exhibits robust generalizability across diverse CL scenarios.
Alex: Twice the accuracy without old photos—that's a clear improvement. But how do they actually build this targeted fix? Walk me through the steps inside their framework.
Sam: After finishing one task, say learning cats, the model first creates a map of where it was looking most—the key patches like fur or eyes. It does this by combining attention spotlights from every layer of the network, multiplying them step by step to get a full-path view from input to output, almost like tracing a river's flow through connected streams. Researchers call this combined map from layer-wise attention rollout. Then, for each image, it sorts those attention scores and finds the sharp drop-off point where focus fades from strong to weak, setting a custom cutoff per picture to separate the vital spots from background. That's adaptive thresholding.
Alex: Okay, so they've got this precise map of old focus areas without storing cat photos. Like outlining the important zones on a blueprint. But then what—how does that stop the drift when dogs come along?
Sam: Exactly. They turn that map into a simple on-off shield—a binary mask where old focus zones are marked "do not change," like taping "Do Not Touch" signs over specific light switches in a control room so new wiring doesn't flip off the old lights. During dog training, when the model calculates adjustments to its weights—those tiny tweaks called gradients that shift how it pays attention—it multiplies by this mask, zeroing out changes in the taped-off cat zones. This blocks updates to query, key, and value projections, the three matching parts that decide attention similarities, preventing drift there while letting dog learning happen elsewhere.
Alex: Hold on—these query, key, value things sound central. Break that down: why zero gradients just there, and not everywhere?
Sam: Think of attention like a matchmaking game: for each patch, it creates a "query" description of what it's seeking, "key" tags on other patches, and "value" info to grab if they match well—like swiping right in a dating app based on profiles. Gradients push those matchmakers toward new task winners, but masking zeros the push in old zones, so cat matches stay stable. One catch: optimizers like Adam don't use raw gradients; they smooth past ones for stable steps. So they scale the final updates to match the masked gradient ratio, keeping everything balanced without big jumps or stalls.
Alex: So it's surgically protecting old gaze without freezing the whole model. Does this work because attention gradients directly tweak those QKV parts?
Sam: Yes—the math shows gradients for QKV weights flow straight from attention changes, so masking there cuts the drift at the source. In class-incremental setups, where each round adds totally new labels like cats then dogs with no overlap, this data-free shield lets the model classify everything seen so far at test time, no task hints needed. The paper suggests this preserves prior concepts effectively, as their visuals show steady focus post-ten tasks.
Alex: That logic tracks—it's not brute force, but precise protection. A meaningful way forward for streaming image tasks. But does the evidence back it up—how does it stack against plain training or other fixes?
Sam: The paper tests it on challenging sequences, like splitting thousands of varied images into ten rounds of new classes each. Their approach reaches about twice the accuracy of sequential fine-tuning alone, while cutting forgetting by nearly half. This holds because masking directly blocks drift in those key attention spots, as visuals confirm steady focus even after ten steps. They fine-tune just those matching parts plus classifiers with standard steps, outperforming fifteen others from recent years by one to four points on average. The direct gap to no-protection training—37% higher accuracy, 46% less forgetting—pins the win to this drift block.
Alex: Twice the accuracy, half the forgetting—without old data. That's a substantial edge. What about tweaks to the pieces, like different ways to build the maps or pick thresholds?
Sam: They check that through careful breakdowns. Using the full-path map from combining layers works better than raw single-layer views or simple sums, lifting accuracy by around one to two points because it traces exactly where focus flows end-to-end. Adaptive cutoffs, spotting the sharp rise in sorted scores like finding the knee in a bent elbow, add another few points over fixed ones by fitting each image's unique hot zones—no manual tuning needed. Protecting those prior focus areas boosts accuracy while random or background masking hurts more, as it lets drift sneak in or blocks useful new learning.
Alex: Makes sense—the map has to be accurate, and the cutoff tailored. So only tweaking the query-key-value matchers in the ViT blocks, keeping the base image features frozen. And it tops recent methods across setups?
Alex: To wrap this up, how does it hold up in tougher, longer runs of tasks?
Sam: The paper pushes it to long sequences, splitting big image sets into 50 or even 100 rounds of new classes. Their method tops the previous best by about three points in accuracy on one key test, trading some controlled forgetting for better overall retention of everything learned so far. Visuals and metrics show it keeps attention drift to just six percent after many steps, versus over 25 for plain training. It works well with varied pre-training weights, like self-supervised ones from DINO or iBOT, gaining around 38 points over baselines and beating others by two to four.
Alex: So a solid trade-off—keeps gaining new skills without total wipeouts.
Sam: A few practical limits emerge in the longer runs, like 100 tasks, where some forgetting builds up despite controls, suggesting it shines more in moderate sequences. It tunes just the attention-matching parts in these models, so other setups might need tweaks, and storing one mask per past task adds a bit of memory overhead—no old data, but those shields take space.
Alex: Fair points—nothing's perfect for endless growth yet. Still, surgically holding gaze on old patterns opens doors for systems that adapt lifelong, say in robotics scanning new scenes or cameras learning fresh threats without full resets.
Sam: Precisely—this data-free way suggests a path to reliable, ongoing image learning without constant retraining or data hoards. The evidence points to meaningful gains in keeping visual know-how stable amid change.
Alex: A clear step for building adaptable AI eyes. That's our look at attention retention in continual learning with Vision Transformers. Thanks for joining, everyone.