Correspondência de Velocidade Terminal

Resumo

Propusemos o Terminal Velocity Matching (TVM), uma generalização do *flow matching* que permite a modelagem generativa de alta fidelidade em uma ou poucas etapas. O TVM modela a transição entre quaisquer dois passos de difusão e regulariza o seu comportamento no tempo terminal, em vez do tempo inicial. Demonstramos que o TVM fornece um limite superior para a distância 2-Wasserstein entre as distribuições de dados e do modelo quando o modelo é Lipschitz-contínuo. No entanto, como os *Diffusion Transformers* não possuem essa propriedade, introduzimos alterações arquitetônicas mínimas que alcançam um treinamento estável e em estágio único. Para tornar o TVM eficiente na prática, desenvolvemos um *kernel* de atenção fundido que suporta passos reversos em Produtos Jacobiano-Vetor, os quais escalam bem com arquiteturas de *transformers*. No ImageNet-256x256, o TVM alcança um FID de 3,29 com uma única avaliação de função (NFE) e 1,99 FID com 4 NFEs. De forma similar, alcança 4,32 FID com 1-NFE e 2,94 FID com 4-NFE no ImageNet-512x512, representando um desempenho de última geração para modelos de uma/poucas etapas treinados do 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.