BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation

Let us assume that $\bm{x}\in{\mathbb{R}^{H\!\times\!W\!\times\!C}}$ denotes the input medical image with a spatial resolution of $H\!\times\!W$ and $C$ number of channels. The ground-truth mask is represented as $\bm{y}_{0}\!\in\!{{\{0,1\}}^{H\times W}}$, where $0$ represents background while $1$ ROI. Inspired by diffusion-based models such as denoising diffusion probabilistic model (DDPM) and DDIM, we propose a novel conditional Bernoulli diffusion model, which can be represented as $p_{\theta}({\bm{y}_{0}}|\bm{x}):=\int p_{\theta}(\bm{y}_{0:T}|\bm{x})\mathrm{d}\bm{y}_{1:T}$, where $\bm{y}_{1},\ldots,\bm{y}_{T}$ are latent variables of the same size as the mask $\bm{y}_{0}$. For medical binary segmentation tasks, the diverse reverse process of our BerDiff starts from the initial Bernoulli noise $\bm{y}_{T}\sim\mathcal{B}(\bm{y}_{T};\frac{1}{2}\!\cdot\!{\mathbf{1}})$ and progresses through intermediate latent variables constrained by the input medical image $\bm{x}$ to produce segmentation masks, where $\mathbf{1}$ denotes an all-ones matrix of the size $H\!\times\!W$.

In previous generation-related diffusion models, Gaussian noise is progressively added with increasing timestep $t$. However, for segmentation tasks, the ground-truth masks are represented by discrete values. To address this, our BerDiff gradually adds more Bernoulli noise using a noise schedule $\beta_{1},\ldots,\beta_{T}$, as shown in Fig. 1. The Bernoulli forward process $q(\bm{y}_{1:T}|\bm{y}_{0})$ of our BerDiff is a Markov chain, which can be represented as:

where $\mathcal{B}$ denotes the Bernoulli distribution with the probability parameters $(1-\beta_{t})\bm{y}_{t-1}+\beta_{t}/2$. Using the notation $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha}_{t}={\textstyle\prod_{\tau=1}^{t}}{\alpha}_{\tau}$, we can efficiently sample $\bm{y}_{t}$ at an arbitrary timestep $t$ as follows:

To ensure that the objective function described in Sec. 2.4 is tractable and easy to compute, we use the sampled Bernoulli noise $\bm{\epsilon}\!\sim\!\mathcal{B}(\bm{\epsilon};\frac{1-\bar{\alpha}_{t}}{2}\!\cdot\!{\mathbf{1}})$ to reparameterize $\bm{y}_{t}$ of Eq. (3) as $\bm{y}_{0}\oplus\bm{\epsilon}$, where $\oplus$ denotes the logical operation of “exclusive or (XOR)”. Additionally, let $\odot$ denote elementwise product, and $\operatorname{Norm}(\cdot)$ denote normalizing the input data along the channel dimension and then returning the second channel. The concrete Bernoulli posterior can be represented as:

where $\theta_{\text{post}}\left(\bm{y}_{t},\bm{y}_{0}\right)=\operatorname{Norm}([\alpha_{t}[1-\bm{y}_{t},\bm{y}_{t}]+\frac{1-\alpha_{t}}{2}]\odot[\bar{\alpha}_{t-1}[1-\bm{y}_{0},\bm{y}_{0}]+\frac{1-\bar{\alpha}_{t-1}}{2}])$.

The diverse reverse process $p_{\theta}(\bm{y}_{0:T})$ can be also viewed as a Markov chain that starts from the Bernoulli noise $\bm{y}_{T}\sim\mathcal{B}(\bm{y}_{T};\frac{1}{2}\cdot{\mathbf{1}})$ and progresses through intermediate latent variables constrained by the input medical image $\bm{x}$ to produce diverse segmentation masks, as shown in Fig. 1. The concrete diverse reverse process of our BerDiff can be represented as:

Specifically, we utilize the estimated Bernoulli noise $\hat{\bm{\epsilon}}(\bm{y}_{t},t,\bm{x})$ of $\bm{y}_{t}$ to parameterize $\hat{\bm{\mu}}(\bm{y}_{t},t,\bm{x})$ via a calibration function $\mathcal{F}_{C}$, as follows:

where $|\cdot|$ denotes the absolute value operation.

Here, we provide an overview of the training and sampling procedure in Algorithms 1 and 2. During the training phase, given an image and mask data pair $\{\bm{x},\bm{y}_{0}\}$, we sample a random timestep $t$ from a uniform distribution $\{1,\dots,T\}$, which is used to sample the Bernoulli noise $\bm{\epsilon}$.

We then use $\bm{\epsilon}$ to sample $\bm{y}_{t}$ from $q(\bm{y}_{t}\!\mid\!\bm{y}_{0})$, which allows us to obtain the Bernoulli posterior $q(\bm{y}_{t-1}\!\mid\!\bm{y}_{t},\bm{y}_{0})$. We pass the estimated Bernoulli noise $\bm{\hat{\epsilon}}(\bm{y}_{t},t,\bm{x})$ through the calibration function $\mathcal{F}_{C}$ to parameterize $p_{\theta}(\bm{y}_{t-1}\!\mid\!\bm{y}_{t},\bm{x})$. Based on the variational upper bound on the negative log-likelihood in previous diffusion models [3], we adopt Kullback-Leibler (KL) divergence and binary cross-entropy (BCE) loss to optimize our BerDiff as follows:

Finally, the overall objective function is presented as:

where $\lambda_{\text{BCE}}$ is set to $1$ in our experiments.

During the sampling phase, our BerDiff first samples the initial latent variable $\bm{y}_{T}$, followed by iterative calculation of the probability parameters of $\bm{y}_{t-1}$ for different $t$. In Algorithm 2, we present two different sampling strategies from DDPM and DDIM for the latent variable $\bm{y}_{t-1}$. Finally, our BerDiff is capable of producing diverse segmentation masks. By taking the mean of these masks, we can further obtain a saliency segmentation mask to highlight salient ROI that can serve as a valuable reference for radiologists. Note that our BerDiff proposes a novel parameterization technology, $\textit{i}.\textit{e}.$ calibration function, to estimate the Bernoulli noise of $\bm{y}_{t}$, which is different from previous discrete state diffusion-based models [3, 11, 22].

We start by conducting ablation experiments to demonstrate the effectiveness of different losses and estimation targets, as shown in Table 1. All experiments are trained for 21,000 training iterations on LIDC-IDRI. We first conduct the ablation study of different losses while estimating Bernoulli noise in the top three rows. We find that the combination of KL divergence and BCE loss can achieve the best performance. Then, we conduct an ablation study of selecting estimation target in the bottom two rows. We observe that estimating Bernoulli noise, instead of directly estimating the ground-truth mask, is more suitable for our binary segmentation task. All of these findings are consistent with previous works [3, 10]. Please refer to Appendix 0.A for extra ablation studies on the sampling strategy and sampled timesteps.

Here, we conduct ablation experiments on our BerDiff with Gaussian or Bernoulli noise, and the results are shown in Table 2. For discrete segmentation tasks, we find that using Bernoulli noise can produce favorable results when training iterations are limited (e.g. 21,000 iterations), and even outperform using Gaussian noise when training iterations are sufficient (e.g. 86,500 iterations). We also provide a more detailed performance comparison between Bernoulli- and Gaussian-based diffusion models over training iterations in Appendix 0.B. In addition, we present a toy example to demonstrate the superiority of Bernoulli diffusion over Gaussian diffusion in Appendix 0.C.