MedFormer-UR: Uncertainty-Routed Transformer for Medical Image Classification

1. Problématique

L'intégration clinique des modèles d'apprentissage profond (Deep Learning) en imagerie médicale se heurte à un défi majeur : la fiabilité et la confiance. Bien que les Vision Transformers (ViT), et en particulier l'architecture hiérarchique MedFormer, obtiennent des performances de pointe en classification et segmentation, ils souffrent de limitations critiques pour un usage médical :

• Manque de quantification de l'incertitude : Les modèles standards produisent des probabilités mal calibrées, souvent trop confiantes (overconfident), même sur des données ambiguës ou hors distribution.
• Boîte noire : L'absence de transparence dans le processus décisionnel rend difficile l'interprétation par les cliniciens.
• Nature des données médicales : Les données cliniques sont bruyantes, déséquilibrées et comportent des ambiguïtés intrinsèques (ex: tissus superposés, faible contraste), ce qui amplifie les risques d'erreurs de diagnostic.

L'objectif est donc de développer un cadre unifié capable de fournir non seulement une haute précision, mais aussi des estimations d'incertitude fiables et des justifications interprétables.

2. Méthodologie : MedFormer-UR

Les auteurs proposent MedFormer-UR, une extension de l'architecture MedFormer qui intègre trois mécanismes clés pour rendre le modèle "conscient de l'incertitude" et basé sur des prototypes.

A. Estimation d'incertitude évidentielle par token

Au lieu d'utiliser une simple sortie Softmax, le modèle utilise une distribution de Dirichlet pour modéliser l'incertitude au niveau de chaque "token" (patch d'image) dans le Transformer.

• Mécanisme : Une tête légère (MLP) prédit des valeurs de "preuves" (evidence) non négatives pour chaque classe. Ces preuves paramètrent une distribution de Dirichlet.
• Résultat : Cela permet de décomposer l'incertitude en aleatoire (bruit des données) et épistémique (manque de connaissances du modèle). L'incertitude est calculée comme l'inverse de la force totale des preuves.
• Avantage : Contrairement aux méthodes bayésiennes (comme Monte Carlo Dropout), cette approche ne nécessite pas d'échantillonnage coûteux lors de l'inférence.

B. Routage et raffinement guidés par l'incertitude (Uncertainty-Guided Routing)

L'incertitude n'est pas seulement une sortie, mais un signal de contrôle actif durant l'entraînement et l'inférence.

• Fonctionnement : Le modèle possède deux chemins parallèles dans les blocs Transformer : un chemin principal (attention DSSA) et un chemin de raffinement (réseau convolutif léger).
• Masque de routage : Un réseau prédit un masque spatial pour déterminer où appliquer le raffinement. Ce masque est modulé par l'incertitude estimée :
  • Si l'incertitude est élevée (zone ambiguë), le raffinement est supprimé ou atténué pour éviter d'amplifier des caractéristiques non fiables.
  • Si l'incertitude est faible, le raffinement se concentre sur les régions pertinentes (ex: contours de lésions).
• Objectif : Filtrer les mises à jour de caractéristiques peu fiables et stabiliser l'apprentissage.

C. Apprentissage basé sur les prototypes (Prototype-Based Learning)

Pour améliorer l'interprétabilité, la tête de classification classique est remplacée ou complétée par une approche basée sur des prototypes.

• Principe : Le modèle apprend un ensemble de prototypes (modèles de référence) spécifiques à chaque classe dans l'espace d'embedding.
• Décision : La prédiction est basée sur la similarité (cosinus) entre les caractéristiques de l'image et ces prototypes appris.
• Avantage : Cela structure l'espace d'embedding et permet aux cliniciens de comprendre la décision en comparant l'image à des exemples de référence visuels, rendant le modèle moins "boîte noire".

3. Contributions Clés

1. Extension de MedFormer : Intégration de têtes évidentielles produisant des paramètres de Dirichlet pour une quantification explicite de l'incertitude par token.
2. Mécanisme de routage dynamique : Utilisation de l'incertitude pour moduler la contribution des tokens, réduisant l'influence des caractéristiques ambiguës durant l'entraînement.
3. Intégration de prototypes : Combinaison de prototypes apprenables avec une perte de contraste pour structurer l'espace de décision et permettre un raisonnement par similarité visuelle.
4. Cadre unifié : Une architecture qui combine estimation d'incertitude, raffinement de caractéristiques et interprétabilité sans modifier le mécanisme d'attention de base (DSSA) du MedFormer.

4. Résultats Expérimentaux

Le modèle a été évalué sur quatre modalités d'imagerie médicale : Mammographie (CBIS-DDSM), Échographie (BUSI), IRM cérébrale et Histopathologie (IDC/BreaKHis).

• Performance de Classification :
  • Des améliorations significatives ont été observées sur l'échographie (BUSI) et l'histopathologie, avec des gains en précision et en AUROC.
  • Sur la mammographie (CBIS-DDSM), tâche particulièrement difficile, les gains en précision sont modestes mais significatifs compte tenu de la complexité des données (passage de 62,9% à 67,3% de précision).
• Calibration et Fiabilité (Résultat Principal) :
  • La réduction de l'Erreur de Calibration Attendue (ECE) est spectaculaire, allant jusqu'à 35% (de 0,1836 à 0,0688 sur CBIS-DDSM avec un paramètre de routage optimal $\beta=0,8$).
  • Cela signifie que les probabilités prédites correspondent beaucoup mieux à la réalité (ex: une prédiction de 80% de confiance est réellement correcte 80% du temps).
• Prédiction Sélective :
  • Le modèle excelle dans la capacité à rejeter les prédictions incertaines. À un taux de couverture de 70%, la précision est nettement supérieure à celle du modèle de base, démontrant une meilleure gestion des cas difficiles.
• Impact du paramètre $\beta$ (Force du routage) :
  • Les résultats montrent une relation non monotone. Un routage trop agressif ($\beta=1,0$) peut améliorer l'AUROC mais dégrader la calibration. Un routage modéré ($\beta \approx 0,7 - 0,8$) offre le meilleur équilibre entre précision et fiabilité.

5. Signification et Conclusion

Ce travail démontre que pour le déploiement clinique, la fiabilité (calibration) est aussi cruciale que la précision brute.

• Sécurité Clinique : En réduisant la surconfiance, MedFormer-UR permet aux médecins de mieux évaluer les risques et de prioriser les cas incertains pour un examen humain.
• Interprétabilité : La combinaison de cartes d'incertitude (montrant où le modèle hésite) et de similarité avec des prototypes (montrant sur quoi le modèle se base) offre un niveau de transparence inédit pour les Transformers médicaux.
• Gestion de l'ambiguïté : Le modèle ne cherche pas seulement à "forcer" une décision, mais à quantifier et localiser l'ambiguïté inhérente aux données médicales, ce qui est essentiel pour les applications à haut risque comme le dépistage du cancer du sein.

En résumé, MedFormer-UR représente une avancée vers une IA de confiance (Trustworthy AI) en transformant l'incertitude d'un défaut en un outil actif d'amélioration des performances et de la sécurité décisionnelle.