Нормализующие траекторные модели

Аннотация

Диффузионные модели разбивают процесс семплирования на множество небольших шагов гауссовского шумоподавления — допущение, которое нарушается при сжатии генерации до нескольких грубых переходов. Существующие методы с малым числом шагов решают эту проблему с помощью дистилляции, обучения согласованности или состязательных целей, но при этом жертвуют вероятностной структурой. Мы представляем Модели нормализующих траекторий (NTM), которые моделируют каждый обратный шаг как выразительный условный нормализующий поток с обучением на основе точного правдоподобия. Архитектурно NTM объединяет мелкие обратимые блоки внутри каждого шага с глубоким параллельным предиктором вдоль траектории, образуя сквозную сеть, обучаемую с нуля или инициализируемую из предварительно обученных моделей согласования потоков. Точное траекторное правдоподобие дополнительно позволяет осуществлять самодистилляцию: легковесный шумоподавитель, обученный на собственном скоре модели, создает высококачественные образцы за четыре шага. На эталонах генерации изображений по текстовому описанию NTM сравнивается или превосходит сильные базовые модели генерации изображений всего за четыре шага семплирования, при этом уникально сохраняя точное правдоподобие на протяжении всей генеративной траектории.

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.