Согласование по конечной скорости

Аннотация

Мы предлагаем метод согласования терминальной скорости (Terminal Velocity Matching, TVM), который обобщает метод согласования потоков и позволяет осуществлять генеративное моделирование высокой точности за один или несколько шагов. TVM моделирует переход между любыми двумя временными шагами диффузии и регулирует поведение модели в конечный момент времени, а не в начальный. Мы доказываем, что TVM обеспечивает верхнюю границу для 2-Вассерштейнова расстояния между распределениями данных и модели, когда модель является липшицевой. Однако, поскольку диффузионные трансформеры не обладают этим свойством, мы вводим минимальные архитектурные изменения, позволяющие достичь стабильного одноэтапного обучения. Для обеспечения практической эффективности TVM мы разрабатываем объединенное ядро внимания, которое поддерживает обратные проходы для вычисления произведений Якобиана на вектор, хорошо масштабирующиеся с архитектурами трансформеров. На ImageNet-256x256 TVM достигает показателя FID 3.29 при одной оценке функции (NFE) и 1.99 FID при 4 NFE. Аналогично, на ImageNet-512x512 метод достигает 4.32 FID при 1 NFE и 2.94 FID при 4 NFE, что представляет собой состояние искусства для одно- и малошаговых моделей, обученных с нуля.

English

We propose Terminal Velocity Matching (TVM), a generalization of flow matching that enables high-fidelity one- and few-step generative modeling. TVM models the transition between any two diffusion timesteps and regularizes its behavior at its terminal time rather than at the initial time. We prove that TVM provides an upper bound on the 2-Wasserstein distance between data and model distributions when the model is Lipschitz continuous. However, since Diffusion Transformers lack this property, we introduce minimal architectural changes that achieve stable, single-stage training. To make TVM efficient in practice, we develop a fused attention kernel that supports backward passes on Jacobian-Vector Products, which scale well with transformer architectures. On ImageNet-256x256, TVM achieves 3.29 FID with a single function evaluation (NFE) and 1.99 FID with 4 NFEs. It similarly achieves 4.32 1-NFE FID and 2.94 4-NFE FID on ImageNet-512x512, representing state-of-the-art performance for one/few-step models from scratch.