Your Pre-trained Diffusion Model Secretly Knows Restoration

Probleemstelling

All-in-One Restauratie (AiOR) richt zich op het herstellen van afbeeldingen of video's die zijn beschadigd door diverse degradaties (zoals mist, regen, sneeuw, bewegingsonscherpte en weinig licht) met één unifyend model. Hoewel vooraf getrainde Diffusiemodellen (zoals Stable Diffusion, FLUX en WAN) sterke generatieve priors bezitten die ideaal lijken voor restauratie, zijn bestaande methoden beperkt:

1. Fine-tuning: Aanpassen van de volledige backbone vereist veel training en kan de sterke, algemene priors van het model verzwakken, wat leidt tot minder robuustheid bij complexe of onbekende degradaties.
2. ControlNet-achtige modules: Deze voegen extra conditionering toe, maar kunnen de vooraf getrainde prior verstoren als de module niet robuust is.
3. Bestaande Prompting-strategieën: Methoden die gebruikmaken van natuurlijke taal prompts of het optimaliseren van tekst-tokens (zoals bij image editing) werken slecht voor restauratie. De modellen neigen er dan toe om alleen ruis te onderdrukken en terug te keren naar de beschadigde invoer, in plaats van de degradatie te verwijderen.

De kernvraag is: Bevatten de priors van grote diffusiemodellen restauratiekennis, en zo ja, hoe kunnen we deze toegankelijk maken zonder de backbone aan te passen?

Methodologie

De auteurs stellen een nieuwe aanpak voor waarbij ze de vooraf getrainde diffusie-backbone volledig bevriezen en alleen lichte "learned prompts" trainen. De methode bestaat uit twee cruciale innovaties:

1. Optimalisatie in Embedding-Space (in plaats van Token-Space)

• Observatie: Het optimaliseren van prompts in de token-ruimte (tekst-tokens) of het gebruik van natuurlijke taal fails om restauratiegedrag op te wekken. Het model blijft vastzitten in de beschadigde toestand.
• Oplossing: De auteurs optimaliseren direct de context-embeddings aan de uitgang van de tekst-encoder (voordat deze de diffusion backbone binnenkomt). Door deze embeddings direct te leren, kunnen ze de frozen backbone dwingen om van een ruizige, beschadigde latent-toestand ($z_{deg}$) naar een schone toestand ($z_{clean}$) te bewegen. Dit onthult dat de restauratiekennis wel degelijk in het model zit, maar niet toegankelijk is via de standaard token-interface.

2. Oplossen van Trajectoïde-Mismatch (Diffusion Bridge)

• Het Probleem: Een naïeve training waarbij de voorwaartse noiseprocessus begint bij de beschadigde afbeelding ($z_{deg}$) en ruis toevoegt, creëert een mismatch met de inferentie. Tijdens inferentie moet het model stap voor stap van $z_{deg}$ naar $z_{clean}$ gaan. De getrainde prompt ziet tijdens training echter alleen toestanden die rond $z_{deg}$ geankerd zijn, wat leidt tot instabiliteit en artefacten tijdens het genereren.
• De Oplossing: De auteurs introduceren een Diffusion Bridge formulering. In plaats van een simpele noiseprocessus, trainen ze de prompt op een pad dat een gladde interpolatie vormt tussen de beschadigde en de schone toestand.
  • Ze gebruiken een specifieke variant genaamd EBR (Energy-Oriented Diffusion Bridge).
  • Dit pad begint met een ruizige versie van de beschadigde input en beweegt monotoon (stelselmatig) naar de schone toestand, waarbij het niveau van ruis afneemt naarmate het signaal schoner wordt.
  • Dit zorgt ervoor dat de prompt tijdens training en inferentie dezelfde reeks tussenliggende toestanden ziet, wat de stabiliteit en prestaties drastisch verbetert.

Implementatie:

• Er wordt één aparte prompt geleerd per degradatietype (bijv. mist, regen).
• Voor gemengde degradaties worden de prompts gecombineerd door hun voorspelde snelheden (velocities) te middelen.
• De prompts worden als een "residual" toegevoegd aan de null-text context in de cross-attention laag, wat efficiënter is dan het volledig herschrijven van de embeddings.

Belangrijkste Bijdragen

1. Onthulling van Restauratiekennis: Het paper toont aan dat vooraf getrainde diffusion-modellen (zoals FLUX en WAN) inherent restauratiekennis bezitten, maar dat deze niet toegankelijk is via natuurlijke taal of token-optimalisatie, maar wel via directe optimalisatie van de embedding-ruimte.
2. Trajectoïde-Alignement: Het identificeren en oplossen van de train-test mismatch door het gebruik van een diffusion bridge (EBR), wat essentieel is voor stabiele prompt-learning bij restauratie.
3. Efficiëntie en Generalisatie: Het introduceren van een framework dat zware fine-tuning en specifieke control-modules overbodig maakt. Het systeem gebruikt slechts lichte prompts om enorme modellen (1.3B en 12B parameters) om te zetten in krachtige restauratiemodellen.

Resultaten

De methode is getest op de FLUX (12B, afbeeldingen) en WAN (1.3B, video) modellen.

• Kwalitatieve Prestaties: De methode overtreft state-of-the-art AiOR-methoden (zoals AutoDIR, PixWizard, FoundIR) in visuele kwaliteit, vooral bij complexe en onbekende degradaties.
• Kwantitatieve Prestaties:
  • Op in-distribution data scoort de methode competitief op zowel vervormingsmetrieken (PSNR, SSIM) als perceptuele metrieken (LPIPS, CLIPIQA, MUSIQ).
  • Op out-of-distribution (OOD) en gemengde degradaties (bijv. regen + mist, weinig licht + onscherpte) presteert de methode consistent beter dan bestaande methoden, wat aantoont dat de generalisatievermogen van de frozen prior behouden blijft.
  • Voor video-restauratie (WAN) behaalt de methode de beste scores op de DOVER video-kwaliteitsmetriek, dankzij de sterke temporale priors van het model.
• Ablatiestudies: Experimenten tonen aan dat naïeve training slecht presteert, dat DDBM-bruggen beter zijn maar nog steeds onderpresteren, en dat de EBR-brug de beste resultaten levert. Ook blijkt dat residual prompt injection efficiënter is dan directe embedding-optimalisatie.

Significantie

Dit werk is significant omdat het een fundamenteel inzicht biedt in hoe vooraf getrainde generatieve modellen kunnen worden gebruikt voor low-level vision taken zonder hun kennis te "vergeten" door zware fine-tuning.

• Het demonstreert dat conditioning de sleutel is: de juiste manier om een model te sturen (via embeddings in plaats van tokens) kan het verschil maken tussen falen en succes.
• Het biedt een lichtgewicht, plug-and-play oplossing voor restauratie die niet afhankelijk is van expliciete degradatiemodellen (zoals in traditionele methoden) of zware training.
• Het opent de deur voor het gebruik van enorme, generieke multimodale modellen voor specifieke, complexe visuele taken met minimale aanpassing, wat de toegankelijkheid en schaalbaarheid van geavanceerde restauratietechnologie vergroot.