Versnel TarFlow-steekproeven met GS-Jacobi-iteratie

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Auteurs: Ben Liu, Zhen Qin

cs.AI

Samenvatting

Beeldgeneratiemodellen hebben brede toepassingen gevonden. Een voorbeeld is het TarFlow-model, dat de transformer-architectuur combineert met Normalizing Flow-modellen en state-of-the-art resultaten behaalt op meerdere benchmarks. Door de causale vorm van aandacht die sequentiële berekening vereist, is het samplingproces van TarFlow echter extreem traag. In dit artikel tonen we aan dat door een reeks optimalisatiestrategieën het TarFlow-sampling aanzienlijk kan worden versneld door gebruik te maken van de Gauss-Seidel-Jacobi (afgekort als GS-Jacobi) iteratiemethode. Specifiek ontdekken we dat blokken in het TarFlow-model variërend belang hebben: een klein aantal blokken speelt een grote rol in beeldgeneratietaken, terwijl andere blokken relatief weinig bijdragen; sommige blokken zijn gevoelig voor beginwaarden en gevoelig voor numerieke overloop, terwijl andere relatief robuust zijn. Op basis van deze twee kenmerken stellen we de Convergence Ranking Metric (CRM) en de Initial Guessing Metric (IGM) voor: CRM wordt gebruikt om te bepalen of een TarFlow-blok "eenvoudig" is (convergeert in weinig iteraties) of "moeilijk" (vereist meer iteraties); IGM wordt gebruikt om te evalueren of de beginwaarde van de iteratie goed is. Experimenten op vier TarFlow-modellen tonen aan dat GS-Jacobi-sampling de samplingefficiëntie aanzienlijk kan verbeteren terwijl de kwaliteit van de gegenereerde beelden (gemeten door FID) behouden blijft, met snelheidsverbeteringen van 4,53x in Img128cond, 5,32x in AFHQ, 2,96x in Img64uncond en 2,51x in Img64cond zonder afbreuk te doen aan FID-scores of samplekwaliteit. Code en checkpoints zijn toegankelijk op 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