Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Accelerate TarFlow Sampling with GS-Jacobi Iteration

Ben Liu, Zhen Qin

TL;DR

This work addresses the slow sampling of TarFlow caused by its autoregressive transformer blocks with causal attention by recasting the sampling step as a diagonalized nonlinear fixed-point system and applying a GS-Jacobi hybrid method. Two metrics, the Initial Guessing Metric ($\mathrm{IGM}$) and the Convergence Ranking Metric ($\mathrm{CRM}$), are proposed to guide blockwise iteration and initialization, enabling adaptive, module-aware GS-Jacobi updates. Empirical results across four TarFlow models show substantial speedups (e.g., $4.53\times$, $5.32\times$, $2.96\times$, $2.51\times$) while maintaining comparable FID scores and image quality. The approach offers a practical path to faster TarFlow sampling and broader applicability of autoregressive normalizing-flow models in high-resolution image generation, with code and checkpoints available online.

Abstract

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

Accelerate TarFlow Sampling with GS-Jacobi Iteration

TL;DR

This work addresses the slow sampling of TarFlow caused by its autoregressive transformer blocks with causal attention by recasting the sampling step as a diagonalized nonlinear fixed-point system and applying a GS-Jacobi hybrid method. Two metrics, the Initial Guessing Metric () and the Convergence Ranking Metric (), are proposed to guide blockwise iteration and initialization, enabling adaptive, module-aware GS-Jacobi updates. Empirical results across four TarFlow models show substantial speedups (e.g., , , , ) while maintaining comparable FID scores and image quality. The approach offers a practical path to faster TarFlow sampling and broader applicability of autoregressive normalizing-flow models in high-resolution image generation, with code and checkpoints available online.

Abstract

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
Paper Structure (18 sections, 1 theorem, 15 equations, 7 figures, 5 tables, 1 algorithm)

This paper contains 18 sections, 1 theorem, 15 equations, 7 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methods
  4. Jacobi Mode Fixed Point Iteration Sampling
  5. Initial Guessing Metric
  6. Convergence Ranking Metric
  7. Modular Guass-Seidel-Jacobi Iteration
  8. Experiment
  9. Initial Guessing Metric
  10. Convergence Ranking Metric
  11. Quantitative Evaluations with FID
  12. Conclusion
  13. Convergence and Error Propagation
  14. Derivation for Convergence Ranking Metric
  15. Components of Convergence Ranking Metric
  16. ...and 3 more sections

Key Result

Proposition 1

For fixed point iteration (fixedpoint), let $\varepsilon^{(k)} = X^{(k)}-X^*$ be the error after $k$ iteration, $e_t$ be its $t$-th component, $f_t^{(k)}=x_t^{(k)}-\sigma_t^{(k)}z_t-u_t^{(k)}$, then:

Figures (7)

  • Figure 1: Simple intuition diagram of GS-Jacobi sampling. First pass forward a small batch of images to compute Initial Guessing Metric (IGM) and Convergence Ranking Metric (CRM) for each block. When sampling, the initial iteration value $X^{(0)}$ is determined by IGM; for blocks whose CRM is non-dominant, parallel Jacobi iterate $X$; for CRM-dominant blocks, segment $X$ into small modules $X_g$, parallel Jacobi iterating within modules, then serially deliver to next module.
  • Figure 2: The distance between $X^{(k)}$ (up to 150 times) and target $X^*$ of all 8 blocks in four models. Most blocks converge within iteration times $<<T$, with each model exhibiting only one or two slowly descending curves.
  • Figure 3: The trace of the sampling in four models. From top to bottom: Img128cond, Img64cond, Img64uncond, AFHQ. From left to right: noise, Block 7-0, denoised image.
  • Figure 4: The influence of different initial value and iteration times of an Img64cond sample. From top to bottom: Set all $X^{(0)}=Z$, Jacobi 30 times; Adaptive initialized according to IGM, Jacobi 20 times; Adaptive by IGM, Jacobi 30 times; GS-Jacobi [0/7-16/8-10/13-6]
  • Figure 5: The distance between GS-Jacobi iteration and target $X_g^*$ of four tough blocks. All modules tend to converge within 30 iterations, and the 1st module suffer a more difficult trace.
  • ...and 2 more figures

Theorems & Definitions (1)

  • Proposition 1: Converge and Error Propagation