Улучшенная методика обучения моделей с латентной согласованностью

Improved Training Technique for Latent Consistency Models

February 3, 2025

Авторы: Quan Dao, Khanh Doan, Di Liu, Trung Le, Dimitris Metaxas

cs.AI

Аннотация

Модели согласованности - это новое семейство генеративных моделей, способных производить высококачественные образцы либо за один шаг, либо за несколько шагов. Недавно модели согласованности продемонстрировали впечатляющую производительность, достигая результатов, сравнимых с моделями диффузии в пространстве пикселей. Однако успех масштабирования обучения согласованности на крупных наборах данных, особенно для задач генерации текста в изображения и видео, определяется производительностью в скрытом пространстве. В данной работе мы анализируем статистические различия между пиксельным и скрытым пространствами, обнаруживая, что скрытые данные часто содержат очень импульсивные выбросы, которые значительно ухудшают производительность iCT в скрытом пространстве. Для решения этой проблемы мы заменяем потери Псевдо-Хьюбера на потери Коши, что эффективно смягчает влияние выбросов. Кроме того, мы вводим потери диффузии на ранних временных шагах и используем оптимальную транспортную связь (OT) для дальнейшего улучшения производительности. Наконец, мы представляем адаптивный планировщик масштабирования-c для управления устойчивым процессом обучения и принимаем Нескалируемую слойную нормализацию в архитектуре для более точного улавливания статистики признаков и снижения влияния выбросов. Благодаря этим стратегиям мы успешно обучаем скрытые модели согласованности, способные к высококачественной генерации с одним или двумя шагами, значительно сокращая разрыв в производительности между скрытой согласованностью и моделями диффузии. Реализация доступна здесь: https://github.com/quandao10/sLCT/

English

Consistency models are a new family of generative models capable of producing high-quality samples in either a single step or multiple steps. Recently, consistency models have demonstrated impressive performance, achieving results on par with diffusion models in the pixel space. However, the success of scaling consistency training to large-scale datasets, particularly for text-to-image and video generation tasks, is determined by performance in the latent space. In this work, we analyze the statistical differences between pixel and latent spaces, discovering that latent data often contains highly impulsive outliers, which significantly degrade the performance of iCT in the latent space. To address this, we replace Pseudo-Huber losses with Cauchy losses, effectively mitigating the impact of outliers. Additionally, we introduce a diffusion loss at early timesteps and employ optimal transport (OT) coupling to further enhance performance. Lastly, we introduce the adaptive scaling-c scheduler to manage the robust training process and adopt Non-scaling LayerNorm in the architecture to better capture the statistics of the features and reduce outlier impact. With these strategies, we successfully train latent consistency models capable of high-quality sampling with one or two steps, significantly narrowing the performance gap between latent consistency and diffusion models. The implementation is released here: https://github.com/quandao10/sLCT/