Simplifying, Stabilizing and Scaling Continuous-time Consistency Models

## Abstract

Abstract Consistency models (CMs) are a powerful class of diffusion-based generative models optimized for fast sampling. Most existing CMs are trained using discretized timesteps, which introduce additional hyperparameters and are prone to discretization errors. While continuous-time formulations can mitigate these issues, their success has been limited by training instability. To address this, we propose a simplified theoretical framework that unifies previous parameterizations of diffusion models and CMs, identifying the root causes of instability. Based on this analysis, we introduce key improvements in diffusion process parameterization, network architecture, and training objectives. These changes enable us to train continuous-time CMs at an unprecedented scale, reaching 1.5B parameters on ImageNet 512×512. Our proposed training algorithm, using only two sampling steps, achieves FID scores of 2.06 on CIFAR-10, 1.48 on ImageNet 64×64, and 1.88 on ImageNet 512×512, narrowing the gap in FID scores with the best existing diffusion models to within 10%.

## 1 Introduction

Figure: Figure 1: Sample quality vs. effective sampling compute (billion parameters $\times$ number of function evaluations during sampling). We compare the sample quality of different models on ImageNet 512$\times$512, measured by FID ($\downarrow$). Our 2-step sCM achieves sample quality comparable to the best previous generative models while using less than 10% of the effective sampling compute.
Refer to caption: x1.png

Diffusion models (Sohl-Dickstein et al., ; Song & Ermon, ; Ho et al., ; Song et al., ) have revolutionized generative AI, achieving remarkable results in image (Rombach et al., ; Ramesh et al., ; Ho et al., ), 3D (Poole et al., ; Wang et al., ; Liu et al., ), audio (Liu et al., ; Evans et al., ), and video generation (Blattmann et al., ; Brooks et al., ). 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, ; Zheng et al., ), adversarial distillation (Wang et al., ; Sauer et al., ), progressive distillation (Salimans & Ho, ), and variational score distillation (VSD) (Wang et al., ; Yin et al., ; ; Luo et al., ; Xie et al., ; Salimans et al., ). 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., ; Song & Dhariwal, ) 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, ; Geng et al., ; Luo et al., ; Xie et al., ) 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., ; Song & Dhariwal, ; Geng et al., ).

Figure: Figure 2: Selected 2-step samples from a continuous-time consistency model trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-intro/54.png

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., ; ) and Flow Matching (Peluchetti, ; Lipman et al., ; Liu et al., ; Albergo et al., ; Heitz et al., ), 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.

## 2 Preliminaries

### 2.1 Diffusion Models

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.

EDM (Karras et al., ; ).
The noising process simply sets $\alpha_{t}=1$ and $\sigma_{t}=t$, with the training objective given by
$\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[w(t)\left\|{\bm{f}}_{\theta}({\bm{x}
}_{t},t)-{\bm{x}}_{0}\right\|_{2}^{2}\right]$,
where $w(t)$ is a weighting function. The diffusion model 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$, and $c_{\text{skip}}$, $c_{\text{out}}$, $c_{\text{in}}$, and $c_{\text{noise}}$ are manually designed coefficients that ensure the training objective has the unit variance across timesteps at initialization. For sampling, EDM solves the probability flow ODE (PF-ODE)  (Song et al., ), defined by
$\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}=[{\bm{x}}_{t}-{\bm{f}}_{\theta}({
\bm{x}}_{t},t)]/t$, starting from ${\bm{x}}_{T}\sim{\mathcal{N}}({\bm{0}},T^{2}{\bm{I}})$ and stopping at ${\bm{x}}_{0}$.

Flow Matching.
The noising process uses differentiable coefficients $\alpha_{t}$ and $\sigma_{t}$, with time derivatives denoted by $\alpha_{t}^{\prime}$ and $\sigma_{t}^{\prime}$ (typically, $\alpha_{t}=1-t$ and $\sigma_{t}=t$). The training objective is given by
$\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[w(t)\left\|{\bm{F}}_{\theta}({\bm{x}
}_{t},t)-(\alpha_{t}^{\prime}{\bm{x}}_{0}+\sigma_{t}^{\prime}{\bm{z}})\right\|
_{2}^{2}\right]$,
where $w(t)$ is a weighting function and ${\bm{F}}_{\theta}$ is a neural network parameterized by $\theta$. The sampling procedure begins at $t=1$ with ${\bm{x}}_{1}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$ and solves the probability flow ODE (PF-ODE), defined by
$\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}={\bm{F}}_{\theta}({\bm{x}}_{t},t)$,
from $t=1$ to $t=0$.

### 2.2 Consistency Models

Figure: Figure 3: Discrete-time CMs (top & middle) vs. continuous-time CMs (bottom). Discrete-time CMs suffer from discretization errors from numerical ODE solvers, causing imprecise predictions during training. In contrast, continuous-time CMs stay on the ODE trajectory by following its tangent direction with infinitesimal steps.
Refer to caption: x2.png

A consistency model (CM) (Song et al., ; Song & Dhariwal, ) 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:

$$ $\mathbb{E}_{{\bm{x}}_{t},t}\left[w(t)d({\bm{f}}_{\theta}({\bm{x}}_{t},t),{\bm{ f}}_{\theta^{-}}({\bm{x}}_{t-\Delta t},t-\Delta t))\right],$ (1) $$

where $\theta^{-}$ denotes $\operatorname{stopgrad}(\theta)$, $w(t)$ is the weighting function, $\Delta t>0$ is the distance between adjacent time steps, and $d(\cdot,\cdot)$ is a metric function; common choices are $\ell_{2}$ loss $d({\bm{x}},{\bm{y}})=\|{\bm{x}}-{\bm{y}}\|_{2}^{2}$, Pseudo-Huber loss $d({\bm{x}},{\bm{y}})=\sqrt{\|{\bm{x}}-{\bm{y}}\|_{2}^{2}+c^{2}}-c$ for $c>0$ (Song & Dhariwal, ), and LPIPS loss (Zhang et al., ). Discrete-time CMs are sensitive to the choice of $\Delta t$, and therefore require manually designed annealing schedules (Song & Dhariwal, ; Geng et al., ) for fast convergence. The noisy sample ${\bm{x}}_{t-\Delta t}$ at the preceding time step $t-\Delta t$ is often obtained from ${\bm{x}}_{t}$ by solving the PF-ODE with numerical ODE solvers using step size $\Delta t$, which can cause additional discretization errors.

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. () show that the gradient of Eq. (1) with respect to $\theta$ converges to

$$ $\nabla_{\theta}\mathbb{E}_{{\bm{x}}_{t},t}\left[w(t){\bm{f}}^{\top}_{\theta}({ \bm{x}}_{t},t)\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)}{\mathrm{d }t}\right],$ (2) $$

where $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)}{\mathrm{d}t}=\nabla_{{
\bm{x}}_{t}}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)\frac{\mathrm{d}{\bm{x}}_{t}}
{\mathrm{d}t}+\partial_{t}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)$ is the tangent of ${\bm{f}}_{\theta^{-}}$ at $({\bm{x}}_{t},t)$ along the trajectory of the PF-ODE $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}$. Notably, continuous-time CMs do not rely on ODE solvers, which avoids discretization errors and offers more accurate supervision signals during training. However, previous work (Song et al., ; Geng et al., ) found that training continuous-time CMs, or even discrete-time CMs with an extremely small $\Delta t$, suffers from severe instability in optimization. This greatly limits the empirical performance and adoption of continuous-time CMs.

Consistency Distillation and Consistency Training.
Both discrete-time and continuous-time CMs can be trained using either *consistency distillation (CD)* or *consistency training (CT)*. In consistency distillation, a CM is trained by distilling knowledge from a pretrained diffusion model. This diffusion model provides the PF-ODE, which can be directly plugged into Eq. (2) for training continuous-time CMs. Furthermore, by numerically solving the PF-ODE to obtain ${\bm{x}}_{t-\Delta t}$ from ${\bm{x}}_{t}$, one can also train discrete-time CMs via Eq. (1). Consistency training (CT), by contrast, trains CMs from scratch without the need for pretrained diffusion models, which establishes CMs as a standalone family of generative models in their own right. Specifically, CT approximates ${\bm{x}}_{t-\Delta t}$ in discrete-time CMs as ${\bm{x}}_{t-\Delta t}=\alpha_{t-\Delta t}{\bm{x}}_{0}+\sigma_{t-\Delta t}{\bm{
z}}$, reusing the same data ${\bm{x}}_{0}$ and noise ${\bm{z}}$ when sampling ${\bm{x}}_{t}=\alpha_{t}{\bm{x}}_{0}+\sigma_{t}{\bm{z}}$. In the continuous-time limit, as $\Delta t\to 0$, this approach yields an unbiased estimate of the PF-ODE $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}\to\alpha_{t}^{\prime}{\bm{x}}_{0}+
\sigma_{t}^{\prime}{\bm{z}}$, leading to an unbiased estimate of Eq. (2) for training continuous-time CMs.

## 3 Simplifying Continuous-Time Consistency Models

Previous consistency models (CMs) adopt the model parameterization and diffusion process formulation in EDM (Karras et al., ). 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., ), 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. ()). 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, ). It closely resembles the trigonometric interpolant proposed by Albergo & Vanden-Eijnden (); Albergo et al. (); Ma et al. (), 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

$$ $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}=\sigma_{d}{\bm{F}}_{\theta}\left( \frac{{\bm{x}}_{t}}{\sigma_{d}},c_{\text{noise}}(t)\right).$ (3) $$

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

$$ ${\mathcal{L}}_{\text{Diff}}(\theta)=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[ \left\|\sigma_{d}{\bm{F}}_{\theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},c_{ \text{noise}}\left(t\right)\right)-{\bm{v}}_{t}\right\|_{2}^{2}\right],$ (4) $$

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

$$ ${\bm{f}}_{\theta}({\bm{x}}_{t},t)=\cos(t){\bm{x}}_{t}-\sin(t)\sigma_{d}{\bm{F} }_{\theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},c_{\text{noise}}(t)\right),$ (5) $$

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

## 4 Stabilizing Continuous-time Consistency Models

Training continuous-time CMs has been highly unstable (Song et al., ; Geng et al., ). 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.

### 4.1 Parameterization and Network Architecture

Key to the training of continuous-time CMs is Eq. (2), which depends on the tangent function $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)}{\mathrm{d}t}$. Under the TrigFlow formulation, this tangent function is given by

$$ $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)}{\mathrm{d}t}=-\cos(t) \left(\sigma_{d}{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t \right)-\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}\right)-\sin(t)\left({\bm{x} }_{t}+\sigma_{d}\frac{\mathrm{d}{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}} {\sigma_{d}},t\right)}{\mathrm{d}t}\right),$ (6) $$

where $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}$ represents the PF-ODE, which is either estimated using a pretrained diffusion model in consistency distillation, or using an unbiased estimator calculated from noise and clean samples in consistency training.

To stabilize training, it is necessary to ensure the tangent function in Eq. (6) is stable across different time steps. Empirically, we found that $\sigma_{d}{\bm{F}}_{\theta^{-}}$, the PF-ODE $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}$, and the noisy sample ${\bm{x}}_{t}$ are all relatively stable. The only term left in the tangent function now is $\sin(t)\frac{\mathrm{d}{\bm{F}}_{\theta^{-}}}{\mathrm{d}t}=\sin(t)\nabla_{{\bm
{x}}_{t}}{\bm{F}}_{\theta^{-}}\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}+\sin(
t)\partial_{t}{\bm{F}}_{\theta^{-}}$. After further analysis, we found $\nabla_{{\bm{x}}_{t}}{\bm{F}}_{\theta^{-}}\frac{\mathrm{d}{\bm{x}}_{t}}{
\mathrm{d}t}$ is typically well-conditioned, so instability originates from the time-derivative $\sin(t)\partial_{t}{\bm{F}}_{\theta^{-}}$, which can be decomposed according to

$$ $\sin(t)\partial_{t}{\bm{F}}_{\theta^{-}}=\sin(t)\frac{\partial c_{\text{noise} }(t)}{\partial t}\cdot\frac{\partial\text{emb}(c_{\text{noise}})}{\partial c_{ \text{noise}}}\cdot\frac{\partial{\bm{F}}_{\theta^{-}}}{\partial\text{emb}(c_{ \text{noise}})},$ (7) $$

where $\text{emb}(\cdot)$ refers to the time embeddings, typically in the form of either positional embeddings (Ho et al., ; Vaswani, ) or Fourier embeddings (Song et al., ; Tancik et al., ) 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 ().

Adaptive Double Normalization.
Song & Dhariwal () found that the AdaGN layer (Dhariwal & Nichol, ), 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, ). 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).

Figure: Figure 4: Stability of different formulations. We show the norms of both terms in $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}=\nabla_{{\bm{x}}}{\bm{f}}_ {\theta^{-}}\cdot\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}+\partial_{t}{\bm{f }}_{\theta^{-}}$ for diffusion models trained with the EDM ($c_{\text{noise}}(t)=\log(\sigma_{d}\tan(t))$) and TrigFlow ($c_{\text{noise}}(t)=t$) formulations using different time embeddings. We observe that large Fourier scales in Fourier embeddings cause instabilities. In addition, the EDM formulation suffers from numerical issues when $t\to\frac{\pi}{2}$, while TrigFlow (using positional embeddings) has stable partial derivatives for both ${\bm{x}}_{t}$ and $t$.
Refer to caption: x5.png

Figure: Figure 5: Comparing different training objectives for consistency distillation. The diffusion models are EDM2 (Karras et al., ) pretrained on ImageNet 512$\times$512. (a) 1-step and 2-step sampling of continuous-time CMs trained by using raw tangents $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$, clipped tangents $\text{clip}(\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t},-1,1)$ and normalized tangents $(\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t})/(\|\frac{\mathrm{d}{\bm{ f}}_{\theta^{-}}}{\mathrm{d}t}\|+0.1)$. (b) Quality of 1-step and 2-step samples from continuous-time CMs trained w/ and w/o adaptive weighting, both are w/ tangent normalization. (c) Quality of 1-step samples from continuous-time CMs vs. discrete-time CMs using varying number of time steps ($N$), trained using all techniques in Sec. 4.
Refer to caption: x8.png

### 4.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

$$ $\nabla_{\theta}\mathbb{E}_{{\bm{x}}_{t},t}\bigg{[}-w(t)\sigma_{d}\sin(t){\bm{F }}^{\top}_{\theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)\frac{\mathrm{ d}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)}{\mathrm{d}t}\bigg{]}.$ $$

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

Tangent Normalization.
As discussed in Sec. 4.1, most gradient variance in CM training comes from the tangent function $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$. We propose to explicitly normalize the tangent function by replacing $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$ with $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}/(\|\frac{\mathrm{d}{\bm{f}
}_{\theta^{-}}}{\mathrm{d}t}\|+c)$, where we empirically set $c=0.1$. Alternatively, we can clip the tangent within $[-1,1]$, which also caps its variance. Our results in Figure 5(a) demonstrate that either normalization or clipping leads to substantial improvements for the training of continuous-time CMs.

Adaptive Weighting.
Previous works (Song & Dhariwal, ; Geng et al., ) design weighting functions $w(t)$ manually for CM training, which can be suboptimal for different data distributions and network architectures. Following EDM2 (Karras et al., ), we propose to train an adaptive weighting function alongside the CM, which not only eases the burden of hyperparameter tuning but also outperforms manually designed weighting functions with better empirical performance and negligible training overhead. Key to our approach is the observation that $\nabla_{\theta}\mathbb{E}[{\bm{F}}_{\theta}^{\top}{\bm{y}}]=\frac{1}{2}\nabla_
{\theta}\mathbb{E}[\|{\bm{F}}_{\theta}-{\bm{F}}_{\theta^{-}}+{\bm{y}}\|_{2}^{2}]$, where ${\bm{y}}$ is an arbitrary vector independent of $\theta$. When training continuous-time CMs using Eq. (2), we have ${\bm{y}}=-w(t)\sigma_{d}\sin(t)\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{
d}t}$. This observation allows us to convert Eq. (2) into the gradient of an MSE objective. We can therefore use the same approach in Karras et al. () to train an adaptive weighting function that minimizes the variance of MSE losses across time steps (details in Appendix D).
In practice, we find that integrating a prior weighting $w(t)=\frac{1}{\sigma_{d}\tan(t)}$ further reduces training variance. By incorporating the prior weighting, we train both the network ${\bm{F}}_{\theta}$ and the adaptive weighting function $w_{\phi}(t)$ by minimizing

$$ $\mathcal{L}_{\text{sCM}}(\theta,\phi)\!\coloneqq\!\mathbb{E}_{{\bm{x}}_{t},t} \left[\!\frac{e^{w_{\phi}(t)}}{D}\left\|{\bm{F}}_{\theta}\left(\frac{{\bm{x}}_ {t}}{\sigma_{d}},t\right)-{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{ \sigma_{d}},t\right)-\cos(t)\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}({\bm{x}}_{t} ,t)}{\mathrm{d}t}\right\|_{2}^{2}-w_{\phi}(t)\right],$ (8) $$

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

Diffusion Finetuning and Tangent Warmup.
For consistency distillation, we find that finetuning the CM from a pretrained diffusion model can speed up convergence, which is consistent with Song et al. (); Geng et al. (). Recall that in Eq. (6), the tangent $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$ can be decomposed into two parts: the first term $\cos(t)(\sigma_{d}{\bm{F}}_{\theta^{-}}-\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{
d}t})$ is relatively stable, whereas the second term $\sin(t)({\bm{x}}_{t}+\sigma_{d}\frac{\mathrm{d}{\bm{F}}_{\theta^{-}}}{\mathrm{
d}t})$ causes instability. To alleviate this problem, we gradually warm up the second term by replacing the coefficient $\sin(t)$ with $r\cdot\sin(t)$, where $r$ linearly increases from 0 0 to $1$ over the first 10k training iterations.

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.

Figure: Figure 6: sCD scales commensurately with teacher diffusion models. We plot the (a) FID and (b) FID ratio against the teacher diffusion model (at the same model size) on ImageNet 64$\times$64 and 512$\times$512. sCD scales better than sCT, and has a constant offset in the FID ratio across all model sizes, implying that sCD has the same scaling property of the teacher diffusion model. Furthermore, the offset diminishes with more sampling steps.
Refer to caption: x11.png

Figure: Figure 7: sCD has higher diversity compared to VSD: Sample quality comparison of the EDM2 (Karras et al., ) diffusion model, VSD (Wang et al., ; Yin et al., ), sCD, and the combination of VSD and sCD, across varying guidance scales. All models are of EDM2-M size and trained on ImageNet 512$\times$512.
Refer to caption: x15.png

**Table 1: Sample quality on unconditional CIFAR-10 and class-conditional ImageNet 64$\times$ 64.**
| Unconditional CIFAR-10 |  |  |
| --- | --- | --- |
| METHOD | NFE ($\downarrow$) | FID ($\downarrow$) |
| Diffusion models & Fast Samplers |  |  |
| Score SDE (deep) (Song et al., ) | 2000 | 2.20 |
| EDM (Karras et al., ) | 35 | 2.01 |
| Flow Matching (Lipman et al., ) | 142 | 6.35 |
| DPM-Solver (Lu et al., ) | 10 | 4.70 |
| DPM-Solver++ (Lu et al., ) | 10 | 2.91 |
| DPM-Solver-v3 (Zheng et al., ) | 10 | 2.51 |
| Joint Training |  |  |
| Diffusion GAN (Xiao et al., ) | 4 | 3.75 |
| Diffusion StyleGAN (Wang et al., ) | 1 | 3.19 |
| StyleGAN-XL (Sauer et al., ) | 1 | 1.52 |
| CTM (Kim et al., ) | 1 | 1.87 |
| Diff-Instruct (Luo et al., ) | 1 | 4.53 |
| DMD (Yin et al., ) | 1 | 3.77 |
| SiD (Zhou et al., ) | 1 | 1.92 |
| Diffusion Distillation |  |  |
| DFNO (LPIPS) (Zheng et al., ) | 1 | 3.78 |
| 2-Rectified Flow (Liu et al., ) | 1 | 4.85 |
| PID (LPIPS) (Tee et al., ) | 1 | 3.92 |
| BOOT (LPIPS) (Gu et al., ) | 1 | 4.38 |
| Consistency-FM (Yang et al., ) | 2 | 5.34 |
| PD (Salimans & Ho, ) | 1 | 8.34 |
|  | 2 | 5.58 |
| TRACT (Berthelot et al., ) | 1 | 3.78 |
|  | 2 | 3.32 |
| CD (LPIPS) (Song et al., ) | 1 | 3.55 |
|  | 2 | 2.93 |
| sCD (ours) | 1 | 3.66 |
|  | 2 | 2.52 |
| Consistency Training |  |  |
| iCT (Song & Dhariwal, ) | 1 | 2.83 |
|  | 2 | 2.46 |
| iCT-deep (Song & Dhariwal, ) | 1 | 2.51 |
|  | 2 | 2.24 |
| ECT (Geng et al., ) | 1 | 3.60 |
|  | 2 | 2.11 |
| sCT (ours) | 1 | 2.97 |
|  | 2 | 2.06 |

**Table 2: Sample quality on class-conditional ImageNet 512$\times$ 512. ^†Our reimplemented teacher diffusion model based on EDM2 (Karras et al., ) but with modifications in Sec. 4.1.**
| METHOD | NFE ($\downarrow$) | FID ($\downarrow$) | #Params |
| --- | --- | --- | --- |
| Diffusion models |  |  |  |
| ADM-G (Dhariwal & Nichol, ) | 250$\times$2 | 7.72 | 559M |
| RIN (Jabri et al., ) | 1000 | 3.95 | 320M |
| U-ViT-H/4 (Bao et al., ) | 250$\times$2 | 4.05 | 501M |
| DiT-XL/2 (Peebles & Xie, ) | 250$\times$2 | 3.04 | 675M |
| SimDiff (Hoogeboom et al., ) | 512$\times$2 | 3.02 | 2B |
| VDM++ (Kingma & Gao, ) | 512$\times$2 | 2.65 | 2B |
| DiffiT (Hatamizadeh et al., ) | 250$\times$2 | 2.67 | 561M |
| DiMR-XL/3R (Liu et al., ) | 250$\times$2 | 2.89 | 525M |
| DIFFUSSM-XL (Yan et al., ) | 250$\times$2 | 3.41 | 673M |
| DiM-H (Teng et al., ) | 250$\times$2 | 3.78 | 860M |
| U-DiT (Tian et al., ) | 250 | 15.39 | 204M |
| SiT-XL (Ma et al., ) | 250$\times$2 | 2.62 | 675M |
| Large-DiT (Alpha-VLLM, ) | 250$\times$2 | 2.52 | 3B |
| MaskDiT (Zheng et al., ) | 79$\times$2 | 2.50 | 736M |
| DiS-H/2 (Fei et al., ) | 250$\times$2 | 2.88 | 900M |
| DRWKV-H/2 (Fei et al., ) | 250$\times$2 | 2.95 | 779M |
| EDM2-S (Karras et al., ) | 63$\times$2 | 2.23 | 280M |
| EDM2-M (Karras et al., ) | 63$\times$2 | 2.01 | 498M |
| EDM2-L (Karras et al., ) | 63$\times$2 | 1.88 | 778M |
| EDM2-XL (Karras et al., ) | 63$\times$2 | 1.85 | 1.1B |
| EDM2-XXL (Karras et al., ) | 63$\times$2 | 1.81 | 1.5B |
| GANs & Masked Models |  |  |  |
| BigGAN (Brock, ) | 1 | 8.43 | 160M |
| StyleGAN-XL (Sauer et al., ) | 1$\times$2 | 2.41 | 168M |
| VQGAN (Esser et al., ) | 1024 | 26.52 | 227M |
| MaskGIT (Chang et al., ) | 12 | 7.32 | 227M |
| MAGVIT-v2 (Yu et al., ) | 64$\times$2 | 1.91 | 307M |
| MAR (Li et al., ) | 64$\times$2 | 1.73 | 481M |
| VAR-$d$36-s (Tian et al., ) | 10$\times$2 | 2.63 | 2.3B |

## 5 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.

### 5.1 Tangent Computation in Large-Scale Models

The common setting for training large-scale diffusion models includes using half-precision (FP16) and Flash Attention (Dao et al., ; Dao, ). As training continuous-time CMs requires computing the tangent $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$ accurately, we need to improve numerical precision and also support memory-efficient attention computation, as detailed below.

JVP Rearrangement.
Computing $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$ involves calculating $\frac{\mathrm{d}{\bm{F}}_{\theta^{-}}}{\mathrm{d}t}=\nabla_{{\bm{x}}_{t}}{\bm{
F}}_{\theta^{-}}\cdot\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}+\partial_{t}{
\bm{F}}_{\theta^{-}}$, which can be efficiently obtained via the Jacobian-vector product (JVP) for ${\bm{F}}_{\theta^{-}}(\frac{\cdot}{\sigma_{d}},\cdot)$ with the input vector $({\bm{x}}_{t},t)$ and the tangent vector $(\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t},1)$. However, we empirically find that the tangent may overflow in intermediate layers when $t$ is near 0 or $\frac{\pi}{2}$. To improve numerical precision, we propose to rearrange the computation of the tangent. Specifically, since the objective in Eq. (8) contains $\cos(t)\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$ and $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}$ is proportional to $\sin(t)\frac{\mathrm{d}{\bm{F}}_{\theta^{-}}}{\mathrm{d}t}$, we can compute the JVP as:

$$ $\cos(t)\sin(t)\frac{\mathrm{d}{\bm{F}}_{\theta^{-}}}{\mathrm{d}t}=\left(\nabla _{\frac{{\bm{x}}_{t}}{\sigma_{d}}}{\bm{F}}_{\theta^{-}}\right)\cdot\left(\cos( t)\sin(t)\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}\right)+\partial_{t}{\bm{F} }_{\theta^{-}}\cdot(\cos(t)\sin(t)\sigma_{d}),$ $$

which is the JVP for ${\bm{F}}_{\theta^{-}}(\cdot,\cdot)$ with the input $(\frac{{\bm{x}}_{t}}{\sigma_{d}},t)$ and the tangent $(\cos(t)\sin(t)\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t},\allowbreak\cos(t)
\sin(t)\sigma_{d})$. This rearrangement greatly alleviates the overflow issues in the intermediate layers, resulting in more stable training in FP16.

JVP of Flash Attention.
Flash Attention (Dao et al., ; Dao, ) 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.

### 5.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, ), ImageNet 64$\times$64 and ImageNet 512$\times$512 (Deng et al., ). We benchmark the sample quality using FID (Heusel et al., ). We follow the settings of Score SDE (Song et al., ) on CIFAR10 and EDM2 (Karras et al., ) 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. () 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, ), we incorporate an additional input $s$ into the model ${\bm{F}}_{\theta}$ to represent the guidance scale (Meng et al., ). We train the model with sCD by uniformly sampling $s\in[1,2]$ 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., ) and VAR (Tian et al., ). 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., ; Yin et al., ) and its multi-step generalization (Xie et al., ; Salimans et al., ) represent another diffusion distillation technique that has demonstrated scalability on high-resolution images (Yin et al., ).
We apply one-step VSD from time $T$ to 0 0 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.

## 6 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.

### Acknowledgements

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.

## 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.

Figure: Algorithm 1 Simplified and Stabilized Continuous-time Consistency Models (sCM).

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

### B.1 Derivations

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}$)(^1^11For 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

$$ ${\mathcal{L}}_{\text{Diff}}(\theta)=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[ w(t)\left\|{\bm{D}}_{\theta}({\bm{x}}_{t},t)-{\bm{x}}_{0}\right\|_{2}^{2} \right],$ (9) $$

where different diffusion model formulation contains four different parts:

- 1.
Parameterization of ${\bm{D}}_{\theta}$, such as score function (Song & Ermon, ; Song et al., ), noise prediction model (Song & Ermon, ; Song et al., ; Ho et al., ), data prediction model (Ho et al., ; Kingma et al., ; Salimans & Ho, ), velocity prediction model (Salimans & Ho, ), EDM (Karras et al., ) and flow matching (Lipman et al., ; Liu et al., ; Albergo et al., ).
- 2.
Noise schedule for $\alpha_{t}$ and $\sigma_{t}$, such as variance preserving process (Ho et al., ; Song et al., ), variance exploding process (Song et al., ; Karras et al., ), cosine schedule (Nichol & Dhariwal, ), and conditional optimal transport path (Lipman et al., ).
- 3.
Weighting function for $w(t)$, such as uniform weighting (Ho et al., ; Nichol & Dhariwal, ; Karras et al., ), weighting by functions of signal-to-noise-ratio (SNR) (Salimans & Ho, ), monotonic weighting (Kingma & Gao, ) and adaptive weighting (Karras et al., ).
- 4.
Proposal distribution for $t$, such as uniform distribution within $[0,T]$ (Ho et al., ; Song et al., ), log-normal distribution (Karras et al., ), SNR sampler (Esser et al., ), and adaptive importance sampler (Song et al., ; Kingma et al., ).

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

#### Step 1: General EDM parameterization.

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

$$ ${\bm{D}}_{\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)),$ (10) $$

and thus the training objective becomes

$$ ${\mathcal{L}}_{\text{Diff}}=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[w(t)c_{ \text{out}}^{2}(t)\left\|{\bm{F}}_{\theta}(c_{\text{in}}(t){\bm{x}}_{t},c_{ \text{noise}}(t))-\frac{c_{\text{skip}}(t)\alpha_{t}{\bm{x}}_{0}-(1-c_{\text{ skip}}(t)\sigma_{t}){\bm{z}}}{c_{\text{out}}(t)}\right\|_{2}^{2}\right].$ (11) $$

To ensure the input data of ${\bm{F}}_{\theta}$ has unit variance, we should ensure $\mathrm{Var}[c_{\text{in}}(t){\bm{x}}_{t}]=1$ by letting

$$ $c_{\text{in}}(t)=\frac{1}{\sigma_{d}\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}}.$ (12) $$

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

$$ $c_{\text{out}}^{2}(t)=\sigma_{d}^{2}c_{\text{skip}}^{2}(t)\alpha_{t}^{2}+(c_{ \text{skip}}(t)\sigma_{t}-1)^{2}.$ (13) $$

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

$$ $\begin{split}c_{\text{skip}}(t)=\frac{\sigma_{t}}{\alpha_{t}^{2}+\sigma_{t}^{2 }},\quad c_{\text{out}}(t)=\pm\frac{\sigma_{d}\alpha_{t}}{\sqrt{\alpha_{t}^{2} +\sigma_{t}^{2}}}.\end{split}$ (14) $$

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

$$ $\displaystyle{\bm{D}}_{\theta}({\bm{x}}_{t},t)$ $\displaystyle=\frac{\sigma_{t}}{\alpha_{t}^{2}+\sigma_{t}^{2}}{\bm{x}}_{t}- \frac{\alpha_{t}}{\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}}\sigma_{d}{\bm{F}}_{ \theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2} }},c_{\text{noise}}(t)\right),$ (15) $\displaystyle{\mathcal{L}}_{\text{Diff}}$ $\displaystyle=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[w(t)\frac{\alpha^{2}_{ t}}{\alpha_{t}^{2}+\sigma_{t}^{2}}\left\|\sigma_{d}{\bm{F}}_{\theta}\left( \frac{{\bm{x}}_{t}}{\sigma_{d}\sqrt{\alpha_{t}^{2}+\sigma_{t}^{2}}},c_{\text{ noise}}(t)\right)-\frac{\alpha_{t}{\bm{z}}-\sigma_{t}{\bm{x}}_{0}}{\sqrt{ \alpha_{t}^{2}+\sigma_{t}^{2}}}\right\|_{2}^{2}\right].$ (16) $$

#### 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

$$ $\displaystyle{\bm{D}}_{\theta}({\bm{x}}_{t},t)$ $\displaystyle=\hat{\sigma}_{t}\hat{{\bm{x}}}_{t}-\hat{\alpha}_{t}\sigma_{d}{ \bm{F}}_{\theta}\left(\frac{\hat{{\bm{x}}}_{t}}{\sigma_{d}},c_{\text{noise}}(t )\right),$ (17) $\displaystyle{\mathcal{L}}_{\text{Diff}}$ $\displaystyle=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[w(t)\hat{\alpha}_{t}^{ 2}\left\|\sigma_{d}{\bm{F}}_{\theta}\left(\frac{\hat{{\bm{x}}}_{t}}{\sigma_{d} },c_{\text{noise}}(t)\right)-(\hat{\alpha}_{t}{\bm{z}}-\hat{\sigma}_{t}{\bm{x} }_{0})\right\|_{2}^{2}\right].$ (18) $$

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

$$ ${\bm{x}}_{t}=\frac{\sigma_{t}}{\sigma_{s}}{\bm{x}}_{s}+\sigma_{t}\int_{\lambda _{s}}^{\lambda_{t}}e^{\lambda}{\bm{D}}_{\theta}({\bm{x}}_{\lambda},\lambda) \mathrm{d}\lambda,$ (19) $$

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

$$ $\hat{{\bm{x}}}_{t}=\frac{\hat{\sigma}_{t}}{\hat{\sigma}_{s}}\hat{{\bm{x}}}_{s} +\hat{\sigma}_{t}\int_{\hat{\lambda}_{s}}^{\hat{\lambda}_{t}}e^{\hat{\lambda}} {\bm{D}}_{\theta}(\hat{{\bm{x}}}_{\hat{\lambda}},\hat{\lambda})\mathrm{d}\hat{ \lambda},$ (20) $$

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., ) 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., ; ) during sampling, not the “straight path” (Lipman et al., ) 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

$$ $\hat{t}\coloneqq\arctan\left(\frac{\hat{\sigma}_{t}}{\hat{\alpha}_{t}}\right)= \arctan\left(\frac{\sigma_{t}}{\alpha_{t}}\right),$ (21) $$

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

$$ ${\mathcal{L}}_{\text{Diff}}=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}}}\left[\int_{0}^{ \frac{\pi}{2}}\underbrace{p\left(\hat{t}\right)w\left(\hat{t}\right)\cos^{2} \left(\hat{t}\right)}_{\text{training weighting}}\left\|\underbrace{\sigma_{d} {\bm{F}}_{\theta}\left(\frac{\hat{{\bm{x}}}_{\hat{t}}}{\sigma_{d}},c_{\text{ noise}}\left(\hat{t}\right)\right)-(\cos\left(\hat{t}\right){\bm{z}}-\sin\left (\hat{t}\right){\bm{x}}_{0})}_{\text{independent from $\alpha_{t}$ and $\sigma _{t}$}}\right\|_{2}^{2}\mathrm{d}\hat{t}\right].$ (22) $$

Therefore, we can always put the influence of different noise schedules into the training weighting of the integral for $\hat{t}$ from 0 0 to $\frac{\pi}{2}$, while the $\ell_{2}$ norm loss at each $\hat{t}$ is independent from the choices of $\alpha_{t}$ and $\sigma_{t}$. As we equivalently convert the noise schedules by trigonometric functions, we name such framework for diffusion models as TrigFlow.

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.

$$ ${\bm{D}}_{\theta}({\bm{x}}_{t},t)=\cos(t){\bm{x}}_{t}-\sin(t)\sigma_{d}{\bm{F} }_{\theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},c_{\text{noise}}(t)\right),$ (23) $$

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

$$ $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}=\sigma_{d}{\bm{F}}_{\theta}\left( \frac{{\bm{x}}_{t}}{\sigma_{d}},c_{\text{noise}}(t)\right).$ (24) $$

#### Training Objective.

$$ ${\mathcal{L}}_{\text{Diff}}(\theta)=\mathbb{E}_{{\bm{x}}_{0},{\bm{z}}}\left[ \int_{0}^{\frac{\pi}{2}}w(t)\left\|\sigma_{d}{\bm{F}}_{\theta}\left(\frac{{\bm {x}}_{t}}{\sigma_{d}},c_{\text{noise}}\left(t\right)\right)-(\cos(t){\bm{z}}- \sin(t){\bm{x}}_{0})\right\|_{2}^{2}\mathrm{d}t\right],$ (25) $$

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., ), 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., ), 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.

#### 1st-order Sampler by DDIM.

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

$$ ${\bm{x}}_{t}=\cos(s-t){\bm{x}}_{t}-\sin(s-t)\sigma_{d}{\bm{F}}_{\theta}\left( \frac{{\bm{x}}_{s}}{\sigma_{d}},c_{\text{noise}}(s)\right)$ (26) $$

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

#### 2nd-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

$$ ${\bm{x}}_{t}=\cos(s-t){\bm{x}}_{s}-\sin(s-t)\sigma_{d}{\bm{F}}_{\theta}\left( \frac{{\bm{x}}_{s}}{\sigma_{d}},c_{\text{noise}}(s)\right)-\frac{\sin(s-t)}{2r _{s}\cos(s)}\left(\bm{\epsilon}_{\theta}({\bm{x}}_{s^{\prime}},s^{\prime})-\bm {\epsilon}_{\theta}({\bm{x}}_{s},s)\right),$ (27) $$

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)}$.

#### 2nd-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

$$ ${\bm{x}}_{t}=\cos(s-t){\bm{x}}_{s}-\sin(s-t)\sigma_{d}{\bm{F}}_{\theta}\left( \frac{{\bm{x}}_{s}}{\sigma_{d}},c_{\text{noise}}(s)\right)+\frac{\sin(s-t)}{2r _{s}\sin(s)}\left({\bm{D}}_{\theta}({\bm{x}}_{s^{\prime}},s^{\prime})-{\bm{D}} ({\bm{x}}_{s},s)\right),$ (28) $$

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

$$ $t=\arctan\left(\frac{\sigma_{t^{\prime}}}{\sigma_{d}\alpha_{t^{\prime}}}\right ),\quad{\bm{x}}_{t}=\frac{\sigma_{d}}{\sqrt{\alpha_{t^{\prime}}^{2}\sigma_{d}^ {2}+\sigma_{t^{\prime}}^{2}}}{\bm{x}}_{t^{\prime}}.$ (29) $$

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.

#### EDM.

As our derivations closely follow the unit variance principle proposed in EDM (Karras et al., ), 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

$$ $t=\arctan\left(\frac{\sigma}{\sigma_{d}}\right),\quad{\bm{x}}_{t}=\cos(t){\bm{ x}}_{\sigma}.$ (30) $$

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., ; Liu et al., ; Albergo et al., ) 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

$$ ${\bm{v}}_{t}=\frac{\mathrm{d}\alpha_{t}}{\mathrm{d}t}{\bm{x}}_{0}+\frac{ \mathrm{d}\sigma_{t}}{\mathrm{d}t}{\bm{z}},$ (31) $$

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

$$ $\mathbb{E}_{{\bm{x}}_{0},{\bm{z}},t}\left[w(t)\left\|{\bm{v}}_{\theta}({\bm{x} }_{t},t)-{\bm{v}}_{t}\right\|_{2}^{2}\right],$ (32) $$

and the final probability flow ODE is defined by

$$ $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}={\bm{v}}_{\theta}({\bm{x}}_{t},t).$ (33) $$

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, ) 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

$$ $\hat{\text{SNR}}(t)=\frac{\alpha_{t}^{2}\sigma_{d}^{2}}{\sigma_{t}^{2}}=\frac{ 1}{\tan^{2}(t)},$ (34) $$

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.

### C.1 Derivations

Assume we have samples ${\bm{x}}_{0}\in\mathbb{R}^{D}$ from data distribution $p_{d}$ with standard deviation $\sigma_{d}$, and define a corresponding forward diffusion process $\{p_{t}\}_{t=0}^{T}$ starting at $p_{0}=p_{d}$ and ending at $p_{T}\approx{\mathcal{N}}({\bm{0}},\hat{\sigma}^{2}{\bm{I}})$, with $p_{t0}({\bm{x}}_{t}|{\bm{x}}_{0})={\mathcal{N}}({\bm{x}}_{t}|\alpha_{t}{\bm{x}
}_{0},\sigma_{t}^{2}{\bm{I}})$. Variational score distillation (VSD) (Wang et al., ; Yin et al., ; ) trains a generator ${\bm{g}}_{\theta}:\mathbb{R}^{D}\rightarrow\mathbb{R}^{D}$ aiming to map noise samples ${\bm{z}}\sim{\mathcal{N}}({\bm{0}},\hat{\sigma}^{2}{\bm{I}})$ to the data distribution, by minimizing

$$ $\min_{\theta}\mathbb{E}_{t}\left[w(t)D_{\text{KL}}\left(q_{t}^{\theta} \parallel p_{t}\right)\right]=\mathbb{E}_{t,{\bm{z}},\bm{\epsilon}}\left[w(t) \left(\log q_{t}^{\theta}(\alpha_{t}{\bm{g}}_{\theta}({\bm{z}})+\sigma_{t}\bm{ \epsilon})-\log p_{t}(\alpha_{t}{\bm{g}}_{\theta}({\bm{z}})+\sigma_{t}\bm{ \epsilon})\right)\right],$ $$

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., ; Wang et al., ; Yin et al., ), 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., ) 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

$$ $\displaystyle\phantom{{}={}}\nabla_{\theta}\mathbb{E}_{t}\left[w(t)D_{\text{KL }}\left(q_{t}^{\theta}\parallel p_{t}\right)\right]$ $\displaystyle=\mathbb{E}_{t,{\bm{z}},\bm{\epsilon}}\left[w(t)\nabla_{\theta} \left(\log q_{t}^{\theta}({\bm{x}}_{t}^{\theta})-\log p_{t}({\bm{x}}_{t}^{ \theta})\right)\right]$ $\displaystyle=\mathbb{E}_{t,{\bm{z}},\bm{\epsilon}}\left[w(t)\left(\partial_{ \theta}\log q_{t}^{\theta}({\bm{x}}_{t})+\left(\nabla_{{\bm{x}}_{t}}\log q_{t} ^{\theta}({\bm{x}}_{t})-\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})\right) \frac{\partial{\bm{x}}_{t}^{\theta}}{\partial\theta}\right)\right]$ $\displaystyle=\underbrace{\mathbb{E}_{t,{\bm{x}}_{t}}\left[w(t)\partial_{ \theta}\log q_{t}^{\theta}({\bm{x}}_{t})\right]}_{=0}+\mathbb{E}_{t,{\bm{z}}, \bm{\epsilon}}\left[w(t)\left(\nabla_{{\bm{x}}_{t}}\log q_{t}^{\theta}({\bm{x} }_{t})-\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})\right)\frac{\alpha_{t} \partial{\bm{g}}_{\theta}({\bm{z}})}{\partial\theta}\right]$ $\displaystyle=\mathbb{E}_{t,{\bm{z}},\bm{\epsilon}}\left[\alpha_{t}w(t)\left( \nabla_{{\bm{x}}_{t}}\log q_{t}^{\theta}({\bm{x}}_{t})-\nabla_{{\bm{x}}_{t}} \log p_{t}({\bm{x}}_{t})\right)\frac{\partial{\bm{g}}_{\theta}({\bm{z}})}{ \partial\theta}\right].$ $$

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

$$ $\min_{\phi}\mathbb{E}_{{\bm{z}},{\bm{z}}^{\prime},t}\left[w(t)\left\|\sigma_{d }{\bm{F}}_{\phi}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)-{\bm{v}}_{t} \right\|_{2}^{2}\right],$ $$

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

$$ $\displaystyle\sigma_{d}{\bm{F}}_{\text{Diff}}({\bm{x}}_{t},t)$ $\displaystyle=\mathbb{E}[{\bm{v}}_{t}|{\bm{x}}_{t}]=\frac{1}{\tan(t)}{\bm{x}}_ {t}-\frac{1}{\sin(t)}\mathbb{E}_{{\bm{x}}_{0}|{\bm{x}}_{t}}\left[{\bm{x}}_{0} \right],$ $\displaystyle\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ $\displaystyle=\mathbb{E}_{{\bm{x}}_{0}|{\bm{x}}_{t}}\left[-\frac{{\bm{x}}_{t}- \cos(t){\bm{x}}_{0}}{\sigma_{d}^{2}\sin^{2}(t)}\right]=-\frac{\cos(t)\sigma_{d }{\bm{F}}_{\text{Diff}}+\sin(t){\bm{x}}_{t}}{\sigma_{d}^{2}\sin(t)}.$ $$

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

$$ $\mathbb{E}_{t,{\bm{z}},{\bm{z}}^{\prime}}\left[\frac{\cos^{2}(t)}{\sigma_{d} \sin(t)}w(t)\left({\bm{F}}_{\text{pretrain}}\left(\frac{{\bm{x}}_{t}}{\sigma_{ d}},t\right)-{\bm{F}}_{\phi}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right) \right)\frac{\partial{\bm{g}}_{\theta}({\bm{z}}^{\prime})}{\partial\theta} \right],$ $$

which is equivalent to the gradient of the following objective:

$$ $\mathbb{E}_{t,{\bm{z}},{\bm{z}}^{\prime}}\left[\frac{\cos^{2}(t)}{\sigma_{d} \sin(t)}w(t)\left\|{\bm{g}}_{\theta}({\bm{z}}^{\prime})-{\bm{g}}_{\theta^{-}}( {\bm{z}}^{\prime})+{\bm{F}}_{\text{pretrain}}\left(\frac{{\bm{x}}_{t}}{\sigma_ {d}},t\right)-{\bm{F}}_{\phi}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right) \right\|_{2}^{2}\right],$ $$

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., ; Yin et al., ) 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. (). 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).

$$ $\displaystyle\min_{\phi}\mathcal{L}_{\text{Diff}}(\phi)$ $\displaystyle\coloneqq\mathbb{E}_{{\bm{z}},{\bm{z}}^{\prime},t}\left[w(t)\left \|\sigma_{d}{\bm{F}}_{\phi}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)-{\bm {v}}_{t}\right\|_{2}^{2}\right],$ $\displaystyle\min_{\theta,\psi}\mathcal{L}_{\text{VSD}}(\theta,\psi)$ $\displaystyle\coloneqq\mathbb{E}_{t,{\bm{z}},{\bm{z}}^{\prime}}\left[\!\frac{e ^{w_{\psi}(t)}}{D}\left\|{\bm{g}}_{\theta}({\bm{z}}^{\prime})-{\bm{g}}_{\theta ^{-}}({\bm{z}}^{\prime})+{\bm{F}}_{\text{pretrain}}\left(\frac{{\bm{x}}_{t}}{ \sigma_{d}},t\right)\!-\!{\bm{F}}_{\phi}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}}, t\right)\right\|_{2}^{2}\!\!-\!w_{\psi}(t)\!\right].$ $$

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

$$ $\min_{\theta,\psi}\mathbb{E}_{t,{\bm{z}},{\bm{z}}^{\prime}}\!\left[\frac{e^{w_ {\psi}(t)}}{D}\!\left\|\sigma_{d}{\bm{F}}_{\theta}\!\left(\frac{{\bm{z}}^{ \prime}}{\sigma_{d}},\frac{\pi}{2}\!\right)\!-\!\sigma_{d}{\bm{F}}_{\theta^{-} }\!\left(\frac{{\bm{z}}^{\prime}}{\sigma_{d}},\frac{\pi}{2}\!\right)\!-\!{\bm{ F}}_{\text{pretrain}}\!\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\!\right)\!+\!{ \bm{F}}_{\phi}\!\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\!\right)\right\|_{2}^{ 2}-w_{\psi}(t)\right].$ $$

## 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

$$ $\nabla_{\theta}\mathcal{L}_{\text{Diff}}(\theta)=\nabla_{\theta}\mathbb{E}_{t, {\bm{x}}_{0},{\bm{z}}}w(t)\left[\|\sigma_{d}{\bm{F}}_{\theta}-{\bm{v}}_{t}\|_{ 2}^{2}\right]=\nabla_{\theta}\mathbb{E}_{t,{\bm{x}}_{0},{\bm{z}}}\left[w(t) \sigma_{d}{\bm{F}}_{\theta}^{\top}(\sigma_{d}{\bm{F}}_{\theta^{-}}-{\bm{v}}_{t })\right],$ $$

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

$$ $\nabla_{\theta}\mathcal{L}_{\text{Diff}}(\theta)=\nabla_{\theta}\mathbb{E}_{t, {\bm{z}},{\bm{z}}^{\prime}}\left[w(t)g_{\theta}({\bm{z}}^{\prime})^{\top}({\bm {F}}_{\text{pretrain}}-{\bm{F}}_{\phi})\right].$ $$

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

$$ $\mathcal{L}_{\text{CM}}(\theta)=\nabla_{\theta}\mathbb{E}_{t,{\bm{x}}_{0},{\bm {z}}}\left[-w(t)\sin(t){\bm{f}}_{\theta}^{\top}\frac{\mathrm{d}{\bm{f}}_{ \theta^{-}}}{\mathrm{d}t}\right],$ $$

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:

$$ $\min_{\theta}\mathbb{E}_{t}\left[{\bm{F}}_{\theta}^{\top}{\bm{y}}\right],$ $$

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

$$ $\nabla_{\theta}\mathbb{E}_{t}\left[\left\|{\bm{F}}_{\theta}-{\bm{F}}_{\theta^{ -}}+{\bm{y}}\right\|_{2}^{2}\right],$ $$

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

$$ $\min_{\phi}\mathbb{E}_{t}\left[\frac{e^{w_{\phi}(t)}}{D}\|{\bm{F}}_{\theta}-{ \bm{F}}_{\theta^{-}}+{\bm{y}}\|_{2}^{2}-w_{\phi}(t)\right].$ $$

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

$$ $\frac{e^{w^{*}(t)}}{D}\mathbb{E}\left[\|{\bm{F}}_{\theta}-{\bm{F}}_{\theta^{-} }+{\bm{y}}\|_{2}^{2}\right]\equiv 1.$ $$

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

- 1.
A prior weighting $\lambda(t)$ for ${\bm{y}}$, which may be helpful for further reducing the variance of ${\bm{y}}$. Then the objective becomes
$\min_{\theta,\phi}\mathbb{E}_{t}\left[\frac{e^{w_{\phi}(t)}}{D}\|{\bm{F}}_{
\theta}-{\bm{F}}_{\theta^{-}}+\lambda(t){\bm{y}}\|_{2}^{2}-w_{\phi}(t)\right].$
e.g., for diffusion models and VSD, since the target is either ${\bm{y}}={\bm{F}}-{\bm{v}}_{t}$ or ${\bm{y}}={\bm{F}}_{\text{pretrain}}-{\bm{F}}_{\phi}$ which are stable across different time steps, we can simply choose $\lambda(t)=1$; while for consistency models, the target ${\bm{y}}=\sin(t)\frac{\mathrm{d}{\bm{f}}}{\mathrm{d}t}$ may have huge variance, we choose $\lambda(t)=\frac{1}{\sigma_{d}\tan(t)}$ to reduce the variance of $\lambda(t){\bm{y}}$, which empirically is critical for better performance.
- 2.
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. () 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.

### E.1 Parameterization and Training Objective

Parameterization.
We also parameterize the CM by TrigFlow:

$$ ${\bm{f}}_{\theta}({\bm{x}}_{t},t)=\cos(t){\bm{x}}_{t}-\sigma_{d}\sin(t){\bm{F} }_{\theta}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right).$ $$

And we denote the pretrained teacher diffusion model as $\frac{\mathrm{d}{\bm{x}}_{t}}{\mathrm{d}t}={\bm{F}}_{\text{pretrain}}(\frac{{
\bm{x}}_{t}}{\sigma_{d}},t)$.

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}$:

$$ ${\bm{x}}_{t^{\prime}}=\cos(t-t^{\prime}){\bm{x}}_{t}-\sigma_{d}\sin(t-t^{ \prime}){\bm{F}}_{\text{pretrain}}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t \right).$ $$

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

$$ ${\bm{f}}_{\theta^{-}}({\bm{x}}_{t^{\prime}},t^{\prime})=\cos(t^{\prime})\cos(t -t^{\prime}){\bm{x}}_{t}-\sigma_{d}\cos(t^{\prime})\sin(t-t^{\prime}){\bm{F}}_ {\text{pretrain}}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)-\sigma_{d}\sin (t^{\prime}){\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t^{\prime}}}{\sigma_{d} },t^{\prime}\right).$ (35) $$

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

$$ $\displaystyle{\bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)$ $\displaystyle=\cos(t){\bm{x}}_{t}-\sigma_{d}\sin(t){\bm{F}}_{\theta^{-}}\left( \frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)$ (36) $\displaystyle=(\cos(t^{\prime})\cos(t-t^{\prime})-\sin(t^{\prime})\sin(t-t^{ \prime})){\bm{x}}_{t}$ $\displaystyle\quad-\sigma_{d}(\sin(t-t^{\prime})\cos(t^{\prime})+\cos(t-t^{ \prime})\sin(t^{\prime})){\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{\sigma _{d}},t\right)$ $$

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

$$ $\begin{split}\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})&\coloneqq\frac{{ \bm{f}}_{\theta^{-}}({\bm{x}}_{t},t)-{\bm{f}}_{\theta^{-}}({\bm{x}}_{t^{\prime }},t^{\prime})}{\sin(t-t^{\prime})}\\ &=-\cos(t^{\prime})\Big{(}\sigma_{d}{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_ {t}}{\sigma_{d}},t\right)-\underbrace{\sigma_{d}{\bm{F}}_{\text{pretrain}} \left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)}_{\frac{\mathrm{d}{\bm{x}}_{t}} {\mathrm{d}t}}\Big{)}\\ &\quad-\sin(t^{\prime})\Big{(}{\bm{x}}_{t}+\underbrace{\frac{\sigma_{d}\cos(t- t^{\prime}){\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right) -\sigma_{d}{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t^{\prime}}}{\sigma_{d}} ,t^{\prime}\right)}{\sin(t-t^{\prime})}}_{\approx\sigma_{d}\frac{\mathrm{d}{ \bm{F}}_{\theta^{-}}}{\mathrm{d}t}}\Big{)}.\end{split}$ (37) $$

Comparing to Eq. (6), it is easy to see that $\lim_{t^{\prime}\to t}\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})=\frac{
\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}({\bm{x}}_{t},t)$. Moreover, when using $d({\bm{x}},{\bm{y}})=\|{\bm{x}}-{\bm{y}}\|_{2}^{2}$, and $\Delta t=t-t^{\prime}$, the training objective of discrete-time CMs in Eq. (1) becomes

$$ $\mathbb{E}_{{\bm{x}}_{t},t}\left[w(t)\|{\bm{f}}_{\theta}({\bm{x}}_{t},t)-{\bm{ f}}_{\theta^{-}}({\bm{x}}_{t-\Delta t},t-\Delta t)\|_{2}^{2}\right],$ $$

which has the gradient of

$$ $\displaystyle\phantom{{}={}}\mathbb{E}_{{\bm{x}}_{t},t}\left[w(t)\nabla_{ \theta}{\bm{f}}^{\top}_{\theta}({\bm{x}}_{t},t)\left({\bm{f}}_{\theta^{-}}({ \bm{x}}_{t},t)-{\bm{f}}_{\theta^{-}}({\bm{x}}_{t-\Delta t},t-\Delta t)\right)\right]$ (38) $\displaystyle=\nabla_{\theta}\mathbb{E}_{{\bm{x}}_{t},t}\left[-w(t)\sin(t-t^{ \prime}){\bm{f}}^{\top}_{\theta}({\bm{x}}_{t},t)\Delta_{\theta^{-}}({\bm{x}}_{ t},t,t^{\prime})\right]$ $\displaystyle=\nabla_{\theta}\mathbb{E}_{{\bm{x}}_{t},t}\left[w(t)\sin(t-t^{ \prime})\sin(t){\bm{F}}^{\top}_{\theta}({\bm{x}}_{t},t)\Delta_{\theta^{-}}({ \bm{x}}_{t},t,t^{\prime})\right]$ $$

Adaptive weighting for discrete-time CMs.
Inspired by the continuous-time consistency models, we can also apply the adaptive weighting technique into discrete-time training objectives in Eq. (38). Specifically, since $\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})$ is a first-order approximation of $\frac{\mathrm{d}{\bm{f}}_{\theta^{-}}}{\mathrm{d}t}({\bm{x}}_{t},t)$, we can directly replace the tangent in Eq. (8) with $\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})$, and obtain the improved objective of discrete-time CMs by:

$$ $\mathcal{L}_{\text{sCM}}(\theta,\phi)\!\coloneqq\!\mathbb{E}_{{\bm{x}}_{t},t} \left[\!\frac{e^{w_{\phi}(t)}}{D}\left\|{\bm{F}}_{\theta}\left(\frac{{\bm{x}}_ {t}}{\sigma_{d}},t\right)-{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{ \sigma_{d}},t\right)-\cos(t)\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime}) \right\|_{2}^{2}-w_{\phi}(t)\right],$ (39) $$

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

$$ ${\bm{g}}_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})\coloneqq\frac{\cos(t)\Delta_{ \theta^{-}}({\bm{x}}_{t},t,t^{\prime})}{\|\cos(t)\Delta_{\theta^{-}}({\bm{x}}_ {t},t,t^{\prime})\|+c},$ $$

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

$$ $\mathcal{L}_{\text{sCM}}(\theta,\phi)\!\coloneqq\!\mathbb{E}_{{\bm{x}}_{t},t} \left[\!\frac{e^{w_{\phi}(t)}}{D}\left\|{\bm{F}}_{\theta}\left(\frac{{\bm{x}}_ {t}}{\sigma_{d}},t\right)-{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{ \sigma_{d}},t\right)-{\bm{g}}_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime})\right\|_ {2}^{2}-w_{\phi}(t)\right],$ $$

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

$$ $\begin{split}\Delta_{\theta^{-}}({\bm{x}}_{t},t,t^{\prime},r)&=-\cos(t^{\prime })\left(\sigma_{d}{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t \right)-\sigma_{d}{\bm{F}}_{\text{pretrain}}\left(\frac{{\bm{x}}_{t}}{\sigma_{ d}},t\right)\right)\\ &\quad-r\cdot\sin(t^{\prime})\left({\bm{x}}_{t}+\frac{\sigma_{d}\cos(t-t^{ \prime}){\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t}}{\sigma_{d}},t\right)- \sigma_{d}{\bm{F}}_{\theta^{-}}\left(\frac{{\bm{x}}_{t^{\prime}}}{\sigma_{d}}, t^{\prime}\right)}{\sin(t-t^{\prime})}\right),\end{split}$ $$

where $r$ linearly increases from 0 0 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.

Figure: Algorithm 2 Simplified and Stabilized Discrete-time Consistency Distillation (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 ().

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

The attention operator (Vaswani, ) needs to compute ${\bm{y}}=\text{softmax}({\bm{x}}){\bm{V}}$ where ${\bm{x}}\in\mathbb{R}^{1\times L},{\bm{V}}\in\mathbb{R}^{L\times D},{\bm{y}}
\in\mathbb{R}^{1\times D}$. Flash Attention (Dao et al., ; Dao, ) computes the output by maintaining three variables $m({\bm{x}})\in\mathbb{R},\ell({\bm{x}})\in\mathbb{R}$, and ${\bm{f}}({\bm{x}})$ with the same dimension as ${\bm{x}}$. The computation is done recursively: for each block, we have

$$ $m({\bm{x}})=\max(e^{{\bm{x}}}),\quad\ell({\bm{x}})=\sum_{i}e^{{\bm{x}}^{(i)}-m ({\bm{x}})},\quad{\bm{f}}({\bm{x}})=e^{{\bm{x}}-m({\bm{x}})}{\bm{V}},$ $$

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

$$ $\displaystyle m({\bm{x}})=\max\left({\bm{x}}^{(a)},{\bm{x}}^{(b)}\right),\quad \ell({\bm{x}})=e^{m({\bm{x}}^{(a)})-m({\bm{x}})}\ell({\bm{x}}^{(a)})+e^{m({\bm {x}}^{(b)})-m({\bm{x}})}\ell({\bm{x}}^{(b)}),$ $\displaystyle{\bm{f}}({\bm{x}})=\left[e^{m({\bm{x}}^{(a)})-m({\bm{x}})}{\bm{f} }({\bm{x}}^{(a)}),e^{m({\bm{x}}^{(b)})-m({\bm{x}})}{\bm{f}}({\bm{x}}^{(b)}) \right],\quad{\bm{y}}=\frac{{\bm{f}}({\bm{x}})}{\ell({\bm{x}})}.$ $$

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.

Denote ${\bm{p}}\coloneqq\text{softmax}({\bm{x}})$. Denote the tangent vector for ${\bm{x}}\in\mathbb{R}^{1\times L},{\bm{p}}\in\mathbb{R}^{1\times L},{\bm{V}}
\in\mathbb{R}^{L\times D},{\bm{y}}\in\mathbb{R}^{1\times D}$ as ${\bm{t}}_{{\bm{x}}}\in\mathbb{R}^{1\times L},{\bm{t}}_{{\bm{p}}}\in\mathbb{R}^
{1\times L},{\bm{t}}_{{\bm{V}}}\in\mathbb{R}^{L\times D},{\bm{t}}_{{\bm{y}}}
\in\mathbb{R}^{1\times D}$, correspondingly. The JVP for attention is computing $({\bm{x}},{\bm{t}}_{{\bm{x}}}),({\bm{V}},{\bm{t}}_{{\bm{V}}})\rightarrow({\bm{
y}},{\bm{t}}_{{\bm{y}}})$, which is

$$ ${\bm{t}}_{{\bm{y}}}={\bm{t}}_{{\bm{p}}}{\bm{V}}+\underbrace{{\bm{p}}{\bm{t}}_{ {\bm{V}}}}_{\text{softmax}({\bm{x}}){\bm{t}}_{{\bm{V}}}},\quad\text{where} \quad{\bm{t}}_{{\bm{p}}}{\bm{V}}=\underbrace{({\bm{p}}\odot{\bm{t}}_{{\bm{x}}} )}_{1\times L}{\bm{V}}-\underbrace{({\bm{p}}{\bm{t}}^{\top}_{{\bm{x}}})}_{1 \times 1}\cdot\underbrace{({\bm{p}}{\bm{V}})}_{{\bm{y}}}.$ $$

Notably, the computation for both ${\bm{p}}{\bm{t}}_{{\bm{V}}}$ and ${\bm{p}}{\bm{V}}$ can be done by the standard Flash Attention with the value matrices ${\bm{V}}$ and ${\bm{t}}_{{\bm{V}}}$. Thus, to compute ${\bm{t}}_{{\bm{y}}}$, we only need to maintain a vector ${\bm{g}}({\bm{x}})\coloneqq({\bm{p}}\odot{\bm{t}}_{{\bm{x}}}){\bm{V}}$ and a scalar $\mu({\bm{x}})\coloneqq{\bm{p}}{\bm{t}}_{{\bm{x}}}^{\top}$ during the Flash Attention computation loop. Moreover, since we do not know ${\bm{p}}$ during the loop, we can reuse the intermediate $m,\ell,{\bm{f}}$ in Flash Attention. Specifically, for each block,

$$ ${\bm{g}}({\bm{x}})=\left(e^{{\bm{x}}-m({\bm{x}})}\odot{\bm{t}}_{{\bm{x}}} \right){\bm{V}},\quad\mu({\bm{x}})=\sum_{i}e^{{\bm{x}}^{(i)}-m({\bm{x}})}{\bm{ t}}_{{\bm{x}}}^{(i)},$ $$

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

$$ $\displaystyle{\bm{g}}({\bm{x}})$ $\displaystyle=\left[e^{m({\bm{x}}^{(a)})-m({\bm{x}})}{\bm{g}}({\bm{x}}^{(a)}), e^{m({\bm{x}}^{(b)})-m({\bm{x}})}{\bm{g}}({\bm{x}}^{(b)})\right],$ $\displaystyle\mu({\bm{x}})$ $\displaystyle=e^{m({\bm{x}}^{(a)})-m({\bm{x}})}\mu({\bm{x}}^{(a)})+e^{m({\bm{x }}^{(b)})-m({\bm{x}})}\mu({\bm{x}}^{(b)}),$ $$

and after obtaining $m,\ell,{\bm{f}},{\bm{g}},\mu$ for the row vector ${\bm{x}}$, the final result of ${\bm{t}}_{{\bm{p}}}{\bm{V}}$ is

$$ ${\bm{t}}_{{\bm{p}}}{\bm{V}}=\frac{{\bm{g}}({\bm{x}})}{\ell({\bm{x}})}-\frac{ \mu({\bm{x}})}{\ell({\bm{x}})}\cdot{\bm{y}}.$ $$

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

### G.1 TrigFlow for Diffusion Models

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., ) architecture (DDPM++). We use the same settings of EDM (Karras et al., ): 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., ) (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 () by

- 1.
Resize the shorter width / height to $64\times 64$ resolution with bicubic interpolation.
- 2.
Center crop the image.
- 3.
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., ) 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 () and Karras et al. () by

- 1.
Resize the shorter width / height to $512\times 512$ resolution with bicubic interpolation.
- 2.
Center crop the image.
- 3.
Disable data augmentation such as horizontal flipping.
- 4.
Encode the images into latents by stable diffusion VAE(^2^22[https://huggingface.co/stabilityai/sd-vae-ft-mse](https://huggingface.co/stabilityai/sd-vae-ft-mse)) (Rombach et al., ; Janner et al., ), 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., ), 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., ). It is because when training the VAE, the images are normalized to $[-1,1]$ before passing to the encoder. However, the channel mean and std used in EDM2 assumes the input images are in $[0,1]$ 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., ) 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., ) (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 (). 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., ) and EDM2 (Karras et al., ). For 2-step sampling, we use the algorithm in Song et al. () 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 0 0 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., ) 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., ) which naturally supports JVP by PyTorch (Paszke et al., ) auto-grad. We use EMA half-life of $0.5$ Mimg (Karras et al., ). 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., ) and iCT (Song & Dhariwal, ). 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., ) and iCT (Song & Dhariwal, ). 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[1,2]$ 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.

**Table 3: Training settings of all models and training algorithms on ImageNet 64$\times$64 dataset.**
|  | Model Size |  |  |  |
| --- | --- | --- | --- | --- |
|  | S | M | L | XL |
| Model details |  |  |  |  |
| Batch size | 2048 | 2048 | 2048 | 2048 |
| Channel multiplier | 192 | 256 | 320 | 384 |
| Time embedding layer | positional | positional | positional | positional |
| noise conditioning $c_{\text{noise}}(t)$ | $t$ | $t$ | $t$ | $t$ |
| adaptive double normalization | ✓ | ✓ | ✓ | ✓ |
| Learning rate decay ($t_{\text{ref}}$) | 35000 | 35000 | 35000 | 35000 |
| Adam $\beta_{1}$ | 0.9 | 0.9 | 0.9 | 0.9 |
| Adam $\beta_{2}$ | 0.99 | 0.99 | 0.99 | 0.99 |
| Adam $\epsilon$ | 1.0e-11 | 1.0e-11 | 1.0e-11 | 1.0e-11 |
| Model capacity (Mparams) | 280.2 | 497.8 | 777.6 | 1119.4 |
| Training details of diffusion models (TrigFlow) |  |  |  |  |
| Training iterations | 1048k | 1486k | 761k | 540k |
| Learning rate max ($\alpha_{\text{ref}}$) | 1.0e-2 | 9.0e-3 | 8.0e-3 | 7.0e-3 |
| Dropout probability | 0% | 10% | 10% | 10% |
| Proposal $P_{\text{mean}}$ | -0.8 | -0.8 | -0.8 | -0.8 |
| Proposal $P_{\text{std}}$. | 1.6 | 1.6 | 1.6 | 1.6 |
| Shared details of consistency models |  |  |  |  |
| Learning rate max ($\alpha_{\text{ref}}$) | 1.0e-4 | 9.0e-5 | 8.0e-5 | 7.0e-5 |
| Proposal $P_{\text{mean}}$ | -1.0 | -1.0 | -1.0 | -1.0 |
| Proposal $P_{\text{std}}$. | 1.6 | 1.6 | 1.6 | 1.6 |
| Tangent normalization constant ($c$) | 0.1 | 0.1 | 0.1 | 0.1 |
| Tangent warm up iterations | 10k | 10k | 10k | 10k |
| EMA length ($\sigma_{\text{rel}}$) of pretrained diffusion | 0.075 | 0.06 | 0.04 | 0.04 |
| Training details of sCT |  |  |  |  |
| Training iterations | 400k | 400k | 400k | 400k |
| Dropout probability for resolution $\leq 16$ | 45% | 45% | 45% | 45% |
| Training details of sCD |  |  |  |  |
| Training iterations | 400k | 400k | 400k | 400k |
| Dropout probability for resolution $\leq 16$ | 0% | 0% | 0% | 0% |

**Table 4: Training settings of all models and training algorithms on ImageNet 512$\times$512 dataset.**
|  | Model Size |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
|  | S | M | L | XL | XXL |
| Model details |  |  |  |  |  |
| Batch size | 2048 | 2048 | 2048 | 2048 | 2048 |
| Channel multiplier | 192 | 256 | 320 | 384 | 448 |
| Time embedding layer | positional | positional | positional | positional | positional |
| noise conditioning $c_{\text{noise}}(t)$ | $t$ | $t$ | $t$ | $t$ | $t$ |
| adaptive double normalization | ✓ | ✓ | ✓ | ✓ | ✓ |
| Learning rate decay ($t_{\text{ref}}$) | 70000 | 70000 | 70000 | 70000 | 70000 |
| Adam $\beta_{1}$ | 0.9 | 0.9 | 0.9 | 0.9 | 0.9 |
| Adam $\beta_{2}$ | 0.99 | 0.99 | 0.99 | 0.99 | 0.99 |
| Adam $\epsilon$ | 1.0e-11 | 1.0e-11 | 1.0e-11 | 1.0e-11 | 1.0e-11 |
| Model capacity (Mparams) | 280.2 | 497.8 | 777.6 | 1119.4 | 1523.4 |
| Training details of diffusion models (TrigFlow) |  |  |  |  |  |
| Training iterations | 1048k | 1048k | 696k | 598k | 376k |
| Learning rate max ($\alpha_{\text{ref}}$) | 1.0e-2 | 9.0e-3 | 8.0e-3 | 7.0e-3 | 6.5e-3 |
| Dropout probability | 0% | 10% | 10% | 10% | 10% |
| Proposal $P_{\text{mean}}$ | -0.4 | -0.4 | -0.4 | -0.4 | -0.4 |
| Proposal $P_{\text{std}}$. | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 |
| Shared details of consistency models |  |  |  |  |  |
| Learning rate max ($\alpha_{\text{ref}}$) | 1.0e-4 | 9.0e-5 | 8.0e-5 | 7.0e-5 | 6.5e-5 |
| Proposal $P_{\text{mean}}$ | -0.8 | -0.8 | -0.8 | -0.8 | -0.8 |
| Proposal $P_{\text{std}}$. | 1.6 | 1.6 | 1.6 | 1.6 | 1.6 |
| Tangent normalization constant ($c$) | 0.1 | 0.1 | 0.1 | 0.1 | 0.1 |
| Tangent warm up iterations | 10k | 10k | 10k | 10k | 10k |
| EMA length ($\sigma_{\text{rel}}$) of pretrained diffusion | 0.025 | 0.03 | 0.015 | 0.02 | 0.015 |
| Training details of sCT |  |  |  |  |  |
| Training iterations | 100k | 100k | 100k | 100k | 100k |
| Dropout probability for resolution $\leq 16$ | 25% | 35% | 35% | 35% | 35% |
| Training details of sCD |  |  |  |  |  |
| Training iterations | 200k | 200k | 200k | 200k | 200k |
| Dropout probability for resolution $\leq 16$ | 0% | 10% | 10% | 10% | 10% |
| Maximum of CFG scale | 2.0 | 2.0 | 2.0 | 2.0 | 2.0 |
| Training details of sCD with adaptive VSD |  |  |  |  |  |
| Training iterations | 20k | 20k | 20k | 20k | 20k |
| Learning rate max ($\alpha_{\text{ref}}$) for ${\bm{F}}_{\phi}$ | 1.0e-4 | 9.0e-5 | 8.0e-5 | 7.0e-5 | 6.5e-5 |
| Dropout probability for ${\bm{F}}_{\phi}$ | 0% | 10% | 10% | 10% | 10% |
| Proposal $P_{\text{mean}}$ for $\mathcal{L}_{\text{Diff}}(\phi)$ | -0.8 | -0.8 | -0.8 | -0.8 | -0.8 |
| Proposal $P_{\text{std}}$. for $\mathcal{L}_{\text{Diff}}(\phi)$ | 1.6 | 1.6 | 1.6 | 1.6 | 1.6 |
| Number of updating of $\phi$ per updating of $\theta$ | 1 | 1 | 1 | 1 | 1 |
| One-step sampling starting time $t_{\text{max}}$ | $\arctan(\frac{80}{\sigma_{d}})$ | $\arctan(\frac{80}{\sigma_{d}})$ | $\arctan(\frac{80}{\sigma_{d}})$ | $\arctan(\frac{80}{\sigma_{d}})$ | $\arctan(\frac{80}{\sigma_{d}})$ |
| Proposal $P_{\text{mean}}$ for $\mathcal{L}_{\text{VSD}}(\theta)$ | 0.4 | 0.4 | 0.4 | 0.4 | 0.4 |
| Proposal $P_{\text{std}}$. for $\mathcal{L}_{\text{VSD}}(\theta)$ | 2.0 | 2.0 | 2.0 | 2.0 | 2.0 |
| Loss weighting $\lambda_{\text{VSD}}$ for $\mathcal{L}_{\text{VSD}}$ | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 |

**Table 5: Evaluation of sample quality of different models on ImageNet 512$\times$512 dataset. Results of EDM2 (Karras et al., ) are with EDM parameterization and the original AdaGN layer. ^†The $\text{FD}_{\text{DINOv2}}$in EDM2 are obtained by tuned EMA rate, which is different from our EMA rates that are tuned for FID scores.**
|  | Model Size |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
|  | S | M | L | XL | XXL |
| Sampling by diffusion models (NFE = 126) |  |  |  |  |  |
| EMA length ($\sigma_{\text{rel}}$) | 0.025 | 0.030 | 0.015 | 0.020 | 0.015 |
| Guidance scale for FID | 1.4 | 1.2 | 1.2 | 1.2 | 1.2 |
| ^†Guidance scale for $\text{FD}_{\text{DINOv2}}$ | 2.0 | 1.8 | 1.8 | 1.8 | 1.8 |
| FID (TrigFlow) | 2.29 | 2.00 | 1.87 | 1.80 | 1.73 |
| FID (EDM2) | 2.23 | 2.01 | 1.88 | 1.85 | 1.81 |
| $\text{FD}_{\text{DINOv2}}$(TrigFlow) | 52.08 | 43.33 | 39.23 | 36.73 | 35.93 |
| ^†$\text{FD}_{\text{DINOv2}}$(EDM2) with $\sigma_{\text{rel}}$ for $\text{FD}_{\text{DINOv2}}$ | 52.32 | 41.98 | 38.20 | 35.67 | 33.09 |
| Sampling by consistency models trained with sCT |  |  |  |  |  |
| Intermediate time $t_{\text{mid}}$ in 2-step sampling | 1.1 | 1.1 | 1.1 | 1.1 | 1.1 |
| 1-step FID | 10.13 | 5.84 | 5.15 | 4.33 | 4.29 |
| 2-step FID | 9.86 | 5.53 | 4.65 | 3.73 | 3.76 |
| 1-step $\text{FD}_{\text{DINOv2}}$ | 278.35 | 192.13 | 169.98 | 147.06 | 146.31 |
| 2-step $\text{FD}_{\text{DINOv2}}$ | 244.41 | 160.66 | 135.80 | 114.65 | 112.69 |
| Sampling by consistency models trained with sCD |  |  |  |  |  |
| Intermediate time $t_{\text{mid}}$ in 2-step sampling | 1.1 | 1.1 | 1.1 | 1.1 | 1.1 |
| Guidance scale for FID, 1-step sampling | 1.5 | 1.3 | 1.3 | 1.3 | 1.3 |
| Guidance scale for FID, 2-step sampling | 1.4 | 1.2 | 1.2 | 1.2 | 1.2 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$, 1-step sampling | 2.0 | 2.0 | 2.0 | 2.0 | 2.0 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$, 2-step sampling | 2.0 | 2.0 | 1.9 | 1.9 | 1.9 |
| 1-step FID | 3.07 | 2.75 | 2.55 | 2.40 | 2.28 |
| 2-step FID | 2.50 | 2.26 | 2.04 | 1.93 | 1.88 |
| 1-step $\text{FD}_{\text{DINOv2}}$ | 104.22 | 83.78 | 76.10 | 70.30 | 67.80 |
| 2-step $\text{FD}_{\text{DINOv2}}$ | 71.15 | 55.70 | 50.63 | 46.66 | 44.97 |
| Sampling by consistency models trained with multistep sCD |  |  |  |  |  |
| Guidance scale for FID | 1.4 | 1.2 | 1.2 | 1.15 | 1.15 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$ | 2.0 | 2.0 | 2.0 | 1.9 | 1.9 |
| FID, $M=2$ | 2.79 | 2.51 | 2.32 | 2.29 | 2.16 |
| FID, $M=4$ | 2.78 | 2.46 | 2.28 | 2.22 | 2.10 |
| FID, $M=8$ | 2.49 | 2.24 | 2.04 | 2.02 | 1.90 |
| FID, $M=16$ | 2.34 | 2.18 | 1.99 | 1.90 | 1.82 |
| $\text{FD}_{\text{DINOv2}}$, $M=2$ | 76.29 | 60.47 | 54.91 | 51.91 | 50.70 |
| $\text{FD}_{\text{DINOv2}}$, $M=4$ | 72.01 | 56.38 | 50.99 | 47.61 | 46.78 |
| $\text{FD}_{\text{DINOv2}}$, $M=8$ | 60.13 | 49.46 | 44.87 | 41.26 | 40.56 |
| $\text{FD}_{\text{DINOv2}}$, $M=16$ | 55.89 | 46.94 | 42.55 | 39.30 | 38.55 |
| Sampling by consistency models trained with sCD + adaptive VSD |  |  |  |  |  |
| Intermediate time $t_{\text{mid}}$ in 2-step sampling | 1.1 | 1.1 | 1.1 | 1.1 | 1.1 |
| Guidance scale for FID, 1-step sampling | 1.2 | 1.0 | 1.0 | 1.0 | 1.0 |
| Guidance scale for FID, 2-step sampling | 1.2 | 1.0 | 1.0 | 1.0 | 1.0 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$, 1-step sampling | 1.7 | 1.5 | 1.6 | 1.5 | 1.5 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$, 2-step sampling | 1.7 | 1.5 | 1.6 | 1.5 | 1.5 |
| 1-step FID | 3.37 | 2.67 | 2.26 | 2.39 | 2.16 |
| 2-step FID | 2.70 | 2.29 | 1.99 | 2.01 | 1.89 |
| 1-step $\text{FD}_{\text{DINOv2}}$ | 72.12 | 54.81 | 50.46 | 48.11 | 45.54 |
| 2-step $\text{FD}_{\text{DINOv2}}$ | 69.00 | 53.53 | 48.54 | 46.61 | 43.93 |

**Table 6: Ablation of adaptive VSD and sCD on ImageNet 512$\times$512 dataset with model size M.**
|  | Method |  |  |
| --- | --- | --- | --- |
|  | VSD | sCD | sCD + VSD |
| EMA length ($\sigma_{\text{rel}}$) | 0.05 | 0.05 | 0.05 |
| Guidance scale for FID, 1-step sampling | 1.1 | 1.3 | 1.0 |
| Guidance scale for FID, 2-step sampling | $\backslash$ | 1.2 | 1.0 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$, 1-step sampling | 1.4 | 2.0 | 1.5 |
| Guidance scale for $\text{FD}_{\text{DINOv2}}$, 2-step sampling | $\backslash$ | 2.0 | 1.5 |
| 1-step FID | 3.02 | 2.75 | 2.67 |
| 2-step FID | $\backslash$ | 2.26 | 2.29 |
| 1-step $\text{FD}_{\text{DINOv2}}$ | 57.19 | 83.78 | 54.81 |
| 2-step $\text{FD}_{\text{DINOv2}}$ | $\backslash$ | 55.70 | 53.53 |

**Table 7: Evaluation of sample quality of different models on ImageNet 64$\times$64 dataset.**
|  | Model Size |  |  |  |
| --- | --- | --- | --- | --- |
|  | S | M | L | XL |
| Sampling by diffusion models (NFE=63) |  |  |  |  |
| EMA length ($\sigma_{\text{rel}}$) | 0.075 | 0.06 | 0.04 | 0.04 |
| FID (TrigFlow) | 1.70 | 1.55 | 1.44 | 1.38 |
| Sampling by consistency models trained with sCT |  |  |  |  |
| Intermediate time $t_{\text{mid}}$ in 2-step sampling | 1.1 | 1.1 | 1.1 | 1.1 |
| 1-step FID | 3.23 | 2.25 | 2.08 | 2.04 |
| 2-step FID | 2.93 | 1.81 | 1.57 | 1.48 |
| Sampling by consistency models trained with sCD |  |  |  |  |
| Intermediate time $t_{\text{mid}}$ in 2-step sampling | 1.1 | 1.1 | 1.1 | 1.1 |
| 1-step FID | 2.97 | 2.79 | 2.43 | 2.44 |
| 2-step FID | 2.07 | 1.89 | 1.70 | 1.66 |

## Appendix H Additional Samples

Figure: Figure 8: Uncurated 1-step samples generated by our sCD-XXL trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-appendix-1step/15_batch.jpg

Figure: Figure 9: Uncurated 2-step samples generated by our sCD-XXL trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-appendix/15_batch.jpg

Figure: Figure 10: Uncurated 1-step samples generated by our sCD-XXL trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-appendix-1step/127_batch.jpg

Figure: Figure 11: Uncurated 2-step samples generated by our sCD-XXL trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-appendix/127_batch.jpg

Figure: Figure 12: Uncurated 1-step samples generated by our sCD-XXL trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-appendix-1step/417_batch.jpg

Figure: Figure 13: Uncurated 2-step samples generated by our sCD-XXL trained on ImageNet 512$\times$512.
Refer to caption: extracted/5926597/figures/samples-appendix/417_batch.jpg

## Peer Reviews

confidence: 5
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. I appreciate the scaling study from ImageNet 64x64 to 512x512, where the former operates directly in the image space, while the latter relies on additional image compression models. The difference between sCT and sCD at the 512x512 resolution is intriguing, as even increasing model size and batch size cannot easily close the gap, while at the 64x64 resolution, scaling is sufficient to compensate for the variance induced by Monte Carlo estimation.
    - Data complexity at 512x512 and the lack of effective mode decomposition definitely contribute to this discrepancy. However, I am curious about the authors' thoughts on the extent to which the increased variance could be caused by the pretrained image encoder/decoder. In other words, could data modes become *more dispersed* in the latent space, making it harder for sCT to learn effectively? Would it be better to train CMs directly in the image space (with some special architecture design and weighting schemes) or to find a latent space that is more suitable for CMs? In some sense, modern autoregressive models focus on tokenizer design for visual generation. Could latent space compression for CMs require properties distinct from those used in DMs? While this is more hypothetical, I would be happy to hear more thoughts.

2. If I understand correctly, there are three weighting functions applied, namely learnable adaptive weighting, tangent normalization (per-sample basis weighting), a prior weighting $w(t) = \frac{1}{\sigma_d \tan(t)} = \frac{1}{\sigma_t}$. To what extent does this prior weighting contribute to variance reduction? Is it helpful to stabilize the learnable adaptive weighting layer? I assume the time embeddings of learnable adaptive weighting could be either positional embedding or Fourier embedding since it is not directly involved in $\frac{\mathrm{d} \boldsymbol{f}}{\mathrm{d}t}$.

As an additional comment, could the authors consider conducting distillation experiments using the data-free formulation in [1]? I am curious how continuous-time CMs would scale and how stable they would be without using an extra dataset. No hurry to complete these experiments during the limited rebuttal period!

[1] Consistency Models Made Easy

### Minor

1. In Line 213 and in Line 227, both diffusion models and consistency models are denoted as $f_\theta(\mathbf{x}_t, t)$ but with different equations.

2. The Adaptive Double Normalization is less explained. Is it the same as local response normalization applied to the modulation layer?
soundness: 4
strengths: This is a very strong paper in analysis, practical techniques, writing, and experiment results. 

- Novelty. This paper's novelty is evident in several aspects.
    - First, it studies an important but less studied problem: consistency models in continuous time, together with the training stability and discretization error of consistency models. 
    - The proposed TrigFlow, as a novel unification of EDM and Flow Matching, substantially simplifies the analysis presented later and the practical techniques.
    - The gradient analysis of continuous-time objective reveals the root cause of instability. To the best of my knowledge, this is the first paper to establish the gradient analysis for CMs.
    - Model architecture modifications are original since existing works are mostly inherited from Diffusion Models' design and focus on the training techniques and formulations, leaving the architectural design underexplored.

- Soundness. Its technical claims are well backed up by both theoretical analysis and empirical results. I particularly appreciate the in-depth investigation into the training dynamics and gradient analysis of continuous-time CMs.

- Presentation. The logical flow of this paper is well structured and smooth. The problem statement is clearly defined, and the explanation of why discretization errors matter for CMs and the motivation toward continuous-time formulation is crystal clear. The gradient analysis into continuous-time CMs is thoughtfully motivated and carefully organized. Even the appendix is well-written, offering useful insights into the proposed techniques.

- Experiments. Proposed techniques allow for training continuous-time Consistency Models (sCMs) at an unprecedented scale. Experiment results are impressive, matching/outperforming adversarial approaches, score distillation, and recent autoregressive models.
    - Gradient variances have been carefully controlled via adaptive weighting and normalization techniques.
    - Comprehensively studying the scaling behaviors of sCMs under continuous-time training. 
    - Comparisons with improved score distillation baseline using many methods developed in this work confirm the mode coverage of CMs.
    - Additionally, the paper discusses efficient and stable implementation strategies for continuous-time CMs.

Given the potential impact of this paper, I strongly recommend acceptance with conference highlights. It was a great pleasure to read through the manuscript!
summary: This paper investigates a fundamental topic in consistency models (CMs), specifically the challenges of discretization errors and the resulting training stability issue. Consistency Models can be trained in discrete or continuous time, either from scratch using a dataset or distilled from pretrained teacher scores. CMs' theoretical foundation elucidates the importance of controlling the discretization error and eventually achieving consistency in continuous time. While continuous-time CMs eliminate the discretization errors present in their discrete-time counterparts, they suffer from training instability, a problem that is not yet well understood in the research community.

This work conducts a comprehensive study into continuous-time CMs, covering forward process parameterization, network architecture, and training techniques. The authors first develop a simplified diffusion process formulation called TrigFlow, which unifies EDM and Flow Matching for the first time. Building upon this foundation, they analyze the gradient flow of continuous-time CMs, identify the root cause of training instability, and mitigate this issue through modifications to time embeddings and adaptive group normalization. Additional training techniques, such as adaptive weighting functions and annealing, further contribute to improved training stability and scalability.

The resulting method, sCT/sCD, allows continuous-time CMs to be trained at an unprecedented scale, scaling up to 1.5B parameters on ImageNet 512x512. These results significantly narrow the performance gap between CMs and state-of-the-art diffusion models to less than 10% in FID, while matching or even surpassing adversarial methods and discrete/continuous autoregressive models in both performance and efficiency.
weaknesses: I did not find any apparent weaknesses in the analysis or experiments (including both ablation studies and performance evaluation). There are research questions worth further investigation, as discussed below.

---

confidence: 4
contribution: 4
initial_rating: 10
presentation: 4
questions: These are mostly from the weakness section:
1. Are there more limitations other than the 10% worse than diffusion SOTA?
2. Is the method really fully stable?
soundness: 4
strengths: The paper is very well grounded mathematically and experimentally. The analysis is based on understanding the causes of training instabilities by decomposing the loss, validating each component experimentally and proposing changes to solve the root causes.
The mathematics while greatly simplified are still pretty complex and the paper shines in its clarity to make the logical reasoning easy to follow.
The experimental results are also outstanding resulting in very significant gains, essentially taking consistency models within 10% of the SOTA for diffusion models.
summary: The authors propose improvements to the consistency models generative paradigm and named their new method sCM. Specifically they -vastly- improve the FID for consistency models with the introduction of several new ideas to both stabilize and simplify continuous consistency models. 

My understanding is the main claim of simplification for sCM comes from the simplification of EDM (Kerras et al.) normalizing design, resulting in $c_{in}=c_{skip}=1$ which in turns simplifies the continuous expression of consistency models. Another simplification is the combination of both EDM and Flow Matching concepts into their method which they call TrigFlow. Yet another simplification, not claimed as such by authors, is the use of vanilla L2 loss compared to Huber/LPIPS use in previous iterations of consistency models. This last simplification has the additional benefit to be more probabilistically grounded.

There are 3 main proposed ideas to stabilize the training of consistency models:
1. Identity-time transformation as a replacement to the log-transformation from EDM
2. Fourier embedding of the time dimension are replaced by positional embeddings
3. AdaGN is modified to also normalize the conditioning inputs for scale and bias.
More ideas are also proposed in the training objective to stabilize training, namely: tangent normalization and tangent warmup. It is my understanding that the adaptative weighing is the same as in EDM.

The paper also provides ample ablations to demonstrate the effects and the reasoning motivating these 3 proposed improvements.
weaknesses: 1. The limitations of the method are not very clear to me besides the fact that it's still 10% worse than SOTA for diffusion.
2. The section on positional embeddings (line 269 and on) lacks details to be fully understandable without having to read another paper. Maybe beefing up that section would make the paper more self-contained.
3. I found all figure very useful with the exception of Figure (3) which I felt did not add much value.

I also noticed a typo (definitely not affecting my score), just leaving it there for authors to fix their manuscript:
- Line 362: "cause instability" => "causes instability"

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: Q1 - In Section 4.1, the authors emphasize the importance of time transformation in mitigating numerical instability. Have the authors considered other potential candidates for time transformation?

Q2 - In Figure 6(a) and Section 5.2, the authors mention that sCT is less effective at higher resolutions. Do the authors have any insights into why this might be?

Q3 - In Figure 6(b), it appears that the performance of sCD-XL under the two-step sCD setting is better than that of sCD-XXL, which contradicts the results reported in Table 2. Could the authors clarify the specific settings used in Figure 6(b)?
soundness: 3
strengths: S1 - The paper provides a comprehensive analysis and set of solutions addressing the numerical instability issues in continuous-time consistency models, significantly improving performance and enabling the model to achieve competitive results on selected benchmarks.

S2 - Many of the enhancements are supported by detailed theoretical justification and experimental results.

S3 - The unified perspective on previous diffusion and flow-matching parameterizations is thorough, complete, and well-grounded, offering novel insights that could benefit the community.

S4 - The paper is well-structured and easy to follow.
summary: The paper presents a unified perspective on diffusion-based and flow-based generative models and introduces a comprehensive set of techniques aimed at improving the training stability and overall performance of continuous-time consistency models for large-scale image generation. The techniques include: 1) enhancing time transformation and embeddings, 2) replacing the AdaGN layer with Adaptive Double Normalization, 3) normalizing the tangent function and applying tangent warm-up, 4) implementing an adaptive weighting function in the training objective, and 5) optimizing forward-mode differentiation. These techniques mitigate the numerical instability issues in continuous-time consistency models and enable the model to achieve highly competitive performance in class-conditioned image generation.
weaknesses: W1 - Several design choices appear arbitrary and lack supporting evidence. 

For example, in Section 4.1, the authors discuss the preference for Adaptive Double Normalization over AdaGN, but there is no experimental evidence supporting this choice. It would be more insightful to add a Figure similar to Figure 5 show experimental comparison between Adaptive Double Norm and AdaGN. 

Similarly, in Section 4.2, the authors propose training with linear warm-up w.r.t the model's time derivative, yet no evidence is provided to demonstrate this choice’s effectiveness. Again, including an ablation study or comparative analysis showing the impact of the linear warm-up on training stability or performance metrics would provide more concrete evidence for the effectiveness of this specific choice

Furthermore, Figure 5(b) suggests that incorporating adaptive weighting in a two-step setting may lead to worse performance, while in the one-step setting, it only yields marginal improvement. Have the authors considered alternative designs for the two step case?

W2 - In Sections 4.1 and 5.2, the paper discusses the training compute of sCM. However, including a comparison of compute efficiency with other models (e.g., ECT [1]) would be more insightful, maybe a table or figure comparing the compute efficiency (e.g., FLOPs or training time) of sCM against ECT and other relevant baselines for a given performance level. 

Additionally, given that the model is trained on a large-scale dataset (ImageNet 512) under latent setting, it would be beneficial to include discussions related to text-to-image generation.

[1] Geng, Zhengyang, Ashwini Pokle, William Luo, Justin Lin, and J. Zico Kolter. "Consistency Models Made Easy." arXiv preprint arXiv:2406.14548 (2024).

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: How do the proposed models compare with [1} and [2]? Even tests on a small dataset using small models of equal size would be informative.. I completely understand however  if such comparisons cannot be made due to the short length of the discussion period. 

Shouldn't $c_{skip}$ and $c_{out}$ be $\sigma_t$ and  $ \alpha_t \sigma_d$, that is for $\sigma_t=t$ and  $\alpha_t=1$: $c_{skip}=t$  and $c_{out}=\sigma_d$ in line 201/202?

Based on the experiments performed, are there any indications that the continuous consistency models will still face instability issues for even larger scales, in particular as compared to diffusion/flow models? This question mostly relates to this part in the paper: "Additionally,
we observe that consistency training is more effective at smaller scales but suffers from increased
variance at larger scales, while consistency distillation shows consistent performance across both
small and large scales."

[1]  Tong et al. 2024. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport

[2]  Kornilov et al 2024. Optimal Flow Matching: Learning Straight Trajectories in Just One Step
soundness: 3
strengths: The contributions of this paper are novel, clear, significant and positioned correctly.

Trigflow simplifies the theoretical analysis of flows by showing that  the mean ($\alpha_t$ ) and noise ($\sigma_t$) schedule in the definition of conditional flows can be normalized while preserving the model/loss formulation and the integrator-generated paths, while simplifying the connection between the scaling parameters of the model and those of the conditional flow. **(novelty, significance)**

Several components of continuous consistency models are studied and probed for potential causes of instability. They include: an inappropriate $c_{noise}(t)$ which is fixed by being set to $t$; inappropriate Fourier scales in the time positional embeddings which are then reduced; the usage  'AdaGN layer', which is then replaced with 'Adaptive Double Normalization'; highly varying target norm alleviated by target normalization; inappropriate weighting mitigated by adaptive weighting and unstable terms solved by slowly introducing such terms into the loss with respect to the number of parameter updates.  **(novelty, significance)**

Strategies for scaling such models to large sizes and datasets are proposed, namely JVP Rearrangement and JVP of Flash Attention. **(quality, significance)**

The proposed method shows competitiveness with the state of the art models, while only using 1 or 2 steps for generation, and outperforms all other tested methods that use 1 to 2 generation steps **(significance, quality)**
summary: This paper tackles the problem of  instability in continuous consistency models and presents many contributions. 

The first novelty is the simplification of model parametrization in the EDM/CM frameworks. It is shown that rescaling the mean ($\alpha_t$ ) and noise ($\sigma_t$) schedule of conditional probability flows with the norm $||(\alpha_t, \sigma_t)||$, produces a second conditional flow, such that the EDM model parametrization and loss in the first are identical to those in the second. Furthermore, the ODE sampling procedure remains unchanged. However, the connection between the new flow schedule parameters and the  EDM scaling parameters is greatly simplified after normalization, facilitating further theoretical analysis, while maintaining the benefits of the EDM model definition (unit variance of input, target and minimizing the scaling of the output). The formulation with $\alpha_t=cos(t)$ and  $\sigma_t=sin(t)$ is named *TrigFlow*.

Then the training objective is studied and probed for causes of instability. Several factors are detected such as: an inappropriate $c_{noise}(t)$ which is fixed by being set to $t$; inappropriate Fourier scales in the time positional embeddings which are then reduced; the usage  'AdaGN layer', which is then replaced with 'Adaptive Double Normalization'; highly varying target norm alleviated by target normalization; inappropriate weighting mitigated by adaptive weighting and unstable terms solved by slowly introducing such terms into the loss with respect to the number of parameter updates. 

In addition, methods for scaling such models to large sizes and datasets are proposed, namely JVP Rearrangement and JVP of Flash Attention.

Finally, the proposed models are compared against the state of the art methods, and results show that continuous consistency models are excellent generators, that only require 1 or 2 sampling steps to generate new quality data from learned distributions with continuous support.
weaknesses: **While the contributions of the paper are numerous, the paper could be strengthened even further by:**

1) Comparing the proposed model with more recent flow models such as [1] and [2] and adding the results for rectified flows with 2 generation steps.

2) Placing the number of parameters and the number of parameter updates (or training time on identical hardware) for each model in Table 1 is very important as using an equal umber of parameters/compute is essential for ensuring a fair comparison.

**Some smaller issues and suggestions:**

a) The paper would be improved if an intuitive explanation is given for the loss in Equation 2. Also it could save readers some time if it is mentioned that the loss is derived in Song et al 2023, *Remark 10*. 

b) The paper would be enriched by adding some generated images with one step.

c) Shouldn't $c_{skip}$ and $c_{out}$ be $\sigma_t$ and  $ \alpha_t \sigma_d$, that is for $\sigma_t=t$ and  $\alpha_t=1$: $c_{skip}=t$  and $c_{out}=\sigma_d$ in line 201/202?

d) In Equation (20) Appendix, shouldn't it be $\hat{D}$ for $D_{\theta}(x_t)=\hat{D}_{\theta}(\hat{x}_t(x_t))$?

e) The paragraph from line 924 to 938 (appendix) needs additional elaboration regarding the implications of having $||(\alpha_t, \sigma_t)||=1$ with respect to the invariance of the geometric set connecting $x_0$ and $z$.



**Typos:**

i) In line 126, a 2 is squared instead of the norm.

ii) in line 122, $z_t$ does not depend on time.


**References**


[1]  Tong et al. 2024. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport

[2]  Kornilov et al 2024. Optimal Flow Matching: Learning Straight Trajectories in Just One Step

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See the weaknesses in the above.
soundness: 3
strengths: - This paper is very well-written and easy to read. 
- This work proposed a new diffusion formulation, called TrigFlow, that unifies EDM and Flowing Matching, and also simplifies the analysis of continuous-time consistency models. 
- It provided a thorough analysis of the training stability of continuous-time consistency models, from the perspective of network architecture, training objective and diffusion process parameterization.
- Experiments on CIFAR-10, ImageNet-64 and ImageNet-512 demonstrate the effectiveness of the proposed method and the scalability of continuous-time consistency models.
summary: This work proposed a set of improved training techniques to stabilize the training of continuous-time consistency models, including new consistency function formulations, new network architectures and new training objectives. With these new training techniques, the proposed method called sCMs outperformed all previous consistency models in terms of one-step and two-step FIDs.
weaknesses: - Although I really like the improvements of continuous-time consistency models, which could fundamentally eliminate the discretization error in discrete-time consistency models, it comes with more time and memory costs related to JVP computation in the loss function. To this end, this work introduces JVP of Flash Attention to reduce the costs, which is great. Still, there may be a considerable gap between the continuous-time and discrete-time consistency models. I wonder if the paper can provide a more detailed comparison between sCMs and the previous discrete-time consistency models - ECMs, in terms of the training convergence and memory cost. 
- There is no explanation for the phenomenon that sCT performs better than sCD on CIFAR-10 and ImageNet-64, but sCTs performs worse on ImageNet-512. Any intuition of why sCT suffers from increased variance at larger scales? 
- There are no ablation study results on “Adaptive Double Normalization” except for claiming it “removes its instability in CM training”.
- In Figure 5b, it looks like “w/o adaptive weighting” achieves better two-step FIDs than “w/ adaptive weighting” and very similar one-step FIDs to “w/ adaptive weighting”. Why do we need adaptive weighting?
- In Figure 5c, do discrete-time CMs have a constant number of time steps $N$ or a timestep schedule up to the maximum number of steps $N$? If it is the former one, it seems to be a bit unfair to discrete-time CMs because the scheduling of time steps is very important to them. Does it make more sense to compare with the best-performing discrete-time CMs?
- In Figure 7, does the paper apply TTUR proposed by DMD2 (Yin et al. 2024a)? From the DMD2 paper, TTUR improves the performance of VSD. Thus, a comparison with VSD + TTUR is more convincing.
- In Figure 7, sCDs condition the consistency network on the guidance scale $s$. I wonder if VSD also condition the generator on the guidance scale, for a consistent evaluation setting? 
- A minor issue: In line 266, should it be $c_{\text{noise}}(t) = \frac{1}{4} \log(\sigma_d \tan t) $?

Meta-review: Consistency models (CMs) are specializations of diffusion models for the purpose of faster sampling. While existing CMs are faster, they are difficult to train, and the quality lags behind other diffusion models. This work provides a close study and several improvements that theoretically re-connect CM and other classes of diffusion, simplify this theory, stabilize training as a result of the theoretical analysis, and then scale up the model, training, and data. The result is a CM that narrows the gap with the state-of-the-art for diffusion on ImageNet 512x512. All told, this work is strong on multiple fronts with theoretical, empirical, and engineering contributions that significantly improve the state-of-the-art for CMs.

This a clear accept with all ratings >= 8 and 3/5 reviewers giving the highest rating. The meta-reviewer agrees because this submission delivers on theoretical and empirical contributions while improving results on a popular and significant topic that is unarguably well-received by the expert reviewers with backgrounds on diffusion modeling, consistency models, and large-scale generative modeling.

Decision: Accept (Oral)