Technische Zusammenfassung: Amortisierte Führung für Bildinpainting mit vortrainierten Diffusionsmodellen (AID)

Problemstellung

Bildinpainting ist eine bedingte Generierungsaufgabe, bei der das Ziel darin besteht, fehlende Bereiche eines Bildes wiederherzustellen und dabei Konsistenz mit den sichtbaren Pixeln zu wahren. Obwohl vortrainierte Diffusionsmodelle diese Fähigkeit erheblich vorangebracht haben, stehen bestehende Ansätze vor einer Dichotomie hinsichtlich Effizienz und Anpassungsfähigkeit:

1. Aufgabenspezifisches Training: Methoden wie LaMa, MAT und BrushNet trainieren dedizierte Modelle oder große Anpassungsmodulen. Zwar amortisiert dies die Inferenzkosten, erfordert jedoch erhebliche Trainingszeit und fügt eine signifikante Modellkomplexität hinzu.
2. Instanzspezifische Optimierung: Methoden wie RePaint, DPS und MCG halten den vortrainierten Diffusions-Backbone fest, passen ihn jedoch separat für jedes maskierte Bild zur Laufzeit an. Dies bewahrt zwar das vortrainierte Modell, verursacht jedoch wiederholte Optimierungskosten während des Einsatzes und behandelt jedes Bild als ein neues Inferenzproblem.

Es besteht eine Lücke für eine „Mittelweg"-Lösung, die die Inferenz amortisiert, ohne ein großes aufgabenspezifisches Modell zu trainieren oder zur Laufzeit eine instanzspezifische Optimierung durchzuführen.

Methodik: Amortisiertes Inpainting mit Diffusion (AID)

Die Autoren schlagen AID vor, ein Framework, das den vortrainierten Diffusions-Backbone festhält und ein kleines, wiederverwendbares Führungsmodul offline trainiert. Dieses Modul wird dann über alle maskierten Bilder zur Laufzeit ohne weitere Optimierung wiederverwendet.

1. Kontrollformulierung

Das Problem wird als deterministisches Kontrollproblem formuliert.

• Dynamik: Der umgekehrte Diffusionsprozess wird durch eine Wahrscheinlichkeitsfluss-ODE gesteuert, die von einer festen vortrainierten Score-Funktion $\hat{S}$ angetrieben wird.
• Führung: Ein aufgabenabhängiger Kontrollterm $u(t, x; \xi)$ wird zum Drift hinzugefügt, wobei $\xi = (M, y)$ die beobachtbare Maske und die sichtbaren Pixel darstellt. Die Kontrolle hängt nicht vom verborgenen sauberen Bild $x^\dagger$ ab.
• Ziel: Das Ziel ist die Minimierung eines terminalen Verlusts $\Psi$ (der die Fidelity zu den sichtbaren Pixeln und der fehlenden Region ausbalanciert) plus eines laufenden Kostenanteils, der die Größe des Führungsfelds bestraft. Dies führt zu einer Hamilton-Jacobi-Bellman-(HJB)-Gleichung, bei der die optimale Kontrolle $u^*$ deterministisch ist.

2. Theoretische Brücke: Deterministisch zu Stochastisch

Die direkte Lösung der hochdimensionalen deterministischen HJB-Gleichung ist unlösbar. Die Autoren führen ein auxiliäres randomisiertes Kontrollproblem unter Verwendung Gaußscher Policies ein, um das Problem durch Continuous-Time Reinforcement Learning (CT-RL) lernbar zu machen.

• Auxiliäres Problem: Anstelle einer deterministischen Kontrolle verwendet das System eine Gaußsche Policy $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$, wobei der Mittelwert $\mu$ das Kandidaten-Führungsgesetz darstellt.
• Optimierer-erhaltende Brücke: Das Papier beweist zwei kritische Theoreme:
  1. Genügsamkeit: Die optimale deterministische Führung $u^*$ kann durch Setzen des Policy-Mittelwerts auf $u^*$ in eine optimale Gaußsche Policy gehoben werden.
  2. Notwendigkeit: Jede optimale Gaußsche Policy für das auxiliäre Problem muss einen Mittelwert haben, der die optimale deterministische Führung $u^*$ exakt wiederherstellt.
• Bedeutung: Diese Brücke stellt sicher, dass das Erlernen einer stochastischen Policy (was gradientenbasierte Updates ermöglicht) das Ziel der deterministischen Führung, die beim Einsatz verwendet wird, nicht verändert. Die Varianz der Gaußschen Policy ist fest und dient nur als Lerninstrument, nicht als Entropie-Regularisierer.

3. Amortisierter Actor-Critic-Algorithmus

Basierend auf der theoretischen Brücke entwickeln die Autoren einen vollständig datengesteuerten Actor-Critic-Algorithmus:

• Parametrisierung: Der Actor (Führungsmittelwert $\hat{\mu}_\phi$) und der Critic (Wertfunktion $\hat{V}_\theta$) werden durch neuronale Netze parametrisiert.
• Lernen: Der Algorithmus verwendet Martingal-Orthogonalitätsbedingungen, die aus der auxiliären HJB-Gleichung abgeleitet sind. Er aktualisiert Actor und Critic mittels stochastischer Approximation basierend auf gesampelten Trajektorien aus der Aufgabenverteilung.
• Einsatz: Sobald trainiert, wird das Führungsmodul deterministisch eingesetzt. Für ein neues maskiertes Bild verwendet das System das gelernte mittlere Führungsmaß $\hat{\mu}_\phi$ innerhalb des zeitumgekehrten Lösers und verwirft das Policy-Rauschen.

Hauptbeiträge

1. Methodik: AID führt ein kontrolliertes Führungs-Framework ein, das die Inferenz amortisiert, indem ein leichtgewichtiges Führungsmodul offline trainiert wird, das dann über Aufgaben hinweg wiederverwendet wird, ohne instanzspezifische Optimierung.
2. Theorie: Das Papier etabliert eine rigorose „Policy-Äquivalenz-Brücke" zwischen dem deterministischen Führungsproblem und einer auxiliären Gaußschen Formulierung. Es beweist, dass das auxiliäre Problem optimierer-erhaltend ist, was die Verwendung stochastischer Actor-Critic-Methoden zum Erlernen eines deterministischen Kontrollfelds rechtfertigt.
3. Effizienz: Der Ansatz fügt weniger als 1 % trainierbarer Parameter relativ zum vortrainierten Backbone hinzu, verbessert gleichzeitig das Verhältnis von Qualität zu Geschwindigkeit erheblich.

Empirische Ergebnisse

Die Autoren bewerteten AID auf AFHQv2, FFHQ (unter Verwendung der Pixelraum-EDM-Pipeline) und ImageNet (unter Verwendung der Latentraum-EDM2-Pipeline).

• Leistung im Vergleich zu Baselines mit festem Backbone: AID übertrifft konsistent starke Baselines mit festem Backbone (Unguided, Replacement, MCG, DPS) über mehrere Metriken (PSNR, SSIM, LPIPS) und Maskentypen (freiform, zentriert, Streifen).
• Leistung im Vergleich zu RePaint: AID übertrifft RePaint (eine starke iterative Optimierungs-Baseline) in der Qualität, während es nur ~10 % des Sampling-Budgets verwendet (z. B. 35 Schritte gegenüber 350 Schritten für RePaint).
• Leistung im Vergleich zu LatentPaint: Auf ImageNet übertrifft AID LatentPaint (eine wiederverwendbare Latentraum-Anpassungsmethode) in allen Metriken, während es ein Führungsmodul mit weniger als 45 % der trainierbaren Parameter verwendet.
• Latenz: AID unterstützt einen Niedrig-Latenz-Betriebspunkt (12 Schritte), der die Inferenzzeit weiter reduziert, während die Leistung besser als bei Nicht-RePaint-Baselines bleibt und mit RePaint wettbewerbsfähig bleibt.
• Generalisierung: Dasselbe trainierte Führungsmodul überträgt sich effektiv auf nicht gesehene Maskentypen und verschiedene Sampling-Budgets ohne Nachtraining.

Bedeutung und Behauptungen

Das Papier behauptet, dass AID das erste amortisierte kontrollierte Führungs-Framework für Bildinpainting ist, das:

1. Den vortrainierten Diffusions-Backbone eingefroren hält.
2. Nur ein kleines, wiederverwendbares Führungsmodul trainiert.
3. Eine exakte optimierer-erhaltende Brücke vom stochastischen Actor-Critic-Lernen zum deterministischen Führungsfeld bietet, das beim Einsatz verwendet wird.

Die Autoren positionieren diese Arbeit als einen Schritt hin zur Nutzung von Kontroll- und RL-basierten Frameworks, um theoretisch fundierte Methoden zur Untersuchung und Verbesserung von Diffusionsmodellen bereitzustellen. Sie betonen, dass Continuous-Time RL sowohl der rigorosen Analyse als auch dem praktischen Algorithmus-Entwurf zugrunde liegt und eine prinzipielle Alternative zu heuristischer instanzspezifischer Optimierung oder schwerem aufgabenspezifischem Nachtraining bietet.