Abbinamento della Velocità Terminale

Abstract

Proponiamo Terminal Velocity Matching (TVM), una generalizzazione del flow matching che abilita la modellazione generativa ad alta fedeltà in uno o pochi passi. TVM modella la transizione tra due qualsiasi step di diffusione e regolarizza il suo comportamento al tempo terminale piuttosto che al tempo iniziale. Dimostriamo che TVM fornisce un limite superiore per la distanza di Wasserstein-2 tra le distribuzioni dei dati e del modello quando il modello è Lipschitz-continuo. Tuttavia, poiché i Diffusion Transformer non possiedono questa proprietà, introduciamo modifiche architetturali minime che consentono un addestramento stabile e a stadio singolo. Per rendere TVM efficiente in pratica, sviluppiamo un kernel di attenzione fusa che supporta i passi all'indietro sui prodotti Jacobiano-Vettore, i quali scalano bene con le architetture a transformer. Su ImageNet-256x256, TVM raggiunge un FID di 3.29 con una singola valutazione di funzione (NFE) e un FID di 1.99 con 4 NFE. Analogamente, raggiunge un FID a 1-NFE di 4.32 e un FID a 4-NFE di 2.94 su ImageNet-512x512, rappresentando lo stato dell'arte per i modelli a uno/pochi passi addestrati da zero.

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.