# Improving Efficiency of Diffusion Models via Multi-Stage Framework and
Tailored Multi-Decoder Architectures

Huijie Zhang  Department of Electrical Engineering & Computer Science,
University of Michigan  Yifu Lu111The first two authors contributed to this
work equally. Department of Electrical Engineering & Computer Science,
University of Michigan  Ismail Alkhouri  Department of Electrical Engineering
& Computer Science, University of Michigan  Department of Computational
Mathematics, Science & Engineering, Michigan State University  Saiprasad
Ravishankar  Department of Computational Mathematics, Science & Engineering,
Michigan State University  Department of Biomedical Engineering, Michigan
State University
Dogyoon Song  Department of Electrical Engineering & Computer Science,
University of Michigan  Qing Qu  Department of Electrical Engineering &
Computer Science, University of Michigan

###### Abstract

Diffusion models, emerging as powerful deep generative tools, excel in various
applications. They operate through a two-steps process: introducing noise into
training samples and then employing a model to convert random noise into new
samples (e.g., images). However, their remarkable generative performance is
hindered by slow training and sampling. This is due to the necessity of
tracking extensive forward and reverse diffusion trajectories, and employing a
large model with numerous parameters across multiple timesteps (i.e., noise
levels). To tackle these challenges, we present a multi-stage framework
inspired by our empirical findings. These observations indicate the advantages
of employing distinct parameters tailored to each timestep while retaining
universal parameters shared across all time steps. Our approach involves
segmenting the time interval into multiple stages where we employ custom
multi-decoder U-net architecture that blends time-dependent models with a
universally shared encoder. Our framework enables the efficient distribution
of computational resources and mitigates inter-stage interference, which
substantially improves training efficiency. Extensive numerical experiments
affirm the effectiveness of our framework, showcasing significant training and
sampling efficiency enhancements on three state-of-the-art diffusion models,
including large-scale latent diffusion models. Furthermore, our ablation
studies illustrate the impact of two important components in our framework:
(i) a novel timestep clustering algorithm for stage division, and (ii) an
innovative multi-decoder U-net architecture, seamlessly integrating universal
and customized hyperparameters.

###### Contents

  1. 1 Introduction
  2. 2 Preliminaries & Related Work
    1. 2.1 Background on Diffusion Models
    2. 2.2 Timestep Clustering
    3. 2.3 Reducing the Sampling Iterations
  3. 3 Proposed Multistage Framework
    1. 3.1 Proposed Multi-stage U-Net Architectures
    2. 3.2 Optimal Denoiser-based Timestep Clustering
  4. 4 Identification of Key Sources of Inefficiency
    1. 4.1 Empirical Observations on the Key Sources of Inefficiency
    2. 4.2 Tackling the Inefficiency via Multistage U-Net Architectures
  5. 5 Experiments
    1. 5.1 Comparison of Image Generation Quality
    2. 5.2 Comparison of Training & Sampling Efficiency
    3. 5.3 Comparison of Different Architectures
    4. 5.4 Comparison of Timestep Clustering Methods
  6. 6 Conclusion
  7. A Proof of Proposition 1
    1. A.1 Background
    2. A.2 Proof
  8. B Generalized Algorithm 1
  9. C Additional Experiments
    1. C.1 Generation Results
    2. C.2 Ablation Study on the Network Parameters
    3. C.3 Ablation Study on the Number of Intervals
    4. C.4 Experimental Setting on Interval Splitting Methods

##  1 Introduction

Figure 1: Overview of three diffusion model architectures: (a) unified, (b)
separate, and (c) our proposed multistage architectures. Compared with (a) and
(b), our approach improves sampling quality, and significantly enhances
training efficiency, as indicated by the FID scores and their corresponding
training iterations (d).

Recently, diffusion models have emerged as powerful deep generative modeling
tools, showcasing remarkable performance in various applications, ranging from
unconditional image generation [1, 2], conditional image generation [3, 4],
image-to-image translation [5, 6, 7], text-to-image generation [8, 9, 10],
inverse problem solving [11, 12, 13, 14], video generation [15, 16], and so
on. These models employ a training process involving continuous injection of
noise into training samples (“diffusion”), which are then utilized to generate
new samples, such as images, by transforming random noise instances through a
reverse diffusion process guided by the “score function222The (Stein) score
function of distribution $p_{t}$ at $\bm{x}_{t}$ is $\nabla_{\bm{x}_{t}}\log
p_{t}(\bm{x}_{t})$. ” of the data distribution learned by the model. Moreover,
recent work demonstrates that those diffusion models enjoy optimization
stability and model reproducibility compared with other types of generative
models [17]. However, diffusion models suffer from slow training and sampling
despite their remarkable generative capabilities, which hinders their use in
applications where real-time generation is desired [1, 2]. These drawbacks
primarily arise from the necessity of tracking extensive forward and reverse
diffusion trajectories, as well as managing a large model with numerous
parameters across multiple timesteps (i.e., diffusion noise levels).

In this paper, we address these challenges based on two key observations: (i)
there exists substantial parameter redundancy in current diffusion models, and
(ii) they are trained inefficiently due to dissimilar gradients across
different noise levels. Specifically, we find that training diffusion models
require fewer parameters to accurately learn the score function at high noise
levels, while larger parameters are needed at low noise levels. Furthermore,
we also observe that when learning the score function, distinct shapes of
distributions at different noise levels result in dissimilar gradients, which
appear to slow down the training process driven by gradient descent.

Building on these insights, we propose a multi-stage framework with two key
components: (1) a multi-decoder U-net architecture, and (2) a new partitioning
algorithm to cluster timesteps (noise levels) into distinct stages. In terms
of our new architecture, we design a multi-decoder U-Net that incorporates one
universal encoder shared across all intervals and individual decoders tailored
to each time stage; see Figure 1 (c) for an illustration. This approach
combines the advantages of both universal and stage-specific architectures,
which is much more efficient that using a single architecture for the entire
training process [2, 1, 18] (Figure 1 (a)). Moreover, compared to previous
approaches that completely separate architectures for each sub-interval [19,
20, 21, 22] (Figure 1 (b)), our approach can effectively mitigate interference
between stages arising from disparate gradient effects, leading to improved
efficiency. On the other hand, when it comes to partitioning the training
stages of our network, we designed an algorithm aimed at grouping the
timesteps. This is achieved by minimizing the functional distance within each
cluster in the training objective and making use of the optimal denoiser
formulation [18]. By integrating these two key components, our framework
enables efficient allocation of computational resources (e.g., U-net
parameters) and stage-tailored parameterization. In Section 5, we showcase via
extensive numerical experiments that our framework effectively improves both
training and sampling efficiency. These experiments are performed on diverse
benchmark datasets, demonstrating significant acceleration by using our
framework when compared to three state-of-the-art (SOTA) diffusion model
architectures.

#### Contributions.

As a summary, the major contributions of this work can be highlighted as
follows:

  * •

Identifying two key sources of inefficiency. We identified two key sources
that cause inefficiencies in training diffusion models across various time
step: (1) a significant variation in the requirement of model capacity, and
(2) the gradient dissimilarity. As such, using a unified network cannot meet
with the changing requirement at different time steps.

  * •

A new multi-stage framework. We introduced a new multi-stage architecture,
illustrated in Figure 1 (c). We tackle the challenges by segmenting the time
interval into multiple stages, where we employ customized multi-decoder U-net
architectures that blends time-dependent models with a universally shared
encoder.

  * •

Improved training and sampling efficiency. With comparable computational
resources for unconditional image generation, we demonstrate that our multi-
stage approach improves the Fréchet Inception Distance (FID) score for all
SOTA methods. For example, on CIFAR-10 dataset [23], our method improves the
FID for DPM-Solver [24] from 2.84 to 2.37, and it improves the FID for EDM
[18] from 2.05 (our training result) to 1.96. Moreover, on the CelebA dataset
[25], our approach significantly reduces the training computation of both EDM
to 82% and the Latent Diffusion Model (LDM) [8] to 30% for obtaining a similar
quality of generation.

#### Organization.

In Section 2, we provide an overview of relevant literature to distinguish us
from previous works. In Section 3, we describe our proposed multistage
framework for diffusion modeling, delineating the two core components. In
Section 4, we present our observations and analysis that motivated the
proposed multistage framework, justifying its development. Finally, in Section
5, we provide the results from our numerical experiments that validate the
effectiveness of the proposed multistage approach.

##  2 Preliminaries & Related Work

In this section, we start by reviewing the basic fundamentals of diffusion
models [1, 2, 18]. Subsequently, we delve into prior endeavors aimed at
improving the training and efficiency of diffusion models through the
partitioning of the timestep interval. Lastly, we revisit prior studies that
significantly decrease the number of required sampling iterations.

###  2.1 Background on Diffusion Models

Let $\bm{x}_{0}\in\mathbb{R}^{n}$ denote a sample from the data distribution
$p_{\text{data}}(\bm{x})$. Diffusion models operate within forward and reverse
processes. The forward process progressively perturbs data $\bm{x}_{0}$ to a
noisy version $\bm{x}_{t\in[0,1]}$ via corrupting with the Gaussian kernel.
This process can be formulated as a stochastic differential equation (SDE) [2]
of the form $\text{d}\bm{x}=\bm{x}_{t}f(t)\text{d}t+g(t)\text{d}\bm{w}_{t}$,
where $f(t)$ and $g(t)$ are the drift and diffusion coefficients,
respectively, that correspond to a pre-defined noise schedule.
$\bm{w}_{t}\in\mathbb{R}^{n}$ is the standard Wiener process. Under the
forward SDE, the perturbation kernel is given by the conditional distribution
defined as

$p_{t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\>,$
(1)

where

$s_{t}=\text{exp}(\int_{0}^{t}f(\xi)\text{d}\xi),\ \ \text{and}\
\sigma_{t}=\sqrt{\int_{0}^{t}\frac{g^{2}(\xi)}{s_{\xi}^{2}}\text{d}\xi}\>.$
(2)

The parameters defined in eq. 2 are designed such that: (i) the data
distribution is approximately estimated when $t=0$, and (ii) a nearly standard
Gaussian distribution is obtained when $t=1$. The objective of diffusion
models is to learn the corresponding reverse SDE, defined as

$\displaystyle\text{d}\bm{x}=\left[f(t)\bm{x}_{t}-g^{2}(t)\nabla_{\bm{x}_{t}}\log
p_{t}(\bm{x}_{t})\right]\text{d}t+g(t)\text{d}\bar{\bm{w}},$ (3)

where $\bar{\bm{w}}\in\mathbb{R}^{n}$ is the standard Wiener process running
backward in time, and $\nabla_{\bm{x}_{t}}\log p_{t}(\bm{x}_{t})$ is the
(Stein) score function. In practice, the score function is approximated using
a neural network
$\bm{\epsilon}_{\bm{\theta}}:\mathbb{R}^{n}\times[0,1]\rightarrow\mathbb{R}^{n}$
parameterized by $\bm{\theta}$, which can be trained by the denoising score
matching technique [26] as

$\min\limits_{\bm{\theta}}\mathbb{E}\left[\omega(t)\|\bm{\epsilon}_{\bm{\theta}}(\bm{x}_{t},t)+s_{t}\sigma_{t}\nabla_{\bm{x}_{t}}\log
p_{t}(\bm{x}_{t}|\bm{x}_{0})\|_{2}^{2}\right],$ (4)

which can also be written as
$\min_{\bm{\theta}}\mathbb{E}[\omega(t)||\bm{\epsilon}_{\bm{\theta}}(\bm{x}_{t},t)-\bm{\epsilon}||^{2}]+C$,
where the expectation is taken over $t\sim[0,1]$, $\bm{x}_{t}\sim
p_{t}(\bm{x}_{t}|\bm{x}_{0})$, $\bm{x}_{0}\sim p_{\text{data}}(\bm{x})$, and
$\bm{\epsilon}\sim\mathcal{N}(\textbf{0},\textbf{I})$. Here, $C$ is a constant
independent of $\bm{\theta}$, and $\omega(t)$ is a scalar representing the
weight of the loss as a function of $t$. In DDPM [1], it is simplified to
$\omega(t)=1$. Once the parameterized score function
$\bm{\epsilon}_{\bm{\theta}}$ is trained, it can be utilized to approximate
the reverse-time SDE using numerical solvers such as Euler-Maruyama.

###  2.2 Timestep Clustering

Diffusion models have demonstrated exceptional performance but face efficiency
challenges in training and sampling. In response, several studies proposed to
cluster the timestep range $t\in[0,1]$ into multiple intervals (e.g.,
$[0,t_{1}),[t_{1},t_{2}),\ldots,[t_{n},1]$). Notably, Choi et al. [19]
reconfigured the loss weights for different intervals to enhance performance.
Balaji et al. [27] introduced “expert denoisers”, which proposed using
distinct architectures for different time intervals in text-to-image diffusion
models. Go et al. [22] further improved the efficiency of these expert
denoisers through parameter-efficient fine-tuning and data-free knowledge
transfer. Lee et al. [21] designed separate architectures for each interval
based on frequency characteristics. Moreover, Go et al. [20] treated different
intervals as distinct tasks and employed multi-task learning strategies for
diffusion model training, along with various timestep clustering methods.

Our approach distinguishes itself from the aforementioned methods in two key
aspects. The first key component is our tailored U-net architecture using a
unified encoder coupled with different decoders for different intervals, while
previous models have either adopted a unified architecture, as seen in [19,
20], or employed separate architectures (referred to as expert denoisers) for
each interval [27, 22, 21]. In comparison, our multistage architecture
surpasses these methodologies, as demonstrated in Section 5.3. Second, we
developed a new timestep clustering method leveraging a general optimal
denoiser (Proposition 1) that showcases superior performance (see Section
5.4). In contrast, prior works rely on (i) a simple timestep-based clustering
cost function [27, 20, 22, 21], (ii) Signal-to-Noise Ratio (SNR) based
clustering [20], or (iii) gradients-based partitioning that uses task affinity
scores [20].

###  2.3 Reducing the Sampling Iterations

Efforts to improve sampling efficiency of diffusion models have led to many
recent advancements in SDE and Ordinary Differential Equation (ODE) samplers
[2]. For instance, the Denoising Diffusion Implicit Model (DDIM) [28]
formulates the forward diffusion as a non-Markovian process with a
deterministic generative path, significantly reducing the number of function
evaluations (NFE) required for sampling (from thousands to hundreds) without
sacrificing generation quality. Generalized DDIM (gDDIM) [29] further
optimized DDIM by modifying the parameterization of the scoring network,
reducing NFE to nearly 27. Furthermore, the works in [24] and [30], termed the
Diffusion Probabilistic Model solver (DPM-solver) and the Diffusion
Exponential Integrator Sampler (DEIS), respectively, introduced fast higher-
order solvers, employing exponential integrators that require even less than
20 NFE for comparable generation quality. Moreover, the consistency model [31]
introduced a novel training loss and parameterization, achieving high-quality
generation with merely 1-2 NFE.

We remark that while the aforementioned methods are indirectly related to our
work, our experiments in Section 5.1 and Section 5.2 show that our approach
can be easily integrated into these techniques, further improving diffusion
models’ overall training and sampling efficiency.

##  3 Proposed Multistage Framework

In this section, we introduce our new multistage framework (as illustrated in
Figure 1 (c)). Specifically, we first introduce the multi-stage U-Net
architecture design in Section 3.1, following a new clustering method for
choosing the optimal interval check points to partition the entire timestep
$[0,1]$ into intervals in Section 3.2. We leave the discussion on the
motivation behind our method to Section 4.

###  3.1 Proposed Multi-stage U-Net Architectures

Most existing diffusion models either employ a unified architecture across all
intervals [1, 2, 18] to share features for all timesteps, or they use
completely separate architectures for different timestep intervals [21, 22,
27] to take advantage of benign properties within different intervals. The
details of the advantages and disadvantages of each architecture are discussed
in section 4.

To harness the advantages of both unified and separate architectures employed
in prior studies, we introduce a multistage U-Net architecture, as illustrated
in Figure 1(c). Specifically, we partition the entire timestep $[0,1]$ into
several intervals, e.g., three intervals $[0,t_{1}),[t_{1},t_{2}),[t_{2},1]$
in Figure 1. For the architecture, we introduce:

  * •

One shared encoder across all time intervals. For each timestep interval, we
implement a shared encoder architecture (plotted in blue in Figure 1 (c)),
which is similar to the architecture employed in the original U-Net framework
[32]. Unlike separate architecture, the shared encoder provide shared
information across all timesteps, preventing models from overfitting (see
Section 5.3 for a discussion).

  * •

Separate decoders for different time intervals. Motivated by the multi-head
structure introduced in the Mask Region-based Convolutional Neural Networks
(Mask-RCNN) method [33], we propose to use multiple distinct decoders (plotted
in colors for different intervals in Figure 1(c)), where each decoder is
tailored to a specific timestep interval. The architecture of each decoder
closely resembles the one utilized in [2], with deliberate adjustments made to
the embedding dimensions to optimize performance.

As we observe, the primary difference in the architecture resides within the
decoder structure. Intuitively, we use a decoder with fewer number of
parameters for intervals closer to the noise, because the learning task is
easier. We use a decoder with larger number of parameters for intervals closer
to the image. In Section 4, we conduct a comprehensive analysis to elucidate
the rationale behind the adoption of this multistage architecture, and the
necessity for varying parameters across different intervals.

###  3.2 Optimal Denoiser-based Timestep Clustering

Next, we discuss how we principally choose the interval partition time points
in practice. For simplicity, we focus on the case where we partition the time
$[0,1]$ into three intervals $[0,t_{1}),[t_{1},t_{2}),[t_{2},1]$, and we
develop a timestep clustering method to choose the optimal $t_{1}$ and $t_{2}$
that we introduce in the following. Of course, our method can be generalized
to multi-stage networks with arbitrary interval numbers. However, in practice,
we find that the choice of three intervals strikes a good balance between
effectiveness and complexity; see our ablation study in the appendices.

To partition the time interval, we employ the optimal denoiser established in
Proposition 1, where we can show the following result.

Algorithm 1 Optimal Denoiser based Timestep Clustering

1:Input: Total samples $K$, optimal denoiser function
$\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x},t)$, thresholds $\alpha$, $\eta$,
dataset $p_{\text{data}}$

2:Output: Timesteps $t_{1}$, $t_{2}$

3:$\mathcal{S}_{0}=\mathcal{S}_{1}=\emptyset$

4:for $k\in\\{1,\dots,K\\}$ do

5: $\bm{y}_{k}\sim
p_{\text{data}},\bm{\epsilon}_{k}\sim\mathcal{N}(0,\mathbf{I}),t_{k}\sim[0,1]$

6:
$\mathcal{S}^{k}_{0}\leftarrow\mathcal{D}(\bm{\epsilon}^{*}_{t_{k}},\bm{\epsilon}^{*}_{0},\bm{y}_{k},\bm{\epsilon}_{k}),\mathcal{S}^{k}_{1}\leftarrow\mathcal{D}(\bm{\epsilon}^{*}_{t_{k}},\bm{\epsilon}^{*}_{1},\bm{y}_{k},\bm{\epsilon}_{k})$

7:
$\mathcal{S}_{0}\leftarrow\mathcal{S}_{0}\cup\\{\left(t_{k},\mathcal{S}^{k}_{0}\right)\\}$

8:
$\mathcal{S}_{1}\leftarrow\mathcal{S}_{1}\cup\\{\left(t_{k},\mathcal{S}^{k}_{1}\right)\\}$

9:end for

10:$t_{1}=\operatorname*{\text{arg~{}max}}\limits_{\tau}\left\\{\tau~{}\Big{|}~{}\dfrac{\sum_{(t_{k},\mathcal{S}^{k}_{0})\in\mathcal{S}_{0}}\left[\mathcal{S}^{k}_{0}\cdot\mathds{1}(t_{k}\leq\tau)\right]}{\sum_{(t_{k},\mathcal{S}^{k}_{0})\in\mathcal{S}_{0}}\left[\mathds{1}(t_{k}\leq\tau)\right]}\geq\alpha\right\\}$

11:$t_{2}=\operatorname*{\text{arg~{}min}}\limits_{\tau}\left\\{\tau~{}\Big{|}~{}\dfrac{\sum_{(t_{k},\mathcal{S}^{k}_{1})\in\mathcal{S}_{0}}\left[\mathcal{S}^{k}_{1}\cdot\mathds{1}(t_{k}\geq\tau)\right]}{\sum_{(t_{k},\mathcal{S}^{k}_{1})\in\mathcal{S}_{0}}\left[\mathds{1}(t_{k}\geq\tau)\right]}\geq\alpha\right\\}$

######  Proposition 1.

Suppose we train a diffusion model denoiser function
$\bm{\epsilon}_{\bm{\theta}}(\bm{x},t)$ with parameters $\bm{\theta}$ using
dataset $\left\\{\bm{y}_{i}\in\mathbb{R}^{n}\right\\}_{i=1}^{N}$ by

$\displaystyle\min_{\bm{\theta}}\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};t)=\mathbb{E}_{\bm{x}_{0},\bm{x}_{t},\epsilon}[||\bm{\epsilon}-\bm{\epsilon}_{\bm{\theta}}(\bm{x}_{t},t)||^{2}],$
(5)

where $\bm{x}_{0}\sim
p_{\text{data}}(\bm{x})=\frac{1}{N}\sum_{i=1}^{N}\delta(\bm{x}-\bm{y}_{i})$,
$\bm{\epsilon}\sim\mathcal{N}(0,\textbf{I})$, and $\bm{x}_{t}\sim
p_{t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})$
with perturbation parameters $s_{t}$ and $\sigma_{t}$ defined in eq. 2. Then,
the optimal denoiser at $t$, defined as
$\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x};t)=\operatorname*{\text{arg~{}min}}_{\bm{\epsilon}_{\bm{\theta}}}\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};t)$,
is given by

$\displaystyle\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x};t)=\frac{1}{s_{t}\sigma_{t}}\left[\bm{x}-s_{t}\frac{\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\bm{y}_{i}}{\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})}\right]\>.$

The proof is provided in appendices, and the result can be generalized from
recent work of Karras et al. [18], extending from a specific kernel
$p_{t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};\bm{x}_{0},\sigma_{t}^{2}\textbf{I})$
to encompassing a broader scope of noise perturbation kernels, given by
$p_{t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})$.
For brevity, we simplify the notation of the optimal denoiser
$\bm{\epsilon}^{*}_{\theta}(\bm{x},t)$ in Proposition 1 as
$\bm{\epsilon}^{*}_{t}(\bm{x})$.

To obtain the optimal interval, our rationale is to _homogenize_ the
regression task as much as possible within each individual time interval. To
achieve this goal, given sampled $\bm{x}_{0}$, $\bm{\epsilon}$, we define the
function distance of the optimal denoiser at any given timestep $t_{a}$,
$t_{b}$ as:

$\mathcal{D}(\bm{\epsilon}^{*}_{t_{a}},\bm{\epsilon}^{*}_{t_{b}},\bm{x}_{0},\bm{\epsilon})=\frac{1}{n}\sum_{i=1}^{n}\mathds{1}(|\bm{\epsilon}^{*}_{t_{a}}(\bm{x}_{t_{a}})-\bm{\epsilon}^{*}_{t_{b}}(\bm{x}_{t_{b}})|_{i}\leq\eta)\>,$

where $\mathds{1}(\cdot)$ is the indicator function, $\eta$ is a pre-specified
threshold,
$\bm{x}_{t_{a}}=s_{t_{a}}\bm{x}_{0}+s_{t_{a}}\sigma_{t_{a}}\bm{\epsilon}$, and
$\bm{x}_{t_{b}}=s_{t_{b}}\bm{x}_{0}+s_{t_{b}}\sigma_{t_{b}}\bm{\epsilon}$.
Consequently, we define the functional similarity of the optimal denoiser at
timesteps $t_{a}$ and $t_{b}$ as:

$\mathcal{S}(\bm{\epsilon}^{*}_{t_{a}},\bm{\epsilon}^{*}_{t_{b}})=\mathbb{E}_{\bm{x}_{0}\sim
p_{\text{data}}}\mathbb{E}_{\bm{\epsilon}\sim\mathcal{N}(0,\textbf{I})}[\mathcal{D}(\bm{\epsilon}^{*}_{t_{a}},\bm{\epsilon}^{*}_{t_{b}},\bm{x}_{0},\bm{\epsilon})]\>.$
(6)

Based upon the definition, we design the following optimization problem to
find the largest $t_{1}$ and smallest $t_{2}$ as:

$\displaystyle t_{1}$
$\displaystyle\leftarrow\operatorname*{\text{arg~{}max}}\limits_{\tau}\left\\{\tau\Big{|}\mathbb{E}_{t\sim[0,\tau)}[\mathcal{S}(\bm{\epsilon}^{*}_{t},\bm{\epsilon}^{*}_{0})]\geq\alpha\right\\},$
(7) $\displaystyle t_{2}$
$\displaystyle\leftarrow\operatorname*{\text{arg~{}min}}\limits_{\tau}\left\\{\tau\Big{|}\mathbb{E}_{t\sim[\tau,1]}[\mathcal{S}(\bm{\epsilon}^{*}_{t},\bm{\epsilon}^{*}_{1})]\geq\alpha\right\\},$
(8)

such that the average functional similarity of $\bm{\epsilon}^{*}_{t}$ (resp.
$\bm{\epsilon}^{*}_{t}$) to $\bm{\epsilon}^{*}_{0}$ (resp.
$\bm{\epsilon}^{*}_{1}$) in $[0,t_{1})$ (resp. $[t_{2},1]$) is larger than or
equal to a pre-defined threshold $\alpha$. As the above optimization problems
are intractable, we propose the procedure outlined in Algorithm 1 to obtain an
approximate solution.333Basically, the algorithm samples $K$ sample pairs
$(\bm{y}_{k},\bm{\epsilon}_{k},t_{k}),k\in\\{1,\ldots,K\\}$ to calculate the
distance
$\mathcal{D}(\bm{\epsilon}^{*}_{t_{k}},\bm{\epsilon}^{*}_{0},\bm{y}_{k},\bm{\epsilon}_{k})$
and
$\mathcal{D}(\bm{\epsilon}^{*}_{t_{k}},\bm{\epsilon}^{*}_{1},\bm{y}_{k},\bm{\epsilon}_{k})$
and then based on those distances to solve the optimization problem define in
the line 10 -11 of Algorithm 1 to get $t_{1}$ and $t_{2}$.

##  4 Identification of Key Sources of Inefficiency

Conventional diffusion model architectures, as exemplified by [2, 1, 18],
treat the training of the diffusion model as a unified process across all
timesteps. Recent research endeavors, such as [19, 20, 21, 22], have
highlighted the benefits of recognizing distinctions between different
timesteps and the potential efficiency gains from treating them as separate
tasks during the training process. However, our experimental results
demonstrate that both unified and separate architectures suffer inefficiency
for training diffusion models, where the infficiency comes from (i)
overfitting and (ii) gradient dissimilarity.

###  4.1 Empirical Observations on the Key Sources of Inefficiency

To illustrate the inefficiency in each interval, we isolate the interval by
using a separate architecture from the rest.

(a) Analysis on interval $[0,t_{1})$

(b) Analysis on interval $[t_{2},1]$

Figure 2: Comparison between separate architecture and unified architecture
w.r.t. the image generation quality in different intervals: (a) analysis on
Interval $[0,t_{1})$; and (b) analysis on Interval $[t_{2},1]$. As illustrated
on top of each figure, we _only_ train separate architectures within specific
intervals for the sampling process in both (a) and (b). For the remaining
period of sampling, we use a well-trained diffusion model
$(\bm{\epsilon}_{\bm{\theta}})^{[0,1]}_{4\times 10^{5}}$ to approximate the
ground truth score function. As shown in the above figure of (a), e.g. for the
separate architecture on interval 1, sampling utilizes trained model
$(\bm{\epsilon}_{\bm{\theta}^{\prime}})^{[0,t_{1})}_{i}$ for interval 0 and
well-trained model  $(\bm{\epsilon}_{\bm{\theta}})^{[0,1]}_{4\times 10^{5}}$
for interval 1 and 2. Notably, for both
$(\bm{\epsilon}_{\bm{\theta}})^{[0,1]}_{i}$ and
$(\bm{\epsilon}_{\bm{\theta}})^{[0,1]}_{4\times 10^{5}}$ , we are using the
model with 108M parameters. For separate architecture, the number in the
parentheses represents the number of parameters of the model
$(\bm{\epsilon}_{\bm{\theta}^{\prime}})^{[a,b]}_{i}$ . For example, for
separate architecture (169M) in (a), the model
$(\bm{\epsilon}_{\bm{\theta}^{\prime}})^{[0,t_{1})}_{i}$ has 169M parameters
for $\bm{\theta}^{\prime}$. The bottom figures in (a-b) illustrate the FID of
the generation from each architecture under different training iterations.

#### Experiment setup.

In our experiments, we consider three-stage training and divide the time steps
into three intervals: $[0,t_{1}),[t_{1},t_{2}),[t_{2},1]$.444Details for
interval clustering can be found in Section 5. Let
$(\bm{\epsilon}_{\bm{\theta}})^{[a,b]}_{i},0\leq a<b\leq 1$ denote a U-Net
architecture with parameter $\bm{\theta}$ trained with $i$ iterations and fed
with data pairs $(\bm{x}_{t},t)$, where $t\in[a,b]$. We then trained models
using two different strategies: a unified architecture with 108M network
parameters for all intervals, i.e.,
$(\bm{\epsilon}_{\bm{\theta}})^{[0,1]}_{i}$, and separate architectures with
varying network parameters (e.g., 47M, 108M, 169M) for each interval; e.g.,
$(\bm{\epsilon}_{\bm{\theta}})^{[0,t_{1})}_{i}$ for the interval $[0,t_{1})$,
etc. It is worth noting that, apart from the differences in network
parameters, we utilize the same network architecture (e.g., U-Net) for both
the unified and the separate approaches. We assessed the training progress of
each model by evaluating image generation quality at different training
iterations. Notably, because some of the models are only trained on one
interval, we need to provide a ground truth score for the other intervals. In
Figure 2, the sampling process is shown above and the experimental results are
shown below.

#### Inefficiency in unified architectures.

From Figure 2, we observe the following:

  * •

Overparameterization and underfitting emerge simultaneously for unified
architectures. In Figure 2(a), we observe that increasing the number of
parameters in Interval 0 can improve the image generation quality (as
indicated by a lower FID score). In contrast, Figure 2(b) reveals that
increasing the number of parameters in Interval 2 has minimal impact on the
quality of image generation. This implies that using a unified architecture
will result in underfitting in Interval 0 and overparameterization in Interval
2. The current unified architecture’s parameter redundancy leaves significant
room for improving its efficiency. To optimize the computational usage, we
should allocate more parameters to Interval 0 while allocate fewer parameters
to Interval 2.

  * •

Gradient dissimilarity hinders training for unified architecture. Quantitative
results from [20] demonstrate dissimilarity in gradients emerges among
different intervals. This can also be observed from our results based upon
both Figure 2(a) and Figure 2(b). For the unified and separate architectures
using the same number of parameters (108M), separate architecture achieves a
significantly lower FID with the same training iterations, implying that
dissimilar gradients among intervals may hinder training when using a unified
architecture. Here, the only difference between training separate and unified
architectures is that the batch gradient for unified architecture is
calculated based on all timesteps while the batch gradient for separate
architecture is calculated only from a specific interval.

#### Inefficiency in existing separate architectures.

Although separated architecture [19, 21, 22] better allocates computational
resources for each interval, it suffers from overfitting. This can be
illustrated based upon training separate architectures (169M) and (108M) in
Interval 0 shown in Figure 2(a), where increasing the number of parameters
will lead to overfitting when we increase the number of training iterations
beyond achieving the best FID. This also happens in Interval 2, when we
compare all separate architectures in Figure 2(b). In comparison, the unified
networks with 108M parameters are less prone to overfit for both Interval 0
and Interval 2. This suggests that we can reduce overfitting by training
shared weights across different intervals together.

###  4.2 Tackling the Inefficiency via Multistage U-Net Architectures

In a unified architecture applied across all timesteps, there is often a dual
challenge: requirements for more parameters (169M) in the interval $[0,t_{1})$
but fewer parameters (47M) in the interval $[t_{2},1]$. This issue is
compounded by the gradient dissimilarity across different timesteps, which can
impede effective training. Alternatively, employing separate architectures for
different intervals might lead to overfitting and a lack of robust early
stopping mechanisms. To address these challenges, our proposed multistage
architecture in Section 3 combines shared parameters to reduce overfitting
with interval-specific parameters to mitigate the impact of gradient
dissimilarity. This tailored approach for each interval ensures improved
adaptability. Furthermore, we conduct an in-depth ablation study in Section
5.3 to showcase the effectiveness of our multi-stage architecture over the
existing models.

##  5 Experiments

For this section, we start with providing the basic experimental setups.
Followed by this, in Section 5.1, we demonstrate the improved image generation
quality of our multistage architecture over the state-of-the art models, under
comparable training and sampling computations. On the other hand, in Section
5.2, we demonstrate the improved training and sampling efficiency of our
methods over existing methods, when compared with similar image generation
quality. Finally, in Section 5.3 and Section 5.4, we conduct comprehensive
ablation studies on timestep clustering methods and multistage architectures,
illustrating the superiority of the choices in our method.

#### Settings of the timestep clustering method.

We implement a Variance Preserving (VP) [1, 18] perturbation kernel for
clustering the time interval. In particular, we use
$s_{t}=\sqrt{e^{\frac{1}{2}\beta_{d}t^{2}+\beta_{\text{min}}-1}},\sigma_{t}=1/\sqrt{e^{\frac{1}{2}\beta_{d}t^{2}+\beta_{\text{min}}}},t\in[\epsilon_{t},1],\beta_{d}=19.9,\beta_{\text{min}}=0.1$,
and $,\epsilon_{t}=10^{-3}$. To obtain $t_{1}$ and $t_{2}$, we utilize the
CIFAR-10 [23] with dataset size $N=5\times 10^{4}$. Specifically, Algorithm 1
is configured with $\eta=\frac{2}{256},\alpha=0.9,$ and $K=5\times 10^{4}$.
This choice of $\eta$ ensures the similarity of measurements in the RGB space.
We divide $t\in[0,1]$ into $10^{4}$ discrete time steps. Consequently, we
employ a grid search to determine optimal values for $t_{1}$ and $t_{2}$ using
the procedure outlined in Algorithm 1. For the CelebA dataset [25], we utilize
the same $t_{1}$ and $t_{2}$. And for Variance Exploding (VE) perturbation
kernel, we calculate an equivalent $\sigma_{1}$ and $\sigma_{2}$.555VE
utilizes $\sigma$ instead of $t$ to represent the timestep. We choose
$\sigma_{1}$ and $\sigma_{2}$ such that
$\text{SNR}_{\text{VE}}(\sigma_{1})=\text{SNR}_{\text{VP}}(t_{1})$ and
$\text{SNR}_{\text{VE}}(\sigma_{2})=\text{SNR}_{\text{VP}}(t_{2})$. Further
details are provided in the appendices.  Results of Section 5.4 verify that
these choices of $t_{1}$ and $t_{2}$ return the best performance when compared
to other clustering method.

#### Multistage architectures.

Our multistage architecture, inspired by the U-Net model [32] used in DDPM++
[1, 18, 2], is modified for interval-specific channel dimensions. The proposed
architecture is adopted to three diffusion models: DPM-Solver [24],
Elucidating Diffusion Model (EDM) [18], and Latent Diffusion Model (LDM) [8].
In particular, for the cases of DPM-Solver and EDM, the encoder’s channel
dimensions are standardized at 128, while the decoders are configured with
192, 128, and 16 channels for intervals $[0,t_{1})$, $[t_{1},t_{2})$, and
$[t_{2},1]$, respectively. In the LDM case, we use 224 channels across the
encoder for all intervals whereas the decoders are configured with 256, 192,
and 128 channels for the respective intervals. To decide the specific number
of parameters for each decoder, we apply ablation studies in the appendices.

#### Datasets, Evaluation Metrics, & Baselines.

We use CIFAR-10 (32 $\times$ 32), CelebA (32 $\times$ 32), and CelebA (256
$\times$ 256) datasets for our experiments. To evaluate the performance of our
multistage diffusion model in terms of the generation quality, we use the
standard Fréchet Inception Distance (FID) metric [34], and the required Number
of Function Evaluations (NFEs). We assess the sampling efficiency using giga-
floating point operations (GFLOPs) per function evaluation. For both separate
architecture and our multistage architecture, equivalent GFLOPs are computed
as a weighted summation (based on the ODE Sampler steps within each interval)
of GFLOPs for each interval. Training efficiency is evaluated using total
training iterations multiplied by the floating point operations per function
evaluation,666Here we simplify it by ignoring the FLOPs for backward
propagation, which is approximately proportional to FLOPs of forward
evaluation. measured by peta-floating point operations (PFLOPs). We choose
DDPM [1], Score SDE [2], Poisson Flow Generative Models (PFGM) [35], DDIM
[28], gDDIM [29], DEIS [30], DPM-solver [24], and EDM [18] as the baseline.

###  5.1 Comparison of Image Generation Quality

In this subsection, we demonstrate the effectiveness of our approach by
comparing the image generation quality (measured by FID) with comparable
training and sampling computations (measured by NFE). Specifically, Table 1
presents FID scores to measure the sampling quality, and NFEs to measure the
number of sampling iterations required using the CIFAR-10 dataset. Our method
is compared to 8 baselines. As observed, our multistage DPM-Solver outperforms
DPM-Solver in terms of the reported FID values while requiring similar
training iterations (both are $4.5\times 10^{5}$), and model GFLOPs (18.65 for
multistage DPM-Solver versus 17.65 for DPM-Solver). A similar observation
holds when we compare our multistage EDM and the vanilla EDM, where we reduce
FID from 2.05 to 1.96 by using the multi-stage architecture. Remarkably,
utilizing only 20 NFE, our Multistage DPM-Solver returns the same FID score as
the one reported for the PFGM method, which requires 147 NFEs. These results
also highlight the adaptability of our framework to higher-order ODE solvers;
see the 8th and last row of Table 1.

METHOD | NFE ($\downarrow$) | FID ($\downarrow$)
---|---|---
DDPM | 1000 | 3.17
Score SDE | 2000 | 2.20
PFGM | 147 | 2.35
DDIM | 100 | 4.16
gDDIM | 20 | 2.97
DEIS | 20 | 2.86
DPM-solver | 20 | 2.73
Multistage DPM-solver (Ours) | 20 | 2.35
EDM | 35 | 2.05
Multistage EDM (Ours) | 35 | 1.96
Table 1: Sampling quality on CIFAR-10 dataset.

###  5.2 Comparison of Training & Sampling Efficiency

Dataset/Method | Training Iterations ($\downarrow$) | GFLOPs ($\downarrow$) | Total Training PFLOPs ($\downarrow$) | FID ($\downarrow$)
---|---|---|---|---
CIFAR-10 $32\times 32$ |  |  |  |
DPM-Solver [18] | $4.5\times 10^{5}$ | 17.65 | 7.94 | 2.73
Multistage DPM-Solver (Ours) | $\textbf{2.5}\times\textbf{10}^{5}\textbf{(56\%)}$ | 18.65 (106%) | 4.66 (59%) | 2.71
CelebA $32\times 32$ |  |  |  |
EDM [18] | $5.7\times 10^{5}$ | 17.65 | 10.06 | 1.55
Multistage EDM (Ours) | $\textbf{4.3}\times\textbf{10}^{5}\textbf{(75\%)}$ | 19.25 (109%) | 8.28 (82%) | 1.44
CelebA $256\times 256$ |  |  |  |
LDM [8] | $4.9\times 10^{5}$ | 88.39 | 43.31 | 8.29
Multistage LDM (Ours) | $\textbf{1.7}\times\textbf{10}^{5}\textbf{(35\%)}$ | 76.19 (86%) | 12.95 (30.0%) | 8.38

Table 2: Training and sampling efficiency on more datasets.

In this subsection, we further demonstrate the superiority of our method by
comparing the training and sampling efficiency under comparable image
generation quality. Specifically, in Table 2, we present the number of
training iterations, GFLOPs, and total training PFLOPs of our approach, DPM-
solver, EDM, and LDM using CIFAR-10 and CelebA datasets. Using the CIFAR-10
dataset, our multistage DPM-solver achieves similar FID scores (2.71 vs 2.73)
while requiring nearly half the training iterations when compared to the
vanilla DPM solver. For the case of EDM (resp. LDM), our approach returns an
FID score of 1.44 (resp. 8.29), requiring $1.4\times 10^{5}$ ($3.2\times
10^{5}$) less iterations when compared to vanilla DPM (resp. LDM). For the
cases of DPM and EDM, we can achieve substantial reduction of training
iterations, which is demonstrated by a marginal increase in the number of
GFLOPs (row 3 vs. row 4 and row 5 vs. row 6). For the LDM case,we also achieve
a significant reduction of training iterations, demonstrated by a reduced
number of GFLOPs (row 8 vs. row 9). These promising results highlight the
significantly improved computational efficiency achieved by using the proposed
multistage framework.

###  5.3 Comparison of Different Architectures

In Section 4, we highlighted the limitations of both unified and separate
diffusion model architectures in terms of training efficiency (see Figure 2).
In this part, we further illustrate these limitations through extensive
experiments as shown in Table 3. Here, we use the U-Net architecture, trained
on the CIFAR-10 dataset, and utilize the DPM-Solver for sampling. For the
unified case, we use a single U-Net model with 128 channels. For the separate
case, three distinct U-Nets are used by which, for each interval, we use a
128-channel decoder. For improved performance of the separate architecture, we
implement two techniques: early stopping (ES) and tailored parameters (TP) to
tackle the overfitting and parameter inefficiency discussed in Section 4.
Under ES, the criteria is to stop training prior to overfitting. For the
tailored parameters, the three U-Nets are configured with 192, 128, and 16
channels decoders for Intervals 0, 1, and 2, respectively.

Method | GFLOPs | FID ($\downarrow$)
---|---|---
Unified | 17.65 | 2.73
Separate | 17.65 | 2.87
Separate (+ ES) | 17.65 | 2.80
Separate (+ TP) | 18.65 | 2.68
Separate (+ ES, TP) | 18.65 | 2.52
Multistage (Ours) | 18.65 | 2.35
Table 3: Sampling quality on CIFAR-10 dataset for different diffusion model
architectures. Here, ES denotes early stopping, and TP denotes tailored
parameters.

Our comparison and analysis in Table 3 reveals notable insights of our network
design. Comparisons between the 2nd and 3rd rows (and between the 4th and 5th
rows) on the separate architectures indicate that early stopping effectively
mitigates overfitting and enhances generation quality. When comparing the 2nd
and 4th rows (as well as the 3rd and 5th rows) on the separate architectures,
we observe that optimizing parameter usage can achieve a significant decrease
in FID under comparable GFLOPs (18.65 vs. 17.65). Most importantly, our
multistage architecture, as shown in the 6th row, benefits from both unified
and separate architectures, achieving the best FID (2.35, compared to 2.73 and
2.52). Comparing the 2nd row and 4th row, the shared encoder not only prevents
overfitting but also improves the convergence of the diffusion model as
reported by the FID scores.

###  5.4 Comparison of Timestep Clustering Methods

Clustering Method | $t_{1}$ | $t_{2}$ | FID ($\downarrow$)
---|---|---|---
Timestep [20, 21] | 0.330 | 0.670 | 3.12
SNR [20] | 0.348 | 0.709 | 2.72
Gradient [20] | 0.360 | 0.660 | 2.75
Optimal Denoiser (Ours) | 0.442 | 0.631 | 2.35
Table 4: Sampling quality on CIFAR-10 dataset for different clustering
methods.

As preciously stated in Section 2, various timestep clustering methods are
proposed including timestep-based, SNR-based, and gradient-based clustering
approaches [20, 21]. In this subsection, we conduct an experiment to
demonstrate the superiority of our clustering method compared to previous
arts. Specifically, we apply the clustering methods in [20, 21] to partition
the interval along with our proposed multistage UNet architecture. The
computed intervals are shown in the Table 4. We use the multistage DPM-Solver
with these different intervals trained on the CIFAR-10 dataset. As observed,
our optimal denoiser-based clustering method achieves the highest FID score,
consistently outperforming all other clustering methods.

##  6 Conclusion

In this study, we significantly enhance the training and sampling efficiency
of diffusion models through a novel multi-stage framework. This method divides
the time interval into several stages, using a specialized multi-decoder U-net
architecture that combines time-specific models with a common encoder for all
stages. Based on our practical findings, this multi-stage approach utilizes
unique parameters for each timestep, along with shared parameters across all
timesteps, grouped according to the most effective denoiser. We conducted
thorough numerical experiments with three leading-edge diffusion models, such
as large-scale latent diffusion models, which confirmed the effectiveness of
our strategy.

In future research, it would be interesting to expand our multi-stage approach
beyond the unconditional diffusion models, such as conditional diffusion
models and diffusion models for solving inverse problems. Our experiments, as
detailed in Section 5.2, demonstrate that training latent diffusion models
within our multi-stage framework requires only 30% of the computational effort
needed for training standard latent diffusion models on the CelebA dataset.
Thus, employing a multi-stage strategy could significantly reduce the
computational demands for training large-scale stable diffusion models, such
as those described in [8], which typically requires significant computational
power.

## Acknowledgment

HJZ, YFL, and QQ acknowledge support from NSF CAREER CCF-2143904, NSF
CCF-2212066, NSF CCF-2212326, NSF IIS 2312842, ONR N00014-22-1-2529, an AWS AI
Award, and a gift grant from KLA. QQ also acknowledges the support from MICDE
Catalyst Grant. IA and SR acknowledge the support from the NSF grants
CCF-2212065 and CCF-2212066. Results presented in this paper were obtained
using CloudBank, which is supported by the NSF under Award #1925001, and the
authors acknowledge efficient cloud management framework SkyPilot [36] for
computing. The authors acknowledge valuable discussions with Prof. Jeffrey
Fessler (U. Michigan), Prof. Liyue Shen (U. Michigan), and Prof. Rongrong Wang
(MSU).

## References

  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in neural information processing systems_ , 33:6840–6851, 2020.
  * Song et al. [2020a] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020a.
  * Dhariwal and Nichol [2021] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in neural information processing systems_ , 34:8780–8794, 2021.
  * Ho and Salimans [2022] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  _arXiv preprint arXiv:2207.12598_ , 2022.
  * Su et al. [2022] Xuan Su, Jiaming Song, Chenlin Meng, and Stefano Ermon.  Dual diffusion implicit bridges for image-to-image translation.  _arXiv preprint arXiv:2203.08382_ , 2022.
  * Saharia et al. [2022] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi.  Palette: Image-to-image diffusion models.  In _ACM SIGGRAPH 2022 Conference Proceedings_ , pages 1–10, 2022.
  * Zhao et al. [2022] Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu.  Egsde: Unpaired image-to-image translation via energy-guided stochastic differential equations.  _Advances in Neural Information Processing Systems_ , 35:3609–3623, 2022.
  * Rombach et al. [2022] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 10684–10695, 2022.
  * Ramesh et al. [2021] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In _International Conference on Machine Learning_ , pages 8821–8831. PMLR, 2021.
  * Nichol et al. [2021] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen.  Glide: Towards photorealistic image generation and editing with text-guided diffusion models.  _arXiv preprint arXiv:2112.10741_ , 2021.
  * Song et al. [2023a] Bowen Song, Soo Min Kwon, Zecheng Zhang, Xinyu Hu, Qing Qu, and Liyue Shen.  Solving inverse problems with latent diffusion models via hard data consistency.  _arXiv preprint arXiv:2307.08123_ , 2023a.
  * Chung et al. [2022] Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye.  Diffusion posterior sampling for general noisy inverse problems.  _arXiv preprint arXiv:2209.14687_ , 2022.
  * Song et al. [2021] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon.  Solving inverse problems in medical imaging with score-based generative models.  _arXiv preprint arXiv:2111.08005_ , 2021.
  * Alkhouri et al. [2023] Ismail Alkhouri, Shijun Liang, Rongrong Wang, Qing Qu, and Saiprasad Ravishankar.  Diffusion-based adversarial purification for robust deep mri reconstruction.  _arXiv preprint arXiv:2309.05794_ , 2023.
  * Ho et al. [2022] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al.  Imagen video: High definition video generation with diffusion models.  _arXiv preprint arXiv:2210.02303_ , 2022.
  * Harvey et al. [2022] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood.  Flexible diffusion modeling of long videos.  _Advances in Neural Information Processing Systems_ , 35:27953–27965, 2022.
  * Zhang et al. [2023] Huijie Zhang, Jinfan Zhou, Yifu Lu, Minzhe Guo, Liyue Shen, and Qing Qu.  The emergence of reproducibility and consistency in diffusion models.  _arXiv preprint arXiv:2310.05264_ , 2023.
  * Karras et al. [2022] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models.  _arXiv preprint arXiv:2206.00364_ , 2022.
  * Choi et al. [2022] Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, and Sungroh Yoon.  Perception prioritized training of diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 11472–11481, 2022.
  * Go et al. [2023a] Hyojun Go, JinYoung Kim, Yunsung Lee, Seunghyun Lee, Shinhyeok Oh, Hyeongdon Moon, and Seungtaek Choi.  Addressing negative transfer in diffusion models.  _arXiv preprint arXiv:2306.00354_ , 2023a.
  * Lee et al. [2023] Yunsung Lee, Jin-Young Kim, Hyojun Go, Myeongho Jeong, Shinhyeok Oh, and Seungtaek Choi.  Multi-architecture multi-expert diffusion models.  _arXiv preprint arXiv:2306.04990_ , 2023.
  * Go et al. [2023b] Hyojun Go, Yunsung Lee, Jin-Young Kim, Seunghyun Lee, Myeongho Jeong, Hyun Seung Lee, and Seungtaek Choi.  Towards practical plug-and-play diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 1962–1971, 2023b.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Lu et al. [2022] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu.  Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.  _arXiv preprint arXiv:2206.00927_ , 2022.
  * Liu et al. [2015] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Deep learning face attributes in the wild.  In _Proceedings of International Conference on Computer Vision (ICCV)_ , December 2015.
  * Vincent [2011] Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011.
  * Balaji et al. [2022] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, Bryan Catanzaro, et al.  ediffi: Text-to-image diffusion models with an ensemble of expert denoisers.  _arXiv preprint arXiv:2211.01324_ , 2022.
  * Song et al. [2020b] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  _arXiv preprint arXiv:2010.02502_ , 2020b.
  * Zhang et al. [2022] Qinsheng Zhang, Molei Tao, and Yongxin Chen.  gddim: Generalized denoising diffusion implicit models.  _arXiv preprint arXiv:2206.05564_ , 2022.
  * Zhang and Chen [2022] Qinsheng Zhang and Yongxin Chen.  Fast sampling of diffusion models with exponential integrator.  _arXiv preprint arXiv:2204.13902_ , 2022.
  * Song et al. [2023b] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever.  Consistency models.  2023b.
  * Ronneberger et al. [2015] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , pages 234–241. Springer, 2015.
  * He et al. [2017] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick.  Mask r-cnn.  In _Proceedings of the IEEE international conference on computer vision_ , pages 2961–2969, 2017.
  * Heusel et al. [2017] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  _Advances in neural information processing systems_ , 30, 2017.
  * Xu et al. [2022] Yilun Xu, Ziming Liu, Max Tegmark, and Tommi Jaakkola.  Poisson flow generative models.  _Advances in Neural Information Processing Systems_ , 35:16782–16795, 2022.
  * Yang et al. [2023] Zongheng Yang, Zhanghao Wu, Michael Luo, Wei-Lin Chiang, Romil Bhardwaj, Woosuk Kwon, Siyuan Zhuang, Frank Sifei Luan, Gautam Mittal, Scott Shenker, et al.  $\\{$SkyPilot$\\}$: An intercloud broker for sky computing.  In _20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23)_ , pages 437–455, 2023.
  * Obukhov et al. [2020] Anton Obukhov, Maximilian Seitzer, Po-Wei Wu, Semen Zhydenko, Jonathan Kyl, and Elvis Yu-Jing Lin.  High-fidelity performance metrics for generative models in pytorch, 2020.  URL https://github.com/toshas/torch-fidelity.  Version: 0.2.0, DOI: 10.5281/zenodo.3786540.

We include proof of Proposition 1 in Appendix A, extension of Algorithm
Algorithm 1 for more intervals in Appendix B, provide generation visualization
in Section C.1, conduct ablation study on network parameters in Section C.2
and number of intervals in Section C.3, and describe experimental settings for
different interval splitting from Section 5.4 in Section C.4.

##  Appendix A Proof of Proposition 1

###  A.1 Background

This section presents the proof of Proposition 1 in Section 3. Our proof
partially follows [18]. As background, let
$p_{t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})$
be the perturbation kernel of the diffusion model, which is a continuous
process gradually adding noise from original image $\bm{x}_{0}$ to
$\bm{x}_{t}$ across $t\in[0,1]$. Both $s_{t}=s(t)$ and $\sigma_{t}=\sigma(t)$
are used interchangeably to denote scalar functions of $t$ that control the
perturbation kernel. It has been shown in [2] that this perturbation kernel is
equivalent to a stochastic differential equation
$\text{d}\bm{x}=\bm{x}_{t}f(t)\text{d}t+g(t)\text{d}\bm{\omega}_{t}$, where
$f(t),g(t)$ are scalar functions of $t$. The relations of $f(t),g(t)$ and
$s_{t},\sigma_{t}$ are:

$s_{t}=\text{exp}(\int_{0}^{t}f(\xi)\text{d}\xi),\ \ \text{and}\
\sigma_{t}=\sqrt{\int_{0}^{t}\frac{g^{2}(\xi)}{s^{2}(\xi)}\text{d}\xi}\>.$ (9)

###  A.2 Proof

Given dataset $\left\\{\bm{y}_{i}\right\\}_{i=1}^{N}$ with $N$ images, we
approximate the original dataset distribution as $p_{\text{data}}$ as multi-
Dirac distribution,
$p_{\text{data}}(\bm{x})=\frac{1}{N}\sum_{i=1}^{N}\delta(\bm{x}-\bm{y}_{i})$.
Then, the distribution of the perturbed image $\bm{x}$ at random timestep $t$
can be calculated as:

$\displaystyle p_{t}(\bm{x})$
$\displaystyle=\int_{\mathbb{R}^{d}}p_{t}(\bm{x}|\bm{x}_{0})p_{\text{data}}(\bm{x}_{0})\text{d}\bm{x}_{0}$
(10)
$\displaystyle=\int_{\mathbb{R}^{d}}p_{\text{data}}(\bm{x}_{0})\mathcal{N}(\bm{x};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\text{d}\bm{x}_{0}$
(11)
$\displaystyle=\int_{\mathbb{R}^{d}}\frac{1}{N}\sum_{i=1}^{N}\delta(\bm{x}_{0}-\bm{y}_{i})\mathcal{N}(\bm{x};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\text{d}\bm{x}_{0}$
(12)
$\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\int_{\mathbb{R}^{d}}\delta(\bm{x}_{0}-\bm{y}_{i})\mathcal{N}(\bm{x};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\text{d}\bm{x}_{0}$
(13)
$\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\>.$
(14)

Let the noise prediction loss, which is generally used across various
diffusion models, be

$\displaystyle\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};t)$
$\displaystyle=\mathbb{E}_{\bm{x}\sim
p_{t}(\bm{x})}[|\bm{\epsilon}-\bm{\epsilon}_{\bm{\theta}}(\bm{x},t)||^{2}]=\int_{\mathbb{R}_{d}}\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})||\bm{\epsilon}-\bm{\epsilon}_{\bm{\theta}}(\bm{x},t)||^{2}\text{d}\bm{x}\>,$
(15)

where $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\textbf{I})$ is defined follow the
perturbation kernel
$p_{t}(\bm{x}|\bm{x}_{0})=\mathcal{N}(\bm{x};s_{t}\bm{x}_{0},s_{t}^{2}\sigma_{t}^{2}\textbf{I})$,
and

$\bm{x}=s_{t}\bm{y}_{i}+s_{t}\sigma_{t}\bm{\epsilon}\quad\Rightarrow\quad\bm{\epsilon}=\frac{\bm{x}-s_{t}\bm{y}_{i}}{s_{t}\sigma_{t}}\>.$
(16)

Here, $\bm{\epsilon}_{\bm{\theta}}$ is a "denoiser" network for learning the
noise $\bm{\epsilon}$. Plugging (16) into (15), the loss can be
reparameterized as:

$\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};t)=\int_{\mathbb{R}_{d}}\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)\text{d}\bm{x}\>,$
(17)

where

$\displaystyle\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)=\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})||\bm{\epsilon}_{\bm{\theta}}(\bm{x},t)-\frac{\bm{x}-s_{t}\bm{y}_{i}}{s_{t}\sigma_{t}}||^{2}\>.$

Eq.(17) means that we can minimize
$\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};t)$ by minimizing
$\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)$ for each $\bm{x}$. As
such, we find the "optimal denoiser" $\bm{\epsilon}^{*}_{\bm{\theta}}$ that
minimize the $\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)$, for every
given $\bm{x}$ and $t$, as:

$\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x};t)=\text{arg}\
\text{min}_{\bm{\epsilon}_{\bm{\theta}}(\bm{x};t)}\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)\>.$
(18)

The above equation is a convex optimization problem by which the solution can
be obtained by setting the gradient of
$\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)$ w.r.t
$\bm{\epsilon}_{\bm{\theta}}(\bm{x};t)$ to zero. This means

$\displaystyle\nabla_{\bm{\epsilon}_{\bm{\theta}}(\bm{x};t)}[\mathcal{L}(\bm{\epsilon}_{\bm{\theta}};\bm{x},t)]=0$
(19) $\displaystyle\Rightarrow$
$\displaystyle\frac{1}{N}\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})[\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x};t)-\frac{\bm{x}-s_{t}\bm{y}_{i}}{s_{t}\sigma_{t}}]=0$
(20) $\displaystyle\Rightarrow$
$\displaystyle\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x};t)=\frac{1}{s_{t}\sigma_{t}}[\bm{x}-s_{t}\frac{\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})\bm{y}_{i}}{\sum_{i=1}^{N}\mathcal{N}(\bm{x};s_{t}\bm{y}_{i},s_{t}^{2}\sigma_{t}^{2}\textbf{I})}]$
(21)

##  Appendix B Generalized Algorithm 1

Algorithm 2 Optimal Denoiser based Timestep Clustering for More Intervals

1:Input: Total samples $K$, optimal denoiser function
$\bm{\epsilon}^{*}_{\bm{\theta}}(\bm{x},t)$, interval number $n$, dataset
$p_{\text{data}}$

2:Output: Timesteps $t_{1}$, $t_{2}$, $\ldots$, $t_{n-1}$

3:$\mathcal{S}\leftarrow\emptyset$

4:for $k\in\\{1,\dots,K\\}$ do

5: $\bm{y}_{k}\sim
p_{\text{data}},\bm{\epsilon}_{k}\sim\mathcal{N}(0,\mathbf{I}),t_{a_{k}}\sim[0,1],t_{b_{k}}\sim[0,1]$

6:
$\mathcal{S}^{k}\leftarrow\mathcal{D}(\bm{\epsilon}^{*}_{t_{a_{k}}},\bm{\epsilon}^{*}_{t_{b_{k}}},\bm{y}_{k},\bm{\epsilon}_{k})$

7:
$\mathcal{S}\leftarrow\mathcal{S}\cup\\{\left(t_{a_{k}},t_{b_{k}},\mathcal{S}^{k}\right)\\}$

8:end for

9:$t_{1},\cdots,t_{i-1}\leftarrow\operatorname*{\text{arg~{}min}}_{t_{1},\cdots,t_{i-1}}\quad\sum\limits_{k=1}^{K}\sum\limits_{i=0}^{n-1}\sum\limits_{t_{i}}^{t_{i+1}}\mathcal{S}^{k}\mathds{1}(t_{a_{k}},t_{b_{k}}\in[t_{i},t_{i+1})),$

10:$\quad\quad\quad\quad\quad\quad\quad\quad\text{s.t.}\quad
0<t_{1}<t_{2}<\cdots<t_{n-1}<1$

In this section, we design a procedure to extend Algorithm 1 for case of
$n\neq 3$. Specifically, we partition the timesteps into $n$ intervals
$[0,t_{1}),\ldots,[t_{n-2},t_{n-1}),[t_{n-1},1]$. Intuitively, the intervals
need to minimize the summation of the total functional distance within each
partition. This corresponds to the following optimization problem.

$\min_{t_{1},\cdots,t_{n+1}}~{}\sum\limits_{k=0}^{n-1}\int\limits_{t_{k}}^{t_{k+1}}\int\limits_{t_{k}}^{t_{k+1}}\mathcal{S}(\bm{\epsilon}^{*}_{t_{a}},\bm{\epsilon}^{*}_{t_{b}})\text{d}t_{a}\text{d}t_{b}\quad\textrm{s.t.}\quad
t_{0}=0,~{}t_{n}=1,~{}0<t_{1}<t_{2}<\cdots<t_{n-1}<1\>.$ (22)

In order to solve (22) numerically, we propose Algorithm 2. A solution of the
program in step 9 is obtained by discretizeing the time index into
$t\in\\{0.001,0.026,\dots,0.951,0.976,1\\}$. We use the CIFAR10 dataset with
$K=50000$ to run the experiment. The following table shows the results of
different stage numbers $n$ and different split intervals:

##  Appendix C Additional Experiments

Figure 3: Sample generations from Multistage LDM (CelebA $256\times 256$) and
Multistage DPM-Solver (CIFAR-10 $32\times 32$).

###  C.1 Generation Results

The generation samples displayed in Figure 3, from the CelebA and CIFAR-10
datasets, demonstrate that our Multistage strategy is capable of producing
high-quality results. Notably, it achieves this with reduced training and
computational requirements in comparison to baseline methods.

###  C.2 Ablation Study on the Network Parameters

In this subsection, we perform a series of experiments to determine the best
number of parameters within the attempted settings for each stage of our
model. For evaluation, we use the best FID scores from all checkpoints. These
evaluations are conducted using the DPM-Solver with $20$ Neural Function
Evaluations (NFE) on $50,000$ samples. Here, the experimental settings are
similar to those used in Section 5.

For the Multistage DPM-solver on CIFAR-10 dataset. The models are trained
until they achieve their best FID scores. For calculating FID, we follow the
methodology described by [28], using the tensorflow_gan
777https://github.com/tensorflow/gan.

The design of the corresponding parameters for encoders and decoders across
different architectures, along with their best FID scores, is detailed in
Table 5. We explore various architectures, aiming to maintain a consistent
total number of parameters while varying the ratio of shared parameters
(encoder parameters) to the total number of parameters. Our ablation study
illustrates that Architecture 1 (our chosen model) indeed exhibits the best
FID performance.

Architecture | Current/Total | Parameters | GFLOPs ($\downarrow$) | FID ($\downarrow$)
---|---|---|---|---
1 | 1/3 | 169M | 29.95 | 2.35
2/3 | 108M | 17.65
3/3 | 47M | 6.31
Shared | 43M | 5.75
Total | 237M | 18.65
2 | 1/3 | 168M | 27.57 | 2.9
2/3 | 138M | 21.73
3/3 | 72M | 9.65
Shared | 68M | 8.98
Total | 242M | 19.7
3 | 1/3 | 189M | 29.43 | 3.1
2/3 | 160M | 23.91
3/3 | 103M | 13.7
Shared | 98M | 12.93
Total | 256M | 22.45
4 | 1/3 | 174M | 26.54 | 3.34
2/3 | 126M | 17.62
3/3 | 103M | 13.7
Shared | 98M | 12.93
Total | 207M | 22.45
Table 5: Ablation study of Multistage DPM-solver on CIFAR-10 dataset.
Current/Total: the number of the stage over the total number of stages.
Shared: the encoder. Total: combined the encoder and decoders. Total GFLOPS:
averaged GFLOPS of all stages weighted by NFE assigned to each stage during
sampling of DPM-Solver.

For Latent Diffusion multistage models on CelebA-HQ, models are trained for
500k iterations, and the best checkpoints are used for evaluation. We compute
the FID following [8] using the torch-fidelity package [37]. The parameter
design and best FID are shown in Table 6. Our ablation study shows the
superiority of utilizing Architecture 4 (our chosen model) in terms of the
reported FID values.

Architecture | Current/Total | Parameters | GFLOPs ($\downarrow$) | FID ($\downarrow$)
---|---|---|---|---
1 | 1/3 | 238M | 80.03 | 9.37
2/3 | 189M | 64.95
3/3 | 161M | 51.87
Shared | 79M | 16.69
Total | 428M | 65.75
2 | 1/3 | 206M | 72.61 | 9.73
2/3 | 158M | 57.85
3/3 | 108M | 34.38
Shared | 55M | 11.6
Total | 361M | 54.37
3 | 1/3 | 274M | 88.39 | 8.43
2/3 | 224M | 72.97
3/3 | 170M | 48.17
Shared | 108M | 22.71
Total | 452M | 69.22
4 | 1/3 | 316M | 105.82 | 8.38
2/3 | 224M | 72.97
3/3 | 170M | 48.17
Shared | 108M | 22.71
Total | 494M | 76.19
Table 6: Ablation study of Multistage LDM on CelebA dataset.

###  C.3 Ablation Study on the Number of Intervals

In this subsection, we conduct ablation studies on the number of intervals of
our multistage architecture. We utilize Multistage DPM-Solver training on the
CIFAR-10 dataset. To extend Algorithm 1, we design the procedure outlined in
Algorithm 2 (Appendix B). The splitting intervals and results are given in
Table 7. We compare the stages from one to five. The parameters designed for
those architectures are based on the best three-stage model architecture in
Table 5. More specifically, we use linear interpolation. Suppose the
partitioning timesteps for the three-stage architecture are $t_{1}$ and
$t_{2}$ and parameter numbers for each interval are $p_{1},p_{2}$, and
$p_{3}$. The parameter $p$ for a random interval $[t_{a},t_{b}]$ is calculated
as:

$p=\frac{1}{L([t_{a},t_{b}])}[\mathcal{L}([t_{a},t_{b}]\cap[0,t_{1}))p_{1}~{}+~{}\mathcal{L}([t_{a},t_{b}]\cap[t_{1},t_{2}))p_{2}~{}+~{}\mathcal{L}([t_{a},t_{b}]\cap[t_{2},1])p_{3}]\>,$
(23)

where $\mathcal{L}$ is an operator that measures the length of an interval. We
design the parameters for all architectures in Table 7 based on (23), and
split intervals for each stage accordingly. As observed, the three-stage
architecture, adopted in the main body of the paper, outperforms all other
settings.

Number of Stages | Current/Total | Interval | Parameters | GFLOPs ($\downarrow$) | FID ($\downarrow$)
---|---|---|---|---|---
1 | 1/1 | [0,1] | 108M | 17.65 | 2.75
2 | 1/2 | [0,0.476) | 136M | 26.59 | 2.56
2/2 | (0.476,1] | 95M | 16.85
| Total |  | 188M | 21.49
3 | 1/3 | [0,0.442) | 169M | 29.95 | 2.35
2/3 | [0.442, 0.630) | 108M | 17.65
3/3 | [0.630,1] | 47M | 6.31
| Total |  | 230M | 18.65
4 | 1/4 | [0, 0.376) | 159M | 29.78 | 2.67
2/4 | [0.376, 0.526) | 95M | 16.85
3/4 | [0.526, 0.726) | 71M | 11.97
4/4 | [0.726, 1] | 32M | 4.41
| Total |  | 284M | 17.33
5 | 1/5 | [0, 0.376) | 154M | 30.38 | 2.88
2/5 | [0.376, 0.476) | 88M | 16.82
3/5 | [0.476, 0.626) | 88M | 16.82
4/5 | [0.626, 0.776) | 27M | 4.42
5/5 | [0.776, 1] | 21M | 3.28
| Total |  | 336M | 17.03
Table 7: Ablation study on the number of stages.

###  C.4 Experimental Setting on Interval Splitting Methods

This section presents the experimental settings of previous timestep
clustering algorithms considered in Table 4. We refer readers to the following
references [20, 21] for implementation details. Specifically, we compare these
methods for the 3 intervals case on CIFAR-10 dataset by training our proposed
Multistage architecture with calculated intervals until convergence. The
timestep and SNR-based clustering algorithms are relatively easy to implement.
For the gradient-based clustering, we collect the gradients of the trained
network by the DPM-Solver for clustering. These gradients are generated every
50k training iterations with a total of 450k iterations. We evaluate the
performances of these models based on their best FID scores across training
iterations.