GS-Jacobi反復法によるTarFlowサンプリングの高速化

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

著者: Ben Liu, Zhen Qin

cs.AI

要旨

画像生成モデルは広範な応用を実現している。一例として、TarFlowモデルはトランスフォーマーアーキテクチャとNormalizing Flowモデルを組み合わせ、複数のベンチマークで最先端の結果を達成している。しかし、因果的な形式のアテンションが逐次計算を必要とするため、TarFlowのサンプリングプロセスは極めて遅い。本論文では、一連の最適化戦略を通じて、Gauss-Seidel-Jacobi（略してGS-Jacobi）反復法を用いることで、TarFlowサンプリングを大幅に加速できることを示す。具体的には、TarFlowモデル内のブロックが異なる重要性を持つことを発見した：少数のブロックが画像生成タスクにおいて主要な役割を果たす一方、他のブロックは比較的寄与が少ない；また、一部のブロックは初期値に敏感で数値オーバーフローを起こしやすいが、他のブロックは比較的頑健である。これら2つの特性に基づき、Convergence Ranking Metric（CRM）とInitial Guessing Metric（IGM）を提案する：CRMはTarFlowブロックが「単純」（少ない反復で収束する）か「困難」（より多くの反復を必要とする）かを識別するために使用され、IGMは反復の初期値が良好かどうかを評価するために使用される。4つのTarFlowモデルでの実験により、GS-Jacobiサンプリングが生成画像の品質（FIDで測定）を維持しながらサンプリング効率を大幅に向上させることが示され、Img128condで4.53倍、AFHQで5.32倍、Img64uncondで2.96倍、Img64condで2.51倍の高速化を達成し、FIDスコアやサンプル品質の低下を招くことなく実現した。コードとチェックポイントは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