Technische Samenvatting: Reconstruction-Shift Discriminatie via Mask-Guided Latent Diffusion voor Medische Anomaliedetectie

Probleemstelling

Ongesuperviseerde medische anomaliedetectie heeft als doel pathologische bevindingen te identificeren en aangetaste regio's te lokaliseren met behulp van uitsluitend gezonde trainingsbeelden. Hoewel reconstructiegebaseerde methoden (bijv. autoencoders, GANs en diffusiemodellen) dominant zijn, lijden ze aan een inherente ambiguïteit bij het gebruik van reconstructieresten als het enige anomaliesignaal. Expressieve modellen kunnen onbedoeld pathologische structuren behouden, wat leidt tot zwakke reacties op werkelijke laesies. Omgekeerd kunnen benigne anatomische variaties, beeldruis en acquisitieverschillen grote reconstructiefouten genereren in gezonde regio's. Het aggregeren van deze resten in een enkel beeldniveau-score leidt vaak tot een vermenging van pathologische veranderingen met benigne reconstructie-artefacten, wat het moeilijk maakt om werkelijke anomalieën te onderscheiden van normale variabiliteit.

Methodologie: Discriminatieve Mask-Guided Diffusion (DMD)

De auteurs stellen Discriminative Mask-Guided Diffusion (DMD) voor, een framework dat de op residu gebaseerde lokalisatie aanvult met een geleerd beeldniveau-signaal afgeleid van reconstructie-shift discriminatie. De methode werkt in twee afzonderlijke stadia:

Stadium 1: Gekwantiseerde Normale Latente Representatie

Een vector-gekwantiseerde autoencoder (VQ-VAE) wordt getraind op normale beelden om een compacte latente representatie te leren.

• Encoder ($E_\phi$) en Decoder ($D_\eta$): Brengen beelden in kaart naar een latente ruimte en reconstrueren deze.
• Vector Kwantisatie: Ruimtelijke latente vectoren worden vervangen door hun dichtstbijzijnde vermeldingen in een leerbare codeboek ($C$), wat een discrete grid $z_q$ creëert.
• Doelstelling: Minimaliseer de reconstructiefout en de commitment loss. Na training worden $E_\phi$, $D_\eta$ en $C$ bevroren om een stabiele latente ruimte te garanderen voor het volgende stadium.

Stadium 2: Mask-Guided Latent Diffusion + Reconstruction-Shift Discriminatie

Dit stadium traint gezamenlijk een latent diffusiemodel en een classifier met behulp van de bevroren encoder/decoder.

1. Mask-Guided Latent Diffusion:

   • Gelokaliseerde vierkante maskers ($M$) worden toegepast op de gekwantiseerde latente grid $z_0$, waarbij specifieke regio's worden geselecteerd voor perturbatie.
   • Een forward diffusieproces voegt Gaussische ruis toe alleen aan de gemaskeerde regio's, waarbij de niet-gemaskeerde regio's als schone context achterblijven.
   • Een denoising netwerk ($\epsilon_\omega$) voorspelt de ruis en reconstrueert de schone latente $\hat{z}_0$.
   • De decoder brengt dit herstelde latente terug naar de beeldruimte, waardoor een "reconstructie tegenhanger" ($x_r$) van het originele normale beeld ($x$) wordt geproduceerd.
   • De diffusie loss ($L_{diff}$) wordt alleen berekend over de gemaskeerde posities.

2. Reconstruction-Shift Discriminatie:

   • Een zelfgesuperviseerde classificatietaak wordt gedefinieerd waarbij het originele normale beeld $x$ wordt gelabeld als klasse 0, en zijn gereconstrueerde tegenhanger $x_r$ (die een "shift" bevat geïnduceerd door de gemaskeerde diffusie) wordt gelabeld als klasse 1.
   • Een classifier ($f_\psi$) leert het onderscheid tussen de originele en de verschoven reconstructie.
   • De classificatie loss ($L_{cls}$) levert discriminatieve feedback aan het denoising netwerk, wat het model aanmoedigt om reconstructies te produceren die onderscheidend zijn van de originele inputs maar anatomisch plausibel blijven.
   • De totale loss is $L_{stage2} = L_{diff} + \lambda L_{cls}$.

Dual-Level Inferentie

Tijdens de testfase biedt DMD twee complementaire outputs:

1. Beeldniveau Detectie: De classifier $f_\psi$ verwerkt het testbeeld direct om een anomalie score $S_{img}$ te produceren, die de waarschijnlijkheid weergeeft dat het beeld lijkt op een "reconstructie shift" in plaats van een normaal origineel.
2. Pixelniveau Lokalisatie: Het testbeeld wordt gecodeerd en gekwantiseerd. Overlappende maskers worden toegepast over de latente grid, en de denoiser reconstrueert het beeld voor elke mask. De pixelniveau anomalie kaart wordt gegenereerd door de residuen (absolute verschillen) tussen de input en de diverse reconstructie tegenhangers te aggregeren, gemiddeld over de overlappende regio's.

Belangrijkste Bijdragen

1. Reconstruction-Shift Discriminatie: De auteurs identificeren en implementeren een zelfgesuperviseerd signaal dat uitsluitend op normale data is gebaseerd, waarmee een beeldniveau anomalie score wordt geleerd van gemaskeerde diffusie-reconstructies. Dit vult traditionele residu-gebaseerde lokalisatie aan zonder dat er abnormale trainingsdata nodig zijn.
2. DMD Framework: Een verenigde architectuur waarbij gelokaliseerde latente masking en latente diffusie zowel beeldniveau detectie als pixelniveau lokalisatie tegelijkertijd aansturen. Het framework gebruikt één enkel restauratieproces om zowel het discriminatieve signaal (voor de classifier) als de ruimtelijke aanwijzingen (voor de residu kaart) te genereren.
3. State-of-the-Art Prestaties: De methode demonstreert superieure prestaties over vijf diverse medische beeldvormingsdatasets (hersen-MRI, borst echografie en thorax radiografie).

Experimentele Resultaten

Het paper evalueert DMD op vijf datasets: BraTS2021, IXI/ATLAS 2.0, WMH, BUSI, en VinDr-CXR.

• Beeldniveau Detectie: DMD behaalde de hoogste AUC, AP en F1 scores over alle datasets. Op BraTS2021 behaalde DMD bijvoorbeeld een AUC van 87,2%, waarmee het de sterkste baseline (MAD-AD) met 1,1% versloeg en de F1 met 3,1% verbeterde. Op VinDr-CXR bereikte het een AUC van 84,3%, terwijl het op BUSI een AUC van 93,8% en een AP van 96,0% bereikte.
• Pixelniveau Lokalisatie: DMD presteerde consistent beter dan klassieke reconstructiemethoden (AE, VAE, f-AnoGAN) en recente diffusie-gebaseerde baselines (AnoDDPM, AutoDDPM, MAD-AD) in termen van AP$_{pix}$ en Dice scores. Op BraTS2021 behaalde DMD een Dice score van 76,3%, aanzienlijk hoger dan de op één na beste methode (MAD-AD op 69,6%).
• Cross-Dataset Robuustheid: Wanneer DMD werd toegepast op de WMH-dataset zonder hertraining (getraind op BraTS2021), behield het sterke prestaties, wat suggereert dat het generaliseert naar verschillende typen laesies en acquisitiekenmerken.
• Kwalitatieve Analyse: Visualisaties toonden aan dat DMD scherpere, meer ruimtelijk coherente anomalie kaarten produceert die beter aansluiten bij de grondwaarheid van de laesiegrenzen, vooral voor subtiele, laag-contrast abnormaliteiten waar andere methoden verspreide of zwakke reacties vertoonden.
• Ablatie Studies: Het verwijderen van de classifier, de discriminatieve feedback, de pre-training fase, de masking strategie of de latente diffusie component leidde allemaal tot prestatievermindering, wat de noodzaak van elk onderdeel bevestigt. Specifiek het verwijderen van latente diffusie veroorzaakte de meest significante daling, wat de rol ervan in het modelleren van anatomische variaties benadrukt.

Betekenis en Claims

Het paper claimt dat DMD de fundamentele ambiguïteit aanpakt van het gebruik van een enkel reconstructieresidu voor zowel detectie als lokalisatie. Door de taken te ontkoppelen — het gebruik van een geleerde classifier voor beeldniveau discriminatie en residuen voor pixelniveau lokalisatie — maakt het framework gebruik van de sterke punten van beide benaderingen.

De auteurs benadrukken dat het beeldniveau-signaal volledig wordt afgeleid van normale trainingsdata via het concept van "reconstructie shift", waardoor de noodzaak voor pathologische annotaties wordt vermeden. Hoewel de methode state-of-the-art resultaten behaalt, merkt het paper bescheiden op dat een beperking is: het beeldniveau-signaal leunt op een vaste masking policy, en de overdraagbaarheid naar onbekende pathologieën en acquisitiecondities vereist verdere brede studie. Toekomstig werk wordt voorgesteld om adaptieve shift constructie en multi-center evaluaties te verkennen.