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

Der Artikel stellt CompDiff vor, ein hierarchisches, compositional Diffusionsframework, das durch eine dedizierte Hierarchische Conditioner Network (HCN) das Problem ungleicher Generierungsqualität über demografische Gruppen hinweg löst und so eine faire sowie zero-shot-fähige Synthese medizinischer Bilder für seltene und nicht im Training enthaltene demografische Schnittstellen ermöglicht.

Das Wesentliche

🎨 Das Problem: Der unfaire Künstler

Stell dir vor, du hast einen sehr talentierten KI-Künstler (ein sogenanntes „Diffusionsmodell"), der medizinische Bilder wie Röntgenaufnahmen oder Augen-Scans malen soll. Dieser Künstler wurde trainiert, um neue Bilder zu erstellen, damit Ärzte mehr Übungsmaterial haben.

Das Problem ist aber: Der Künstler hat nur viele Bilder von bestimmten Menschen gesehen (z. B. junge weiße Männer). Er hat kaum Bilder von seltenen Kombinationen gesehen (z. B. alte asiatische Frauen mit einer bestimmten Krankheit).

Wenn du ihn jetzt bittest, ein Bild von einer „alten asiatischen Frau" zu malen, macht er das zwar, aber das Ergebnis sieht oft schief, verzerrt oder einfach schlecht aus. Er versucht, es zu erraten, weil er keine echten Beispiele dafür hat. Das nennen die Autoren das „unausgewogene Generator-Problem".

Bisherige Lösungen waren wie ein strenger Lehrer, der dem Künstler sagt: „Mach bei den seltenen Bildern mehr Druck!" (das nennt man „Reweighting" oder Gewichtung). Aber das hilft wenig, wenn der Künstler gar keine echten Beispiele in seiner Erinnerung hat, auf die er sich stützen kann. Es ist, als würdest du jemanden bitten, ein Rezept für ein Gericht zu kochen, das er noch nie gekocht hat, und ihm nur sagen: „Sei bitte besonders fleißig dabei." Das Ergebnis wird trotzdem nicht schmecken.

💡 Die Lösung: CompDiff – Der Baumeister mit Legosteinen

Die Forscher haben eine neue Methode namens CompDiff entwickelt. Statt den Künstler zu drängen, ändern sie, wie er denkt und lernt.

Stell dir vor, der Künstler lernt nicht, ganze Bilder auswendig zu memorieren. Stattdessen lernt er, wie man Legosteine kombiniert.

1. Die einzelnen Steine: Der Künstler lernt sehr gut, wie ein „alter" Stein aussieht, wie ein „asiatischer" Stein aussieht und wie ein „weiblicher" Stein aussieht.
2. Die Kombination: Wenn er nun ein Bild von einer „alten asiatischen Frau" malen soll, baut er es einfach aus diesen drei gut gelernten Steinen zusammen. Er muss das Bild nicht auswendig gelernt haben; er kann es zusammensetzen (komponieren).

Das ist wie beim Sprechen: Du kennst die Wörter „Haus", „rot" und „groß". Du hast vielleicht nie den Satz „ein großes rotes Haus" gehört, aber du kannst ihn trotzdem verstehen und bilden, weil du die Regeln der Sprache (die Grammatik) kennst. CompDiff gibt dem KI-Künstler genau diese Grammatik für menschliche Merkmale.

🏗️ Wie funktioniert das technisch? (Die Hierarchische Leiter)

Die Forscher haben dem KI-Künstler eine spezielle Bauleiter (ein sogenanntes „Hierarchical Conditioner Network") gebaut:

• Unten (Die Basis): Die KI lernt die einzelnen Merkmale (Alter, Geschlecht, Herkunft) einzeln.
• Mitte (Die Paare): Sie lernt, wie diese Merkmale zusammenwirken (z. B. wie sich „Alter" und „Geschlecht" gegenseitig beeinflussen).
• Oben (Das Ganze): Am Ende werden alle diese Teile zu einem einzigen, perfekten Bauplan für die gewünschte Person verschmolzen.

Dadurch kann die KI auch Bilder von Menschen erstellen, für die es keine einzigen Trainingsbilder gab (sogenannte „Zero-Shot"-Fälle). Sie baut sie einfach aus den bekannten Teilen neu zusammen.

🏆 Was bringt das? (Die Ergebnisse)

Die Forscher haben CompDiff an echten medizinischen Daten getestet (Röntgenbilder der Lunge und Bilder vom Augenhintergrund). Das Ergebnis war beeindruckend:

1. Bessere Bilder: Die generierten Bilder sahen viel natürlicher aus als bei alten Methoden (bessere Qualität, gemessen am FID-Score).
2. Fairer: Die Qualität war für alle Gruppen gleich gut. Ob die KI ein Bild von einer jungen Frau oder einer alten asiatischen Frau malte – das Ergebnis war immer hochwertig.
3. Bessere Ärzte: Als sie mit diesen neuen, fairen Bildern trainierte KI-Ärzte (Klassifikatoren), machten diese weniger Fehler und waren fairer bei der Diagnose für alle Patientengruppen.

🚀 Fazit in einem Satz

Statt den KI-Künstler zu zwingen, mehr zu arbeiten, haben die Forscher ihm eine intelligere Art des Denkens beigebracht: Statt ganze Bilder auswendig zu lernen, lernt er, sie aus kleinen, fairen Bausteinen zu zusammensetzen. So entstehen für alle Menschen – auch die seltenen – gleich gute medizinische Bilder.

Technische Zusammenfassung

1. Problemstellung: Das „Imbalanced Generator"-Problem

Die Autoren identifizieren ein kritisches, oft übersehenes Problem bei der Anwendung von generativen Modellen (insbesondere Diffusionsmodellen) in der medizinischen Bildgebung: Die Annahme, dass Generatoren für alle demografischen Gruppen gleich hochwertige Bilder produzieren.

• Ursache: Modelle, die auf unausgewogenen Daten trainiert werden, leiden unter einem „Imbalanced Generator"-Problem. Während sie im Durchschnitt gute Ergebnisse liefern, sinkt die Synthesequalität für seltene Untergruppen drastisch.
• Das Intersektions-Problem: Das Problem verschärft sich bei demografischen Schnittstellen (Intersektionalität). Ein Datensatz kann ältere Patienten, asiatische Patienten und Frauen enthalten, aber keine einzigen Beispiele für die Schnittmenge „ältere asiatische Frauen" mit einer bestimmten Pathologie.
• Limitierung bestehender Ansätze: Herkömmliche Methoden wie Loss-Reweighting (z. B. FairDiffusion) operieren auf der Optimierungsebene. Sie können keine Lernsignale für Kombinationen erzeugen, die im Trainingsdatensatz gar nicht existieren. Sie scheitern also bei „Zero-Shot"-Szenarien für seltene demografische Kombinationen.

2. Methodik: CompDiff Framework

Die vorgeschlagene Lösung, CompDiff, adressiert das Problem auf der Repräsentationsebene (Representation Level) statt auf der Optimierungsebene. Der Kernansatz ist die kompositionelle Demografie-Conditioning.

A. Grundlegende Architektur

Im Gegensatz zu Standard-Diffusionsmodellen, die demografische Attribute als Teil des Text-Prompts codieren (und damit mit klinischen Tokens um den begrenzten Token-Budget von CLIP konkurrieren), trennt CompDiff die Verarbeitung:

• Klinische Daten: Werden über CLIP in Text-Embeddings ($E_{text}$) kodiert.
• Demografische Daten: Werden separat durch ein dediziertes Hierarchical Conditioner Network (HCN) verarbeitet.
• Fusion: Das HCN erzeugt einen dedizierten demografischen Token ($c$), der mit den CLIP-Embeddings konkateniert wird ($E_{combined}$) und als Cross-Attention-Kontext für das Diffusions-UNet dient.

B. Das Hierarchical Conditioner Network (HCN)

Das HCN führt eine strukturierte induktive Verzerrung (Inductive Bias) ein, indem es die demografische Conditionierung in hierarchische Komponenten zerlegt, analog zur Komposition von Wörtern in Sätzen:

1. Single-Attribute Embeddings („Großeltern"): Einzelne Attribute (Alter, Geschlecht, Ethnie) werden in einen gemeinsamen latenten Raum eingebettet ($e_a, e_s, e_r$).
2. Pairwise Interactions („Eltern"): Um nicht-additive Beziehungen zu erfassen, werden alle Paarkombinationen (z. B. Alter-Geschlecht) über dedizierte MLPs modelliert ($h_{a,s}, h_{a,r}, h_{s,r}$). Dies verhindert Overfitting bei sehr seltenen Gruppen, indem es die Komplexität begrenzt.
3. Full Composition („Kind"): Die finalen Interaktionen werden durch einen weiteren MLP ($g$) kombiniert, um die vollständige demografische Repräsentation ($h_{demo}$) zu erhalten.
4. Latent Mapping: $h_{demo}$ wird in eine diagonale Gauß-Verteilung ($\mu, \sigma$) projiziert. Ein latenter Vektor $z$ wird per Reparameterisierung gesampelt (Training) oder auf $\mu$ gesetzt (Inferenz) und dann auf die Cross-Attention-Dimension projiziert ($c$).

C. Trainingsziel (Loss Function)

Das Gesamt-Training minimiert eine kombinierte Loss-Funktion:
$$L = L_{diff} + \lambda_{comp}L_{comp} + \lambda_{aux}L_{aux} + \lambda_{KL}L_{KL}$$

• $L_{diff}$: Standard Diffusions-Loss.
• $L_{comp}$: Kompositionelle Konsistenz. Dies wirkt als „weicher Anker", der die komplexe Hierarchie an eine einfache additive Basis ($e_a + e_s + e_r$) bindet, um das Training zu stabilisieren, erlaubt aber dennoch nicht-additive Interaktionen.
• $L_{aux}$: Auxiliary Classification Loss. Wichtig ist hier, dass dieser Loss auf dem projizierten Token $c$ (dem Input für das UNet) angewendet wird, nicht auf dem latenten Mittelwert $\mu$. Dies stellt sicher, dass die demografische Information für das UNet erhalten bleibt.
• $L_{KL}$: Regularisierung des latenten Raums.

3. Wichtige Beiträge

1. Paradigmenwechsel: Der Übergang von optimierungsbasierten Lösungen (Reweighting) zu repräsentationsbasierten Lösungen (strukturelle Zerlegung der Conditionierung).
2. Zero-Shot Generalisierung: Durch die kompositionelle Struktur kann das Modell Repräsentationen für demografische Kombinationen generieren, die im Training nie gesehen wurden (z. B. „80+ asiatische Frauen"), indem es gelernte Einzelattribute und Paar-Interaktionen kombiniert.
3. Effizienz: Das HCN fügt nur einen minimalen Overhead hinzu (0,19% mehr trainierbare Parameter) und erfordert keine Änderungen an den Diffusions-Timesteps oder Sampling-Prozeduren.

4. Ergebnisse und Evaluation

Die Methode wurde auf zwei medizinischen Modalitäten evaluiert: Brust-Röntgenaufnahmen (MIMIC-CXR) und Fundusbilder (FairGenMed).

• Bildqualität (FID): CompDiff erzielt auf beiden Modalitäten die besten FID-Werte (z. B. 64,3 für Röntgen vs. 75,1 bei FairDiffusion).
• Fairness (ES-FID): CompDiff zeigt die geringste Diskrepanz in der Bildqualität zwischen verschiedenen demografischen Untergruppen (Sex, Race, Alter).
• Zero-Shot Intersektionalität: In Experimenten, bei denen 5 seltene Schnittstellen komplett aus dem Training entfernt wurden, übertraf CompDiff Baseline und FairDiffusion deutlich (bis zu 21% FID-Verbesserung). FairDiffusion scheiterte hier oft, da es keine Trainingsdaten für diese Kombinationen hatte.
• Downstream Utility: Klassifikatoren, die auf den von CompDiff generierten Daten trainiert wurden, zeigten höhere AUROC-Werte und reduzierte demografische Verzerrungen (z. B. geringere Unterdiagnose-Raten) auf echten Testdaten.

Ablationsstudien zeigten, dass die hierarchische Struktur entscheidend ist (flache MLPs scheitern bei der Kontrolle) und dass die Auxiliary Loss auf dem projizierten Token angewendet werden muss, um wirksam zu sein.

5. Bedeutung und Fazit

CompDiff demonstriert, dass die architektonische Gestaltung der demografischen Conditionierung ein entscheidender, aber oft unterschätzter Faktor für faire medizinische Bildgenerierung ist.

• Wissenschaftlicher Impact: Die Arbeit zeigt, dass strukturelle Induktiv-Bias (hierarchische Komposition) effektiver ist als reine Daten- oder Optimierungsstrategien, um Generalisierung auf unsichtbare demografische Gruppen zu erreichen.
• Praktische Relevanz: Die Methode ermöglicht die Erzeugung hochwertiger, fairer synthetischer Daten für seltene Patientengruppen, was direkt die Leistung und Fairness nachgelagerter diagnostischer KI-Systeme verbessert.
• Limitationen: Die Fairness-Evaluation basiert noch auf quantitativen Metriken statt klinischer Expertenbewertungen. Zudem ist die Methode auf strukturierte demografische Attribute angewiesen und weniger geeignet für kontinuierliche oder unstrukturierte Attribute.

Zusammenfassend bietet CompDiff einen robusten Weg, um die „Imbalanced Generator"-Problematik zu lösen und KI-Systeme in der Medizin inklusiver und fairer zu gestalten.