Diet-Seg: Dynamic Hardness-Aware Learning for Enhanced Brain Tumor Segmentation

Abstract

Accurate brain tumor segmentation in magnetic resonance imaging (MRI) remains a critical challenge due to complex tumor heterogeneity, fuzzy boundaries, and significant inter-patient variability. In this study, we propose Diet-Seg (Difficulty-Informed Edge-enhanced Tiny Segmentation), a novel segmentation framework that integrates entropy-based pixel-wise hardness estimation into the training process via a dynamic learning rate modulation strategy. Specifically, we employ a pretrained 3D U-Net information model to quantify voxel-level prediction uncertainty, which is then used to guide the optimization of the main segmentation model. Diet-Seg is further enhanced by an RWKV-based U-Net backbone to capture global spatial dependencies and an EdgeNet module to preserve tumor boundaries through edge-aware fusion. Extensive experiments on the BraTS2018-2021 datasets demonstrate that Diet-Seg consistently outperforms state-of-the-art baselines across all tumor subregions. Notably, Diet-Seg achieves superior generalization when trained on one dataset and validated across multiple years. Moreover, the hardness maps offer interpretable insights into segmentation difficulty, potentially enabling human-AI collaboration in clinical practice. These results highlight the promise of entropy-aware training as a general strategy for robust and efficient medical image segmentation. The work is implemented in the open-source project available on GitHub (https://github.com/ManuelTurner/Diet-Seg)