Accelerate TarFlow Sampling with GS-Jacobi Iteration
Authors: name: Ben Liu, name: Zhen Qin

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