Interpretable Medical Image Classification using Prototype Learning and Privileged Information

Dit paper introduceert Proto-Caps, een interpreteerbaar model voor medische beeldclassificatie dat capsule-netwerken, prototype-learning en voorrechtelijke informatie combineert om op het LIDC-IDRI-dataset zowel de voorspellingsnauwkeurigheid voor longknopen significant te verhogen als visueel verifieerbare redeneringen te bieden.

De kern

Stel je voor dat een computerprogramma een foto van een longknobbeltje bekijkt en zegt: "Dit is kanker." Voor een arts is dat vaak niet genoeg. Ze willen weten: Waarom denk je dat? Is het omdat de randen scherp zijn? Omdat de vorm bol is? Of omdat het eruit ziet als een specifieke, bekende kwaadaardige vorm?

De meeste slimme computerprogramma's (diep neurale netwerken) zijn als een "zwarte doos": ze geven het juiste antwoord, maar kunnen niet goed uitleggen hoe ze daar gekomen zijn. Dit artikel introduceert een nieuwe methode, genaamd Proto-Caps, die de computer niet alleen slim maakt, maar ook eerlijk en begrijpelijk.

Hier is hoe het werkt, vertaald in alledaagse taal:

1. De "Gevangenis" van de Zwarte Doos

Normaal gesproken werkt een AI als een genie dat in een donkere kamer zit. Het ziet een foto, denkt hard na, en roept het antwoord. Maar als je vraagt "Waarom?", zegt het: "Ik weet het niet precies, het voelt gewoon zo." In de geneeskunde is dat gevaarlijk. Artsen moeten kunnen controleren of de redenering klopt.

2. De Oplossing: Een Leermeester en een Verzameling

De auteurs van dit artikel hebben een slimme truc bedacht die twee dingen combineert:

• De "Privilege Informatie" (De Geheime Leermeester):
  Tijdens het trainen van de computer krijgen de programmeurs extra hulp van menselijke experts (radiologen). Deze experts zeggen niet alleen "kanker" of "geen kanker", maar geven ook details: "Deze knobbeltje heeft een scherpe rand, is bolvormig en heeft een bepaalde textuur."

  • De analogie: Stel je voor dat je een student leert wiskunde. Je geeft hem niet alleen de antwoorden op de toets, maar ook de stap-voor-stap uitleg van de leraar. De student leert zo niet alleen wat het antwoord is, maar ook hoe je er komt.
  • Na het trainen is deze extra informatie weg (de "privilege" is voorbij), maar de computer heeft de logica al geleerd.

• Prototype Learning (De Verzameling van Voorbeelden):
  In plaats van alleen abstracte getallen te gebruiken, leert de computer om nieuwe foto's te vergelijken met een verzameling van perfecte voorbeelden (prototypen).

  • De analogie: Stel je een detective voor die een nieuwe verdachte ziet. Hij zegt niet: "Hij lijkt op 85% op een crimineel." Hij zegt: "Kijk, deze verdachte lijkt precies op deze foto van een bekende inbreker die we in ons archief hebben, en op deze foto van een andere."
  • De computer toont dus echte, originele foto's van eerdere patiënten die lijken op de huidige patiënt. De arts kan dan zelf kijken: "Ah, ja, die randen zijn inderdaad hetzelfde als bij die andere kwaadaardige tumor."

3. De "Capsule" (De Slimme Doos)

De techniek gebruikt iets dat een "Capsule Network" heet.

• De analogie: Denk aan een oude doos met vakjes. In plaats van dat alles door elkaar loopt, heeft elk vakje een specifieke taak. Eén vakje kijkt alleen naar de vorm, een ander alleen naar de randen, en weer een ander naar de textuur.
• De computer vult deze vakjes met informatie. Als het vakje "randen" zegt dat ze scherp zijn, en het vakje "vorm" zegt dat ze bol zijn, dan komt het eindoordeel (kanker of niet) daaruit voort. Omdat elk vakje een duidelijke taak heeft, kunnen we zien welke eigenschap het belangrijkst was voor het oordeel.

4. Wat leverde dit op?

De onderzoekers hebben hun systeem getest op een grote database van longfoto's (LIDC-IDRI).

• Beter dan de rest: Het systeem was niet alleen makkelijker te begrijpen, maar ook preciezer dan andere geavanceerde systemen. Het had een nauwkeurigheid van 93% bij het voorspellen van kanker, wat meer dan 6% beter was dan de vorige beste "uitlegbare" methoden.
• De "Check": Als de computer zegt dat een knobbeltje kanker is, toont hij de vergelijkbare voorbeelden. Als de arts kijkt naar die voorbeelden en denkt: "Nee, die voorbeelden hebben een heel andere vorm dan deze patiënt," dan weet de arts: "Oké, de computer heeft hier waarschijnlijk een fout gemaakt." Dit maakt het systeem betrouwbaarder.

5. De Grootte van de Stap

Het mooiste aan deze methode is dat je niet alle details van elke patiënt nodig hebt om het systeem te laten werken. Zelfs als je maar voor 10% van de patiënten die extra, gedetailleerde beschrijvingen hebt (de "privilege informatie"), werkt het systeem nog steeds uitstekend.

Kort samengevat:
Deze paper introduceert een AI die niet alleen slim is, maar ook een uitleg geeft. Het doet dit door te leren van menselijke experts tijdens het trainen en door nieuwe gevallen te vergelijken met een verzameling van bekende, perfecte voorbeelden. Het is alsof je een arts krijgt die niet alleen een diagnose stelt, maar ook laat zien: "Kijk, dit is waarom ik dit denk, en hier zijn drie eerdere gevallen die er precies zo uitzagen."

Technische samenvatting

Probleemstelling

Diepe leersystemen presteren uitstekend in medische beeldanalyse, maar hun complexiteit maakt het moeilijk om de besluitvorming te verklaren (het "black-box"-probleem). In de medische beeldvorming is interpreteerbaarheid echter essentieel voor vertrouwen en klinische adoptie. Bestaande methoden voor Explainable AI (XAI) voegen vaak uitleg toe bovenop een bestaand model, wat kan leiden tot een daling in prestaties of onnauwkeurige uitleggen. Een veelbelovende aanpak is het gebruik van Privileged Information (informatie die alleen tijdens het trainen beschikbaar is, zoals radioloog-gedefinieerde kenmerken van longknobbels), maar eerdere methoden vertrouwen vaak op voorspelde scores die niet direct te valideren zijn. Daarnaast bieden bestaande prototype-netwerken wel visuele voorbeelden, maar niet altijd de reden waarom een bepaald kenmerk relevant is.

Methodologie: Proto-Caps

De auteurs stellen Proto-Caps voor, een innovatieve architectuur die Capsule Networks, Prototype Learning en Privileged Information combineert om een model te creëren dat per ontwerp interpreteerbaar is.

1. Architectuur (Capsule Network):

   • Het model gebruikt een "backbone" van een Capsule Network. De invoer bestaat uit 2D-slices van CT-scans ($1 \times 32 \times 32$).
   • Het netwerk bestaat uit een convolutielaag, een primaire capsule-laag en een finale dichte capsule-laag.
   • In plaats van de klassieke methode waarbij de lengte van de capsule-vector de aanwezigheid van een kenmerk aangeeft, gebruiken de auteurs een Attribute Head. Elke capsule vertegenwoordigt een specifiek, vooraf gedefinieerd kenmerk (zoals sfericiteit, margin, spiculation, etc.).
   • Er zijn drie trainingsheads:
     • Target Head: Voorspelt de maligniteit (kwaadaardigheid) van de knobbels.
     • Reconstruction Head: Reconstructie van het segmentatiemasker van de knobbels (om de structuur te behouden).
     • Attribute Head: Voorspelt de scores voor de specifieke kenmerken.

2. Privileged Information:

   • Tijdens het trainen worden de kenmerken (attributen) gebruikt als "privilege information". De radiologen hebben deze waarden (bijv. mate van spiculation) al gemarkeerd.
   • Het model leert de relatie tussen de capsule-encodings en deze menselijke attributen.

3. Prototype Learning voor Validering:

   • Om de voorspellingen van de attributen te valideren, leert het model prototypen (representatieve voorbeelden) voor elke attribuutklasse.
   • Er worden twee prototypen geleerd per mogelijke attribuutwaarde.
   • Tijdens de inferentie wordt de voorspelling voor een attribuut niet alleen gebaseerd op de geleerde dichte laag, maar wordt deze gekoppeld aan de dichtstbijzijnde prototype in de latent space. De originele afbeelding van dat prototype wordt getoond aan de gebruiker.
   • Dit stelt de arts in staat om visueel te verifiëren of de voorspelling logisch is: als het prototype sterk afwijkt van de invoer, kan de voorspelling betwijfeld worden.

4. Verliesfunctie:

   • De totale verliesfunctie ($L$) is een gewogen som van:
     • $L_{mal}$: Kullback-Leibler-divergentie voor maligniteit (om rekening te houden met onzekerheid tussen radiologen).
     • $L_{recon}$: Reconstructieverlies (MSE) voor het segmentatiemasker.
     • $L_{attr}$: Verlies voor attribuutvoorspelling.
     • $L_{clu}$ en $L_{sep}$: Verliesfuncties om prototypes dicht bij hun juiste klasse te brengen en ver weg van andere klassen te houden.

Belangrijkste Bijdragen

• Proto-Caps: De eerste methode (volgens de auteurs) die privileged information en prototype learning combineert voor medische classificatie, wat leidt tot een verhoogde verklarende kracht.
• Hybride Architectuur: Een capsule-netwerk gebaseerd op attributen, uitgebreid met een prototype-laag voor case-based redenering.
• Superieure Prestaties: Een oplossing die niet alleen interpreteerbaar is, maar ook de state-of-the-art prestaties (zowel voor interpreteerbare als niet-interpreteerbare methoden) overtreft op de LIDC-IDRI dataset.
• Open Source: De code en het model zijn beschikbaar gesteld op GitHub.

Resultaten

De methode werd getest op de LIDC-IDRI dataset (1018 thoracale CT-scans) met 5-voudige kruisvalidatie.

• Prestaties:
  • De voorspelling van maligniteit bereikte een Within-1-Accuracy van 93,0%. Dit is meer dan 6% hoger dan de beste bestaande interpreteerbare baseline (X-Caps, 86,4%) en verbetert ook de prestaties van niet-interpreteerbare methoden.
  • Het model presteerde bovendien uitstekend in het voorspellen van de specifieke kenmerken (attributen), met hoge nauwkeurigheden (bijv. 99,8% voor subtiliteit, 95,4% voor sfericiteit).
• Robuustheid (Data Reductie):
  • Experimenten waarbij attributen slechts in 10% van de trainingsdata beschikbaar waren, toonden aan dat de prestaties nauwelijks afnamen (92,4% vs 93,0%). Dit suggereert dat het model robuust is en dat privileged information vooral dient om de interpreteerbaarheid te verhogen zonder de prestaties te schaden.
• Ablatie-studie:
  • Het gebruik van prototypes tijdens inferentie bleek geen significante daling in nauwkeurigheid te veroorzaken ten opzichte van modellen zonder prototypes, maar wel een groot voordeel in interpreteerbaarheid.
• Kwalitatieve Analyse:
  • Visuele inspectie toonde aan dat discrepanties tussen een invoerbeeld en het dichtstbijzijnde prototype (bijv. bij margin of lobulatie) kunnen wijzen op mogelijke fouten in de maligniteitsvoorspelling. Dit biedt een mechanisme voor "case-based reasoning".

Significantie en Conclusie

Proto-Caps bewijst dat het mogelijk is om een model te bouwen dat per ontwerp interpreteerbaar is zonder in te leveren op prestaties. Door privileged information (radioloog-kenmerken) te combineren met prototype learning, biedt het model niet alleen een voorspelling, maar ook een visuele onderbouwing die door een arts kan worden gecontroleerd.

De studie toont aan dat:

1. Interpreteerbaarheid en prestatie geen tegenstelling hoeven te zijn.
2. Privileged Information effectief kan worden gebruikt om menselijke redenering te simuleren en te valideren.
3. Het model robuust is, zelfs als de hoeveelheid gelabelde attribuutdata beperkt is, wat de toepasbaarheid op andere datasets met minder annotatie-uitgaven vergroot.

De auteurs concluderen dat deze aanpak veelbelovend is voor hoog-risico toepassingsdomeinen en een nieuwe richting opent voor verder onderzoek in Explainable AI, met name door het gebruik van andere vormen van privileged information (zoals medische rapporten) en geavanceerdere 3D-Capsule-netwerken.