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 모델의 블록들은 다양한 중요성을 지니고 있음을 발견하였다: 소수의 블록들이 이미지 생성 작업에서 주요 역할을 하는 반면, 다른 블록들은 상대적으로 적은 기여를 한다; 일부 블록들은 초기값에 민감하고 수치적 오버플로우가 발생하기 쉬운 반면, 다른 블록들은 상대적으로 견고하다. 이러한 두 가지 특성을 바탕으로, 우리는 수렴 순위 지표(Convergence Ranking Metric, CRM)와 초기 추정 지표(Initial Guessing Metric, IGM)를 제안한다: CRM은 TarFlow 블록이 "단순한"(적은 반복으로 수렴)지 "까다로운"(더 많은 반복이 필요)지를 식별하는 데 사용되며, IGM은 반복의 초기값이 좋은지 평가하는 데 사용된다. 네 가지 TarFlow 모델에 대한 실험 결과, GS-Jacobi 샘플링은 생성된 이미지의 품질(FID로 측정)을 유지하면서 샘플링 효율성을 크게 향상시켰으며, Img128cond에서는 4.53배, AFHQ에서는 5.32배, Img64uncond에서는 2.96배, Img64cond에서는 2.51배의 속도 향상을 달성하였다. 코드와 체크포인트는 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