Diffusion-Refined Segmentation and Vision-Language Interpretation for Pediatric Brain Tumor MRI

Abstract: Accurate pediatric brain tumor segmentation remains challenging due to limited annotated data, heterogeneous imaging phenotypes, diffuse tumor boundaries, and class imbalance across tumor subregions. Here, we present a two-stage deep learning framework for improving multi-modal pediatric brain MRI segmentation and clinical interpretation. First, we evaluate 3D Res U-Net and Swin-UNETR baselines on BraTS-PEDs MRI scans, using four co-registered modalities to predict tumor core, whole tumor, and enhancing tumor regions. Second, we introduce diffusion-based refinement models conditioned on coarse Swin-UNETR predictions, including a 3D DDPM refiner and MedSegDiff. Conditioning substantially improves diffusion stability and performance, particularly for enhancing tumor boundary segmentation. Conditioned MedSegDiff achieves the strongest boundary agreement with the lowest HD95. Finally, predicted tumor volumes and representative segmentation overlays are integrated with a multimodal LLM to generate structured radiology-style reports. Together, our results suggest that coarse-to-refined diffusion segmentation can improve pediatric tumor boundary delineation and support end-to-end interpretable AI-assisted neuro-oncology workflows.

• The paper introduces a two-stage pipeline combining U-Net based coarse segmentation with conditional diffusion refinement to improve pediatric brain tumor delineation.
• The approach leverages transformer-based architectures and pretrained weights to address class imbalance and achieve robust Dice scores.
• The integration of a multimodal language model translates technical segmentation outputs into actionable radiology reports.

Diffusion-Refined Segmentation and Vision-Language Interpretation for Pediatric Brain Tumor MRI

Introduction and Motivation

Accurate segmentation of pediatric brain tumor MRI is critical for diagnosis, treatment planning, and monitoring. Pediatric tumors differ markedly from adult gliomas in anatomical distribution, imaging appearance, and boundary heterogeneity, as illustrated in the comparative visualization (Figure 1). Manual segmentation is time-consuming and highly variable, and conventional deep learning models trained on adult datasets perform inadequately on pediatric cases due to limited annotated data and pronounced class imbalance. These complexities necessitate robust, pediatric-focused approaches to achieve reliable and reproducible results in clinical workflows.

Figure 1: Difference between adult and pediatric brain tumor. Pediatric tumors present more diffuse and poorly defined boundaries compared to adults.

The paper presents a two-stage deep learning pipeline: initial coarse segmentation via U-Net-based architectures, followed by boundary refinement using conditional diffusion models. Integration with a multimodal LLM (MLLM) enables automated, interpretable clinical reporting, thereby bridging the gap between algorithmic segmentation and actionable clinical diagnostics.

Model Architectures and Segmentation Pipeline

Baseline Segmentation Models

The foundational segmentation architectures include a 3D Residual U-Net and the Swin-UNETR transformer model, both adapted from MONAI. The Res U-Net incorporates residual units to improve gradient propagation and adopts Soft Dice Loss to address severe class imbalance. The Swin-UNETR leverages efficient shifted-window self-attention on 3D patch embeddings, initialized with pretrained weights to accelerate convergence and mitigate overfitting risks. Both models process four co-registered MRI modalities and output voxel-wise masks for tumor subregions.

Diffusion-Based Refinement

Segmentation refinement employs generative denoising diffusion models: an unconditional 3D U-Net DDPM and MedSegDiff. The models cast segmentation as a conditional denoising problem. A prior coarse mask from Swin-UNETR is used as structural guidance, substantially improving stability and performance, particularly for small, ambiguous regions such as enhancing tumor (ET). Conditioning sharpens boundaries and mitigates noise, as demonstrated by training and validation dynamics (Figure 2).

Figure 3: Overview of the model and training setup. The pipeline consists of coarse segmentation followed by conditional diffusion-based refinement.

Data and Preprocessing

The BraTS-PEDs 2023 dataset comprises 99 multi-parametric pediatric MRI subjects, split into training/validation/test cohorts (Figure 4). Each subject's features span T1, T1CE, T2, and T2-FLAIR modalities. Rigorous preprocessing, including foreground cropping, channel-wise normalization, and extensive geometric/intensity augmentation, prevents overfitting due to limited data and enhances robustness to heterogeneous clinical imaging variations.

Figure 4: Dataset split, anatomical labels, and features for pediatric brain tumor segmentation.

Numerical Performance and Boundary Analysis

Quantitative comparisons highlight a marked hierarchy among models. The Res U-Net baseline is inadequate (Dice Avg: 0.48). The Swin-UNETR (pretrained) achieves robust accuracy (Dice Avg: 0.77, Dice WT: 0.87). Unconditioned diffusion models are unstable and prone to over-segmentation (MedSegDiff-scratch Dice Avg: 0.62), but conditional variants show substantial performance jumps (MedSegDiff-conditioned Dice Avg: 0.75, HD95: 9.35), matching or surpassing CNN baselines in boundary precision.

Figure 2: Training loss and validation Dice for the six models show the decisive impact of coarse conditioning on diffusion model stability and convergence.

Boundary refinement is particularly salient in conditioned diffusion models. While DDPM (conditioned) achieves strong region-wise Dice, its boundary noise remains elevated (HD95: 33.54). MedSegDiff (conditioned) attains the lowest HD95 among all models, indicating its utility for precisely delineating tumor margins beyond the coarse prior. Qualitative case comparisons validate these improvements, with conditioned models tracking ground truth morphology and outperforming unconditional diffusion (Figure 5).

Figure 5: Qualitative comparison of segmentation overlays for necrotic core, edema, and enhancing tumor. Conditioned models most closely follow ground truth.

Vision-Language Clinical Interpretation

To transform voxel-level masks into actionable outputs, a multimodal LLM (Gemini 2.5 Pro) ingests both volumetric metrics and representative imaging overlays, synthesizing structured radiology-style reports (Figure 6). The pipeline correlates geometric features and imaging phenotype, generating plausible diagnostic impressions and differential hypotheses. For example, Gemini's inference on cases with large necrotic cores and minimal edema accurately highlights aggressive lesion morphology, closely paralleling expert reasoning.

Figure 6: End-to-end AI-assisted clinical radiology report generation. Inputs include both volumetric metrics and segmentation overlays, with Gemini 2.5 Pro producing final diagnostic interpretations.

Implications, Limitations, and Future Directions

Practical and Theoretical Implications

The results underscore the value of conditional diffusion models in boundary refinement for pediatric tumor segmentation, with conditioned MedSegDiff achieving the strongest agreement in challenging regions. This suggests that generative refinement is advantageous for ambiguous or diffuse tumor margins, which are common in pediatric cases. The clinical integration of MLLMs enables seamless translation from segmentation to interpretative reporting, potentially reducing diagnostic workload and standardizing clinical evaluations.

Limitations

The workflow encounters several critical bottlenecks:

• Severe class imbalance, especially for small ET regions, remains problematic.
• Transformer and diffusion models are computationally expensive, constraining batch sizes and hindering real-time deployment.
• Iterative sampling in diffusion models prolongs inference times.
• Clinical report generation via foundation models is vulnerable to hallucination and error propagation from upstream segmentation inaccuracies.

Future Directions

Acceleration of diffusion inference (e.g., DDIM, Consistency Trajectory) is necessary for clinical adoption. Memory-efficient attention and patch-wise pretraining could further address class imbalance and scalability. Robust retrieval-augmented generation approaches and human-in-the-loop validation are recommended to minimize diagnostic hallucination and increase clinical reliability. Larger pediatric datasets and benchmarked, lesion-wise evaluation metrics would clarify comparative efficacy.

Conclusion

This study demonstrates that a coarse-to-refined pipeline, combining transformer-based segmentation with conditional diffusion refinement, achieves state-of-the-art boundary accuracy for pediatric brain tumor MRI. Integration with MLLMs bridges technical outputs and clinical interpretation. Addressing computational, morphological, and interpretative challenges will further enhance practical deployment and reliability in neuro-oncology AI applications.