PaperlayerView exact revision on arXiv

Source-linked AI summary

Improved Training Technique for Latent Consistency Models

Quan Dao, Khanh Doan, Di Liu, Trung Le, Dimitris Metaxas

arXiv:2502.01441v2cs.CVcs.LG

TL;DR

Consistency training’s performance in latent space remained uncertain, despite its strong pixel-space results and the importance of latent space for large-scale generation. The authors analyze latent outliers and introduce robust training changes, producing high-fidelity latent consistency images that narrow the gap with latent diffusion.

  • Problem

    Consistency training had reached strong pixel-space performance, but its effectiveness in latent space—important for large-scale applications—remained unresolved.

    Source: §1 · p. 2

  • Method

    The authors improve latent consistency training with Cauchy loss, early-timestep diffusion loss, optimal-transport coupling, adaptive scaling-c scheduling, and Non-scaling LayerNorm.

    Source: §1 · p. 2

  • Results

    The resulting latent consistency training generates high-fidelity images and narrows the gap with latent diffusion models.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Consistency training is highly sensitive to training-data statistics, and impulsive noise in latent data is associated with instability and poor performance.

    Source: §6 · p. 10

Abstract

from arXiv · show

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/

1 INTRODUCTION

Diffusion models generate high-quality images but require many evaluations, while consistency models offer one- or multistep sampling. This work investigates whether consistency training can perform effectively in latent space, where large-scale applications operate, and develops techniques to improve it.

1 INTRODUCTION

  • Diffusion sampling can require thousands of function evaluations, making it computationally expensive and hindering real-world application.Prior efforts address this through new training frameworks or distillation.Source: §1 · p. 1
  • Consistency models support high-quality single- or multistep sampling, but consistency distillation requires a pretrained diffusion model and consistency training initially lagged behind it.Later training improvements brought consistency training closer to diffusion performance in pixel space.Source: §1 · p. 2
  • Large-scale applications such as text-to-image and video generation operate in latent space, while prior consistency-training advances focused on pixel space.The authors identify latent-space performance as essential for scaling consistency models to large datasets.Source: §1 · p. 2
  • Latent data contains highly impulsive outliers that produce large temporal-difference values, undermining Pseudo-Huber loss and degrading latent consistency training.The authors propose Cauchy loss, early-timestep diffusion loss, optimal-transport matching, adaptive scaling-c scheduling, and Non-scaling LayerNorm to improve performance.Source: §1 · p. 2

2 RELATED WORKS

Prior consistency-training methods improve performance or efficiency through added discriminators, multistep training, or pretrained diffusion initialization. Latent consistency distillation methods extend these approaches to latent space, with differing sampling quality and training designs.

2 RELATED WORKS

  • ACT (Kong et al., 2023) and CTM (Kim et al., 2023) add a GAN discriminator, while MCM (Heek et al., 2024) uses 2–8 sampling steps to trade training efficiency for image quality.The added discriminator may require careful hyperparameter tuning.Source: §2 · p. 2
  • ECM (Geng et al., 2024) initializes from pretrained diffusion and fine-tunes with consistency training, but requires the same architecture as that pretrained model.It improves training time while maintaining good generation performance.Source: §2 · p. 2
  • LCM applies consistency techniques in latent space, but its one- or two-step images remain blurry.Source: §2 · p. 3
  • TCD (Zheng et al., 2024) uses CTM instead of consistency distillation, and Hyper-SD (Ren et al., 2024) applies TCD to PF-ODE segments before progressively merging them.Hyper-SD also integrates human-feedback learning and score distillation to optimize one-step generation.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries define the diffusion process and the consistency function, then describe consistency distillation and training. They also summarize iCT’s training improvements, which established strong pixel-space performance.

3 PRELIMINARIES

  • Forward diffusion adds Gaussian noise with increasing standard deviation σ(t), from σmin to σmax; setting σ(t) = t defines the probability-flow ODE.Source: §3 · p. 3
  • The consistency function maps noisy samples to clean data, and the model is trained to approximate that mapping.The mapping is defined through the probability-flow ODE and connects noisy and data distributions.Source: §3 · p. 3
  • The model parameterization combines a time-scaled skip connection with a neural network, with boundary conditions on cskip and cout.At σmin, cskip equals 1; at σmax, cout equals 0.Source: §3 · p. 3
  • Consistency training and distillation discretize the probability-flow ODE into an ordered sequence of noise levels.Source: §3 · p. 3
  • Consistency distillation uses a pretrained diffusion score model to construct a noise-level pair for its training loss.The pair is formed by perturbing the same clean sample with Gaussian noise and estimating the lower-noise sample using the pretrained score.Source: §3 · p. 3
  • Consistency training compares online and target networks on adjacent noise levels built from the same clean sample and noise draw.The target parameters use an exponential moving average; the loss also uses a timestep weighting function and predefined metric.Source: §3 · pp. 3–4
  • iCT improved consistency training with EMA settings, embedding and dropout tuning, Pseudo-Huber losses, an exponential discretization curriculum, inverse-gap weighting, and lognormal timestep sampling.These changes brought consistency training to performance on par with diffusion models in pixel space.Source: §3 · p. 4

4 METHOD

The method addresses weaker latent-space consistency training by targeting impulsive outliers and unstable temporal differences with changes to the loss, training objectives, scaling schedule, and normalization. It combines Cauchy loss, small-noise diffusion loss, OT matching, an adaptive c scheduler, and Non-scaling LayerNorm.

4.1 ANALYSIS OF LATENT SPACE

  • Latent temporal differences contain wider impulsive outliers than pixel-space differences, and Pseudo-Huber loss may not sufficiently stabilize latent training.The reported latent TD outliers range from −3.2 to 5, compared with −1.5 to 1.7 for pixel TD.Source: §4.1 · p. 4

4.2 CAUCHY LOSS AGAINST IMPULSIVE OUTLIER

  • Cauchy loss grows logarithmically for large residuals, damping extreme outliers without the zero-gradient behavior that could discard useful latent information.The authors choose Cauchy over Geman-McClure for latent consistency models.Source: §4.2 · p. 5

4.3 DIFFUSION LOSS AT SMALL TIMESTEP

  • An additional diffusion loss regularizes consistency training at small noise levels, while the authors avoid applying it at large noise levels.They heuristically set the diffusion threshold to r = 0.25; large-noise diffusion loss could mislead training because its L2 loss is high.Source: §4.3 · p. 6

4.4 OT MATCHING REDUCES THE VARIANCE

  • OT matching reassigns noise-data pairs within a training batch to reduce variance in consistency training.The reported minibatch OT overhead is around 0.93% of training time, with performance improvement reported in table 2a.Source: §4.4 · p. 6

4.5 ADAPTIVE c SCHEDULER

  • The adaptive c scheduler scales c down as discretization steps increase and TD variance falls, while impulsive outliers persist.The scheduler is fitted heuristically to the reduction rate of TD variance.Source: §4.5 · pp. 6–7

4.6 NON-SCALING LAYERNORM

  • Non-scaling LayerNorm fixes its scaling term at 1 and disables its gradient update to capture feature statistics while reducing sensitivity to outliers.The authors suspect LayerNorm’s shared scaling term can amplify outliers across features.Source: §4.6 · p. 7

5 EXPERIMENT

The proposed technique is evaluated on three 256×256 datasets against latent consistency and diffusion baselines, with component ablations on CelebA-HQ. The ablations test robust loss, diffusion-loss timing, OT coupling, adaptive scaling, and normalization choices.

5.1 PERFORMANCE OF OUR TRAINING TECHNIQUE

  • CelebA-HQ, LSUN Church, and FFHQ are tested at 256×256 using 32×32×4 VAE latents, a shared OpenAI UNet/VE process, and FID/Recall on 50K decoded samples.Source: §5.1 · pp. 7–8
  • 1-NFE FIDs are 7.27 (CelebA-HQ), 8.87 (LSUN Church), and 8.29 (FFHQ); the method beats iLCT and matched LDM-8†, but not LDM/LFM.Its 1-NFE sampling outperforms LDM-8† using the same noise scheduler and architecture, while LDM and LFM have lower FIDs but require multiple timesteps.Source: Table 1 · p. 8, §5.1 · p. 8

5.2 ABLATION OF PROPOSED FRAMEWORK

  • On CelebA-HQ 1-NFE, replacing Pseudo-Huber with Cauchy lowers FID from 37.15 to 13.02; subsequent components reduce it to 9.11, 8.89, 7.76, and 7.27.Small-timestep diffusion loss, OT coupling, adaptive scaling c, and NsLN correspond to the successive reductions.Source: §5.2 · p. 9
  • Huber and Geman-McClure losses underperform Cauchy in latent space: Huber remains sensitive to extreme outliers, while Geman-McClure can ignore them entirely.Source: §5.2 · p. 9
  • Diffusion loss at every timestep underperforms selective use at smaller timesteps such as r = 0.25; at larger timesteps, it may conflict with consistency loss.Source: §5.2 · pp. 9–10
  • As discretization steps increase and TD values scale down, impulsive outliers persist; scaling c down with the schedule performs better than a fixed large c.Source: §5.2 · p. 10
  • NsLN removes LayerNorm’s shared scaling term, capturing feature statistics while resisting outliers; it outperforms the second-best GroupNorm by 0.5 FID.Source: §5.2 · p. 10

6 CONCLUSION

The conclusion links latent-space CT instability to impulsive outliers and presents Cauchy loss with improved training strategies as a route to high-fidelity latent consistency generation. It also identifies normalization changes and integration with CTM as future directions.

6 CONCLUSION

  • Cauchy loss and improved training strategies yield high-fidelity images from latent CT, bridging the gap between latent diffusion and consistency models.The conclusion also notes that CT is highly sensitive to training-data statistics.Source: §6 · p. 10
  • Future work could investigate normalization methods that reduce outlier impact, including removing scaling from group or instance normalization.Source: §6 · p. 10
  • The authors also propose integrating this technique with CTM, which has demonstrated improved performance over traditional consistency models.Source: §6 · p. 10

A APPENDIX

The appendix provides additional uncurated samples across three datasets, with one-step and two-step examples, plus CelebA-HQ samples trained with two alternative losses.

A APPENDIX

Related papers

Loading 2502.01441v2…