Terminale Snelheid Afstemming

Samenvatting

Wij stellen Terminal Velocity Matching (TVM) voor, een generalisatie van flow matching die hoogwaardige generatieve modellering in één of enkele stappen mogelijk maakt. TVM modelleert de overgang tussen twee willekeurige diffusietijdstappen en regulariseert het gedrag ervan op het eindtijdstip in plaats van op het begintijdstip. Wij bewijzen dat TVM een bovengrens biedt voor de 2-Wasserstein-afstand tussen data- en modelverdelingen wanneer het model Lipschitz-continu is. Omdat Diffusion Transformers deze eigenschap echter missen, introduceren we minimale architecturale aanpassingen die een stabiele, enkelvoudige training mogelijk maken. Om TVM in de praktijk efficiënt te maken, ontwikkelden we een gefuseerde attention-kernel die backward passes op Jacobiaan-Vectorproducten ondersteunt, wat goed schaalt met transformer-architecturen. Op ImageNet-256x256 behaalt TVM een FID van 3,29 met een enkele functie-evaluatie (NFE) en een FID van 1,99 met 4 NFE's. Evenzo behaalt het een 1-NFE FID van 4,32 en een 4-NFE FID van 2,94 op ImageNet-512x512, wat state-of-the-art prestaties vertegenwoordigt voor modellen met één/weinig stappen die from scratch zijn getraind.

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.