Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models

Introduction

Diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020; Song et al., 2021b) have revolutionized generative AI, achieving remarkable results in image (Rombach et al., 2022; Ramesh et al., 2022; Ho et al., 2022), 3D (Poole et al., 2022; Wang et al., 2024; Liu et al., 2023b), audio (Liu et al., 2023a; Evans et al., 2024), and video generation (Blattmann et al., 2023; Brooks et al., 2024). Despite their success, a significant drawback is their slow sampling speed, often requiring dozens to hundreds of steps to generate a single sample. Various diffusion distillation techniques have been proposed, including direct distillation (Luhman & Luhman, 2021; Zheng et al., 2023b), adversarial distillation (Wang et al., 2022; Sauer et al., 2023), progressive distillation (Salimans & Ho, 2022), and variational score distillation (VSD) (Wang et al., 2024; Yin et al., 2024b; a; Luo et al., 2024; Xie et al., 2024b; Salimans et al., 2024). However, these methods come with challenges: direct distillation incurs extensive computational cost due to the need for numerous diffusion model samples; adversarial distillation introduces complexities associated with GAN training; progressive distillation requires multiple training stages and is less effective for one or two-step generation; and VSD can produce overly smooth samples with limited diversity and struggles at high guidance levels.

Consistency models (CMs) (Song et al., 2023; Song & Dhariwal, 2023) offer significant advantages in addressing these issues. They eliminate the need for supervision from diffusion model samples, avoiding the computational cost of generating synthetic datasets. CMs also bypass adversarial training, sidestepping its inherent difficulties. Aside from distillation, CMs can be trained from scratch with consistency training (CT), without relying on pre-trained diffusion models. Previous work (Song & Dhariwal, 2023; Geng et al., 2024; Luo et al., 2023; Xie et al., 2024a) has demonstrated the effectiveness of CMs in few-step generation, especially in one or two steps. However, these results are all based on discrete-time CMs, which introduces discretization errors and requires careful scheduling of the timestep grid, potentially leading to suboptimal sample quality. In contrast, continuous-time CMs avoid these issues but have faced challenges with training instability (Song et al., 2023; Song & Dhariwal, 2023; Geng et al., 2024).

In this work, we introduce techniques to simplify, stabilize, and scale up the training of continuous-time CMs. Our first contribution is TrigFlow, a new formulation that unifies EDM (Karras et al., 2022; 2024) and Flow Matching (Peluchetti, 2022; Lipman et al., 2022; Liu et al., 2022; Albergo et al., 2023; Heitz et al., 2023), significantly simplifying the formulation of diffusion models, the associated probability flow ODE and CMs. Building on this foundation, we analyze the root causes of instability in CM training and propose a complete recipe for mitigation. Our approach includes improved time-conditioning and adaptive group normalization within the network architecture. Additionally, we re-formulate the training objective for continuous-time CMs, incorporating adaptive weighting and normalization of key terms, and progressive annealing for stable and scalable training.

With these improvements, we elevate the performance of consistency models in both consistency training and distillation, achieving comparable or better results compared to previous discrete-time formulations. Our models, referred to as sCMs, demonstrate success across various datasets and model sizes. We train sCMs on CIFAR-10, ImageNet 64$\times$64, and ImageNet 512$\times$512, reaching an unprecedented scale with 1.5 billion parameters—the largest CMs trained to date (samples in Figure 2). We show that sCMs scale effectively with increased compute, achieving better sample quality in a predictable way. Moreover, when measured against state-of-the-art diffusion models, which require significantly more sampling compute, sCMs narrow the FID gap to within 10% using two-step generation. In addition, we provide a rigorous justification for the advantages of continuous-time CMs over discrete-time variants by demonstrating that sample quality improves as the gap between adjacent timesteps narrows to approach the continuous-time limit. Furthermore, we examine the differences between sCMs and VSD, finding that sCMs produce more diverse samples and are more compatible with guidance, whereas VSD tends to struggle at higher guidance levels.

Preliminaries

Given a training dataset, let $p_{d}$ denote its underlying data distribution and $\sigma_{d}$ its standard deviation. Diffusion models generate samples by learning to reverse a noising process that progressively perturbs a data sample ${\bm{x}}_{0}\sim p_{d}$ into a noisy version ${\bm{x}}_{t}=\alpha_{t}{\bm{x}}_{0}+\sigma_{t}{\bm{z}}_{t}$, where ${\bm{z}}_{t}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$ is standard Gaussian noise. This perturbation increases with $t\in[0,T]$, where larger $t$ indicates greater noise.

We consider two recent formulations for diffusion models.

2 Consistency Models

A consistency model (CM) (Song et al., 2023; Song & Dhariwal, 2023) is a neural network ${\bm{f}}_{\theta}({\bm{x}}_{t},t)$ trained to map the noisy input ${\bm{x}}_{t}$ directly to the corresponding clean data ${\bm{x}}_{0}$ in one step, by following the sampling trajectory of the PF-ODE starting at ${\bm{x}}_{t}$. A valid ${\bm{f}}_{\theta}$ must satisfy the boundary condition, ${\bm{f}}_{\theta}({\bm{x}},0)\equiv{\bm{x}}$. One way to meet this condition is to parameterize the consistency model as ${\bm{f}}_{\theta}({\bm{x}}_{t},t)=c_{\text{skip}}(t){\bm{x}}_{t}+c_{\text{out}}(t){\bm{F}}_{\theta}(c_{\text{in}}(t){\bm{x}}_{t},c_{\text{noise}}(t))$ with $c_{\text{skip}}(0)=1$ and $c_{\text{out}}(0)=0$. CMs are trained to have consistent outputs at adjacent time steps. Depending on how nearby time steps are selected, there are two categories of consistency models, as described below.

Discrete-time CMs. The training objective is defined at two adjacent time steps with finite distance:

Continuous-time CMs. When using $d({\bm{x}},{\bm{y}})=\|{\bm{x}}-{\bm{y}}\|_{2}^{2}$ and taking the limit $\Delta t\to 0$, Song et al. (2023) show that the gradient of Eq. (1) with respect to $\theta$ converges to

Simplifying Continuous-Time Consistency Models

Previous consistency models (CMs) adopt the model parameterization and diffusion process formulation in EDM (Karras et al., 2022). Specifically, the CM is parameterized as ${\bm{f}}_{\theta}({\bm{x}}_{t},t)=c_{\text{skip}}(t){\bm{x}}_{t}+c_{\text{out}}(t){\bm{F}}_{\theta}(c_{\text{in}}(t){\bm{x}}_{t},c_{\text{noise}}(t))$, where ${\bm{F}}_{\theta}$ is a neural network with parameters $\theta$. The coefficients $c_{\text{skip}}(t)$, $c_{\text{out}}(t)$, $c_{\text{in}}(t)$ are fixed to ensure that the variance of the diffusion objective is equalized across all time steps at initialization, and $c_{\text{noise}}(t)$ is a transformation of $t$ for better time conditioning. Since EDM diffusion process is variance-exploding (Song et al., 2021b), meaning that ${\bm{x}}_{t}={\bm{x}}_{0}+t{\bm{z}}_{t}$, we can derive that $c_{\text{skip}}(t)=\sigma_{d}^{2}/(t^{2}+\sigma_{d}^{2})$, $c_{\text{out}}(t)=\sigma_{d}\cdot t/\sqrt{\sigma_{d}^{2}+t^{2}}$, and $c_{\text{in}}(t)=1/\sqrt{t^{2}+\sigma_{d}^{2}}$ (see Appendix B.6 in Karras et al. (2022)). Although these coefficients are important for training efficiency, their complex arithmetic relationships with $t$ and $\sigma_{d}$ complicate theoretical analyses of CMs.

To simplify EDM and subsequently CMs, we propose TrigFlow, a formulation of diffusion models that keep the EDM properties but satisfy $c_{\text{skip}}(t)=\cos(t)$, $c_{\text{out}}(t)=\sin(t)$, and $c_{\text{in}}(t)\equiv 1/\sigma_{d}$ (proof in Appendix B). TrigFlow is a special case of flow matching (also known as stochastic interpolants or rectified flows) and v-prediction parameterization (Salimans & Ho, 2022). It closely resembles the trigonometric interpolant proposed by Albergo & Vanden-Eijnden (2023); Albergo et al. (2023); Ma et al. (2024), but is modified to account for $\sigma_{d}$, the standard deviation of the data distribution $p_{d}$. Since TrigFlow is a special case of flow matching and simultaneously satisfies EDM principles, it combines the advantages of both formulations while allowing the diffusion process, diffusion model parameterization, the PF-ODE, the diffusion training objective, and the CM parameterization to all have simple expressions, as provided below.

Diffusion Process. Given ${\bm{x}}_{0}\sim p_{d}({\bm{x}}_{0})$ and ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$, the noisy sample is defined as ${\bm{x}}_{t}=\cos(t){\bm{x}}_{0}+\sin(t){\bm{z}}$ for $t\in[0,\frac{\pi}{2}]$. As a special case, the prior sample ${\bm{x}}_{\frac{\pi}{2}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$.

Diffusion Models and PF-ODE. We parameterize the diffusion model as ${\bm{f}}_{\theta}({\bm{x}}_{t},t)={\bm{F}}_{\theta}({\bm{x}}_{t}/\sigma_{d},c_{\text{noise}}(t))$, where ${\bm{F}}_{\theta}$ is a neural network with parameters $\theta$, and $c_{\text{noise}}(t)$ is a transformation of $t$ to facilitate time conditioning. The corresponding PF-ODE is given by

Diffusion Objective. In TrigFlow, the diffusion model is trained by minimizing

where ${\bm{v}}_{t}=\cos(t){\bm{z}}-\sin(t){\bm{x}}_{0}$ is the training target.

Consistency Models. As mentioned in Sec. 2.2, a valid CM must satisfy the boundary condition ${\bm{f}}_{\theta}({\bm{x}},0)\equiv{\bm{x}}$. To enforce this condition, we parameterize the CM as the single-step solution of the PF-ODE in Eq. (3) using the first-order ODE solver (see Appendix B.1 for derivations). Specifically, CMs in TrigFlow take the form of

where $c_{\text{noise}}(t)$ is a time transformation for which we defer the discussion to Sec. 4.1.

Stabilizing Continuous-time Consistency Models

Training continuous-time CMs has been highly unstable (Song et al., 2023; Geng et al., 2024). As a result, they perform significantly worse compared to discrete-time CMs in prior works. To address this issue, we build upon the TrigFlow framework and introduce several theoretically motivated improvements to stabilize continuous-time CMs, with a focus on parameterization, network architecture, and training objectives.

where $\text{emb}(\cdot)$ refers to the time embeddings, typically in the form of either positional embeddings (Ho et al., 2020; Vaswani, 2017) or Fourier embeddings (Song et al., 2021b; Tancik et al., 2020) in the literature of diffusion models and CMs.

Below we describe improvements to stabilize each component from Eq. (7) in turns.

Identity Time Transformation ($\bm{c_{\text{noise}}(t)=t}$). Most existing CMs use the EDM formulation, which can be directly translated to the TrigFlow formulation as described in Appendix B.2. In particular, the time transformation becomes $c_{\text{noise}}(t)=\log(\sigma_{d}\tan t)$. Straightforward derivation shows that with this $c_{\text{noise}}(t)$, $\sin(t)\cdot\partial_{t}c_{\text{noise}}(t)=1/\cos(t)$ blows up whenever $t\to\frac{\pi}{2}$. To mitigate numerical instability, we propose to use $c_{\text{noise}}(t)=t$ as the default time transformation.

Positional Time Embeddings. For general time embeddings in the form of $\text{emb}(c)=\sin(s\cdot 2\pi\omega\cdot c+\phi)$, we have $\partial_{c}\text{emb}(c)=s\cdot 2\pi\omega\cos(s\cdot 2\pi\omega\cdot c+\phi)$. With larger Fourier scale $s$, this derivative has greater magnitudes and oscillates more vibrantly, causing worse instability. To avoid this, we use positional embeddings, which amounts to $s\approx 0.02$ in Fourier embeddings. This analysis provides a principled explanation for the observations in Song & Dhariwal (2023).

Adaptive Double Normalization. Song & Dhariwal (2023) found that the AdaGN layer (Dhariwal & Nichol, 2021), defined as ${\bm{y}}=\text{norm}({\bm{x}})\odot{\bm{s}}(t)+{\bm{b}}(t)$, negatively impacts CM training. Our modification is adaptive double normalization, defined as ${\bm{y}}=\text{norm}({\bm{x}})\odot\text{pnorm}({\bm{s}}(t))+\text{pnorm}({\bm{b}}(t))$, where $\text{pnorm}(\cdot)$ denotes pixel normalization (Karras, 2017). Empirically we find it retains the expressive power of AdaGN for diffusion training but removes its instability in CM training.

As shown in Figure 4, we visualize how our techniques stabilize the time-derivates for CMs trained on CIFAR-10. Empirically, we find that these improvements help stabilize the training dynamics of CMs without hurting diffusion model training (see Appendix G).

2 Training Objectives

Using the TrigFlow formulation in Sec. 3 and techniques proposed in Sec. 4.1, the gradient of continuous-time CM training in Eq. (2) becomes

Below we propose additional techniques to explicitly control this gradient for improved stability.

where $D$ is the dimensionality of ${\bm{x}}_{0}$, and we sample $\tan(t)$ from a log-Normal proposal distribution (Karras et al., 2022), that is, $e^{\sigma_{d}\tan(t)}\sim\mathcal{N}(P_{\text{mean}},P_{\text{std}}^{2})$ (details in Appendix G).

With all techniques in place, the stability of both discrete-time and continuous-time CM training substantially improves. We provide detailed algorithms for discrete-time CMs in Appendix E, and train continuous-time CMs and discrete-time CMs with the same setting. As demonstrated in Figure 5(c), increasing the number of discretization steps $N$ in discrete-time CMs improves sample quality by reducing discretization errors, but degrades once $N$ becomes too large (after $N>1024$) to suffer from numerical precision issues. By contrast, continuous-time CMs significantly outperform discrete-time CMs across all $N$’s which provides strong justification for choosing continuous-time CMs over discrete-time counterparts. We call our model sCM (s for simple, stable, and scalable), and provide detailed pseudo-code for sCM training in Appendix A.

Scaling up Continuous-Time Consistency Models

Below we test all the improvements proposed in previous sections by training large-scale sCMs on a variety of challenging datasets.

JVP of Flash Attention. Flash Attention (Dao et al., 2022; Dao, 2023) is widely used for attention computation in large-scale model training, providing both GPU memory savings and faster training. However, Flash Attention does not compute the Jacobian-vector product (JVP). To fill this gap, we propose a similar algorithm (detailed in Appendix F) that efficiently computes both softmax self-attention and its JVP in a single forward pass in the style of Flash Attention, significantly reducing GPU memory usage for JVP computation in attention layers.

2 Experiments

To test our improvements, we employ both consistency training (referred to as sCT) and consistency distillation (referred to as sCD) to train and scale continuous-time CMs on CIFAR-10 (Krizhevsky, 2009), ImageNet 64$\times$64 and ImageNet 512$\times$512 (Deng et al., 2009). We benchmark the sample quality using FID (Heusel et al., 2017). We follow the settings of Score SDE (Song et al., 2021b) on CIFAR10 and EDM2 (Karras et al., 2024) on both ImageNet 64$\times$64 and ImageNet 512$\times$512, while changing the parameterization and architecture according to Section 4.1. We adopt the method proposed by Song et al. (2023) for two-step sampling of both sCT and sCD, using a fixed intermediate time step $t=1.1$. For sCD models on ImageNet 512$\times$512, since the teacher diffusion model relies on classifier-free guidance (CFG) (Ho & Salimans, 2021), we incorporate an additional input $s$ into the model ${\bm{F}}_{\theta}$ to represent the guidance scale (Meng et al., 2023). We train the model with sCD by uniformly sampling $s\in$ and applying the corresponding CFG to the teacher model during distillation (more details are provided in Appendix G). For sCT models, we do not test CFG since it is incompatible with consistency training.

Training compute of sCM. We use the same batch size as the teacher diffusion model across all datasets. The effective compute per training iteration of sCD is approximately twice that of the teacher model. We observe that the quality of two-step samples from sCD converges rapidly, achieving results comparable to the teacher diffusion model using less than 20% of the teacher training compute. In practice, we can obtain high-quality samples after only 20k finetuning iterations with sCD.

Benchmarks. In Tables 1 and 2, we compare our results with previous methods by benchmarking the FIDs and the number of function evaluations (NFEs). First, sCM outperforms all previous few-step methods that do not rely on joint training with another network and is on par with, or even exceeds, the best results achieved with adversarial training. Notably, the 1-step FID of sCD-XXL on ImageNet 512$\times$512 surpasses that of StyleGAN-XL (Sauer et al., 2022) and VAR (Tian et al., 2024a). Furthermore, the two-step FID of sCD-XXL outperforms all generative models except diffusion and is comparable with the best diffusion models that require $63$ sequential steps. Second, the two-step sCM model significantly narrows the FID gap with the teacher diffusion model to within 10%, achieving FIDs of 2.06 on CIFAR-10 (compared to the teacher FID of 2.01), 1.48 on ImageNet 64×64 (teacher FID of 1.33), and 1.88 on ImageNet 512×512 (teacher FID of 1.73). Additionally, we observe that sCT is more effective at smaller scales but suffers from increased variance at larger scales, while sCD shows consistent performance across both small and large scales.

Scaling study. Based on our improved training techniques, we successfully scale continuous-time CMs without training instability. We train various sizes of sCMs using EDM2 configurations (S, M, L, XL, XXL) on ImageNet 64$\times$64 and 512$\times$512, and evaluate FID under optimal guidance scales, as shown in Fig. 6. First, as model FLOPs increase, both sCT and sCD show improved sample quality, showing that both methods benefit from scaling. Second, compared to sCD, sCT is more compute efficient at smaller resolutions but less efficient at larger resolutions. Third, sCD scales predictably for a given dataset, maintaining a consistent relative difference in FIDs across model sizes. This suggests that the FID of sCD decreases at the same rate as the teacher diffusion model, and therefore sCD is as scalable as the teacher diffusion model. As the FID of the teacher diffusion model decreases with scaling, the absolute difference in FID between sCD and the teacher model also diminishes. Finally, the relative difference in FIDs decreases with more sampling steps, and the sample quality of the two-step sCD becomes on par with that of the teacher diffusion model.

Comparison with VSD. Variational score distillation (VSD) (Wang et al., 2024; Yin et al., 2024b) and its multi-step generalization (Xie et al., 2024b; Salimans et al., 2024) represent another diffusion distillation technique that has demonstrated scalability on high-resolution images (Yin et al., 2024a). We apply one-step VSD from time $T$ to to finetune a teacher diffusion model using the EDM2-M configuration and tune both the weighting functions and proposal distributions for fair comparisons. As shown in Figure 7, we compare sCD, VSD, a combination of sCD and VSD (by simply adding the two losses), and the teacher diffusion model by sweeping over the guidance scale. We observe that VSD has artifacts similar to those from applying large guidance scales in diffusion models: it increases fidelity (as evidenced by higher precision scores) while decreasing diversity (as shown by lower recall scores). This effect becomes more pronounced with increased guidance scales, ultimately causing severe mode collapse. In contrast, the precision and recall scores from two-step sCD are comparable with those of the teacher diffusion model, resulting in better FID scores than VSD.

Conclusion

Our improved formulations, architectures, and training objectives have simplified and stabilized the training of continuous-time consistency models, enabling smooth scaling up to 1.5 billion parameters on ImageNet 512$\times$512. We ablated the impact of TrigFlow formulation, tangent normalization, and adaptive weighting, confirming their effectiveness. Combining these improvements, our method demonstrated predictable scalability across datasets and model sizes, outperforming other few-step sampling approaches at large scales. Notably, we narrowed the FID gap with the teacher model to within 10% using two-step generation, compared to state-of-the-art diffusion models that require significantly more sampling steps.

We would like to thank Allan Jabri, Aaron Lou, Alex Nichol, Huiwen Chang, Heewoo Jun, and Ishaan Gulrajani for technical discussions, and Mingyu Ye for assistance with plots and diagrams. We also appreciate the support from Mark Chen and Prafulla Dhariwal for this research project.

References

Appendix

We include additional derivations, experimental details, and results in the appendix. The detailed training algorithm for sCM, covering both sCT and sCD, is provided in Appendix A. We present a comprehensive discussion of the TrigFlow framework in Appendix B, including detailed derivations (Section B.1) and its connections with other parameterization (Section B.2). We introduce a new algorithm called adaptive variational score distillation in Appendix C, which eliminates the need for manually designed training weighting. Furthermore, we elaborate on a general framework for adaptive training weighting in Appendix D, applicable to diffusion models, consistency models, and variational score distillation. As our improvements discussed in Sec. 4 are also applicable for discrete-time consistency models, we provide detailed derivations and the training algorithm for discrete-time consistency models in Appendix E, incorporating all the improved techniques of sCM. We also provide a complete description of the Jacobian-vector product algorithm for Flash Attention in Appendix F. Finally, all experimental settings and evaluation results are listed in Appendix G, along with additional samples generated by our sCD-XXL model trained on ImageNet at 512×512 resolution in Appendix H.

Appendix A Training Algorithm of sCM

We provide the detailed algorithm of sCM in Algorithm 1, where we refer to consistency training of sCM as sCT and consistency distillation of sCM as sCD.

Appendix B TrigFlow: A Simple Framework Unifying EDM, Flow Matching and Velocity Prediction

Denote the standard deviation of the data distribution $p_{d}$ as $\sigma_{d}$. We consider a general forward diffusion process at time $t\in[0,T]$ with ${\bm{x}}_{t}=\alpha_{t}{\bm{x}}_{0}+\sigma_{t}{\bm{z}}$ for the data sample ${\bm{x}}_{0}\sim p_{d}$ and the noise sample ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$ (note that the variance of ${\bm{z}}$ is the same as that of the data ${\bm{x}}_{0}$)For any diffusion process with ${\bm{x}}_{t}=\alpha^{\prime}_{t}{\bm{x}}_{0}+\sigma^{\prime}_{t}\bm{\epsilon}$ where $\bm{\epsilon}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$, we can always equivalently convert it to ${\bm{x}}_{t}=\alpha^{\prime}_{t}{\bm{x}}_{0}+\frac{\sigma^{\prime}_{t}}{\sigma_{d}}\cdot(\sigma_{d}\bm{\epsilon})$ and let ${\bm{z}}\coloneqq\sigma_{d}\bm{\epsilon},\alpha_{t}\coloneqq\alpha^{\prime}_{t},\sigma_{t}\coloneqq\frac{\sigma^{\prime}_{t}}{\sigma_{d}}$. So the assumption for ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$ does not result in any loss of generality., where $\alpha_{t}>0,\sigma_{t}>0$ are noise schedules such that $\alpha_{t}/\sigma_{t}$ is monotonically decreasing w.r.t. $t$, with $\alpha_{0}=1,\sigma_{0}=0$. The general training loss for diffusion model can always be rewritten as

where different diffusion model formulation contains four different parts:

Parameterization of ${\bm{D}}_{\theta}$, such as score function (Song & Ermon, 2019; Song et al., 2021b), noise prediction model (Song & Ermon, 2019; Song et al., 2021b; Ho et al., 2020), data prediction model (Ho et al., 2020; Kingma et al., 2021; Salimans & Ho, 2022), velocity prediction model (Salimans & Ho, 2022), EDM (Karras et al., 2022) and flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo et al., 2023).

Noise schedule for $\alpha_{t}$ and $\sigma_{t}$, such as variance preserving process (Ho et al., 2020; Song et al., 2021b), variance exploding process (Song et al., 2021b; Karras et al., 2022), cosine schedule (Nichol & Dhariwal, 2021), and conditional optimal transport path (Lipman et al., 2022).

Weighting function for $w(t)$, such as uniform weighting (Ho et al., 2020; Nichol & Dhariwal, 2021; Karras et al., 2022), weighting by functions of signal-to-noise-ratio (SNR) (Salimans & Ho, 2022), monotonic weighting (Kingma & Gao, 2024) and adaptive weighting (Karras et al., 2024).

Proposal distribution for $t$, such as uniform distribution within $[0,T]$ (Ho et al., 2020; Song et al., 2021b), log-normal distribution (Karras et al., 2022), SNR sampler (Esser et al., 2024), and adaptive importance sampler (Song et al., 2021a; Kingma et al., 2021).

Below we show that, under the unit variance principle proposed in EDM (Karras et al., 2022), we can obtain a general but simple framework for all the above four parts, which can equivalently reproduce all previous diffusion models.

We consider the parameterization for ${\bm{D}}_{\theta}$ as the same principle in EDM (Karras et al., 2022) by

To ensure the training target of ${\bm{F}}_{\theta}$ has unit variance, we have

To reduce the error amplification from ${\bm{F}}_{\theta}$ to ${\bm{D}}_{\theta}$, we should ensure $c_{\text{out}}(t)$ to be as small as possible, which means we should take $c_{\text{skip}}(t)$ by letting $\frac{\partial c_{\text{out}}}{\partial c_{\text{skip}}}=0$, which results in

Though equivalent, we choose $c_{\text{out}}(t)=-\frac{\sigma_{d}\alpha_{t}}{\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}}$ which can simplify some derivations below.

In summary, the parameterization and objective for the general diffusion noise schedule are

Step 2: All noise schedules can be equivalently transformed.

One nice property of the unit variance principle is that the $\alpha_{t},\sigma_{t}$ in both the parameterization and the objective are homogenous, which means we can always assume $\alpha_{t}^{2}+\sigma_{t}^{2}=1$ without loss of generality. To see this, we can apply a simple change-of-variable of $\hat{\alpha}_{t}=\frac{\alpha_{t}}{\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}}$, $\hat{\sigma}_{t}=\frac{\sigma_{t}}{\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}}$ and $\hat{{\bm{x}}}_{t}=\frac{{\bm{x}}_{t}}{\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}}=\hat{\alpha}_{t}{\bm{x}}_{0}+\hat{\sigma}_{t}{\bm{z}}$, thus we have

As for the sampling procedure, according to DPM-Solver++ (Lu et al., 2022b), the exact solution of diffusion ODE from time $s$ to time $t$ satisfies

where $\lambda_{t}=\log\frac{\alpha_{t}}{\sigma_{t}}$, so the sampling procedure is also homogenous for $\alpha_{t},\sigma_{t}$. To see this, we can use the fact that $\frac{{\bm{x}}_{t}}{\sigma_{t}}=\frac{\hat{{\bm{x}}}_{t}}{\hat{\sigma}_{t}}$ and $\lambda_{t}=\log\frac{\hat{\alpha}_{t}}{\hat{\sigma}_{t}}\coloneqq\hat{\lambda}_{t}$, thus the above equation is equivalent to

which is exactly the sampling procedure of the diffusion process $\hat{{\bm{x}}_{t}}$, which means noise schedules of diffusion models won’t affect the performance of sampling. In other words, for any diffusion process $(\alpha_{t},\sigma_{t},{\bm{x}}_{t})$ at time $t$, we can always divide them by $\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}$ to obtain the diffusion process $(\hat{\alpha}_{t},\hat{\sigma}_{t},\hat{{\bm{x}}}_{t})$ with $\hat{\alpha}_{t}^{2}+\hat{\sigma}_{t}^{2}=1$ and all the parameterization, training objective and sampling procedure can be equivalently transformed. The only difference is the corresponding training weighting $w(t)\sigma^{2}_{d}\hat{\alpha}_{t}^{2}$ in Eq. (18), which we will discuss in the next step.

A straightforward corollary is that the “optimal transport path” (Lipman et al., 2022) in flow matching with $\alpha_{t}=1-t,\sigma_{t}=t$ can be equivalently converted to other noise schedules. The reason of its better empirical performance is essentially due to the different weighting during training and the lack of advanced diffusion sampler such as DPM-Solver series (Lu et al., 2022a; b) during sampling, not the “straight path” (Lipman et al., 2022) itself.

Step 3: Unified framework by TrigFlow.

As we showed in the previous step, we can always assume $\hat{\alpha}_{t}^{2}+\hat{\sigma}_{t}^{2}=1$. An equivalent change-of-variable of such constraint is to define

so $\hat{t}\in[0,\frac{\pi}{2}]$ is a monotonically increasing function of $t\in[0,T]$, thus there exists a one-one mapping between $t$ and $\hat{t}$ to convert the proposal distribution $p(t)$ to the distribution of $\hat{t}$, denoted as $p\left(\hat{t}\right)$. As $\hat{\alpha}_{t}=\cos\left(\hat{t}\right),\hat{\sigma}_{t}=\sin\left(\hat{t}\right)$, the training objective in Eq. (18) is equivalent to

For simplicity and with a slight abuse of notation, we omit the $\hat{t}$ and denote the whole training weighting as a single $w(t)$, we summarize the diffusion process, parameterization, training objective and samplers of TrigFlow as follows.

Diffusion Process.

${\bm{x}}_{0}\sim p_{d}({\bm{x}}_{0})$, ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$, ${\bm{x}}_{t}=\cos(t){\bm{x}}_{0}+\sin(t){\bm{z}}$ for $t\in[0,\frac{\pi}{2}]$.

Parameterization.

where $c_{\text{noise}}(t)$ is the conditioning input of the noise levels for ${\bm{F}}_{\theta}$, which can be arbitrary one-one mapping of $t$. Moreover, the parameterized diffusion ODE is defined by

Training Objective.

where $w(t)$ is the training weighting, which we will discuss in details in Appendix D.

As for the sampling procedure, although we can directly solve the diffusion ODE in Eq. (24) by Euler’s or Heun’s solvers as in flow matching (Lipman et al., 2022), the parameterization for $\sigma_{d}{\bm{F}}_{\theta}$ may be not the optimal parameterization for reducing the discreteization errors. As proved in DPM-Solver-v3 (Zheng et al., 2023c), the optimal parameterization should cancel all the linearity of the ODE, and the data prediction model ${\bm{D}}_{\theta}$ is an effective approximation of such parameterization. Thus, we can also apply DDIM, DPM-Solver and DPM-Solver++ for TrigFlow by rewriting the coefficients into the TrigFlow notation, as listed below.

st-order Sampler by DDIM.

Starting from ${\bm{x}}_{s}$ at time $s$, the solution ${\bm{x}}_{t}$ at time $t$ is

One good property of TrigFlow is that the 1st-order sampler can naturally support zero-SNR sampling (Lin et al., 2024) by letting $s=\frac{\pi}{2}$ without any numerical issues.

nd-order Sampler by DPM-Solver.

Starting from ${\bm{x}}_{s}$ at time $s$, by reusing a previous solution ${\bm{x}}_{s^{\prime}}$ at time $s^{\prime}$, the solution ${\bm{x}}_{t}$ at time $t$ is

where $\bm{\epsilon}_{\theta}({\bm{x}}_{t},t)=\sin(t){\bm{x}}_{t}+\cos(t)\sigma_{d}{\bm{F}}_{\theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},c_{\text{noise}}(t)\right)$ is the noise prediction model, and $r_{s}=\frac{\log\tan(s)-\log\tan(s^{\prime})}{\log\tan(s)-\log\tan(t)}$.

nd-order Sampler by DPM-Solver++.

Starting from ${\bm{x}}_{s}$ at time $s$, by reusing a previous solution ${\bm{x}}_{s^{\prime}}$ at time $s^{\prime}$, the solution ${\bm{x}}_{t}$ at time $t$ is

where $r_{s}=\frac{\log\tan(s)-\log\tan(s^{\prime})}{\log\tan(s)-\log\tan(t)}$.

B.2 Relationship with other parameterization

As previous diffusion models define the forward process with ${\bm{x}}_{t^{\prime}}=\alpha_{t^{\prime}}{\bm{x}}_{0}+\sigma_{t^{\prime}}\bm{\epsilon}=\alpha_{t^{\prime}}{\bm{x}}_{0}+\frac{\sigma_{t^{\prime}}}{\sigma_{d}}(\sigma_{d}\bm{\epsilon})$ for $\bm{\epsilon}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$, we can obtain the relationship between $t^{\prime}$ and TrigFlow time steps $t\in[0,\frac{\pi}{2}]$ by

Thus, we can always translate the notation from previous noise schedules to TrigFlow notations. Moreover, below we show that TrigFlow unifies different current frameworks for training diffusion models, including EDM, flow matching and velocity prediction.

As our derivations closely follow the unit variance principle proposed in EDM (Karras et al., 2022), our parameterization can be equivalently converted to EDM notations. Specifically, the transformation between TrigFlow $({\bm{x}}_{t},t)$ and EDM $({\bm{x}}_{\sigma},\sigma)$ is

The reason why TrigFlow notation is much simpler than EDM is just because we define the end point of the diffusion process as ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$ with the same variance as the data distribution. Thus, the unit variance principle can ensure that all the intermediate ${\bm{x}}_{t}$ does not need to multiply other coefficients as in EDM.

Flow Matching.

Flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo et al., 2023) defines a stochastic path between two samples ${\bm{x}}_{0}$ from data distribution and ${\bm{z}}$ from a tractable distribution which is usually some Gaussian distribution. For a general path ${\bm{x}}_{t}=\alpha_{t}{\bm{x}}_{0}+\sigma_{t}{\bm{z}}$ with $\alpha_{0}=1,\alpha_{T}=0,\sigma_{0}=0,\sigma_{T}=1$, the conditional probability path is

and it learns a parameterized model ${\bm{v}}_{\theta}({\bm{x}}_{t},t)$ by minimizing

and the final probability flow ODE is defined by

As TrigFlow uses $\alpha_{t}=\cos(t)$ and $\sigma_{t}=\sin(t)$, it is easy to see that the training objective and the diffusion ODE of TrigFlow are also the same as flow matching with ${\bm{v}}_{\theta}({\bm{x}}_{t},t)=\sigma_{d}{\bm{F}}_{\theta}(\frac{{\bm{x}}_{t}}{\sigma_{d}},c_{\text{noise}}(t))$. To the best of our knowledge, TrigFlow is the first framework that unifies EDM and flow matching for training diffusion models.

Velocity Prediction.

The velocity prediction parameterization (Salimans & Ho, 2022) trains a parameterization network with the target $\alpha_{t}{\bm{z}}-\sigma_{t}{\bm{x}}_{0}$. As TrigFlow uses $\alpha_{t}=\cos(t),\sigma_{t}=\sin(t)$, it is easy to see that the training target in TrigFlow is also the velocity.

Discussions on SNR.

Another good property of TrigFlow is that it can define a data-variance-invariant SNR. Specifically, previous diffusion models define the SNR at time $t$ as $\text{SNR}(t)=\frac{\alpha^{2}_{t}}{\sigma_{t}^{2}}$ for ${\bm{x}}_{t}=\alpha_{t}{\bm{x}}_{0}+\sigma_{t}\bm{\epsilon}$ with $\bm{\epsilon}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$. However, such definition ignores the influence of the variance of ${\bm{x}}_{0}$: if we rescale the data ${\bm{x}}_{0}$ by a constant, then such SNR doesn’t get rescaled correspondingly, which is not reasonable in practice. Instead, in TrigFlow we can define the SNR by

which is data-variance-invariant and also simple.

Appendix C Adaptive Variational Score Distillation in TrigFlow Framework

In this section, we propose the detailed derivation for variational score distillation (VSD) in TrigFlow framework and an improved objective with adaptive weighting.

where $\bm{\epsilon}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$, $q_{t}^{\theta}$ is the diffused distribution at time $t$ with the same forward diffusion process as $p_{t}$ while starting at $q^{\theta}_{0}$ as the distribution of ${\bm{g}}_{\theta}({\bm{z}})$, $w(t)$ is an ad-hoc training weighting (Poole et al., 2022; Wang et al., 2024; Yin et al., 2024b), and $t$ follows a proposal distribution such as uniform distribution. It is proved that the optimum of $q_{t}^{\theta}$ satisfies $q_{0}=p_{d}$ (Wang et al., 2024) and thus the distribution of the generator matches the data distribution.

Moreover, by denoting ${\bm{x}}_{t}^{\theta}\coloneqq\alpha_{t}{\bm{g}}_{\theta}({\bm{z}})+\sigma_{t}\bm{\epsilon}$ and taking the gradient w.r.t. $\theta$, we have

Therefore, we need to approximate the score functions $\nabla_{{\bm{x}}_{t}}\log q_{t}^{\theta}({\bm{x}}_{t})$ for the generator and $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ for the data distribution. VSD trains a diffusion model for samples from ${\bm{g}}_{\theta}({\bm{z}})$ to approximate $\nabla_{{\bm{x}}_{t}}\log q_{t}^{\theta}({\bm{x}}_{t})$ and uses a pretrained diffusion model to approximate $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$.

In this work, we train the diffusion model in TrigFlow framework, with $\alpha_{t}=\cos(t)$, $\sigma_{t}=\sigma_{d}\sin(t)$, $\hat{\sigma}=\sigma_{d}$, $T=\frac{\pi}{2}$. Specifically, assume we have a pretrained diffusion model ${\bm{F}}_{\text{pretrain}}$ parameterized by TrigFlow, and we train another diffusion model ${\bm{F}}_{\phi}$ to approximate the diffused generator distribution, by

where ${\bm{x}}_{t}=\cos(t){\bm{x}}_{0}+\sin(t){\bm{z}}$, ${\bm{v}}_{t}=\cos(t){\bm{z}}-\sin(t){\bm{x}}_{0}$, ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$, ${\bm{x}}_{0}={\bm{g}}_{\theta}({\bm{z}}^{\prime})$ with ${\bm{z}}^{\prime}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$. Moreover, the relationship between the ground truth diffusion model ${\bm{F}}_{\text{Diff}}({\bm{x}}_{t},t)$ and the score function $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ is

Thus, we train the generator ${\bm{g}}_{\theta}$ by the following gradient w.r.t. $\theta$:

which is equivalent to the gradient of the following objective:

where ${\bm{g}}_{\theta^{-}}({\bm{z}}^{\prime})$ is the same as ${\bm{g}}_{\theta}({\bm{z}}^{\prime})$ but stops the gradient for $\theta$. Note that the weighting functions used in previous works (Wang et al., 2024; Yin et al., 2024b) is proportional to $\frac{\sin^{2}(t)}{\cos(t)}$, thus the prior weighting is proportional to $\sin(t)\cos(t)$, which has a U-shape similar to the log-normal distribution used in Karras et al. (2022). Thus, we can instead use a log-normal proposal distribution and apply the adaptive weighting by training another weighting network $w_{\psi}(t)$. We refer to Appendix D for detailed discussions about the learnable adaptive weighting. Thus we can obtain the training objective, as listed below.

C.2 Training Objective

Training Objective of Adaptive Variational Score Distillation (aVSD).

And we also choose a proportional distribution of $t$ for estimating ${\mathcal{L}}_{\text{VSD}}(\theta,\psi)$ by $\log(\tan(t)\sigma_{d})\sim{\mathcal{N}}(P_{\text{mean}},P_{\text{std}}^{2})$ and tune these two hyperparameters (note that they may be different from the proposal distribution for training ${\mathcal{L}}_{\text{Diff}}(\phi)$, as detailed in Appendix G.

In addition, for consistency models ${\bm{f}}_{\theta}({\bm{x}}_{t},t)$, we choose ${\bm{z}}^{\prime}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$ and ${\bm{g}}_{\theta}({\bm{z}}^{\prime})\coloneqq{\bm{f}}_{\theta}({\bm{z}}^{\prime},\frac{\pi}{2})=-\sigma_{d}{\bm{F}}_{\theta}(\frac{{\bm{z}}^{\prime}}{\sigma_{d}},\frac{\pi}{2})$, and thus the corresponding objective is

Appendix D Adaptive Weighting For Diffusion Models, Consistency Models and Variational Score Distillation

We first list the objectives of diffusion models, consistency models and variational score distillation (VSD). For diffusion models, as shown in Eq. (25), the gradient of the objective is

where ${\bm{F}}_{\theta^{-}}$ is the same as ${\bm{F}}_{\theta}$ but stops the gradient w.r.t. $\theta$. For VSD, the gradient of the objective is

And for continuous-time CMs parameterized by TrigFlow, the objective is

where ${\bm{f}}_{\theta^{-}}$ is the same as ${\bm{f}}_{\theta}$ but stops the gradient w.r.t. $\theta$. Interestingly, all these objectives can be rewritten into a form of inner product between a neural network and a target function which has the same dimension (denoted as $D$) as the output of the neural network. Specifically, assume the neural network is ${\bm{F}}_{\theta}$ parameterized by $\theta$, we study the following objective:

where we do not compute the gradients w.r.t. $\theta$ for ${\bm{y}}$. In such case, the gradient will be equivalent to

where ${\bm{F}}_{\theta^{-}}$ is the same as ${\bm{F}}_{\theta}$ but stops the gradient w.r.t. $\theta$. In such case, we can balance the gradient variance w.r.t. $t$ by training an adaptive weighting network $w_{\phi}(t)$ to estimate the loss norm, i.e., minimizing

This is the adaptive weighting proposed by EDM2 (Karras et al., 2024), which balances the loss variance across different time steps, inspired by the uncertainty estimation of Sener & Koltun (2018). By taking the partial derivative w.r.t. $w$ in the above equation, it is easy to verify that the optimal $w^{*}(t)$ satisfies

Therefore, the adaptive weighting reduces the loss variance across different time steps. In such case, all we need to do is to choose

A prior weighting $\lambda(t)$ for ${\bm{y}}$, which may be helpful for further reducing the variance of ${\bm{y}}$. Then the objective becomes

A proposal distribution for sampling the training $t$, which determines which part of $t$ we should focus on more. For diffusion models, we generally need to focus on the intermediate time steps since both the clean data and pure noise cannot provide precise training signals. Thus, the common choice is to choose a normal distribution over the log-SNR of time steps, which is proposed by Karras et al. (2022) and also known as log-normal distribution.

In this way, we do not need to manually choose the weighting functions, significantly reducing the tuning complexity of training diffusion models, CMs and VSD.

Appendix E Discrete-time Consistency Models with Improved Training Objectives

Note that the improvements proposed in Sec. 4 can also be applied to discrete-time consistency models (CMs). In this section, we discuss the improved version of discrete-time CMs for consistency distillation.

Parameterization. We also parameterize the CM by TrigFlow:

Reference sample by DDIM. Assume we sample ${\bm{x}}_{0}\sim p_{d}$, ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\sigma_{d}^{2}{\bm{I}})$, and ${\bm{x}}_{t}=\cos(t){\bm{x}}_{0}+\sin(t){\bm{z}}$, we need a reference sample ${\bm{x}}_{t^{\prime}}$ at time $t^{\prime}<t$ to guide the training of the CM, which can be obtained by one-step DDIM from $t$ to $t^{\prime}$:

Thus, the output of the consistency model at time $t^{\prime}$ is

Original objective of discrete-time CMs. The consistency model at time $t$ can be rewritten into

Therefore, by computing the difference between Eq. (35) and Eq. (36), we define

where $w_{\phi}(t)$ is the adaptive weighting network.

Tangent normalization for discrete-time CMs. We apply the simliar tangent normalization method as continuous-time CMs by defining

where $c>0$ is a hyperparameter, and then the objective in Eq. (39) becomes

Tangent warmup for discrete-time CMs. We replace the $\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})$ with the warmup version:

where $r$ linearly increases from to $1$ over the first 10k training iterations.

We provide the detailed algorithm of discrete-time sCM (dsCM) in Algorithm 2, where we refer to consistency distillation of discrete-time sCM as dsCD.

E.2 Experiments of Discrete-time sCM

We use the algorithm in Algorithm 2 to train discrete-time sCM, where we split $[0,\frac{\pi}{2}]$ into $N$ intervals by EDM sampling spacing. Specifically, we first obtain the EDM time step by $\sigma_{i}=(\sigma_{\text{min}}^{1/\rho}+\frac{i}{M}(\sigma_{\text{max}}^{1/\rho}-\sigma_{\text{min}}^{1/\rho}))^{\rho}$ with $\rho=7,\sigma_{\text{min}}=0.002$ and $\sigma_{\text{max}}=80$, and then obtain $t_{i}=\arctan(\sigma_{i}/\sigma_{d})$ and set $t_{0}=0$. During training, we sample $t$ with a discrete categorical distribution that splits the log-normal proposal distribution as used in continuous-time sCM, similar to Song & Dhariwal (2023).

As demonstrated in Figure 5(c), increasing the number of discretization steps $N$ in discrete-time CMs improves sample quality by reducing discretization errors, but obviously degrades once $N$ becomes too large (after $N>1024$) to suffer from numerical precision issues. By contrast, continuous-time CMs significantly outperform discrete-time CMs across all $N$’s which provides strong justification for choosing continuous-time CMs over discrete-time counterparts.

Appendix F Jacobian-Vector Product of Flash Attention

However, to the best of knowledge, there does not exist an algorithm for computing the Jacobian-Vector product of the attention operator in the Flash Attention style for faster computation and memory saving. We propose a recursive algorithm for the JVP computation of Flash Attention below.

and for combining two blocks ${\bm{x}}=[{\bm{x}}^{(a)},{\bm{x}}^{(b)}]$, we merge their corresponding ${\bm{g}}$ and $\mu$ by

Therefore, we can use a single loop to obtain both the output ${\bm{y}}$ and the JVP output ${\bm{t}}_{{\bm{y}}}$, which accesses the memory for the attention matrices only once and avoids saving the intermediate activations, thus saving the GPU memory.

Appendix G Experiment Settings and Results

We train the teacher diffusion models on CIFAR-10, ImageNet 64$\times$64 and ImageNet 512$\times$512 with the proposed improvements of parameterization and architecture, including TrigFlow parameterization, positional time embedding and adaptive double normalization layer. We list the detailed settings below.

CIFAR-10. Our architecture is based on the Score SDE (Song et al., 2021b) architecture (DDPM++). We use the same settings of EDM (Karras et al., 2022): dropout rate is 0.13, batch size is 512, number of training iterations is 400k, learning rate is 0.001, Adam $\epsilon=10^{-8}$, $\beta_{1}=0.9$, $\beta_{2}=0.999$. We use 2nd-order single-step DPM-Solver (Lu et al., 2022a) (DPM-Solver-2S) with Heun’s intermediate time step with 18 steps (NFE=35), which is exactly equivalent to EDM Heun’s sampler. We obtain FID of 2.15 for the teacher model.

ImageNet 64$\times$64. We preprocess the ImageNet dataset following Dhariwal & Nichol (2021) by

Resize the shorter width / height to $64\times 64$ resolution with bicubic interpolation.

Disable data augmentation such as horizontal flipping.

Except for the TrigFlow parameterization, positional time embedding and adaptive double normalization layer, we follow exactly the same setting in EDM2 config G (Karras et al., 2024) to train models with sizes of S, M, L, and XL, while the only difference is that we use Adam $\epsilon=10^{-11}$.

ImageNet 512$\times$512. We preprocess the ImageNet dataset following Dhariwal & Nichol (2021) and Karras et al. (2024) by

Resize the shorter width / height to $512\times 512$ resolution with bicubic interpolation.

Disable data augmentation such as horizontal flipping.

Encode the images into latents by stable diffusion VAEhttps://huggingface.co/stabilityai/sd-vae-ft-mse (Rombach et al., 2022; Janner et al., 2022), and rescale the latents by channel mean $\mu_{c}=[1.56,-0.695,0.483,0.729]$ and channel std $\sigma_{c}=[5.27,5.91,4.21,4.31]$. We keep the $\sigma_{d}=0.5$ as in EDM2 (Karras et al., 2024), so for each latent we substract $\mu_{c}$ and multiply it by $\sigma_{d}/\sigma_{c}$.

When sampling from the model, we redo the scaling of the generated latents and then run the VAE decoder. Notably, our channel mean and channel std are different from those in EDM2 (Karras et al., 2024). It is because when training the VAE, the images are normalized to $$before passing to the encoder. However, the channel mean and std used in EDM2 assumes the input images are in$$ range, which mismatches the training phase of the VAE. We empirically find that it is hard to distinguish the reconstructed samples by human eyes of these two different normalization, while it has non-ignorable influence for training diffusion models evaluated by FID. After fixing this mismatch, our diffusion model slightly outperforms the results of EDM2 at larger scales (XL and XXL). More results are provided in Table 5.

Except for the TrigFlow parameterization, positional time embedding and adaptive double normalization layer, we follow exactly the same setting in EDM2 config G (Karras et al., 2024) to train models with sizes of S, M, L, XL and XXL, while the only difference is that we use Adam $\epsilon=10^{-11}$. We enable label dropout with rate $0.1$ to support classifier-free guidance. We use 2nd-order single-step DPM-Solver (Lu et al., 2022a) (DPM-Solver-2S) with Heun’s intermediate time step with 32 steps (NFE=63), which is exactly equivalent to EDM Heun’s sampler. We find that the optimal guidance scale for classifier-free guidance and the optimal EMA rate are also the same as EDM2 for all model sizes.

G.2 Continuous-time Consistency Models

In all experiments, we use $c=0.1$ for tangent normalization, and use $H=10000$ for tangent warmup. We always use the same batch size as the teacher diffusion training, which is different from Song & Dhariwal (2023). During sampling, we start at $t_{\text{max}}=\arctan\left(\frac{\sigma_{\text{max}}}{\sigma_{d}}\right)$ with $\sigma_{\text{max}}=80$ such that it matches the starting time of EDM (Karras et al., 2022) and EDM2 (Karras et al., 2024). For 2-step sampling, we use the algorithm in Song et al. (2023) with an intermediate $t=1.1$ for all the experiments. We always initialize the CM from the EMA parameters of the teacher diffusion model. For sCD, we always use the ${\bm{F}}_{\text{pretrain}}$ of the teacher diffusion model with its EMA parameters during distillation.

We empirically find that the proposal distribution should have small $P_{\text{mean}}$, i.e. close to the clean data, to ensure the training stability and improve the final performance. Intuitively, this is because the training signal of CMs only come from the clean data, so we need to reduce the training error for $t$ near to to further reduce the accumulation errors.

CIFAR-10. For both sCT and sCD, we initialize from the teacher diffusion model trained with the settings in Section G.1, and use RAdam optimizer (Liu et al., 2019) with learning rate of $0.0001$, $\beta_{1}=0.9$, $\beta_{2}=0.99$, $\epsilon=10^{-8}$, and without learning rate schedulers. proposal distribution of $P_{\text{mean}}=-1.0,P_{\text{std}}=1.4$. For the attention layers, we use the implementation in (Karras et al., 2022) which naturally supports JVP by PyTorch (Paszke et al., 2019) auto-grad. We use EMA half-life of $0.5$ Mimg (Karras et al., 2022). We use dropout rate of $0.20$ for sCT and disable dropout for sCD.

ImageNet 64$\times$64. We only enable dropout at the resolutions equal to or less than $16$, following Simple Diffusion (Hoogeboom et al., 2023) and iCT (Song & Dhariwal, 2023). We multiply the learning rate of the teacher diffusion model by $0.01$ for both sCT and sCD. We train the model with half precision (FP16), and use the flash attention jvp proposed in Appendix F for computing the tangents of flash attention layers. Other training settings are the same as the teacher diffusion models. More details of training and sampling are provided in Table 3 and Table 7. During sampling, we always use EMA length $\sigma_{\text{rel}}=0.05$ for sampling from CMs.

ImageNet 512$\times$512. We only enable dropout at the resolutions equal to or less than $16$, following Simple Diffusion (Hoogeboom et al., 2023) and iCT (Song & Dhariwal, 2023). We multiply the learning rate of the teacher diffusion model by $0.01$ for both sCT and sCD. We train the model with half precision (FP16), and use the flash attention jvp proposed in Appendix F for computing the tangents of flash attention layers. Other training settings are the same as the teacher diffusion models. More details of training and sampling are provided in Table 4 and Table 5. During sampling, we always use EMA length $\sigma_{\text{rel}}=0.05$ for sampling from CMs.

We add an additional input in ${\bm{F}}_{\theta}(\frac{{\bm{x}}_{t}}{\sigma_{d}},t,s)$ where $s$ represents the CFG guidance scale of the teacher model, where $s$ is embedded by positioinal embedding layer and an additional linear layer, and the embedding is added to the embedding of $t$, similar to the label conditioning. During training, we uniformly sample $s\in$ and apply CFG with guidance scale $s$ to the teacher diffusion model to get ${\bm{F}}_{\text{pretrain}}$.

VSD experiments. We do not use EMA for ${\bm{F}}_{\phi}$ in VSD, instead we always use the original model for ${\bm{F}}_{\phi}$ for stabilizing the training. The learning rate of ${\bm{F}}_{\phi}$ is the same as the learning rate of CMs. More details and results are provided in Tables 4, 5 and 6.

Appendix H Additional Samples