SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weijian Luo $^{1*}$ , Tianyang Hu $^{2\dagger}$ , Shifeng Zhang $^{2}$ , Jiacheng Sun $^{2}$ , Zhenguo Li $^{2}$ , Zhihua Zhang $^{1}$
$^{1}$ Peking University,  $^{2}$ Huawei Noah's Ark Lab
Due to the ease of training, ability to scale, and high sample quality, diffusion models (DMs) have become the preferred option for generative modeling, with numerous pre-trained models available for a wide variety of datasets. Containing intricate information about data distributions, pre-trained DMs are valuable assets for downstream applications. In this work, we consider learning from pre-trained DMs and transferring their knowledge to other generative models in a data-free fashion. Specifically, we propose a general framework called Diff-Instruct to instruct the training of arbitrary generative models as long as the generated samples are differentiable with respect to the model parameters. Our proposed Diff-Instruct is built on a rigorous mathematical foundation where the instruction process directly corresponds to minimizing a novel divergence we call Integral Kullback-Leibler (IKL) divergence. IKL is tailored for DMs by calculating the integral of the KL divergence along a diffusion process, which we show to be more robust in comparing distributions with misaligned supports. We also reveal non-trivial connections of our method to existing works such as DreamFusion [54], and generative adversarial training. To demonstrate the effectiveness and universality of Diff-Instruct, we consider two scenarios: distilling pre-trained diffusion models and refining existing GAN models. The experiments on distilling pre-trained diffusion models show that Diff-Instruct results in state-of-the-art single-step diffusion-based models. The experiments on refining GAN models show that the Diff-Instruct can consistently improve the pre-trained generators of GAN models across various settings. Our official code is released through https://github.com/pkulwj1994/diff_instruct.
Over the last decade, the field of deep generative models has made significant strides across various domains such as data generation [30, 32, 51, 53, 24, 54, 25, 34], density estimation [35, 7], image-editing [46, 8] and others. Notably, recent advancements in text-driven high-resolution image generation [60, 59, 58] have pushed the limits of using generative models for Artificial Intelligence Generated Content (AIGC). Behind the empirical success are fruitful developments of a wide variety of deep generative models, among which, diffusion models (DMs) are the most prominent. DMs leverage the diffusion processes and model the data across a wide spectrum of noise levels. Their ease of training, ability to scale, and high sample quality have made DMs the preferred option for generative modeling, with numerous pre-trained models available for a wide variety of datasets and applications. The trained DMs contain intricate information about the data distribution, making them valuable assets for downstream applications.
*Email: luoweijian@stu.pku.edu.cn.
† Corresponding to: Tianyang Hu (hutianyang1@huawei.com)
Compared with training from scratch, extracting knowledge from a zoo of pre-trained models enables us to learn more efficiently. For instance, [15] employed a variety of pre-trained feature extractors and significantly boosted the state-of-the-art performance on domain generalization benchmarks. [9, 52, 16] exploited the rich multi-modal information stored in off-the-shelf CLIP models [55] for efficient text-guided image generation. Currently, we are witnessing a rising trend of learning from models, especially when accessing large amounts of high-quality data is difficult. Such a model-driven learning scheme can be particularly appealing for handling new tasks by providing a solid base model, which can be further improved by additional training data. While this research direction has been extensively investigated for discriminative models and supervised learning tasks [71, 38, 3, 72, 57, 56], its application to generative models remains largely unexplored. To this end, we are motivated to study the following question.
(Q1): Can we transfer knowledge from pre-trained DMs to other generative models instead of learning from original training data?
The seminal work DreamFusion [54] demonstrated the feasibility of such a quest for text-to-3D generation. Without a large text-labeled 3D dataset, [54] took advantage of the rich text-to-image knowledge stored in a large-scale diffusion model, the Imagen [23], to learn the 3D Neural Radiance Fields (NeRF) and achieved surprisingly good performance without using any 3D data. DreamFusion is a rare success and for general scenarios, the task can be difficult due to the vast difference among generative models. DMs represent a class of explicit generative models wherein the data's score function is modeled. Conversely, in various downstream applications, implicit generative models are favored due to their inherent flexibility and efficiency. An implicit model typically learns a neural transformation (i.e., a generator) that maps from a latent space to the data space, such as in generative adversarial networks (GANs), thereby enabling expeditious generation.
By exploring diverse architectural designs and latent space configurations, implicit models can readily adapt to structural constraints (e.g., molecules must be chemically valid, etc. [14, 10, 63]), assimilate prior knowledge [6, 13, 76, 45], and exhibit other advantageous properties. For implicit models that lack explicit score information, how to receive supervision from DM's multi-level score network is technically challenging, which greatly limits the potential use cases of pre-trained DMs. Therefore, we would like to further address the following question:
(Q2): Can we tackle this challenge so that knowledge from DMs can be more broadly transferred?
In this work, we give affirmative answers to (Q2) and propose a universal framework, Diff-Instruct (DI), to leverage pre-trained DMs to instruct the training of arbitrary implicit generative models as long as the generated samples from the implicit model are differentiable with respect to model parameters. When applied to single-step generation models such as GANs, Diff-Instruct provides an alternative non-adversarial training scheme. When the student model is a U-Net (with a fixed time), our method enters as a strong contender in the diffusion distillation literature [43, 61, 66], providing extreme acceleration for sampling from DMs with even one single step.
Our proposed Diff-Instruct is built on a rigorous mathematical foundation where the instruction process directly corresponds to minimizing a novel divergence we call Integral Kullback-Leibler (IKL) divergence. IKL is tailored for DMs by calculating the integral of the KL divergence along a diffusion process, which we show to be more robust in comparing distributions with misaligned supports (Section 3.2). We also reveal non-trivial connections of our method to existing works such as DreamFusion (Section 3.3.1) and generative adversarial training (Section 3.3.2). Interestingly, we show that the SDS objective can be seen as a special case of our Diff-Instruct on the scenario that the generator outputs a Dirac's Delta distribution.
To demonstrate the effectiveness and universality of Diff-Instruct, we consider two scenarios mentioned earlier: distilling pre-trained diffusion models to single step (Section 4.1) and improving pre-trained GAN generators (Section 4.2). The experiments on distilling pre-trained diffusion models on the ImageNet dataset of a resolution of  $64 \times 64$  show that Diff-Instruct results in state-of-the-art single-step diffusion-based models over both diffusion distillation, such as the consistency distillation [66] and direct training methods [66, 83, 73]. The experiments on improving GAN generators show that the Diff-Instruct can consistently improve the pre-trained generators across various settings.
Assume we observe data from the underlying distribution  $p_d(\pmb{x})$ . In generative modeling, we want to generate new samples  $\pmb{x} \sim p_d(\pmb{x})$ , where there are mainly two approaches, explicit and implicit. Currently, DMs are the most powerful explicit models while GANs are the most powerful implicit models.
Diffusion models. The forward diffusion process of DM transforms any initial distribution  $p^{(0)}$  towards some simple noise distribution,
$$
d \boldsymbol {x} _ {t} = \boldsymbol {F} (\boldsymbol {x} _ {t}, t) \mathrm {d} t + G (t) d \boldsymbol {w} _ {t}, \tag {2.1}
$$
where  $\pmb{F}$  is a pre-defined drift function,  $G(t)$  is a pre-defined scalar-value diffusion coefficient, and  $\pmb{w}_t$  denotes an independent Wiener process. A multiple-level or continuous-indexed score network  $s_{\phi}(\pmb{x}, t)$  is usually employed in order to approximate marginal score functions of the forward diffusion process (2.1). The learning of marginal score functions is achieved by minimizing a weighted denoising score matching objective [69, 65],
$$
\mathcal {L} _ {D S M} (\phi) = \int_ {t = 0} ^ {T} w (t) \mathbb {E} _ {\boldsymbol {x} _ {0} \sim p ^ {(0)}, \boldsymbol {x} _ {t} | \boldsymbol {x} _ {0} \sim p _ {t} (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0})} \| \boldsymbol {s} _ {\phi} (\boldsymbol {x} _ {t}, t) - \nabla_ {\boldsymbol {x} _ {t}} \log p _ {t} (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0}) \| _ {2} ^ {2} d t. \tag {2.2}
$$
Here the weighting function  $w(t)$  controls the importance of the learning at different time levels and  $p_t(\pmb{x}_t|\pmb{x}_0)$  denotes the conditional transition of the forward diffusion (2.1). High-quality samples from a DM can be drawn by simulating SDE which is implemented by learned score network [65]. However, the simulation of an SDE is significantly slower than that of other models such as implicit models.
Generative adversarial networks. GANs are representative implicit generative models [30, 75, 67, 77]. They leverage neural networks (generators) to map an easy-to-sample latent vector to generate a sample. Therefore they are efficient. However, the training of GANs is challenging, particularly because of the reliance on adversarial training. To train a GAN model, a neural discriminator  $h$  is optimized to distinguish the data and generated samples. This leads to the creation of a surrogate probability metric  $D_{h}(\cdot ,\cdot)$  between  $p_{g}(\pmb {x})$  and  $p_d(\pmb {x})$ . The generator is updated based on this metric, with the aim of improving the quality of the generated samples [19, 2, 44]. The objective of a most commonly used GAN [19] can be written as
$$
\mathcal {L} _ {h} = - \mathbb {E} _ {x \sim p _ {d}} [ \log h (\boldsymbol {x}) ] - \mathbb {E} _ {z \sim p _ {z}} [ \log (1 - h (g (\boldsymbol {z}))) ], \quad \mathcal {L} _ {g} = - \mathbb {E} _ {z \sim p _ {z}} [ \log h (g (\boldsymbol {z})) ],
$$
where the training alternates between minimizing  $\mathcal{L}_h$  and  $\mathcal{L}_g$  with the other part fixed. For a fixed  $g$ , the optimal  $h$  should recover the density ratio  $p_d(\boldsymbol{x}) / (p_d(\boldsymbol{x}) + p_g(\boldsymbol{x}))$ , and in turn,  $D_h$  is the Jensen-Shannon divergence. There are variants of objectives that minimize other divergences. For instance, if the  $\mathcal{L}_g^{(KL)} = \mathbb{E}_{z \sim p_z}[\log \frac{1 - h(g(z))}{h(g(z))}]$ , the objective aims to minimize the KL divergence between generator and data distribution. In Section 3.3.2, we establish the equivalence of our Diff-Instruct with the adversarial training that aims to minimize the KL divergence.
Neural radiance fields. The neural radius field (NeRF) [48] is a kind of 3D object model that uses a multi-layer-perceptron (MLP) to map coordinates of a mesh grid to volume properties such as color and density. Given the camera parameters, a rendering algorithm can output a 2D image that is a view projection of the 3D NeRF. The rendering algorithm is usually differentiable to learnable parameters of NeRF's MLP, this makes the NeRF can be updated through proper instructions on the rendered 2D image.
3 Diff-Instruct
The main goal of Diff-Instruct is to transfer the knowledge of a pre-trained DM to other generative models. To demonstrate the universality of our approach, we consider the more general case where the student model is an implicit model, i.e., a generator.
Problem setup. Recall our setting that we have a pre-trained diffusion model with the multi-level score net denoted as  $s_{p^{(t)}}(\boldsymbol{x}_t) \coloneqq \nabla_{\boldsymbol{x}_t} \log p^{(t)}(\boldsymbol{x}_t)$  where  $p^{(t)}(\boldsymbol{x}_t)$ 's are the underlying distributions diffused at time  $t$  according to (2.1). Assume the pre-trained diffusion model provides a sufficiently good approximation of data distribution, i.e.,  $p^{(0)} \approx p_d$ . For ease of mathematical treatment, we use
Figure 1: Illustration of our Diff-Instruct pipeline. The generator accepts instructions from all diffusion time levels to calculate the gradient of the Integral KL divergence. The gradient is used to update its parameters. This data-free learning scheme enables us to employ pre-trained DMs as teachers to instruct a wide variety of generative models.
$p^{(0)}, p_d$  interchangeably. The goal of our Diff-Instruct is to train an implicit model  $g_{\theta}$  without any training data, such that the distribution of the generated samples, denoted as  $p_g$ , matches that of the pre-trained DM. The instruction process involves minimizing certain probability divergences between the implicit distribution and the data distribution.
Instruction criterion. In order to receive supervision from the multi-level score functions  $s_{p^{(t)}}(\boldsymbol{x}_t)$ , introducing the same diffusion process to the generated samples seems inevitable. Consider diffusing  $p_g$  along the same forward process as the instructor DM and let  $q^{(t)}$  be the corresponding densities at time  $t$ . Let  $s_{q^{(t)}}(\boldsymbol{x}_t) \coloneqq \nabla_{\boldsymbol{x}_t} \log q^{(t)}(\boldsymbol{x}_t)$  be the marginal score functions. At each time level, how to design the instruction criterion and how to combine different time levels are of critical importance. To this end, we consider integrating the Kullback-Leibler divergence along the forward diffusion process with a proper weighting function, such as in (2.2). The resulting Integral Kullback-Leibler (IKL) divergence is a valid probability divergence with two important properties: 1) IKL is more robust than KL in comparing distributions with misaligned supports; 2) The gradient of IKL with respect to the generator's parameters only requires the marginal score functions of the diffusion process, making it a suitable divergence for incorporating the scoring network of pre-trained diffusion models.
In the following sections, we first formally define the IKL and then go into detail about the mathematical ground of our Diff-Instruct algorithms. We then establish the connections of Diff-Instruct to existing methods and discuss in detail a novel application of Diff-Instruct on data-free diffusion distillation, together with a comparison to existing diffusion distillation approaches.
3.1 Integral Kullback-Leibler divergence
The IKL is tailored to incorporate knowledge of pre-trained DMs in multiple time levels. It generalizes the concept of KL divergence to involve all time levels of the diffusion process.
Definition 3.1 (Integral KL divergence). Given a diffusion process (2.1) and a proper weighting function  $w(t) > 0, t \in [0,T]$ , the IKL divergence between two distributions  $p,q$  is defined as
$$
\mathcal {D} _ {I K L} ^ {[ 0, T ]} (q, p) := \int_ {t = 0} ^ {T} w (t) \mathcal {D} _ {K L} \left(q ^ {(t)}, p ^ {(t)}\right) \mathrm {d} t := \int_ {t = 0} ^ {T} w (t) \mathbb {E} _ {\boldsymbol {x} _ {t} \sim q ^ {(t)}} \left[ \log \frac {q ^ {(t)} \left(\boldsymbol {x} _ {t}\right)}{p ^ {(t)} \left(\boldsymbol {x} _ {t}\right)} \right] \mathrm {d} t, \tag {3.1}
$$
where  $q^{(t)}$  and  $p^{(t)}$  denote the marginal densities of the diffusion process (2.1) at time  $t$  initialized with  $q^{(0)} = q$  and  $p^{(0)} = p$  respectively.
The IKL divergence integrates the KL divergence along a diffusion process, which enables us to amalgamate knowledge from pre-trained DM at multiple diffusion times. For simplicity, we use the notation  $\mathcal{D}_{IKL}(p,q)$  to represent  $\mathcal{D}_{IKL}^{[0,\infty)}(p,q)$  if the integral exists. Since the KL divergence is well-defined, the proposed IKL divergence, as the integral of KL divergence, is also well-defined.
Proposition 3.2. The  $\mathcal{D}_{IKL}(q,p)$  satisfies that  $\mathcal{D}_{IKL}(q,p) \geq 0, \forall q,p$ . Furthermore, the equality holds if and only if  $q = p$ , almost everywhere under measure  $p$ .
Algorithm 1: Diff-Instruct Algorithm
Input: pre-trained DM  $s_{p^{(t)}}$ , generator  $g_{\theta}$ , prior distribution  $p_z$ , DM  $s_{\phi}$ ; forward diffusion (2.1). while not converge do
update  $\phi$  using SGD with gradient
$$
\operatorname {Grad}(\phi) = \frac{\partial}{\partial\phi}\int_{t = 0}^{T}w(t)\mathbb{E}_{\substack{\boldsymbol {z}\sim p_{z},\boldsymbol{x}_{0} = g_{\theta}(\boldsymbol {z}),\\ \boldsymbol{x}_{t}|\boldsymbol{x}_{0}\sim p_{t}(\boldsymbol{x}_{t}|\boldsymbol{x}_{0})}}\| \boldsymbol{s}_{\phi}(\boldsymbol{x}_{t},t) - \nabla_{\boldsymbol{x}_{t}}\log p_{t}(\boldsymbol{x}_{t}|\boldsymbol{x}_{0})\|_{2}^{2}\mathrm{d}t.
$$
update  $\theta$  using SGD with the gradient
$$
\operatorname {Grad}(\theta) = \int_{t = 0}^{T}w(t)\mathbb{E}_{\substack{\boldsymbol {z}\sim p_{z},\boldsymbol{x}_{0} = g_{\theta}(\boldsymbol {z}),\\ \boldsymbol{x}_{t}|\boldsymbol{x}_{0}\sim p_{t}(\boldsymbol{x}_{t}|\boldsymbol{x}_{0})}}\left[ \boldsymbol{s}_{\phi}(\boldsymbol{x}_{t},t) - \boldsymbol{s}_{p^{(t)}}(\boldsymbol{x}_{t})\right]\frac{\partial\boldsymbol{x}_{t}}{\partial\theta}\mathrm{d}t.
$$
return  $\theta ,\phi$
One of the advantages of using IKL instead of KL is its robustness. For instance, with proper weighting function, the IKL is well-defined even when the vanilla KL divergence degenerates to infinity. This demonstrates that the IKL divergence is more robust than the KL divergence for two distributions with misaligned supports. We consider a famous example in [2] where the generator distribution  $p_{g}$  and target distribution  $p_{d}$  have disjoint support. In this case, the KL divergence between  $p_{g}$  and  $p_{d}$  degenerates to positive infinity, while the IKL divergence has a finite value for all generator parameters and unique minima that match the generator and data distribution. Check Appendix A.1 for details.
3.2 Instruct algorithm
Let  $g_{\theta}$  be the generator of the implicit model. Let  $q^{(0)}$  denote the implicit distribution for samples which are obtained by  $\pmb{x}_0 = g_{\theta}(\pmb{z}), \pmb{z} \sim p_z$  and denote  $q^{(t)}$  as the marginal distribution of the forward SDE ((2.1)) initialized with  $q^{(0)}$ . Let  $\{p^{(t)}\}_{t \in [0, \infty]}$  and  $\{s_{p^{(t)}}(\cdot)\}_{t \in [0, \infty]}$  represent the marginal densities and score functions of the pre-trained diffusion model. Our Diff-Instruct aims to minimize the IKL between  $q^{(0)}$  and  $p^{(0)}$  so as to update the generator's parameters. Following the notations of definition (3.1), we give a non-trivial gradient formula for minimizing the IKL in Theorem 3.3 that includes only the score functions.
Theorem 3.3. The gradient of the IKL in (3.1) between  $q^{(0)}$  and  $p^{(0)}$  is
$$
\operatorname {G r a d} (\theta) = \int_ {t = 0} ^ {T} w (t) \mathbb {E} _ {\substack {\boldsymbol {z} \sim p _ {z}, \boldsymbol {x} _ {0} = g _ {\theta} (\boldsymbol {z}), \\ \boldsymbol {x} _ {t} | \boldsymbol {x} _ {0} \sim p _ {t} (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0})}} \left[ \boldsymbol {s} _ {\phi} (\boldsymbol {x} _ {t}, t) - \boldsymbol {s} _ {p ^ {(t)}} (\boldsymbol {x} _ {t}) \right] \frac {\partial \boldsymbol {x} _ {t}}{\partial \theta} \mathrm {d} t. \tag{3.2}
$$
Theorem 3.3 gives an explicit gradient to minimize the IKL divergence w.r.t. the parameter of the generator. Note that the gradient estimation only requires the marginal score functions  $s_{p^{(t)}}$  and  $s_{q^{(t)}}$ . If the marginal score functions of the implicit distribution can be approximated by another diffusion model  $s_{\phi}(\boldsymbol{x}_t,t)$ , we can utilize the gradient formula (3.2) to update the generator's parameter  $\theta$ .
Now we formally propose Diff-Instruct as in Algorithm 1, which trains the implicit model through two alternative phases between learning the marginal score functions  $s_{\phi}$ , and updating the implicit model with gradient (3.2). The former phase follows the standard DM learning procedure, i.e., minimizing loss function (2.2), with a slight change that the data is generated from the generator. The resulting  $s_{\phi}(\boldsymbol{x}_t,t)$  provides an estimation of  $s_{q^{(t)}}(\boldsymbol{x}_t)$ . The latter phase updates the generator's parameter  $\theta$  using gradient from (3.2), where two needed functions are provided by pre-trained DM  $s_{p^{(t)}}(\boldsymbol{x}_t)$  and learned DM  $s_{\phi}(\boldsymbol{x}_t,t)$ . When the algorithm converges,  $s_{\phi}(\boldsymbol{x}_t,t) \approx s_{p^{(t)}}(\boldsymbol{x}_t)$ , and the gradient (3.2) is approximately zero.
3.3 Connections to existing methods
In this section, we establish the connections of Diff-Instruct to two typical methods, the score distillation sampling proposed in DreamFusion [54], and the generative adversarial training in Goodfellow et al. [19].
3.3.1 Connection to score distillation sampling
The score distillation sampling (SDS) algorithm was proposed by Poole et al. [54] to distill the knowledge of a large-scale text-to-image diffusion model into a 3D NeRF model. The idea of SDS has been applied in various contexts, including the text-to-3D NeRF generation based on text-to-2D diffusion models [54, 39, 47], and the text-guided image editing [20].
It turns out that the SDS algorithm is a special case of our Diff-Instruct when the generator's output is a Dirac's Delta distribution with learnable parameters. More precisely, we find that Diff-Instruct's gradient formula (3.2) will degenerate to the gradient formula of SDS under the assumption that the generator outputs a Delta distribution.
Corollary 3.4. If the generator's output is a Dirac's Delta distribution with learnable parameters, i.e.  $q(\pmb{x}_0) = \delta_{g(\theta)}(\pmb{x}_0)^3$ . Then the gradient formula (3.2) becomes
$$
\operatorname {G r a d} (\theta) = \int_ {t = 0} ^ {T} w (t) \mathbb {E} _ {\substack {\boldsymbol {x} _ {0} = g (\theta), \\ \boldsymbol {x} _ {t} | \boldsymbol {x} _ {0} \sim p _ {t} (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0})}} \left[ \nabla_ {\boldsymbol {x} _ {t}} \log p _ {t} (\boldsymbol {x} _ {t} | \boldsymbol {x} _ {0}) - \boldsymbol {s} _ {p ^ {(t)}} (\boldsymbol {x} _ {t}) \right] \frac {\partial \boldsymbol {x} _ {t}}{\partial \theta} \mathrm {d} t. \tag{3.3}
$$
(3.3) does not depend on another diffusion model  $s_{q^{(t)}}$  as in (3.2). So under the assumption of Corollary 3.4, there is no need for using another DM to estimate the generator's marginal score functions. This is because when the generator outputs a Delta distribution, there is no randomness in  $x_0$ . So the marginal score functions are only determined by  $p_t(x_t|x_0)$ . The gradient (3.3) is equivalent to the score distillation sampling proposed in DreamFusion [54].