Ускорение выборки TarFlow с использованием итерации GS-Jacobi

Accelerate TarFlow Sampling with GS-Jacobi Iteration

May 19, 2025

Авторы: Ben Liu, Zhen Qin

cs.AI

Аннотация

Модели генерации изображений получили широкое распространение. В качестве примера, модель TarFlow объединяет архитектуру трансформера с моделями нормализующих потоков, достигая передовых результатов на множестве бенчмарков. Однако из-за причинной формы внимания, требующей последовательных вычислений, процесс сэмплирования в TarFlow крайне медленный. В данной работе мы показываем, что с помощью ряда стратегий оптимизации сэмплирование в TarFlow можно значительно ускорить, используя метод итераций Гаусса-Зейделя-Якоби (сокращённо GS-Jacobi). В частности, мы обнаружили, что блоки в модели TarFlow имеют разную важность: небольшое количество блоков играет основную роль в задачах генерации изображений, в то время как другие блоки вносят относительно небольшой вклад; некоторые блоки чувствительны к начальным значениям и склонны к числовому переполнению, тогда как другие относительно устойчивы. На основе этих двух характеристик мы предлагаем метрику ранжирования сходимости (Convergence Ranking Metric, CRM) и метрику начального предположения (Initial Guessing Metric, IGM): CRM используется для определения, является ли блок TarFlow "простым" (сходится за небольшое количество итераций) или "сложным" (требует больше итераций); IGM используется для оценки того, насколько хорошим является начальное значение итерации. Эксперименты на четырёх моделях TarFlow демонстрируют, что сэмплирование с использованием GS-Jacobi может значительно повысить эффективность сэмплирования, сохраняя качество генерируемых изображений (измеряемое по FID), достигая ускорений в 4.53 раза для Img128cond, 5.32 раза для AFHQ, 2.96 раза для Img64uncond и 2.51 раза для Img64cond без ухудшения показателей 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