Joint Flashback Adaptation for Forgetting-Resistant Instruction Tuning

1. Problem Statement

Large Language Models (LLMs) face the catastrophic forgetting problem when incrementally learning new tasks. When fine-tuned on new instruction-following tasks, models tend to lose previously acquired knowledge. Existing solutions suffer from significant limitations in real-world scenarios:

• Replay-based methods: Require access to original training data, which is often unavailable due to privacy concerns or proprietary restrictions.
• Regularization-based methods: Often lead to under- or over-constraint issues or implicitly require past data.
• Parameter-Efficient Fine-Tuning (PEFT) & Adapters: May require task-specific identifiers (task-IDs), reducing flexibility, or fail to fully mitigate knowledge conflicts.

The goal is to adapt an existing, well-aligned LLM to new tasks while preserving performance on old tasks, without relying on replay data or task differentiation.

2. Methodology: Joint Flashback Adaptation (JFA)

The authors propose Joint Flashback Adaptation (JFA), a task-agnostic framework that combines limited "flashback" prompts with a novel joint task learning mechanism.

A. Flashback Adaptation Framework

Instead of replaying full datasets, JFA uses a small set of flashbacks:

• Definition: A limited number of prompts (without labels) sampled from old tasks (or manually crafted).
• Constraint: The model is constrained to generate outputs similar to a reference model (the original pre-trained model) when processing these flashback prompts.
• Loss Function:
  • $L_{SFT}$: Standard supervised fine-tuning loss on new tasks.
  • $L_{DIV}$: Divergence loss (dual KL divergence) between the current model and the reference model on flashback prompts.
  • Total Objective: $L = L_{SFT} + \alpha \cdot L_{DIV}$.
• Gradient Projection: To resolve conflicts between minimizing new task loss and minimizing forgetting, the authors employ PCGrad (Gradient Projection) to orthogonalize conflicting gradients.

B. Joint Task Learning (JTL)

To address the data sparsity of flashbacks and facilitate knowledge sharing, JFA introduces latent tasks as interpolations between flashbacks and new tasks.

• Latent Task Representation: Each latent task is a tuple $(e_j, \Delta\theta_j)$, where $e_j$ is a vector key and $\Delta\theta_j$ is a weight increment (compiled via LoRA-style low-rank matrices).
• Initialization: $H$ orthogonal keys are initialized. Weight increments are initialized using LoRA (random normal for $B$, zeros for $A$).
• Retrieval: For any input (new task or flashback), the model encodes the prompt using RoBERTa. It retrieves the top-$k$ most similar latent tasks based on cosine similarity between the input encoding and the latent keys.
• Re-parameterization: The model parameters for a specific sample are dynamically adjusted by interpolating the retrieved weight increments:
  $$\Theta'_{sample} = \Theta' + \sum_{j=1}^{k} D(q_i, e_j) \Delta\theta_j$$
  where $D$ is the similarity weight.
• Training: The model learns the weight increments ($\Delta\theta$) jointly with the base parameters, allowing knowledge to flow smoothly between old and new tasks via the latent space.

3. Key Contributions

1. Novel Framework (JFA): A continual learning method that requires only a few flashback prompts (no labels, no replay data) and is task-agnostic (no task IDs needed).
2. Joint Task Learning Mechanism: Introduces latent tasks to interpolate between flashbacks and new tasks, alleviating data sparsity and enabling effective knowledge sharing.
3. Empirical Validation: Extensive experiments demonstrating superior performance in generalization on new tasks and retention of old task capabilities compared to state-of-the-art baselines.

4. Experimental Results

• Setup: Evaluated on Vicuna-13B and Llama-3.1-8B.
• Datasets:
  • New Task: Super Natural Instructions (SNI) training set.
  • Old Tasks (Evaluation): GSM-8K, SVAMP (Arithmetic Reasoning); ARC-Challenge, MMLU-Pro, BIG-Bench-Hard (General Reasoning).
• Baselines: Vanilla LLM, SFT, PACE, Replay (full data), FROMP, and SLM.
• Performance:
  • Generalization: JFA achieved the best performance on the new task (SNI) across all metrics (BLEU, ROUGE), outperforming SFT, PACE, and other CL methods.
  • Forgetting Mitigation: JFA consistently delivered the best or second-best results on old tasks (e.g., GSM-8K, BBH), significantly outperforming standard SFT which suffered severe forgetting.
  • Efficiency: JFA matched the performance of "Replay" (which uses thousands of extra training samples) using only 30 flashback prompts (30 per dataset type).
• Ablation Studies:
  • Joint Task Learning: Removing JTL significantly degraded performance on both new and old tasks.
  • Flashbacks: Using 0 prompts caused severe forgetting; 10–20 prompts were sufficient to stabilize performance.
  • Hyperparameters: The method is robust to changes in the balance weight $\alpha$ and latent task group sizes ($C, Q, k$).

5. Significance and Limitations

Significance:

• Practicality: JFA solves the data privacy and storage bottlenecks of continual learning by eliminating the need for replay buffers.
• Scalability: It enables LLMs to adapt to thousands of new tasks incrementally without catastrophic forgetting, using minimal overhead (few prompts + latent task vectors).
• Generalization: The approach is task-agnostic, making it suitable for dynamic, real-world deployment where task definitions may not be known in advance.

Limitations:

• Model Dependency: Validated primarily on strong, aligned models; performance on weaker models with low-quality reference targets is unverified.
• Granularity: The method relies on global knowledge sharing and a scalar weight $\alpha$ to balance learning vs. retention. It cannot yet fine-grainly distinguish specific capabilities that need updating versus those that must be preserved, which is a direction for future work.