Coincidencia de Velocidad Terminal

Resumen

Proponemos Terminal Velocity Matching (TVM), una generalización de flow matching que permite un modelado generativo de alta fidelidad en uno o pocos pasos. TVM modela la transición entre dos pasos de difusión cualesquiera y regulariza su comportamiento en el tiempo terminal en lugar de en el tiempo inicial. Demostramos que TVM proporciona una cota superior para la distancia 2-Wasserstein entre las distribuciones de datos y del modelo cuando este es continuo según Lipschitz. Sin embargo, dado que los Diffusion Transformers carecen de esta propiedad, introducimos cambios arquitectónicos mínimos que logran un entrenamiento estable en una sola etapa. Para hacer que TVM sea eficiente en la práctica, desarrollamos un kernel de atención fusionado que admite pasos hacia atrás en Productos Jacobiano-Vector, los cuales escalan bien con arquitecturas de transformadores. En ImageNet-256x256, TVM logra un FID de 3.29 con una única evaluación de función (NFE) y un FID de 1.99 con 4 NFE. De manera similar, alcanza un FID de 4.32 con 1-NFE y un FID de 2.94 con 4-NFE en ImageNet-512x512, lo que representa un rendimiento de vanguardia para modelos de uno/pocos pasos entrenados desde cero.

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.