Your Pre-trained Diffusion Model Secretly Knows Restoration

1. Problemstellung

Das Ziel von All-in-One Restoration (AiOR) ist die Wiederherstellung von Bildern oder Videos, die durch diverse Degradationen (z. B. Dunst, Regen, Schnee, Bewegungsunschärfe, schlechte Beleuchtung) beschädigt sind, mithilfe eines einzigen einheitlichen Modells.

Obwohl vortrainierte Diffusionsmodelle (wie Stable Diffusion, FLUX, WAN) starke Priors für natürliche Bilder besitzen und für Bildbearbeitungsaufgaben hervorragend funktionieren, scheitern sie bei der direkten Übertragung auf Restaurierungsaufgaben:

• Fehlschlagen von Text-Prompts: Natürlichsprachliche Prompts (z. B. „Entferne den Dunst") oder optimierte Token-Prompts (Textual Inversion) führen dazu, dass das Modell zwar Rauschen entfernt, aber die Degradation selbst nicht beseitigt. Das Ergebnis ist oft ein entrauschtes, aber immer noch verzerrtes Eingabebild.
• Ineffizienz bestehender Ansätze: Bisherige Methoden nutzen entweder Fine-Tuning des gesamten Backbones (rechenintensiv, Risiko des „Catastrophic Forgetting" der generativen Priors) oder ControlNet-ähnliche Module (erfordern zusätzliche Architektur und Training).
• Trajektorien-Mismatch: Naives Prompt-Learning führt zu einer Diskrepanz zwischen dem Trainingsprozess (wo der Zustand vom degradierten Bild aus mit Rauschen angereichert wird) und dem Inferenzprozess (wo das Modell vom verrauschten Zustand zum sauberen Bild denoisen muss). Dies führt zu Instabilität und Artefakten.

2. Methodik

Die Autoren schlagen einen Ansatz vor, der keine Anpassung des Diffusions-Backbones erfordert, sondern stattdessen leichte, erlernte Prompts in einem spezifischen Embedding-Raum nutzt.

A. Optimierung im Embedding-Raum statt im Token-Raum

Das zentrale Erkenntnis des Papers ist, dass Restaurierungswissen im vortrainierten Diffusionsprior vorhanden, aber über den Standard-Text-Encoder nicht zugänglich ist.

• Token-Space (versagt): Die Optimierung von Eingabe-Token-Embeddings (vor dem Text-Encoder) führt nicht zur gewünschten Restaurierung.
• Embedding-Space (erfolgreich): Die Autoren optimieren direkt den Kontext-Embedding-Vektor ($p$) am Ausgang des Text-Encoders (vor der Cross-Attention im Backbone). Dieser Vektor wird als feste Eingabe in das eingefrorene Diffusionsmodell eingespeist. Dies zwingt das Modell, den verrauschten degradierten Zustand ($z_{deg}$) direkt auf das saubere Ziel ($z_{clean}$) zu lenken.

B. Lösung des Trajektorien-Mismatches (Diffusion Bridge)

Ein naives Training auf degradierten Bildern mit Rauschaddition führt zu einem Mismatch zwischen Trainings- und Inferenzpfaden. Um dies zu lösen, führen die Autoren eine Diffusion-Bridge-Formulierung ein:

• Statt nur vom degradierten Bild zu Rauschen zu gehen, konstruieren sie einen Pfad, der monoton vom verrauschten degradierten Zustand zum sauberen Zustand übergeht.
• Sie nutzen eine EBR-Bridge (Energy-Oriented Diffusion Bridge), die Zwischenzustände der Form $z_t = \mu(t) + s(t)\epsilon$ definiert, wobei $\mu(t)$ eine Interpolation zwischen $z_{deg}$ und $z_{clean}$ ist.
• Dies stellt sicher, dass der gelernte Prompt während des Trainings und der Inferenz auf ähnlichen statistischen Zuständen operiert, was die Stabilität und die Korrektureffizienz drastisch verbessert.

C. Architektur und Training

• Backbones: Das Framework wird auf dem WAN 1.3B (Video) und FLUX 12B (Bild) Modell implementiert.
• Prompt-Design: Statt den vollen Kontext zu lernen, wird ein residualer Prompt gelernt, der auf einen „Null-Text"-Kontext aufgesetzt wird. Dies ist parameter-effizienter und performanter.
• Multi-Degradation: Für verschiedene Degradationstypen (Haze, Rain, etc.) werden separate Prompts gelernt. Bei gemischten Degradationen werden die Vorhersagen der Prompts während des Sampling-Prozesses gemittelt.

3. Hauptbeiträge

1. Entdeckung versteckter Fähigkeiten: Es wird gezeigt, dass große Diffusionsmodelle Restaurierungswissen besitzen, das jedoch nicht über natürliche Sprache oder Token-Optimierung, sondern nur durch direkte Optimierung der Konditionierungs-Embeddings im Text-Encoder-Ausgangsräumen zugänglich ist.
2. Identifikation und Lösung des Trajektorien-Mismatches: Die Autoren identifizieren die Instabilität des naiven Prompt-Learnings als Folge von Trainings-Inferenz-Diskrepanzen und lösen dies durch eine Bridge-basierte Trainingsstrategie (EBR), die einen kohärenten Denoising-Pfad erzwingt.
3. Leichtgewichtiges, hochleistungsfähiges Framework: Durch die Kombination aus Embedding-Optimierung und Bridge-Training werden state-of-the-art Restaurierungsergebnisse erreicht, ohne das teure Fine-Tuning des Backbones oder ControlNet-Module zu benötigen.

4. Ergebnisse

Die Methode wurde auf dem FLUX-Modell (Bilder) und dem WAN-Modell (Videos) evaluiert und mit State-of-the-Art-Methoden (wie AutoDIR, PixWizard, DFPIR, FoundIR) verglichen.

• Quantitative Leistung:
  • Die Methode erreicht wettbewerbsfähige Ergebnisse in Bezug auf Verzerrungsmetriken (PSNR, SSIM) und insbesondere in perzeptuellen Metriken (CLIPIQA, MUSIQ, DISTS), oft führend bei der visuellen Qualität.
  • Auf Out-of-Distribution (OOD) Datensätzen (z. B. reale Regen- oder Dunstszenen) und bei gemischten/unbekannten Degradationen zeigt die Methode eine starke Generalisierungsfähigkeit, die viele spezialisierte Modelle übertrifft.
  • Im Video-Bereich (WAN) übertrifft die Methode spezialisierte Video-Restaurierungsmodelle (wie ViWS-Net, AverNet) in der temporalen Konsistenz (DOVER-Metrik), da die starken temporalen Priors des vortrainierten WAN-Modells erhalten bleiben.
• Qualitative Leistung:
  • Visuelle Vergleiche zeigen, dass die Methode Artefakte vermeidet, die bei naiven Ansätzen oder SDEdit-ähnlichen Verfahren auftreten (z. B. fehlende Entfernung von Dunst oder unscharfe Strukturen).
  • Die Methode bewahrt Details und Texturen besser als Fine-Tuning-Ansätze, die oft zu einer Überanpassung an den Trainingsdatensatz neigen.

5. Bedeutung und Fazit

Das Paper liefert einen fundamentalen Einblick in die Funktionsweise vortrainierter Diffusionsmodelle für Low-Level-Vision-Aufgaben. Es widerlegt die Annahme, dass diese Modelle für Restaurierung nicht geeignet seien, und zeigt, dass das Problem nicht im Backbone, sondern in der Art und Weise liegt, wie Konditionierung erfolgt.

Die technische Bedeutung liegt in der Demonstration, dass:

1. Konditionierung ist entscheidend: Die Fähigkeit zur Restaurierung ist im Prior vorhanden, muss aber durch die richtige Konditionierung (Embedding-Space) aktiviert werden.
2. Effizienz: Hochleistungs-Restaurierung ist ohne teures Fine-Tuning oder zusätzliche Control-Net-Module möglich, was Speicher- und Rechenressourcen spart.
3. Robustheit: Der Ansatz bietet eine robuste Lösung für komplexe, reale Degradationen, da er die generativen Stärken der großen Modelle nutzt, ohne deren Priors durch übermäßiges Training zu zerstören.

Zusammenfassend stellt die Arbeit einen Paradigmenwechsel dar: Statt das Modell neu zu trainieren, wird es durch intelligente, leichte Prompt-Optimierung „entschlüsselt", um seine verborgenen Restaurierungsfähigkeiten freizulegen.