SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ying Fan\*,1,2, Olivia Watkins3, Yuqing Du3, Hao Liu3, Moonkyung Ryu1, Craig Boutilier1, Pieter Abbeel3, Mohammad Ghavamzadeh†,4, Kangwook Lee2, Kimin Lee\*,†,5
*Equal technical contribution †Work was done at Google Research
<sup>1</sup>Google Research <sup>2</sup>University of Wisconsin-Madison <sup>3</sup>UC Berkeley <sup>4</sup>Amazon <sup>5</sup>KAIST
Learning from human feedback has been shown to improve text-to-image models. These techniques first learn a reward function that captures what humans care about in the task and then improve the models based on the learned reward function. Even though relatively simple approaches (e.g., rejection sampling based on reward scores) have been investigated, fine-tuning text-to-image models with the reward function remains challenging. In this work, we propose using online reinforcement learning (RL) to fine-tune text-to-image models. We focus on diffusion models, defining the fine-tuning task as an RL problem, and updating the pre-trained text-to-image diffusion models using policy gradient to maximize the feedback-trained reward. Our approach, coined DPOK, integrates policy optimization with KL regularization. We conduct an analysis of KL regularization for both RL fine-tuning and supervised fine-tuning. In our experiments, we show that DPOK is generally superior to supervised fine-tuning with respect to both image-text alignment and image quality. Our code is available at https://github.com/google-research/google-research/tree/master/dpok.
Recent advances in diffusion models [10, 37, 38], together with pre-trained text encoders (e.g., CLIP [27], T5 [28]) have led to impressive results in text-to-image generation. Large-scale text-to-image models, such asImagen [32], Dalle-2 [29], and Stable Diffusion [30], generate high-quality, creative images given novel text prompts. However, despite these advances, current models have systematic weaknesses. For example, current models have a limited ability to compose multiple objects [6, 7, 25]. They also frequently encounter difficulties when generating objects with specified colors and counts [12, 17].
Learning from human feedback (LHF) has proven to be an effective means to overcome these limitations [13, 17, 42, 43]. Lee et al. [17] demonstrate that certain properties, such as generating objects with specific colors, counts, and backgrounds, can be improved by learning a reward function from human feedback, followed by fine-tuning the text-to-image model using supervised learning. They show that simple supervised fine-tuning based on reward-weighted loss can improve the reward scores, leading to better image-text alignment. However, supervised fine-tuning often induces a deterioration in image quality (e.g., over-saturated or non-photorealistic images). This is likely due to the model being fine-tuned on a fixed dataset that is generated by a pre-trained model (Figure 1(a)).
In this work, we explore using online reinforcement learning (RL) for fine-tuning text-to-image diffusion models (Figure 1(b)). We show that optimizing the expected reward of a diffusion model's image output is equivalent to performing policy gradient on a multi-step diffusion model under certain
(a) Supervised fine-tuning
Figure 1: Illustration of (a) reward-weighted supervised fine-tuning and (b) RL fine-tuning. Both start with the same pre-trained model (the blue rectangle). In supervised fine-tuning, the model is updated on a fixed dataset generated by the pre-trained model. In contrast, the model is updated using new samples from the previously trained model during online RL fine-tuning.
(b) RL fine-tuning
regularity assumptions. We also incorporate Kullback-Leibler (KL) divergence with respect to the pre-trained model as regularization in an online manner, treating this as an implicit reward.
In our experiments, we fine-tune the Stable Diffusion model [30] using ImageReward [43], an open-source reward model trained on a large dataset comprised of human assessments of (text, image) pairs. We show that online RL fine-tuning achieves strong text-image alignment while maintaining high image fidelity by optimizing its objective in an online manner. Crucially, online training allows evaluation of the reward model and conditional KL divergence beyond the (supervised) training dataset. This offers distinct advantages over supervised fine-tuning, a point we demonstrate empirically. In our empirical comparisons, we also incorporate the KL regularizer in a supervised fine-tuning method for a fair comparison.
- We frame the optimization of the expected reward (w.r.t. an LHF-reward) of the images generated by a diffusion model given text prompts as an online RL problem. Moreover, we present DPOK: Diffusion Policy Optimization with KL regularization, which utilizes KL regularization w.r.t. the pre-trained text-to-image model as an implicit reward to stabilize RL fine-tuning.
- We study incorporating KL regularization into supervised fine-tuning of diffusion models, which can mitigate some failure modes (e.g., generating over-saturated images) in [17]. This also allows a fairer comparison with our RL technique.
- We discuss the key differences between supervised fine-tuning and online fine-tuning of text-to-image models (Section 4.3).
- Empirically, we show that online fine-tuning is effective in optimizing rewards, which improves text-to-image alignment while maintaining high image fidelity.
Text-to-image diffusion models. Diffusion models [10, 37, 39] are a class of generative models that use an iterative denoising process to transform Gaussian noise into samples that follow a learned data distribution. These models have proven to be highly effective in a range of domains, including image generation [4], audio generation [14], 3D synthesis [26], and robotics [3]. When combined with large-scale language encoders [27, 28], diffusion models have demonstrated impressive performance in text-to-image generation [29, 30, 32]. However, there are still many known weaknesses of existing text-to-image models, such as compositionality and attribute binding [6, 20] or text rendering [21].
Learning from human feedback. Human assessments of (or preferences over) learned model outcomes have been used to guide learning on a variety of tasks, ranging from learning behaviors [16] to language modeling [1, 23, 19, 41]. Recent work has also applied such methods to improve the alignment of text-to-image models. Human preferences are typically gathered at scale by asking annotators to compare generations, and a reward model is trained (e.g., by fine-tuning a vision-language model such as CLIP [27] or BLIP [18]) to produce scalar rewards well-aligned with the human feedback [43, 13]. The reward model is used to improve text-to-image model quality by fine-tuning a pre-trained generative model [17, 42]. Unlike prior approaches, which typically focus on reward-filtered or reward-weighted supervised learning, we develop an online fine-tuning framework with an RL-based objective.
RL fine-tuning of diffusion models. Fan & Lee [5] first introduced a method to improve pre-trained diffusion models by integrating policy gradient and GAN training [8]. They used policy gradient with reward signals from the discriminator to update the diffusion model and demonstrated that the fine-tuned model can generate realistic samples with few diffusion steps with DDPM sampling [10] on relatively simple domains (e.g., CIFAR [15] and CelebA [22]). In this work, we explore RL fine-tuning especially for large-scale text-to-image models using human rewards. We also consider several design choices like adding KL regularization as an implicit reward, and compare RL fine-tuning to supervised fine-tuning.
Concurrent and independent from our work, Black et al. [2] have also investigated RL fine-tuning to fine-tune text-to-image diffusion models. They similarly frame the fine-tuning problem as a multi-step decision-making problem, and demonstrate that RL fine-tuning can outperform supervised fine-tuning with reward-weighted loss [17] in optimizing the reward, which aligns with our own observations. Furthermore, our work analyzes KL regularization for both supervised fine-tuning and RL fine-tuning with theoretical justifications, and shows that adopting KL regularization is useful in addressing some failure modes (e.g., deterioration in image quality) of fine-tuned models. For a comprehensive discussion of prior work, we refer readers to Appendix D.
In this section, we describe our basic problem setting for text-to-image generation of diffusion models.
Diffusion models. We consider the use of denoising diffusion probabilistic models (DDPMs) [10] for image generation and draw our notation and problem formulation from [10]. Let  $q_{0}$  be the data distribution, i.e.,  $x_{0} \sim q_{0}(x_{0})$ ,  $x_{0} \in \mathbb{R}^{n}$ . A DDPM approximates  $q_{0}$  with a parameterized model of the form  $p_{\theta}(x_0) = \int p_{\theta}(x_{0:T})dx_{1:T}$ , where  $p_{\theta}(x_{0:T}) = p_T(x_T)\prod_{t = 1}^T p_{\theta}(x_{t - 1}|x_t)$  and the reverse process is a Markov chain with the following dynamics:
$$
p \left(x _ {T}\right) = \mathcal {N} (0, I), \quad p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right) = \mathcal {N} \left(\mu_ {\theta} \left(x _ {t}, t\right), \Sigma_ {t}\right). \tag {1}
$$
A unique characteristic of DDPMs is the exploitation of an approximate posterior  $q(x_{1:T}|x_0)$ , known as the forward or diffusion process, which itself is a Markov chain that adds Gaussian noise to the data according to a variance schedule  $\beta_1,\ldots ,\beta_T$ :
$$
q \left(x _ {1: T} \mid x _ {0}\right) = \prod_ {t = 1} ^ {T} q \left(x _ {t} \mid x _ {t - 1}\right), \quad q \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {N} \left(\sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} I\right). \tag {2}
$$
Let  $\alpha_{t} = 1 - \beta_{t}$ ,  $\bar{\alpha}_{t} = \prod_{s=1}^{t} \alpha_{s}$ , and  $\tilde{\beta}_{t} = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_{t}} \beta_{t}$ . Ho et al. [10] adopt the parameterization  $\mu_{\theta}(x_{t}, t) = \frac{1}{\sqrt{\alpha_{t}}} \left(x_{t} - \frac{\beta_{t}}{\sqrt{1 - \bar{\alpha}_{t}}} \epsilon_{\theta}(x_{t}, t)\right)$ .
Training a DDPM is performed by optimizing a variational bound on the negative log-likelihood  $\mathbb{E}_q\bigl [-\log p_\theta (x_0)\bigr ]$ , which is equivalent to optimizing:
$$
\mathbb {E} _ {q} \left[ \sum_ {t = 1} ^ {T} \operatorname {K L} \left(q \left(x _ {t - 1} \mid x _ {t}, x _ {0}\right) \| p _ {\theta} \left(x _ {t - 1} \mid x _ {t}\right)\right) \right]. \tag {3}
$$
Note that the variance sequence  $(\beta_{t})_{t = 1}^{T}\in (0,1)^{T}$  is chosen such that  $\bar{\alpha}_{T}\approx 0$  , and thus,  $q(x_{T}|x_{0})\approx$ $\mathcal{N}(0,I)$  . The covariance matrix  $\Sigma_t$  in (1) is often set to  $\sigma_t^2 I$  , where  $\sigma_t^2$  is either  $\beta_{t}$  or  $\tilde{\beta}_t$  , which is not trainable. Unlike the original DDPM, we use a latent diffusion model [30], so  $x_{t}$  's are latent.
Text-to-image diffusion models. Diffusion models are especially well-suited to conditional data generation, as required by text-to-image models: one can plug in a classifier as guidance function [4], or can directly train the diffusion model's conditional distribution with classifier-free guidance [9].
Given text prompt  $z \sim p(z)$ , let  $q(x_0|z)$  be the data distribution conditioned on  $z$ . This induces a joint distribution  $p(x_0, z)$ . During training, the same noisy process  $q$  is used regardless of input  $z$ , and both the unconditional  $\epsilon_{\theta}(x_t, t)$  and conditional  $\epsilon_{\theta}(x_t, t, z)$  denoising models are learned. For data sampling, let  $\bar{\epsilon}_{\theta} = w\epsilon_{\theta}(x_t, t, z) + (1 - w)\epsilon_{\theta}(x_t, t)$ , where  $w \geqslant 1$  is the guidance scale. At test time, given a text prompt  $z$ , the model generates conditional data according to  $p_{\theta}(x_0|z)$ .
4 Fine-tuning of Diffusion Models
In this section, we describe our approach for online RL fine-tuning of diffusion models. We first propose a Markov decision process (MDP) formulation for the denoising phase. We then use this MDP and present a policy gradient RL algorithm to update the original diffusion model. The RL algorithm optimizes an objective consisting of the reward and a KL term that ensures the updated model is not too far from the original one. We also present a modified supervised fine-tuning method with KL regularization and compares it with the RL approach.
4.1 RL Fine-tuning with KL Regularization
Let  $p_{\theta}(x_{0:T}|z)$  be a text-to-image diffusion model where  $z$  is some text prompt distributed according to  $p(z)$ , and  $r(x_0,z)$  be a reward model (typically trained using human assessment of images).
MDP formulation: The denoising process of DDPMs can be modeled as a  $T$ -horizon MDP:
$$
s _ {t} = (z, x _ {T - t}), \qquad a _ {t} = x _ {T - t - 1}, \qquad P _ {0} (s _ {0}) = \big (p (z), \mathcal {N} (0, I) \big), \qquad P (s _ {t + 1} \mid s _ {t}, a _ {t}) = (\delta_ {z}, \delta_ {a _ {t}}),
$$
$$
R \left(s _ {t}, a _ {t}\right) = \left\{ \begin{array}{l l} r \left(s _ {t + 1}\right) = r \left(x _ {0}, z\right) & \text {i f} t = T - 1, \\ 0 & \text {o t h e r w i s e .} \end{array} , \quad \pi_ {\theta} \left(a _ {t} \mid s _ {t}\right) = p _ {\theta} \left(x _ {T - t - 1} \mid x _ {T - t}, z\right), \right. \tag {4}
$$
in which  $s_t$  and  $a_t$  are the state and action at time-step  $t$ ,  $P_0$  and  $P$  are the initial state distribution and the dynamics,  $R$  is the reward function, and  $\pi_{\theta}$  is the parameterized policy. As a result, optimizing policy  $\pi_{\theta}$  in (4) is equivalent to fine-tuning the underlying DDPM. Finally, we denote by  $\delta_z$  the Dirac distribution at  $z$ .
It can be seen from the MDP formulation in (4) that the system starts by sampling its initial state  $s_0$  from the Gaussian distribution  $\mathcal{N}(0, I)$ , similar to the first state of the dinoising process  $x_T$ . Given the MDP state  $s_t$ , which corresponds to state  $x_{T - t}$  of the denoising process, the policy takes the action at time-step  $t$  as the next denoising state, i.e.,  $a_t = x_{T - t - 1}$ . As a result of this action, the system transitions deterministically to a state identified by the action (i.e., the next state of the denoising process). The reward is zero, except at the final step in which the quality of the image at the end of the denoising process is evaluated w.r.t. the prompt, i.e.,  $r(x_0, z)$ .
A common goal in re-training/fine-tuning the diffusion models is to maximize the expected reward of the generated images given the prompt distribution, i.e.,
$$
\min  _ {\theta} \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\theta} \left(x _ {0} \mid z\right)} [ - r (x _ {0}, z) ]. \tag {5}
$$
The gradient of this objective function can be obtained as follows:
Lemma 4.1 (A modification of Theorem 4.1 in [5]). If  $p_{\theta}(x_{0:T}|z)r(x_0,z)$  and  $\nabla_{\theta}p_{\theta}(x_{0:T}|z)r(x_0,z)$  are continuous functions of  $\theta$ , then we can write the gradient of the objective in (5) as
$$
\nabla_ {\theta} \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\theta} (x _ {0} | z)} [ - r (x _ {0}, z) ] = \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\theta} (x _ {0: T} | z)} \left[ - r (x _ {0}, z) \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log p _ {\theta} (x _ {t - 1} | x _ {t}, z) \right]. \quad (6)
$$
Proof. We present the proof in Appendix A.1.
<sup>1</sup>We keep the covariance in DDPM as constant,  $\Sigma_{t} = \Sigma$ , and only train  $\mu_{\theta}$  (see Eq. 1). This would naturally provide a stochastic policy for online exploration.
Algorithm 1 DPOK: Diffusion policy optimization with KL regularization
Input: reward model  $r$ , pre-trained model  $p_{\mathrm{pre}}$ , current model  $p_{\theta}$ , batch size  $m$ , text distribution  $p(z)$
Initialize  $p_{\theta} = p_{\mathrm{pre}}$
while  $\theta$  not converged do
Obtain  $m$  i.i.d. samples by first sampling  $z \sim p(z)$  and then  $x_{0:T} \sim p_{\theta}(x_{0:T}|z)$
Compute the gradient using Eq. (9) and update  $\theta$
Output: Fine-tuned diffusion model  $p_{\theta}$
Equation (6) is equivalent to the gradient used by the popular policy gradient algorithm, REINFORCE, to update a policy in the MDP (4). The gradient in (6) is estimated from trajectories  $p_{\theta}(x_{0:T}|z)$  generated by the current policy, and then used to update the policy  $p_{\theta}(x_{t-1}|x_t,z)$  in an online fashion.
Note that REINFORCE is not the only way to solve (5). Alternatively, one could compute the gradient through the trajectories to update the model; but the multi-step nature of diffusion models makes this approach memory inefficient and potentially prone to numerical instability. Consequently, scaling it to high-resolution images becomes challenging. For this reason, we adopt policy gradient to train large-scale diffusion models like Stable Diffusion [30].
Adding KL regularization. The risk of fine-tuning purely based on the reward model learned from human or AI feedback is that the model may overfit to the reward and discount the "skill" of the initial diffusion model to a greater degree than warranted. To avoid this phenomenon, similar to [23, 41], we add the KL between the fine-tuned and pre-trained models as a regularizer to the objective function. Unlike the language models in which the KL regularizer is computed over the entire sequence/trajectory (of tokens), in text-to-image models, it makes sense to compute it only for the final image, i.e.,  $\mathrm{KL}\big(p_{\theta}(x_0|z)\| p_{\mathrm{pre}}(x_0|z)\big)$ . Unfortunately,  $p_{\theta}(x_0|z)$  is a marginal (see the integral in Section 3) and its closed-form is unknown. As a result, we propose to add an upper-bound of this KL-term to the objective function.
Lemma 4.2. Suppose  $p_{pre}(x_{0:T}|z)$  and  $p_{\theta}(x_{0:T}|z)$  are Markov chains conditioned on the text prompt  $z$  that both start at  $x_{T} \sim \mathcal{N}(0, I)$ . Then, we have
$$
\mathbb {E} _ {p (z)} \left[ \mathrm {K L} \left(p _ {\theta} \left(x _ {0} \mid z\right)\right) \| p _ {p r e} \left(x _ {0} \mid z\right)\right) ] \leqslant \mathbb {E} _ {p (z)} \left[ \sum_ {t = 1} ^ {T} \mathbb {E} _ {p _ {\theta} \left(x _ {t} \mid z\right)} \left[ \mathrm {K L} \left(p _ {\theta} \left(x _ {t - 1} \mid x _ {t}, z\right) \| p _ {p r e} \left(x _ {t - 1} \mid x _ {t}, z\right)\right) \right] \right]. \tag {7}
$$
We report the proof of Lemma 4.2 in Appendix A.2. Intuitively, this lemma tells us that the divergence between the two distributions over the output image  $x_0$  is upper-bounded by the sum of the divergences between the distributions over latent  $x_t$  at each diffusion step.
Using the KL upper-bound in (7), we propose the following objective for regularized training:
$$
\left. \right. \mathbb {E} _ {p (z)} \left[ \alpha \mathbb {E} _ {p _ {\theta} \left(x _ {0: T} | z\right)} \left[ - r \left(x _ {0}, z\right)\right] + \beta \sum_ {t = 1} ^ {T} \mathbb {E} _ {p _ {\theta} \left(x _ {t} | z\right)} \left[ \mathrm {K L} \left(p _ {\theta} \left(x _ {t - 1} \mid x _ {t}, z\right) \| p _ {\text {p r e}} \left(x _ {t - 1} \mid x _ {t}, z\right)\right)\right]\right], \tag {8}
$$
where  $\alpha, \beta$  are the reward and KL weights, respectively. We use the following gradient to optimize the objective (8):
$$
\mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\theta} \left(x _ {0: T} | z\right)} \left[ - \alpha r \left(x _ {0}, z\right) \sum_ {t = 1} ^ {T} \nabla_ {\theta} \log p _ {\theta} \left(x _ {t - 1} \mid x _ {t}, z\right) + \beta \sum_ {t = 1} ^ {T} \nabla_ {\theta} \mathrm {K L} \left(p _ {\theta} \left(x _ {t - 1} \mid x _ {t}, z\right) \| p _ {\text {p r e}} \left(x _ {t - 1} \mid x _ {t}, z\right)\right) \right]. \tag {9}
$$
Note that (9) has one term missing from the exact gradient of (8) (see Appendix A.3). Removing this term is for efficient training. The pseudo-code of our algorithm, which we refer to as DPOK, is summarized in Algorithm 1. To reuse historical trajectories and be more sample efficient, we can also use importance sampling and clipped gradient, similar to [36]. We refer readers to Appendix A.6 for these details.
4.2 Supervised Learning with KL Regularization
We now introduce KL regularization into supervised fine-tuning (SFT), which allows for a more meaningful comparison with our KL-regularized RL algorithm (DPOK). We begin with a supervised
fine-tuning objective similar to that used in [17], i.e.,
$$
\mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\text {p r e}} (x _ {0} | z)} [ - r (x _ {0}, z) \log p _ {\theta} (x _ {0} | z) ]. \tag {10}
$$
To compare with RL fine-tuning, we augment the supervised objective with a similar KL regularization term. Under the supervised learning setting, we consider  $\mathrm{KL}(p_{\mathrm{pre}}(x_0|z)\| p_\theta (x_0|z))$  , which is equivalent to minimizing  $\mathbb{E}_{p_{\mathrm{pre}}(x_0|z)}[-\log p_\theta (x_0|z)]$  , given any prompt  $z$  . Let  $q(x_{t})$  be the forward process used for training. Since the distribution of  $x_0$  is generally not tractable, we use approximate upper-bounds in the lemma below (see derivation in Appendix A.4).
Lemma 4.3. Let  $\gamma$  be the regularization weight and  $\tilde{\mu}_t(x_t,x_0)\coloneqq \frac{\sqrt{\bar{\alpha}_{t - 1}}\beta_t}{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t - 1})}{1 - \bar{\alpha}_t} x_t$ . Assume  $r(x_0,z) + \gamma >0$ , for all  $x_0$  and  $z$ . Then, we have
$$
\begin{array}{l} \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {p r e} (x _ {0} | z)} [ - (r (x _ {0}, z) + \gamma) \log p _ {\theta} (x _ {0} | z) ] \\ \leqslant \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {p r e} (x _ {0} | z)} [ (r (x _ {0}, z) + \gamma) \sum_ {t > 1} \mathbb {E} _ {q \left(x _ {t} \mid x _ {0}, z\right)} [ \frac {1}{2 \sigma_ {t} ^ {2}} \left| \left| \tilde {\mu} _ {t} \left(x _ {t}, x _ {0}\right) - \mu_ {\theta} \left(x _ {t}, t, z\right) \right| ^ {2} \right] ] + C _ {1}. \tag {11} \\ \end{array}
$$
Moreover, we also have another weaker upper-bound in which  $C_1$  and  $C_2$  are two constants:
$$
\begin{array}{l} \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {p r e} (x _ {0} | z)} [ - (r (x _ {0}, z) + \gamma) \log p _ {\theta} (x _ {0} | z) ] \\ \leqslant \mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {p r e} (x _ {0} | z)} [ \sum_ {t > 1} \mathbb {E} _ {q \left(x _ {t} \mid x _ {0}, z\right)} [ \frac {r \left(x _ {0} , z\right) \left| \left| \tilde {\mu} _ {t} \left(x _ {t} , x _ {0}\right) - \mu_ {\theta} \left(x _ {t} , t , z\right) \right| \right| ^ {2}}{2 \sigma_ {t} ^ {2}} \tag {12} \\ \left. \left. + \frac {\gamma \left(\left| \left| \mu_ {p r e} \left(x _ {t} , t , z\right) - \mu_ {\theta} \left(x _ {t} , t , z\right) \right| \right| ^ {2}\right)}{2 \sigma_ {t} ^ {2}} \right] \right] + C _ {2}. \\ \end{array}
$$
In Lemma 4.3, we introduce KL regularization (Eq. (11)) for supervised fine-tuning, which can be incorporated by adjusting the original reward with a shift factor  $\gamma$  in the reward-weighted loss, smoothing the weighting of each sample towards the uniform distribution. We refer to this regularization as KL-D since it is based only on data from the pre-trained model. The induced supervised training objective for KL-D is as follows:
$$
\mathbb {E} _ {p (z)} \mathbb {E} _ {p _ {\text {p r e}} (x _ {0} | z)} \left[ (r (x _ {0}, z) + \gamma) \sum_ {t > 1} \mathbb {E} _ {q \left(x _ {t} \mid x _ {0}, z\right)} \left[ \frac {\left| \left| x _ {0} - f _ {\theta} \left(x _ {t} , t , z\right) \right| \right| ^ {2}}{2 \sigma_ {t} ^ {2}} \right] \right]. \tag {13}
$$
We also consider another KL regularization presented in Eq. (12). This KL regularization can be implemented by introducing an additional term in the reward-weighted loss. This extra term penalizes the  $L_{2}$ -distance between the denoising directions derived from the pre-trained and current models. We refer to it as KL-O because it also regularizes the output from the current model to be close to that from the pre-trained model. The induced supervised training objective for KL-O is as follows: