SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shenghuan Sun
shenghuan.sun@ucsf.edu
Gregory M. Goldgof*
Memorial Sloan Kettering Cancer Center
goldgofg@mskcc.org
Atul Butte
atul.butte@ucsf.edu
Ahmed M. Alaa
UC Berkeley and UCSF
amalaa@berkeley.edu
Generative models capable of capturing nuanced clinical features in medical images hold great promise for facilitating clinical data sharing, enhancing rare disease datasets, and efficiently synthesizing annotated medical images at scale. Despite their potential, assessing the quality of synthetic medical images remains a challenge. While modern generative models can synthesize visually-realistic medical images, the clinical validity of these images may be called into question. Domain-agnostic scores, such as FID score, precision, and recall, cannot incorporate clinical knowledge and are, therefore, not suitable for assessing clinical sensibility. Additionally, there are numerous unpredictable ways in which generative models may fail to synthesize clinically plausible images, making it challenging to anticipate potential failures and manually design scores for their detection. To address these challenges, this paper introduces a pathologist-in-the-loop framework for generating clinically-plausible synthetic medical images. Starting with a diffusion model pretrained using real images, our framework comprises three steps: (1) evaluating the generated images by expert pathologists to assess whether they satisfy clinical desiderata, (2) training a reward model that predicts the pathologist feedback on new samples, and (3) incorporating expert knowledge into the diffusion model by using the reward model to inform a finetuning objective. We show that human feedback significantly improves the quality of synthetic images in terms of fidelity, diversity, utility in downstream applications, and plausibility as evaluated by experts. We also show that human feedback can teach the model new clinical concepts not annotated in the original training data. Our results demonstrate the value of incorporating human feedback in clinical applications where generative models may struggle to capture extensive domain knowledge from raw data alone.
Diffusion models have recently shown incredible success in the conditional generation of high-fidelity natural, stylized and artistic images [1-6]. The generative capabilities of these models can be leveraged to create synthetic data in application domains where obtaining large-scale annotated datasets is challenging. The medical imaging field is one such domain, where there is often a difficulty in obtaining high-quality labeled datasets [7]. This difficulty may stem from the regulatory hurdles that impede data sharing [8], the costs involved in getting experts to manually annotate images [9], or the natural scarcity of data in rare diseases [10]. Generative (diffusion) models may provide a partial solution to these problems by synthesizing high-fidelity medical images that can be easily shared among researchers to replace or augment real data in downstream modeling applications [11, 12].
Step 1. Pathologist feedback collection
Step 2. Clinical plausibility reward modeling
Step 3. Clinically-informed finetuning
Figure 1: Overview of our pathologist-in-the-loop synthetic data generation framework. (1) In Step 1, a synthetic dataset is sampled from a generative model pretrained using a dataset of real medical images. The dataset is then inspected by a pathologist who examines each image to determine its plausibility based on a set of criteria. For each image, the pathologist provides binary feedback, labeling a synthetic image as "1" if it fails to meet all of the plausibility criteria. (2) In Step 2, the synthetic images and pathologist feedback obtained in Step 1 are used to train a reward model that predicts human feedback on new images. (3) Finally, the generative model is finetuned via an objective function that uses the reward model to incentivize the generation of clinically plausible images.
What sets medical image synthesis apart from image generation in other fields? Unlike mainstream generative modeling applications that prioritize visually realistic or artistically expressive images, synthetic medical images require a different approach. They must be grounded in objective clinical and biological knowledge, and as such, they leave no room for creative or unrestricted generation. Given that the ultimate goal of synthetic medical images is to be used in downstream modeling and analysis, they must faithfully reflect nuanced features that represent various clinical concepts, such as cell types [13], disease subtypes [7], and anatomies [14]. Off-the-shelf image generation models are not capable of recognizing or generating clinical concepts, rendering them unsuitable for generating plausible medical images without further adaptation [15, 16]. Therefore, our aim is to develop a framework for generating synthetic medical images that not only exhibit visual realism but also demonstrate biological plausibility and alignment with clinical expertise.
One way to generate synthetic medical images is to finetune a pretrained "foundation" vision model, such as Stable Diffusion, that has been trained on billions of natural images (such as the LAION-5B dataset [17]), using real medical images. With a sufficiently large set of medical images, we can expect the finetuned model to capture the clinical knowledge encoded in medical images. However, the sample sizes of annotated medical images are typically limited to a few thousand. When a large vision model is finetuned on such a dataset using generic objective functions (such as the likelihood function), the model may capture only the generic features that make the medical images appear visually realistic, but it may miss nuanced features that make them biologically plausible and compliant with clinical domain knowledge (see examples in the next Section). Designing domain-specific objective functions for finetuning that ensure a generative model adheres to clinical knowledge is challenging. The difficulty arises from the numerous unpredictable ways in which these models can generate images that lack clinical plausibility. As a result, it is impractical to anticipate every possible failure scenario and manually construct a loss function that penalizes such instances.
Summary of contributions. In this paper, we develop a pathologist-in-the-loop framework for synthesizing medical images that align with clinical knowledge. Our framework is motivated by the success of reinforcement learning with human feedback (RLHF) in aligning the outputs of large language models (LLMs) with human preference [18, 19], and is directly inspired by [20], where human feedback was used to align the visual outputs of a generative model with input text prompts. To generate clinically-plausible medical images, our framework (outlined in Figure 1) comprises 3 steps:
Step 1: We train a (conditional) diffusion model using real medical images. We then sample a synthetic dataset from the model to be evaluated by a pathologist. Each image is carefully examined, and the pathologist provides feedback on whether it meets the necessary criteria for clinical plausibility.
Step 2: We collate a dataset of synthetic images paired with pathologist feedback and train a reward model to predict the pathologist feedback, i.e., clinical plausibility, on new images.
Step 3: Finally, the reward model in Step 2 is utilized to incorporate expert knowledge into the generative model. This is achieved by finetuning the diffusion model using a reward-weighted loss function, which penalizes the generation of images that the pathologist considers clinically implausible.
Throughout this paper, we apply the steps above to the synthetic generation of bone marrow image patches, but the same conceptual framework can generalize to any medical imaging modality. We gathered pathologist feedback on thousands of synthetic images of various cell types generated by a conditional diffusion model. Then, we analyzed the impact of this feedback on the quality of the finetuned synthetic images. Our findings suggest that incorporating pathologist feedback significantly enhances the quality of synthetic images in terms of all existing quality metrics such as fidelity, accuracy of downstream predictive models, and clinical plausibility as evaluated by experts. Additionally, it also improves qualities that are not directly addressed in the pathologist evaluation, such as the diversity of synthetic samples. Furthermore, we show that human feedback can teach the generative model new clinical concepts, such as more refined identification of cell types, that are not annotated in the original training data. These results demonstrate the value of incorporating human feedback in clinical applications where generative models may not be readily suited to capture intricate and extensive clinical domain knowledge from raw data alone.
Figure 2: Samples for biologically-implausible synthetic bone marrow image patches across four different cell types. On the bottom panels, we show pathologist evaluations detailing the reasons for their implausibility.
2 Pathologist-in-the-Loop Generation of Synthetic Medical Images
In this Section, we provide a detailed description of our synthetic medical image generation framework. We will use a running example pertaining to single-cell images extracted from bone marrow aspirate whole slide images. Details of this setup and the dataset used in our study are provided in Section 4.
2.1 Step 1: Pathologist feedback collection
Generative model pretraining. The first step in our framework starts with training a generative model to synthesize medical images through the standard training procedure. As we discuss in more detail in Section 4, we utilized a dataset of 2,048 bone marrow image patches to train a conditional diffusion model [6]. The model was trained to generate class-conditional images, where an image class corresponds to a cell type. We conducted an exploratory analysis where we found that neither latent diffusion nor Stable Diffusion models yielded superior results compared to a customized diffusion model that we employed in this study (See Appendix A). We opted for using a class-conditional model rather than a text-conditional model as we found that existing pretrained vision-language models were not fit for capturing the scientific jargon related to bone marrow cell types.
Given a dataset of real images  $\mathcal{D}_r = \{(x^i, c^i)\}_{i=1}^{n_r}$ , where  $x$  is a medical image and  $c \in \mathcal{C}$  is an image class, we train a diffusion model to generate class-conditional images through the forward process:
$$
x _ {t + 1} = x _ {t} - \alpha_ {t} \cdot \nabla_ {x} \log \left(p _ {\theta} \left(x _ {t} \mid x, c\right)\right) + \epsilon_ {t}, \tag {1}
$$
where  $\epsilon_t \sim \mathcal{N}(0, \rho^2)$  is the noise term at time-step  $t$ ,  $x_t$  is the data point at time-step  $t$ ,  $\alpha_t$  is the step size at  $t$ ,  $\theta$  is the model parameters and  $\nabla_x \log(p(x_t | x, c))$  is the gradient of the log probability distribution with respect to  $x$ , conditioned on the original data  $x$  and class  $c$ . Once the model is pretrained, we sample a synthetic dataset  $\mathcal{D}_s = \{(\widetilde{x}^j, \widetilde{c}^j)\}_{j=1}^{n_s}$  by first sampling a class  $\widetilde{c}$  from  $\mathcal{C}$ , and then sampling a medical image  $\widetilde{x}$  conditioned on the class  $\widetilde{c}$  through the reverse diffusion process.
**Pathologist evaluation.** Each image in the synthetic dataset  $\mathcal{D}_s = \{(\widetilde{x}^j, \widetilde{c}^j)\}_{j=1}^{n_s}$  generated by the pretrained model is inspected by an expert pathologist to assess its clinical plausibility. The objective of this evaluation is to identify the specific inaccuracies in the synthetic data that can only be identified by an expert, and provide feedback for the model to refine its synthesized images in the finetuning step.
When a model is trained with only a modest number of real image samples, it may generate bone marrow image patches that look visually appealing but are not biologically plausible. In Figure 2, we present 8 synthetic images sampled from the conditional diffusion model, which correspond to four different cell types. Each of these images achieves high precision and fidelity scores individually, but they also have biological implausibilities such as inaccurate cell coloring or nucleus shapes. Therefore, models that prioritize visual features without considering biological knowledge may miss important clinical features required for synthetic images to be useful for downstream analysis. Generic evaluation scores (e.g., [21, 22]) cannot diagnose these failures because they also lack biological domain knowledge. By incorporating feedback from pathologists, we can refine the generative model by identifying biological information that is missed by the pretrained model.
The expert pathologist examined each synthetic image and provided a feedback score on its biological plausibility. The evaluation typically involved inspecting the image and checking 7 aspects that contribute to its perceived plausibility (Table 1). These aspects pertain to the consistency of the shapes, sizes, patterns and colors of the contents of a synthetic bone marrow image  $\widetilde{x}$  with the cell type  $\widetilde{c}$ .
Clinical plausibility criteria
(1) Cell size	(5) Chromatin pattern
(2) Nucleus shape and size	(6) Inclusions
(3) Nucleus-to-cytoplasm ratio	(7) Granules
(4) Cytoplasm color and consistency
Table 1: Pathologist evaluation criteria.
Among the determinants of plausibility is the cell size—different cells have different sizes, e.g., Lymphocytes are generally smaller than Monocytes or Neutrophils. Nucleus shape and size also depend on the cell type, e.g., Band Neutrophils have a horseshoe-shaped nucleus, whereas Segmented Neutrophils have a multi-lobed nucleus. Chromatin patterns within the nucleus are dense and clumped in Lymphocytes, while in Myeloid cells they are diffuse and fine. The number, size and color of granules also contribute to plausibility. Detailed explanation of all criteria is provided in the Appendix.
Note that we have full control over the number of synthetic images  $n_{s}$ , i.e., we can sample an arbitrary number of synthetic images from the conditional diffusion model. The key limiting factor on  $n_{s}$  is the time-consuming nature of the feedback collection process. To enable scalable feedback collection, we limited our study to binary feedback, i.e., the pathologist flagged a synthetic image as "implausible" if they found a violation of any of the criteria in Table 1 upon visual inspection. We collected these binary signals and did not pursue a full checklist on all plausibility criteria for each synthetic image.
The output of Step 1 is an annotated dataset  $\mathcal{D}_s = \{(\widetilde{x}^j,\widetilde{c}^j,\widetilde{y}^j)\}_{j = 1}^{n_s}$ , where  $\widetilde{y}^j\in \{0,1\}$  is the pathologist feedback on the  $j$ -th synthetic image, where  $\widetilde{y}^j = 1$  means that the image is implausible.
2.2 Step 2: Clinical plausibility reward modeling
We conceptualize the pathologist as a "labelling function"  $\Gamma : \mathcal{X} \times \mathcal{C} \to \{0,1\}$  that maps the observed synthetic image  $\widetilde{x}$  and declared cell type (class)  $\widetilde{c}$  to a binary plausibility score. In Step 2, we model the "pathologist" by learning their labelling function  $\Gamma$  on the basis of their feedback annotations.
To train a model  $\Gamma$  that estimates the pathologist labelling function, we construct a training dataset that comprises a mixture of real and synthetic images as follows:
Synthetic images
We construct a dataset  $\mathcal{D}_s^\Gamma = \{(\widetilde{x}^j,\widetilde{c}^j,\widetilde{y}^j)\}_{j = 1}^{n_s}$  comprising the synthetic images and corresponding pathologist feedback collected in Step 1.
We build a dataset  $\mathcal{D}_r^\Gamma = \{(x^i,\widetilde{c}^i,y^i)\}_{i = 1}^{n_r}$  comprising the real images and pseudo-labels defined as  $y^{i} = 1\{c^{i}\neq \widetilde{c}^{i}\}$ , where  $\widetilde{c}^i\sim \mathrm{Uniform}(1,2,\ldots ,|\mathcal{C}|)$ .
We combine both datasets  $\mathcal{D}^{\Gamma} = \mathcal{D}_r^\Gamma \cup \mathcal{D}_s^\Gamma$  to construct a training dataset for the model  $\Gamma$ . The real dataset  $\mathcal{D}_r^\Gamma$  is built by randomly permuting the image class and assigning an implausibility label of 1 if the permuted class does not coincide with the true class. We use the real dataset to augment the synthetic dataset with the annotated pathologist feedback. By augmenting the datasets  $\mathcal{D}_r^\Gamma$  and  $\mathcal{D}_s^\Gamma$ , we teach the model  $\Gamma$  to recognize two forms of implausibility in image generation: (i) instances where the synthetic image looks clinically plausible but belongs to a wrong cell type (i.e., training examples
in  $\mathcal{D}_r^\Gamma$ ), and (ii) instances where the synthetic image is visually consistent with the correct cell type but fails to meet some of the plausibility criteria in Table 1 (i.e., subset of the training examples in  $\mathcal{D}_s^\Gamma$ ). We call the resulting model  $\Gamma$  a clinical plausibility reward model. Using the augmented feedback dataset  $\mathcal{D}^\Gamma$ , we train the reward model by minimizing the mean square error as follows:
$$
L _ {\Gamma} (\phi) = \sum_ {j \in \mathcal {D} _ {s} ^ {\Gamma}} \left(\tilde {y} ^ {j} - \Gamma_ {\phi} \left(\tilde {x} ^ {j}, \tilde {c} ^ {j}\right)\right) ^ {2} + \lambda_ {r} \sum_ {i \in \mathcal {D} _ {r} ^ {\Gamma}} \left(y ^ {i} - \Gamma_ {\phi} \left(x ^ {i}, \tilde {c} ^ {i}\right)\right) ^ {2}, \tag {2}
$$
where  $\lambda_r$  is a hyper-parameter that controls the contribution of real images in training the reward functions, and  $\phi$  is the parameter of the reward model.
2.3 Step 3: Clinically-informed finetuning
In the final step, we refine the diffusion model by leveraging the pathologist feedback. Specifically, we incorporate domain knowledge into the model by utilizing the reward model  $\Gamma$  in the finetuning objective. Following [20], we use a reward-weighted negative log-likelihood (NLL) objective, i.e.,
$$
L (\theta , \widehat {\phi}) = \mathbb {E} _ {(\widetilde {x}, \widetilde {c}) \sim \mathcal {D} _ {s}} \left[ - \Gamma_ {\widehat {\phi}} (\widetilde {x}, \widetilde {c}) \cdot \log \left(p _ {\theta} (\widetilde {x} | \widetilde {c})\right) \right] + \beta_ {r} \cdot \mathbb {E} _ {(x, c) \sim \mathcal {D} _ {r}} [ - \log \left(p _ {\theta} (x | c)\right) ], \tag {3}
$$
to finetune the conditional diffusion model, where  $\beta_{\tau}$  is a hyper-parameter and  $\widehat{\phi}$  is the reward model parameters obtained by minimizing (2) in Step 2. The finetuning objective in (3) incorporates the pathologist knowledge through the reward model, which predicts the pathologist evaluation of the synthetic images that the model generates as it updates its parameters  $\theta$ . The reward-weighted objective penalizes the generation of images that do not align with the pathologist preferences, hence we expect that the finetuned model will be less likely to generate clinically implausible synthetic images.
2.4 Bonus step: Feedback-driven generation of new clinical concepts
Besides refining generative models for clinical plausibility, pathologist feedback can be used to incorporate novel clinical concepts into the generative process that were not initially labeled in the real dataset. This could allow generative models to continuously adapt in changing clinical environments. For instance, in our bone marrow image generation setup, pathologist feedback can refine image generation by introducing new sub-types of the original cell types in  $\mathcal{C}$ , as illustrated in Figure 3.
Instead of collecting pathologist feedback that is limited to clinical plausibility, we also collect their annotation of new cell sub-types (e.g., segmented and band variants of Neutrophil cell types). Next, we train an auxiliary model  $\Gamma_{\rho}(x)$  with parameter  $\rho$  to classify the new sub-types based on the pathologist annotations. Finally, we finetune the conditional diffusion model through a combined loss function, i.e.,  $L(\theta ,\widehat{\phi}) + L(\theta ,\widehat{\rho})$ , that incorporates the two forms of pathologist feedback, i.e., annotations of new cell types and clinical plausibility.
Figure 3: Generation of refined cell sub-types.
2.5 Pathologist feedback vs. Automated feedback
To assess the added value of human feedback, we consider a baseline where the generative model is supplemented with automatically generated feedback on clinical plausibility. To this end, we implement a baseline based on classifier-guided diffusion [23], where a classifier serves as an automatic feedback signal that deems a synthetic image implausible if it does not match the corresponding cell type. For this baseline, we train an auxiliary classifier to predict the cell type  $c$  of an image  $x$ , and then we incorporate the gradient of the log-likelihood of this classifier in the training objective as described in [23] (See the Supplementary material for implementation details).
Before delving into the experimental results, we discuss three strands of literature that are related to our synthetic data generation framework. These include: generative modeling for synthetic clinical data, evaluation of generative models and learning from human feedback.
Generative modeling of synthetic medical images. The dominant approach for synthesizing medical images is to train or finetune a generative model, such as a Variational Autoencoder (VAE) [24], a Generative Adversarial Network (GAN) [25], or a diffusion model [6, 23], using a sufficiently large sample of images from the desired modality. Owing to their recent success in achieving state-of-the-art results in high-fidelity image synthesis [23], diffusion probabilistic models have become the model of choice for medical image synthesis applications [12, 26-29]. In [12], the Stable Diffusion model—an open-source pretrained diffusion model—was used to generate synthetic X-ray images, and in [28] it was shown that diffusion models can synthesize high-quality Magnetic Resonance Images (MRI) and Computed Tomography (CT) images. [27] used latent diffusion models to generate synthetic images from high-resolution 3D brain images. All of these models are trained with the standard likelihood objective and the synthetic images are typically evaluated through downstream classification tasks or generic, domain-agnostic metrics for image fidelity. To the best of our knowledge, none of the previous studies have explored a human-AI collaboration approach to synthetic image generation or incorporated clinical knowledge into generative models of medical images.
Evaluation of synthetic images in the medical domain and beyond. Unlike discriminative modeling (i.e., predictive modeling) where model accuracy can be straightforwardly evaluated by comparing the model predictions with ground-truth labels in a testing set, evaluating the quality of generative models can be quite challenging since we do not have a "ground-truth" for defining what makes a synthetic sample is of high or low quality. Devising a generic score to evaluate a generative model can be tricky since there are many potential modes of failure [22]. Consequently, it is essential to design robust multidimensional scores that capture the most relevant failure modes for a given application.