Accelerate Autoregressive Normalizing Flows Sampling with GS-Jacobi Iteration

ICLR 2026 Conference Submission19585 Authors

Keywords: Normalize Flow, Generation model

TL;DR: Accelerate TarFlow Sampling with Gauss-Seidel-Jacobi Iteration, speed up 5×

Abstract:

AutoRegressive Normalizing Flows (abbreviated as AR Flow) enjoy extensive applications in tasks such as density estimation and image generation. However, due to the causal affine coupling blocks requiring sequential computation, the sampling process is extremely slow.
In this paper, we demonstrate that through a series of optimization strategies, such AR Flows sampling can be greatly accelerated by using the Gauss-Seidel-Jacobi (abbreviated as GS-Jacobi) iteration method.
Specifically, we find that blocks in AR Flows have varying importance: a small number of blocks play a major role in image generation, 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 Flow 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.
The TarFlow was chosen as the main experimental subject in our study owing to its SOTA performance on several benchmarks.
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.53× in Img128cond, 5.32× in AFHQ, 2.96× in Img64uncond, and 2.51× in Img64cond without degrading FID scores or sample quality.

Primary Area: generative models

Submission Number: 19585