SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Marco Aversa
University of Glasgow and Dotphoton
Glasgow, United Kingdom
marco.aversa@glasgow.ac.uk
Gabriel Nobis
Fraunhofer HHI
g Gabriel.nobis@hhi.fraunhofer.de
Miriam Hagele
miriam.haegele@aignostics.com
Kai Standvoss
kai.standvoss@aIGNostics.com
Mihaela Chirica
Institute of Pathology, LMU Munich
mihaela.chirica@med.uni-muenchen.de
Roderick Murray-Smith
roderick.murray-smith@glasgow.ac.uk
Ahmed Alaa
Lukas Ruff
lukas.ruff@aignostics.com
Daniela Ivanova
daniela.ivanova@glasgow.ac.uk
Wojciech Samek
Fraunhofer HHI and TU Berlin
wojciech.samek@hhi.fraunhofer.de
Frederick Klauschen
f.klauschen@lmu.de
Bruno Sanguinetti
Dotphoton
Zug, Switzerland
bruno.sanguinetti@dotphoton.com
Luis Oala
luis.oala@dotphoton.com
We present DiffInfinite, a hierarchical diffusion model that generates arbitrarily large histological images while preserving long-range correlation structural information. Our approach first generates synthetic segmentation masks, subsequently used as conditions for the high-fidelity generative diffusion process. The proposed sampling method can be scaled up to any desired image size while only requiring
small patches for fast training. Moreover, it can be parallelized more efficiently than previous large-content generation methods while avoiding tiling artifacts. The training leverages classifier-free guidance to augment a small, sparsely annotated dataset with unlabelled data. Our method alleviates unique challenges in histopathological imaging practice: large-scale information, costly manual annotation, and protective data handling. The biological plausibility of DiffInfinite data is evaluated in a survey by ten experienced pathologists as well as a downstream classification and segmentation task. Samples from the model score strongly on anti-copying metrics which is relevant for the protection of patient data.
b)
Figure 1: a) Examples of synthetic and real  $2048 \times 2048$  images. b) Pairs of  $512 \times 512$  synthetic tiles (top) with the closest real images found with Inception-v3 near-neighbour (bottom).
Deep learning (DL) models are promising auxiliary tools for medical diagnosis [1-3]. Applications like segmentation and classification have been refined and pushed to the limit on natural images [4]. However, these models trained on rich datasets still have limited applications in medical data. While segmentation models rely on sharp object contours when applied to natural data, in medical imaging, the model struggles to detect a specific feature because it has a "limited ability to handle objects with missed boundaries" and often "miss tiny and low-contrast objects" [5, 6]. Therefore, task-specific medical applications require their own specialised and fine-grained annotation. Data labelling is arguably one of the most critical bottlenecks in healthcare machine learning (ML) applications. In histopathology, pathologists examine the histological slide at multiple levels, usually starting with a lower magnification to analyse the tissue architecture and cellular arrangement and gradually proceeding to a higher magnification to examine cell morphology and subcellular features, such as the appearance and number of nucleoli, chromatin density and cytoplasm appearance. Annotating features within gigapixel whole slide images (WSIs) with this level of detail demands effort and time, often leading to sparse, limited annotated data. In addition, due to privacy regulations and ethics [7, 8], having access to medical data can be challenging since it has been shown that it is possible to extract patients' sensitive information [9] from this data.
In histopathology, state-of-the-art ML models require the context of the entire WSIs, with features at different scales, in order to distinguish between different tumor sub-types, grades and stages [10]. Despite the demonstrated effectiveness of diffusion models (DMs) in generating natural images compared to other approaches, they still have rarely been applied in medical imaging. Existing
generative models in histopathology can generate images of relatively small resolution compared to WSIs. To give a few examples, the application of Generative Adversarial Networks (GANs) in cervical dysplasia detection [11], glioma classification [12], and generating images of breast and colorectal cancer [13], generate images with  $256 \times 128$  px,  $384 \times 384$  px and  $224 \times 224$  px, respectively. In spite of their current limitations in generating images at scales necessary to fully address all medical concerns, the use of synthetic data in medical imaging can provide a valuable solution to the persistent issue of data scarcity [14-17]. Models generally improve after data augmentation and synthetic images are equally informative as real images when added to the training set [18, 19]. Data augmentation could also help with the underrepresentation in data sets of rare cancer subtypes. By adding synthetic images to the training set, Chen et al. [20] demonstrated that their model had better accuracy in detecting chromophobe renal cell carcinoma, which is a rare subtype of renal cell carcinoma. Furthermore, Doleful et al. [21] showed how synthetic histological images could be used for educational purposes for pathology residents. Regarding the challenges highlighted before, we present a novel sampling method to generate large histological images with long-range pixel correlation (see Fig. 1), aiming to extend up to the resolution of the WSI.
Our contributions are as follows: 1) We introduce DiffInfinite, a hierarchical generative framework that generates arbitrarily large images, paired with their segmentation masks. 2) We introduce a fast outpainting method that can be efficiently parallelized. 3) The quality of DiffInfinite data is evaluated by ten experienced pathologists as well as downstream machine learnings tasks (classification and segmentation) and anti-duplication metrics to assess the leakage of patient data from the training set.
Large-content image generation can be reduced to inpainting/outpainting tasks. Image inpainting is the problem of reconstructing unknown or unwanted areas within an image. A closely related task is image outpainting, which aims to predict visual content beyond the boundaries of an image. In both cases, the newly in- or outpainted image regions have to be visually indistinguishable with respect to the rest of the image. Such image completion approaches can help utilise models trained on smaller patches for the purpose of generating large images, by initially generating the first patch, followed by its extension outward in the desired direction.
Traditional approaches Traditional methods for image region completion rely on repurposing known image features, necessitating costly nearest neighbour searches for suitable pixels or patches [22-26]. Such methods often falter with complex or large regions [24]. In contrast, DL enables novel, realistic image synthesis for inpainting and outpainting. Some methods like Deep Image Prior [27] condition new image areas on the existing image, while others aim to learn natural image priors for realistic generation [28, 29].
Generative modelling for conditional image synthesis GANs have dominated image-to-image translation tasks like inpainting and outpainting for years [28-42]. Recently, DMs have surpassed GANs in various image generation tasks [43]. Palette [44] was the first to apply DMs to tasks like inpainting and outpainting. RePaint [45] and ControlNet [46] demonstrate resampling and masking techniques for conditioning using a pre-trained diffusion model. SinDiffusion [47] and DiffCollage [48] offer state-of-the-art outpainting solutions using DMs trained with overlapping patches. In parallel to our work, Bond-Taylor and Willcocks [49] developed a related approach called  $\infty$ -Diff which trains on random coordinates, allowing the generation of infinite-resolution images during sampling. However, in contrast to our approach the method does not involve image compression in a latent space.
Synthetic data assessment The authenticity of synthetic data produced by DMs, trained on vast paired labelled datasets [50], remains contentious. Ethical implications necessitate distinguishing if generated images are replicas of training data [51, 52]. The task is complicated due to subjective visual similarities and diverse dataset ambiguities. Various metrics have been proposed for quantifying data replication, including information theory distances from real data [53], consistency measurements using downstream models [54, 55], comparison with inpainted areas [52], and detection of "forgotten" examples [56].
Figure 2: DiffInfinite generation method. a) Large-scale context mask generation. A diffusion model conditioned on a large-scale conditional prompt (e.g. Adenocarcinoma subtype) generates a low-resolution mask. The mask is upsampled via linear interpolation to the desired image size. b) Diffusion steps on large images. Given a random position, we select a sub-tile with its segmentation mask. A diffusion model generates in parallel the next step conditioned on each conditional label, or prompt, found in the mask. The outputs are masked individually with the corresponding label. The next step is the union of all the sub-patches. c) Tracking time steps pixel-wise. We keep track of the time step of each pixel in the large image. The model evolves only the pixels with the higher time step on each iteration.
Diffusion Models DMs [57-59] represent a class of parameterized Markov chains that effectively optimize the lower variational bound associated with the likelihood function of the unknown data distribution. By iteratively adding small amounts of noise until the image signal is destroyed and then learning to reverse this process, DMs can approximate complex distributions much more faithfully than GANs [60]. The increased diversity of samples while preserving sample fidelity comes at the cost of training and sampling speed, with DMs being much slower than GANs [43]. The universally adopted solution to this problem is to encode the images from pixel space into a lower dimensional latent space via a Vector Quantised-Variational AutoEncoder (VQ-VAE), and perform the diffusion process over the latents, before decoding back to pixel space [61]. Pairing this with the Denoising Diffusion Implicit Models (DDIMs) sampling method [62] leads to faster sampling while preserving the DM objective
$$
z _ {t - 1} = \sqrt {\alpha_ {t - 1}} \left(\frac {z _ {t} - \sqrt {1 - \alpha_ {t}} \epsilon_ {\theta} (z _ {t} , t)}{\sqrt {\alpha_ {t}}}\right) + \sqrt {1 - \alpha_ {t - 1} - \sigma_ {t} ^ {2}} \epsilon_ {\theta} (z _ {t}, t) + \sigma_ {t} \epsilon_ {t}, \tag {1}
$$
where  $z_{t}$  is the latent variable at time step  $t$  in the VQ-VAE latent space,  $\alpha_{t}$  is the noise scheduler,  $\epsilon_{\theta}$  is the noise learned by the model and  $\epsilon_{t}$  is random noise. Conditioning can be achieved either by specifically feeding the condition with the noised data [44, 63], by guiding an unconditional model using an external classifier [64, 65] or by classifier-free guidance [66] used in this work, where the
convex combination
$$
\tilde {\epsilon} _ {\theta} \left(z _ {t}, c\right) = (1 + \omega) \epsilon_ {\theta} \left(z _ {t}, c\right) - \omega \epsilon_ {\theta} \left(z _ {t}, \varnothing\right), \tag {2}
$$
of a conditional diffusion model  $\epsilon_{\theta}(z_t, c)$  and an unconditional model  $\epsilon_{\theta}(z_t, \emptyset)$  is used for noise estimation. The parameter  $\omega$  controls the tradeoff between conditioning and diversity, since  $\omega > 0$  introduces more diversity in the generated data by considering the unconditional model while  $\omega = 0$  uses only the conditional model.
4 Infinite Diffusion
The DiffInfinite approach we present here<sup>1</sup>, is a generative algorithm to generate arbitrarily large images without imposing conditional independence, allowing for long-range correlation structural information. The method overcomes this limitation of DMs for large-content generation by deploying multiple realizations of a DM on smaller patches. In this section, we first define a mathematical description of this hierarchical generation model and then describe the sampling method paired with a masked conditioned generation process.
4.1 The Method
Let  $X \sim \mathcal{X}$  be a large-content generating random variable taking values in  $\mathbb{R}^{KD}$ . Using the approach of latent diffusion models [61], the high-dimensional content is first mapped to the latent space  $\mathbb{R}^D$  by  $\Phi(X) = Y \sim \mathcal{Y}_{\Phi}$ . For simplicity, we assume throughout this work the existence of an ideal encoder-decoder pair  $(\Phi, \Psi)$  such that  $\Psi(\Phi(X)) = X$  is the identity on  $\mathbb{R}^{KD}$ . Assume further, to have a reverse time model  $(SM_{\theta}, \epsilon_{\theta})$  at hand consisting of a sampling method  $SM_{\theta}$  and a learned model  $\epsilon_{\theta}$  trained on small patches  $Z \sim \mathcal{Z}_{\Phi}$  taking values in  $\mathbb{R}^d$ . The reverse time model transforms  $z_T \sim \mathcal{N}(0, I_d)$  over the time steps  $t \in \{T, T-1, \dots, 1\}$  recursively by
$$
z _ {t - 1} = S M _ {\theta} \left(z _ {t}\right) \tag {3}
$$
to an approximate instance of  $\mathcal{Z}_{\Phi}$ . We aim to sample instances from  $\mathcal{V}_{\Phi}$  by deploying multiple realizations of the reverse time model  $(SM_{\theta},\epsilon_{\theta})$ . Towards that goal, define the set of projections
$$
\mathcal {C} := \left\{\operatorname {p r o j} _ {I}: \mathbb {R} ^ {D} \rightarrow \mathbb {R} ^ {d} \mid I \subset \mathbb {N} \text {c o r r e s p o n d t o} d \text {i n d i c e s o f c o n n e c t e d p i x e l s i n} \mathbb {R} ^ {D} \right\}, \tag {4}
$$
where  $proj \in \mathcal{C}$  models a crop  $proj( Y ) \in \mathbb { R } ^ { d }$  of  $d$  connected pixels from the latent image  $Y$ . Since the model  $\epsilon_{\theta}$  is trained on images taking values in  $\mathbb{R}^d$  the standing assumption is
Assumption 1 Any projection  $\operatorname{proj} \in \mathcal{C}$  maps  $Y$  to the same distribution  $\operatorname{proj}(Y) \sim \mathcal{Z}_{\Phi}$  in  $\mathbb{R}^d$ .
Since the goal is to approximate an instance of  $\mathcal{V}_{\Phi}$ , we initialize the sampling method by  $y_{T} \sim \mathcal{N}(0, I_{D})$  and proceed in the following way: Given  $y_{t}$ , randomly choose  $proj_{I_{1}}, \dots, proj_{I_{m}} \in \mathcal{C}$  independent of the state  $y_{t}$  such that  $proj_{I_{1}}, \dots, proj_{I_{m}}$  are non equal crops that cover all latent pixels in  $\mathbb{R}^{D}$ . To be more precise, for every  $i \in \{1, \dots, D\}$  we find at least one  $j \in \{1, \dots, m\}$  with  $i \in I_{j}$ . For every projection  $proj_{I_{1}}, \dots, proj_{I_{m}}$  we calculate the crop  $z_{t}^{j} = proj_{I_{j}}(y_{t})$  of the current state  $y_{t}$  and perform one step of the reverse time model following the sampling scheme
$$
z _ {t - 1} ^ {j} = S M _ {\theta} \left(z _ {t} ^ {j}\right), \quad j \in \{1, \dots , m \}. \tag {5}
$$
This results in overlapping estimates  $z_{t-1}^{1}, \ldots, z_{t-1}^{m}$  of the subsequent state  $t-1$  and we simply assign to every pixel in the latent space the first value computed for this pixel such that
$$
\left[ y _ {t - 1} \right] _ {i} = \left[ z _ {t - 1} ^ {j} \right] _ {l}, \quad \text {w h e r e} j = \min  \left\{j ^ {\prime} \mid i \in I _ {j ^ {\prime}} \right\} \tag {6}
$$
and  $l$  refers to the entry in  $z_{t - 1}^{j}$  corresponding to  $i$  with  $[proj_{I_j}(y_{t - 1})]_l = [y_{t - 1}]_i$ . Hence, starting from  $y_{T}\sim \mathcal{N}(0,I_{D})$  we sample in the first step from a distribution
$$
y _ {T - 1} \sim p _ {T - 1, \theta} (y \mid y _ {T}, p r o j _ {I _ {1}}, \dots , p r o j _ {I _ {m}}). \tag {7}
$$
Using Bayes' theorem, this distribution simplifies to
$$
p _ {T - 1, \theta} (y \mid y _ {T}, p r o j _ {I _ {1}}, \dots , p r o j _ {I _ {m}}) = p _ {T - 1, \theta} (y \mid y _ {T}), \tag {8}
$$
$^{1}$ Code available at https://github.com/marcoaversa/diffinfinite
since we sample the projections independently from  $y_{T}$ . Repeating the argument, we sample in every step from a distribution  $y_{t - 1} \sim p_{t - 1,\theta}(y|y_t, \dots, y_T)$  over  $\mathbb{R}^D$  instead of sampling from  $z_{t - 1} \sim q_{t - 1,\theta}(z|z_t, \dots, z_T)$  over  $\mathbb{R}^d$ . Hence, we approximate the true latent distribution  $\mathcal{V}_{\Phi}$  by the approximate distribution with density  $p_{0,\theta}(y|y_1, \dots, y_T)$ . In contrast to [48], our method does not use the assumption of conditional independence and the method can be applied to a wide range of DMs, without an adjustment of the training method. As the authors of [48] point out in their section on limitations, the assumption of conditional independence is not well-suited in cases of a data distribution with long-range dependence. For image generation in the medical context, we aim to circumvent this assumption as we do not want to claim that the density of a given region depends only on one neighboring region. The drawback of dropping the assumption is that we only approximate the reverse time model of the latent image distribution  $\mathcal{V}_{\Phi}$  indirectly, by multiple realizations of a reverse time model that approximates  $\mathcal{Z}_{\Phi}$ .
4.2 Semi-supervised Guidance
In order to generate diverse high-fidelity data, DMs require lots of training data. Perhaps, training on a few samples still extracts significant features but it lacks variability, resulting in simple replicas. Here, we show how to enhance synthetic data diversity using classifier-free guidance as a semi-supervised learning method. In the classifier-free guidance [66], a single model is trained conditionally and unconditionally on the same dataset. We adapt the training scheme using two separate datasets. The model is guided by a small and sparsely annotated dataset  $q_{1}$ , used for the conditional training step, while extracts features by the large unlabelled dataset  $q_{0}$ , used on the unconditional training step (see Alg.1)
$$
\left(z _ {0}, c\right) = \left\{ \begin{array}{l l} \left(z _ {0}, \varnothing\right) \sim q _ {0} \left(z _ {0}\right) & \text {i f} u \geqslant p _ {u n c} \\ \left(z _ {0}, c\right) \sim q _ {1} \left(z _ {0}, c\right) & \text {o t h e r w i s e} \end{array} , \right. \tag {9}
$$
where  $u$  is sampled from a uniform distribution in [0,1],  $p_{unc}$  is the probability of switching from the conditional to the unconditional setting and  $\varnothing$  is a null label. During the sampling, a tradeoff between conditioning and diversity is controlled via the parameter  $\omega$  in eq.2.
4.3 Sampling
High-level content generation The outputs of DMs have pixel consistency within the training image size. Outpainting an area with a generative model might lead to unrealistic and odd artifacts due to poor long-range spatial correlations. Here, we show how to predict pixels beyond the image's boundaries by generating a hierarchical mapping of the data. The starting point is the generation of the highest-level representation of the data. In our case, it is the sketch of the cellular arrangement in the WSI (see Figure 2a). Since higher-frequency details are unnecessary at this stage, we can downsample the masks until the clustering pattern is still recognizable. The diffusion model, conditioned on the context prompt (e.g. Adenocarcinoma subtype), learns the segmentation masks which contain the cellular macro-structures information.
Figure 3: Comparison of sampling speed for DiffCollage and DiffInfinite, measuring diffusion steps required for image sampling. Demonstrating increased efficiency of DiffInfinite for larger images.
Random patch diffusion Given a segmentation mask  $M$ , we can proceed with the large image sampling according to Section 4.1 in the latent space  $\mathbb{R}^D$  of  $Y = \Phi(X)$  (see Alg.2). Since we trained a conditional diffusion model with conditions  $c_1, \ldots, c_N$ , the learned model takes the form  $\epsilon_{\theta}(z_t, t) = (\epsilon_{\theta}(z_t, t|c_1), \dots, \epsilon_{\theta}(z_t, t|c_N))$ . Given  $y_t$ , we first sample projections  $proj_{I_1}, \dots, proj_{I_m} \in \mathcal{C}$ , corresponding to different crops of  $d$  connected pixels up to the  $m$ -th projection with  $\cup_{j=1}^{m} I_j = \{1, \dots, D\}$  and  $\cup_{j=1}^{m} I_j \setminus \cup_{j=1}^{m-1} I_j \neq \emptyset$  (see the left-hand-side of Figure 2b). Note that  $m$  is not fixed, but varies over the sampling steps and is upper bounded by the number of possible crops of  $d$  connected pixels. The random selection of
the projection is implemented such that regions with latent pixels of low projection coverage are more likely. Secondly, we calculate for every projection  $j \in \{1, \dots, m\}$  the crop  $proj_{I_j}(y_t)$  and perform one step of the DDIM sampling procedure using the classifier-free guidance model  $(1 + \omega)\epsilon_\theta(z_t^j, t, c) - \omega \epsilon_\theta(z_t^j, t, \emptyset)$ , where  $\epsilon_\theta$  is the learned model and  $z_t^j = proj_{I_j}(y_t)$ . This results for every pixels  $i \in I_j$  in  $N$  values  $DDIM_{\theta, c_1}(proj_{I_j}(y_t)), \dots, DDIM_{\theta, c_N}(proj_{I_j}(y_t))$ , one for every condition  $c_i$  (see the right hand-side of Figure 2b). If  $i \notin I_{j'}$  for all  $j'$ , the pixel  $i$  has not been considered yet and we assign  $i$  the value  $[y_{t-1}]_i = [DDIM_{\theta, M_i}(proj_{I_j}(y_t))]_l$ , where  $l$  corresponds to the pixel  $i$  under the projection  $I_j$  and  $M_i$  is the value of  $i$  in the mask  $M$ . Since we are updating random projections of the overall image, in the  $t$ -th step pixels either have the time index  $t$  or  $t+1$ , resulting in a reversed diffusion process of differing time states. We initialize a tensor  $L_t$ , with the same size  $D$  as the latent variable, to keep track of the time index for each pixel. Each element is set to  $L_T \equiv T$ . In the  $j$ -th iteration of the  $t$ -th step we only update the pixels that have not been considered in one of the previous iterations of the  $t$ -th diffusion step, hence all the pixels in  $i \in I_j$  with  $proj_{I_j}(L_t)_i = t+1$ , similarly to the inpainting mask in the Repaint sampling method [45].
To restore the pixels that already received an update, i.e. every pixel  $i \in I_j$  with  $proj_{I_j}(L_t)_i = t$ , we store a replica of the previous diffusion step for every pixel. Finally, we update all the time states in  $L_t$  that received an update in the  $j$ -th iteration to  $t$  resulting in  $proj_{I_j}(L_t)_i = t$  for all  $i \in I_j$ . See the top row of Fig.2c for an illustration of the evolution of  $L_t$ . The random patch diffusion can also be applied to mask generation, where the only condition is the context prompt. This method can generate segmentation masks of arbitrary sizes with the correlation length bounded by two times the training mask image size.
Parallelization The sampling method proposed has several advantages. In Zhang et al. [48] each sequential patch is outpainted from the previous one with  $50\%$  of the pixels shared. Here, the randomization eventually leads to every possible overlap with the neighboring patches. This introduces a longer pixel correlation across the whole generated image, avoiding artifacts due to tiling. In Figure 3, we show that the number of steps in the whole large image generation process is drastically reduced with the random patching method with respect to the sliding window one. Moreover, in the sliding sampling method, the model can be paralleled only 2 or 4 times, depending if we are outpainting the image horizontally or on both axis. In our approach, we can parallelize the sampling up to the computational resource limit.
Algorithm 1 DiffInfinite Training
1: Randomly train on labelled or unlabelled data with probability  $p_{unc}$
2: Sample random time step