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

1. Problema Identificado

O uso de Grandes Modelos de Linguagem (LLMs) na medicina para tarefas como documentação clínica e geração de relatórios tem crescido. No entanto, há um desafio crítico conhecido como esquecimento catastrófico: quando um LLM de propósito geral é ajustado (fine-tuned) especificamente para um domínio médico (usando dados clínicos), ele tende a "esquecer" sua capacidade de seguir instruções (instruction-following), que foi desenvolvida durante o treinamento prévio com técnicas como RLHF (Reinforcement Learning from Human Feedback).

A adaptação tradicional via Supervised Fine-Tuning (SFT) em grandes conjuntos de dados médicos é cara, requer muitos dados rotulados e frequentemente degrada a habilidade do modelo de obedecer a comandos complexos, limitando sua utilidade prática em ambientes clínicos.

2. Metodologia

O estudo propõe um framework de fusão de modelos (Model Merging) no espaço de pesos como uma alternativa eficiente e de baixo custo computacional para adaptar LLMs ao domínio médico sem perder a capacidade de seguir instruções.

• Modelos Envolvidos:

  • Modelo Clínico (Base): GatorTronLlama, um modelo baseado em Llama-3.1-8B-Base, submetido a pré-treinamento contínuo (CPT) com 282 bilhões de tokens de texto clínico (literatura do PubMed, diretrizes e notas de saúde da UF Health).
  • Modelo de Instrução (Base): Llama-3.1-8B-Instruct, um modelo de propósito geral otimizado para seguir instruções via SFT e RLHF.
  • Baseline de Comparação: Uma versão do GatorTronLlama ajustada via SFT usando o dataset UltraMedical.

• Técnicas de Fusão:
  Os autores fundiram os pesos do modelo clínico ($\theta_{clínico}$) com os do modelo de instrução ($\theta_{instrução}$) utilizando duas estratégias de interpolação:

  1. Interpolação Linear: Uma média ponderada simples dos pesos: $\theta_{fusão} = (1 - \alpha)\theta_{clínico} + \alpha\theta_{instrução}$.
  2. Interpolação Linear Esférica (SLERP): Realiza uma interpolação ao longo de uma hipersfera no espaço de pesos, preservando melhor a magnitude dos vetores de parâmetros quando os modelos não estão perfeitamente alinhados.

• Seleção de Receitas (Merge Recipes):
  Foi realizada uma busca em grade (grid search) variando o peso de interpolação ($\alpha$ ou $t$) de 0,0 a 1,0. Os pontos ótimos foram selecionados com base na fronteira de Pareto, buscando o melhor equilíbrio entre o desempenho em benchmarks de conhecimento médico e a capacidade de seguir instruções.

• Avaliação:

  • Zero-shot: Avaliação em benchmarks de conhecimento médico (MMLU, MedMCQA, MedQA, PubMedQA) e de instruções (IFEval).
  • Ajuste Fino Pós-Fusão (Post-merge SFT): Os melhores modelos fundidos foram submetidos a SFT em cinco tarefas de geração clínica (resumo de relatórios de radiologia, resumo de alta hospitalar, criação de listas de problemas, etc.) usando poucos exemplos (few-shot) e dados completos.

3. Principais Contribuições

• Solução para Esquecimento Catastrófico: Demonstra que a fusão de modelos pode transferir a capacidade de seguir instruções de um LLM geral para um modelo clínico especializado, mitigando o esquecimento típico do SFT tradicional.
• Eficiência Computacional: A fusão de modelos é uma solução "sem treinamento" (training-free) que leva apenas 40 minutos para gerar checkpoints, comparado a 18 horas e 20 minutos necessários para o ajuste fino tradicional em um único GPU B200.
• Eficiência de Dados: O modelo fundido (SLERP) demonstrou uma eficiência de dados superior, alcançando desempenho comparável ao modelo base treinado com 256 exemplos usando apenas 64 exemplos em tarefas de sumarização de radiologia.

4. Resultados Chave

• Desempenho Zero-shot:

  • O modelo fundido SLERP (t=0.4) superou o modelo clínico puro em tarefas médicas, alcançando uma média médica de 0,7039 (uma melhoria de 1,43% em relação ao GatorTronLlama).
  • Ao mesmo tempo, recuperou 98,3% da capacidade de seguir instruções do modelo geral (Llama-3.1-8B-Instruct), com pontuação de 0,5166 no IFEval (vs. 0,5253 do modelo puro).
  • A relação entre a taxa de fusão e o desempenho não é monotônica; existe um "ponto ideal" (sweet spot) onde a injeção de parâmetros de alinhamento melhora o desempenho clínico sem diluir excessivamente a expertise especializada.

• Desempenho Pós-SFT (Tarefas Clínicas):

  • Após o ajuste fino em cinco tarefas clínicas, o modelo fundido via SLERP obteve consistentemente as melhores pontuações compostas (média de ROUGE, BLEU, METEOR, BERTScore), superando tanto o modelo clínico base quanto a baseline ajustada via SFT.
  • Exemplo: Na tarefa de geração de lista de problemas (Problem List), o modelo SLERP com 64 exemplos superou o modelo clínico base treinado com 512 exemplos.

5. Significado e Conclusão

O estudo valida que a fusão de modelos no espaço de pesos é uma estratégia altamente escalável e eficiente para adotar LLMs de código aberto em aplicações clínicas.

• Impacto Prático: Permite que instituições de saúde com recursos computacionais limitados utilizem modelos de ponta sem a necessidade de custosos processos de RLHF ou grandes conjuntos de dados de instruções médicas rotuladas.
• Viabilidade: A abordagem resolve o dilema entre especialização clínica e flexibilidade de instruções, oferecendo um caminho viável para a implantação segura e eficaz de IAs generativas em ambientes de saúde, facilitando a automação de fluxos de trabalho clínicos e reduzindo o burnout médico.

O trabalho sugere que a fusão de modelos não é apenas uma técnica de otimização, mas uma ferramenta fundamental para a adaptação de domínio em cenários onde a coleta de dados instrucionais de alta qualidade é proibitivamente cara.