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

Probleemstelling

Grote Taalmodellen (LLM's) worden steeds vaker ingezet in de medische sector om de werkdruk op clinici te verlichten, bijvoorbeeld door het genereren van ontslagrapporten of samenvattingen van radiologieverslagen. Een kritieke uitdaging bij het toepassen van algemene LLM's in deze domeinen is het fenomeen van catastrophic forgetting (catastrofaal vergeten).

Wanneer een algemeen LLM (zoals Llama) wordt fijngefineerd (fine-tuned) op een specifiek medisch dataset om domeinkennis te verwerven, verliest het vaak zijn vermogen om instructies correct te volgen. Dit instructievolgingsvermogen is doorgaans ontwikkeld tijdens de post-trainingfase (via RLHF of SFT) van het algemene model. Het creëren van nieuwe, hoogwaardige medische instructiedatasets is echter zeer duur en tijdrovend. Traditionele methoden zoals Supervised Fine-Tuning (SFT) of Reinforcement Learning from Human Feedback (RLHF) vereisen bovendien aanzienlijke rekenkracht en grote hoeveelheden gelabelde data.

Methodologie

De auteurs stellen een model-merging framework voor als een rekenkundig efficiënt alternatief om algemene LLM's aan te passen aan de medische domein zonder het instructievolgingsvermogen te verliezen. De studie omvat de volgende stappen:

1. Model Selectie en Voorbereiding:

   • Klinisch Basismodel: GatorTronLlama, een Llama-3.1-8B-Base model dat is voortgezet gepretraind (CPT) op 282 miljard tokens aan medische data (PubMed, klinische richtlijnen en geanonimiseerde notities van UF Health).
   • Algemeen Instructiemodel: Llama-3.1-8B-Instruct, een model dat is geoptimaliseerd voor instructievolging via SFT en RLHF.
   • Referentiekader: Een SFT-versie van het klinische model (GatorTronLlama_SFT) getraind op het UltraMedical-dataset voor vergelijking.

2. Merging Strategieën:
   De auteurs combineren de gewichten van het klinische model ($\theta_c$) en het instructiemodel ($\theta_i$) tot een nieuw model ($\theta_m$) zonder extra training, gebruikmakend van twee interpolatiemethoden:

   • Lineaire Interpolatie (Linear): $\theta_m = (1 - \alpha)\theta_c + \alpha\theta_i$, waarbij $\alpha$ de mengverhouding bepaalt.
   • Sferische Lineaire Interpolatie (SLERP): Een methode die interpolatie uitvoert langs een hypersfeer in de gewichtenruimte, wat de grootte (norm) van de parametervectoren beter behoudt dan lineaire menging.

3. Zoeken naar Pareto-optimale checkpoints:
   Er werd een grid-search uitgevoerd over mengverhoudingen ($\alpha$ of $t$ van 0.0 tot 1.0). De selectie van de beste modellen gebeurde op basis van een Pareto-frontier: modellen die een optimale balans bieden tussen klinische kennis (gemeten via benchmarks zoals MMLU-medisch, MedQA, PubMedQA) en instructievolging (gemeten via IFEval).

4. Evaluatie en Fine-tuning:
   De geselecteerde merged-modellen werden geëvalueerd in een zero-shot setting en vervolgens onderworpen aan Supervised Fine-Tuning (SFT) op vijf klinische generatietaken (o.a. samenvatting van radiologieverslagen, ontslagrapporten, en probleemlijsten). Er werd getest op data-efficiëntie door training met beperkte datasets (few-shot: 8 tot 512 voorbeelden) vergeleken met volledige training.

Belangrijkste Bijdragen

• Oplossing voor Catastrophic Forgetting: Het bewijzen dat gewichtsruimte-merging (weight-space merging) een effectieve manier is om domeinspecifieke kennis te behouden terwijl het instructievolgingsvermogen van een algemeen model wordt behouden.
• Rekenkundige Efficiëntie: De methode vereist geen gradiënt-berekeningen of extra trainingstijd voor het creëren van het merged-model (alleen 40 minuten voor 9 checkpoints), in tegenstelling tot 18+ uur voor traditionele fine-tuning.
• Data-efficiëntie: Het aantonen dat merged-modellen aanzienlijk minder gelabelde data nodig hebben om vergelijkbare prestaties te bereiken als volledig getrainde basismodellen.

Resultaten

• Zero-shot Prestaties:

  • Het pure klinische model (GatorTronLlama) had sterke medische kennis maar slechte instructievolging (IFEval: 0.22).
  • Het pure instructiemodel (Llama-Instruct) had goede instructievolging maar lagere medische kennis.
  • De SLERP-gecombineerde modellen (bijv. met $t=0.4$) behaalden een verbetering van 1,43% in medische taken ten opzichte van het basismodel, terwijl ze 98,3% van het instructievolgingsvermogen behielden.
  • Er werd een niet-monotoon verband gevonden: een gematigde mengverhouding (rond 0,4–0,6) gaf de beste resultaten, wat suggereert dat een beperkte injectie van uitlijningsparameters de domeinexpertise kan versterken zonder deze te verdunnen.

• Fine-tuning en Data-efficiëntie:

  • Na SFT presteerden de SLERP-gecombineerde modellen consistent beter dan zowel het oorspronkelijke klinische basismodel als het klinisch getrainde SFT-basismodel op alle vijf de taken.
  • Few-shot learning: Op de IU-Xray dataset (radiologie) behaalde het SLERP-model getraind op slechts 64 voorbeelden dezelfde prestaties als het basismodel getraind op 256 voorbeelden.
  • Bij het genereren van probleemlijsten overtrof het SLERP-model (64 voorbeelden) zelfs het basismodel dat getraind was op het volledige dataset (512 voorbeelden).

Betekenis en Conclusie

De studie concludeert dat modelmerging een schaalbare en kosteneffectieve oplossing is voor het implementeren van open-source LLM's in de gezondheidszorg. Het overbrugt de kloof tussen medische expertise en het vermogen om instructies te volgen, wat essentieel is voor betrouwbare klinische toepassingen.

De belangrijkste implicaties zijn:

1. Verlaagde drempel: Ziekenhuizen en onderzoekers kunnen open-source modellen aanpassen zonder de enorme kosten van het verzamelen van duizenden medische instructiedata of het hebben van enorme GPU-resources voor RLHF/SFT.
2. Betrouwbaarheid: Het verminderen van "catastrophic forgetting" zorgt ervoor dat modellen niet alleen medisch accuraat zijn, maar ook veilig en bruikbaar in interactieve scenario's waar duidelijke instructievolging noodzakelijk is.
3. Toekomstperspectief: Hoewel de huidige resultaten veelbelovend zijn (gebaseerd op 8B modellen), is toekomstig onderzoek nodig om de effectiviteit te valideren op grotere modelgroottes en met klinische evaluatie door experts (in plaats van alleen automatische metrics) om hallucinaties en feitelijke juistheid te garanderen.