Beschleunigung der TarFlow-Stichprobenentnahme mit GS-Jacobi-Iteration

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Autoren: Ben Liu, Zhen Qin

cs.AI

Zusammenfassung

Bildgenerierungsmodelle haben weitverbreitete Anwendungen gefunden. Ein Beispiel ist das TarFlow-Modell, das die Transformer-Architektur mit Normalizing Flow-Modellen kombiniert und dabei state-of-the-art Ergebnisse auf mehreren Benchmarks erzielt. Aufgrund der kausalen Form der Aufmerksamkeit, die eine sequenzielle Berechnung erfordert, ist der Sampling-Prozess von TarFlow jedoch extrem langsam. In diesem Artikel zeigen wir, dass durch eine Reihe von Optimierungsstrategien das TarFlow-Sampling mithilfe der Gauss-Seidel-Jacobi (abgekürzt als GS-Jacobi) Iterationsmethode erheblich beschleunigt werden kann. Insbesondere stellen wir fest, dass Blöcke im TarFlow-Modell unterschiedliche Bedeutung haben: Eine kleine Anzahl von Blöcken spielt eine Hauptrolle bei Bildgenerierungsaufgaben, während andere Blöcke relativ wenig beitragen; einige Blöcke sind empfindlich gegenüber Anfangswerten und anfällig für numerische Überläufe, während andere relativ robust sind. Basierend auf diesen beiden Eigenschaften schlagen wir das Convergence Ranking Metric (CRM) und das Initial Guessing Metric (IGM) vor: CRM wird verwendet, um festzustellen, ob ein TarFlow-Block „einfach“ ist (konvergiert in wenigen Iterationen) oder „schwierig“ (erfordert mehr Iterationen); IGM wird verwendet, um zu bewerten, ob der Anfangswert der Iteration gut ist. Experimente an vier TarFlow-Modellen zeigen, dass GS-Jacobi-Sampling die Sampling-Effizienz erheblich steigern kann, während die Qualität der generierten Bilder (gemessen durch FID) erhalten bleibt, und dabei Beschleunigungen von 4,53x in Img128cond, 5,32x in AFHQ, 2,96x in Img64uncond und 2,51x in Img64cond erreicht, ohne die FID-Werte oder die Sample-Qualität zu verschlechtern. Code und Checkpoints sind auf https://github.com/encoreus/GS-Jacobi_for_TarFlow verfügbar.

English

Image generation models have achieved widespread applications. As an instance, the TarFlow model combines the transformer architecture with Normalizing Flow models, achieving state-of-the-art results on multiple benchmarks. However, due to the causal form of attention requiring sequential computation, TarFlow's sampling process is extremely slow. In this paper, we demonstrate that through a series of optimization strategies, TarFlow sampling can be greatly accelerated by using the Gauss-Seidel-Jacobi (abbreviated as GS-Jacobi) iteration method. Specifically, we find that blocks in the TarFlow model have varying importance: a small number of blocks play a major role in image generation tasks, while other blocks contribute relatively little; some blocks are sensitive to initial values and prone to numerical overflow, while others are relatively robust. Based on these two characteristics, we propose the Convergence Ranking Metric (CRM) and the Initial Guessing Metric (IGM): CRM is used to identify whether a TarFlow block is "simple" (converges in few iterations) or "tough" (requires more iterations); IGM is used to evaluate whether the initial value of the iteration is good. Experiments on four TarFlow models demonstrate that GS-Jacobi sampling can significantly enhance sampling efficiency while maintaining the quality of generated images (measured by FID), achieving speed-ups of 4.53x in Img128cond, 5.32x in AFHQ, 2.96x in Img64uncond, and 2.51x in Img64cond without degrading FID scores or sample quality. Code and checkpoints are accessible on https://github.com/encoreus/GS-Jacobi_for_TarFlow