Technical Summary: Attribution-Guided Continual Learning for Large Language Models

Problem Statement

Large Language Models (LLMs) trained on static, broad-domain data often suffer from catastrophic forgetting when adapted to new tasks via continual learning. While existing mitigation strategies—such as data replay, parameter regularization, and freezing specific layers—have shown some success, they lack a mechanistic understanding of how knowledge and task-specific behaviors are distributed across the internal parameters of LLMs. Consequently, these methods cannot effectively distinguish which parameters should be preserved to retain prior knowledge and which should remain plastic to accommodate new tasks. This limitation often leads to interference, where updates for a new task inadvertently modify parameters critical to previous tasks.

Methodology

The authors propose an attribution-guided continual fine-tuning framework that leverages parameter importance estimation to modulate gradient updates. The core methodology consists of two main phases:

1. Element-wise Parameter Importance Estimation

The framework adapts Layer-wise Relevance Propagation (LRP) to attribute the predicted next-token logit to individual model parameters within each Transformer layer.

• Adaptation to Transformers: Standard LRP attributes relevance to input features. The authors extend this to quantify the contribution of model parameters (weights and biases) in Feed-Forward Networks (FFN) and Multi-Head Attention (MHA) modules.
• Bilinear Attribution: To handle matrix multiplications inherent in Transformers (e.g., $QK^T$ and $WV$), the method employs a bilinear attribution principle (inspired by AttnLRP). This treats matrix multiplication as a bilinear operation, distributing relevance scores between the input activations and the weight matrices while satisfying a conservation property.
• Derivation:
  • FFN: Relevance is propagated through the two linear layers and the activation function, deriving scores for weights $W_1, W_2$ and biases $b_1, b_2$.
  • MHA: Relevance is propagated through the attention mechanism, deriving scores for query, key, value, and output projection weights ($W_Q, W_K, W_V, W_O$).
  • LoRA Support: The method is also formulated for LoRA-adapted models, computing relevance for both the frozen backbone and the trainable LoRA matrices ($A, B$).

2. Attribution-Guided Fine-Tuning

The estimated importance scores are used to construct a task-specific prior that guides the optimization process for new tasks.

• Prior Construction: For a new task $T_t$, the method first estimates a task-specific importance map $R(W^{(l,t)})$ by aggregating normalized relevance scores from correctly generated samples.
• Historical Aggregation: During the learning of task $T_t$, the framework aggregates importance from all previous tasks ($T_1, \dots, T_{t-1}$) by taking the element-wise maximum of their importance scores:
  $$\bar{R}(W^{(l,t)}) = \left[ \max_{\tau < t} R(W^{(l,\tau)}) \right]_+$$
• Gradient Modulation: The aggregated importance map acts as a gradient gate. The standard gradient $\nabla W$ is modulated as follows:
  $$\tilde{\nabla} W^{(l,t)} = \left( 1 - \bar{R}(W^{(l,t)}) \right) \odot \nabla W^{(l,t)}$$
  • Parameters with high historical relevance (important to past tasks) receive a gate value close to 0, resulting in smaller updates.
  • Parameters with low historical relevance receive a gate value close to 1, allowing them to be freely updated to learn the new task.

Key Contributions

1. Mechanistic Insight: The paper provides a mechanistic analysis showing that in sequential continual learning, the sets of parameters deemed "important" for different tasks become highly overlapped in the final model, unlike in independent single-task fine-tuning where they remain distinct. This overlap explains the interference causing catastrophic forgetting.
2. Attribution Framework: It introduces a novel application of LRP to attribute relevance directly to LLM parameters rather than just inputs, deriving specific propagation rules for Transformer components (FFN and MHA).
3. Gradient Modulation Strategy: The proposed method offers a dynamic, element-wise approach to continual learning that selectively constrains updates based on learned parameter importance, without requiring data replay or rigid parameter freezing.

Experimental Results

The framework was evaluated on continual learning benchmarks using Large Language Models (specifically Llama-3.2-Instruct-3B) under both full supervised fine-tuning and LoRA-based adaptation settings.

• Performance: The method consistently outperformed existing continual learning baselines (including regularization, replay, and freezing methods).
• Retention vs. Plasticity: It achieved superior retention of performance on previously learned tasks while maintaining competitive performance on new tasks.
• Observation: The experiments confirmed that the proposed method effectively reduces the interference between tasks by preserving the specific parameter subsets critical to prior knowledge.

Significance and Claims

The paper claims that by moving beyond heuristic strategies (like random replay or layer freezing) and instead utilizing attribution-based importance estimation, continual learning for LLMs can be made more efficient and effective. The significance lies in:

• Semantic Awareness: The method explicitly accounts for the internal distribution of knowledge within the model, distinguishing between parameters that should be preserved and those that can be adapted.
• Generalizability: The framework is applicable to both full-parameter fine-tuning and parameter-efficient fine-tuning (LoRA) scenarios.
• Theoretical Grounding: It bridges the gap between explainability (attribution methods) and optimization (continual learning), offering a principled way to mitigate catastrophic forgetting based on the model's own learned representations.