PaperlayerView exact revision on arXiv

Source-linked AI summary

Multi-Stage Prompt-Guided Feature Modulation for Generalizable Brain Tumor Segmentation

Mohammad Mahdi Danesh Pajouh, Sara Saeedi

arXiv:2608.23745v1eess.IVcs.CV

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.

    Source: §1 · p. 2

  • Method

    Multi-Stage Dynamic Prompt nnU-Net adds image-dependent prompt embeddings and FiLM-based feature modulation at three deep encoder stages.

    Source: §1 · p. 2, §2.1 · p. 3, Multi-Stage Prompt Conditioning · p. 5

  • Results

    Overall improvements were observed across the majority of lesion-wise overlap, boundary, and detection metrics on the BraTS GOAT validation dataset.

    Source: §4 · p. 12

  • Takeaways & Limitations

    Adaptive prompt-guided feature modulation is presented as an effective strategy for robust, generalizable brain tumor segmentation.

    Source: §4 · p. 12

    Alternative prompt sizes, additional stages, and longer training schedules did not yield consistent validation improvements, constraining the selected configuration.

    Source: Hardware and Optimization · p. 7

Abstract

from arXiv · show

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.Source: §1 · p. 1
  • Tumor morphology and MRI acquisition differences create domain shifts that can reduce performance outside the training distribution.Source: §1 · p. 2
  • Dynamic feature conditioning adapts internal representations to image-specific characteristics across heterogeneous tumor domains.Source: §1 · p. 2
  • Multi-Stage Dynamic Prompt nnU-Net inserts three dynamic prompt modules into the deepest encoder stages to generate FiLM parameters for adaptive feature conditioning.Source: §1 · p. 2
  • The proposed model improves lesion-wise Dice, NSD, and HD95 across the majority of evaluated tumor subregions compared with nnU-Net.Source: §1 · p. 2

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.Source: §2.1 · p. 3
  • The module pools each skip feature, uses a Selector Network to predict routing weights, aggregates a learnable Prompt Bank, and applies FiLM modulation.Source: Figure 2 · p. 4, §2.2 · p. 4
  • Standard global average pooling can destroy spatial awareness, which is important because anatomical location carries semantic clues in brain imaging.Source: Context-Aware Prompt Weighting · p. 4
  • Adaptive average pooling compresses features to a 2 × 2 × 2 grid, preserving coarse spatial information across eight brain-volume octants.Source: Context-Aware Prompt Weighting · p. 4
  • A lightweight two-layer selector network processes the spatial descriptor to generate prompt-selection logits.Source: Context-Aware Prompt Weighting · p. 4
  • Temperature-scaled softmax converts logits into a prompt-selection distribution whose sharpness controls routing behavior.Source: Context-Aware Prompt Weighting · p. 5

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.

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.

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.Source: §2.3 · p. 6

Segmentation Loss

  • Soft Dice Loss and Focal Loss with γ = 2 form the composite segmentation loss for handling class imbalance.Source: Segmentation Loss · p. 6
  • λ begins at 0.6 for Dice Loss and gradually decreases to 0.5, increasing Focal Loss contribution during training.Source: Segmentation Loss · p. 6
  • The weighting schedule was motivated by fewer false-negative predictions under the final balanced loss weighting.Source: Segmentation Loss · p. 6

Multi-Scale Deep Supervision

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.Source: §2.4 · p. 7

Two-Stage Training Curriculum

Training uses a two-phase curriculum that first promotes broad prompt exploration and then encourages specialized prompt selection.

Phase 1 (Warm-up and Exploration)

Phase 2 (Specialization)

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.Source: Hardware and Optimization · p. 7
  • The study evaluated alternative prompt banks, larger prompt dimensionality, additional encoder stages, and longer schedules without consistent validation gains.Source: Hardware and Optimization · p. 7
  • Optimization used SGD with Nesterov momentum, polynomial decay, and separate gradient clipping limits for backbone and prompt parameters.Source: Hardware and Optimization · p. 7
  • The baseline and proposed models used the same preprocessed images, nnU-Net configuration, training duration, 3D full-resolution backbone, and input patch configuration.Source: Baseline Configuration · pp. 7–8
  • The 3D backbone used six encoder stages with widths 32, 64, 128, 256, 320, and 320 channels and patches of 128 × 160 × 112.Source: Baseline Configuration · p. 8

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.Source: §3.1 · p. 9

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.Source: §3.1 · p. 9

Robustness

  • Median Dice and NSD improved for ET and WT and remained comparable or slightly reduced for TC across the validation cohort.Source: §3.1 · p. 9

Computational complexity

  • 39.95 M total trainable parameters: prompt-related components added approximately 9.16 M parameters to the 30.79 M baseline.Source: §3.2 · pp. 9–10
  • 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.Source: §3.2 · p. 10

Prompt behavior analysis

  • Prompt routing was non-uniform, with distinct utilization patterns across the three encoder stages and validation subjects.Source: Figure 3 · p. 10, §3.3 · p. 11

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.Source: §4 · p. 11

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.Source: §4 · p. 12

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.Source: §4 · p. 12

Related papers

Loading 2608.23745v1…