Improved Training Technique for Latent Consistency Models

2025-02-03

AI Summary

To address training instability and degraded generation quality in latent consistency models (LCMs) caused by outlier interference, this paper proposes a robust training framework enabling high-fidelity single-step and two-step text-to-image and text-to-video generation. Our key contributions are: (1) the first adoption of Cauchy loss—replacing pseudo-Huber loss—to enhance robustness against outliers in latent-space optimization; (2) coupling early-diffusion loss with an optimal transport (OT) objective to improve alignment between predicted and target latent distributions; and (3) introducing an adaptive scaling-c scheduler and non-scaled LayerNorm to stabilize training dynamics. Experiments demonstrate that our method significantly narrows the performance gap between LCMs and diffusion models, achieving state-of-the-art fidelity and sampling efficiency across multi-scale text-conditioned generation tasks.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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/

Problem

• Consistency Model Improvement
• Robustness to Anomalies
• Diffusion Model Performance Gap

Innovation

• Optimized Computational Flow
• Diffusion Computation
• Smart Adjustment Strategy