SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 a 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 as 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 a summary of our findings on 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 FLOPs. Figure 3 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^{31}$  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-optimal models [37] with substantially better parameter efficiency. For example, at  $2^{31}$  training GFLOPs, Figure 3(b) suggests that we 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 a 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.
Comparison 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 1-step distilled student model. Using the hyperparameters reported in Salimans and Ho [88], PD with DDIM model needs 13 passes of
Table 1: Generative performance on unconditional CIFAR-10.
Method	NFE ↓	FID ↓	IS ↑
Diffusion Models
DDPM [35]	1000	3.17	9.46
Score SDE [96]	2000	2.2	9.89
DDIM [94]	10	13.36	-
EDM [48]	35	2.04	9.84
GANs
StyleGAN2 [47]	1	8.32	9.18
StyleGAN-XL [91]	1	1.85	-
Diffusion Distillation
KD [66]	1	9.36	8.36
PD [88]	1	9.12	-
DFNO [111]	1	4.12	-
TRACT-EDM [11]	1	4.17	-
PD-EDM [88, 97]	1	8.34	8.69
CD-EDM (LPIPS) [97]	1	3.55	9.48
Consistency Models
CT [97]	1	8.70	8.49
CT [97]	2	5.83	8.85
Ours
GET-Base	1	6.91	9.16
Table 2: Generative performance of GETs on unconditional CIFAR-10.
Models	Params	NFE ↓	FID ↓	IS ↑
GET-Tiny	8.6M	1	15.19	8.37
GET-Mini	19.2M	1	10.72	8.69
GET-Small	37.2M	1	8.00	9.03
GET-Base	62.2M	1	7.42	9.16
GET-Base+	83.5M	1	7.19	9.09
More Training
GET-Tiny-4×Iters	8.9M	1	11.47	8.64
GET-Base-2×BS	62.2M	1	6.91	9.16
Table 3: Generative performance on class-conditional CIFAR-10.  $w$  indicates the level of classifier guidance.
Method	NFE ↓	FID ↓	IS ↑
GANs
BigGAN [12]	1	14.73	9.22
StyleGAN2-ADA [46]	1	2.42	10.14
Diffusion Distillation
Guided Distillation (w=0) [68]	1	8.34	8.63
Guided Distillation (w=0.3) [68]	1	7.34	8.90
Guided Distillation (w=1) [68]	1	8.62	9.21
Guided Distillation (w=2) [68]	1	13.23	9.23
Ours
GET-Base	1	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 in all metrics. Results are benchmarked on 4 A6000 GPUs using a batch size of 128, 800k iterations, and PyTorch [74] distributed training protocol. Training Mem stands for training memory consumed per GPU.  $\mathcal{O}(1)$  symbolizes the  $\mathcal{O}(1)$  training memory mode, which differs only in training memory and speed.
Model	FID↓	IS↑	Params↓	FLOPs↓	Training Mem↓	Training Speed↑
ViT-Base	11.49	8.61	85.2M	23.0G	10.1GB	4.83 iter/sec
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  $50\mathrm{K}$  iterations, and the last pass uses  $100\mathrm{K}$  iterations. Each step of PD uses 2 teacher model NFEs. Thus, the overall number of teacher model NFEs can be evaluated as  $2 \times 128$  (batch size)  $\times$  (12 passes  $\times 50\mathrm{K} + 100\mathrm{K}) = 179\mathrm{M}$  samples. The number of NFEs of the teacher model increases to 1.433B if we assume that each of 8 TPUs use a batch size of 128. Consistency distillation [97] needs 409.6M teacher model NFEs (512 batch size  $\times 800\mathrm{K}$  iterations  $= 409.6\mathrm{M}$ ). 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 in Table 1 and Table 3, respectively. GET outperforms a much more complex distillation procedure—PD with classifier-free guidance—in class-conditional image generation. GET also outperforms PD and KD in 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-art distillation techniques. We have outlined key differences in the experimental setup between our 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 in 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 ↑	BS	Training Phases	#Models	Trajectory	Teacher
KD [66]†	9.36	-	4×	1	1	×	DDIM
PD [88]	9.12	-	1×	log2(T)	2	✓	DDIM
DFNO [111]†	4.12	-	2×	1	1	✓	DDIM
TRACT [11]	14.40	-	2×	1	1	✓	DDIM
TRACT [11]	4.17	-	2×	2	1	✓	EDM
PD-EDM [88, 97]	8.34	8.69	4×	log2(T)	2	✓	EDM
CD▲ [97]	3.55	9.48	4×	1	3	✓	EDM
Ours†	7.42	9.16	1×	1	1	×	EDM
Ours†	6.91	9.16	2×	1	1	×	EDM
Guided Distillation [68]	7.34	8.90	4×	log2(T) + 1	3	✓	DDIM
Ours†	6.25	9.40	2×	1	1	×	EDM
Distillation techniques for diffusion models. Knowledge distillation (KD) [66] proposed to distill a multi-step DDIM [94] sampler into the pretrained UNet by training the student model on 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-conditional image generation.
Progressive distillation (PD) [88] proposes a strategy for online distillation to distill a  $T$ -step teacher DDIM [94] diffusion model into a 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 of training phases. Consistency models [97] achieve online distillation in a 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-free guided diffusion models into few-step generative models by first distilling a combined conditional and unconditional model, and then progressively distilling the resulting model for faster generation.
Fast sampler for diffusion models. While distillation is a 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-free 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 Euler 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.
Our method for offline distillation relies on deterministic samplers to ensure a unique mapping between initial noise  $\mathbf{e}$  and image  $\mathbf{x}$ . As a result, it cannot be directly applied to stochastic samplers which do 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].
We propose a simple yet effective approach to distill diffusion models into generative models capable 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 pretrained 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, a small GET can generate higher quality images than a  $5 \times$  larger ViT, sampling faster while using less training memory and fewer compute FLOPs, demonstrating its effectiveness.
Zhengyang Geng and Ashwini Pokle are supported by grants from the Bosch Center for Artificial Intelligence.
[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 B Grosse. Path independent equilibrium models can better exploit test-time computation. Advances in Neural Information Processing Systems, 35:7796-7809, 2022. 18
[3] Uri M Ascher and Linda R Petzold. Computer methods for ordinary differential equations and differential-algebraic equations, volume 61. Siam, 1998. 3, 5
[4] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.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 on Machine Learning (ICML), 2021. 18
[8] Shaojie Bai, Zhengyang Geng, Yash Savani, and J Zico Kolter. Deep equilibrium optical flow estimation. In 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 covariance with imperfect mean in diffusion probabilistic models. arXiv preprint arXiv:2206.07309, 2022. 9
[10] Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint arXiv: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-distillation, 2023. 1, 8, 9
[12] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809.11096, 2018. 8, 18
[13] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, 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
[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 on 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 to 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 arXiv:2109.01652, 2021. 18
[105] Ezra Winston and J Zico Kolter. Monotone operator equilibrium networks. Advances in neural information processing systems, 33:10718-10728, 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. In International Conference on Machine Learning (ICML), 2020. 4
[107] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022. 1
[108] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2023. 9
[109] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 586-595, 2018. 9
[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 Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. arXiv preprint arXiv:2211.13449, 2022. 1, 8, 9
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 2) injection with adaptive layer normalization (AdaLN-Zero) as used in DiT [75]. Despite using almost the same parameters as unconditional ViT-B, the class-conditional 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 an FID of 17.19 at 200k iterations. Another surprising observation is that ViT-B w/ AdaLN-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.
Table 6: Ablation on class conditioning.
Model	FID↓	IS↑	Params↓
ViT-Uncond	15.20	8.27	85.2M
ViT-AdaLN-Zero	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_{i}$  and  $L_{e}$  denote the number of transformer blocks in 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 \times D$ . For the injection transformer, we always adopt an expanding factor of 4.
Table 7: Details of configuration for GET architectures.
Model	Params	Li	Le	D	E
GET-Tiny	8.9M	6	3	256	6
GET-Mini	19.2M	6	3	384	6
GET-Small	37.2M	6	3	512	6
GET-Base	62.2M	1	3	768	12
GET-Base+	83.5M	6	3	768	8
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,  $L$  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	L	D
ViT-B	85.2M	12	768
ViT-L	302.6M	24	1024
Transformers. Transformers were first proposed by Vaswani et al. [99] for machine translation and since then have been widely applied in many domains like natural language processing [21, 55, 79, 84], reinforcement learning [16, 73], self-supervised learning [15], vision [25, 62], and generative
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 (ViT) 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 positional information. This sequence of embedding vectors is fed into the standard transformer architecture. Diffusion Transformers (DiT) are based on 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 use as a backbone for diffusion models and help them scale better with increasing model size, including adaptive Layer Normalization (AdaLN-Zero) [12, 22, 45, 76] for time and class embedding, and zero-initialization for the final convolution layer [32].
Deep equilibrium models. Deep Equilibrium models (DEQs) [5] solve for a fixed point in the forward pass. Specifically, given the input  $\mathbf{x}$  and the equilibrium function  $f_{\theta}$ , DEQ models approach the infinite-depth representation of  $f_{\theta}$  by solving for its fixed point  $\mathbf{z}^{\star}$ :  $\mathbf{z}^{\star} = f_{\theta}(\mathbf{z}^{\star}, \mathbf{x})$ . For the backward pass, one can differentiate analytically through  $z^{\star}$  by the implicit function theorem. The training dynamics of DEQ models can be unstable for certain model designs [7]. As a result, recent efforts focus on addressing these issues by designing variants of DEQs with provable guarantees [83, 105], or through optimization techniques such as Jacobian regularization [7], and fixed-point correction [8]. 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], out-of-distribution generalization [2], language modelling [5], input optimization [33], unsupervised learning [98], and generative models [64, 77].