Acelere a Amostragem TarFlow com Iteração GS-Jacobi

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Autores: Ben Liu, Zhen Qin

cs.AI

Resumo

Modelos de geração de imagens alcançaram aplicações generalizadas. Como exemplo, o modelo TarFlow combina a arquitetura transformer com modelos de Normalizing Flow, obtendo resultados de ponta em múltiplos benchmarks. No entanto, devido à forma causal da atenção, que requer computação sequencial, o processo de amostragem do TarFlow é extremamente lento. Neste artigo, demonstramos que, por meio de uma série de estratégias de otimização, a amostragem do TarFlow pode ser significativamente acelerada utilizando o método de iteração Gauss-Seidel-Jacobi (abreviado como GS-Jacobi). Especificamente, descobrimos que os blocos no modelo TarFlow possuem importância variável: um pequeno número de blocos desempenha um papel principal nas tarefas de geração de imagens, enquanto outros blocos contribuem relativamente pouco; alguns blocos são sensíveis a valores iniciais e propensos a overflow numérico, enquanto outros são relativamente robustos. Com base nessas duas características, propomos a Métrica de Classificação de Convergência (CRM) e a Métrica de Estimativa Inicial (IGM): a CRM é usada para identificar se um bloco do TarFlow é "simples" (converge em poucas iterações) ou "difícil" (requer mais iterações); a IGM é usada para avaliar se o valor inicial da iteração é bom. Experimentos em quatro modelos TarFlow demonstram que a amostragem GS-Jacobi pode melhorar significativamente a eficiência de amostragem enquanto mantém a qualidade das imagens geradas (medida pelo FID), alcançando acelerações de 4,53x no Img128cond, 5,32x no AFHQ, 2,96x no Img64uncond e 2,51x no Img64cond, sem degradar as pontuações FID ou a qualidade das amostras. Código e checkpoints estão disponíveis em 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