Consistency Models Made Easy

Introduction

Diffusion Models (DMs) , or Score-based Generative Models (SGMs) , have vastly changed the landscape of visual content generation with applications in images , videos , and 3D objects . DMs progressively transform a data distribution to a known prior distribution (e.g. Gaussian noise) according to a stochastic differential equation (SDE) and train a model to denoise noisy observations. Samples can be generated via a reverse-time SDE that starts from noise and uses the trained model to progressively denoise it. However, sampling from a DM naively requires hundreds to thousands of model evaluations due to the curvature of the diffusion sampling trajectory , making the entire generative process slow. Many approaches have been proposed to address this issue, including training-based techniques such as distillation , adaptive compute architectures for the backbone model , as well as training-free methods such as fast samplers or interleaving small and large backbone models during sampling . However, the speedup achieved by these sampling techniques usually comes at the expense of the quality of generated samples.

Consistency Models (CMs) are a new family of generative models, closely related to diffusion models, that have demonstrated promising results as faster generative models. These models learn a mapping between noise and data that all the points of the sampling trajectory map to the same initial data point. Owing to this condition, these models are capable of generating high-quality samples in 1-2 model evaluations. The best such models so far, built using improved Consistency Training (iCT) , have pushed the quality of images generated by 1-step CMs trained from scratch to a level comparable with SoTA DMs using thousands of steps for generation. Unfortunately, CMs remain time-consuming and practically challenging to train: the best practice so far takes many times longer than similar-quality diffusion models while involving complex hyperparameter choices in the training process. In total, this has substantially limited the uptake of CMs within the community.

In this work, we propose a vastly more efficient scheme for building consistency models. Specifically, we first rewrite the consistency condition, through the lens of dynamical systems, as a finite-difference approximation of a particular differential equation, which we call the differential consistency condition. Directly training consistency models using this condition as a loss function raises several challenges; we propose several strategies (including choosing a specific discretization curriculum and using a particular weighting function) to address them. Most importantly, using this formalism, we observe that DMs can be seen as a special case of CMs with a loose form of differential consistency condition. This lets us smoothly interpolate from DM to CM by progressively tightening the consistency condition, gradually making the time discretization smaller as training progresses. Combining these strategies results in a streamlined training recipe, called Easy Consistency Tuning (ECT), that lets us build a consistency model starting from a pretrained diffusion model, and thereby significantly decreasing the training cost.

ECT achieves better 2-step sample quality on CIFAR-10 and ImageNet 64$\times$64 than the prior art, while using less than $\mathbf{2\%}$ of the training FLOPs; even if we factor in the pre-training budget of the corresponding diffusion model, ECT only requires $\nicefrac{{1}}{{4}}$ the time of the state-of-the-art iCT . ECT also significantly reduces the inference cost by $\nicefrac{{1}}{{1000}}$ compared to Pretrained Score SDE/DMs while maintaining comparable generation quality.

Preliminaries

Let $p_{\text{data}}({\mathbf{x}}_{0})$ denote the data distribution. Diffusion models (DMs) perturb this distribution by adding monotonically increasing i.i.d. Gaussian noise with standard deviation $\sigma(t)$ from $t=0$ to $T$ such that $p_{t}({\mathbf{x}}_{t}|{\mathbf{x}}_{0})={\mathcal{N}}({\mathbf{x}}_{0},\sigma^{2}(t){\bm{I}})$, and $\sigma(t)$ is chosen such that $\sigma(0)=\sigma_{\text{min}}$ and $\sigma(T)=\sigma_{\text{max}}$. This process is described by the following SDE

where $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ denotes the score function. Prior works set $\sigma(t)=t$ which gives

where $f({\mathbf{x}}_{t},t)$ denotes a denoising function that predicts clean image ${\mathbf{x}}_{0}$ given noisy image ${\mathbf{x}}_{t}$. We will follow this parametrization in the rest of this paper. Note that time is same as noise level with this parametrization, and we will use these two terms interchangeably.

Consistency Models.

CMs are built upon the PF-ODE in Eq. 3, which establishes a bijective mapping between data distribution and noise distribution. CMs learn a consistency function $f({\mathbf{x}}_{t},t)$ that maps the noisy image ${\mathbf{x}}_{t}$ back to the clean image ${\mathbf{x}}_{0}$

Note that the consistency function needs to satisfy the boundary condition at $t=0$. Prior works impose this boundary condition by parametrizing the CM as

where $\theta$ is the model parameter, $F_{\theta}$ is the network to train, and $c_{\text{skip}}(t)$ and $c_{\text{out}}(t)$ are time-dependent scaling factors such that $c_{\text{skip}}(0)=1$, $c_{\text{out}}(0)=0$. This parameterization guarantees the boundary condition by design. We discuss specific choices of $c_{\text{skip}}(t)$ and $c_{\text{out}}(t)$ in Appendix C.

During training, CMs first discretize the PF-ODE into $N-1$ subintervals with boundaries given by $t_{\text{min}}=t_{1}<t_{2}<\ldots<t_{N}=T$. The model is trained on the following CM loss, which minimizes a metric between adjacent points on the sampling trajectory

In the seminal work, Song et al. use Learned Perceptual Similarity Score (LPIPS) as a metric function, set $w(t_{i})=1$ for all $t_{i}$, and sample $t_{i}$ according to the sampling scheduler by Karras et al. : $t_{i}=\left(t_{\text{max}}^{1/\rho}+\frac{i}{N-1}(t_{\text{min}}^{1/\rho}-t_{\text{max}}^{1/\rho})\right)^{\rho}$ for $i\in{\mathcal{U}}[1,N-1]$ and $\rho=0.7$. Further, the score function $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t})$ can either be estimated from a pretrained diffusion model, which results in Consistency Distillation (CD), or can be estimated with an unbiased score estimator

which results in consistency training (CT).

The follow-up work, iCT , introduces several improvements that significantly improve training efficiency as well as the performance of CMs. First, the LPIPS metric, which introduces undesirable bias in generative modeling, is replaced with a Pseudo-Huber metric. Second, the network $f_{\theta^{-}}$ does not maintain an EMA of the past values of $f_{\theta}$. Third, iCT replaces the uniform weighting scheme $w(t_{i})=1$ with $w(t_{i})=\frac{1}{t_{i+1}-t_{i}}$. Further, the scaling factors of noise embeddings and dropout are carefully selected. Fourth, iCT introduces a complex discretization curriculum during training:

where $m$ is the current number of iterations, $M$ is the total number of iterations, $\sigma_{\max}$ and $\sigma_{\min}$ is the largest and smallest noise level for training, $s_{0}=10$ and $s_{1}=1280$ are hyperparameters. Finally, during training, iCT samples $i\sim p(i)\propto\text{erf}\left(\frac{\log(t_{i+1})-P_{\text{mean}}}{\sqrt{2}P_{\text{std}}}\right)-\text{erf}\left(\frac{\log(t_{i})-P_{\text{mean}}}{\sqrt{2}P_{\text{std}}}\right)$ from a discrete Lognormal distribution, where $P_{\text{mean}}=-1.1$ and $P_{\text{std}}=2.0$.

Easy Consistency Tuning (ECT)

We will first introduce the differential consistency condition that forms the core of our method to train CMs and derive a loss objective based on this condition. Next, we will analyze this loss objective and highlight some challenges of training CMs with it. Based on this analysis, we present our method, Easy Consistency Tuning (ECT). ECT is a simple, principled approach to efficiently train CMs to meet the (differential) consistency condition. The resulting CMs can generate high-quality samples in 1 or 2 sampling steps.

As previously stated in Sec. 2, CMs learn a consistency function $f({\mathbf{x}}_{t},t)$ that maps the noisy image ${\mathbf{x}}_{t}$ back to the clean image ${\mathbf{x}}_{0}$: $f({\mathbf{x}}_{t},t)={\mathbf{x}}_{0}$. Instead, we argue that CMs can be defined through the time derivative of the model output, given by the differential form

Since the clean image ${\mathbf{x}}_{0}$ is independent of the noise level $t$, its time derivative is zero.

This boundary condition $f({\mathbf{x}}_{0},0)={\mathbf{x}}_{0}$ ensures that the model output matches the clean image when the noise level is zero. Together, the differential form in Eq. 9 and the boundary condition define the differential consistency condition in Eq. 10, or consistency condition in short.

To satisfy the consistency condition, we minimize the distance between $f_{\theta}({\mathbf{x}}_{t})$ and $f_{\theta}({\mathbf{x}}_{r})$. For the boundary condition ($t>r=0$), we optimize $f_{\theta}({\mathbf{x}}_{r})$ to align with the clean image ${\mathbf{x}}_{0}$. For higher noise levels ($t>r>0$), we freeze $f_{\operatorname{sg}(\theta)}({\mathbf{x}}_{r})$ using the stop-gradient operator $\operatorname{sg}$ and optimize $f_{\theta}({\mathbf{x}}_{t})$ to align with $f_{\operatorname{sg}(\theta)}({\mathbf{x}}_{r})$.

Loss Function.

Given the discretization, the training objective for CMs can be formulated as:

where we can extract a weighting function $w(t,r)=\frac{1}{t-r}$, ${\mathbf{x}}_{t}={\mathbf{x}}_{0}+t\cdot\bm{\epsilon}$, ${\mathbf{x}}_{r}={\mathbf{x}}_{0}+r\cdot\bm{\epsilon}$ using a shared noise direction $\bm{\epsilon}\sim p(\bm{\epsilon})$, and ${\mathbf{d}}(\cdot,\cdot)$ is a metric function. In practice, we set ${\mathbf{d}}(\cdot,\cdot)$ to the squared $L_{2}$ metric. Further, we discuss an interesting relation between the Pseudo-Huber metric proposed in Song and Dhariwal and the weighting function in Sec. 3.3. To regularize the model, dropout masks should also be consistent across $f_{\theta}({\mathbf{x}}_{t})$ and $f_{\theta}({\mathbf{x}}_{r})$ by setting a shared random seed, ensuring that the noise levels $t$ and $r$ are the only varying factors.

2 The "Curse of Consistency" and its Implications

The objective in Eq. 12 can be challenging to optimize when $\Delta t\to 0$. This is because the prediction errors from each discretization interval accumulate, leading to slow training convergence or, in the worst case, divergence. To further elaborate, consider a large noise level $T$. We first follow the notation of CT and split the noise horizon $[0,T]$ into $N$ smaller consecutive subintervals. We can error bound $1$-step prediction of CM as

where $r_{1}=0<t_{1}=r_{2}<\cdots<t_{N-2}=r_{N}<t_{N}=T$, and $e_{\max}=\max_{i}\|f_{\theta}({\mathbf{x}}_{t_{i}})-f_{\theta}({\mathbf{x}}_{r_{i}})\|$, for $i=1,\cdots,N$. Ideally, we want both $N$ and $e_{\max}$ to be small so that this upper bound is small, but in practice, there is a trade-off between these two terms. As $\Delta t_{i}=(t_{i}-r_{i})\to 0$, $e_{\max}$ decreases because ${\mathbf{x}}_{t_{i}}$ and ${\mathbf{x}}_{r_{i}}$ will be close, and it is easier to predict both $f_{\theta}({\mathbf{x}}_{r_{i}})$ and $f_{\theta}({\mathbf{x}}_{t_{i}})$. However, $N$ will increase as $\Delta t_{i}\to 0$. In contrast, for a large $\Delta t_{i}$, vice-versa holds true. It is difficult to theoretically estimate the rate at which $e_{\max}$ decreases when $\Delta t_{i}=t_{i}-r_{i}\to 0$, as it depends on the optimization process—specifically, how effectively we train the model to minimize the consistency error in each interval. If $e_{\max}$ decreases more slowly than $N$, the product of the two can increase instead, resulting in a worse prediction error $\|f_{\theta}({\mathbf{x}}_{T})-{\mathbf{x}}_{0}\|$.

3 Tackling the "Curse of Consistency"

The discussion in Sec. 3.2 highlighted the training instability issue that might arise when we train CMs from scratch with the loss objective in Eq. 12 while directly following $\Delta t\approx 0$. However, these issues can be mitigated with the right design choices. In this section, we list several strategies that address the aforementioned training instability issues and help improve the efficiency of CMs.

Drawing inspiration from iCT , we start ECT with a large $\Delta t$, and gradually shrink $\Delta t\to 0$. In our problem setup, Since $t>r\geq 0$, we have the largest possible $\Delta t=t$ with $r=0$, which yields

Continuous-time training schedule.

Choice of metric.

As discussed in Sec. 2, iCT uses the Pseudo-Huber metric to mitigate the perceptual bias caused by the LPIPS metric . When taking a careful look as this metric, we reveal that this metric offers adaptive per-sample scaling of the gradients, which reduces the variance of gradients during training. Specifically, the differential of the Pseudo-Huber loss can be decomposed into two terms: an adaptive scaling factor $w(\Delta)$ and the differential of the squared $L_{2}$ loss. We discuss this in more detail in Appendix A. Therefore, we retain the squared $L_{2}$ metric used in DMs, and explore varying adaptive weighting terms.

Weighting function.

Weighting functions usually lead to a substantial difference in performance in DMs, and the same holds true for CMs. Substituting the mapping function $p(r|t)$ in Eq. 15 into the weighting function in Eq. 12, we get $w(t,r)=\nicefrac{{1}}{{(t-r)}}=q^{a}\cdot\nicefrac{{1}}{{tn(t)}}$, where $q^{a}$ is the scaling factor of the loss when $\Delta t\to 0$, and $\nicefrac{{1}}{{tn(t)}}$ is the weighting akin to DMs. This couples the weighting function with the mapping function. Instead, we consider decoupled weighting functions. Motivated by the adaptive scaling factor that appears in Pseudo-Huber loss (See Appendix A for more details), we rewrite the weighting function as

where $\Delta=f({\mathbf{x}}_{t})-f({\mathbf{x}}_{r})$. We define $\bar{w}(t)$ as regular weighting, while $w(\Delta)$ as adaptive weighting. The adaptive weighting aids in improved training efficiency with the $L_{2}$ metric because when $\Delta\to 0$ (usually happens at $t\to 0$), this weighting $w(\Delta)$ can upscale gradients to avoid vanishing gradient in learning fine-grained features, while $c$ avoids potential numerical issues at $\Delta\approx 0$. We direct the reader to Tab. 2 and Tab. 3 in Appendix B for a detailed overview of various choices of $\bar{w}(t)$ and $w(\Delta)$ considered in this work. In general, we notice that CMs’ generative capability greatly benefits from weighting functions that control the variance of the gradients across different noise levels.

Experiments

We measure the efficiency and scalability of ECT on two datasets: CIFAR-10 and ImageNet $64\times 64$ . We evaluate the sample quality using Fréchet Inception Distance (FID) and Fréchet Distance under the DINOv2 model ($\text{FD}_{\text{DINOv2}}$ ) and measure sampling efficiency using the number of function evaluations (NFEs). We also indicate the relative training costs of each of these methods. The extensive implementation details can be found in Appendix C.

We compare CMs trained with ECT (denoted as ECM) against state-of-the-art diffusion models, diffusion models with advanced samplers, distillation methods such as consistency distillation (CD) , and improved Consistency Training (iCT) . The key results are summarized in Fig. 2. We show the training FLOPs, inference cost, and generative performance of the four training schemes.

We compare ECMs against Score SDE , EDM , and EDM with DPM-Solver-v3 . $2$-step ECM, which has been fine-tuned for only $100$k iterations, matches Score SDE-deep, with $2\times$ model depth and $2000$ NFEs, in terms of FID. As noted in Fig. 2, ECM only requires $\nicefrac{{1}}{{1000}}$ of its inference cost and latency to achieve the same sample quality. $2$-step ECM fine-tuned for $100$k iterations outperforms EDM with advanced DPM-Solver-v3 (NFE=$10$).

Diffusion Distillation.

We compare ECT against Consistency Distillation (CD) , a SoTA approach that distills a pretrained DM into a CM. As shown in Tab. 1, ECM significantly outperforms CD on both CIFAR-10 and ImageNet 64$\times$64. We note that ECT is free from the errors of teacher DM and does not incur any additional cost of running teacher DM. $2$-step ECM outperforms $2$-step CD (with LPIPS ) in terms of FID ($2.20$ vs $2.93$) on CIFAR-10 while using around $\nicefrac{{1}}{{3}}$ of training/fine-tuning compute of CD.

Consistency training from scratch.

Improved Consistency Training (iCT) is the SoTA recipe for training a consistency model from scratch without inferring the diffusion teacher. Compared to training from scratch, ECT rivals iCT-deep using $\nicefrac{{1}}{{32}}$ of the training compute and $\nicefrac{{1}}{{2}}$ of the model size as shown in Fig. 2 and Tab. 1.

ECT unlocks state-of-the-art few-step generative abilities through a simple yet principled approach. With a negligible tuning cost, ECT demonstrates strong results while benefiting from the scaling in training FLOPs to enhance its few-step generation capability.

2 Scaling Laws of ECT

We investigate the scaling behavior of ECT, focusing on three aspects: training compute, model size, and model FLOPs. When computing resources are not a bottleneck, ECT scales well and follows the classic power law.

Model Size & FLOPs.

Related Work

Consistency models are a new family of generative models designed for efficient generation with 1 or 2 model steps without the need for adversarial training. CMs do not rely on a pretrained diffusion model (DM) to generate training targets but instead leverage an unbiased score estimator. CMs have been extended to multi-step sampling , latent space models , ControlNet , and combined with extra adversarial loss . Despite their sampling efficiency, CMs are typically more challenging to train and require significantly more compute resources compared to their diffusion counterparts. Our work substantially improves the training efficiency of CMs, reducing the cost of future research and deployment on CMs.

Diffusion Distillation

Drawing inspiration from knowledge distillation , diffusion distillation is the most widespread training-based approach to accelerate the diffusion sampling procedure. In diffusion distillation, a pretrained diffusion model (DM), which requires hundreds to thousands of model evaluations to generate samples, acts as a teacher. A student model is trained to match the teacher model’s sample quality, enabling it to generate high-quality samples in a few steps.

There are two main lines of work in this area. The first category involves trajectory matching, where the student learns to match points on the teacher’s sampling trajectory. Methods in this category include offline distillation , which require an offline synthetic dataset generated by sampling from a pretrained DM to distill a teacher model into a few-step student model; progressive distillation , and TRACT , which require multiple training passes or offline datasets to achieve the same goal; and BOOT , Consistency Distillation (CD), and Imagine-Flash, which minimize the difference between the student predictions at carefully selected points on the sampling trajectory.

CD is closely related to our method, as it leverages a teacher model to generate pairs of adjacent points and enforces the student predictions at these points to map to the initial data point. However, it limits the quality of consistency models to that of the pretrained diffusion model, and the LPIPS-based metric used in CD loss introduces undesirable bias when evaluated on metrics such as FID.

The second category minimizes the probabilistic divergence between data and model distributions, i.e., distribution matching . Methods such as DreamFusion and ProfilicDreamer are also used for 3D object generation. These methods use score distillation or adversarial loss to distill an expensive teacher model into an efficient student model. However, they can be challenging to train in a stable manner due to the alternating updating schemes from either adversarial or score distillation.

A drawback of training-based approaches is that they usually require extra training procedures after pretraining to distill an efficient student. For a detailed discussion on the recent progress of diffusion distillation, we refer to .

Fast Samplers for Diffusion Models.

Fast samplers are usually training-free and use advanced solvers to simulate the diffusion stochastic differential equation (SDE) or ordinary differential equation (ODE) to reduce the number of sampling steps. These methods reduce the discretization error during sampling by analytically solving a part of SDE or ODE , by using exponential integrators and higher order polynomials for better approximation of the solution , using higher order numerical methods , using better approximation of noise levels during sampling , correcting predictions at each step of sampling and ensuring that the solution of the ODE lies on a desired manifold . Another orthogonal strategy is to achieve acceleration through parallelizing the sampling process . A drawback of these fast samplers is that the quality of samples drastically reduces as the number of sampling steps goes below a threshold such as $10$ steps.

Limitations

One of the major limitations of ECT is that it requires a dataset to tune DMs to CMs, while recent works developed data-free approaches for diffusion distillation. The distinction between ECT and data-free methods is that ECT learns the consistency condition on a given dataset via the self-teacher, while data-free diffusion distillations acquire knowledge from a frozen diffusion teacher. This feature of ECT can be a potential limitation since the training data of some models are unavailable to the public. However, we hold an optimistic view on tuning CMs using datasets different from pretraining. The data composition and scaling for consistency tuning will be a valuable research direction.

Conclusion

We propose a simple yet efficient scheme for training consistency models. First, we derive the differential consistency condition and formulate a loss objective based on this condition. Next, we identify the "curse of consistency" of learning CMs and propose several strategies to mitigate these issues. We find that diffusion models can be considered a special case of consistency models with a loose discretization of the consistency condition. Based on these observations and strategies, we propose a streamlined training recipe called Easy Consistency Tuning (ECT). The resulting models, ECMs, obtained through ECT, unlock state-of-the-art few-step generative capabilities at a minimal tuning cost and are able to benefit from scaling. We hope this work will aid in the wider adoption of consistency models within the research community.

Acknowledgments

Zhengyang Geng and Ashwini Pokle are supported by funding from the Bosch Center for AI. Zico Kolter gratefully acknowledges Bosch funding for the lab.

Appendix A Design Choices for ECT

In this section, we expand upon our motivation behind the design decisions for the mapping function, metric, and weighting function used for ECT.

We first assume that $\Delta t$ is approximately proportional to $t$. Let $0<c\leq 1$ be this constant of proportionality, then we can write:

As training progresses, the mapping function should gradually shrink $\Delta t\to 0$. However, the above parameterization does not achieve this. An alternative parameterization is to exponentially decrease $\Delta t$. We can rewrite the ratio between $r$ and $t$ as:

Finally, we adjust the mapping function to balance the prediction difficulties across different noise levels:

For $n(t)$, we choose $n(t)=1+k,\sigma(-b,t)=1+\frac{k}{1+e^{bt}}$, using the sigmoid function $\sigma$. Since $r\geq 0$, we also clamp $r$ to satisfy this constraint after the adjustment.

The intuition behind this mapping function is that the relative difficulty of predicting $f({\mathbf{x}}_{r})$ from ${\mathbf{x}}_{t}$ can vary significantly across different noise levels $t$ when using a linear mapping between $t$ and $r$.

Consider $\nicefrac{{r}}{{t}}=0.9$. At small values of $t$, ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{r}$ are close, making the alignment of $f({\mathbf{x}}_{t})$ with $f({\mathbf{x}}_{r})$ relatively easy. In contrast, at larger $t$, where ${\mathbf{x}}_{t}$ and ${\mathbf{x}}_{r}$ are relatively far apart, the distance between the predictions $f({\mathbf{x}}_{t})$ and $f({\mathbf{x}}_{r})$ can be substantial. This leads to imbalanced gradient flows across different noise levels, impeding the training dynamics.

Therefore, we downscale $\nicefrac{{r}}{{t}}$ when $t$ is near through the mapping function, balancing the gradient flow across varying noise levels. This prevents the gradient at any noise level from being too small or too large relative to other noise levels, thereby controlling the variance of the gradients.

We direct the reader to Appendix B for details of how to set $q^{\lceil\text{iters}/d\rceil}$.

Choice of Metric.

As discussed in Sec. 2, iCT uses pseudo-Huber metric to mitigate the perceptual bias caused by the LPIPS metric ,

This metric indeed improves the performance of CMs over the classic squared $L_{2}$ loss. When taking a careful look as this metric, we reveal that one of the reasons for this improvement is that this metric is more robust to the outliers compared to the $L_{2}$ metric due to its adaptive per-sample scaling of the gradients. Let $\Delta={\mathbf{x}}-{\mathbf{y}}$, then the differential of the pseudo-Huber metric can be written as

where we have decomposed the differential of pseudo-Huber loss into an adaptive weighting term and the differential of the squared $L_{2}$ loss. Therefore, we retain the squared $L_{2}$ metric used in DMs, and explore varying adaptive weighting terms.

Distinction between the training schedules of ECT and iCT.

As noted in Sec. 2, iCT employs a discrete-time curriculum given by Eq. 8. This curriculum divides the noise horizon $[0,T]$ into $N$ smaller consecutive subintervals to apply the consistency loss, characterized by non-overlapping segments $[t_{i},t_{i+1}]$, and gradually increases the number of intervals $N=10\to 1280$. However, the "boundary" condition of this schedule is to start with the number of intervals to $N=1$, learning a model solely mapping samples at noise levels $T_{\text{max}}$ to the clean data ${\mathbf{x}}_{0}$, largely distinct from the classic diffusion models training. We instead investigate the design principles of a continuous-time schedule whose "boundary" condition yields diffusion pretraining, i.e., constructing training pairs of $r=0$ for all $t$ at the beginning.

Appendix B Exploring the Design Space of Consistency Models

Due to ECT’s efficiency, we can explore the design space of CMs at minimal cost. We specifically examine the weighting function, training schedule, and regularization for CMs.

Our most significant finding is that controlling gradient variances and balancing the gradients across different noise levels are fundamental to CMs’ training dynamics. Leveraging the deep connection between CMs and DMs, we also improve the diffusion pretraining using our insights.

Forward processes with different noise schedules and model parameterizations can be translated into each other at the cost of varying weighting functions . From our experiments on a wide range of weighting schemes, we learn three key lessons.

$(1)$ There is no free lunch for weighting function, i.e., there is likely no universal regular weighting $\bar{w}(t)$ that can outperform all other candidates for different datasets, models, and target metrics for both 1-step and 2-step generation.

$(2)$ The adaptive weighting $w(\Delta)$ achieves better results by controlling gradient variance. The adaptive weighting $w(\Delta)$ on a per-sample basis shows uniform improvements on both CIFAR-10 and ImageNet 64$\times$64. See Tab. 3 for the ablation study.

Beyond ECT, we further investigate the role of adaptive weighting $w(\Delta)$ in pretraining on a toy SwissRoll dataset using flow matching and a simple MLP network.

where $p=0$ corresponds to no adaptive weighting. We set $\epsilon=10^{-6}$ and control the strength of gradient normalization by varying $p$ from to $1$.

As we increase the strength of adaptive weighting, flow models become easier to sample from. Surprisingly, even $p=1$ demonstrates strong few-step sampling results when pretraining the flow model. See Fig. 6 for visualization.

$(3)$ The regular weighting $\bar{w}(t)$ and the adaptive weighting $w(\Delta)$ are compatible. This compatibility is how we build ECMs in this work, using adaptive weighting $w(\Delta)$ to achieve strong 1-step sampling capabilities and regular weighting $\bar{w}(t)$ to improve 2-step sampling.

Mapping Function.

We compare the constant mapping function with $n(t)\equiv 1$ in Eq. 17 and mapping function equipped with the sigmoid $n(t)$ in Eq. 18. We use $k=8$ and $b=1$ for all the experiments, which transfers well from CIFAR-10 to ImageNet 64$\times$64 and serves as a baseline in our experiments. Though $b=2$ can further improve the 1-step FIDs on ImageNet 64$\times$64, noticed post hoc, we don’t rerun our experiments.

On CIFAR-10, the constant mapping function with $n(t)\equiv 1$ achieves 1-step FID of 4.06 at $200$k iterations, worse than the 1-step FID of 3.86 by $n(t)=1+\frac{k}{1+e^{bt}}$. Under our forward process (${\mathbf{x}}_{t}={\mathbf{x}}_{0}+t\cdot\epsilon$) and model parameterization (EDM ), the constant mapping function also incurs training instability on ImageNet 64$\times$64, likely due to the imbalanced gradient flow.

The role of the mapping function, regarding training, is to balance the difficulty of learning consistency condition across different noise levels, avoiding trivial consistency loss near $t\to 0$. For model parameterizations and forward processes different from ours, for example, flow matching , we advise readers to start from the constant mapping function due to its simplicity.

Dropout.

We find that CMs benefit significantly from dropout . On CIFAR-10, we apply a dropout of $0.20$ for models evaluated on FID and a dropout of $0.30$ for models evaluated on $\text{FD}_{\text{DINOv2}}$. On ImageNet 64$\times$64, when increasing the dropout rate from $0.10$ to $0.40$, the 2-step FID decreases from $4.53$ to $3.24$. Increasing the dropout rate further can be helpful for 1-step FID under certain weighting functions, but the 2-step FID starts to deteriorate. In general, we optimize our model configurations for efficient 2-step generation and, therefore, choose the dropout rate of $0.40$ for ECM-S.

Finally, we note that the dropout rate tuned at a given weighting function $w(t)$ transfers well to the other weighting functions, thereby reducing the overall cost of hyperparameter tuning. On ImageNet 64$\times$64, the dropout rate can even transfer to different model sizes. We apply a dropout rate of $0.50$ for all the model sizes of ECM-M/L/XL.

Using this feature of ECT, we demonstrate its efficiency by training ECMs to surpass previous Consistency Distillation, which took hundreds of GPU hours, using one hour on a single A100 GPU.

Training Generative Models in 1 GPU Hour.

Deep generative models are typically computationally expensive to train. Unlike training a classifier on CIFAR-10, which usually completes within one GPU hour, leading generative models on CIFAR-10 as of 2024 require days to a week to train on 8 GPUs. Even distillation from pretrained diffusion models can take over a day on 8 GPUs or even more, equivalently hundreds of GPU hours.

Pretraining using ECT.

Compared with the EDM baseline, ECT brings a convergence acceleration over $2\times$ regarding $\text{FD}_{\text{DINOv2}}$, matching EDM’s final performance using less than half of the pretraining budget and largely outperforming it at the full pretraining budget.

EDM pretrained by ECT achieves $\text{FD}_{\text{DINOv2}}$ of $150.39$ for unconditional generation and $121.05$ for class-conditional generation, considerably better than the EDM baseline’s $\text{FD}_{\text{DINOv2}}$ of $168.17$ for unconditional generation and $145.20$ for class-conditional generation, when using the same training budget and inference steps (NFE=35).

Influence of Pretraining Quality.

ECM from the ECT pretraining surpasses SoTA GANs in 1 sampling step and advanced DMs in 2 sampling steps, only slightly falling behind our pretrained models and setting up a new SoTA for modern metric $\text{FD}_{\text{DINOv2}}$. See Tab. 4 for details.

On ImageNet 64$\times$64, ECM-M, initialized from EDM2-M , deviates from the power law scaling and achieves better generative performance than the log-linear trend. (See Fig. 4 Right). We speculate that it is due to sufficient pretraining, in which EDM2-M was pretrained by a budget of $2\times$ training images of other model sizes (S/L/XL).

Comparison with iCT.

We offer an ablation study to compare ECT and the state-of-the-art iCT . We interpolated back to iCT by removing components from ECT, including dropout rate for 2-step models, continuous-time training schedule through mapping function, and diffusion pretraining. We show the FIDs of 1-step and 2-step sampling and the training iterations in Tab. 5.

Differences between 1-step and 2-step Generation.

Our empirical results suggest that the training recipe for the best 1-step generative models can differ from the best few-step generative models in many aspects, including weighting function, dropout rate, and EMA rate/length. Fig. 7 shows an example of how FIDs from different numbers of function evaluations (NFEs) at inference vary with dropout rates.

Starting from a proper model size, the benefits from 2-step sampling seem larger than increasing the model size by a factor of two in our setups. In the prior works, iCT employs $2\times$ deeper model, but the 1-step generative performance can still be inferior to the 2-step results from ECT. This finding is consistent with recent theoretical analysis , which indicates a tighter bound on the sample quality for the 2-step generation compared to the 1-step generation.

Pareto Frontier & Scaling Law.

The Pareto Frontier reveals a seemingly power law scaling behavior. Training configurations not optimized for the current compute budget, i.e., not on the Pareto Frontier, deviate from this scaling. Simply scaling up the training compute without adjusting other parameters may result in suboptimal performance. In our compute scaling experiments, we increased the batch size and enabled the smoothing factor $c$ in the adaptive weighting to maintain this trend.

Appendix C Experimental Details

We follow Karras et al. , Song et al. and set $c_{\text{skip}}(t)=\nicefrac{{\sigma_{\text{data}}^{2}}}{{(t^{2}+\sigma_{\text{data}}^{2})}}$ and $c_{\text{out}}(t)=\nicefrac{{t\sigma_{\text{data}}}}{{\sqrt{t^{2}+\sigma_{\text{data}}^{2}}}}$, where $\sigma_{\text{data}}^{2}$ is the variance of (normalized) data, and set to $0.5$ for both CIFAR-10 and ImageNet 64$\times$64.

Computational Cost.

ECT is computationally efficient. On ImageNet 64$\times$64, ECT requiring only $\mathbf{0.39\%}$ of the iCT training budget, and $\mathbf{0.60\%}$ to $\mathbf{1.91\%}$ of the EDM2 pretraining budget depending on the model sizes. The exact computational resources required to train each individual model is shown in Tab. 6.

Training Details.

We use RAdam optimizer for experiments on CIFAR-10 and Adam optimizer for experiments on ImageNet 64$\times$64. We set the $\beta$ to ($0.9$, $0.999$) for CIFAR-10 and ($0.9$, $0.99$) for ImageNet 64$\times$64. All the hyperparameters are indicated in Tab. 6. We do not use any learning rate decay, weight decay, or warmup on CIFAR-10. We follow EDM2 to apply an inverse square root learning rate decay schedule on ImageNet 64$\times$64.

On CIFAR-10, we employ the traditional Exponential Moving Average (EMA). To better understand the influence of the EMA rate, we track three Power function EMA models on ImageNet 64$\times$64, using EMA lengths of $0.01$, $0.05$, and $0.10$. The multiple EMA models introduce no visible cost to the training speed. Considering our training budget is much smaller than the diffusion pretraining stage, we didn’t perform the Post-Hoc EMA search as in EDM2 .

Sampling Details.

We apply stochastic sampling for 2-step generation. For 2-step sampling, we follow Song and Dhariwal and set the intermediate $t=0.821$ for CIFAR-10, and $t=1.526$ for ImageNet 64$\times$64.

Evaluation Details.

For both CIFAR-10 and ImageNet 64$\times$64, FID and $\text{FD}_{\text{DINOv2}}$ are computed using $50$k images sampled from ECMs. As suggested by recent works , $\text{FD}_{\text{DINOv2}}$ aligns better with human evaluation. We use dgm-evalhttps://github.com/layer6ai-labs/dgm-evalto calculate $\text{FD}_{\text{DINOv2}}$ to ensure align with previous practice.

Visualization Setups.

Image samples in Fig. 3 are from class bubble (971), class flamingo (130), class golden retriever (207), class space shuttle (812), classs Siberian husky (250), classs ice cream (928), class oscilloscope (688), class llama (355), class tiger shark (3).

Scaling Laws of Training Compute.

For the results on scaling laws for training compute on CIFAR-10 shown in Fig. 4 (Left), we train 6 class-conditional ECMs, each with varying batch size and number of training iterations. All ECMs in this experiment are initialized from the pretrained class-conditional EDM.

The first two points of $2^{0}$ and $2^{1}$ on Fig. 4 (Left) use a batch size of $128$ for $100$k and $200$k iterations, respectively. The third point of $2^{2}$ corresponds to ECM trained with batch sizes of $256$ for $200$k iterations. The final three points of $2^{3}$, $2^{4}$, and $2^{5}$ correspond to ECM trained with a batch size of $512$ for $200$k, $400$k, and $800$k iterations, respectively, with the smoothing factor $c=0.03$ enabled in the adaptive weighting $w(\Delta)$. We use $\bar{w}(t)=\nicefrac{{1}}{{t}}$ as the regular weighting function to train all these models as this $\bar{w}(t)$ achieves good performance on $\text{FD}_{\text{DINOv2}}$.

Scaling Laws of Model Size and Model FLOPs.

We have included details of model capacity as well as FLOPs in Tab. 6 to replicate this plot on ImageNet 64$\times$64.

Appendix D Broader Impacts

We propose Easy Consistency Tuning (ECT) that can efficiently train consistency models as state-of-the-art few-step generators, using only a small fraction of the computational requirements compared to current CMs training and diffusion distillation methods. We hope that ECT will democratize the creation of high-quality generative models, enabling artists and creators to produce content more efficiently. While this advancement can aid creative industries by reducing computational costs and speeding up workflows, it also raises concerns about the potential misuse of generative models to produce misleading, fake, or biased content. We conduct experiments on academic benchmarks, whose resulting models are less likely to be misused. Further experiments are needed to better understand these consistency model limitations and propose solutions to address them.

Appendix E Qualitative Results

We provide some randomly generated 2-step samples from ECMs trained on CIFAR-10 and ImageNet-$64\times 64$ in Fig. 8 and Fig. 9, respectively.

References