Uncertainty Estimation in Pathology Foundation Models via Deep Mutual Learning

Résumé Technique : Estimation de l'incertitude dans les modèles de fondation en pathologie via l'apprentissage mutuel profond

Énoncé du problème
Les modèles de fondation en pathologie (PFM) ont démontré de fortes capacités de généralisation pour l'extraction de représentations à partir d'images de lames entières (WSI) pour des tâches telles que l'analyse de la survie, la découverte de biomarqueurs et la sous-typologie tumorale. Cependant, leur traduction dans la pratique clinique de routine est entravée par deux limitations critiques :

1. Absence d'estimations de confiance : Les pipelines de PFM existants produisent généralement des prédictions sans mesures explicites de confiance ou d'incertitude. Dans les contextes cliniques, les modèles doivent non seulement être précis, mais aussi être capables d'indiquer quand leurs prédictions sont incertaines afin de soutenir une prise de décision sécurisée.
2. Performance dépendante de la tâche : Aucun PFM unique n'est universellement supérieur à travers toutes les tâches, tous les jeux de données et tous les contextes cliniques. La sélection du modèle optimal pour le déploiement est difficile, et s'appuyer sur un seul modèle risque d'entraîner des performances sous-optimales ou une surconfiance dans des cas sujets à l'échec.

Les méthodes de fusion actuelles améliorent souvent les performances mais ne fournissent pas d'estimations d'incertitude explicites et agnostiques à la tâche. Inversement, les méthodes de quantification de l'incertitude comme les ensembles profonds (deep ensembles) ou le dropout de Monte Carlo ne sont pas spécifiquement conçues pour les ensembles de PFM ou sont spécifiques à une tâche (par exemple, limitées au diagnostic du glioblastome).

Méthodologie : Cadre DICE
Les auteurs proposent DICE (Disagreement-Informed Coordination of Experts - Coordination d'experts informée par le désaccord), un cadre « plug-and-play » qui constitue un ensemble de $K$ PFM gelés pour estimer l'incertitude et améliorer les performances en aval. L'intuition centrale est que le désaccord entre des PFM hétérogènes sert de signal informatif pour l'incertitude, à condition que ce désaccord soit structuré plutôt que dominé par le bruit d'optimisation.

Le cadre fonctionne comme suit :

1. Architecture : Le système utilise $K$ PFM gelés (par exemple, Virchow2, UNI2-h, H-optimus-1, CONCHv1.5, Hibou-L). Chaque PFM agit comme un « expert » $E^{(k)}$, associé à un agrégateur d'apprentissage multi-instances (MIL) indépendant (spécifiquement ABMIL) et une tête de classification.

2. Objectifs d'entraînement : Pour garantir que le désaccord de l'ensemble reflète une incertitude significative plutôt que du bruit, les experts sont couplés pendant l'entraînement via trois objectifs :

   • Perte supervisée ($L_{sup}$) : Perte d'entropie croisée standard pour assurer une classification précise au niveau de la lame.
   • Apprentissage mutuel profond ($L_{DML}$) : Un objectif d'alignement dans l'espace des prédictions où chaque expert minimise la divergence de Kullback–Leibler (KL) entre sa distribution prédictive et celles de ses $K-1$ pairs. Cela encourage le consensus tout en maintenant l'alignement avec la vérité terrain.
   • Régularisation de Gram ($L_{Gram}$) : Un objectif d'alignement dans l'espace des représentations basé sur la mesure de Gram. Il encourage les plongements (embeddings) de la lame des experts à couvrir un sous-espace de faible volume, empêissant un désalignement géométrique qui pourrait être confondu avec de l'incertitude.
   • Perte totale : L'objectif d'entraînement combine ces termes : $L = L_{sup} + L_{DML} + \lambda_{Gram}(t) L_{Gram}$, où $\lambda_{Gram}$ suit un programme de montée en puissance (warm-up).

3. Inférence et estimation de l'incertitude :

   • Proxy d'incertitude : À l'inférence, la prédiction de l'ensemble est la moyenne des distributions des $K$ experts. L'incertitude prédictive totale est décomposée en composantes aléatoires (données) et épistémiques (modèle). Les auteurs prouvent théoriquement que l'incertitude épistémique (le désaccord résiduel) est bornée supérieurement par l'agrégat de la perte DML. Ainsi, minimiser la perte DML pendant l'entraînement contracte la borne supérieure de l'incertitude du modèle.
   • Localisation : La localisation des lésions au niveau des patchs est réalisée de manière non supervisée en calculant le consensus des poids d'attention à travers tous les experts.

Contributions clés

1. Premier cadre agnostique à la tâche : DICE est le premier cadre qui exploite les ensembles de PFM spécifiquement pour l'estimation de l'incertitude agnostique à la tâche en pathologie computationnelle.
2. Fondement théorique : L'article fournit une preuve théorique que la minimisation de l'objectif d'alignement DML lors de l'entraînement contracte une borne supérieure sur l'incertitude épistémique de l'ensemble, validant l'utilisation du désaccord résiduel comme un proxy d'incertitude fiable.
3. Performance empirique : Des expériences approfondies démontrent que DICE non seulement signale efficacement les échantillons sujets à l'échec (en distribution et hors distribution), mais égale ou surpasse les PFM uniques de pointe et les bases de fusion en termes de classification, de calibration et de localisation non supervisée des lésions.

Résultats expérimentaux
Le cadre a été évalué sur trois benchmarks de WSI sous supervision faible : PANDA (graduation du cancer de la prostate), CAMELYON16 et CAMELYON17 (détection de métastases ganglionnaires mammaires).

• Classification et Calibration : DICE a obtenu les scores F1 les plus élevés et les plus faibles NLL (Negative Log-Likelihood) sur la plupart des jeux de données. Sur PANDA, il a amélioré le score F1 de 3,8 points de pourcentage par rapport au meilleur PFM unique et de 2,8 points par rapport à la meilleure base de fusion.
• Estimation de l'incertitude : Dans les expériences de prédiction sélective, DICE a efficacement différé les lames à haute incertitude, résultant en les taux d'erreur retenus les plus bas. Il a démontré une séparation supérieure entre les prédictions correctes et incorrectes (mesurée par le $\Delta$ de Cliff) par rapport au Dropout de Monte Carlo et aux autres méthodes de fusion. Cette capacité s'est maintenue même sous des changements de cohorte (par exemple, entraînement sur CAMELYON16 et test sur CAMELYON17).
• Localisation : En agrégeant les cartes d'attention, DICE a fourni une localisation par patch non supervisée qui était compétitive ou supérieure aux modèles uniques, particulièrement en concentrant l'attention dans les régions tumorales annotées plutôt que dans le stroma environnant.
• Études d'ablation : Le cadre s'est montré robuste au choix de l'agrégateur MIL (DSMIL, TransMIL) et a montré que l'augmentation du nombre d'experts produit des rendements décroissants une fois que l'espace de caractéristiques est suffisamment couvert, soulignant l'importance de la complémentarité des experts plutôt que de la simple taille de l'ensemble.

Signification et revendications
L'article positionne DICE comme une étape cruciale vers la transition des PFM vers des systèmes d'aide à la décision sensibles à l'incertitude. En établissant une approche théoriquement fondée pour fusionner les PFM, DICE répond à la nature de « boîte noire » des modèles actuels, offrant un mécanisme pour identifier les cas sujets à l'erreur et différer en toute sécurité vers les cliniciens. Les auteurs affirment que leur méthode rend les PFM plus dignes de confiance dans les contextes cliniques en fournissant des estimations d'incertitude fiables sans nécessiter de supervision explicite pour la localisation ou la calibration de la confiance.

Les auteurs reconnaissent des limites, notamment le coût computationnel de l'extraction de caractéristiques à partir de $K$ modèles gelés (bien que les coûts d'entraînement restent inférieurs à un réglage fin (fine-tuning) de bout en bout) et la nécessité de balancer soigneusement l'alignement et la diversité pour éviter de supprimer le désaccord même requis pour l'estimation de l'incertitude. Ils concluent que, bien que le nombre d'experts soit actuellement fixé, le cadre ouvre la voie à des modèles de pathologie plus fiables et conscients de l'incertitude, intégrables dans les flux de travail cliniques.