SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhengyang Geng*  
Regie Mellon University
zgeng2@cs.cmu.edu
Ashwini Pokle* Carnegie Mellon University apokle@cs.cmu.edu
Bosch Center for AI
zkolter@cs.cmu.edu
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 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.
Diffusion models [35, 93, 95, 96] have demonstrated remarkable performance on a wide range of generative tasks such as 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 a 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 a more stable training procedure, which in turn allows diffusion models to 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-constrained scenarios.
In an 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 a more efficient few-step or single-step process. However, these techniques often come with their
*Equal Contribution. Correspondence to Zhengyang Geng and Ashwini Pokle
own set of challenges. The distillation targets must be carefully designed to successfully transfer knowledge from the larger model to the smaller one. Further, distilling a long sampling process into a 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 to 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 a multi-step diffusion process into a 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-tied transformer layers, which solve for a fixed point in the forward pass. This architectural choice allows for the adaptive application of these layers in the forward pass, striking a balance between inference speed and sample quality. Furthermore, we incorporate an almost parameter-free 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 to 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 us 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 a  $5 \times$  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 transformer 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.
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:
$$
\mathbf {z} ^ {[ i + 1 ]} = f _ {\theta} ^ {[ i ]} \left(\mathbf {z} ^ {[ i ]}; \mathbf {x}\right) \quad \text {f o r} i = 0, \dots , L - 1 \tag {1}
$$
where  $\mathbf{x} \in \mathbb{R}^{n_x}$  is the input injection,  $\mathbf{z}^{[i]} \in \mathbb{R}^{n_z}$  is the hidden state of  $i^{th}$  layer, and  $f_{\theta}^{[i]}: \mathbb{R}^{n_x \times n_z} \mapsto \mathbb{R}^{n_z}$  is the feature transformation of  $i^{th}$  layer, parametrized by  $\theta$ . If the above model is weight-tied, i.e.,  $f_{\theta}^{[i]} = f_{\theta}, \forall i$ , then in the limit of infinite depth, the output  $\mathbf{z}^{[i]}$  of this network approaches a fixed point  $\mathbf{z}^{\star}$ :
$$
\lim  _ {i \rightarrow \infty} f _ {\theta} \left(\mathbf {z} ^ {[ i ]}; \mathbf {x}\right) = f _ {\theta} \left(\mathbf {z} ^ {\star}; \mathbf {x}\right) = \mathbf {z} ^ {\star} \tag {2}
$$
Deep equilibrium (DEQ) models [5] directly solve for this fixed point  $\mathbf{z}^{\star}$  using black-box root finding algorithms like Broyden's method [14], or Anderson acceleration [1] in the forward pass. DEQs utilize implicit differentiation to differentiate through the fixed point analytically. Let  $g_{\theta}(\mathbf{z}^{\star};\mathbf{x}) = f_{\theta}(\mathbf{z}^{\star};\mathbf{x}) - \mathbf{z}^{\star}$ , then the Jacobian of  $\mathbf{z}^{\star}$  with respect to the model weights  $\theta$  is given by
$$
\frac {\partial \mathbf {z} ^ {\star}}{\partial \theta} = - \left(\frac {\partial g _ {\theta} \left(\mathbf {z} ^ {\star} , \mathbf {x}\right)}{\partial \mathbf {z} ^ {\star}}\right) ^ {- 1} \frac {\partial f _ {\theta} \left(\mathbf {z} ^ {\star} ; \mathbf {x}\right)}{\partial \theta} \tag {3}
$$
Figure 1: Generative Equilibrium Transformer (GET). (Left) GET consists of two major components: 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 of transformer blocks in 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 high-dimensional 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] or score-based generative models [95, 96] progressively perturb images with an increasing amount of Gaussian noise and then reverse this process through sequential denoising to generate images. Specifically, consider a dataset of i.i.d. samples  $p_{\mathrm{data}}$ , then the diffusion process  $\{\mathbf{x}(t)\}_{t=0}^{T}$  for  $t \in [0, T]$  is given by an Ito SDE [96]:
$$
\mathrm {d} \mathbf {x} = \mathbf {f} (\mathbf {x}, t) \mathrm {d} t + g (t) \mathrm {d} \mathbf {w} \tag {4}
$$
where  $\mathbf{w}$  is the standard Wiener process,  $\mathbf{f}(\cdot ,t):\mathbb{R}^d\to \mathbb{R}^d$  is the drifting coefficient,  $g(\cdot):\mathbb{R}\rightarrow \mathbb{R}$  is the diffusion coefficient, and  $\mathbf{x}(0)\sim p_{\mathrm{data}}$  and  $\mathbf{x}(T)\sim \mathcal{N}(0,I)$ . All diffusion processes have a 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:
$$
\mathrm {d} \mathbf {x} = - \dot {\sigma} (t) \sigma (t) \nabla_ {\mathbf {x}} \log p (\mathbf {x}, \sigma (t)) \mathrm {d} t \tag {5}
$$
where  $\sigma(t)$  is the noise schedule of diffusion process, and  $\nabla_{\mathbf{x}} \log p(\mathbf{x}, \sigma(t))$  represents the score function. Karras et al. [48] show that the optimal choice of  $\sigma(t)$  in Eq. (5) is  $\sigma(t) = t$ . Thus, the PF-ODE can be simplified to  $\mathrm{d}\mathbf{x} / \mathrm{d}t = -t\nabla_{\mathbf{x}} \log p(\mathbf{x}, \sigma(t)) = (\mathbf{x} - D_{\theta}(\mathbf{x}; t)) / t$ , where  $D_{\theta}(\cdot, t)$  is a denoiser function parametrized with a neural network that minimizes the expected  $L_2$  denoising error for samples drawn from  $p_{\mathrm{data}}$ . Samples can be efficiently generated from this ODE through numerical methods like Euler's method, Runge-Kutta method, and Heun'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  $\mathbf{e}$  and optional class labels  $\mathbf{y}$  to images  $\tilde{\mathbf{x}}$ . The major components of GET include an injection transformer (InjectionT, Eq. (7)) and an equilibrium transformer (EquilibriumT, Eq. (8)). The InjectionT transforms tokenized noise embedding  $\mathbf{h}$  to an intermediate representation  $\mathbf{n}$  that serves as the input injection for the equilibrium transformer. The EquilibriumT, which is the equilibrium layer, solves for the fixed point  $\mathbf{z}^{\star}$  by taking in the noise injection  $\mathbf{n}$  and an optional class embedding  $\mathbf{c}$ . Finally, this fixed point  $\mathbf{z}^{\star}$  is decoded and rearranged to generate an image sample  $\tilde{\mathbf{x}}$  (Eq. (9)). Figure 1 provides an overview of the GET architecture. Note that because we are directly distilling the entire generative process, there is no need for a time embedding  $t$  as is common in standard diffusion models.
$$
\mathbf {h}, \mathbf {c} = \operatorname {E m b} (\mathbf {e}), \operatorname {E m b} (\mathbf {y}); \text {i f} \mathbf {y} \notin \emptyset \tag {6}
$$
$$
\mathbf {n} = \operatorname {I n j e c t i o n T} (\mathbf {h}, \mathbf {c}) \tag {7}
$$
$$
\mathbf {z} ^ {\star} = \text {E q u i l i b r i u m T} (\mathbf {z} ^ {\star}, \mathbf {n}, \mathbf {c}) \tag {8}
$$
$$
\tilde {\mathbf {x}} = \operatorname {D e c o d e r} \left(\mathbf {z} ^ {\star}\right) \tag {9}
$$
Noise Embedding. GET first converts an input noise  $\mathbf{e} \in \mathbb{R}^{H \times W \times C}$  into a sequence of 2D patches  $\mathbf{p} \in \mathbb{R}^{N \times (P^2 \cdot C)}$ , 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 / P^2$  is the resulting number of patches. Let  $D = P^2 \cdot C$  denote the width of the network. We follow ViT to use a linear layer to project the  $N$  patches to  $D$  dimensional embedding. We add standard sinusoidal position encoding [99] to produce the noise embedding  $\mathbf{h}$ . Position encoding plays a crucial role in capturing the spatial structure of patches by encoding their relative positional information.
InjectionT & EquilibriumT. Both InjectionT and EquilibriumT are composed of a sequence of Transformer blocks. InjectionT is called only once to produce the noise injection  $\mathbf{n}$ , while EquilibriumT defines the function  $f_{\theta}$  of the implicit layer  $\mathbf{z}^{\star} = f_{\theta}(\mathbf{z}^{\star}, \mathbf{n}, \mathbf{c})$  that is called multiple times—creating a weight-tied computational graph—until convergence. A linear layer is added at the end of InjectionT to compute the noise injection  $\mathbf{n}_l \in \mathbb{R}^{N \times 3D}$ ,  $l \in [L_e]$ , for each of the  $L_e$  GET blocks in EquilibriumT. For convenience, we overload the notation  $\mathbf{n}_l$  and  $\mathbf{n}$ , in the subsequent paragraphs.
Transformer Block. GET utilizes a near-identical block design for the noise injection (InjectionT) and the equilibrium layer (EquilibriumT), differing only at the injection interface. Specifically, the transformer block is built upon the standard Pre-LN transformer block [25, 75, 106], as shown below:
$$
\mathbf {z} = \mathbf {z} + \text {A t t e n t i o n} (\mathrm {L N} (\mathbf {z}), \mathbf {u})
$$
$$
\mathbf {z} = \mathbf {z} + \text {F F N} (\ln (\mathbf {z}))
$$
Here,  $\mathbf{z} \in \mathbb{R}^{N \times D}$  represents the latent token,  $\mathbf{u} \in \mathbb{R}^{N \times 3D}$  is the input injection, LN, FFN, and Attention stand for Layer Normalization [4], a 2-layer Feed-Forward Network with a hidden dimension of size  $D \times E$ , and an attention [99] layer with an injection interface, respectively.
For blocks in the injection transformer,  $\mathbf{u}$  is equal to the class embedding token  $\mathbf{c} \in \mathbb{R}^{1 \times 3D}$  for conditional image generation, i.e.,  $\mathbf{u} = \mathbf{c}$  for conditional models, and  $\mathbf{u} = \mathbf{0}$  otherwise. In contrast, for blocks in the equilibrium transformer,  $\mathbf{u}$  is the broadcast sum of noise injection  $\mathbf{n} \in \mathbb{R}^{N \times 3D}$  and class embedding token  $\mathbf{c} \in \mathbb{R}^{1 \times 3D}$ , i.e.,  $\mathbf{u} = \mathbf{n} + \mathbf{c}$  for conditional models and  $\mathbf{u} = \mathbf{n}$  otherwise.
We modify the standard transformer attention layer to incorporate an additive injection interface before the query  $\mathbf{q} \in \mathbb{R}^{N \times D}$ , key  $\mathbf{k} \in \mathbb{R}^{N \times D}$ , and value  $\mathbf{v} \in \mathbb{R}^{N \times D}$ ,
$$
\mathbf {q}, \mathbf {k}, \mathbf {v} = \mathbf {z} \mathbf {W} _ {i} + \mathbf {u}
$$
$$
\mathbf {z} = \operatorname {M H A} (\mathbf {q}, \mathbf {k}, \mathbf {v})
$$
$$
\mathbf {z} = \mathbf {z} \mathbf {W} _ {o}
$$
where  $\mathbf{W}_i\in \mathbb{R}^{D\times 3D}$ ,  $\mathbf{W}_o\in \mathbb{R}^{D\times D}$ . The injection interface enables interactions between the latent tokens and the input injection in the multi-head dot-product attention (MHA) operation,
$$
\mathbf {q} \mathbf {k} ^ {\top} = \left(\mathbf {z} \mathbf {W} _ {q} + \mathbf {u} _ {q}\right) \left(\mathbf {z} \mathbf {W} _ {k} + \mathbf {u} _ {k}\right) ^ {\top} = \mathbf {z} \mathbf {W} _ {q} \mathbf {W} _ {k} ^ {\top} \mathbf {z} ^ {\top} + \mathbf {z} \mathbf {W} _ {q} \mathbf {u} _ {k} ^ {\top} + \mathbf {u} _ {q} \mathbf {W} _ {k} ^ {\top} \mathbf {z} ^ {\top} + \mathbf {u} _ {q} ^ {\top} \mathbf {u} _ {k}, \tag {10}
$$
where  $\mathbf{W}_q$ ,  $\mathbf{W}_k \in \mathbb{R}^{D \times D}$  are slices from  $\mathbf{W}_i$ , and  $\mathbf{u}_q$ ,  $\mathbf{u}_k \in \mathbb{R}^{N \times D}$  are slices from  $\mathbf{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.
Figure 2: Data and Parameter Efficiency of GET:(a) (Left) GET outperforms PD and a  $5 \times$  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 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.
Image Decoder. 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  $\bar{\mathbf{p}}\in \mathbb{R}^{N\times D}$ . The resulting patches  $\bar{\mathbf{p}}$  are rearranged back to the resolution of the input noise  $\mathbf{e}$  to produce the image sample  $\hat{\mathbf{x}}\in \mathbb{R}^{H\times W\times C}$ . Thus, the decoder maps the features back to the image space.
We evaluate the effectiveness of our proposed Generative Equilibrium Transformer (GET) in offline distillation of diffusion models through a series of experiments on 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 ViTs and GETs of varying scales on 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.
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 config and training specifics, please refer to the Appendix.
Data Collection. For unconditional image generation on CIFAR-10 [52], we generate 1M noise/image 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 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 4 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 noise-label/image pairs from the conditional VP-EDM Karras et al. [48] using the same settings. This dataset is denoted as EDM-Cond-1M. Both the datasets will be released for future studies.
Offline Distillation. We distill a pretrained EDM [48] into ViTs and GETs by training on a dataset  $\mathcal{D}$  with noise/image pairs sampled from the teacher diffusion model using a reconstruction loss:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {\mathbf {e}, \mathbf {x} \sim \mathcal {D}} \| \mathbf {x} - G _ {\theta} (\mathbf {e}) \| _ {1}
$$
where  $\mathbf{x}$  is the desired ground truth image,  $G_{\theta}(\cdot)$  is unconditional ViT/GET with parameters  $\theta$ , and  $\mathbf{e}$  is the initial Gaussian noise. To train a class-conditional GET, we also use class labels  $\mathbf{y}$  in addition to noise/image pairs:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {\mathbf {e}, \mathbf {y}, \mathbf {x} \sim \mathcal {D}} \| \mathbf {x} - G _ {\theta} ^ {c} (\mathbf {e}, \mathbf {y}) \| _ {1}
$$
Figure 3: (a) (Left) Sampling speed of GET: GET can sample faster than large ViTs, while achieving better FID scores. The size of each individual circle is proportional to the model size. For GETs, we 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 to Table 2 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 a class in CIFAR-10.
where  $G_{\theta}^{c}(\cdot)$  is class-conditional ViT/GET with parameters  $\theta$ . As is the standard practice, we also maintain an 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 AdamW [63] optimizer with a learning rate of 1e-4, a batch size of 128 (denoted as  $1 \times \mathrm{BS}$ ), and 800k training iterations, which are identical to Progressive Distillation (PD) [88]. For conditional models, we adopt a batch size of 256 ( $2 \times \mathrm{BS}$ ). No warm-up, weight decay, or learning rate decay is applied. We convert input noise to patches of size  $2 \times 2$ . We use 6 steps of fixed point iterations in the forward pass of GET-DEQ and differentiate through it. For the  $\mathcal{O}(1)$  memory mode, we utilize gradient checkpoint [17] for DEQ'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.
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 1M 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 FID score of a  $5 \times$  large ViT, suggesting substantial parameter efficiency.