Accelerare il Campionamento TarFlow con l'Iterazione GS-Jacobi

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Autori: Ben Liu, Zhen Qin

cs.AI

Abstract

I modelli di generazione di immagini hanno raggiunto un'ampia gamma di applicazioni. Ad esempio, il modello TarFlow combina l'architettura transformer con i modelli Normalizing Flow, ottenendo risultati all'avanguardia su molteplici benchmark. Tuttavia, a causa della forma causale dell'attenzione che richiede un calcolo sequenziale, il processo di campionamento di TarFlow è estremamente lento. In questo articolo, dimostriamo che, attraverso una serie di strategie di ottimizzazione, il campionamento di TarFlow può essere notevolmente accelerato utilizzando il metodo di iterazione Gauss-Seidel-Jacobi (abbreviato come GS-Jacobi). Nello specifico, abbiamo scoperto che i blocchi nel modello TarFlow hanno un'importanza variabile: un piccolo numero di blocchi svolge un ruolo principale nei compiti di generazione di immagini, mentre altri blocchi contribuiscono relativamente poco; alcuni blocchi sono sensibili ai valori iniziali e soggetti a overflow numerico, mentre altri sono relativamente robusti. Sulla base di queste due caratteristiche, proponiamo la Convergence Ranking Metric (CRM) e la Initial Guessing Metric (IGM): la CRM viene utilizzata per identificare se un blocco TarFlow è "semplice" (converge in poche iterazioni) o "complesso" (richiede più iterazioni); la IGM viene utilizzata per valutare se il valore iniziale dell'iterazione è buono. Esperimenti condotti su quattro modelli TarFlow dimostrano che il campionamento GS-Jacobi può migliorare significativamente l'efficienza del campionamento mantenendo la qualità delle immagini generate (misurata tramite FID), ottenendo accelerazioni di 4.53x in Img128cond, 5.32x in AFHQ, 2.96x in Img64uncond e 2.51x in Img64cond senza degradare i punteggi FID o la qualità dei campioni. Codice e checkpoint sono accessibili su 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