Source-linked AI summary
Multi-Stage Prompt-Guided Feature Modulation for Generalizable Brain Tumor Segmentation
Mohammad Mahdi Danesh Pajouh, Sara Saeedi
TL;DR
Domain shifts from heterogeneous tumors and MRI settings challenge generalizable brain tumor segmentation. The paper extends nnU-Net with multi-stage dynamic prompt conditioning, and reports overall improvements across most evaluated metrics and tumor subregions on BraTS GOAT validation data.
Problem
Tumor and MRI variability across clinical domains can substantially reduce segmentation performance outside the training distribution.
Method
Multi-Stage Dynamic Prompt nnU-Net adds image-dependent prompt embeddings and FiLM-based feature modulation at three deep encoder stages.
Results
Overall improvements were observed across the majority of lesion-wise overlap, boundary, and detection metrics on the BraTS GOAT validation dataset.
Takeaways & Limitations
Adaptive prompt-guided feature modulation is presented as an effective strategy for robust, generalizable brain tumor segmentation.
Takeaways & Limitations
Alternative prompt sizes, additional stages, and longer training schedules did not yield consistent validation improvements, constraining the selected configuration.
Abstract
from arXiv · showhide
Accurate brain tumor segmentation from magnetic resonance imaging (MRI) is essential for diagnosis, treatment planning, surgical guidance, and disease monitoring. However, developing automated segmentation models that generalize across diverse tumor characteristics, imaging protocols, acquisition sites, and patient populations remains challenging. Variations in tumor morphology and imaging distributions can substantially degrade performance outside the training domain. Consequently, improving the robustness and generalization of deep learning-based segmentation models has become a key objective in medical image analysis. To improve segmentation robustness, we propose Multi-Stage Dynamic Prompt nnU-Net, a prompt-conditioned extension of nnU-Net. Three independent dynamic prompt modules are inserted into the deepest encoder stages. Each module contains a learnable bank of ten 256-dimensional prompt vectors and uses globally pooled encoder features to generate image-specific prompt representations. These representations are projected into feature-wise scaling $(γ)$ and shifting $(β)$ parameters that modulate encoder feature maps through Feature-wise Linear Modulation (FiLM), enabling adaptive feature conditioning at multiple semantic levels. Evaluation on the BraTS GOAT validation dataset demonstrated that the proposed Multi-Stage Dynamic Prompt nnU-Net outperformed the baseline nnU Net across the majority of evaluated metrics and tumor subregions. The proposed model achieved average lesion-wise Dice scores of 76.16% (ET), 80.04% (TC), and 86.42% (WT), compared with 74.38%, 78.14% and 84.01% for the baseline model. The results demonstrate that multi-stage dynamic prompt conditioning improves segmentation accuracy and boundary delineation for brain tumor segmentation.
1 Introduction
Brain tumor segmentation models must handle substantial variation across tumors, imaging conditions, institutions, and patient populations. Multi-Stage Dynamic Prompt nnU-Net addresses this challenge by adaptively conditioning encoder features and improves most evaluated metrics over nnU-Net on BraTS GOAT validation data.
- Clinical brain tumor segmentation supports diagnosis, treatment planning, surgical guidance, radiotherapy planning, and longitudinal monitoring.
- Tumor morphology and MRI acquisition differences create domain shifts that can reduce performance outside the training distribution.
- Dynamic feature conditioning adapts internal representations to image-specific characteristics across heterogeneous tumor domains.
- Multi-Stage Dynamic Prompt nnU-Net inserts three dynamic prompt modules into the deepest encoder stages to generate FiLM parameters for adaptive feature conditioning.
- The proposed model improves lesion-wise Dice, NSD, and HD95 across the majority of evaluated tumor subregions compared with nnU-Net.
2 Methodology
The method adds input-dependent prompt embeddings to nnU-Net and uses spatially aware feature summaries to select and aggregate prompts. The resulting context-aware prompt conditions feature modulation for heterogeneous inputs.
- The Dynamic Multi-Prompt framework adds instance-aware prompt embeddings to nnU-Net without drastically increasing computational footprint.
- The module pools each skip feature, uses a Selector Network to predict routing weights, aggregates a learnable Prompt Bank, and applies FiLM modulation.
- Standard global average pooling can destroy spatial awareness, which is important because anatomical location carries semantic clues in brain imaging.
- Adaptive average pooling compresses features to a 2 × 2 × 2 grid, preserving coarse spatial information across eight brain-volume octants.
- A lightweight two-layer selector network processes the spatial descriptor to generate prompt-selection logits.
- Temperature-scaled softmax converts logits into a prompt-selection distribution whose sharpness controls routing behavior.
Residual Feature-wise Linear Modulation (FiLM)
The aggregated prompt generates bounded channel-wise scaling and shifting parameters that modulate encoder features through a residual FiLM operation. Near-zero initialization preserves the base representation at training onset.
- The aggregated prompt is projected into channel-wise scaling γ and shifting β parameters for FiLM feature modulation.
- Hyperbolic tangent activations bound γ and β between -1 and 1 to constrain modulation strength and improve training stability.
- Residual modulation preserves semantic features from the base nnU-Net while incorporating prompt-derived changes.
- When projections are initialized near zero, the modulation approaches an identity mapping, preventing early forward-pass destabilization.
Multi-Stage Prompt Conditioning
Separate Dynamic Prompt Modules condition the three deepest encoder skip connections. This provides the decoder with complementary adaptively modulated features at multiple semantic scales.
- A separate Dynamic Prompt Module is attached to each of the three deepest encoder skip connections.
- Conditioning multiple depths supplies complementary representations spanning different levels of semantic abstraction to the decoder.
2.3 Training Strategy and Loss Formulation
The training strategy combines region-based segmentation losses, deep supervision, and orthogonality regularization, with loss weighting adjusted during training.
- 2.3 Training Strategy and Loss Formulation: The network combines region-based segmentation objectives, deep supervision, and orthogonality regularization.
- Segmentation Loss: Soft Dice Loss and Focal Loss with γ = 2 form the composite segmentation loss for handling class imbalance.
- Segmentation Loss: λ begins at 0.6 for Dice Loss and gradually decreases to 0.5, increasing Focal Loss contribution during training.
- Segmentation Loss: The weighting schedule was motivated by fewer false-negative predictions under the final balanced loss weighting.
- Multi-Scale Deep Supervision: Deep supervision evaluates segmentation loss at multiple decoder resolutions using the default nnU-Net weighting scheme.
- Multi-Scale Deep Supervision: Deep-supervision weights decay exponentially at lower resolutions and are normalized to sum to one.
- Multi-Scale Deep Supervision: Orthogonality regularization constrains each stage’s L2-normalized prompt matrix toward an identity matrix to prevent prompt collapse.
- Multi-Scale Deep Supervision: The final objective function consolidates the segmentation, deep-supervision, and regularization components.
2.4 Implementation Details Prompt Configuration
Each Dynamic Prompt Module uses a compact, distinct learnable prompt bank with fixed dimensionality.
- Each Dynamic Prompt Module maintains N=10 learnable prompt vectors, each with dimensionality 256.
Two-Stage Training Curriculum
Training uses a two-phase curriculum that first promotes broad prompt exploration and then encourages specialized prompt selection.
- Training was divided into two phases to support stable convergence and prevent prompt collapse.
- Phase 1 (Warm-up and Exploration): During Phase 1, T=0.5 produces softer routing so gradients update prompts more uniformly while model components co-adapt.
- Phase 2 (Specialization): During Phase 2, T=0.09 sharpens routing to force more confident, specialized prompt selections.
- Phase 2 (Specialization): Phase 2 also reduces learning rates to 0.01× for base U-Net parameters and 0.1× for the selector network.
Hardware and Optimization
The model was trained under constrained hardware conditions and compared against a baseline using matched preprocessing, architecture configuration, and training duration.
- The proposed architecture was trained for 675 epochs on a single NVIDIA GTX 1080 GPU with 8 GB VRAM, batch size 1, and AMP.
- The study evaluated alternative prompt banks, larger prompt dimensionality, additional encoder stages, and longer schedules without consistent validation gains.
- Optimization used SGD with Nesterov momentum, polynomial decay, and separate gradient clipping limits for backbone and prompt parameters.
- The baseline and proposed models used the same preprocessed images, nnU-Net configuration, training duration, 3D full-resolution backbone, and input patch configuration.
- The 3D backbone used six encoder stages with widths 32, 64, 128, 256, 320, and 320 channels and patches of 128 × 160 × 112.
3 Experiments and Results
On the BraTS GOAT validation set, DPM nnU-Net improved most segmentation metrics over baseline nnU-Net, with strongest gains for whole-tumor segmentation. The model added parameters but increased inference time only modestly, while prompt routing varied across stages and subjects.
- Overall performance: The proposed DPM nnU-Net improved most segmentation metrics over baseline nnU-Net, with only minor degradations for some individual measurements.
- Tumor-region results: The proposed method improved enhancing-tumor Dice, NSD, F1-score, and HD95, while tumor-core Dice and F1-score improved despite slightly better baseline HD95 and NSD.
- Robustness: Median Dice and NSD improved for ET and WT and remained comparable or slightly reduced for TC across the validation cohort.
- Computational complexity: 39.95 M total trainable parameters: prompt-related components added approximately 9.16 M parameters to the 30.79 M baseline.
- Computational complexity: Inference time increased from 3.528 s to 3.649 s per volume, while the framework introduced approximately 29.7% more trainable parameters and a 3.4% time increase.
- Prompt behavior analysis: Prompt routing was non-uniform, with distinct utilization patterns across the three encoder stages and validation subjects.
4 Conclusion
The proposed Multi-Stage Dynamic Prompt nnU-Net adds instance-aware feature modulation to nnU-Net for heterogeneous brain tumor segmentation. On BraTS GOAT validation, it improved segmentation quality across most evaluated metrics.
- Method: Dynamic Prompt Modules generate image-specific prompt embeddings and recalibrate semantic encoder features through FiLM before decoding.The framework uses spatially-aware feature pooling and soft routing to adapt representations for each input volume.
- Results: Overall improvements were observed across the majority of lesion-wise overlap, boundary, and detection metrics on the BraTS GOAT validation dataset.These results support improved segmentation quality under heterogeneous tumor conditions.
- Implication: Adaptive prompt-guided feature modulation is presented as an effective strategy for robust, generalizable segmentation across heterogeneous clinical neuro-oncology data.The conclusion contrasts adaptive conditioning with static feature extraction.