SlowGuess/NeurIPS_2023_OCR huggingface dataset:
TaeHo Yoon $^{1}$  Kibeom Myoung $^{1}$  Keon Lee $^{4}$  Jaewoong Cho $^{4}$  Albert No $^{3}$  Ernest K. Ryu $^{1,2}$
$^{1}$ Department of Mathematical Science, Seoul National University  $^{2}$ Interdisciplinary Program in Artificial Intelligence, Seoul National University  $^{3}$ Department of Electronic and Electrical Engineering, Hongik University  $^{4}$ KRAFTON
Diffusion models have recently shown remarkable success in high-quality image generation. Sometimes, however, a pre-trained diffusion model exhibits partial misalignment in the sense that the model can generate good images, but it sometimes outputs undesirable images. If so, we simply need to prevent the generation of the bad images, and we call this task censoring. In this work, we present censored generation with a pre-trained diffusion model using a reward model trained on minimal human feedback. We show that censoring can be accomplished with extreme human feedback efficiency and that labels generated with a mere few minutes of human feedback are sufficient. Code available at: https://github.com/tetrzim/diffusion-human-feedback.
Diffusion probabilistic models [19, 12, 42] have recently shown remarkable success in high-quality image generation. Much of the progress is driven by scale [35, 36, 38], and this progression points to a future of spending high costs to train a small number of large-scale foundation models [4] and deploying them, sometimes with fine-tuning, in various applications. In particular use cases, however, such pre-trained diffusion models may be misaligned with goals specified before or after the training process. An example of the former is text-guided diffusion models occasionally generating content with nudity despite the text prompt containing no such request. An example scenario of the latter is deciding that generated images should not contain a certain type of concepts (for example, human faces) even though the model was pre-trained on images with such concepts.
Fixing misalignment directly through training may require an impractical cost of compute and data. To train a large diffusion model again from scratch requires compute costs of up to hundreds of thousands of USD [30, 29]. To fine-tune a large diffusion model requires data size ranging from 1,000 [28] to 27,000 [25].<sup>1</sup> We argue that such costly measures are unnecessary when the pre-trained model is already capable of sometimes generating "good" images. If so, we simply need to prevent the generation of "bad" images, and we call this task censoring. (Notably, censoring does not aim to improve the "good" images.) Motivated by the success of reinforcement learning with human feedback (RLHF) in language domains [9, 50, 43, 33], we perform censoring using human feedback.
In this work, we present censored generation with a pre-trained diffusion model using a reward model trained on extremely limited human feedback. Instead of fine-tuning the pre-trained diffusion model, we train a reward model on labels generated with a few minutes of human feedback and perform guided generation. By not fine-tuning the diffusion model (score network), we reduce both compute
<sup>1</sup>The prior work [28] fine-tunes a pre-trained diffusion model on a new dataset of size 1k using a so-called adapter module while [25] improves text-to-image alignment using 27k human-feedback data.
(a) Baseline: MNIST class "7"
(c) Baseline: LSUN Church with LDM
(e) Baseline: ImageNet class "tench" (fish)
(g) Baseline: LSUN bedroom
Figure 1: Uncensored baseline vs. censored generation. Setups are precisely defined in Section 5. Due to space constraints, we present selected representative images here. Full sets of non-selected samples are shown in the appendix.
(b) Censored: Crossed 7
(d) Censored: Stock photo watermarks
(f) Censored: Human faces
(h) Censored: Broken images
(i) Baseline: Stable Diffusion, "A photo of a human"
(j) Censored: Embedded text
Figure 1: (Continued) Uncensored baseline vs. censored generation. Setups are precisely defined in Section 5. Due to space constraints, we present selected representative images here. Full sets of non-selected samples are shown in the appendix.
and data requirements for censored generation to negligible levels. (Negligible compared to any amount of compute and man-hours an ML scientist would realistically spend building a system with a diffusion model.) We conduct experiments within multiple setups demonstrating how minimal human feedback enables removal of target concepts. The specific censoring targets we consider are: A handwriting variation ("crossed 7"s) in MNIST [11]; Watermarks in the LSUN [47] church images; Human faces in the ImageNet [10] class "tench"; "Broken" images in the generation of LSUN bedroom images.
Contribution. Most prior work focuses on training new capabilities into diffusion models, and this inevitably requires large compute and data. Our main contribution is showing that a very small amount of human feedback data and computation is sufficient for guiding a pre-trained diffusion model to do what it can already do while suppressing undesirable behaviors.
1.1 Background on diffusion probabilistic models
Due to space constraints, we defer the comprehensive review of prior works to Appendix D. In this section, we briefly review the standard methods of diffusion probabilistic models (DPM) and set up the notation. For the sake of simplicity and specificity, we only consider the DPMs with the variance preserving SDE.
Consider the variance preserving (VP) SDE
$$
d X _ {t} = - \frac {\beta_ {t}}{2} X _ {t} d t + \sqrt {\beta_ {t}} d W _ {t}, \quad X _ {0} \sim p _ {0} \tag {1}
$$
for  $t\in [0,T]$ , where  $\beta_{t} > 0$ ,  $X_{t}\in \mathbb{R}^{d}$ , and  $W_{t}$  is a  $d$ -dimensional Brownian motion. The process  $\{X_{t}\}_{t\in [0,T]}$  has the marginal distributions given by
$$
X _ {t} \stackrel {{\mathcal {D}}} {{=}} \sqrt {\alpha_ {t}} X _ {0} + \sqrt {1 - \alpha_ {t}} \varepsilon_ {t}, \quad \alpha_ {t} = e ^ {- \int_ {0} ^ {t} \beta_ {s} d s}, \varepsilon_ {t} \sim \mathcal {N} (0, I)
$$
for  $t \in [0, T]$  [39, Chapter 5.5]. Let  $p_t$  denote the density for  $X_t$  for  $t \in [0, T]$ . Anderson's theorem [1] tells us that the reverse-time SDE by
$$
d \overline {{X}} _ {t} = \beta_ {t} \left(- \nabla \log p _ {t} (\overline {{X}} _ {t}) - \frac {1}{2} \overline {{X}} _ {t}\right) d t + \sqrt {\beta_ {t}} d \overline {{W}} _ {t}, \qquad \overline {{X}} _ {T} \sim p _ {T},
$$
where  $\{\overline{W}_t\}_{t\in [0,T]}$  is a reverse-time Brownian motion, satisfies  $\overline{X}_t\stackrel {\mathcal{D}}{=}X_t\sim p_t$
In DPMs, the initial distribution is set as the data distribution, i.e.,  $p_0 = p_{\mathrm{data}}$  in (1), and a score network  $s_\theta$  is trained so that  $s_\theta(X_t, t) \approx \nabla \log p_t(X_t)$ . For notational convenience, one often uses the error network  $\varepsilon_\theta(X_t, t) = -\sqrt{1 - \alpha_t} s_\theta(X_t, t)$ . Then, the reverse-time SDE is approximated by
$$
d \bar {X} _ {t} = \beta_ {t} \left(\frac {1}{\sqrt {1 - \alpha_ {t}}} \varepsilon_ {\theta} (\bar {X} _ {t}, t) - \frac {1}{2} \bar {X} _ {t}\right) d t + \sqrt {\beta_ {t}} d \bar {W} _ {t}, \quad \bar {X} _ {T} \sim \mathcal {N} (0, I)
$$
for  $t\in [0,T]$
When an image  $X$  has a corresponding label  $Y$ , classifier guidance [40, 12] generates images from
$$
p _ {t} \left(X _ {t} \mid Y\right) \propto p _ {t} \left(X _ {t}, Y\right) = p _ {t} \left(X _ {t}\right) p _ {t} \left(Y \mid X _ {t}\right)
$$
for  $t\in [0,T]$  using
$$
\hat {\varepsilon} _ {\theta} \left(\bar {X} _ {t}, t\right) = \varepsilon_ {\theta} \left(\bar {X} _ {t}, t\right) - \omega \sqrt {1 - \alpha_ {t}} \nabla \log p _ {t} (Y \mid \bar {X} _ {t})
$$
$$
d \overline {{X}} _ {t} = \beta_ {t} \left(\frac {1}{\sqrt {1 - \alpha_ {t}}} \hat {\varepsilon} _ {\theta} (\overline {{X}} _ {t}, t) - \frac {1}{2} \overline {{X}} _ {t}\right) d t + \sqrt {\beta_ {t}} d \overline {{W}} _ {t}, \qquad \overline {{X}} _ {T} \sim \mathcal {N} (0, I),
$$
where  $\omega > 0$ . This requires training a separate time-dependent classifier approximating  $p_t(Y \mid X_t)$ . In the context of censoring, we use a reward model with binary labels in place of a classifier.
2 Problem description: Censored sampling with human feedback
Informally, our goal is:
Given a pre-trained diffusion model that is partially misaligned in the sense that generates both "good" and "bad" images, fix/modify the generation process so that only good images are produced.
The meaning of "good" and "bad" depends on the context and will be specified through human feedback. For the sake of precision, we define the terms "benign" and "malign" to refer to the good and bad images: A generated image is malign if it contains unwanted features to be censored and is benign if it is not malign.
Our assumptions are: (i) the pre-trained diffusion model does not know which images are benign or malign, (ii) a human is willing to provide minimal ( $\sim$  3 minutes) feedback to distinguish benign and malignant images, and (iii) the compute budget is limited.
Mathematical formalism. Suppose a pre-trained diffusion model generates images from distribution  $p_{\mathrm{data}}(x)$  containing both benign and malign images. Assume there is a function  $r(x) \in (0,1)$  representing the likelihood of  $x$  being benign, i.e.,  $r(x) \approx 1$  means image  $x$  is benign and should be considered for sampling while  $r(x) \approx 0$  means image  $x$  is malign and should not be sampled. We mathematically formalize our goal as: Sample from the censored distribution
$$
p _ {\mathrm {c e n s o r}} (x) \propto p _ {\mathrm {d a t a}} (x) r (x).
$$
Human feedback. The definition of benign and malign images are specified through human feedback. Specifically, we ask a human annotator to provide binary feedback  $Y \in \{0,1\}$  for each image  $X$  through a simple graphical user interface shown in Appendix E. The feedback takes 1-3 human-minutes for the relatively easier censoring tasks and at most 10-20 human-minutes for the most complex task that we consider. Using the feedback data, we train a reward model  $r_{\psi} \approx r$ , which we further detail in Section 3.
Evaluation. The evaluation criterion of our methodology are the human time spent providing feedback, quantified by direct measurement, and sample quality, quantified by precision and recall.
In this context, precision is the proportion of benign images, and recall is the sample diversity of the censored generation. Precision can be directly measured by asking human annotators to label the final generated images, but recall is more difficult to measure. Therefore, we primarily focus on precision for quantitative evaluation. We evaluate recall qualitatively by providing the generated images for visual inspection.
3 Reward model and human feedback
Let  $Y$  be a random variable such that  $Y = 1$  if  $X$  is benign and  $Y = 0$  if  $X$  is malignant. Define the time-independent reward function as
$$
r (X) = \mathbb {P} (Y = 1 \mid X).
$$
Algorithm 1 Reward model ensemble
Require: Images: malign  $\{X^{(1)},\ldots ,X^{(N_M)}\}$  , benign  $\{X^{(N_M + 1)},\dots ,X^{(N_M + N_B)}\} (N_M <   N_B)$  for  $k = 1,\ldots ,K$  do Randomly select with replacement  $N_{M}$  benign samples  $X^{(N_M + i_1)},\ldots ,X^{(N_M + i_{N_M})}$  Train reward model  $r_{\psi_k}^{(k)}$  with  $\{X^{(1)},\dots ,X^{(N_M)}\} \cup \{X^{(N_M + i_1)},\dots ,X^{(N_M + i_{N_M})}\}$  end for return  $r_{\psi} = \prod_{k = 1}^{K}r_{\psi_k}^{(k)}$
As we later discuss in Section 4, time-dependent guidance requires a time-dependent reward function. Specifically, let  $X \sim p_{\mathrm{data}}$  and  $Y$  be its label. Let  $\{X_t\}_{t \in [0,T]}$  be images corrupted by the VP SDE (1) with  $X_0 = X$ . Define the time-dependent reward function as
$$
r _ {t} \left(X _ {t}\right) = \mathbb {P} (Y = 1 \mid X _ {t}) \quad \text {f o r} t \in [ 0, T ].
$$
We approximate the reward function  $r$  with a reward model  $r_{\psi}$ , i.e., we train
$$
r _ {\psi} (X) \approx r (X) \qquad \text {o r} \qquad r _ {\psi} (X _ {t}, t) \approx r _ {t} (X _ {t}),
$$
using human feedback data  $(X^{(1)},Y^{(1)}),\ldots ,(X^{(N)},Y^{(N)})$  . (So the time-dependent reward model uses  $(X_{t}^{(n)},Y^{(n)})$  as training data.) We use weighted binary cross entropy loss. In this section, we describe the most essential components of the reward model while deferring details to Appendix F.
The main technical challenge is achieving extreme human-feedback efficiency. Specifically, we have  $N < 100$  in most setups we consider. Finally, we clarify that the diffusion model (score network) is not trained or fine-tuned. We use relatively large pre-trained diffusion models [12, 36], but we only train the relatively lightweight reward model  $r_{\psi}$ .
3.1 Reward model ensemble for benign-dominant setups
In some setups, benign images constitute the majority of uncensored generation. Section 5.2 considers such a benign-dominant setup, where  $11.4\%$  of images have stock photo watermarks and the goal is to censor the watermarks. A random sample of images provided to a human annotator will contain far more benign than malign images.
To efficiently utilize the imbalanced data in a sample-efficient way, we propose an ensemble method loosely inspired by ensemble-based sample efficient RL methods [23, 6]. The method trains  $K$  reward models  $r_{\psi_1}^{(1)}, \ldots, r_{\psi_K}^{(K)}$ , each using a shared set of  $N_M$  (scarce) malign images joined with  $N_M$  benign images randomly subsampled bootstrap-style from the provided pool of  $N_B$  (abundant) benign data as in Algorithm 1. The final reward model is formed as  $r_\psi = \prod_{k=1}^K r_{\psi_k}^{(k)}$ . Given that a product becomes small when any of its factor is small,  $r_\psi$  is effectively asking for unanimous approval across  $r_{\psi_1}^{(1)}, \ldots, r_{\psi_K}^{(K)}$ .
In experiments, we use  $K = 5$ . We use the same neural network architecture for  $r_{\psi_1}^{(1)} \ldots, r_{\psi_K}^{(K)}$ , whose parameters  $\psi_1, \ldots, \psi_K$  are either independently randomly initialized or transferred from the same pre-trained weights as discussed in Section 3.3. We observe that the ensemble method significantly improves the precision of the model without perceivably sacrificing recall.
3.2 Imitation learning for malign-dominant setups
In some setups, malign images constitute the majority of uncensored generation. Section 5.3 considers such a malign-dominant setup, where  $69\%$  of images are tench (fish) images with human faces and the goal is to censor the images with human faces. Since the ratio of malign images starts out high, a single round of human feedback and censoring may not sufficiently reduce the malign ratio.
Therefore, we propose an imitation learning method loosely inspired by imitation learning RL methods such as DAgger [37]. The method collects human feedback data in multiple rounds and improves the reward model over the rounds as described in Algorithm 2. Our experiment of Section 5.3 indicates that 2-3 rounds of imitation learning dramatically reduce the ratio of malign images. Furthermore,
Algorithm 2 Imitation learning of reward model
Require: Pre-trained  $\varepsilon_{\theta}$  . Initialize  $\mathcal{D} = \emptyset$  Sample  $X^{(1)},\ldots ,X^{(N_1)}$  using  $\varepsilon_{\theta}$  and no censoring. Receive  $Y^{(1)},\dots ,Y^{(N_1)}$  from human feedback. Add data to buffer:  $\mathcal{D}\gets \{(X^{(i)},Y^{(i)})\}_{i = 1}^{N_1}$  Train reward model  $r_{\psi}$  with  $\mathcal{D}$  for  $r = 2,\dots ,R$  do Sample  $X^{(1)},\ldots ,X^{(N_r)}$  using  $\varepsilon_{\theta}$  and censoring with  $r_{\psi}$  Receive  $Y^{(1)},\ldots ,Y^{(N_r)}$  from human feedback. Add data to buffer:  $\mathcal{D}\gets \{(X^{(i)},Y^{(i)})\}_{i = 1}^{N_r}$  Train reward model  $r_{\psi}$  with  $\mathcal{D}$  end for return  $r_{\psi}$
imitation learning is a practical model of an online scenario where one continuously trains and updates the reward model  $r_{\psi}$  while the diffusion model is continually deployed.
Ensemble vs. imitation learning. In the benign-dominant setup, imitation learning is too costly in terms of human feedback since acquiring sufficiently many ( $\sim 10$ ) malign labels may require the human annotator to go through too many benign labels ( $\sim 1000$ ) for the second round of human feedback and censoring. In the malign-dominant setup, one can use a reward model ensemble, where reward models share the benign data while bootstrap-subsampling the malign data, but we empirically observe this to be ineffective. We attribute this asymmetry to the greater importance of malign data over benign data; the training objective is designed so as our primary goal is to censor malign images.
3.3 Transfer learning for time-independent reward
To further improve human-feedback efficiency, we use transfer learning. Specifically, we take a ResNet18 model [17, 18] pre-trained on ImageNet1k [10] and replace the final layer with randomly initialized fully connected layers which have 1-dimensional output features. We observe training all layers to be more effective than training only the final layers. We use transfer learning only for training time-independent reward models, as pre-trained time-dependent classifiers are less common. Transfer learning turns out to be essential for complex censoring tasks (Sections 5.2, 5.4 and 5.5), but requires guidance techniques other than the simple classifier guidance (see Section 4).
4 Sampling
In this section, we describe how to perform censored sampling with a trained reward model  $r_{\psi}$ . We follow the notation of Section 1.1.
Time-dependent guidance. Given a time-dependent reward model  $r_{\psi}(X_t, t)$ , our censored generation follows the SDE
$$
\begin{array}{l} \hat {\varepsilon} _ {\theta} (\bar {X} _ {t}, t) = \varepsilon_ {\theta} (\bar {X} _ {t}, t) - \omega \sqrt {1 - \alpha_ {t}} \nabla \log r _ {t} (\bar {X} _ {t}) \\ d \bar {X} _ {t} = \beta_ {t} \left(\frac {1}{\sqrt {1 - \alpha_ {t}}} \hat {\varepsilon} _ {\theta} (\bar {X} _ {t}, t) - \frac {1}{2} \bar {X} _ {t}\right) d t + \sqrt {\beta_ {t}} d \bar {W} _ {t}, \quad \bar {X} _ {T} \sim \mathcal {N} (0, I) \tag {2} \\ \end{array}
$$
for  $t \in [0, T]$  with  $\omega > 0$ . From the standard classifier-guidance arguments [42, Section I], it follows that  $X_0 \sim p_{\mathrm{censor}}(x) \propto p_{\mathrm{data}}(x)r(x)$  approximately when  $\omega = 1$ . The parameter  $\omega > 0$ , which we refer to as the guidance weight, controls the strength of the guidance, and it is analogous to the "gradient scale" used in prior works [12]. Using  $\omega > 1$  can be viewed as a heuristic to strengthen the effect of the guidance, or it can be viewed as an effort to sample from  $p_{\mathrm{censor}}^{(\omega)} \propto p_{\mathrm{data}}r^{\omega}$ .
Time-independent guidance. Given a time-independent reward model  $r_{\psi}(X_t)$ , we adopt the ideas of universal guidance [2] and perform censored generation via replacing the  $\hat{\varepsilon}_{\theta}$  of (2) with
$$
\begin{array}{l} \hat {\varepsilon} _ {\theta} \left(\bar {X} _ {t}, t\right) = \varepsilon_ {\theta} \left(\bar {X} _ {t}, t\right) - \omega \sqrt {1 - \alpha_ {t}} \nabla \log r \left(\hat {X} _ {0}\right), \quad \text {w h e r e} \\ \hat {X} _ {0} = \mathbb {E} \left[ X _ {0} \mid X _ {t} = \bar {X} _ {t} \right] = \frac {\bar {X} _ {t} - \sqrt {1 - \alpha_ {t}} \varepsilon_ {\theta} (\bar {X} _ {t} , t)}{\sqrt {\alpha_ {t}}} \tag {3} \\ \end{array}
$$
(a) MNIST: Censoring "crossed 7"
(b) LSUN church: Censoring watermarks
Figure 2: Mean proportion of malign images after censoring with standard deviation over 5 trials, each measured with 500 samples. Reward ensemble outperforms non-ensemble models, and the universal guidance components further improve the results. Left: Censoring "crossed 7" from MNIST. Before censoring, the proportion is  $11.9\%$ . The mean values of each point are:  $1.30\%$ ,  $0.98\%$ ,  $0.60\%$ , and  $0.42\%$ . Right: Censoring watermarks from LSUN Church. Before censoring, the proportion is  $11.4\%$ . The mean values of each point are:  $3.02\%$ ,  $3.84\%$ ,  $1.36\%$ , and  $0.76\%$ .
for  $t \in [0, T]$  with  $\omega > 0$ . To clarify,  $\nabla$  differentiates through  $\hat{X}_0$ . While this method has no mathematical guarantees, prior work [2] has shown strong empirical performance in related setups.
Backward guidance and recurrence. The prior work [2] proposes backward guidance and self-recurrence to further strengthen the guidance. We find that adapting these methods to our setup improves the censoring performance. We provide the detailed description in Appendix G.
We now present the experimental results. Precision (censoring performance) was evaluated with human annotators labeling generated images. The human feedback time we report includes annotation of training data for the reward model  $r_{\psi}$ , but does not include the annotation of the evaluation data.
5.1 MNIST: Censoring 7 with a strike-through cross
In this setup, we censor a handwriting variation called "crossed 7", which has a horizontal stroke running across the digit, from an MNIST generation, as shown in Figure 1a. We pre-train our own diffusion model (score network). In this benign-dominant setup, the baseline model generates about  $11.9\%$  malign images.
We use 10 malign samples to perform censoring. This requires about 100 human feedback labels in total, which takes less than 2 minutes to collect. We observe that such minimal feedback is sufficient for reducing the proportion of crossed 7s to  $0.42\%$  as shown in Figure 1b and Figure 2a. Further details are provided in Appendix H.
Ablation studies. We achieve our best results by combining the time-dependent reward model ensemble method described in Section 3.1 and the universal guidance components (backward guidance with recurrence) detailed in Appendix G. We verify the effectiveness of each component through an ablation study, summarized in Figure 2a. Specifically, we compare the censoring results using a reward model ensemble (labeled "Ensemble" in Figure 2a) with the cases of using (i) a single reward model within the ensemble (trained on 10 malign and 10 benign images; labeled "Single") and (ii) a standalone reward model separately trained on the union of all training data (10 malign and 50 benign images; labeled "Union") used in ensemble training. We also show that the backward and recurrence components do provide an additional benefit (labeled "Ensemble+Universal").
2If we simply perform time-dependent guidance with a time-independent reward function  $r_{\psi}(X)$ , the observed performance is poor. This is because  $r_{\psi}(X)$  fails to provide meaningful guidance when the input is noisy, and this empirical behavior agrees with the prior observations of [32, Section 2.4] and [2, Section 3.1].
5.2 LSUN church: Censoring watermarks from latent diffusion model
In the previous experiment, we use a full-dimensional diffusion model that reverses the forward diffusion (1) in the pixel space. In this experiment, we demonstrate that censored generation with minimal human feedback also works with latent diffusion models (LDMs) [46, 36], which perform diffusion on a lower-dimensional latent representation of (variational) autoencoders. We use an  $\mathrm{LDM}^3$  pre-trained on the  $256\times 256$  LSUN Churches [36] and censor the stock photo watermarks. In this benign-dominant setup, the baseline model generates about  $11.4\%$  malign images.
Training a time-dependent reward model in the latent space to be used with an LDM would introduce additional complicating factors. Therefore, for simplicity and to demonstrate multiple censoring methods, we train a time-independent reward model ensemble and apply time-independent guidance as outlined in Section 4. To enhance human-feedback efficiency, we use a pre-trained ResNet18 model and use transfer learning as discussed in Section 3.3. We use 30 malign images, and the human feedback takes approximately 3 minutes. We observe that this is sufficient for reducing the proportion of images with watermarks to  $0.76\%$  as shown in Figure 1d and Figure 2b. Further details are provided in Appendix I.