Normalização de Modelos de Trajetória

Resumo

Modelos baseados em difusão decompõem a amostragem em várias etapas pequenas de remoção de ruído gaussiano — uma suposição que colapsa quando a geração é comprimida em algumas transições grosseiras. Métodos existentes de poucas etapas abordam isso por meio de destilação, treinamento de consistência ou objetivos adversariais, mas sacrificam o arcabouço de verossimilhança no processo. Apresentamos os Modelos de Trajetória Normalizante (MTN), que modelam cada etapa reversa como um fluxo normalizante condicional expressivo com treinamento de verossimilhança exata. Arquiteturalmente, o MTN combina blocos invertíveis rasos dentro de cada etapa com um preditor paralelo profundo ao longo da trajetória, formando uma rede ponta a ponta treinável do zero ou inicializável a partir de modelos de correspondência de fluxo pré-treinados. Sua verossimilhança exata da trajetória permite ainda a autodestilação: um removedor de ruído leve treinado no próprio escore do modelo produz amostras de alta qualidade em quatro etapas. Em benchmarks de texto para imagem, o MTN iguala ou supera fortes linhas de base de geração de imagem em apenas quatro etapas de amostragem, ao mesmo tempo que retém exclusivamente a verossimilhança exata sobre a trajetória generativa.

English

Diffusion-based models decompose sampling into many small Gaussian denoising steps -- an assumption that breaks down when generation is compressed to a few coarse transitions. Existing few-step methods address this through distillation, consistency training, or adversarial objectives, but sacrifice the likelihood framework in the process. We introduce Normalizing Trajectory Models (NTM), which models each reverse step as an expressive conditional normalizing flow with exact likelihood training. Architecturally, NTM combines shallow invertible blocks within each step with a deep parallel predictor across the trajectory, forming an end-to-end network trainable from scratch or initializable from pretrained flow-matching models. Its exact trajectory likelihood further enables self-distillation: a lightweight denoiser trained on the model's own score produces high-quality samples in four steps. On text-to-image benchmarks, NTM matches or outperforms strong image generation baselines in just four sampling steps while uniquely retaining exact likelihood over the generative trajectory.