Konsistenz-Trajektorie-Modelle: Lernen der Wahrscheinlichkeitsstrom-ODE-Trajektorie von Diffusion

Konsistenzmodelle (CM) (Song et al., 2023) beschleunigen die Sampling-Prozedur von scorebasierten Diffusionsmodellen auf Kosten der Generationsqualität und verfügen über keinen natürlichen Mechanismus, um Qualität gegen Geschwindigkeit auszutauschen. Um diese Einschränkung zu überwinden, schlagen wir das Konsistenz-Trajektorien-Modell (CTM) vor, eine Verallgemeinerung, die sowohl CM als auch scorebasierte Modelle als Spezialfälle umfasst. CTM trainiert ein einzelnes neuronales Netzwerk, das in einer einzigen Vorwärtsdurchlauf – unabhängig von der Anfangs- und Endzeit – sowohl Scores (d. h. Gradienten der Log-Dichte) berechnen als auch beliebige Durchquerungen entlang der Lösungstrajektorien des Wahrscheinlichkeitsfluss-Ordinären Differentialgleichungssystems (ODE) in einem Diffusionsprozess ermöglicht. CTM erlaubt die effiziente Kombination von adversarieller Trainingsstrategie und Noise-Reduktion-Score-Matching-Verlust, wodurch die Leistung verbessert wird. Es erreicht neue State-of-the-Art-FID-Werte für Ein-Schritt-Sampling von Diffusionsmodellen auf CIFAR-10 (FID 1,73) und ImageNet bei 64×64 Auflösung (FID 1,92). Zudem ermöglicht CTM eine neue Familie von Sampling-Schemata, sowohl deterministischer als auch stochastischer Natur, die längere Sprünge entlang den ODE-Lösungstrajektorien erlauben. Es verbessert die Generationsqualität konsistent mit steigendem Rechenaufwand und vermeidet die Qualitätsdegradation, die bei CM auftritt. Im Gegensatz zu CM ermöglicht CTM zudem den direkten Zugriff auf die Score-Funktion, was die Integration etablierter Methoden zur kontrollierten bzw. bedingten Generierung aus der Diffusions-Community erleichtert. Darüber hinaus erlaubt dieser Zugriff die Berechnung der Likelihood. Der Quellcode ist unter https://github.com/sony/ctm verfügbar.