Discrepancy-based Diffusion Models for Lesion Detection in Brain MRI

TL;DR

This work tackles brain MRI lesion detection under limited pixel-level annotations by introducing Discrepancy Distribution Medical Diffusion (DDMD). DDMD combines discrepancy feature generation from image-level labels with a diffusion process that conditions on a concatenated prior to produce pixel-level lesion masks, optimizing the simple loss $L_{\text{simple}}=E_{t,\mathbf{x}_0,\boldsymbol{\epsilon}}\left[\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_{\theta}(\mathbf{X}_t,t)\right\|^2\right]$ and using a diffusion horizon of $T=1000$ steps. The method leverages two auto-encoder modules to obtain inter- and intra-discrepancy features $(\mathcal{X},\mathcal{Y})$, which are fused with the input image as priors $\mathbf{X}=\mathbf{b}\oplus\mathcal{X}\oplus\mathcal{Y}\oplus\mathbf{x}_b$ for segmentation. On BRATS2020, DDMD variants—especially DDMD-light—achieve superior Dice, Miou, and PA scores compared with Unet, SegNet, DeepLabv3+, and CIMD, demonstrating improved lesion localization and robustness across multi-modal MRI data.

Abstract

Diffusion probabilistic models (DPMs) have exhibited significant effectiveness in computer vision tasks, particularly in image generation. However, their notable performance heavily relies on labelled datasets, which limits their application in medical images due to the associated high-cost annotations. Current DPM-related methods for lesion detection in medical imaging, which can be categorized into two distinct approaches, primarily rely on image-level annotations. The first approach, based on anomaly detection, involves learning reference healthy brain representations and identifying anomalies based on the difference in inference results. In contrast, the second approach, resembling a segmentation task, employs only the original brain multi-modalities as prior information for generating pixel-level annotations. In this paper, our proposed model - discrepancy distribution medical diffusion (DDMD) - for lesion detection in brain MRI introduces a novel framework by incorporating distinctive discrepancy features, deviating from the conventional direct reliance on image-level annotations or the original brain modalities. In our method, the inconsistency in image-level annotations is translated into distribution discrepancies among heterogeneous samples while preserving information within homogeneous samples. This property retains pixel-wise uncertainty and facilitates an implicit ensemble of segmentation, ultimately enhancing the overall detection performance. Thorough experiments conducted on the BRATS2020 benchmark dataset containing multimodal MRI scans for brain tumour detection demonstrate the great performance of our approach in comparison to state-of-the-art methods.

Figures (5)

• Figure 1: The framework of the proposed DDMD method. It consists of two key processes: a) the discrepancy feature generation process and b) the diffusion process.
• Figure 2: Forward and reverse diffusion processes in DPM.
• Figure 3: Qualitative comparison between different brain lesion detection methods. Three examples are showcased here. For each example, the first row illustrates the input brain images at four different modalities (i.e., T1, T1ce, T2 and FLAIR) along with the ground-truth segmentation mask, and the second row displays the segmentation mask achieved by each method.
• Figure 4: Histograms of the inter-discrepancy and intra-discrepancy scores obtained from the normal and abnormal images in the test dataset for different data modalities. The first and second rows show the results across the individual data modalities (i.e., T1, T1ce, T2 and FlAIR), and the third row displays the results under the whole data modalities. Scores are normalized to [0, 1].
• Figure 5: Diffusion results of different sampling methods at different time steps. The first row shows the input MRI images with T1, T1ce, T2 and FLair modalities and the corresponding ground-truth segmentation mask. The second to fifth rows present intermediate results of the CIMD, DDMD-mini, DDMD, and DDMD-light sampling processes, respectively.