Accélération de l'échantillonnage TarFlow avec l'itération GS-Jacobi

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Auteurs: Ben Liu, Zhen Qin

cs.AI

Résumé

Les modèles de génération d'images ont connu des applications généralisées. Par exemple, le modèle TarFlow combine l'architecture Transformer avec les modèles de Normalizing Flow, obtenant des résultats de pointe sur plusieurs benchmarks. Cependant, en raison de la forme causale de l'attention qui nécessite un calcul séquentiel, le processus d'échantillonnage de TarFlow est extrêmement lent. Dans cet article, nous démontrons qu'à travers une série de stratégies d'optimisation, l'échantillonnage de TarFlow peut être grandement accéléré en utilisant la méthode d'itération de Gauss-Seidel-Jacobi (abrégée en GS-Jacobi). Plus précisément, nous constatons que les blocs du modèle TarFlow ont une importance variable : un petit nombre de blocs jouent un rôle majeur dans les tâches de génération d'images, tandis que d'autres blocs contribuent relativement peu ; certains blocs sont sensibles aux valeurs initiales et sujets à des débordements numériques, tandis que d'autres sont relativement robustes. Sur la base de ces deux caractéristiques, nous proposons la Convergence Ranking Metric (CRM) et l'Initial Guessing Metric (IGM) : la CRM est utilisée pour identifier si un bloc TarFlow est "simple" (converge en peu d'itérations) ou "difficile" (nécessite plus d'itérations) ; l'IGM est utilisée pour évaluer si la valeur initiale de l'itération est bonne. Les expériences sur quatre modèles TarFlow démontrent que l'échantillonnage GS-Jacobi peut significativement améliorer l'efficacité de l'échantillonnage tout en maintenant la qualité des images générées (mesurée par le FID), atteignant des accélérations de 4,53x pour Img128cond, 5,32x pour AFHQ, 2,96x pour Img64uncond et 2,51x pour Img64cond sans dégrader les scores FID ou la qualité des échantillons. Le code et les points de contrôle sont accessibles sur https://github.com/encoreus/GS-Jacobi_for_TarFlow.

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