One-Step Diffusion Distillation via Deep Equilibrium Models
Authors: Zhengyang Geng
Ashwini Pokle
J Zico Kolter
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep Equilibrium Models
Diffusion Models
Distillation
Generative Models

Abstract: Diffusion models excel at producing high-quality samples but naively require hundreds of iterations, prompting multiple attempts to distill the generation process into a faster network. However, many existing approaches suffer from a variety of challenges: the process for distillation training can be complex, often requiring multiple training stages, and the resulting models perform poorly when utilized in single-step generative applications. In this paper, we introduce a simple yet effective means of distilling diffusion models *directly* from the initial noise to the resulting image. Of particular importance to our approach is to leverage a new Deep Equilibrium (DEQ) model as the distilled architecture: the Generative Equilibrium Transformer (GET). Our method enables fully offline training with just noise/image pairs from the diffusion model while achieving superior performance compared to existing one-step methods on comparable training budgets. We demonstrate that the DEQ architecture is crucial to this capability, as GET matches a $5\times$ larger ViT in terms of FID scores while striking a critical balance of computational cost and image quality. Code, checkpoints, and datasets are available [here](https://github.com/locuslab/get).

One-Step Diffusion Distillation via Deep Equilibrium Models
Zhengyang Geng". Carnegie Mellon University zgeng20cs. cmu. edu
Ashwini Pokle" Carnegie Mellon University apokle@cs. cmu. edu
J. Zico Kolter
Carnegie Mellon University Bosch Center for AI zkolterêcs cmu edu
Abstract
Diffusion models excel at producing high-quality samples but naively require hundreds of iterations, prompting multiple attempts to distill the generation process into faster network. However, many existing approaches suffer from a variety of challenges: the process for distillation training can be complex, often requiring multiple training stages, and the resulting models perform poorly when utilized in single-step generative applications. In this paper, we introduce a simple yet effective means of distilling diffusion models directly from initial noise to the resulting image. Of particular importance to our approach is to leverage a new Deep Equilibrium (DEQ) model the distilled architecture: the Generative Equilibrium Transformer (GET). Our method enables fully offline training with just noise/image pairs from the diffusion model while achieving superior performance compared existing one-step methods on comparable training budgets. We demonstrate that the DEQ architecture is crucial this capability, as GET matches 5x larger ViT in terms of FID scores while striking critical balance of computational cost and image quality. Code, checkpoints, and datasets are available here.
Introduction
Diffusion models [35, 93, 95, 96] have demonstrated remarkable performance on a wide range of generative tasks such high-quality image generation [71, 82, 85, 87] and manipulation [19, 67, 71, 85, 86], audio synthesis [39, 40, 51, 58], video [36, 92], 3D shape [44, 78], text [30, 56). and molecule generation [18, 38, 107]. These models are trained with denoising objective derived from score matching [42, 95, 96], variational inference [35, 49, 93], or optimal transport [57, 60), enabling them to generate clean data samples by progressively denoising the initial Gaussian noise during the inference process. Unlike adversarial training, the denoising objective leads to more stable training procedure, which in turn allows diffusion models scale up effectively [82, 85, 87]. Despite the promising results, one major drawback of diffusion models is their slow generative process, which often necessitates hundreds to thousands of model evaluations [22, 35, 96]. This computational complexity limits the applicability of diffusion models in real-time or resource-constrainee scenarios.
In effort to speed up the slow generative process of diffusion models, researchers have proposed distillation methods [11, 68, 88, 97, 111] aimed at distilling the multi-step sampling process into more efficient few-step or single-step process. However, these techniques often come with their
Equal Contribution. Correspondence Zhengyang Geng and Ashwini Pokle
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

own set of challenges. The distillation targets must be carefully designed to successfully transfer knowledge from the larger model the smaller one. Further, distilling a long sampling process into few-step process often calls for multiple training passes. Most of the prevalent techniques for online distillation require maintaining dual copies of the model, leading to increased memory and computing requirements. As a result, there is a clear need for simpler and more efficient approaches that address the computational demands of distilling diffusion models without sacrificing the generative capabilities.
In this work, our objective is streamline the distillation of diffusion models while retaining the perceptual quality of the images generated by the original model. To this end, we introduce a simple and effective technique that distills multi-step diffusion process into single-step generative model, using solely noise/image pairs. At the heart of our technique is the Generative Equilibrium Transformer (GET), a novel Deep Equilibrium (DEQ) model [5] inspired by the Vision Transformer (ViT) [25, 75]. GET can be interpreted as an infinite depth network using weight-ried transformer layers, which solve for fixed point in the forward pass. This architectural choice allows for the adaptive application of these layers in the forward pass, striking balance between inference speed and sample quality. Furthermore, we incorporate an almost parameter- class conditioning mechanism in the architecture, expanding its utility to class-conditional image generation.
Our direct approach for distillation via noise/image pairs generated by a diffusion model, can, in fact, be applied to both VIT and GET architectures. Yet. in our experiments, we show that the GET architecture, in particular, is able achieve substantially better quality results with smaller models. Indeed, GET delivers perceptual image quality on par with or superior to other complex distillation techniques, such as progressive distillation [68, 88], in the context of both conditional and unconditional image generation. This leads to explore the potential of GETs further. We preliminarily investigate the scaling law of GETs-how its performance evolves as model complexity, in terms of parameters and computations, increases. Notably, GET exhibits significantly better parameter and data efficiency compared to architectures like ViT, as GET matches the FID scores of 5x larger ViT, underscoring the transformative potential of GET in enhancing the efficiency of generative models.
To summarize, we make the following key contributions:
We propose Generative Equilibrium Transformer (GET), a deep equilibrium vision trans- former that is well-suited for single-step generative models.
We streamline the diffusion distillation by training GET directly on noise/image pairs sampled from diffusion models, which turns out to be a simple yet effective strategy for producing one-step generative models in both unconditional and class-conditional cases.
For the first time, we show that implicit models for generative tasks can outperform classic networks in terms of performance, model size, model compute, training memory, and speed.
2 Preliminaries
Deep Equilibrium Models. Deep equilibrium models [5] compute internal representations by solving for a fixed point in their forward pass. Specifically, consider a deep feedforward model with L layers:
5+ 1] (21 (211);
for i = o...,, L-
where x E R== is the input injection, e e R== is the hidden state of 5th layer, and f! R Xnx + RT is the feature transformation of 7th layer, parametrized by 0. If the above model is weight-tied, i.e., ( So, Vi, then in the limit of infinite depth, the output 2 of this network approaches fixed point zt:
lim fa(z! =
6-000
Deep equilibrium (DEQ) models [5] directly solve for this fixed point z using black-box root finding
algorithms like Broyden's 's method [14], or Anderson acceleration [1] in the forward pass. DEQs
utilize implicit differentiation to differentiate through the fixed point analytically. Let ;
fz x)) z², then the Jacobian of zt with respect the model weights is given by
Oz*
(gel (2" (zz x)
(3)
00
Oz
00

---

m.o)
m.o)
MLP
Xx
bxx
LN
LN
fole m)
Inj
DEQ
Equilibrium Transformer
NHA
MHA
+
LN
LN
Emb
Injection
Transformer
Figure I: Generative Equilibrium Transformer (GET), (Left) GET consists of two major compo- nents: Injection transformer and Equilibrium transformer. The Injection transformer transforms noise embeddings into an input injection for the Equilibrium transformer. The Equilibrium transformer is the equilibrium layer that takes in noise input injection and an optional class embedding and solves for the fixed point. (Right) Details transformer blocks the Injection transformer (Inj) and Equilibrium transformer (DEQ), respectively. Blue dotted boxes denote optional class label inputs.
Computing the inverse of Jacobian can quickly become intractable as we deal with ligendiimnsional feature maps. One can replace the inverse-Jacobian term with cheap approximations [27-29] without sacrificing the final performance.
Diffusion Models. Diffusion models [22, 35, 93, 94] score-based based generative models [95, 96] progressively perturb images with an increasing amount Gaussian noise and then reverse this process through sequential denoising to generate images. Specifically, consider dataset of i.i.d. samples Pdata, then the diffusion process [x(t))T= for t E [0, T) is given by an Itô SDE [96]:
dx = fxx=f(xttt)
g(t)dw
where w is the standard Wiener process, t) Rdd + Rd is the drifting coefficient, g(.) R + IR the diffusion coefficient, and x(0) ~ Pdata and x(T) ~ NO, I). All diffusion processes have corresponding deterministic process known as the probability flow ODE (PF-ODE) [96] whose trajectories share the same marginal probability densities as the SDE. This ODE can be written as:
dx - (t)a v log p(x, at))dtt
(5)
where (t) is the noise schedule of diffusion process, and Vx log p(x, (t)) represents the score function. Karras et al. [48] show that the optimal choice of (t) in Eq. (5) is (t) = t. Thus, the PF-ODE can be simplified to dx/dt ~ log p(x, (t)) (x (xxxtt))+ where Do(:,)) t) is denoiser function parametrized with neural network that minimizes the expected L, denoising error for samples drawn from Pdata- Samples can be efficiently generated from this ODE through numerical methods like Euler's method, Minge-Kutta method, and Heun's 's second-order solver [3].
3 Generative Equilibrium Transformer
We introduce the Generative Equilibrium Transformer (GET), a Deep Equilibrium (DEQ) vision transformer designed to distill diffusion models into generative models that are capable of rapidly sampling images using only a single model evaluation. Our approach builds upon the key components and best practices of the classic transformer [99], the Vision transformer (ViT) [25], and the Diffusion transformer (DiT) [75]. We will now describe each component of the GET in detail.

---

GET. Generative Equilibrium Transformer (GET) directly maps Gaussian noises and optional class labels y to images x. The major components of GET include an injection transformer (InjectionT. Eq. (7)) and an equilibrium transformer (Equilibrium'T, Eq. (8)). The Injection'I transforms tokenized noise embedding h an intermediate representation n that serves as the input injection for the equilibrium transformer. The Equilibrium',, which is the equilibrium layer, solves for the fixed point z" by taking in the noise injection n and an optional class embedding c. Finally, this fixed point z² is decoded and rearranged to generate an image sample x (Eq. (9)). Figure provides overview of the GET architecture. Note that because we are directly distilling the entire generative process, there is no need for time embedding as is common in standard diffusion models.
h, Emb (e), Emb (y); if y &
(6)
n Injection'T' (h, c)
(7)
z Equilibrium'T (z" n.
(8)
X Decoder (z*)
(9)
Noise Embedding. GET first converts input noise E RRHWWCC into a sequence of 2D patches
p E (PP2..)) where C is the number of channels, P is the size of patch, H and W denotes height and width of the original image, and N = HW/pp² is the resulting number of patches. Let D = p2_ C denote the width of the network. We follow ViT patches to D dimensional embedding. We add standard sinusoidal position encoding [99] to produce the noise embedding h. Position encoding plays crucial role in capturing the spatial structure of patches by encoding their relative positional information.
linear layer project the N
InjectionT & EquilibriumTT. Both InjectionTT and EquilibriumTT composed of sequence of Transformer blocks. InjectionT is called only once to produce the noise injection n, while EquilibriumT defines the function So of the implicit layer z - fo(z*, n, c) that is called multiple times-creating a weight-tied computational graph-until convergence. A linear layer is added at the end of Injection'T to compute the noise injection ny R
[Le]. for each of the L
GET blocks in EquilibriumT. For convenience, we overload the notation ny and n, in the subsequent paragraphs.
Transformer Block, GET utilizes a near-identical block design for the noise injection (Injection') and the equilibrium layer (Equilibrium differing only the injection interface. Specifically, the transformer block is built upon the standard Pre-LN transformer block [25, 75, 106], as shown below:
+ Atlention (LN (z), u)
= z + FFN (LN (z))
Here, z E RD represents the latent token, E NR333 is the input injection, LN, FFN, and Atten- tion stand for Layer Normalization [4], a 2-layer Feed-Forward Network with a hidden dimension of size D x E, and an attention [99] layer with an injection interface, respectively.
For blocks in the injection transformer, u is equal to the class embedding token c E RRx3D for conditional image generation, i.e., for blocks in the equilibrium transformer, u is the broadcast sum of noise injection n E R N x and class embedding token E R 1 i.e.,
for conditional models, and u = 0 otherwise. In contrast,
+ for conditional models and n otherwise.
We modify the standard transformer attention layer to incorporate an additive injection interface before the query E NR, key k E NxD and value v E RND.
q, k, zW, +u
Z = MHA (q. v)
z zW.
where W, E R Dx3D W. E PxD, The injection interface enables interactions between the latent tokens and the input injection the multi-head dot- product attention (MHA) operation,
qk (zW, + (z + )
T ZW,WZ z T + zW + u. W ZJ u ug.
(10)
where W q, W E DxD are slices from W. and Ug Un E NxD are slices from u. This scheme adds no more computational cost compared to the standard MHA operation, yet it achieves a similar effect as cross-attention and offers good stability during training.

---

FID Ittrationn
Parameter Scaling
Progressive Distillation (PD)
GET-200k
PD-EDM-4XB5
GET-400K
VIT-B-1xBS
20
GET-800k
FID Scores
VIT-L-1xBS (502M)
FID Scores
VIT-200k
20
GET-B-1XB5 (62M)
VIT-400k
ViT-800k
10
10
200
400
600
800
80
160
240
Iterations (1k)
Parameters (M)
Figure 2: Data and Parameter Efficiency of GET:(a) (Left) GET outperforms PD and a 5x larger ViT in fewer iterations, yielding better FID scores. Additionally, longer training times lead to improved FID scores. (b) (Right) Smaller GETs can achieve better FID scores than larger ViTs, demonstrating DEQ's s parameter efficiency. Each curve in this plot connects models of different sizes within the same model family at identical training iterations, as indicated by the numbers after the model names in the legend.
The output of the GET-DEQ is first normalized with Layer Normalization [4]. The normalized output is then passed through another linear layer to generate patches p NR The resulting patches p are rearranged back to the resolution of the input noise e to produce the image sample "x xxx Thus, the decoder maps the features back to the image space.
Image Decoder.
4
Experiments
We evaluate the effectiveness of our proposed Generative Equilibrium Transformer (GET) in offline distillation of diffusion models through series of experiments single-step class-conditional and unconditional image generation. Here, we use "single-step" to refer to the use of a single model evaluation while generating samples. We train and evaluate VTT's and GETss of varying scales these tasks. GETs exhibit substantial data and parameter efficiency in offline distillation compared to the strong VIT baseline. Note that owing to the computational resources required to fully evaluate models, we report all our results on CIFAR-10 [52]; extensions to the ImageNet-scale [20] are possible, but would require substantially larger GPU resources.
4.1 Experiment setup
We will first outline our data collection process, followed by an in-depth discussion of our offline distillation procedure. We also include a brief summary of training details and evaluation metrics. For detailed network configs and training specifics, please refer to the Appendix.
Data Collection. For unconditional image generation on CIFAR-10 [52], we generate IM noise/im- age pairs from the pretrained unconditional EDM Karras et al. [48]. This dataset is denoted as EDM-Uncond- 1M. As in EDM, we sample 1M images using Heun's 's second-order deterministic solver [3]. Generating a batch of images takes 18 steps or 35 NFEs (Number of Function Evaluations). Overall, this dataset takes up around 29 GB of disk space. The entire process of data generation takes about hours on 4 NVIDIA A6000 GPUs using Pytorch [74] Distributed Data Parallel (DDP) and a batch size of 128 per GPU. In addition to unconditional image generation, we sample 1M labe/laabl/ pairs from the conditional VP-EDM Karras et [48] using the same settings. This dataset is denoted as EDM-Cond-111 1M. Both the datasets will be released for future studies.
Offline Distillation. We distill pretrained EDM [48] into ViTs and GETs by training on a dataset DD with noise/image pairs sampled from the teacher diffusion model using a reconstruction loss:
(( = Ee.x~>l /| - Ge
where x is the desired ground truth image, Go(.) is unconditional VIT/GET with parameters 0, and is the initial Gaussian noise. To train a class-conditions GET, we also use class labels in addition to noise/image pairs:
[()) Ee |x Ge,
5

---

FID Sampling Speed
FID GFLOPs
VIT-B
15
VIT-L
FID Scores
FID Scores
32
GET-Mini
GET-Tiny
10
GET-Small
16
GET-Baset
GET-Mini
GET-Small
VIT-B
VIT-L
8
20
40
60
2
2
2
2
Sampling Speed (sec/50k)
Training Compute (GFLOPs)
Figure 3: (a) (Leff)) Sampling speed of GET: GET can sample faster than large ViTs, while achieving better FIID scores. The size of each individual circle is proportional to the model size. For GET', vary the number of iterations in the Equilibrium transformer (2 to 6 iterations). The trends indicate that GETs can improve their FID scores by using more compute. (b) (Right) Compute efficiency of GET: Larger GET models use training compute more efficiently. For a given GET, the training budget is calculated from training iterations. Refer Table for the exact size of GET models.
Figure 4: Uncurated CIFAR-10 image samples generated by (Left) (a) unconditional GET and (Right) (b) class-conditional GET. Each row corresponds to class in CIFAR-10.
where (.) is class-conditional VITIGET with parameters 6. As is the standard practice, we also maintain exponential moving average (EMA) of weights of the model, which in turn is used at inference time for sampling.
Training Details and Evaluation Metrics. We use Adam W [63] optimizer with learning rate of lc-4, a batch size of 128 (denoted as 1 BS), and 800k training iterations, which are identical to Progressive Distillation (PD) [88]. For conditional models. we adopt batch size of 256 (2 BS). No warm-up, weight decay, or learning rate decay is applied. We convert input noise to patches of size 2 2. We use 6 steps of fixed point iterations in the forward pass of GET-DEQ and differentiate through it. For the (1)) memory mode, we utilize gradient checkpoint [17] for DEQ's 's computational graph. We measure image sample quality for all our experiments via Frechet inception distance (FID) [34] of 50k samples. We also report Inception Score (IS) [89] computed on 50k images. We include other relevant metrics such as FLOPs, training speed, memory, sampling speed, and the Number of Function Evaluations (NFEs), wherever necessary.
4.2 Experiment Results
We aim to answer the following questions through extensive experiments: 1) Can offline distillation match online distillation for diffusion models using GETs? 2) What is the scaling behavior of GET

---

as the model size and training compute increase? 3) How does GET compare to existing one-step generative models in terms of image quality and diversity?
Efficiency. Models trained with offline distillation require high data efficiency to make optimal use of limited training data sampled from pretrained diffusion models. DEQs have a natural regularization mechanism due to weight-tying, which allows us to efficiently fit significantly compact data-efficient models even in limited data settings. In Figure 2(a), we observe that even with a fixed and limited offline data budget of IM samples, GET achieves parity with online distilled EDM [48, 88, 97] while using only half the number of training iterations. For comparison. PD, TRACT, and CM use a much larger data budget of 96M, 256M, and 409.6M samples, respectively. Moreover, GET is able to match the FII score of a 5 x large VIT, suggesting substantial parameter efficiency.
Sampling Speed. Figure 3(a) illustrates the sampling speed of both ViT and GET. A smaller GET (37.2M) can achieve faster sampling than a larger ViT (302.6M) while achieving lower FID scores. GET can also improve its FID score by increasing its test-time iterations in the Equilibrium transformer at the cost of speed. Note that despite this trade-off, GET still outperforms larger VIT in terms of both sampling speed and sample quality.
Why Scaling Laws for Implicit Models? As prospective study, we preliminarily investigate the scaling properties of Deep Equilibrium models using GET. The scaling law is an attractive property, as it enables us to predict models' performance at extremely large compute based on the performance of tiny models. This predictive capability allows us to select the most efficient model given the constraints of available training budget [13, 37, 72). While the scaling law for explicit networks has been extensively studied, its counterpart for implicit models remains largely unexplored. Implicit models are different from explicit models they utilize more computation through weight-tying under similar parameters and model designs. Therefore, it is natural to question whether their scaling laws align with those of their explicit counterparts.
Scaling Model Size. We conduct extensive experiments to understand the trends of sample quality as we scale the model size of GET. Table 2 provides summary of our findings single-step unconditional image generation. We find that even small GET models with 10-20M parameters can generate images with sample quality on par with NAS-derived AutoGAN [31]. In general, sample quality improves with the increase in model size.
Scaling Training Compute. Our experimental results support the findings of Peebles and Xie [75] for explicit models (DiT) and extend them to implicit models. Specifically, for both implicit and explicit models, larger models are better at exploiting training FI.OPs. Figure shows that larger models eventually outperform smaller models when the training compute increases. For implicit models, there also exists a "sweet spot" in terms of model size under a fixed training budget, e.g. GET-Small outperforms both smaller and larger GETs at 2-1 training GFLOPs. Furthermore, because of the internal dynamics of implicit models, they can match a much larger explicit model in terms of performance while using fewer parameters. This underscores the potential of implicit models as candidates for compute-optimet models [37] with substantially better parameter efficiency. For example, at 231 training GFLOPs, Figure 3(b) suggests that should choose GET-Small (31.2M) among implicit models for the best performance, which is much more parameter efficient and faster in sampling than the best-performing explicit model, ViT-L (302M), at this training budget.
Benchmarking GET against VIT. Table 4 summarizes key metrics for unconditional image generation for ViT and GET. Our experiments indicate that smaller GET (19.2M) can generate higher-quality images faster than a much larger ViT (85.2M) while utilizing less training memory and fewer FLOPs. GET also demonstrates substantial parameter efficiency over ViTs as shown in Figure 2(b) where smaller GETs achieve better FID scores than larger ViTs.
Comparizon of NFEs of teacher model. Offline distillation requires significantly fewer number of function evaluations (NFEs) for the teacher network compared to other online distillation methods. In the experimental setup used in this paper, GET requires 35M overall NFEs for the teacher model, as we train on 1M data samples, and use 35 NFEs to generate each data sample with EDM. In contrast, progressive distillation requires 179M NFES to get -step distilled student model. Using the hyperparameters reported in Salimans and Ho [88], PD with DDIM model needs 13 passes of
7

---

Table 2: Generative performance of GET' conditional CIFAR-10.
Table I: Generative performance on uncondi-
tional CIFAR-10.
Models
Params
NFE
FID
IS
Method
NFE 4
FID 4
f
GET-Timy
8.6M
15.19
8.37
Ditfusion Models
GET-Mini
19.2M
10.72
8.69
DDPM [35]
3.17
GET-Small
37.2M
8.00
9.03
Score SDE [96]
2000
GET-Base
62.2M
7.42
9.16
GET-Base-
83.5M
7.19
9.09
DDIM [94]
13.36
EDM [48]
2.04
More Training
GANs
GET-Tiny--4xxe xters
8.9M
11.47
8.64
GET-Base-2xB3.
62.2M
6.91
9.16
StyleGAN2 [47]
8.32
9.18
StyleGAN XXL [91]
1.85
Table 3: Generative performance on class-
Diffusion Distillation
conditional CIFAR-IO. w indicates the level of classifier guidance.
KD [66]
9.36
8.36
PD [88]
9.12
DENO [111]
4.12
Method
NFE
FID
IS
TRACT-EDM
4.17
GANs
PD-EDM [88, 97]
8.34
8.69
BigGAN [12]
i
14.73
9.22
CD-EDM (LPIPS) (97]
3.55
9.48
StyleGAN2-ADD [46]
2.42
10.14
Consistency Models
Diffusion Distillation
CT [97]
8.70
8.49
Guided Distillation (w [68]
8.34
8.63
CT 1971
5.83
8.85
Guided Distillation (w = 0..3) [68]
7.34
8.90
Guided Distillation (w -1)168]
Ours
Guided Distillation (w 2) [68]
13.23
9.23
GET-Base
6.91
9.16
Ours
GET-Base
6.25
9.40
Table 4: Benchmarking GET against VIT on unconditional image generation on CIFAR-10. For the first time, implicit models for generative tasks strictly surpass explicit models all metrics. Results are benchmarked on 4 A6000 GPUs using a batch size of 128, 800k iterations, and PyTTorch [74] distributed training protocol. Training Mem stands for training memory consumed per GPU. O()) symbolizes the (1) training memory mode, which differs only in training memory and speed.
Model
FID
ISt
Params|
FLOPs,
Training Mem.
Training Speed
ViT-Base
11.49
8.61
85.2M
23.0G
10.1GB
4.83 iter/see
GET-Mini
10.72
8.69
19.2M
15.2G
9.2GB
5.79 iter/sec
GET-Mini-O(1)
5.0GB
4.53 iter/sec
distillation. The initial 12 passes use 50K iterations, and the last pass uses 100K iterations. Each step of PD uses 2 teacher model NFEs. Thus, the overall number of teacher model NFTS can be evaluated as x 128 (batch size) x (12 passes x50K + 100K) = 1 19M samples. The number of NFEs of the teacher model increases 1.433B if we assume that each of 8 TPUs use batch size of 128. Consistency distillation [97] needs 409.6M teacher model NFEs (512 batch size x 800K iterations 409.6M). In addition, the perceptual loss requires double NFEs as the teacher model.
One-Step Image Generation. We provide results for unconditional and class-conditional image generation on CIFAR- 10 Table and Table 3, respectively. GET outperforms a much more complex image generation.
distillation procedure- PD with classifier-free guidance-in in class-conditions GET also outperforms PD and KD terms of FID score for unconditional image generation. This effectiveness is intriguing, given that our approach for offline distillation is relatively simpler when compared to other state-of-the-ar distillation techniques. We have outlined key differences in the experimental setup between approach and other distillation techniques in Table 5.
We also visualize random CIFAR-10 [52] samples generated by GET for both unconditional and class-conditional cases in Figure 4. GET can learn rich semantics and world knowledge from the dataset, as depicted the images. For instance, GET has learned the symmetric layout of dog faces solely using reconstruction loss in the pixel space, as shown in Figure 4(b).

---

Table 5: Comparison of relevant training and hyperparameter settings for common distillation
techniques. GET requires neither multiple training phases nor any trajectory information. We only count the number of models involved in the forward pass and exclude EMA in #Models. indicates offline distillation techniques. For CD, we count the VGG network used in the perceptual loss [109].
Model
FID
IS f
BS
Training Phases
#Models
Trajectory
Teacher
KD [66]
9.36
4x
1
DDIM
PD [88]
9.12
1x
logg(T)
2
DDIM
DENO IIII
4.12
2x
1
DDIM
TRACT [11]
14.40
2x
1
1
DDIM
TRACT [11]
4.17
2x
2
1
V
EDM
PD-EDM [88, 971
8.34
8.69
4x
logg (T)
2
V
EDM
CD [97]
3,55
9.48
4x
3
EDM
Ours
7.42
9.16
1x
1
1
EDM
Ours
6.91
9.16
2x
1
1
X
EDM
Guided Distillation [68]
7.34
8.90
4x
loga (T) + 1
3
/
DDIM
Ours
6.25
9.40
2x
1
1
EDM
Related Work
Distillation techniques for diffusion models. Knowledge distillation (KD) [66] proposed to distill multi-step DDIM [94] sampler into the pretrained UNet by training the student model synthetic image samples. There are several key differences from this work: Our approach does not rely on temporal embeddings or generative pretrained weights and predicts images instead of noises. Further, GET is built upon ViT [25], unlike the UNet in KD. Additionally, we demonstrate the effectiveness of our approach on both unconditional and class consiccnn image generation.
Progressive distillation (PD) [88] proposes strategy for online distillation to distill 7 7'-step teacher DDIM [94] diffusion model into new T/2 step student DDIM model, repeating this process until one-step models are achieved. Transitive closure time-distillation (TRACT) [11] generalizes PD to distill N > 2 steps together at once, reducing the overall number training phases. Consistency models [97] achieve online distillation in single pass by taking advantage of a carefully designed teacher and distillation loss objective.
Diffusion Fourier neural operator (DFNO) [111] maps the initial Gaussian distribution to the solution trajectory of the reverse diffusion process by inserting the temporal Fourier integral operators in the pretrained U-Net backbone. Meng et al. [68] propose a two-stage approach to distill classifier-fres guided diffusion models into few-step generative models by first distilling combined conditional and unconditional model, and then progressively distilling the resulting model for faster generation.
Fast sampler for diffusion models. While distillation is predominant approach to speed up the sampling speed of existing diffusion models, there are alternate lines of work to reduce the length of sampling chains by considering alternate formulations of diffusion model [48, 50, 94, 96, 103], correcting bias and truncation errors in the denoising process [9, 10, 90], and through training- fast samplers at inference [24, 43, 50, 59, 65, 108). Several works like Improved DDPM [70], SGM- CLD [23], EDM [48] modify or optimize the forward diffusion process so that the reverse denoising process can be made more efficient. Diffusion Exponential Integrator Sampler (DEIS) [108] uses an exponential integrator over the Buler method to minimize discretization error while solving SDE. DPM-Solver [65], and GENIE [24] are higher-order ODE solvers that generate samples in a few steps.
6 Limitations
Our method for offline distillation relies on deterministic samplers to ensure a unique mapping between initial noise and image x. As result, it cannot be directly applied to stochastic samplers which not satisfy this requirement. However, this limitation also applies to many other distillation techniques, as they cannot maintain their fidelity under stochastic trajectories [11, 66, 88, 97].

---

7 Conclusion
We propose simple yet effective approach distill diffusion models into generative models ca- pable of sampling with just a single model evaluation. Our method involves training a Generative Equilibrium Transformer (GET) architecture directly on noise/image pairs generated from a pre- trained diffusion model, eliminating the need for trajectory information and temporal embedding. GET demonstrates superior performance over more complex online distillation techniques such as progressive distillation [68, 88] in both class conditional and unconditional settings. In addition, small GET can generate higher quality images than a 5x larger ViT, sampling faster while using less training memory and fewer compute FLOPs, demonstrating its effectiveness.
8 Acknowledgements
Zhengyang Geng and Ashwini Pokle are supported by grants from the Bosch Center for Artificial Intelligence.
References
[1] Donald G. Anderson. Iterative procedures for nonlinear integral equations. Journal of the
ACM (JACM), 12(4))547-560 October 1965. 2
[2] Cem Anil, Ashwini Pokle, Kaiqu Liang. Johannes Treutlein, Yuhuai Wu, Shaojie Bai, J Zico Kolter, and Roger Grosse. Path independent equilibrium models better exploit test-time computation. Advances Neural Information Processing Systems, 35:7796-7809, 2022. 18
[3] Uri Ascher and Linda R Petzold. Computer methods for ordinary differential equations and differential-algebruii equations, volume 61. Siam, 1998. 3,5
[4] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv::607.00 007,606607.06450, 2016. 4, 5
[5] Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In Neural
Information Processing Systems (NeurIPS), 2019. 2, 18
[6] Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Multiscale deep equilibrium models. Advances in Neural Information Processing Systems, 33:5238-5250, 2020. 18
[7] Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Stabilizing Equilibrium Models by Jacobian Regularization. In International Conference Machine Learning (ICML), 2021. 18
[8] Shaojie Bai, Zhengyang Geng, Yash Savani, and Zico Kolter. Deep equilibrium optical flow estimation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 620-630, 2022. 18
[9] Fan Bao, Chongxuan Li, Jiacheng Sun, Jun Zhu, and Bo Zhang. Estimating the optimal covari- ance with imperfect mean in diffusion probabilistic models. arXiv preprint arXiv:2206.07339, 2022. 9
[10] Fan Bao, Chongxuan Li, Jun Zhu, and Zhang. Analytie-dpm: an analytic estimate the optimal reverse variance in diffusion probabilistic models. arXiv preprint urXiv:2201.06503 2022. 9
[11] David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap. Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time -ime-iitillattin, 2023. 1, 8,
[12] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint trXiv:/809: /096, 2018. 8, 18
[13] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020. 7, 18
10

---

[14] Charles G Broyden. A Class of Methods for Solving Nonlinear Simultaneous Equations. Mathematics of computation, 19(92):577-533, 1965. 2
[15] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In IEEE International Conference on Computer Vision (ICCV), 2021. 17
[16] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084- 15097, 2021. 17
[17] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.00 1116604.66144 06177, 2016. 6
[18] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210.01776.
2022.
[19] Guillaume Couairon, Jakob Verbeek, Holger Schwenk, and Matthieu Cord. Diffedit: Diffusion- based semantic image editing with mask guidance. arXiv preprint arXiv:2210.1 1427, 2022.
[20] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 245-255. leee, 2009.
[21] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT:
Pre-training of deep bidirectional transformers for language understanding. In Annual Conference of the North American Chapter of the Association for Computational Linguistics (NAACL), 2019. 17, 18
[22] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Neural Information Processing Systems (NeurIPS), 2021. 1,3, 18
[23] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with
critically-damped langevin diffusion. arXiv preprint arXiv:2112.07066 2021. 9
[24] Tim Dockhorn, Arash Vahdat, and Karsten Kreis. GENIE: Higher- Denoising Diffusion Solvers. In Advances in Neural Information Processing Systems, 2022.
[25] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021. 2, 3, 4, 9, 17, 18
[26] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12873-12883, 2021. 18
[27] Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley J. Osher, and Wotao Yin. Fixed Point Networks: Implicit Depth Models with Jacobian Free Backprop. arXiv preprint arXiv:2103.12800 2021.
[28] Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is Attention Better Than Matrix Decomposition? In International Conference on Learning Representations (ICLR), 2021.
[29] Zhengyang Geng, Xin- Yu Zhang, Shaojie Bai, Yisen Wang, and Zhouchen Lin. On training implicit models. In Neural Information Processing Systems (NeurIPS), 2021. 3
[30] Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models, arXiv preprint arXiv:2210.089) 2022.
11

---

[31] Xinyu Gong, Shiyu Chang, Yifan Jiang, and Zhangyang Wang. Autogan: search for generative adversarial networks. In IEEE International Conference on Computer Vision (ICCV), 2019. 7
Neural architecture
[32] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: XXXxv:/00..00 170670066777777 2017.
[33] Swaminathan Gurumurthy, Shaojie Bai, Zachary Manchester, and J 7ico Kolter. Joint inference and input optimization in equilibrium networks. Neural Information Processing Systems (NeurlPS), 2021. 18
[34] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[35] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Information Processing Systems (NeurIPS), 2020. 1, 1.3.7 8
Denoising diffusion probabilistic models. Neural
[36] Jonathan Ho, Tim Salimans, Alexey A Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. In Neural Information Processing Systems (NeurIPS),
2022. 1
[37] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022. 7
[38] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning,
pages 8867-8887. PMLR, 2022. 1
[39] Qingqing Huang, Daniel S Park, Tao Wang, Timo Denk, Andy Ly, Nanxin Chen, Zhengdong Zhang. Zhishuai Zhang, Jiahui Yu, Christian Frank, et al. Moise2music: Text-conditioned music generation with diffusion models. arXiv preprint arXiv:2302.03917 2023.
[40] Rongjie Huang, Max WY Lam, Jun Wang, Dan Su, Dong Yu, Yi Ren, and Zhou Zhao Fastdiff: A fast conditional diffusion model for high-quality speech synthesis. arXiv preprint arXiv:2204.00 09934, 2022.
[41] Drew A Hudson and Larry Zitnick. Generative adversarial transformers. In International Conference on Machine Learning (ICML), 2021. 18
[42] Aapo Hyvärinen and Peter Dayan. Estimation of non nonmmoomal statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[43] Alexia Gotta go fast when generating data with score- score-based models. arXiv preprint 110:22005.14080 2021. 9
Ke Li, Rémi Taillefe,,,,,,, Tal Kachman, and Ioannis Mitliagkas.
[44] Heewoo Jun and Alex Nichol. Shap-e: Generating conditional 3d implicit functions, 2023. 1
[45] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IÉEE/CVF conference on computer vision and pattern recognition, pages 4401-4410 2019. 18
[46] Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. Advances in neural information processing systems, 12:1200--2114 2020.
[47] Tero Karras, Samuli I Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEF/CVF conference computer vision and pattern recognition, pages $110-8119, 2020. 8
12

---

[48] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-baasd based generative models. In Neural Information Processing Systems (NeurIPS).
2022. 3,5, 7, 8, 8,9
[49] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Neural Information Processing Systems (NeurIPS), 2021.
[50] Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabiliistic models, arXiv preprint wrX1=22006.0012,, 2021.
[51] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.0976/ 2020. 1
[52] Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. 5, 8
[53] Guillaume Lample and Alexis Conneau. Cross-lingual language model pretraining, arXiv preprint arXiv: 07091, 2019. 18
[54] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: lite bert for self-supervised learning of language representations. arXiv preprint arXiv: 1942, 2019. 18
[55] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising sequence-to-sequenne pre-training for natural language generation, translation, and comprehension. arXiv preprint 0tti://910.3466/0 2019. 17
[56] Xiang Li, John Thickstun, Ishaan Guirajani, Percy S Liang, and Tatsunori Hashimoto. Diffusion-lm improves controllable text generation. Advances in Neural Information Process- ing Systems, 35:4328-4345 2022.
[57] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference Learning Representations (ICLR), 2023.
[58] Jinglin Liu, Chengxi Li, Yi Ren, Feiyang Chen, and Zhou Zhao. Diffsinger: Singing voice synthesis via shallow diffusion mechanism. Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages I 1020-11028, 2022.
[59] Luping Liu, Yi Ren, Zhijic Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09776 2022. 9
[60] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. In International Conference on Learning Representations (ICLR), 2023.
[61] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:
2019. 18
[62] 7e Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In IEEE International Conference on Computer Vision (ICCV), 2021. 17
[63] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 0rXiv:1711.051000 2017. 6
[64] Cheng Lu, Jianfei Chen, Chongxuan Li, Qiuhao Wang, and Jun Zhu. Implicit normalizing flows. arXiv preprint a 2021. 18
[65] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling around 10 steps. arXiv preprint arXiv:2206 wrXix:2206.00927, 2022. 9
13

---

[66] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02336 0228888,, 2021. 8,
[67] Chenlin Meng, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Image synthesis and editing with stochastic differential equations. arXiv preprint arXiv:2108.0000 2021.
[68] Chenlin Meng, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans, On distillation of guided diffusion models. arXiv preprint arXiv:2210003112, 2022.
1, 2, 8, 9, 10
[69] Paul Micaelli, Arash Vahdat, Hongxu Yin, Jan Kautz, and Pavlo Molchanov. Recurrence without recurrence: Stable video landmark detection with deep equilibrium models. arXiv preprint arXivv2304.00600, 2023. 18
[70] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning (ICML), 2021.
[71] Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Megrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. In International Conference on Machine Learning (ICMI.), 2022.
[72] OpenAL. Gpt-4 technical report. ArXiv, abs/2303.08774, 2023. 7
[73] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning
(ICML), 2020. 17
[74] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-performance Deep Learning Library. In Neural Information Processing Sys- tems (NeurIPS), 2019. 5,8
[75] William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09745. 2022. 2, 3, 4, 7, 17, 18
[76] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In Association for the Advancement of Artificial Intelligence (AAAI), 2018. 18
[77] Ashwini Pokle, Zhengyang Geng, and Zico Kolter. Deep equilibrium approaches to diffusion models. Advances in Neural Information Processing Systems, 35:37975-37990, 2022. 18
[78] Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv, 2022.
[79] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018. 17, 18
[80] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 18
[81] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, (55555585-5555 2020. 18
[82] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25, 2022. 1,
18
14

---

[83] Max Revay, Ruigang Wang, and Ian R Manchester. Lipschitz bounded equilibrium networks. arXiv preprint arXiv:2010.01732. 2020. 18
[84] Adam Roberts, Colin Raffel, Katherine Lee, Michael Matena, Noam Shazeer, Peter Liu, Sharan Narang, Wei Li, and Yanqi Zhou. Exploring the limits of transfer learning with unified text-to-text transformer. 2019. 17
[85] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022. 1, 18
[86] Chitwan Saharia, William Chan, Huiwen Chang, Chris Lee, Jonathan Ho, Tim Salimans, David Fleet, and Mohammad Norouzi. Palette: image-to-imaae diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, pages 1-10, 2022.
[87] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang. Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Neural Information Processing Systems (NewrIPS), 2022.
[88] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005/2 2022. 1, 1, 6, 7, 8, 9, 10
[89] Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung. Alec Radford, and Xi Chen. Improved techniques for training gans, Advances in neural information processing systems, 29, 2016. 6
[90] Robin San-Roman, Eliya Nachmani, and Lior Wolf. Noise estimation for generative diffusion models. arXiv preprint arXiv:2 02600, 2021. 9
[91] Axel Sauer, Katja Schwarz, and Andreas Geiger. Stylegan-xl: Scaling stylegan to large diverse datasets. In ACM SIGGRAPH 2022 conference proceedings, pages 1-10, 2022. 8
[92] Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, al. Make-a-video: Text-to-video generation without text-video data. arXiv preprint arXiv:2209.1479) 2022. 1
[93] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (ICML), 2015. 1,3
[94] Jiaming Song, Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR), 2021. 3, 8,
[95] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Neural Information Processing Systems (NeurIPS). 2019. 1,3
[96] Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. scsre-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR). 2021. 3, 8,
[97] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.00 01690, 2023. 1, 7.
[98] Russell Tsuchida and Cheng Soon Ong. Deep equilibrium models estimators for continuous latent variables. In International Conference on Artificial Intelligence and Statistics, pages 1646-1671. PMLR, 2023. 18
[99] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Neural Information Processing Systems (NeurIPS), 2017. 3, 4, 17, 18
[100] Shuai Wang, Yao Teng, and Limin Wang. Deep equilibrium object detection. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6296-6306, 2023. 18
15

---

[101] Thomas Wang, Adam Roberts, Daniel Hesslow, Teven Le Scao, Hyung Won Chung, Iz Beltagy. Julien Launay, and Colin Raffel. What language model architecture and pretraining objective works best for zero-shot generalization? In International Conference Machine Learning.
pages 22964-22984. PMLR, 2022. 18
[102] Tiancai Wang, Xiangyu Zhang, and Jian Sun. Implicit Feature Pyramid Network for Object Detection. arXiv preprint arXiv:2012.13563, 2020. 18
[103] Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning efficiently sample from diffusion probabilistic models. arXiv preprint arXiv:2106.03802 2021. 9
[104] Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan
Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint aww.ix2200900603300 2021. 18
[105] Ezra Winston and J Zico Kolter. Monotone operator equilibrium networks. Advances in neural information processing systems, 10700788-0728, 2020. 18
[106] Ruibin Xiong, Yunchang Yang, Di He. Kai Zheng. Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. International Conference on Machine Learning (ICML), 2020. 4
[107] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geod- iff: A geometric diffusion model for molecular conformation generation. arXiv preprini arXiv:2203.02923, 2022.
[108] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. The Eleventh International Conference on Learning Representations, 2023.
[109] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unrea- sonable effectiveness of deep features perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 586-595, 2018.
[110] Sai Zhang, Liangjia Zhu, and Yi Gao. An efficient deep equilibrium model for medical image segmentation. Computers in Biology and Medicine, 148: 105831, 2022. 18
[111] Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Arizzadenesheli, and Anima Anandkumar.
Fast sampling of diffusion models via operator learning. arXiv preprint arX1:22211.1344,
2022. 1, 8,
16

---

A Additional Experiments
Class Conditioning. As both GET and VIT share the same class injection interface, we perform an ablation study on ViT. We consider two types of input injection schemes for class labels: 1) additive injection scheme injection with adaptive layer normalization (AdaLN-Zero) Despite using almost the same parameters as unconditional ViT-B, the class-conditions ViT-B using additive injection interface has an FID of 12.43 at 200k, while the ViT-B w/ AdaLN-Zero class embedding [75] set up FID of 17.19 200k iterations. Another surprising observation is that ViT-B w/ AdaL---Zero class embedding performs worse than unconditional VIT in terms of FID score. Therefore, it seems that adaptive layer normalization might not be useful when used only with class embedding.
used in DiT [75].
Table 6: Ablation on class conditioning.
Model
FID
ISt
Params.
ViT-Uncond
15.20
8.27
85.2M
-iTeed
17.19
8.38
128.9M
ViT-Inj-Interface
12.43
8.69
85.2M
B Model Configuration
We set the EMA momentum to 0.9999 for all the models.
The configuration of different GET architectures are listed in Table 7. Here, L; and Le denote the number of transformer blocks the Injection transformer and Equilibrium transformer, respectively. D denotes the width of the network. E corresponds to the expanding factor of the FFN layer in the Equilibrium transformer, which results in the hidden dimension of E x D. For the injection transformer, we always adopt expanding factor of 4.
Table 7: Details of configuration for GET architectures.
Model
Params
L;
Le
D
E
GET-Tiny
8.9M
256
GET-Mini
19.2M
6
384
GET-Small
37.2M
512
GET-Base
62.2M
1
768
12
GET-Base-
83.5M
768
We have listed relevant model configuration details of VIT in Table 8. The model configurations are adopted from DiT [75], whose effectiveness was tested for learning diffusion models. In this table, denotes the number of transformer blocks in ViT. D stands for the width of the network. We always adopt an expanding factor of 4 following the common practice [25, 75, 99].
Table 8: Details of configuration for VIT architectures.
Model
Params
D
ViT-B
85.2M
768
VIT-L
302.6M
24
1024
C
Related Work
Transformers.
Transformers were first proposed by Vaswani et al. [99] for machine translation and since then have been widely applied many domains like natural language processing [21, 55, 79, 84], reinforcement learning [16, 73], self-supervised learning [15], vision [25, 62), and generative
17

---

modeling [26, 41, 75, 82]. Many design paradigms for transformer architectures have emerged over the years. Notable ones include encoder-only [21, 54, 61], decoder-only [13. 79, 80, 101, 104], and encoder-decoder architectures [53, 81, 99]. We are interested in scalable transformer architectures for generative modeling. Most relevant to this work are two encoder-only transformer architectures: Vision Transformer (ViT) [25] and Diffusion Transformer (DiT) [75]. Vision Transformer (VTT) closely follows the original transformer architecture. It first converts 2D images into patches that are flattened and projected into an embedding space. 2D Positional encoding is added to the patch embedding to retain positionna information. This sequence embedding vectors is fed into the standard transformer architecture. Diffusion Transformers (DiT) based VIT architecture and operate on sequences of patches of an image that are projected into a latent space through an image encoder [85]. In addition, DiTs adapt several architectural modifications that enable their as a backbone for diffusion models and help them scale better with increasing model size, including adaptive Layer Normalization (-Zero) [12, 22, 45, 76] for time and class embedding, and peroiinitialization for the final convolution layer [32]-
Deep equilibrium models. Deep Equilibrium models (DEQs) [5] solve for fixed point in the for- ward pass. Specifically, given the input x and the equilibrium function fa. DEQ models approach the infinite-depth representation of So by solving for its fixed point z pass, one can differentiate analytically through =" by the implicit function theorem. The training dynamics of DEQ models can be unstable for certain model designs [7]. As result, recent efforts focus on addressing these issues by designing variants of DEQs with provable guarantees [83, 105]. or through optimization techniques such Jacobian regularization [7], and fixed-point correction [8].
fo(2", For the backward
DEQs have been successfully applied on a wide range of tasks such as image classification [6], semantic segmentation [6. 110], optical flow estimation [8], object detection [100, 102]. landmark detection [69], nut-of-distribun generalization [2], language modelling [5], input optimization [33]. unsupervised learning [98], and generative models [64, 77].
18


Reviews:
[{"review_id": "W7OtB34quI", "reviewer": "Reviewer_BXuM", "summary": "Distillation of diffusion models into smaller models that require fewer steps for generation is an important topic in current research. The authors propose to use a deep equilibrium model as the student. In particular, the student model is a Generative Equilibrium Transformer (GET) that consists of two main modules -- (a) an Injection Transformer that outputs a noise embedding and the (b) Equilibrium Transformer that is trained using the fixed point iteration where the fixed point is the latent embedding of a clean image, which when passed through an image decoder outputs the image output by the diffusion model. GET is trained using a simple L1 loss between the output of GET and the image generated by the diffusion model. At test time, a few steps of performing forward passes through the equilibrium transformer seem sufficient to generate a clean image. The paper also performs a series of experiment demonstrating advantages in terms of data and parameter efficiency, sampling speed etc. compared to existing methods.", "questions": "Please see weaknesses above.", "limitations": "Yes, limitations have been adequately addressed. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is easy to understand at the implementation level and well-written. \n\n2. The authors perform a good set of experiments to understand the trade-offs between data and parameter efficiency, sampling speed, training compute and generation quality. These experiments put the proposed method favorably compared to existing methods. \n\n3.  In particular, with only about three forward passes through the distilled model, with a fraction of the parameters and sampling time compared to a ViT, better FID scores are observed using the proposed method.", "weaknesses": "1. I am not sure if the proposed method can be called one-step generation. My understanding is that some steps through the equilibrium transformer are needed before getting close to convergence to the fixed point. This is also shown in Figure 3. Perhaps some clarification is needed as to what the one-step means in the title. The equilibrium model is trained to produce an image from noise directly, but it still needs multiple passes to get a good quality image.\n\n2. CIFAR-10 is used for the experiments. While the generated images look okay, it is hard to determine their quality as they are small. Evaluation on just CIFAR-10 is weak in my opinion. A dataset with larger images should make the results stronger. \n\n3. Some of the design choices seem arbitrary. For example, why is a separate Injection Transformer necessary? Why the student model has to be a Deep Equilibrium Model is also not clear from the paper. What about training to produce the fixed point is important for distillation? These questions can be partly answered with ablation analysis, but some intuition should also be provided. \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your encouraging feedback! We have tried our best to address all your questions and concerns below.\n\n__Clarification about one-step generation__: We define one-step generation as the ability to generate an image directly from Gaussian noise in a _single_ forward pass through the network.  There can be multiple iterations over the equilibrium transformer during this forward pass but it is important to note that all of these iterations contribute to the depth of the network. In Figure 3(a), we demonstrate how the performance changes with variation in the effective depth of GET. Standard neural network architectures (e.g., ResNet, Diffusion Transformer (DiT) ) increase the depth of the network by applying a block multiple times, increasing the number of parameters in the process. In contrast, GET weight-ties the repeated block, preserving the overall number of parameters. \n\n__Evaluation on larger datasets__: We agree that large scale experiments will provide more comprehensive insights into the capabilities of GET. However, scaling up to ImageNet would require significantly more computing resources. For comparison, consistency distillation uses 64 Nvidia A100 GPUs for experiments on ImageNet-64X64 and LSUN 256X256. Progressive Distillation reports use of 64 TPUv4 chips for their large scale experiments on ImageNet-64X64 and LSUN 256X256. \n\n__Clarification about design choices within GET__: \n\n- __Necessity of the Injection Transformer__: The reviewer has a valid concern about the role of the injection transformer. To address this concern, we trained a model that does not have the injection transformer, but is otherwise architecturally similar to GET. We report the results in the table below. \n\n| Model | Params | FLOPS | FID  |  IS  |\n|----------|:-------------:|:---:|:----:|:----:|\n| GET-M   \t  | 19.2M  | 15.2G  | 10.72 | 8.69 |\n| GET-M-no-inj  |  10.7M | 15.5G  | 14.10 | 8.55 | \n\n\nThus the injection transformer improves GET’s performance.\n\n- __Motivation of the student model to be a DEQ__: Our motivation to model the student network as a DEQ stems from the observation that the relatively complex process of distilling diffusion models has an element of fixed point process, as indicated by recent works like Consistency Models (Song et al. 2023) and TRACT (Berthelot et al. 2023). Ideally, we want a network architecture that has the ability to adapt its compute requirements to tradeoff perceptual quality for compute. DEQs possess both of these capabilities. The empirical results in our paper indeed indicate that modeling the student network (i.e. GET) as a DEQ is vital to achieve good performance in our set-up of one-step distillation.\nWe also include new empirical results that we report in the table below which show that a weight-tied GET outperforms the non-weight tied counterpart in terms of both FID and IS, despite having significantly more parameters. We note that this non-weight-tied model has the same FLOPs as the original weight-tied model. This outcome further underscores that weight-tying within GET indeed contributes to improved performance. \n\n| Model | Params |  FID  |  IS  |\n|----------|:-------------:|:----:| :----:|\n| GET-B   \t    | 62.2M | 7.42 | 9.16 |\n| GET-B-non-WT | 310M  | 7.55 | 9.08 |\n\n- __Why train for fixed point iterations?__ As mentioned in the previous point, modeling GET as a DEQ that solves for the fixed point helps to achieve superior one-step distillation compared to its non-weight tied counterpart. Figure 3 in the main paper also corroborates this, as it indicates that fixed point iterations lead to improved FID scores. \n\nWe are happy to answer any follow up questions and address any concerns you may have.\n"}, {"review_id": "8MXLo3PHkC", "reviewer": "Reviewer_UV2u", "summary": "This paper proposes a new model architecture based on deep equilibrium models and applies it for distilling pretrain diffusion models, by generating samples offline and only utilizing the noise-sample pairs generated by the diffusion models. The architecture consists of embedding, injection transformer, equilibrium transformer and decoder. Empirical results show that the proposed approach leads to comparable performance to the other online distillation approaches, and the architecture is more efficient than ViT architecture in terms of training efficiency, model capacity and inference speed. ", "questions": "- How is z initialized at the start of fix point iteration. Is it sensitive to the initialization of z? \n- In table 4, would be great to include the sampling speed as well. \n", "limitations": "Yes. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is clearly written and easy to follow. \n- The paper shows that deep equilibrium model can be a competitive model class in the context of distilling diffusion models. \n- Extensive comparison with ViT-based models demonstrates the effectiveness of the GET architecture. ", "weaknesses": "- It's not clear to me why the work proposes to apply GET to this specific problem of distilling diffusion models. In principle, this architecture can serve as training a diffusion model directly similar to DiT. And I'm not convinced that offline distillation of diffusion models is a promising direction to work with, given the expensive sampling time the diffusion models need to generate the offline data (and that's why we need distillation). It would be more convincing if the work can show the effectiveness of the proposed architecture in a broader context, or explaining why the proposed architecture is especially useful for distillation. \n- in l.228 it is claimed that other distillation approaches use much larger data budget but I don't think this comparison is fair enough, as the other approaches only require one functional evaluation for generating one sample, while this approach requires go through the whole sampling process of the diffusion model. I'd like to see the comparison of how many functional evaluations of the teacher model is required for different approaches. \n-  I'd like to see the comparison of training/sampling speed with other distillation approaches as increasing sampling speed is the main motivation for doing distillation. I'm a bit concerned that the fixed point iteration would hinder the training/sampling speed compared to the original U-Net architecture. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and good questions. Please find our responses to your questions and concerns below.\n\n__Motivation to use DEQ for distillation__: Our motivation to model the student network as a DEQ stems from the observation that the relatively complex process of distilling diffusion models has an element of fixed point process, as indicated by recent works like Consistency Models (Song et al. 2023) and TRACT (Berthelot et al. 2023). Ideally, we want a network architecture that has the ability to adapt its compute requirements to tradeoff perceptual quality for compute. DEQs possess both of these capabilities. The empirical results in our paper indeed indicate that modeling the student network (i.e. GET) as a DEQ is vital to achieve good performance in our set-up of one-step distillation. \n\n__Why offline distillation?__ Our primary goal is to demonstrate a relatively simple technique for distillation that does not use any trajectory information but can still achieve comparable performance to much more complex distillation techniques like progressive distillation. We generate data (noise/image pairs) in an offline manner to reduce the overhead due to sampling  during training. Further, this data generation needs to be done only once, and the generated  data can be reused later. Offline distillation with GET also requires fewer teacher evaluations compared to progressive distillation and consistency distillation, as we demonstrate in the next point.\n\n__Comparison of overall number of functional evaluations of teacher model__:\nProgressive distillation requires a significantly higher number of function evaluations for the teacher network than offline distillation. We state the exact number of function evaluations of teacher model below:\n\n**GET**: 1M (data samples) * 35 (NFEs for EDM sampling) = 35M  \n**Progressive distillation** (assuming batch size of 128 for each of 8 TPUs):  2 DDIM steps * 128 (batch size) * 8 (TPUs) * (12 passes * 50k (iterations) + 100K (iterations for the last pass of 2 step to 1 step)) = 1.433B  \n**Progressive distillation** (assuming batch size of 128 is shared across 8 TPUs):  2 DDIM steps * 128 (batch size) * (12 passes * 50k + 100K) = 179M  \n**Consistency distillation**:  512 (batch size) * 800k (iterations)= 409.6M. In addition, the perceptual loss requires the same number of network evaluations as the teacher model.\n\nNote that progressive distillation uses DDIM with 8192 steps for CIFAR10.\n\n__Broader experiments with GET as a general architecture__: We acknowledge that GET is a general architecture that might have potential advantages in other applications. Our application of GET in distilling diffusion models was motivated by the inherent fixed point solving mechanism of the distillation process. We empirically demonstrated advantages of GET architecture in distillation over the non-weight tied counterparts. We leave the investigation of performance benefits of GET in other potential applications as future work. \n\n__GET as a backbone architecture for diffusion models__: While both GET and DiT are inspired by ViT, a primary architectural difference lies in the absence of temporal embeddings in GET which makes it tricky to use the current GET as a backbone architecture to train diffusion models from scratch. While it is feasible to incorporate temporal embeddings into GET and train it using denoising loss objective, conventional training of diffusion backbone architectures with this loss objective involves denoising images at a single time step. As a result, the advantages of the fixed point solving mechanism used in DEQs might be less pronounced in this case. In contrast, the process of distilling diffusion trajectories inherently involves a fixed point solving mechanism, which allows the benefits of fixed point solving mechanism to be better highlighted in this specific application.\n\n__Comparison of training and sampling speed with other distillation methods__: \n\nTo ensure a fair comparison between GET and PD, we need to benchmark GET against PD baseline that uses ViT/DiT architecture. However, we do not have access to a ViT/DiT model that is distilled via PD on CIFAR-10. As a proxy for PD, we report the sampling time of EDM (50k images, 1NFE). A single forward pass through EDM requires 27s, whereas GET-Tiny completes sampling in 14.7s, GET-Mini takes 23.3s, and GET-Small needs 35.7s (all under 6 iterations).\n\nConsidering that we use significantly fewer teacher model evaluations than PD (35M vs 179M (lower bound) or 1.433B (upper bound)), the performance can potentially be improved by using more training data and by increasing training iterations. As an example, the GET-Tiny with 8.9M parameters and 14.7s sampling time can reach FID of 11.47 and IS of 8.64 (while PD-EDM has IS of 8.69 at > 410M teacher evaluations), given 4$\\times$ training iterations on *the same data* (35M teacher evaluations).\nIn case, there is a publicly available model that the reviewer would like us to benchmark against, we are happy to evaluate it. \n\n> In Table 4, it would be great to include the sampling speed as well.\n\n The sampling speed of ViT-B is 21.66 secs, and that of GET-Mini is 17.18 secs (4 iterations), 20.16 secs (5 iterations) and 23.27 secs (6 iterations), respectively. We will update Table 4 to include these sampling speeds. \n\n__Initialization of z at the start of fixed point iteration__ : We initialize z at the beginning of fixed point iterations by sampling it from the standard Gaussian distribution. We expect the fixed point iterations to show limited sensitivity to the initial value of z, provided that this value is reasonably set. \n\nPlease reach out to us in case you have any follow up questions/concerns."}, {"review_id": "5w7DGFBJ5r", "reviewer": "Reviewer_GPxt", "summary": "This paper proposes a new model, called Generative Equilibrium Transformer (GET). GET is a deep equilibrium model, trained as an implicit model, to match noise/image pairs generated with a pretrained diffusion model, and thereby distill that pretrained (multi-step) diffusion model into a fast, single-step approach. Experiments on CIFAR-10 demonstrate that this approach is able to distill diffusion models into a fast, efficient architecture that strictly outperforms an explicit implementation.", "questions": "- Can the DEQ approach be used to distill a pretrained UNet (like Stable Diffusion) into a DET?\n- Why no experiments on convolutional architectures, which are very widely used in the community? For example, could a DEQ be initilaized from a pretrained Unet?", "limitations": "adressed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This is a nice paper, with promising results on distillation of diffusion models into very fast, efficient models. Although only conducted on a relatively small-scale dataset, experiments hint at scalability, efficiency; strictly outperforming an explicit counterpart (implemented as a ViT).", "weaknesses": "Overall, this is a strong paper with thorough experiments. A few points, however, remain unaddressed:\n- How does the approach depend on the number of noise/data samples synthesized with the pretrained diffusion model? For conditional models, do we need to adapt the sampling strategy? \n- While experiments on scalability of the model size are conducted on a small scale, the paper would benefit massively from applying the method to some larger pretrained model, ideally on a more demanding task like text-to-image (although intermediate steps on either (i) higher-resolution images, e.g. LSUN, or FFHQ, or more complex tasks like class-conditional ImageNet could suffice).  \n\nFor further comments, please refer to the \"Questions\" section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are encouraged to know that the reviewer feels that this paper is strong with thorough experiments! We have tried our best to answer your questions below.\n\n__Variation of performance with number of samples__: \nGiven this is a supervised learning set up, we anticipate that GET’s generalization will improve as we increase the number of noise/image pairs, up to a certain threshold, beyond which saturation is likely to occur.  To test this hypothesis empirically, we sampled an additional 1M data pairs for GET in the conditional setting and trained GET-B using the total 2M data. The resulting model achieved an FID of **5.66** and IS of **9.63**, while the model trained on 1M data has an FID of 6.23 and IS of 9.42. We believe that further scaling up of training data is certainly possible. GET might also potentially benefit from better supervision through perceptual loss.\n\n__Sampling strategy for conditional models__: For class-conditional sampling, we need to provide the class label in addition to the initial Gaussian noise as inputs to GET. GET generates images in a single forward pass for both class-conditional and unconditional cases. As GET learns a direct mapping from the initial noise and conditioning to the resulting images, we don’t need to adapt a special sampling strategy to generate class-conditioned images. This is a potential advantage of GET over distillation methods like PD in conditional settings.\n\n__Large scale scalability experiments__: We agree that conducting experiments on larger datasets would indeed provide more comprehensive insights into the scaling capabilities of GET. However, we lack computational resources to perform such large scale experiments. For instance, consistency distillation uses 64 Nvidia A100 GPUs for experiments on ImageNet-64x64 and LSUN 256x256. Progressive distillation reports use of 64 TPUv4 chips for their large scale experiments on ImageNet-64X64 and LSUN 256X256.\n\n__Can the DEQ approach be used to distill a pretrained UNet (like Stable Diffusion) into a GET?__ We can distill any pretrained diffusion model regardless of the network architecture into GET as long as we have noise/image pairs from the pretrained diffusion model. \n\n__Experiments on convolutional architectures like UNet__: Offline distillation for convolutional architectures has been explored by previous works like knowledge distillation (Luhman & Luhman, 2021). Knowledge distillation distills DDIM sampling chains with 1000 steps and uses UNet architecture. We outperform knowledge distillation on metrics like FID Score and Inception Score.\n\n__Initializing GET with pretrained UNet__: While we cannot initialize GET network with weights of a pretrained UNet, we can use noise/image pairs from (multi-step) UNet to distill diffusion trajectories into GET.\n\nPlease let us know if you have any follow up questions and other concerns. We are happy to answer them! \n"}, {"review_id": "qpsMgVcDrN", "reviewer": "Reviewer_FiXo", "summary": "The submission proposes the Generative Equilibrium Transformer (GET), a lightweight refinement of vision transformer that is well-suited as an efficient single-step student model for diffusion distillation. The author empirically shows that the GET outperforms classic networks in terms of performance, model size and inference efficiency. Overall, the introduced architecture improvements shed light on further studies for improving student models of diffusion distillation.", "questions": "See the weakness above.", "limitations": "The author has discussed the limitations.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The introduction of DEQ for efficient students is novel: The idea of introducing the DEQ model into the design of a more efficient student model is novel. Till now, most diffusion distillation works take the same architecture of the student model as the teacher model. Few of them have investigated the design of an efficient student model, which is potentially important for applications with strict requirements on inference efficiency. The author brings the DEQ model to design a ViT-based encoder-decoder student network and shows promising performance on CIFAR10 and ImageNet64 data. \n2. Good Writing: The writing of this paper is clear and easy to follow.", "weaknesses": "Insufficient evaluations: the paper proposes an architecture-level improvement, the GET, that is related to ViT. The readers will be more convinced that the GET should be compared with the ViT-based diffusion model on modeling tasks on a large scale, such as the DiT on ImageNet128 and 256. I think the pros and cons of the proposed GET can be observed more clearly for large-scale benchmarks. Besides, the one-step student model shows an FID of 6.91 on the CIFAR10 dataset, which is significantly worse than Consistency Distillation which has an FID of 3.55. I am curious about the possibility to combine GET with other distillation methods such as CD to obtain stronger performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your encouraging feedback! Please find our responses to your questions and concerns below.\n\n__Extensive Large Scale Evaluations__: It is certainly possible to scale up to larger datasets like ImageNet but this would require significantly more computing resources. For example, even training Progressive Distillation on ImageNet 64 would consume 5 days on 64 TPUv4 chips, which is equivalent to/more than 7680 A100 hours. We have tried our best to demonstrate the scaling capabilities of GET in terms of training FLOPs and parameters (See Figure 3a/b) to support the possibility of transferring GET to more complex settings.\n\n__Comparison with CD__: Thanks for pointing out that CD is a strong baseline. However, CD uses far more complex distillation settings compared to GET. As shown in our paper, \n- _Complexity_: CD leverages trajectory information from the teacher diffusion model and employs perceptual loss (LPIPS), which results in forward passes through 3 models during the distillation process. \n- _Efficiency_: On CIFAR10, CD requires significantly more teacher models evaluations  (512 (batch size) * 800k iterations = 409.6M evaluations) compared to our method which consumes 1M samples * 35 NFE = 35M teacher model evaluations. \n- _Data requirement_: CD assumes access to the dataset on which the original diffusion model was trained. In contrast, GET can be trained fully on synthetic data. This is beneficial in cases where we do not have access to the original dataset on which the diffusion model was trained.\n- _Auxiliary loss_: CD uses LPIPS perceptual loss as an auxiliary loss during training. In the absence of this loss, CD’s performance drops considerably (even doubles). Please see Figure 4 in CD’s paper.\n\n__Combining GET with other distillation methods__: We agree that combining GET with other distillation techniques like consistency distillation is a promising direction. However, pursuing this avenue necessitates training GET as a diffusion model with denoising loss, and followed by distilling it into a few-step generative model using CD or PD. This process will require multiple model training phases, rendering the overall process more complex. This complexity runs  contrary to our initial motivation of using a DEQ-based architecture for diffusion distillation in a single-forward pass.\n\nWe are happy to answer any follow up questions that you might have.\n"}, {"review_id": "hEYdWuCcrw", "reviewer": "Reviewer_sYNQ", "summary": "The paper proposes a simple approach to distill diffusion models into generative models capable of sampling with just a single model evaluation. The method involves training a Generative Equilibrium Transformer (GET) architecture directly on noise/image pairs generated from a pre-trained diffusion model, eliminating the need for trajectory information and temporal embedding. The experiments verify its effectiveness.", "questions": "- Could you explain more about the motivation and advantages of using DEQ? \n- Following the prior question, about the comparison in Table 4: have you tried the diffusion transformer (DiT), which may be more parameter-efficient for denoising process? The same question for the scaling experiments.\n- In Figure 2(Left), the performance of GET falls off quickly with less iterations in the Equilibrium transformer. Compared to other distillation methods such as PD and CD with single step, what are the advantages? Are there any results about the sampling speed using different iterations between PD(CD) and GET? \n- Instead of 6 iterations, have you tried two GET models for two stage of denoising steps and used 2 iterations for two GET models sperately?($t$ from 0 to 0.5 and $t$ from 0.5 to 1.0). This may be more efficient than using 6 iterations.\n- Does GET drop condition during training and support classifier-free guidance? Can this technique improve the performance?", "limitations": "The authors have adequately addressed the limitations.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow.\n- The idea of using DEQ for distillation is interesting.", "weaknesses": "To verify the effectiveness, the results of ImageNet are neccessary, since ImageNet is another mainstream benchmark. Some experimental design are not convincing enough. More details please refer to the questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for well-thought questions and valuable feedback. We have tried our best to answer your questions.\n\n__Motivation and Advantages of DEQ__: \nOur motivation to model the student network as a DEQ stems from the observation that the relatively complex process of distilling diffusion models has an element of fixed point process, as indicated by recent works like Consistency Models (Song et al. 2023) and TRACT (Berthelot et al. 2023).  In this work, we sought a network architecture capable of adapting its compute requirements according to the complexity of image being generated. DEQs possess both of these capabilities. The empirical results in our paper indeed indicate that modeling the student network (i.e. GET) as a DEQ is vital to achieve good performance in our set-up of one-step distillation. \n\n__Benchmarking against DiT__: In Table. 4, the ViT baseline was simplified from the diffusion transformer (DiT). Because our distillation strategy does not require the temporal embedding, we remove this layer from DiT, leading to the ViT baseline shown in Table. 4. \n\n__Clarification about one-step generation__: We define one-step generation as the ability to generate an image directly from Gaussian noise in a _single_ forward pass through the network.  There can be multiple iterations over the equilibrium transformer during this forward pass but it is important to note that all of these iterations contribute to the depth of the network. In Figure 3(a), we demonstrate how the performance changes with variation in the effective depth of GET. Standard neural network architectures (e.g., ResNet, Diffusion Transformer (DiT)) increase the depth of the network by applying a block multiple times, increasing the number of parameters in the process. In contrast, GET weight-ties the repeated block, preserving the overall number of parameters. \n\n__Variation of GET’s performance with iterations__: As shown in Figure 3a (left), GET’s performance improves with more iterations because of the increase in effective depth of the network. However, this computation still constitutes a single forward pass through the network, and thus one-step generation. \n\n__Advantages of GET over other distillation methods such as PD and CD with single step__: \n\n__Advantages over PD__: GET has several advantages over Progressive Distillation (PD).\n- _Complexity_: On CIFAR10, PD uses 13 training passes to distill 8192-step DDIM chains to 1 step. In addition, PD also carefully selects distillation targets in each step.\n- _Performance_: One-step GET (that uses a single forward pass) outperforms PD for both class-conditional and unconditional cases on CIFAR-10 in terms of FID score and Inception Score. \n- _Efficiency_: On CIFAR10, PD consumed 179M teacher model evaluations (2 DDIM steps * 128 (batch size) * (12 passes * 50k + 100K) = 179M), while GET only consumed 1M samples * 35 NFE = 35M teacher model evaluations.\n\n__Advantages over CD__: CD uses far more complex distillation settings compared to GET. As shown in our paper, \n- _Complexity_: CD leverages trajectory information from the teacher diffusion model and employs perceptual loss (LPIPS), which results in forward passes through 3 models during the distillation process. \n- _Efficiency_: On CIFAR10, CD requires significantly more teacher models evaluations  (512 (batch size) * 800k iterations = 409.6M evaluations) compared to our method which consumes 1M samples * 35 NFE = 35M teacher model evaluations. \n- _Data requirement_: CD assumes access to the dataset on which the original diffusion model was trained. In contrast, GET can be trained fully on synthetic data. This is beneficial in cases where we do not have access to the original dataset on which the diffusion model was trained.\n- _Auxiliary loss_: CD uses LPIPS perceptual loss as an auxiliary loss during training. In the absence of this loss, CD’s performance drops considerably (even doubles). Please see Figure 4 in CD’s paper.\n\n__Sampling speed of GET against PD and CD__: \nTo ensure a fair comparison between GET and PD, we need to benchmark GET against PD baseline that uses ViT/DiT architecture. However, we do not have access to a ViT/DiT model that is distilled via PD on CIFAR-10. Similar concerns hold true for CD. As a proxy, we report the sampling time of EDM (50k images, 1 NFE). A single forward pass through EDM requires 27s, whereas GET-Tiny completes sampling in 14.7s, GET-Mini takes 23.3s, and GET-Small needs 35.7s (all under 6 iterations).\n\nConsidering that we use significantly fewer teacher model evaluations than PD (35M vs 179M (lower bound) or 1.433B (upper bound)), the performance can potentially be improved by using more training data and by increasing training iterations. As an example, the GET-Tiny with 8.9M parameters and 14.7s sampling time can reach FID of 11.47 and IS of 8.64 (while PD-EDM has IS of 8.69 at > 410M teacher evaluations), given 4$\\times$ training iterations on *the same data* (35M teacher evaluations).\nIn case, there is a publicly available model that the reviewer would like us to benchmark against, we are happy to evaluate it.\n\n__Two stage denoising with GET__: GET learns a direct mapping between noise and image pairs without any use of temporal embeddings. As a result, we cannot stack two GET models to get two-state denoising. Further, in our paper, we demonstrate empirically that GET benefits from its DEQ architecture where we solve for a fixed point. Reducing the number of iterations in the Equilibrium transformer to 3 will lead to poor convergence of this fixed point, and might adversely affect the performance of GET.\n\n__Classifier-free guidance while training GET__: Thanks for this suggestion! It is certainly possible to train GET with classifier guidance. This is an interesting experiment that we will try in the future.\n"}, {"review_id": "yShWDHBBP9", "reviewer": "Reviewer_MHac", "summary": "This paper applies Deep Equilibrium Models to distillation of (conditioned) diffusion models. The key part of the architecture consists of a repeated application of weight-tied block of layers on the internal activations (theoretically until convergence to a fixed point, but in practice a few iterations). The model is trained to mimic the mapping from noise to images induced by the diffusion sampling chain.", "questions": "Beyond the baselines discussed above, the role of the fixed point remains somewhat unclear either way. Does the iteration converge to a fixed point in six iterations? Could it even diverge? If one were to evaluate it all the way to a fixed point, how much extra performance could be squeezed out? Is six steps enough in training to reap these benefits?\n\nMiscellaneous observation: are you sure you are citing the correct FID/IS for Progressive Distillation [61] in unconditional CIFAR (Table 1)? If I understand correctly, the number you cite is the conditional model (Table 4 in appendix of [50]), which benefits from the conditioning info despite lack of guidance boost. However, for the unconditional model, Table 2 in [61] cites a FID of 9.12. Wouldn’t this be the correct number (not sure)? Related, I’m not sure why this model is referred to as PD-EDM, and cited in combination with Consistency Models?", "limitations": "There is a brief discussion of limitations, but it does not go into much depth.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The premise of the paper makes sense. One important reason for the success of diffusion models might be that the chain of repeated (though time-varying) neural steps is, in aggregate, able to represent extremely complicated mappings, with near-discontinuous regions packed in complex high-dimensional shapes, etc. – something that is amply present in the latent-to-image mapping induced by the diffusion flow. Capturing this complexity with typical single-pass networks of modest layer counts is challenging, as evidenced by many attempts in the literature. DEQ seems like an appealing and natural approach here, as it has the similar character of iteratively shaping an initially simple distribution into a very complex one. There is also the promise of parameter count savings: effectively they implement a more complex function than just a single layer block, but parametrized on the same amount of weights.\n\nThe proposed network architectures seem well justified for this task.\n\nThe evaluation makes a reasonably convincing case that this can be beneficial, and that the promised benefits of DEQs are realized up to a point. There is a decent amount of useful experimentation and attempt to understand scaling around model size and other parameters. ", "weaknesses": "The paper is scarce on details about the DEQ training and inference. The impression I’m left with (based on line 213 and the fact that no details are provided), is the authors do not actually use almost any of the DEQ machinery, such as fixed point solvers and related backpropagation shortcuts. Instead, the model seems to be trained as though it were a regular neural network, where the same block of layers is stacked in sequence six times (but with tied weights between them). This may work well in practice but raises some questions about whether DEQs are as relevant to the story as claimed, and whether some of their potential advantages are left on the table. There does seem to be fixed-point like behavior and capability though, as illustrated in Figure 3a.\n\nNevertheless, this raises the question about some baselines. What if the weights were _not_ tied? Wouldn’t the inference and training remain equal number of flops to what you currently do, but with some modest extra memory consumption for storing the weight values (but importantly, no more or less activations to store)? By default, wouldn’t one expect this to strictly extend the capability of the model? It’s possible that the tying has some implicit regularization effects like discussed on line 224, but this claim is speculative. It might also make the training more effective (as in, converge with fewer iterations), but this is also not clear-cut. The evaluation does not currently explore or rule out the possibility that weight tying could even be harming an otherwise higher performing model. Such a baseline would distinguish whether the benefit comes from the particular architecture and depth, or from the DEQ-like aspects. At present, I am somewhat uncertain about which it is – the current ViT baseline seems sufficiently different that it might underperform for unrelated reasons.\n\nThis would of course increase the parameter count, but note that parameter count is often taken as a proxy for evaluation time and memory consumption, neither of which would be seriously affected here. That is, unless my understanding is wrong, and the paper actually does use DEQ training mechanisms. In this case, it is missing a significant amount of important detail about how those are implemented and how the associated difficulties are dealt with (this should be clarified in either case).\n\nThe water is also a bit muddy in terms of how many iterations any given distilled model requires. One could argue that this is in spirit a six-step model (flanked by pre and post-processing to embed in a latent space), or conversely, that e.g. a four-step progressive distillation model is actually just a one-step model that internally calls some layers four times. The more relevant question comes down to flops and wall-clock time rather than the semantics of what we take the iterations to mean. Accordingly, including something like 2 or 4 step PD (which beat the metrics here) in the comparisons would be reasonable, unless it can be convincingly argued that their runtime is significantly slower due to the sub-stepping (I am unsure of whether this is the case).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your extremely thoughtful feedback and suggestions. We have tried our best to answer your questions and concerns below.\n\n> Does the iteration converge to a fixed point in six iterations? \n\nWe report the relative fixed point error $\\frac{\\|| f(z) - z \\||} {\\|| f(z) \\||}$ and the resulting FID score as we vary the number of iterations of the equilibrium transformer within GET-M at inference time.\n\n| Iterations | Relative Error  |  FID |\n|----------|:-------------:|------:|\n| 1 |  1.028 |  42.96 |\n| 2 |  0.334 |  26.70 |\n| 3 |  0.215 |  20.27 |\n| 4 |  0.148 |  17.21 |\n| 5 |  0.107 |  15.77 |\n| 6 |  0.082 |  15.18 |\n| 7 |  0.066 |  15.04 |\n| 8 |  0.055 |  15.18 |\n| 9 |  0.048 |  15.46 |\n| 10 |  0.043 | 15.85 |\n\nBased on our empirical results, we have indeed converged to a fixed point with a relatively low relative fixed point error of 0.082 within 6 iterations. As a comparison, DEQ-transformer [1] trained with gradients computed via implicit function theorem requires 30 iterations to achieve a relative error of 0.10 on language modeling.\n\nWe also observe improved FID scores as we increase the depth to reach the depth at training time. The FID score eventually flattens out and marginally increases as we further increase the depth. As the reviewer has noted, it is certainly possible to get improved performance from GET by marginally increasing the depth at inference time. For instance, GET-M’s FID score improves to 15.04 when we increase the iterations at inference time to 7. \n\n__Does weight-tying harm GET’s performance?__: \n\nThe reviewer has valid concerns on whether weight tying harms GET’s performance. To address this concern, we trained and evaluated a non-weight tied architecture. In this architecture, we unroll the equilibrium transformer in GET, effectively using six separate GET blocks, instead of iterating over a single GET block within the equilibrium transformer. This architecture is otherwise exactly similar to GET’s architecture, and thus has the same number of FLOPs. We trained this model for 800k iterations, similar to the training duration of the weight-tied model, and found this model was 16%  slower in training, despite having the same FLOPs. Further, This model also achieves an inferior FID score and Inception score compared to the weight-tied model. This outcome underscores that weight-tying within GET indeed contributes to improved performance. \n\n| Model | Params |  FID  |  IS  |\n|----------|:-------------:|:----:| :----:|\n| GET-B   \t    | 62.2M | 7.42 | 9.16 |\n| GET-B-non-WT | 310M  | 7.55 | 9.08 | \n\n__Comparison of runtime of GET against PD__: \n\nTo ensure a fair comparison between GET and PD, we need to benchmark GET against PD baseline that uses ViT/DiT architecture. However, we do not have access to a ViT/DiT model that is distilled via PD on CIFAR-10. As a proxy for PD, we report the sampling time of EDM (50k images, 1NFE). A single forward pass through EDM requires 27s, whereas GET-Tiny completes sampling in 14.7s, GET-Mini takes 23.3s, and GET-Small needs 35.7s (all under 6 iterations).\n\nConsidering that we use significantly fewer teacher model evaluations than PD (35M vs 179M (lower bound) or 1.433B (upper bound)), the performance can potentially be improved by using more training data and by increasing training iterations. As an example, the GET-Tiny with 8.9M parameters and 14.7s sampling time can reach FID of 11.47 and IS of 8.64 (while PD-EDM has IS of 8.69 at > 410M teacher evaluations), given 4$\\times$ training iterations on *the same data* (35M teacher evaluations).\n\n[1] Bai, Shaojie, J. Zico Kolter, and Vladlen Koltun. \"Deep equilibrium models.\" Advances in Neural Information Processing Systems 32 (2019).\n\n> Are you sure you are citing the correct FID/IS for Progressive Distillation [61] in unconditional CIFAR (Table 1)? \n\nThank you for requesting this clarification. In Table 1, we have cited PD [61] at FID 9.12 (second row in Diffusion Distillation). We have also cited guided distillation [50] in Table 2, which is the performance for class-conditional models. The results for PD-EDM in Table 1 have been reported by Consistency Models [70] using EDM and stronger training settings, and we include it as a relevant baseline.\n\nWe are happy to answer any further questions or concerns that you may have.\n"}]

Author rebuttal:
We thank all the reviewers for their thoughtful comments and suggestions. Here, we address some of the common concerns raised by the reviewers.

__What does one-step generation mean?__ We define one-step generation as the ability to generate an image directly from Gaussian noise in a _single_ forward pass through the network. While there can be multiple iterations within the equilibrium transformer in this single forward pass, all of these iterations contribute to the “effective depth” of the network.

__Evaluation on larger datasets__: We agree with the reviewers that experiments on larger datasets like ImageNet will significantly strengthen the results. However, scaling up to ImageNet would require significantly more computing resources. For comparison, consistency distillation uses 64 Nvidia A100 GPUs for experiments on ImageNet-64$\times$64 and LSUN 256$\times$256. Progressive distillation needs 5 days on 64 TPUv4 chips for experiments on ImageNet-64$\times$64 and LSUN 256$\times$256. This is equivalent to/more than 7680 A100 GPU hours.  We have tried our best to demonstrate the scaling capabilities of GET in terms of training FLOPs and parameters (See Figure 3a/b) to support the possibility of transferring GET to more complex settings.

__New result on teacher network evaluations__: Offline distillation with GET requires significantly fewer teacher model evaluations (35M) compared to progressive distillation (179M (lower bound) or 1.433B (upper bound)) and consistency distillation (409.6M).



Meta-review: This paper proposes a Deep Equilibrium Models (DEM) based architecture for offline one-step distillation of diffusion models. Experimentally, it is shown that such an architecture is favorable compared to a standard ViT, and other distillation methods such as PD. The paper has received an average rating of 5.33, which is a borderline score, but all reviewers are leaning towards acceptance. The AC agrees that there is merit to the proposed method, but the paper could be further strengthened with more evaluations as suggested by the reviewers. In addition to what's being discussed by the reviewers, the AC believes that it would be beneficial to conduct a direct comparison to knowledge distillation (Luhman & Luhman, 2021), while evaluating both with the same teacher model. The AC recommend accepting the paper, but strongly encourage the authors to incorporate the suggestions from the reviewers in the final version.
Decision: Accept (poster)