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

1. Problématique : L'Incohérence des Modèles VLM Médicaux

Les modèles de vision-langage (VLM) médicaux, tels que MedGemma-4B, souffrent d'un problème critique de sécurité : ils peuvent fournir des réponses différentes (ex. "Oui" vs "Non") à des reformulations sémantiquement équivalentes d'une même question clinique.

• Contexte : Un radiologue peut demander "Y a-t-il des signes de pneumothorax ?" ou "Ceci montre-t-il un pneumothorax ?". La réponse devrait être identique.
• Observation : Sur le jeu de données PSF-Med (basé sur MIMIC-CXR), le taux de "flip" (changement de réponse binaire) est de 14,6 % pour le modèle de base.
• Mesure d'instabilité : Même lorsque la réponse binaire reste la même, la marge de confiance (différence entre les logits "Oui" et "Non") varie considérablement (moyenne de 1,63 logits), ce qui indique une instabilité interne dangereuse pour la prise de décision clinique.

2. Méthodologie

L'approche proposée combine l'interprétabilité mécaniste et l'adaptation fine paramétrique efficace (LoRA).

A. Analyse Mécaniste et Interprétabilité

Les auteurs utilisent des Autoencodeurs Sparse (SAE) pré-entraînés sur le modèle de base Gemma (Gemma Scope 2) et valident leur transfert vers MedGemma-4B (avec un $R^2 \approx 0,997$).

• Construction de FlipBank : Un ensemble de 158 paires de questions où le modèle échoue systématiquement (réponses contradictoires).
• Identification de la cause : L'analyse des deltas de features révèle que la Feature 3818 (au niveau 17) est fortement activée par le "registre" de la question (présence vs exclusion) et non par le contenu clinique.
• Validation Causale (Activation Patching) : En modifiant artificiellement l'activation de la Feature 3818 lors du passage d'une paraphrase à l'originale, les auteurs récupèrent partiellement (28 %) la marge initiale, prouvant que cette feature influence causalement la décision.

B. Entraînement avec LoRA et Fonction de Perte Combinée

Pour corriger ce problème sans dégrader les performances, les auteurs entraînent des adaptateurs LoRA (Low-Rank Adaptation) sur les couches 15 à 19.

• Le défi de l'effondrement de mode (Mode Collapse) : Un entraînement purement axé sur la consistance (minimiser la divergence entre paraphrases) conduit le modèle à prédire toujours la même réponse (ex. toujours "Oui") pour minimiser l'erreur, détruisant ainsi sa capacité discriminative.
• Solution : Perte Combinée ($L$) : Ils introduisent une fonction de perte hybride :
  $$L = L_{consistency} + \lambda L_{accuracy}$$
  • $L_{consistency}$ : Divergence KL symétrique entre les distributions de probabilité des paraphrases.
  • $L_{accuracy}$ : Perte d'entropie croisée supervisée par la vérité terrain (Yes/No).
  • $\lambda = 1,0$ : Poids égal pour éviter l'effondrement de mode tout en forçant la cohérence.

3. Résultats Clés

Les expériences ont été menées sur MIMIC-CXR (n=158) et généralisées sur PadChest (n=250).

Sur MIMIC-CXR (Questions binaires) :

• Taux de Flip : Réduction de 14,6 % à 4,4 % (réduction relative de 69,6 %, $p=0,002$).
• Différence de Marge : Réduction de 1,63 à 0,33 logits (baisse de 79,5 %), indiquant une stabilité interne accrue.
• Précision (Accuracy) : Reste stable (84,2 % à 82,3 %), la différence n'étant pas statistiquement significative. Cela confirme que le modèle n'a pas perdu sa capacité de diagnostic.

Généralisation sur PadChest (Données externes) :

• Le taux de flip passe de 13,6 % à 7,8 %.
• La précision augmente même de 66,4 % à 69,4 %, suggérant que l'entraînement à la cohérence améliore la robustesse hors domaine.

Ablation des Couches :

Une étude sur l'emplacement des adaptateurs LoRA révèle une surprise :

• Bien que le mécanisme (Feature 3818) soit détecté au niveau 17, l'intervention sur les couches précoces (0 à 10) est la plus efficace pour réduire la différence de marge (86 % de réduction vs 80 % pour les couches 15-19).
• Hypothèse : Intervenir en amont empêche l'instabilité de se développer, plutôt que de tenter de la corriger après coup.

4. Contributions Principales

1. Caractérisation Systématique : Distinction claire entre le taux de flip (changement de réponse) et l'instabilité de la marge (changement de confiance) dans les VLM médicaux.
2. Transfert de SAE : Validation que les SAEs de Gemma Scope 2 fonctionnent sur un modèle médical fine-tuné (MedGemma).
3. Étude de Cas Mécaniste : Identification de la Feature 3818 comme étant sensible au registre linguistique (présence vs exclusion) et validation de son rôle causal.
4. Méthode d'Entraînement Innovante : Développement d'une fonction de perte combinée (Cohérence + Précision) qui évite l'effondrement de mode, un problème majeur des approches de consistance pure.

5. Signification et Implications

Ce travail démontre que l'interprétabilité mécaniste peut guider efficacement l'amélioration des modèles de sécurité en IA médicale.

• Sécurité Clinique : Réduire l'incohérence des paraphrases est essentiel pour la confiance des cliniciens et le déploiement sécurisé.
• Efficacité : L'utilisation de LoRA permet d'obtenir ces gains avec une très faible augmentation de paramètres (0,10 % du modèle total), rendant la solution économiquement viable.
• Paradigme d'Intervention : L'article souligne que le lieu où un mécanisme est observé (niveau 17) n'est pas nécessairement le lieu optimal pour l'intervention (couches 0-10), offrant un nouvel angle pour la recherche en débogage de modèles.

En résumé, cette étude propose une méthode robuste pour rendre les assistants IA médicaux plus fiables face aux variations linguistiques, sans sacrifier leur précision diagnostique.