Mechanistically Guided LoRA Improves Paraphrase Consistency in Medical Vision-Language Models

Probleemstelling

Medische Vision-Language Models (VLM's), zoals MedGemma-4B, vertonen een kritieke onbetrouwbaarheid: ze kunnen verschillende antwoorden geven op semantisch equivalente herschrijvingen (parafases) van dezelfde klinische vraag. In de medische praktijk is dit gevaarlijk, omdat verschillende artsen dezelfde vraag op verschillende manieren kunnen formuleren (bijv. "Is er bewijs voor een pneumothorax?" versus "Toont dit een pneumothorax?"), maar een consistent antwoord verwachten.

De auteurs bestuderen dit fenomeen met behulp van de benchmark PSF-Med. Op een dataset van 158 binaire vragen (Ja/Nee) afkomstig van MIMIC-CXR werd een flip-rate (het percentage keren dat het antwoord verandert bij herschrijving) van 14,6% vastgesteld. Daarnaast was de gemiddelde marge-differentie (de absolute verandering in log-odds tussen Ja en Nee) 1,63 logits, wat aangeeft dat de interne representaties van het model gevoelig zijn voor oppervlakkige taalkundige variaties die geen invloed zouden moeten hebben op klinische beslissingen.

Methodologie

De aanpak van de auteurs combineert mechanistische interpretabiliteit met parameter-efficiënt fine-tuning.

1. Mechanistische Analyse met Sparse Autoencoders (SAE's)

• Validatie: De auteurs valideren eerst of Gemma Scope 2 Sparse Autoencoders (oorspronkelijk getraind op basis-Gemma) effectief kunnen worden overgebracht naar de fijngefineerde MedGemma-4B. Ze vinden een reconstructie-kwaliteit ($R^2$) van ongeveer 0,997, wat aantoont dat de SAE's geldig blijven voor het medische model.
• FlipBank: Om de oorzaak van de inconsistentie te onderzoeken, construeren ze "FlipBank", een gecureerde dataset van 158 gevallen waarin het model duidelijk van antwoord verandert bij een parafase.
• Feature Detectie: Door delta-analyse toe te passen op de activaties, identificeren ze Feature 3818 in laag 17 als een kandidaat-mechanisme. Deze feature reageert sterk op het "register" van de vraag (bijv. aanwezigheid vs. uitsluiting van een bevinding) en niet op oppervlakkige formaliteit.
• Causale Validatie: Via activatie-patching (het manipuleren van de feature-activatie tijdens de forward pass) tonen ze aan dat Feature 3818 een causale invloed heeft op de beslissing. Het patchen van deze feature herstelde in een voorbeeldgeval 28% van de marge-verschuiving, wat significant meer was dan bij controle-features.

2. LoRA Fine-Tuning met Gecombineerde Loss
Om het probleem op te lossen, trainen de auteurs Low-Rank Adaptation (LoRA) adapters.

• Architectuur: LoRA wordt toegepast op lagen 15 tot 19 (geïnspireerd door de locatie van Feature 3818), zowel op de attention- als MLP-modules. De visuele encoder blijft bevroren.
• Het Gecombineerde Loss-functie: Een puur consistentieloss (die probeert de uitkomsten voor parafases gelijk te maken) leidt tot mode collapse: het model leert om voor elke vraag "Ja" te zeggen om de divergentie te minimaliseren, wat de nauwkeurigheid vernietigt.
  • Om dit te voorkomen, gebruiken ze een loss-functie die consistentie (symmetrische KL-divergentie tussen de distributies van originele en herschreven vragen) combineert met nauwkeurigheid (cross-entropy loss tegen de ground truth labels).
  • Formule: $L = L_{consistency} + \lambda L_{accuracy}$ (waarbij $\lambda = 1,0$).

Kernresultaten

De resultaten tonen aan dat de methode de consistentie drastisch verbetert zonder de prestaties op de taak zelf te verstoren.

1. Prestaties op MIMIC-CXR (PSF-Med, n=158)

• Flip-rate: Daalt van 14,6% naar 4,4% (een reductie van 69,6%, $p=0,002$).
• Marge-differentie: Daalt van 1,63 naar 0,33 (een reductie van 79,5%).
• Nauwkeurigheid: Blijft stabiel, met een kleine, niet-significante daling van 84,2% naar 82,3%.

2. Generalisatie naar PadChest (n=250)

• Zelfs getraind op MIMIC-CXR, verbetert het model op het onafhankelijke PadChest-dataset (Spaanse borstfoto's).
• Flip-rate daalt van 13,6% naar 7,8%.
• Marge-differentie daalt van 1,08 naar 0,35.
• Opvallend genoeg stijgt de nauwkeurigheid hier van 66,4% naar 69,4%.

3. Ablatie-onderzoek (Laagselectie)

• Hoewel Feature 3818 in laag 17 werd gevonden, bleek dat het toepassen van LoRA in vroege lagen (0-10) de beste resultaten opleverde voor marge-reductie (86% verbetering), beter dan de mechanistisch geselecteerde middenlagen (15-19, 80% verbetering).
• Dit suggereert dat het onderdrukken van gevoeligheid in vroege verwerkingstappen effectiever is dan het corrigeren ervan nadat het mechanisme zich heeft gemanifesteerd.

Belangrijkste Bijdragen

1. Systematische Karakterisering: Een gedetailleerde analyse van parafase-gevoeligheid in MedGemma-4B, waarbij onderscheid wordt gemaakt tussen flip-rate en marge-instabiliteit.
2. Validatie van SAE-overdracht: Het bewijs dat SAE's van de basis-Gemma-modelfamilie effectief kunnen worden gebruikt voor geanalyseerde medische VLM's.
3. Mechanistisch Casestudy: Identificatie van Feature 3818 als een register-gevoelige component die causaal bijdraagt aan inconsistente beslissingen.
4. Nieuwe Trainingstrategie: Een gecombineerde loss-functie voor LoRA-training die mode collapse voorkomt en tegelijkertijd consistentie en nauwkeurigheid waarborgt.

Significantie en Conclusie

Dit werk is significant omdat het een praktische oplossing biedt voor een veiligheidsrisico in medische AI: de onbetrouwbaarheid bij variaties in vraagstelling. Door mechanistische interpretatie te combineren met gerichte fine-tuning, tonen de auteurs aan dat het mogelijk is om de "gedragstabiliteit" van een model te verbeteren zonder de klinische bruikbaarheid (nauwkeurigheid) te offeren.

De bevinding dat vroege lagen effectiever zijn voor interventie dan de lagen waar het probleem mechanistisch zichtbaar is, biedt een belangrijke inzichten voor toekomstig onderzoek naar robuustheid in grote taalmodellen: het is vaak effectiever om gevoeligheid bij de bron te onderdrukken dan om de gevolgen later te corrigeren. De code en methodologie zijn beschikbaar voor replicatie en verdere ontwikkeling in de medische AI-veiligheid.