Average Calibration Losses for Reliable Uncertainty in Medical Image Segmentation

Probleemstelling

Diepe neurale netwerken (DNN's) voor medische beeldsegmentatie vertonen vaak oververtrouwen (overconfidence). Dit betekent dat de voorspelde kansen (confidentie) niet overeenkomen met de werkelijke waarschijnlijkheid van correctheid (nauwkeurigheid). In veiligheidskritieke domeinen zoals medische diagnose en stralingsplanning kan dit leiden tot ernstige risico's, waarbij artsen worden misleid door hoge confidenties bij verkeerde beslissingen.

Bestaande oplossingen zijn vaak post-hoc (na het trainen), zoals Temperature Scaling. Deze methoden hebben beperkingen: ze benutten niet de volledige capaciteit van het model om kalibratie direct te leren, en ze passen vaak globale correcties toe die geen rekening houden met instance- of class-specifieke fouten. Bovendien zijn traditionele kalibratiemetrics zoals de Expected Calibration Error (ECE) vaak gevoelig voor de verdeling van steekproeven in bins en negeren ze minder frequente, maar klinisch kritieke, onzekere voorspellingen (zoals weefselgrenzen).

Methodologie

De auteurs stellen een nieuwe aanpak voor waarbij kalibratie direct in het trainingsproces wordt geïntegreerd (train-time calibration) door een differentieerbare mL1-ACE (marginal L1 Average Calibration Error) als auxiliaire loss-functie toe te voegen.

Kerncomponenten van de methode:

1. mL1-ACE Loss: In plaats van ECE, gebruiken de auteurs ACE. ECE weegt bins af op basis van het aantal samples (wat leidt tot dominantie van zeer zeker voorspellingen), terwijl ACE elke bin gelijk weegt. Dit is cruciaal voor medische segmentatie waar onzekere voorspellingen vaak voorkomen in kritieke gebieden.
2. Per-Image Berekening: Omdat segmentatie per pixel/voxel voorspellingen genereert, kunnen betrouwbaarheidsdiagrammen en kalibratiemetrics per afbeelding worden berekend, in plaats van alleen op dataset-niveau.
3. Binning Strategieën: De auteurs vergelijken twee varianten voor het discretiseren van de voorspelde kansen in bins:
   • Hard-binning (hL1-ACE): Gebruikt een vierkante kernel (square wave). Elke voxel behoort tot precies één bin.
   • Soft-binning (sL1-ACE): Gebruikt een driehoekige kernel (triangular kernel). Voxels kunnen bijdragen aan meerdere aangrenzende bins, wat zorgt voor soepelere gradiënten.
4. Dataset Reliability Histograms: Een nieuwe visualisatietechniek die aggregatie van per-image betrouwbaarheidsdiagrammen toont. Dit geeft inzicht in de variabiliteit van de kalibratie over de hele dataset.
5. Experimenteel Opzet: De methode wordt getest op vier datasets (ACDC, AMOS, BraTS, KiTS) met een SegResNet architectuur. De baselines zijn DSC + Cross-Entropy (CE), CE alleen, en post-hoc Temperature Scaling. De loss-functie is een combinatie van DSC, CE en de mL1-ACE term (met gewicht $\lambda$).

Belangrijkste Bijdragen

• Differentieerbare mL1-ACE Loss: De eerste toepassing van een differentieerbare ACE-loss specifiek voor semantische segmentatie, die per-afbeelding berekeningen mogelijk maakt zonder proxy-losses of zachte relaxaties (hoewel soft-binning wel wordt gebruikt voor gradiënten).
• Dataset Reliability Histograms: Introductie van een nieuwe visualisatiemethode om kalibratietrends en variabiliteit binnen een dataset beter te begrijpen.
• Uitgebreide Evaluatie: Een grondige vergelijking van hard- versus soft-binning op vier diverse medische datasets, inclusief gevoeligheidsanalyses voor het aantal bins en loss-gewichten.
• Controle over de Trade-off: Het bieden van een expliciete manier voor praktici om de afweging tussen segmentienauwkeurigheid (DSC) en kalibratie te sturen.

Resultaten

De experimenten tonen aan dat het toevoegen van mL1-ACE de kalibratiefouten aanzienlijk verlaagt ten opzichte van de DSC+CE baseline, terwijl de segmentienauwkeurigheid grotendeels behouden blijft.

• Kalibratieverbetering:
  • Soft-binned (sL1-ACE): Bereikt de grootste verbetering in kalibratie (gemiddeld 33% reductie in macro-ACE en 51% in micro-ACE). Het reduceert oververtrouwen effectiever dan post-hoc Temperature Scaling.
  • Hard-binned (hL1-ACE): Levert een solide verbetering (gemiddeld 16% reductie in macro-ACE) en presteert vergelijkbaar met Temperature Scaling.
• Segmentienauwkeurigheid (DSC):
  • Hard-binned: Behoudt de DSC-score bijna exact op het niveau van de baseline (gemiddeld +0.03% verandering).
  • Soft-binned: Leidt tot een kleine, maar statistisch significante daling in DSC (gemiddeld -0.7%), wat de afweging tussen kalibratie en nauwkeurigheid illustreert.
• Vergelijking met Baselines: Beide varianten presteren beter in kalibratie dan alleen CE-training (die wel beter gekalibreerd is dan DSC+CE, maar veel slechtere DSC-scores heeft) en beter dan post-hoc Temperature Scaling in termen van maximale kalibratieverbetering.
• Sensitiviteit: Hard-binning is robuust tegen variaties in het aantal bins en loss-gewichten. Soft-binning is gevoeliger; hogere gewichten en meer bins leiden tot sterkere kalibratie maar grotere DSC-dalingen.

Betekenis en Conclusie

Deze studie toont aan dat het direct optimaliseren voor kalibratie tijdens het trainen van medische beeldsegmentatiemodellen haalbaar en effectief is.

• Klinische Relevantie: Door oververtrouwen te verminderen, kunnen artsen beter vertrouwen op de AI-voorspellingen en onzekere gebieden (zoals tumorgrenzen) beter identificeren voor menselijke controle.
• Flexibiliteit: De methode biedt een "knop" voor ontwikkelaars. Als segmentienauwkeurigheid het allerbelangrijkst is, kan hard-binning worden gekozen. Als betrouwbare onzekerheidsschattingen cruciaal zijn (bijv. voor actieve learning of veiligheid), kan soft-binning worden gebruikt, mits een kleine daling in DSC acceptabel is.
• Metrische Vooruitgang: De auteurs pleiten voor het gebruik van ACE in plaats van ECE voor medische segmentatie, omdat ACE beter presteert in het detecteren van fouten in kritieke, minder frequente bins.

Samenvattend biedt deze work een praktische en theoretisch onderbouwde oplossing om diepe neurale netwerken in de medische beeldvorming betrouwbaarder en veiliger te maken voor klinische integratie.