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

2026-01-13

AI Summary

This study addresses the boundary noise commonly produced by existing CNN- and U-Net-based models in glioma segmentation, particularly within tumor infiltration zones, which can adversely affect clinical decision-making. To mitigate this issue, the authors propose UAMSA-UNet, a novel architecture that, for the first time, integrates a variance-penalizing loss into a Bayesian U-Net framework with Monte Carlo Dropout to learn smoothness priors in a data-driven manner, thereby enhancing spatial consistency. The model further incorporates a multi-scale attention mechanism to improve segmentation accuracy while reducing computational overhead. Experimental results demonstrate consistent performance gains: on BraTS2023, Dice and mIoU scores increase by 3.3% and 2.7%, respectively, with further improvements of 4.5% and 4.0% on BraTS2024. Additionally, the proposed method achieves a 42.5% reduction in FLOPs compared to U-Net++.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Multimodal Learning
• Multiagent Systems: Adversarial Agents

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

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.

Problem

• brain tumour segmentation
• noisy boundaries
• tumor infiltration
• spatial coherence
• segmentation quality

Innovation

• Variance-Penalized MC-Dropout
• Learned Smoothing Prior
• Multi-Scale Attention
• Uncertainty-Aware Segmentation
• Bayesian U-Net