Aceleración del Muestreo de TarFlow con Iteración GS-Jacobi

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Autores: Ben Liu, Zhen Qin

cs.AI

Resumen

Los modelos de generación de imágenes han logrado aplicaciones generalizadas. Como ejemplo, el modelo TarFlow combina la arquitectura transformer con modelos de Normalizing Flow, alcanzando resultados de vanguardia en múltiples benchmarks. Sin embargo, debido a la forma causal de la atención que requiere cálculo secuencial, el proceso de muestreo de TarFlow es extremadamente lento. En este artículo, demostramos que, mediante una serie de estrategias de optimización, el muestreo de TarFlow puede acelerarse significativamente utilizando el método de iteración Gauss-Seidel-Jacobi (abreviado como GS-Jacobi). Específicamente, encontramos que los bloques en el modelo TarFlow tienen una importancia variable: un pequeño número de bloques desempeña un papel principal en las tareas de generación de imágenes, mientras que otros bloques contribuyen relativamente poco; algunos bloques son sensibles a los valores iniciales y propensos a desbordamientos numéricos, mientras que otros son relativamente robustos. Basándonos en estas dos características, proponemos la Métrica de Clasificación de Convergencia (CRM) y la Métrica de Estimación Inicial (IGM): CRM se utiliza para identificar si un bloque de TarFlow es "simple" (converge en pocas iteraciones) o "complejo" (requiere más iteraciones); IGM se utiliza para evaluar si el valor inicial de la iteración es adecuado. Los experimentos en cuatro modelos TarFlow demuestran que el muestreo GS-Jacobi puede mejorar significativamente la eficiencia del muestreo mientras mantiene la calidad de las imágenes generadas (medida por FID), logrando aceleraciones de 4.53x en Img128cond, 5.32x en AFHQ, 2.96x en Img64uncond y 2.51x en Img64cond sin degradar las puntuaciones FID o la calidad de las muestras. El código y los checkpoints están disponibles en 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