Improving Generalization of Deep Learning for Brain Metastases Segmentation Across Institutions

Abstract: Background: Deep learning has demonstrated significant potential for automated brain metastases (BM) segmentation; however, models trained at a singular institution often exhibit suboptimal performance at various sites due to disparities in scanner hardware, imaging protocols, and patient demographics. The goal of this work is to create a domain adaptation framework that will allow for BM segmentation to be used across multiple institutions. Methods: We propose a VAE-MMD preprocessing pipeline that combines variational autoencoders (VAE) with maximum mean discrepancy (MMD) loss, incorporating skip connections and self-attention mechanisms alongside nnU-Net segmentation. The method was tested on 740 patients from four public databases: Stanford, UCSF, UCLM, and PKG, evaluated by domain classifier's accuracy, sensitivity, precision, F1/F2 scores, surface Dice (sDice), and 95th percentile Hausdorff distance (HD95). Results: VAE-MMD reduced domain classifier accuracy from 0.91 to 0.50, indicating successful feature alignment across institutions. Reconstructed volumes attained a PSNR greater than 36 dB, maintaining anatomical accuracy. The combined method raised the mean F1 by 11.1% (0.700 to 0.778), the mean sDice by 7.93% (0.7121 to 0.7686), and reduced the mean HD95 by 65.5% (11.33 to 3.91 mm) across all four centers compared to the baseline nnU-Net. Conclusions: VAE-MMD effectively diminishes cross-institutional data heterogeneity and enhances BM segmentation generalization across volumetric, detection, and boundary-level metrics without necessitating target-domain labels, thereby overcoming a significant obstacle to the clinical implementation of AI-assisted segmentation.

• The paper introduces a novel VAE-MMD unsupervised domain adaptation framework integrated with nnU-Net to harmonize multi-site imaging data.
• The methodology yields significant improvements in segmentation performance, increasing average F1 by 11.1% and reducing HD95 by 65.5%.
• The approach enables robust cross-institutional deployment without target labels, enhancing clinical viability in neuro-oncology segmentation.

Improving Generalization of Deep Learning for Brain Metastases Segmentation Across Institutions

Introduction

Cross-institutional deployment of deep neural networks for medical imaging is fundamentally hindered by domain shift, particularly in the context of brain metastases (BM) segmentation. Heterogeneous scanner hardware, diverse patient demographics, and site-specific acquisition protocols induce significant alignment and generalization challenges, as evidenced by characteristic domain gaps in lesion morphology and imaging appearances among major datasets—Stanford (multiple miliary lesions), UCSF (vascular enhancement and size variability), UCLM (solitary, well-circumscribed lesions), and PKG (large tumors from lung cancer primaries).

Figure 1: Representative T1-weighted post-contrast MRI slices across four institutions demonstrating heterogeneity in lesion characteristics and contributing to substantial cross-domain variance.

Conventional transfer learning and direct mixed training do not reliably solve these issues, often resulting in catastrophic forgetting and poor transfer performance. This work addresses these limitations by introducing a Variational Autoencoder with Maximum Mean Discrepancy (VAE-MMD) based unsupervised domain adaptation framework integrated with nnU-Net to harmonize domain biases and enhance generalization without access to target-domain labels (2604.00397).

Methodology

Datasets and Preprocessing

The study utilizes 740 patient MRIs from four public datasets (Stanford, UCSF, UCLM, PKG), employing only T1-weighted contrast-enhanced (T1CE) sequences to ensure modality uniformity and maximize cross-institutional comparability. Images were resampled, normalized, and standardized according to nnU-Net guidelines, before further normalization to $[-1, 1]$ and patch extraction for VAE-based transformation.

VAE-MMD Architecture

The VAE-MMD module comprises a deep 3D convolutional encoder-decoder with residual connections and self-attention in deep layers to model long-range dependencies. The latent space (512-D) is regularized via KL-divergence and aligned across sites using a multi-scale RBF MMD loss. Skip connections from encoder to decoder are critical to recover fine anatomical detail, minimizing the inherent information bottleneck of vanilla VAE models. The final reconstructions are leveraged as harmonized inputs for the downstream nnU-Net segmenter.

Figure 2: (a) VAE-MMD model structure; (b) VAE-MMD processing for a Stanford subject—domain-specific intensity statistics and stylistic variability are suppressed in the reconstructions, which display high PSNR and low mean squared error, with minimal difference from input anatomical structures.

The total loss includes an L2+L1+SSIM composite (for pixel fidelity and perceptual structure), KL-divergence, strong MMD for domain alignment, and adversarial loss for photorealistic reconstructions.

Domain Alignment Metrics

Domain shift ablation is evaluated by t-SNE embedding and logistic regression-based institution classifier accuracy on latent vectors, where a drop in classifier accuracy quantifies enhanced feature alignment and decorrelation from acquisition site.

Figure 3: (a) t-SNE plots across two and four centers pre- and post-VAE-MMD, showing strong post-alignment mixing; (b) confusion matrices for the domain classifier, evidencing near-random (chance) performance after VAE-MMD adaptation.

Segmentation and Evaluation

nnU-Net is trained using original and VAE-MMD reconstructed volumes across standard cross-validation splits. Performance is assessed via sensitivity, F1, F2, surface Dice (sDice), and 95th percentile Hausdorff distance (HD95), with metric computation stratified by institution.

Results

Domain Adaptation

VAE-MMD yields a marked reduction in discernibility of site-specific feature signatures: the institution identification classifier accuracy falls from 91.0% to 50.0%—a 45.1% reduction—demonstrating the effective elimination of domain-aligned information from the latent space, consistent with visual evidence from t-SNE plots.

Quantitative Segmentation Gains

VAE-MMD preprocessing consistently improves segmentation performance across all evaluation centers:

• Average F1 increases by 11.1% (0.700 → 0.778)
• Mean sDice rises 7.93% (0.7121 → 0.7686)
• Mean HD95 is reduced by 65.5% (11.33 mm → 3.91 mm)

Figure 4: Across Stanford, UCSF, UCLM, and PKG, VAE-MMD preprocessing consistently elevates sensitivity, F1, and sDice, while producing substantial reductions in HD95—highlighting improved lesion boundary detection and geometric accuracy.

The most substantial HD95 improvements are seen for Stanford (22.46 → 2.12 mm), underscoring the model's ability to recover missed small miliary lesions that drive cross-institutional failure modes. UCLM, which presents the strongest baseline domain gap, achieves F1 = 0.874 post-adaptation.

Qualitative Analysis

Figure 5: Qualitative overlay of baseline vs. VAE-MMD segmentations across four institutions, showing enhanced alignment with expert annotations and improved lesion boundary adherence.

Discussion

Design Contributions

• Latent-space MMD alignment is superior to pixel-aligned normalization, as it enables adaptation at the semantic level, filtering out both global and fine-grained site artifacts.
• Multi-scale MMD kernels robustly capture both fine and coarse domain discrepancies.
• Residual and attention modules, together with skip connections, preserve micro-anatomical structures necessary for accurate BM segmentation, minimizing the typical trade-off between domain invariance and loss of diagnostic cues.
• The architecture demonstrates that unsupervised adaptation without target labels is feasible and effective, substantially lowering the barrier to out-of-distribution deployment in real-world hospital settings.

Clinical Implications

The increased sensitivity, especially in datasets with subtle or small lesions, directly impacts the clinical missed-lesion rate in stereotactic radiosurgery planning, with reductions in HD95 further supporting the clinical acceptability of predicted contours relative to the stringent margin demands for radiotherapeutic planning.

The unsupervised nature enables rapid scaling of segmentation tools to new clinical sites without costly annotation campaigns or retraining, thus facilitating wider adoption of AI in neuro-oncology.

Limitations and Future Outlook

• Restriction to T1CE modalities may limit applicability to multi-sequence MRI protocols; multimodal VAE-MMD extensions are a natural next step.
• Prospective, federated, and privacy-preserving adaptations are suggested for exploring secure, decentralized model deployment.

Conclusion

This study provides rigorous evidence that VAE-MMD driven unsupervised domain adaptation can bridge substantial inter-institutional gaps in BM segmentation without recourse to labeled target-domain data. By injecting robust domain invariance into training protocols while maintaining boundary-level anatomical precision, this approach lays the groundwork for scalable, transferable AI segmentation in heterogeneous clinical environments.

Reference: "Improving Generalization of Deep Learning for Brain Metastases Segmentation Across Institutions" (2604.00397)