Countering Catastrophic Forgetting of Large Language Models for Better Instruction Following via Weight-Space Model Merging

1. Problem Statement

The integration of Large Language Models (LLMs) into the medical domain faces a critical challenge known as catastrophic forgetting.

• The Dilemma: General-purpose LLMs (e.g., Llama-3.1) possess strong instruction-following capabilities but lack specific clinical knowledge. Conversely, clinical foundation models (e.g., GatorTron) are trained on medical corpora to understand clinical context but often lose their ability to follow complex instructions when fine-tuned on task-specific medical datasets.
• Current Limitations: Traditional adaptation methods, such as Supervised Fine-Tuning (SFT) or Reinforcement Learning from Human Feedback (RLHF), are computationally expensive and require large, high-quality, specialized instruction datasets that are difficult and costly to curate in the medical field. Furthermore, these methods often exacerbate the forgetting of instruction-following abilities.
• Goal: The study aims to develop a resource-efficient method to adapt open-source general-purpose LLMs to the medical domain while retaining their instruction-following capabilities and acquiring clinical expertise, without the heavy computational cost of full fine-tuning.

2. Methodology

The authors propose a weight-space model merging framework that combines a clinical foundation model with a general instruct model without additional gradient-based training.

Models Involved

1. Clinical Foundation Model: GatorTronLlama (based on Llama-3.1-8B-Base), continuously pre-trained on 282 billion tokens of medical literature and de-identified clinical notes from UF Health.
2. General Instruct Model: Llama-3.1-8B-Instruct, optimized for instruction following via SFT and RLHF.
3. Baseline: GatorTronLlama_SFT, a version of the clinical model further fine-tuned on the UltraMedical dataset for comparison.

Merging Strategies

The study employs two interpolation methods to merge the parameters ($\theta_{clinical}$ and $\theta_{instruct}$) into a new checkpoint ($\theta_{merged}$):

• Linear Interpolation: A simple weighted average: $\theta_{merged} = (1 - \alpha)\theta_{clinical} + \alpha\theta_{instruct}$.
• Spherical Linear Interpolation (SLERP): Interpolates along a hypersphere in weight space to better preserve the magnitude (norm) of parameter vectors, potentially offering better alignment than linear blending.

Experimental Design

• Phase 1 (Zero-Shot Search): A grid search was performed over interpolation weights ($\alpha$ or $t \in [0.0, 1.0]$) to identify Pareto-optimal checkpoints. These are models that balance clinical knowledge (measured by MMLU-med, MedMCQA, etc.) and instruction following (measured by IFEval).
• Phase 2 (Downstream Adaptation): The best merged checkpoints were subjected to few-shot and full-parameter SFT on five clinical generation tasks:
  1. Radiology report summarization (IU X-ray).
  2. Brief hospital course summarization (MIMIC-IV).
  3. Dialogue-to-note generation (MTS-Dialog, ACI-Bench).
  4. Patient problem list generation (BioNLP).

3. Key Contributions

1. Counteracting Catastrophic Forgetting: Demonstrated that model merging effectively bridges the gap between clinical domain expertise and general instruction-following, preventing the loss of the latter during domain adaptation.
2. Training-Free Efficiency: The merging process is training-free, requiring only 40 minutes to generate 9 checkpoints on a single B200 GPU, compared to ~18 hours for standard instruction fine-tuning.
3. Data Efficiency: Showed that merged models require significantly fewer labeled examples to reach competitive performance. For instance, a merged model trained on 64 shots matched or exceeded the performance of a foundation model trained on 256 shots.
4. Discovery of Non-Monotonic Dynamics: Identified that performance does not scale linearly with the merge ratio. There is a "sweet spot" (e.g., SLERP at $t=0.4$) where a limited injection of instruction weights boosts clinical performance without sacrificing instruction adherence.

4. Key Results

Zero-Shot Performance

• Baseline Comparison: GatorTronLlama had high clinical knowledge (Medical Avg: 0.6896) but poor instruction following (IFEval: 0.2244). Llama-3.1-8B-Instruct had high instruction following (0.5253) but lower clinical knowledge (0.6845).
• Merged Success: The SLERP (t=0.4) model achieved a Medical Avg of 0.7039 (a 1.43% improvement over the clinical baseline) while retaining 98.3% of the instruction-following capability of the general instruct model (IFEval: 0.5166).

Downstream Fine-Tuning (SFT)

• Composite Scores: The SLERP-merged model consistently achieved the highest composite scores (BLEU, ROUGE, METEOR, BERTScore) across all five clinical tasks after full SFT, outperforming both the pure clinical baseline and the instruction-tuned clinical baseline.
  • Example: On the ACI-Bench task, SLERP achieved 0.4099, compared to 0.3685 for the clinical baseline.
• Few-Shot Learning:
  • Radiology (IU-Xray): SLERP trained on 64 examples (Score: 0.5167) matched the performance of the foundation model trained on 256 examples (Score: 0.5157).
  • Problem List: SLERP trained on 64 examples (Score: 0.2587) slightly surpassed the foundation model trained on the full 512 examples (Score: 0.2581).

5. Significance and Implications

• Scalability for Healthcare: This approach offers a highly scalable solution for deploying open-source LLMs in resource-constrained healthcare environments. It eliminates the need for expensive RLHF pipelines and massive instruction datasets.
• Cost Reduction: By reducing the required labeled data by a factor of 4x to 8x (e.g., 64 shots vs. 256/512 shots), the method significantly lowers the cost and time associated with clinical model adaptation.
• Robustness: The study proves that weight-space merging is a viable alternative to gradient-based fine-tuning, preserving the "best of both worlds" (clinical knowledge + instruction following) and mitigating the catastrophic forgetting that plagues traditional domain adaptation.

6. Limitations and Future Work

• Evaluation Metrics: The study relies heavily on automatic metrics (ROUGE, BLEU). Future work requires rigorous clinician adjudication to assess factual correctness, hallucination rates, and patient safety.
• Search Space: The hyperparameter grid was relatively constrained; broader searches might yield further optimizations.
• Model Scale: Experiments were conducted on 8B parameter models; future research is needed to verify if these dynamics hold for larger-scale models.