IR-Flow: Bridging Discriminative and Generative Image Restoration via Rectified Flow

Probleemstelling

In het domein van beeldherstel (image restoration) bestaan er twee dominante paradigma's met elk hun eigen beperkingen:

1. Discriminatieve methoden: Deze gebruiken vaak enkelvoudige stap-mappingen (single-step mappings) die zijn getraind om de verwachting (expectation) te minimaliseren (bijv. via $L_2$-verlies). Hoewel dit leidt tot lage vervorming (distortion), resulteren de outputbeelden vaak in een te gladde weergage die perceptuele realisme en fijne details mist.
2. Generatieve methoden (zoals Diffusion Models): Deze benaderingen leren complexe data-distributies en kunnen hoge-frequentie details en textuur genereren. Echter, ze lijden vaak onder inefficiënte multi-stap sampling (veel iteraties nodig) en een koppeling tussen ruis en residu (noise-residual coupling), wat de inferentie vertraagt en de optimalisatie bemoeilijkt.

Er is een gebrek aan een uniek raamwerk dat de efficiëntie van discriminatieve methoden combineert met de perceptuele kwaliteit van generatieve methoden, zonder de nadelen van complexe ruis-schedules.

Methodologie: IR-Flow

De auteurs stellen IR-Flow voor, een unificerend raamwerk gebaseerd op Rectified Flow (gecorrigeerde stroming). In plaats van een proces te modelleren waarbij een beeld geleidelijk wordt verstoord tot witte ruis en vervolgens weer wordt gereconstrueerd, bouwt IR-Flow een directe, lineaire transportstroom op tussen het verstoide beeld ($X_1$) en het schone doelbeeld ($X_0$).

De kerncomponenten van de methode zijn:

1. Multilevel Data Distribution Flows:

   • In plaats van alleen te werken met extreme degradaties, worden er kunstmatige tussenliggende staten gecreëerd met gecontroleerde degradatieniveaus.
   • Dit helpt het model om een bredere reeks degradatiescenario's te leren en verrijkt de lokale structuur van de data-distributie, wat leidt tot betere aanpassing aan verschillende degradatieniveaus.

2. Cumulatieve Veldsnelheid (Cumulative Velocity Field - CVF):

   • Standaard Rectified Flow leert de instantane raakvector (tangent direction) $v(X_t, t) = E[X_1 - X_0 | X_t]$. Dit vereist kleine stapgroottes voor nauwkeurige integratie.
   • IR-Flow introduceert een cumulatieve snelheid: $\hat{v}(X_t, t) = E[X_t - X_0 | X_t]$.
   • Deze vector wijst direct van de huidige tussenliggende staat naar het doel ($X_0$). Dit vermindert de directionele afwijking en vereist minder "energie" (transmissie-energie) om het doel te bereiken. Wiskundig is de transportenergie van CVF slechts $1/3$ van die van de standaard snelheid, wat leidt tot snellere convergentie en stabielere training.

3. Multi-stap Consistentie Training (MCT Loss):

   • Om ervoor te zorgen dat de gelearnde trajecten lineair blijven en coherent zijn tijdens de multi-stap oplossing, wordt een consistentieverlies toegevoegd.
   • Dit verlies dwingt de ODE-oplosser (Ordinary Differential Equation) om, na 2 tot 10 discretisatiestappen, het originele schone beeld te reconstrueren.
   • De totale loss functie is: $L = L_S + \lambda_{MCT} L_{MCT}$, waarbij $L_S$ de standaard snelheidsmatching is en $L_{MCT}$ de consistentie van de trajecten garandeert.

4. Inferentie:

   • Het model gebruikt een eerste-orde ODE-oplosser (Euler-integratie).
   • Door de lineaire aard van de stroom en de consistentie-beperkingen, kan het model hoogwaardige resultaten behalen met slechts 1 tot 4 sampling-stappen, in tegenstelling tot de tientallen of honderden stappen die vaak nodig zijn bij diffusion-modellen.

Belangrijkste Bijdragen

• Unificatie van Paradigma's: IR-Flow overbrugt de kloof tussen discriminatieve en generatieve restauratie door een directe transportstroom te definiëren, waardoor het de voordelen van beide werelden combineert.
• Efficiënte Snelheidsleer: De introductie van het Cumulatieve Veldsnelheid (CVF) versnelt de convergentie van het leerproces aanzienlijk en vermindert de modelcomplexiteit.
• Traject-Consistentie: De MCT-loss zorgt ervoor dat de tussenliggende staten op een rechte lijn blijven, wat de balans tussen vervorming (distortion) en perceptuele kwaliteit optimaliseert met zeer weinig stappen.
• Robuustheid: Het model toont sterke generalisatievermogen voor out-of-distribution (OOD) degradaties dankzij de multilevel distributiestromen.

Resultaten

De auteurs evalueren IR-Flow op diverse taken, waaronder regenverwijdering (deraining), ruisreductie (denoising) en regendruppelverwijdering.

• Kwalitatieve en Kwantitatieve Prestaties:
  • IR-Flow bereikt state-of-the-art (SOTA) resultaten op distortion-metrics (PSNR, SSIM) en perceptuele metrics (LPIPS).
  • Op de Rain100H dataset voor regenverwijdering behaalt het een PSNR van 32.69 met slechts 1 stap (NFE), terwijl concurrerende diffusion-methoden zoals IR-SDE vaak 100 stappen nodig hebben voor vergelijkbare resultaten.
  • Bij ruisreductie (denoising) op het McMaster-dataset overtreft het bestaande methoden zoals DnCNN en FFDNet, en benadert het de prestaties van IR-SDE maar dan met een factor 100 snellere inferentie.
• Efficiëntie:
  • De inferentie is 100x sneller dan IR-SDE en aanzienlijk sneller dan Resfusion en RDDM, terwijl de modelgrootte (parameters en MACs) vergelijkbaar blijft met de basisarchitectuur.
• Generalisatie:
  • Het model presteert goed op OOD-data (bijv. getraind op zware regen, getest op lichte regen of echte internetbeelden), wat aantoont dat het effectief omgaat met variërende degradatieniveaus.

Betekenis en Impact

IR-Flow is een doorbraak omdat het de traditionele afweging (trade-off) tussen efficiëntie en kwaliteit in beeldherstel doorbreekt.

• Het lost het probleem op van trage inferentie bij generatieve modellen zonder in te leveren op de perceptuele kwaliteit.
• Het biedt een fysiek interpreteerbaar raamwerk (lineaire transportstromen) dat makkelijker te trainen en te optimaliseren is dan complexe SDE-gebaseerde benaderingen met gekoppelde ruis-residu's.
• De methode is breed toepasbaar en kan potentieel worden uitgebreid naar andere restauratietaken zoals super-resolutie en ontwaziging, zoals aangetoond in de experimenten.

Kortom, IR-Flow levert een efficiënt, flexibel en hoogpresterend raamwerk dat de toekomst van real-time, hoogkwalitatieve beeldherstelling vormgeeft.