SANA I2I: A Text Free Flow Matching Framework for Paired Image to Image Translation with a Case Study in Fetal MRI Artifact Reduction

Probleemstelling

Huidige generatieve modellen voor beeld-naar-beeld vertaling (Image-to-Image of I2I), zoals SanaControlNet, zijn vaak afhankelijk van tekstuele prompts (conditioning) naast beeldinvoer. Dit creëert beperkingen in scenario's waar tekstannotaties ontbreken, dubbelzinnig zijn of onnodig zijn, zoals bij medische beeldverwerking. In taken zoals ruisreductie, artefactverwijdering of modality-vertaling is het gewenste resultaat volledig bepaald door de invoerbeeld, niet door een tekstuele beschrijving. Het afhankelijke zijn van tekst kan leiden tot ongewenste stochasticiteit of inconsistenties.

Daarnaast is het in de medische beeldvorming, specifiek bij foetale MRI, zeer moeilijk om echte gepaarde datasets te verkrijgen (d.w.z. een "corrupt" beeld en het bijbehorende "schoon" beeld van exact dezelfde foetus op hetzelfde moment). Foetale en maternale bewegingen maken het onmogelijk om twee beelden zonder beweging te nemen. Bestaande methoden vertrouwen vaak op ongepaarde leermethodes die de anatomische trouw kunnen schaden.

Methodologie

De auteurs stellen SANA-I2I voor, een tekstvrij, hoog-resolutie framework voor gepaarde beeld-naar-beeld vertaling dat het SANA-gezinsmodel uitbreidt door tekstconditioning volledig te verwijderen.

Kerncomponenten:

1. Flow Matching in Latente Ruimte: Het model leert een conditionele snelheidsveld (velocity field) dat een verdeling van doelbeelden naar een andere verdeling mapt. Dit gebeurt in de latente ruimte, gebruikmakend van de Deep Compression Autoencoder (DC-AE) voor efficiënte compressie van hoge-resolutie beelden.
2. Architecturale Aanpassing:
   • Het model is afgeleid van SanaControlNet, maar verwijdert de tekstuele embeddinglevel (y als tekst) volledig.
   • De invoer wordt gedefinieerd als (x, t, y), waarbij y nu een beeld- of latente representatie is die dient als de enige conditionele signaal (in plaats van tekst).
   • Voor architecturale compatibiliteit wordt y behandeld als het control-signaal. Beide worden verwerkt via dezelfde PatchEmbed netwerken, wat zorgt voor een gedeelde latente ruimte voor ruis, conditionering en controle.
3. Classifier-Free Guidance (CFG): Het model behoudt het CFG-mechanisme. Tijdens het training wordt de conditionele variabele y met een bepaalde waarschijnlijkheid (p_drop) op nul gezet, terwijl het controle-pad actief blijft. Dit stelt het model in staat om zowel conditionele als unconditional gedrag te leren zonder opnieuw te hoeven trainen.
4. Synthetische Datageneratie: Om gepaarde training mogelijk te maken voor foetale MRI, gebruiken de auteurs een strategie van Duffy et al. om realistische bewegingsartefacten synthetisch toe te voegen aan schone beelden. Hierdoor ontstaat een gepaarde dataset (schone beelden vs. synthetisch vervormde beelden) met pixel-nauwkeurige alignering. Ze focussen op zware artefacten (SSIM tussen 0.6 en 0.9) om robuustheid te garanderen.

Varianten:

• Primary Approach: Alleen de conditionele variabele y wordt gedropt voor CFG. Dit is de hoofdmethode voor gepaarde I2I-taken.
• Bis Approach (Exploratorisch): Zowel y als het controle-signaal worden gedropt, wat volledige ongeconditioneerde generatie mogelijk maakt. Deze variant presteerde echter slechter voor artefactreductie.

Belangrijkste Bijdragen

• Tekstvrije I2I-extensie: Introductie van een SanaControlNet-variant die volledig is ontworpen voor supervised gepaarde beeldvertaling zonder tekstuele prompts.
• Efficiëntie: Het bewijzen dat het verwijderen van tekstconditioning het model beter laat aansluiten bij beeldtransformatietaken en het aantal benodigde inferentiestappen drastisch reduceert (tot slechts 2-5 stappen voor hoge kwaliteit).
• Schalbaarheid: Behoud van de efficiëntie en schaalvoordelen van de SANA-familie voor hoog-resolutie generatie.
• Medische Validatie: Validatie op de uitdagende taak van bewegingsartefactreductie in foetale MRI, met een strategie voor synthetische datageneratie om gepaarde training mogelijk te maken.

Resultaten

De methode is geëvalueerd op synthetische data (gepaard) en echte foetale MRI-data (ongepaard).

1. Synthetische Evaluatie (C → B):
   • Op traditionele pixel-gebaseerde metrics zoals SSIM en MAE scoorde de voorgaande methode [16] iets hoger.
   • Analyse: Dit komt doordat SANA-I2I de globale luminantie corrigeert (wat de pixel-waarden verandert en de fout vergroot), terwijl de oude methode de oorspronkelijke helderheid behoudt maar artefacten laat staan. De auteurs betogen dat de lagere SSIM/MAE hier een kunstmatig nadeel is; visueel is het resultaat van SANA-I2I schoner en anatomisch correcter.
2. Realistische Evaluatie (A → B):
   • Op distributie-gebaseerde metrics (FID en KID), die de statistische gelijkenis met schone beelden meten, scoorde SANA-I2I significant beter (FID: 51.74 vs 60.14).
   • Dit bevestigt dat de gegenereerde beelden statistisch dichter bij de verdeling van hoogwaardige foetale MRI's liggen.
3. Kwalitatieve Analyse:
   • Visuele inspectie toont aan dat SANA-I2I bewegingsartefacten effectiever onderdrukt en anatomische grenzen beter behoudt dan de vorige methoden, zonder onwaarschijnlijke structuren in te brengen.
4. Ablatie-studie:
   • Stappen: Hoogwaardige reconstructies zijn mogelijk met slechts 2 tot 5 inferentiestappen.
   • Guidance Scale: Een waarde tussen 1.0 en 1.1 biedt de beste balans tussen artefactverwijdering en behoud van anatomische details. Hogere waarden leiden tot over-agressieve correctie en detailverlies.

Betekenis en Conclusie

SANA-I2I demonstreert dat flow-matching modellen zeer geschikt zijn voor medische beeldverwerkingstaken waar tekstconditioning niet nodig is of zelfs schadelijk kan zijn. Door zich te focussen op gepaarde beeldinvoer, bereikt het model:

• Superieure prestaties in het verwijderen van bewegingsartefacten in foetale MRI.
• Extreme snelheid door het gebruik van slechts een paar inferentiestappen.
• Robuustheid door training op synthetisch gegenereerde, zware artefacten.

De studie benadrukt dat traditionele pixel-gebaseerde metrics (SSIM/MAE) soms misleidend kunnen zijn bij medische beeldherstel taken waar globale intensiteitscorrecties noodzakelijk zijn, en dat distributie-gebaseerde metrics (FID/KID) een betrouwbaarder maatstaf zijn voor de perceptuele en statistische kwaliteit. De methode biedt een veelbelovende basis voor andere toepassingen in de medische beeldvorming die gepaarde data vereisen.