Uncertainty Estimation in Pathology Foundation Models via Deep Mutual Learning

Technische Samenvatting: Onzekerheidsschatting in Pathologie Fundatiemodellen via Deep Mutual Learning

Probleemstelling
Pathologie Fundatiemodellen (PFM's) hebben sterke generalisatiecapaciteiten aangetoond bij het extraheren van representaties uit whole-slide afbeeldingen (WSI's) voor taken zoals overlevingsanalyse, biomarkerontdekking en tumorsubtypering. De vertaling naar de dagelijkse klinische praktijk wordt echter gehinderd door twee kritieke beperkingen:

1. Gebrek aan vertrouwensschattingen: Bestaande PFM-pipelines leveren doorgaans voorspellingen zonder expliciete mate van vertrouwen of onzekerheid. In een klinische setting moeten modellen niet alleen accuraat zijn, maar ook in staat zijn om aan te geven wanneer hun voorspellingen onzeker zijn om veilige besluitvorming te ondersteunen.
2. Taakafhankelijke prestaties: Geen enkel PFM is universeel superieur over alle downstream-taken, datasets en klinische contexten heen. Het selecteren van het optimale model voor implementatie is uitdagend, en het vertrouwen op een enkel model brengt het risico met zich mee op suboptimale prestaties of overmatige zelfverzekerdheid in gevallen die gevoelig zijn voor fouten.

Huidige fusiemethoden verbeteren vaak de prestaties, maar falen in het bieden van expliciete, taak-agnostische onzekerheidsschattingen. Omgekeerd zijn onzekerheidskwantificatiemethoden zoals deep ensembles of Monte Carlo dropout niet specifiek ontworpen voor PFM-ensembles of zijn ze taakspecifiek (bijv. beperkt tot glioblastoomdiagnose).

Methodologie: DICE Framework
De auteurs stellen DICE voor (Disagreement-Informed Coordination of Experts), een plug-and-play framework dat $K$ bevroren PFMs ensembleert om onzekerheid te schatten en downstream-prestaties te verbeteren. De kernintuïtie is dat de onenigheid tussen heterogene PFMs een informatieve signaal voor onzekerheid vormt, mits deze onenigheid gestructureerd is en niet wordt gedomineerd door optimalisatieruis.

Het framework werkt als volgt:

1. Architectuur: Het systeem maakt gebruik van $K$ bevroren PFMs (bijv. Virchow2, UNI2-h, H-optimus-1, CONCHv1.5, Hibou-L). Elk PFM fungeert als een "expert" $E^{(k)}$, gekoppeld aan een onafhankelijke Multiple Instance Learning (MIL) aggregator (specifiek ABMIL) en een classificatiekop.

2. Trainingsdoelstellingen: Om ervoor te zorgen dat de onenigheid van het ensemble betekenisvolle onzekerheid reflecteert in plaats van ruis, worden de experts tijdens de training gekoppeld via drie doelstellingen:

   • Supervised Loss ($L_{sup}$): Standaard cross-entropy loss om nauwkeurige slide-niveau classificatie te garanderen.
   • Deep Mutual Learning ($L_{DML}$): Een doelstelling voor alignment in de voorspellingsruimte waarbij elke expert de Kullback–Leibler divergentie tussen zijn eigen voorspellende distributie en die van zijn $K-1$ gelijken minimaliseert. Dit stimuleert consensus terwijl de uitlijning met de grondwaarheid behouden blijft.
   • Gramian Regularization ($L_{Gram}$): Een doelstelling voor alignment in de representatieruimte gebaseerd op de Gramiaanse maat. Het moedigt de slide-niveau embeddings van de experts aan om een subspace met een laag volume te beslaan, wat geometrische misuitlijning voorkomt die zich als onzekerheid zou kunnen voordoen.
   • Totale Loss: De trainingsdoelstelling combineert deze termen: $L = L_{sup} + L_{DML} + \lambda_{Gram}(t) L_{Gram}$, waarbij $\lambda_{Gram}$ een warm-up schema volgt.

3. Inference en Onzekerheidsschatting:

   • Onzekerheidsproxie: Tijdens de inference is de ensemblevoorspelling het gemiddelde van de $K$ expert-distributies. De totale voorspellende onzekerheid wordt ontleed in aleatorische (data) en epistemische (model) componenten. De auteurs bewijzen theoretisch dat de epistemische onzekerheid (residuele onenigheid) begrensd wordt door de geaggregeerde DML-loss. Daarom verkleint het minimaliseren van de DML-loss tijdens de training de bovengrens op de modelonzekerheid.
   • Lokalisatie: Patch-niveau laesielokalisatie wordt ongesuperviseerd bereikt door de consensus van aandachtsgewichten (attention weights) over alle experts te berekenen.

Belangrijkste Bijdragen

1. Eerste Taak-Agnostisch Framework: DICE is het eerste framework dat PFM-ensembles specifiek inzet voor taak-agnostische onzekerheidsschatting in computationele pathologie.
2. Theoretische Fundering: Het artikel biedt een theoretisch bewijs dat het minimaliseren van de trainingstijd DML-alignment doelstelling de bovengrens op de epistemische onzekerheid van het ensemble inkrimpt, wat het gebruik van residuele onenigheid als een betrouwbare onzekerheidsproxie valideert.
3. Empirische Prestaties: Uitgebreide experimenten tonen aan dat DICE niet alleen foutgevoelige samples (zowel in- als out-of-distribution) accuraat signaleert, maar ook de state-of-the-art (SOTA) single PFMs en fusie-baselines evenaart of overtreft in classificatie, kalibratie en ongesuperviseerde laesielokalisatie.

Experimentele Resultaten
Het framework werd geëvalueerd op drie zwak gesuperviseerde WSI-benchmarks: PANDA (prostaatcarceroom gradering), CAMELYON16 en CAMELYON17 (borst lymfeklier metastase detectie).

• Classificatie & Kalibratie: DICE behaalde de hoogste F1-scores en de laagste Negative Log-Likelihood (NLL) over de meeste datasets. Op PANDA verbeterde het de F1 met 3,8 procentpunt ten opzichte van de beste single PFM en met 2,8 punten ten opzichte van de beste fusie-baseline.
• Onzekerheidsschatting: In experimenten met selectieve voorspelling slaagde DICE erin om slides met een hoge onzekerheid uit te stellen, wat resulteerde in de laagste retentie van fouten. Het vertoonde een superieure scheiding tussen correcte en incorrecte voorspellingen (gemeten met de Cliff's $\Delta$) vergeleken met Monte Carlo Dropout en andere fusiemethoden. Deze capaciteit bleef behouden zelfs onder cohortverschuivingen (bijv. trainen op CAMELYON16 en testen op CAMELYON17).
• Lokalisatie: Door aandachtskaarten (attention maps) te aggregeren, bood DICE ongesuperviseerde patch-niveau lokalisatie die competitief of superieur was aan single modellen, met name door de aandacht te concentreren binnen geannoteerde tumorregio's in plaats van het omliggende stromaal weefsel.
• Ablatie-studies: Het framework bleek robuust tegen de keuze van de MIL-aggregator (DSMIL, TransMIL) en toonde aan dat het verhogen van het aantal experts een afnemend rendement oplevert zodra de feature-ruimte voldoende is gedekt, wat het belang van de complementariteit van experts benadrukt boven de loutere omvang van het ensemble.

Betekenis en Claims
Het artikel positioneert DICE als een cruciale stap naar de vertaling van PFMs naar onzekerheidsbewuste beslissingsondersteunende systemen. Door een theoretisch gefundeerde aanpak te bieden voor het fuseren van PFMs, adresseert DICE de "black box"-natuur van huidige modellen, door een mechanisme te bieden om foutgevoelige gevallen te identificeren en veilig over te dragen aan clinici. De auteurs beweren dat hun methode PFMs betrouwbaarder maakt in klinische settings door betrouwbare onzekerheidsschattingen te bieden zonder dat expliciete supervisie voor lokalisatie of vertrouwen-kalibratie vereist is.

De auteurs erkennen beperkingen, waaronder de computationele kosten voor het extraheren van features uit $K$ bevroren modellen (hoewel de trainingskosten lager blijven dan end-to-end fine-tuning) en de noodzaak om alignment en diversiteit zorgvuldig in balans te houden om te voorkomen dat de onenigheid die nodig is voor onzekerheidsschatting wordt onderdrukt. Ze concluderen dat hoewel het aantal experts momenteel vaststaat, het framework de weg vrijmaakt voor betrouwbaardere, onzekerheidsbewuste pathologiemodellen geïntegreerd in klinische workflows.