An Uncertainty-Aware Loss Function Incorporating Fuzzy Logic: Application to MRI Brain Image Segmentation

Probleemstelling

Het segmenteren van hersenweefsel uit magnetische resonantiebeeldvorming (MRI) is cruciaal voor de diagnose en behandeling van neurologische aandoeningen. Hoewel diepe leermodellen, zoals U-Net en U-Net++, uitstekende prestaties leveren, kampen ze met een fundamenteel probleem: onzekerheid.

• Aard van de onzekerheid: Medische beelden bevatten inherent onzekerheid door factoren zoals ruis, de complexe aard van weefsels, en het "partial volume effect" (waarbij één pixel meerdere weefsels bevat).
• Beperking van bestaande methoden: Traditionele loss-functies, zoals de Categorical Cross-Entropy (CCE), opereren op basis van klassieke verzamelingstheorie (crisp classifications). Ze behandelen pixels als strikt tot één klasse behorend (waar/onwaar), wat niet rekening houdt met de vaagheid en de gradaties van toewijzing die vaak voorkomen in medische beeldanalyse. Dit kan leiden tot minder robuuste modellen en onnauwkeurige segmentaties aan de randen van weefsels.

Methodologie

De auteurs stellen een nieuwe loss-functie voor: de Fuzzy Categorical Cross-Entropy (FCCE). Deze functie combineert de standaard Categorical Cross-Entropy met Fuzzy Entropy, gebaseerd op fuzzy-logica, om onzekerheid expliciet te modelleren tijdens het trainingsproces.

De kerncomponenten van de methode zijn:

1. Fuzzy C-Means (FCM) Algorithm:
   • Om de "lidmaatschapsgraad" (degree of membership) te berekenen, wordt het FCM-algoritme gebruikt. In tegenstelling tot harde clustering, geeft FCM elke pixel een lidmaatschapsgraad ($\mu_{ij}$) tussen 0 en 1 voor elke klasse.
   • Dit creëert een zachte classificatie die de onzekerheid in de data weerspiegelt.
2. Fuzzy Entropy ($L_{FE}$):
   • Afgeleid van de Shannon-entropy, meet deze term de onzekerheid of "vaagheid" binnen de fuzzy verzamelingen.
   • Formule: $L_{FE} = -\sum \sum u_{ij} \log(u_{ij})$. Een hogere entropy wijst op meer onzekerheid.
3. De Gecombineerde FCCE Loss-functie:
   • De totale loss-functie is de som van de Categorical Cross-Entropy ($L_{CCE}$) en de Fuzzy Entropy ($L_{FE}$):
     $$L_{FCCE} = L_{CCE} + L_{FE}$$
   • Werking: Tijdens het trainen berekent het model niet alleen de fout tussen de voorspelling en de ground truth (via $L_{CCE}$), maar straft het ook de onzekerheid in de lidmaatschapsgraden (via $L_{FE}$). Hierdoor leert het model om onzekerheid te minimaliseren en robuustere voorspellingen te doen, vooral bij pixels met vaagke grenzen.
4. Architectuur:
   • De loss-functie is geïmplementeerd in twee populaire segmentatie-architecturen: U-Net en U-Net++.
   • U-Net++ wordt gebruikt vanwege de verbeterde skip-connections en diepe supervisie, wat de prestaties verder kan optimaliseren.

Belangrijkste Bijdragen

• Novel Loss-functie: Introductie van FCCE, die fuzzy-logica-integratie mogelijk maakt in het trainingsproces van diepe neurale netwerken voor medische beeldsegmentatie.
• Onzekerheidsbeheer: Het vermogen om de inherente onzekerheid in MRI-pixels (door ruis en partial volume effect) expliciet te modelleren en te minimaliseren, in plaats van deze te negeren.
• Hybride Aanpak: Effectieve combinatie van toezicht (gelabelde data via CCE) en ontoezicht (ongelabelde data/structuur via FCM en Fuzzy Entropy) binnen één loss-functie.
• Robuustheid: Het verbeteren van de modelrobustheid en segmentatie-accuraatheid, met name in complexe gebieden met vaagke grenzen.

Resultaten

De methode is geëvalueerd op twee publieke datasets: IBSR (Internet Brain Segmentation Repository) en OASIS (Open Access Series of Imaging Studies). De prestaties zijn vergeleken met modellen die alleen de standaard CCE loss-functie gebruiken.

• Kwantitatieve Resultaten:
  • IBSR Dataset: Zowel U-Net als U-Net++ presteerden beter met FCCE dan met CCE.
    • U-Net: Verbetering in Dice Coëfficiënt (DC) van 0.9980 naar 0.9986 en IoU van 0.9960 naar 0.9972.
    • U-Net++: Toonde de grootste verbetering in validatiemetrics (bijv. DCval steeg van 0.9717 naar 0.9881).
  • OASIS Dataset: Ook hier overtrof FCCE de CCE in de meeste training-metrics (Accuracy, DC, IoU).
    • Opmerking: Bij U-Net op de OASIS dataset waren de validatiemetrics iets lager dan bij CCE, wat suggereert dat overfitting een rol kan spelen, maar U-Net++ met FCCE leverde hier wel de beste validatieresultaten op (DCval van 0.9762 naar 0.9839).
• Kwalitatieve Resultaten:
  • Visuele inspectie toont aan dat FCCE betere resultaten levert bij de segmentatie van cerebrospinaal vocht (CSF), grijze stof (GM) en witte stof (WM).
  • Het model is beter in staat om de onzekerheid aan de randen van weefsels te hanteren, wat leidt tot scherpere en nauwkeurigere grenzen in de segmentatie.

Significantie en Conclusie

Dit onderzoek demonstreert dat het integreren van fuzzy-logica in de loss-functie van diepe leermodellen een effectieve strategie is om de uitdagingen van onzekerheid in medische beeldsegmentatie aan te pakken.

• Betrouwbaarheid: De FCCE loss-functie verbetert niet alleen de numerieke metrics, maar verhoogt ook de betrouwbaarheid van de modelvoorspellingen in vaagke situaties.
• Toekomstperspectief: De methode biedt een nieuwe richting voor het ontwikkelen van segmentatiemodellen die minder gevoelig zijn voor ruis en partial volume effecten, wat essentieel is voor klinische toepassingen waar nauwkeurigheid levensbelangrijk is.

Kortom, door de onzekerheid niet als ruis te zien maar als een te modelleren eigenschap via fuzzy entropy, kunnen neurale netwerken betere en meer betrouwbare diagnoses ondersteunen.