Совместите вашу касательную: обучение более качественных моделей согласованности через касательные, выровненные по многообразию

Align Your Tangent: Training Better Consistency Models via Manifold-Aligned Tangents

October 1, 2025

Авторы: Beomsu Kim, Byunghee Cha, Jong Chul Ye

cs.AI

Аннотация

С появлением диффузионных моделей и моделей согласования потоков, достигших наилучших показателей в генерации, внимание сообщества теперь сосредоточено на сокращении времени вывода без ущерба для качества образцов. Модели согласованности (Consistency Models, CMs), которые обучаются быть согласованными на траекториях диффузии или обыкновенных дифференциальных уравнений вероятностного потока (PF-ODE), позволяют выполнять выборку потока или диффузии за один или два шага. Однако CMs обычно требуют длительного обучения с большими размерами пакетов для достижения конкурентоспособного качества образцов. В данной работе мы исследуем динамику обучения CMs вблизи сходимости и обнаруживаем, что касательные CMs — направления обновления выходов CMs — являются довольно осцилляторными, в том смысле, что они движутся параллельно многообразию данных, а не в направлении к нему. Для смягчения осцилляторных касательных мы предлагаем новую функцию потерь, называемую расстоянием по признакам многообразия (Manifold Feature Distance, MFD), которая обеспечивает касательные, выровненные по многообразию и направленные к нему. В результате наш метод, названный Align Your Tangent (AYT), может ускорить обучение CMs на порядки и даже превзойти метрику схожести воспринимаемых изображений (LPIPS). Кроме того, мы обнаруживаем, что наша функция потерь позволяет обучать модели с крайне малыми размерами пакетов без ущерба для качества образцов. Код: https://github.com/1202kbs/AYT

English

With diffusion and flow matching models achieving state-of-the-art generating performance, the interest of the community now turned to reducing the inference time without sacrificing sample quality. Consistency Models (CMs), which are trained to be consistent on diffusion or probability flow ordinary differential equation (PF-ODE) trajectories, enable one or two-step flow or diffusion sampling. However, CMs typically require prolonged training with large batch sizes to obtain competitive sample quality. In this paper, we examine the training dynamics of CMs near convergence and discover that CM tangents -- CM output update directions -- are quite oscillatory, in the sense that they move parallel to the data manifold, not towards the manifold. To mitigate oscillatory tangents, we propose a new loss function, called the manifold feature distance (MFD), which provides manifold-aligned tangents that point toward the data manifold. Consequently, our method -- dubbed Align Your Tangent (AYT) -- can accelerate CM training by orders of magnitude and even out-perform the learned perceptual image patch similarity metric (LPIPS). Furthermore, we find that our loss enables training with extremely small batch sizes without compromising sample quality. Code: https://github.com/1202kbs/AYT