Technische Samenvatting: Diffusiemodel-gebaseerde Beeldrestauratie via Interactief Leren van Determinisme en Stochastiek

Probleemstelling

Beeldrestauratie (IR) streeft naar het reconstrueren van latente visuele informatie uit gedegradeerde observaties, een taak die wordt bemoeilijkt door de significante distributieverschillen tussen verschillende typen degradatie (bijv. ruis, regen, schaduwen, lage lichtintensiteit) in de feature-ruimte. Bestaande deep learning-methoden, inclusief CNN's, Transformers en GAN's, hebben vaak moeite om te generaliseren over diverse IR-taken zonder taakspecifieke herontwerpen. Hoewel diffusiemodellen een hoge getrouwheid en interpreteerbaarheid bieden, kampen huidige benaderingen met twee primaire beperkingen:

1. Inefficiëntie in IR: Standaard diffusiemodellen starten het omgekeerde proces vanuit zuivere ruis, wat inefficiënt is voor IR-taken waarbij de input al structurele informatie bevat.
2. Statische Modellering van Determinisme en Stochastiek: Eerdere pogingen om deterministische (residuele) en stochastische (ruis-gebaseerde) diffusie te verenigen, zoals het Residual Denoising Diffusion Model (RDDM), maken gebruik van onafhankelijke U-Nets voor elk pad. Deze statische scheiding slaagt er niet in om de onderscheidende waarschijnlijkheidsverdelingen van verschillende typen degradatie te vangen en kan excessieve diversiteit introduceren die nauwkeurige reconstructie belemmert.

Methodologie

De auteurs stellen een nieuw diffusiekader voor dat Interactive Learning of Determinism and Stochasticity (ILDS) gebruikt om deze beperkingen aan te pakken. De kernarchitectuur maakt gebruik van een dual-path netwerk gebaseerd op twee eenvoudige U-Nets.

1. Constructie van het Forward Diffusion Proces

In tegen tegenstelling tot standaard diffusie, die alleen ruis toevoegt, incorporeert het voorgestelde forward proces zowel residuen (determinisme) als ruis (stochastiek). De transitie wordt gedefinieerd als:
$$x_t = x_{t-1} + \alpha_t x_{res} + \beta_t \epsilon_{t-1}$$
waarbij $x_{res} = x_{in} - x_0$ het degradatie-residue vertegenwoordigt, en $\epsilon$ Gaussische ruis is. Dit stelt de forward traject mogelijk om degradatie-gerelateerde structuren te behouden, waardoor het reverse proces kan starten vanuit toestanden die dichter bij het doel liggen dan zuivere ruis, wat het vereiste aantal sampling-stappen vermindert.

2. Interactief Leermechanisme

Om de beperkingen van onafhankelijke modellering in RDDM te overwinnen, introduceren de auteurs een interactieve leerstrategie binnen de U-Net architectuur:

• Dual-Path Predictie: Twee U-Nets worden gebruikt om respectievelijk de ruis ($\epsilon_\theta$) en het residu ($x_{res}^\theta$) te voorspellen.
• Feature Fusie: In plaats van geïsoleerd te opereren, delen de encoders van de twee U-Nets features. Specifiek worden na de eerste downsampling-block de features van beide paden geconcateneerd en verwerkt via $1 \times 1$ convoluties ($\tilde{h}_{res}$ en $\tilde{h}_\epsilon$).
• Dynamische Tijd-Afhankelijke Modellering: De timestep-inputs voor de twee paden worden dynamisch gemoduleerd door parameters die de temporele accumulatie ($\bar{\alpha}_t$) en ruisintensiteit ($\bar{\beta}_t$) controleren. Dit stelt het model in staat om de bijdrage van deterministische en stochastische componenten adaptief aan te passen op basis van de specifieke feature-waarschijnlijkheidsverdeling van het type degradatie.
• Gedeelde Context: De geïnteracteerde features bieden gedeelde cross-path contextuele informatie tijdens de decoding-fase, wat ervoor zorgt dat de deterministische sturing en de stochastische diversiteit elkaar beïnvloeden.

3. Training en Generatie

Het model wordt getraind om de Kullback-Leibler (KL) divergentie tussen de ware posterior en de geleerde reverse transitie te minimaliseren. De uiteindelijke loss-functie is de som van de kwadratische fouten voor zowel de ruis- als de residu-voorspellingen:
$$L = L_{res}(\theta) + L_{\epsilon}(\theta)$$
Opvallend genoeg zijn deze componenten ongewogen, aangezien het interactieve mechanisme een zelf-adaptieve evenwicht tijdens de training faciliteert. Tijdens de inferentie genereert het model afbeeldingen van hoge kwaliteit door iteratief te samplen uit de conditionele distributie, waarbij het proces volledig deterministisch wordt wanneer de ruisschaalfactor $\eta = 0$ is.

Belangrijkste Bijdragen

1. Nieuw Diffusiekader: Het ontwerp van een framework dat interactief zowel deterministische als stochastische aspecten van beelddegradatie modelleert, waardoor het diffusieproces kan adapteren aan taken die ofwel hoge getrouwheid (precisie) of diversiteit vereisen.
2. Dynamische Tijd-Afhankelijke Modellering: De incorporatie van temporele accumulatie en ruisintensiteit parameters om timestep-inputs voor duale diffusiepaden dynamisch te moduleren, wat de generalisatie over diverse IR-scenario's verbetert.
3. Efficiënte Dual-Path Architectuur: De implementatie van de dual-path diffusie met gebruik van slechts twee eenvoudige U-Nets met gedeelde encoder-features. Dit ontwerp bereikt competitieve prestaties met minimale parameter-overhead en een hoge schaalbaarheid.

Experimentele Resultaten

De voorgestelde methode werd geëvalueerd op vier representatieve IR-taken: Beeldruisverwijdering (Denoising), Schaduwverwijdering (Shadow Removal), Verlichting verbeteren (Low-Light Enhancement) en Beeldontregeningsverwijdering (Deraining).

• Prestaties: De methode behaalde superieure of competitieve prestaties over alle taken:
  • Denoising (SIDD): Bereikte de beste PSNR (41.67 dB) en SSIM (0.9385), waarmee het RDDM en andere baselines overtrof.
  • Shadow Removal (ISTD): Behaalde de beste PSNR (29.29 dB) en MAE (8.776), met de op één na beste SSIM.
  • Deraining (Raindrop): Bereikte de beste PSNR (31.30 dB) en SSIM (0.9460), waarbij regenstrepen succesvol werden verwijderd waar andere methoden artefacten achterlieten.
  • Low-Light Enhancement (LOL): Behaalde de op één na beste prestatie (PSNR 24.04, SSIM 0.9118), waarbij het qua metrieken net achter LLFormer bleef maar sterke structurele preservatie behield.
• Efficiëntie: Het interactieve leren verbeterde de sampling-efficiëntie aanzienlijk. Het model had slechts 1–4 sampling-stappen nodig om resultaten van hoge kwaliteit te bereiken, vergeleken met honderden of duizenden stappen voor andere diffusie-gebaseerde methoden (bijv. DiffusionVI vereiste 1000 stappen).
• Ablatie-studies:
  • Het verwijderen van de interactie-module resulteerde in prestatievermindering over alle taken, wat het voordeel van feature fusie bevestigt.
  • Het gebruik van dynamische tijd-afhankelijke modellering ($\bar{\alpha}_t, \bar{\beta}_t$) presteerde beter dan statische timestep-modellering ($t$).
  • Optimale sampling-stappen varieerden per taak (1 voor low-light, 2 voor denoising, 3 voor shadow removal, 4 voor deraining).

Betekenis en Claims

Het artikel claimt dat het voorgestelde framework een verenigde oplossing biedt voor multi-task beeldrestauratie die een balans vindt tussen reconstructienauwkeurigheid en diversiteit. Door deterministische en stochastische modellering te integreren via interactief leren, biedt de aanpak een sterkere interpreteerbaarheid dan GAN-gebaseerde methoden en een betere generalisatie dan taakspecifieke CNN's of Transformers.

De auteurs erkennen beperkingen, met name bij low-light enhancement taken met grote illuminatie-gradiënten, waarbij de helderheidsrestauratie minder adequaat kan zijn vergeleken met gespecialiseerde baselines. Ze merken ook op dat de huidige methode afhankelijk is van gepaarde clean-degraded data voor supervised training. Toekomstig werk is gericht op het optimaliseren van diffusievergelijkingen voor illuminatie-gevoelige taken en het ontwikkelen van unsupervised trainingstrategieën om de afhankelijkheid van data te verminderen.

Al met al demonstreert de studie dat een eenvoudig, interactief dual-path diffusie-architectuur effectief kan omgaan met de heterogeniteit van beelddegradatieprocessen, en een schaalbare en efficiënte alternatief biedt voor diverse beeldrestauratie-uitdagingen.