CompDiff: Hierarchical Compositional Diffusion for Fair and Zero-Shot Intersectional Medical Image Generation

Het paper introduceert CompDiff, een hiërarchisch compositional diffusion-framework dat het probleem van ongelijke generatiekwaliteit bij zeldzame demografische groepen in medische beeldvorming oplost door conditionering op representatieniveau te decomponeren, wat leidt tot verbeterde eerlijkheid, zero-shot generalisatie en prestaties van downstream-klassificatoren.

De kern

CompDiff: De "Lego-meester" voor eerlijke medische foto's

Stel je voor dat je een kunstenaar bent die foto's maakt van mensen voor een medische studie. Je wilt dat je kunstwerk zo eerlijk mogelijk is: foto's van oude mensen, jonge mensen, mannen, vrouwen, en mensen van alle verschillende achtergronden moeten er allemaal even goed uitzien.

Het probleem is dat de kunstenaars (de computerprogramma's) die we vandaag de dag gebruiken, vaak een slechte gewoonte hebben. Ze zijn getraind op foto's waar veel witte mannen op staan, maar heel weinig foto's van bijvoorbeeld een oudere Aziatische vrouw. Als je de computer vraagt om een foto van die specifieke vrouw te maken, komt hij in de problemen. Hij probeert het, maar het resultaat ziet er wazig, raar of zelfs onherkenbaar uit.

De auteurs van dit paper noemen dit het "ongelijkheids-probleem": de kunstenaar is zelf niet eerlijk, dus hij kan geen eerlijke foto's maken voor iedereen.

De oude manier: "Meer oefenen" werkt niet

Vroeger probeerden onderzoekers dit op te lossen door de computer te dwingen om meer aandacht te besteden aan de zeldzame groepen. Het was alsof je een leerling die slecht wiskunde doet, dwingt om 10 uur per dag te oefenen, terwijl hij al 10 uur per dag wiskunde doet. Het helpt niet echt als de leerling de basisconcepten nog niet snapt.

De nieuwe manier: CompDiff (De Lego-benadering)

De auteurs van dit paper, Mahmoud Ibrahim en zijn team, hebben een slimme nieuwe manier bedacht die ze CompDiff noemen. In plaats van de computer te dwingen om meer te oefenen, geven ze hem een nieuwe manier van denken.

Stel je voor dat je een enorme doos met Lego-blokken hebt.

• Je hebt blokken voor "Leeftijd" (oud, jong).
• Je hebt blokken voor "Geslacht" (man, vrouw).
• Je hebt blokken voor "Herkomst" (Aziatisch, Europees, etc.).

De oude computers probeerden een foto te maken alsof ze een compleet, nieuw Lego-kasteel moesten bouwen voor elke combinatie. Als ze nooit een "Oud-Aziatisch-Vrouwelijk" kasteel hadden gezien, wisten ze niet hoe ze die moesten bouwen.

CompDiff werkt als een echte Lego-meester.
Deze computer heeft een speciale module (de Hierarchical Conditioner Network) die de blokken eerst loskoppelt. Hij leert eerst perfect hoe je een "Oud"-blok bouwt, hoe je een "Aziatisch"-blok bouwt en hoe je een "Vrouwelijk"-blok bouwt.

Als je nu vraagt om een foto van een "Oude Aziatische Vrouw", pakt de computer gewoon die drie losse, goed gebouwde blokken en plakt ze aan elkaar. Omdat hij de losse onderdelen al perfect kent, kan hij ook de combinatie maken die hij nog nooit eerder heeft gezien. Dit noemen ze compositional generalization (samenstellende generalisatie).

Wat levert dit op?

1. Beter voor iedereen: De foto's die de computer maakt zijn scherper en realistischer, niet alleen voor de grote groepen, maar ook voor de zeldzame groepen.
2. Eerlijkheid: De kwaliteit van de foto's is nu voor iedereen gelijk. Er is geen "wazige foto" meer voor de zeldzame groepen.
3. Beter voor artsen: Als artsen deze gegenereerde foto's gebruiken om hun diagnose-apparatuur te trainen, maken die apparaten minder fouten. Ze zien bijvoorbeeld ziektes bij vrouwen van verschillende achtergronden net zo goed als bij mannen.

De conclusie in één zin

In plaats van de computer te dwingen om meer te oefenen op moeilijke gevallen, geven we hem een slimme bouwset (Lego) zodat hij zelfstandig nieuwe combinaties kan maken. Hierdoor worden medische foto's niet alleen mooier, maar ook eerlijker voor iedereen, ongeacht hun leeftijd, geslacht of afkomst.

Het is een stap in de richting van een toekomst waarin kunstmatige intelligentie niemand uitsluit, zelfs niet als de data daarvoor schaars is.

Technische samenvatting

Hieronder volgt een gedetailleerde technische samenvatting van het paper "CompDiff: Hierarchical Compositional Diffusion for Fair and Zero-Shot Intersectional Medical Image Generation" in het Nederlands.

Probleemstelling: Het "Imbalanced Generator"-probleem

Generatieve modellen, zoals diffusion-modellen, worden steeds vaker gebruikt om medische beeldvormingsdatasets aan te vullen voor het trainen van eerlijkere AI-systemen. Een cruciale, maar vaak over het hoofd geziene aanname is echter dat deze generators zelf even hoogwaardige beelden produceren voor alle demografische groepen.

Wanneer modellen worden getraind op onbalans data, erven ze deze onbalans over. Dit leidt tot:

1. Gedegradeerde synthese-kwaliteit voor zeldzame subgroepen.
2. Moeite met demografische intersecties die niet in de trainingsdata voorkomen (bijv. een "80+ Aziatische vrouw" met een specifieke pathologie, terwijl er in de dataset wel oude patiënten, Aziatische patiënten en vrouwen apart zijn, maar geen combinatie van alle drie).

Bestaande oplossingen, zoals FairDiffusion, werken op het optimalisatieniveau (bijv. herweging van de loss-functie). Dit helpt echter niet wanneer er geen trainingsvoorbeelden zijn voor bepaalde combinaties; je kunt geen leersignaal genereren voor data die niet bestaat. Het paper noemt dit het "imbalanced generator problem".

Methodologie: CompDiff

Het auteurs stellen CompDiff voor, een hiërarchisch compositional diffusion-framework dat het probleem oplost op het representatieniveau in plaats van op het optimalisatieniveau. De kerninzicht is dat demografische identiteit compositional is: een zeldzame intersectie kan worden samengesteld uit goed geleerde enkelvoudige attributen en gepaarde interacties.

De architectuur bestaat uit de volgende componenten:

1. Hierarchical Conditioner Network (HCN):
   In plaats van demografische attributen (leeftijd, geslacht, ras) als tekst-tokens in de CLIP-prompt te stoppen (waar ze concurreren met klinische tokens), worden ze verwerkt via een dedicated HCN.

   • Enkelvoudige attributen ("Grootouders"): Elke attribut (bijv. leeftijd $a$, geslacht $s$, ras $r$) wordt ingebed in een gedeelde latente ruimte.
   • Paarsgewijze interacties ("Ouders"): Om niet-additieve relaties te vangen, worden MLP's gebruikt om interacties tussen paren te modelleren (bijv. $f_{a,s}$ voor leeftijd en geslacht).
   • Volledige compositie ("Kind"): Een laatste MLP combineert deze paarsgewijze interacties tot een finale demografische representatie ($h_{demo}$).

2. Latente Projectie en Conditionering:
   De representatie $h_{demo}$ wordt gemap naar een diagonale Gaussische verdeling (met mean $\mu$ en variance $\sigma$). Tijdens training wordt er gesampled via reparameterization, en tijdens inferentie wordt $\mu$ gebruikt. Deze latente vector wordt geprojecteerd naar een demografisch token $c$, dat vervolgens wordt geconcateneerd met de CLIP-embeddings van de klinische tekst. Deze gecombineerde vector dient als cross-attention context voor de diffusion UNet.

3. Trainingsdoel (Loss Function):
   De totale loss bestaat uit vier delen:

   • $L_{diff}$: Standaard diffusion loss.
   • $L_{KL}$: Regularisatie van de variational latent naar een standaard normale verdeling.
   • $L_{comp}$: Een compositional consistentie term die de hiërarchische output regulariseert ten opzichte van een simpele additieve baseline (om training te stabiliseren).
   • $L_{aux}$: Een hulp-classificatieterm die direct wordt toegepast op het geprojecteerde token $c$ (niet op $\mu$) om ervoor te zorgen dat de demografische informatie behouden blijft voor de UNet.

Belangrijkste Bijdragen

• Representatie-architectuur: Introductie van een hiërarchische conditioner die demografische attributen expliciet factoriseert, wat parameter-sharing tussen subgroepen stimuleert.
• Zero-Shot Generalisatie: Het vermogen om beelden te genereren voor demografische intersecties die volledig afwezig zijn in de trainingsdata, door deze te "componeren" uit bekende attributen.
• Eerlijkheid op representatieniveau: Het tonen aan dat het aanpassen van de input-structuur effectiever is voor eerlijkheid dan het aanpassen van de loss-weights.

Resultaten

De methode is geëvalueerd op twee datasets: MIMIC-CXR (borst-Röntgenfoto's) en FairGenMed (fundusfoto's). Vergelijkingen zijn gemaakt met standaard fine-tuning en FairDiffusion.

1. Beeldkwaliteit: CompDiff behaalde de beste FID-scores (Fréchet Inception Distance) op beide modaliteiten (64.3 voor CXR vs. 75.1 voor FairDiffusion).
2. Eerlijkheid (Equity): CompDiff toonde de laagste ES-FID (Equity-Scaled FID), wat betekent dat de kwaliteitsverschillen tussen demografische subgroepen het kleinst waren.
3. Zero-Shot Generalisatie: Bij het volledig verwijderen van bepaalde intersecties uit de trainingsdata, presteerde CompDiff aanzienlijk beter dan de baselines (tot 21% verbetering in FID). FairDiffusion faalde hier vaak omdat er geen trainingsdata was om de loss te herwegen.
4. Downstream Utility: Classificatoren getraind op gegenereerde data van CompDiff presteerden beter op echte testdata (hogere AUROC) en vertoonden minder demografische bias (lagere onderdiagnose-rates en gelijkheid in odds).

Betekenis en Conclusie

CompDiff demonstreert dat de architecturale inrichting van demografische conditionering een kritieke, maar onderbelichte factor is in eerlijke medische beeldgeneratie. Door te focussen op compositional generalization via een hiërarchische structuur, kan het model leren van zeldzame groepen en generaliseren naar ongezette combinaties zonder de kwaliteit van de meerderheidsgroepen te verstoren.

Hoewel het een grote stap voorwaarts is, blijven beperkingen bestaan: de methodiek is afhankelijk van gestructureerde demografische attributen (niet continu of ongestructureerd) en de evaluatie is voornamelijk kwantitatief. Toekomstig werk richt zich op geavanceerdere conditioneringsmechanismen, zoals graf-based interactiemodellering.