NNDM: NN_UNet Diffusion Model for Brain Tumor Segmentation

TL;DR

This work tackles brain tumor segmentation in MRI by addressing generalization and boundary precision. It introduces NNDM, a hybrid framework that uses an NN_UNet backbone for initial segmentation and a diffusion-based residual refinement module to model and denoise the residual error $e = y - \hat{y}$ through a denoising process with schedule $\{\beta_t\}$. The method achieves state-of-the-art Dice and boundary metrics on the BraTS 2021 dataset, with improved robustness across modalities and tumor subregions, demonstrating the value of coupling deterministic segmentation with stochastic diffusion. The approach advances automated brain tumor analysis by combining the strengths of deterministic and generative models to produce more accurate and reliable segmentations in clinical settings.

Abstract

Accurate detection and segmentation of brain tumors in magnetic resonance imaging (MRI) are critical for effective diagnosis and treatment planning. Despite advances in convolutional neural networks (CNNs) such as U-Net, existing models often struggle with generalization, boundary precision, and limited data diversity. To address these challenges, we propose NNDM (NN\_UNet Diffusion Model)a hybrid framework that integrates the robust feature extraction of NN-UNet with the generative capabilities of diffusion probabilistic models. In our approach, the diffusion model progressively refines the segmentation masks generated by NN-UNet by learning the residual error distribution between predicted and ground-truth masks. This iterative denoising process enables the model to correct fine structural inconsistencies and enhance tumor boundary delineation. Experiments conducted on the BraTS 2021 datasets demonstrate that NNDM achieves superior performance compared to conventional U-Net and transformer-based baselines, yielding improvements in Dice coefficient and Hausdorff distance metrics. Moreover, the diffusion-guided refinement enhances robustness across modalities and tumor subregions. The proposed NNDM establishes a new direction for combining deterministic segmentation networks with stochastic diffusion models, advancing the state of the art in automated brain tumor analysis.