Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

Technische Samenvatting: Flow Matching met Geoptimaliseerde Subklasse-priors voor Medische Beeldaugmentatie

Probleemstelling
Medische beeldvormingsdatasets worden gekenmerkt door ernstige langstaart-ongelijkheden, waarbij veelvoorkomende aandoeningen domineren terwijl zeldzame, klinisch kritieke pathologieën (bijvoorbeeld kleine pneumothoraxen, bronchiëctasieën) ondervertegenwoordigd zijn. Deze onbalans zorgt ervoor dat discriminatieve classifiers overfitten op hoofdklassen en falen bij zeldzame aandoeningen. Hoewel generatieve data-augmentatie een mogelijke oplossing biedt, staan standaard conditionele flow matching (FM) en diffusiemodellen voor twee versterkende moeilijkheden wanneer ze worden getraind op grove ziekte-indicatoren:

1. Multi-modale Aggregatie: Een enkele diagnostische label vaak diverse subtypes, opnameprotocollen en comorbiditeiten. Dit creëert multi-modale condities die generatoren bevooroordeeld naar dominante submodi, wat effectief zeldzame maar klinisch geldige variaties afsnijdt.
2. Inefficiënte Transportgeometrie: Het trekken van alle generatietrajecten uit een gedeelde standaard-Gaussische bron dwingt specifieke subpopulaties door onevenredig lange transportpaden. Dit vergroot de regressieschwierigheid en verslechtert de fideliteit van de generatie van staartklassen.

Methodologie
De auteurs stellen een offline strategie voor die informatieve priors op twee niveaus introduceert om het generatieve proces te verfijnen zonder de kerntrainingslus van het flow-matching-model te wijzigen.

1. Subklasse-inductie via Gaussische Mixturer Modeling (GMM):

   • De methode partitioneert elke grove label $c$ in coherente submodi (subklassen $k$) binnen de latente ruimte van het generatieve model.
   • Met behulp van een bevroren VAE (FLUX1.dev) worden deterministische latente variabelen berekend. Voor elke klasse $c$ wordt een diagonale Gaussische Mixturer Model aangepast aan de residu-latente variabelen (na aftrekken van het klassencentrum $\mu_c$).
   • Het aantal componenten $K_c$ wordt geselecteerd met behulp van het Extended Bayesian Information Criterion (EBIC). Steekproeven krijgen harde toewijzingen aan subklassen, waardoor een partitie met lagere conditionele variantie ontstaat.
   • Theoretisch bewijzen de auteurs dat het conditioneren op deze verfijnde labels de Bayes-optimale flow-matching-risico niet kan verhogen; de verbetering komt overeen met de verwachte variabiliteit tussen subklassen van de conditionele gemiddelde verplaatsing.

2. Subklasse-afhankelijke Bronoptimalisatie:

   • In plaats van een gedeelde Gaussische bron leert de methode een specifieke bronverdeling $\pi_0(\cdot | c, k)$ voor elke subklasse. Deze bronnen worden geparametriseerd als leerbare Gaussians $N(\mu_{c,k}, \text{diag}(\sigma^2_{c,k}))$ die de startverdeling opnieuw centreren en herschalen.
   • Dit verkort typische transporttrajecten en vermindert de dispersie binnen subklassen.
   • Geometrische Controle: Om degeneratieve oplossingen (zoals norminflatie of overmatig geconcentreerde paden) te voorkomen, optimaliseren de auteurs een offline proxy-doelstelling die dwingt op:
     • Directionele Compactheid: Het aanmoedigen dat genormaliseerde verplaatsingsrichtingen matig aligneren met leerbare prototypes.
     • Padlengte-demping: Het opleggen van een vaste straf voor uitbijters met behulp van een subklasse-specifieke cap ($r_{cap}$), gedefinieerd als het 99e percentiel van afstanden tot het klassencentrum. Dit voorkomt dat het model de loss omzeilt door verplaatsingsnormen op te blazen.
     • Schaalregularisatie: Het regulariseren van de log-schalen van de bronverdelingen naar een eenheidsreferentie.

Belangrijkste Bijdragen

1. Subklasse-conditionele Formulering: Een nieuwe aanpak voor langstaart-medische beeldgeneratie die multi-modale condities partitioneert via GMM-geïnduceerde harde toewijzingen, wat effectief conditionele ambiguïteit vermindert.
2. Geleerde Subklasse-afhankelijke Bronnen: De introductie van informatieve priors die conditioneel transport herschikken door de bronverdeling per subklasse opnieuw te centreren en te herschalen, waardoor de dispersie binnen subklassen wordt verminderd.
3. Geometrisch Geregulariseerde Doelstelling: Een praktische proxy-doelstelling die een evenwicht zoekt tussen matige hoekige compactheid en expliciete controle van padlengte-uitbijters, waardoor shortcuts voor norminflatie en directionele ineenstorting worden voorkomen.
4. Empirische Validatie: Demonstratie van verbeterde staartdekking en downstream prestaties bij langstaart-classificatie op benchmarks voor borst-röntgenfoto's (MIMIC-LT, NIH-LT) en CT (CT-RATE).

Experimentele Resultaten
De methode werd geëvalueerd op drie langstaart-medische beeldvormingsbenchmarks:

• Downstream Classificatie: De voorgestelde methode behaalde consistent de hoogste gebalanceerde nauwkeurigheid (bAcc) en macro-F1-scores over alle drie de datasets in vergelijking met niet-augmenteerde baselines en andere augmentatiestrategieën (Vanilla FM, CBFM, CPD). Opmerkelijk was dat het de enige strategie was die betrouwbaar prestaties verbeterde over zowel röntgen- als CT-modi.
• Generatieve Kwaliteit: Op MIMIC-LT behaalde de methode de beste Fréchet Inception Distance (FID) en Inception Ratio Score (IRS), wat gelijktijdige winsten in fideliteit en diversiteit aangeeft. Op CT-RATE, hoewel de FID iets hoger was dan een baseline, behield de methode een hoge IRS en superieure downstream classificatie, wat wijst op betere distributie-dekking in plaats van louter memorisatie.
• Ablatiestudies: Resultaten bevestigden dat beide pijplijnfasen (subklassificatie en bronoptimalisatie) noodzakelijk zijn. Subklassificatie alleen verbeterde de prestaties, maar de combinatie leverde de beste resultaten op. Optimalisatie zonder subklassificatie degradeerde soms de prestaties door te middelen over conflicterende modi.
• Subklasse-geldigheid: Analyse toonde aan dat geïnduceerde subklassen pathologie-relevante structuur vangen in plaats van alleen opname-metadata. Verwarrende probes gaven aan dat metadata detecteerbaar was maar niet dominant, en de kNN-reinheid van geleerde subklassen overschreed significant willekeurige partities.

Betekenis en Claims
Het artikel claimt dat bronontwerp en conditionele verfijning effectieve hefbomen zijn voor robuuste generatieve augmentatie in klinische langstaart-situaties. Door heterogene labels te partitioneren en subklasse-afhankelijke bronnen te leren, vermindert de methode conditionele ambiguïteit en stabiliseert transporttrajecten tegen een verwaarloosbare rekenkosten (het hele proces is offline en duurt minder dan één GPU-uur per dataset). De auteurs positioneren dit als een principiële strategie die de fideliteit en diversiteit van staartklasse-generatie verbetert, wat leidt tot betrouwbaardere downstream classifiers voor zeldzame medische aandoeningen. Het werk benadrukt dat naïeve conditionele generatie vaak faalt bij zeldzame subpopulaties vanwege problemen met transportgeometrie, die kunnen worden verzacht door gestructureerde priors.