Technische Zusammenfassung: Metropolis-Adjustierte Diffusionsmodelle (MADM)

Problemstellung

Das Sampling aus score-basierten Diffusionsmodellen ist inhärent mit einer Verzerrung behaftet, die aus zwei Quellen resultiert: dem Approximationsfehler der gelernten Score-Funktion und dem durch numerische Integratoren eingeführten Zeit-Diskretisierungsfehler. Im standardmäßigen Predictor-Corrector-Rahmenwerk bewegt ein Predictor-Schritt (oft eine Probability Flow ODE) eine Stichprobe zwischen Rauschpegeln, gefolgt von einem Corrector-Schritt (typischerweise der Unadjusted Langevin Algorithmus, ULA), um die Stichprobe auf dem aktuellen Rauschpegel zu verfeinern.

Während ULA die Score-Funktion nutzt, um Bewegungen vorzuschlagen, ist es selbst ein verzerrter Sampler, da es einen kontinuierlichen Langevin-Diffusionsprozess diskretisiert. Folglich führt der „Corrector" im Predictor-Corrector-Schema seine eigene Diskretisierungsverzerrung ein und lässt die Zielrandverteilung nicht invariant. Standard-Metropolis-Hastings (MH) oder Barker-Akzeptanz-Ablehnungs-Mechanismen, die eine solche Verzerrung in Markov-Chain-Monte-Carlo (MCMC) korrigieren, können nicht direkt auf Diffusionsmodelle angewendet werden. Dies liegt daran, dass MH das Verhältnis der Zielwahrscheinlichkeitsdichten, $p_t(\tilde{x})/p_t(x)$, bewerten muss, was in Diffusionsmodellen nicht handhabbar ist, in denen nur die Score-Funktion $s_\phi(x, t) \approx \nabla_x \log p_t(x)$ verfügbar ist.

Methodik

Die Autoren schlagen Metropolis-Adjustierte Diffusionsmodelle (MADM) vor, ein Rahmenwerk, das verzerrte ULA-Correctors durch akzeptanz-ablehnungs-korrigierte Langevin-Schritte ersetzt. Die Kerninnovation ist eine Methode zur Berechnung der notwendigen Akzeptanzwahrscheinlichkeit ausschließlich unter Verwendung der Score-Funktion, wodurch die Notwendigkeit des nicht handhabbaren Dichteverhältnisses umgangen wird.

1. Dichteverhältnis-Identität

Die Autoren zeigen, dass der Logarithmus des Ziel-Dichteverhältnisses als Linienintegral der Score-Funktion entlang des Pfades zwischen dem aktuellen Zustand $x$ und dem Vorschlag $\tilde{x}$ ausgedrückt werden kann:
$$\log \left( \frac{p_t(\tilde{x})}{p_t(x)} \right) = \int_0^1 \langle s(x + u(\tilde{x} - x), t), \tilde{x} - x \rangle \, du$$
Diese Identität ermöglicht die Schätzung des Dichteverhältnisses unter Verwendung ausschließlich von Score-Auswertungen.

2. Exakte Korrektur mittels Bernoulli-Fabrik

Um eine exakte Korrektur ohne Bewertung der Dichte zu erreichen, schlagen die Autoren einen Algorithmus vor, der auf einer zwei-Münzen-Bernoulli-Fabrik (Gonçalves et al., 2017) basiert.

• Annahme: Es existiert eine Schranke $C(x, \tilde{x})$, sodass der Betrag des Integranden im Linienintegral durch $C$ beschränkt ist.
• Mechanismus: Der Algorithmus konstruiert einen erwartungstreuen Schätzer $W$ für das Exponential des Integrals (skaliert mit der Schranke) unter Verwendung einer Poisson-abgeschnittenen Potenzreihe.
• Verfahren: Der Algorithmus verwendet ein „zwei-Münzen"-Verfahren, um eine Barker-Akzeptanz-Ablehnungs-Entscheidung zu simulieren. Zuerst wird mit einer festen Wahrscheinlichkeit $\alpha'$ abgelehnt; falls nicht abgelehnt, wird eine Münze mit der Wahrscheinlichkeit $W$ geworfen. Wenn die Münze erfolgreich ist, wird der Vorschlag akzeptiert; andernfalls beginnt die Runde von neuem.
• Ergebnis: Dieses Verfahren ergibt einen exakten Barker-adjustierten Langevin-Corrector, der die Zielverteilung $p_t$ invariant lässt, sofern die Score-Funktion exakt ist und die Schrankenannahme gilt.

3. Effiziente Approximation mittels Simpson-Regel

In Anbetracht dessen, dass die exakte Bernoulli-Fabrik-Methode eine zufällige und potenziell große Anzahl von Score-Auswertungen erfordert, schlagen die Autoren eine skalierbare Approximation vor.

• Methode: Das Linienintegral wird deterministisch mittels der Simpson-1/3-Regel (eine Newton-Cotes-Quadraturformel) approximiert. Dies erfordert die Auswertung der Score-Funktion am aktuellen Zustand, am Vorschlag und am Mittelpunkt.
• Akzeptanz: Die resultierende Schätzung wird in die Standard-Metropolis-Hastings-Akzeptanzwahrscheinlichkeit eingesetzt.
• Eigenschaften: Obwohl dies eine kleine Verzerrung durch Quadratur-Abschneidung einführt, erreicht die Approximation eine Genauigkeit der Ordnung $O(h^{5/2})$ in der Schrittweite $h$. Sie ist rechnerisch günstig und erfordert nur eine zusätzliche Score-Auswertung am Mittelpunkt pro Vorschlag.

Hauptbeiträge

1. MADM-Rahmenwerk: Einführung eines Metropolis-adjustierten Corrector-Rahmenwerks für Diffusionsmodelle, das systematisch verzerrte ULA-Schritte ersetzt.
2. Exakte Methode: Die erste exakte Methode zur Anpassung von Langevin-Korrekturen in Diffusionsmodellen unter Verwendung einer zwei-Münzen-Bernoulli-Fabrik, die nur die Score-Funktion und eine Schranke für das Linienintegral erfordert.
3. Skalierbare Approximation: Eine praktische Simpson-Regel-Approximation, die eine Genauigkeit hoher Ordnung ($O(h^{5/2})$) mit nahezu marginalen Kosten (eine zusätzliche Score-Auswertung) erreicht.
4. Empirische Validierung: Demonstration verbesserter Stichprobenqualität auf synthetischen und Bilddatensätzen.

Ergebnisse

Die Autoren bewerten MADM auf synthetischen 2D-Datensätzen und Standard-Bild-Benchmarks (CIFAR-10, FFHQ, AFHQv2, ImageNet-64) unter Verwendung vortrainierter Modelle von Karras et al. (2022).

• Synthetische Datensätze: Auf 2D-Datensätzen (Spirale, Trichter, Sierpinski-Dreieck, Windmühle) erzeugte der Standard-Ancestral-Sampler (entspricht einem Predictor mit einem ULA-Corrector) Stichproben mit signifikanten Ausreißern außerhalb des Supports der wahren Datenverteilung. MADM-Corrector-Schritte bewegten diese Stichproben erfolgreich zurück in den Daten-Support und eliminierten Ausreißer.
• Bilddatensätze:
  • Für unbedingte Modelle (CIFAR-10, FFHQ, AFHQv2) führte die Anwendung eines MADM-Corrector-Schritts pro Predictor-Schritt zu konsistenten, wenn auch bescheidenen, Verbesserungen der Fréchet Inception Distance (FID) im Vergleich zu deterministischen Baselines (Heun- und Euler-ODE-Löser).
  • Für das klassenbedingte ImageNet-64-Modell waren die Gewinne ausgeprägter. Die Verwendung von MADM mit einem Euler-Predictor erzielte größere FID-Verbesserungen als die Verwendung eines höherordentlichen Heun-Predictors ohne Korrektur. Dies legt nahe, dass MH-Anpassungen effektiver sind, um Diskretisierungsfehler zu mildern, als einfach die Ordnung des ODE-Lösers zu erhöhen, wenn die Anzahl der Schritte groß ist.
• Rechenkosten: Die exakte Methode (Algorithmus 2) erwies sich aufgrund der zufälligen Anzahl erforderlicher Score-Auswertungen als rechenintensiv. Folglich nutzten die Experimente hauptsächlich die Simpson-Regel-Approximation (Algorithmus 3), die konsistente Gewinne bei überschaubarem Overhead bot.

Bedeutung und Behauptungen

Die Arbeit behauptet, die erste exakte Methode zur Anpassung von Langevin-Korrekturen in Diffusionsmodellen zu liefern und adressiert die inhärente Verzerrung von ULA-Correctors. Durch die Nutzung der Score-Funktion zur Schätzung von Dichteverhältnissen über Linienintegrale ermöglicht MADM die Verwendung rigoroser MCMC-Akzeptanzkriterien in der generativen Modellierung, ohne Zugriff auf die nicht handhabbare Datendichte zu benötigen.

Die Autoren positionieren MADM als natürliche Erweiterung des Predictor-Corrector-Rahmenwerks, die die Stichprobenqualität verbessert. Sie stellen fest, dass die exakte Methode zwar theoretisch bedeutsam ist, die Simpson-Regel-Approximation jedoch eine praktische, skalierbare Lösung bietet, die die FID auf Bild-Benchmarks konsistent verbessert. Die Arbeit geht von einer perfekten Score-Schätzung aus und erkennt dies als inhärente Einschränkung von Diffusionsmodellen an, nicht als spezifischen Fehler ihrer Methode. Als zukünftige Arbeit wird die Entwicklung effizienter Varianten der exakten Methode oder ihrer Subroutinen vorgeschlagen.