Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Variance-Penalized MC-Dropout as a Learned Smoothing Prior for Brain Tumour Segmentation

Satyaki Roy Chowdhury, Golrokh Mirzaei

TL;DR

Brain tumor segmentation from MRI often suffers from jagged boundaries that hinder clinical utility. The authors present UAMSA-UNet, a Bayesian UNet variant that leverages Monte Carlo dropout to learn a data-driven smoothing prior and a Multi-Scale Attention Module to fuse local and global features. A smoothing-regularized loss adds a variance term across $T$ stochastic forward passes, acting as a learned spatial prior to suppress spurious fluctuations. On BraTS2023 and BraTS2024, the method achieves Dice gains up to 4.5% and IoU gains up to 4.0%, while also reducing FLOPs relative to strong baselines, demonstrating improved accuracy and efficiency with a flexible path toward transformer-based integration.

Abstract

Brain tumor segmentation is essential for diagnosis and treatment planning, yet many CNN and U-Net based approaches produce noisy boundaries in regions of tumor infiltration. We introduce UAMSA-UNet, an Uncertainty-Aware Multi-Scale Attention-based Bayesian U-Net that in- stead leverages Monte Carlo Dropout to learn a data-driven smoothing prior over its predictions, while fusing multi-scale features and attention maps to capture both fine details and global context. Our smoothing-regularized loss augments binary cross-entropy with a variance penalty across stochas- tic forward passes, discouraging spurious fluctuations and yielding spatially coherent masks. On BraTS2023, UAMSA- UNet improves Dice Similarity Coefficient by up to 3.3% and mean IoU by up to 2.7% over U-Net; on BraTS2024, it delivers up to 4.5% Dice and 4.0% IoU gains over the best baseline. Remarkably, it also reduces FLOPs by 42.5% rel- ative to U-Net++ while maintaining higher accuracy. These results demonstrate that, by combining multi-scale attention with a learned smoothing prior, UAMSA-UNet achieves both better segmentation quality and computational efficiency, and provides a flexible foundation for future integration with transformer-based modules for further enhanced segmenta- tion results.

Variance-Penalized MC-Dropout as a Learned Smoothing Prior for Brain Tumour Segmentation

TL;DR

Brain tumor segmentation from MRI often suffers from jagged boundaries that hinder clinical utility. The authors present UAMSA-UNet, a Bayesian UNet variant that leverages Monte Carlo dropout to learn a data-driven smoothing prior and a Multi-Scale Attention Module to fuse local and global features. A smoothing-regularized loss adds a variance term across stochastic forward passes, acting as a learned spatial prior to suppress spurious fluctuations. On BraTS2023 and BraTS2024, the method achieves Dice gains up to 4.5% and IoU gains up to 4.0%, while also reducing FLOPs relative to strong baselines, demonstrating improved accuracy and efficiency with a flexible path toward transformer-based integration.

Abstract

Brain tumor segmentation is essential for diagnosis and treatment planning, yet many CNN and U-Net based approaches produce noisy boundaries in regions of tumor infiltration. We introduce UAMSA-UNet, an Uncertainty-Aware Multi-Scale Attention-based Bayesian U-Net that in- stead leverages Monte Carlo Dropout to learn a data-driven smoothing prior over its predictions, while fusing multi-scale features and attention maps to capture both fine details and global context. Our smoothing-regularized loss augments binary cross-entropy with a variance penalty across stochas- tic forward passes, discouraging spurious fluctuations and yielding spatially coherent masks. On BraTS2023, UAMSA- UNet improves Dice Similarity Coefficient by up to 3.3% and mean IoU by up to 2.7% over U-Net; on BraTS2024, it delivers up to 4.5% Dice and 4.0% IoU gains over the best baseline. Remarkably, it also reduces FLOPs by 42.5% rel- ative to U-Net++ while maintaining higher accuracy. These results demonstrate that, by combining multi-scale attention with a learned smoothing prior, UAMSA-UNet achieves both better segmentation quality and computational efficiency, and provides a flexible foundation for future integration with transformer-based modules for further enhanced segmenta- tion results.
Paper Structure (10 sections, 9 equations, 3 figures, 1 table)

This paper contains 10 sections, 9 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Methodology
  3. MULTI SCALE ATTENTION MODULE
  4. MONTE CARLO DROPOUT AS AN APPROXIMATE BAYESIAN INFERENCE
  5. SMOOTHING REGULARIZED LOSS
  6. EXPERIMENTS AND RESULTS
  7. DATASET AND IMPLEMENTATION DETAILS
  8. RESULTS AND DISCUSSIONS
  9. Conclusion
  10. Compliance with Ethical Standards

Figures (3)

  • Figure 1: Overview of the proposed Bayesian U-Net with Multi-Scale Attention Module (MSAM). The model follows a UNet style encoder-decoder architecture, with convolutional blocks and Monte Carlo dropout for Bayesian uncertainty estimation. The encoder integrates MSAM at each layer to enhance feature refinement by emphasizing on the most relevant regions while preserving spatial information and passing the information to the decoder through skip connections.
  • Figure 2: Qualitative comparison on BraTS2023 (row one and two) and BraTS2024 (row three and row four) dataset for T2W images. The brighter areas represent region where the model is less confident.
  • Figure 3: Histograms and boxplots of Dice (blue) and IoU (green) for T2W test images (BraTS2024). High central tendency indicates strong performance; outliers suggest rare segmentation errors.