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

Titolo

Contrasto dell'Oblio Catastrofico nei Modelli Linguistici di Grandi Dimensioni (LLM) per un Migliore Seguito delle Istruzioni tramite Fusione di Modelli nello Spazio dei Pesi.

1. Il Problema

L'adozione dei Large Language Models (LLM) generici nel dominio medico incontra una sfida critica nota come oblio catastrofico (catastrophic forgetting).

• Contesto: Gli LLM generici (come Llama) sono eccellenti nel seguire le istruzioni, ma mancano di conoscenze cliniche specifiche. Per adattarli alla medicina, vengono spesso sottoposti a fine-tuning su dataset medici (ad esempio, note cliniche).
• La Sfida: Quando un modello generico viene addestrato su dati medici specifici, tende a "dimenticare" le capacità di instruction following (seguire le istruzioni) acquisite durante la fase di pre-addestramento e RLHF (Reinforcement Learning from Human Feedback).
• Conseguenza: I modelli medici risultanti sono spesso esperti in ambito clinico ma incapaci di seguire prompt complessi o di generare output strutturati come richiesto, limitandone l'utilità pratica in scenari reali.
• Limiti delle Soluzioni Esistenti: Creare dataset di istruzioni cliniche di alta qualità è costoso e il Supervised Fine-Tuning (SFT) su tali dati è computazionalmente oneroso e spesso non risolve il problema dell'oblio.

2. Metodologia

Gli autori propongono un framework di fusione di modelli (Model Merging) nello spazio dei pesi come alternativa efficiente e priva di addestramento (training-free) per integrare conoscenze cliniche e capacità di seguire le istruzioni.

• Modelli Coinvolti:

  1. Modello Clinico (GatorTronLlama): Un modello basato su Llama-3.1-8B-Base, sottoposto a pre-addestramento continuo (CPT) su 282 miliardi di token di testi clinici (letteratura PubMed, linee guida, note UF Health).
  2. Modello Istruttivo Generale (Llama-3.1-8B-Instruct): Un modello generico ottimizzato per il seguito delle istruzioni tramite SFT e RLHF.
  3. Baseline SFT: GatorTronLlama sottoposto a SFT su dataset UltraMedical per confronto.

• Tecnica di Fusione:
  I pesi del modello clinico ($\theta_c$) e del modello istruttivo ($\theta_i$) vengono fusi per creare un modello intermedio ($\theta_m$) utilizzando due strategie di interpolazione:

  • Interpolazione Lineare: $\theta_m = (1 - \alpha)\theta_c + \alpha\theta_i$.
  • Interpolazione Sferica Lineare (SLERP): Un'interpolazione lungo una ipersfera nello spazio dei pesi, che preserva meglio la norma dei vettori rispetto alla linea retta, utile quando i checkpoint non sono ben allineati.
  • Ricerca dell'Ipertrofia: È stata eseguita una ricerca a griglia sul parametro di fusione $\alpha$ (o $t$) da 0.0 a 1.0 per identificare i checkpoint ottimali sul fronte di Pareto, bilanciando prestazioni cliniche e capacità di seguire le istruzioni.

• Valutazione:

  • Zero-shot: Valutazione su benchmark di conoscenza medica (MMLU-med, MedQA, PubMedQA) e di istruzioni (IFEval).
  • Post-Fusione SFT: I checkpoint fusi ottimali sono stati sottoposti a Supervised Fine-Tuning su 5 task clinici di generazione (sintesi di rapporti radiologici, note di dimissione, liste di problemi attivi, ecc.) con diverse configurazioni di dati (few-shot da 8 a 512 esempi e full-data).

3. Risultati Chiave

• Efficienza Computazionale: La fusione dei modelli richiede solo 40 minuti per generare 9 checkpoint su una singola GPU B200, contro le 18 ore e 20 minuti richieste per un addestramento SFT completo su dataset medici. Non richiede dati etichettati aggiuntivi per la fase di fusione.
• Mitigazione dell'Oblio Catastrofico:
  • Il modello clinico puro (GatorTronLlama) ha un punteggio IFEval (seguire istruzioni) di 0.2244, mentre il modello generico (Llama-Instruct) ha 0.5253.
  • I modelli fusi (SLERP con $t=0.7$) hanno recuperato il 98.3% della capacità di seguire le istruzioni (IFEval: 0.5166) mantenendo prestazioni cliniche superiori al modello clinico puro.
  • SLERP con $t=0.4$ ha migliorato le prestazioni cliniche del 1.43% rispetto al modello clinico di base, dimostrando che l'iniezione di pesi "istruttivi" può migliorare anche la conoscenza di dominio.
• Relazione Non Monotona: È stata osservata una relazione non monotona tra il rapporto di fusione e le prestazioni. Le prestazioni cliniche raggiungono un picco a pesi di interpolazione moderati (0.4-0.6), suggerendo l'esistenza di un "punto dolce" (sweet spot) dove l'equilibrio è ottimale.
• Efficienza dei Dati (Few-Shot Learning):
  • I modelli fusi richiedono significativamente meno dati per raggiungere prestazioni competitive.
  • Nel task di sintesi dei rapporti radiologici (IU-Xray), il modello SLERP fuso addestrato su 64 esempi ha eguagliato le prestazioni del modello clinico di base addestrato su 256 esempi.
  • Nel task di generazione delle liste di problemi, il modello SLERP con 64 esempi ha superato le prestazioni del modello di base con 512 esempi.
• Prestazioni Post-SFT: Dopo il fine-tuning completo, i modelli fusi (in particolare SLERP) hanno ottenuto i punteggi compositi più alti su tutti e cinque i task clinici, superando sia il modello clinico di base che la baseline SFT clinica.

4. Contributi Principali

1. Soluzione all'Oblio Catastrofico: Dimostrazione che la fusione di modelli nello spazio dei pesi è un metodo efficace per trasferire capacità di instruction following ai modelli medici senza perdere la conoscenza di dominio.
2. Efficienza delle Risorse: Offerta di una strategia scalabile e a basso costo computazionale per l'adattamento dei LLM open-source al settore sanitario, riducendo la dipendenza da dataset di istruzioni cliniche costosi e dall'addestramento su larga scala.
3. Nuova Strategia di Adattamento: Identificazione che l'interpolazione sferica (SLERP) con parametri specifici offre un punto di partenza superiore per il fine-tuning clinico rispetto all'inizializzazione standard.

5. Significato e Implicazioni

Questo studio offre una via praticabile per l'implementazione di LLM generici in ambienti sanitari con risorse limitate.

• Scalabilità: La fusione dei modelli permette di adottare rapidamente nuovi modelli open-source nel dominio medico senza i costi proibitivi del RLHF o dell'SFT massiccio.
• Sicurezza e Affidabilità: Mantenere le capacità di seguire le istruzioni è cruciale per l'interazione uomo-macchina in ambito clinico, garantendo che il modello esegua compiti specifici (es. "riassumi questo rapporto") senza allucinazioni o deviazioni.
• Futuro: Sebbene la valutazione si sia basata su metriche automatiche, il lavoro apre la strada a futuri studi che includano la validazione clinica umana e l'estensione a modelli di dimensioni maggiori, confermando che l'interpolazione dei pesi è una tecnica fondamentale per l'evoluzione dei sistemi di intelligenza artificiale in medicina.