Source-linked AI summary
MS-Net: Multi-Site Network for Improving Prostate Segmentation with Heterogeneous MRI Data
Quande Liu, Qi Dou, Lequan Yu, Pheng Ann Heng
TL;DR
Prostate MRI segmentation needs multi-site data because single-site samples are scarce, yet scanner and protocol differences make aggregation difficult. MS-Net combines site-specific batch normalization with multi-site-guided knowledge transfer to learn robust shared representations, and experiments show consistent improvements across three heterogeneous datasets. The method is most applicable when the testing site's identity is known and represented during training.
Problem
Single-site medical imaging data are scarce, while heterogeneous MRI distributions across sites complicate effective joint training for prostate segmentation.
Method
MS-Net combines domain-specific batch normalization with auxiliary branches that transfer multi-site knowledge to shared network kernels.
Results
MS-Net consistently improves performance on three heterogeneous datasets and outperforms Parallel-Adapter by 0.78% in Dice score and 0.07mm in Average symmetric distance overall.
Takeaways & Limitations
The experiments support MS-Net as a multi-site prostate-segmentation strategy for learning robust representations from heterogeneous clinical MRI data.
Takeaways & Limitations
Testing requires the sample's site prior, and the method's generalization is limited for unseen sites because their statistics and trainable variables are unknown.
Abstract
from arXiv · showhide
Automated prostate segmentation in MRI is highly demanded for computer-assisted diagnosis. Recently, a variety of deep learning methods have achieved remarkable progress in this task, usually relying on large amounts of training data. Due to the nature of scarcity for medical images, it is important to effectively aggregate data from multiple sites for robust model training, to alleviate the insufficiency of single-site samples. However, the prostate MRIs from different sites present heterogeneity due to the differences in scanners and imaging protocols, raising challenges for effective ways of aggregating multi-site data for network training. In this paper, we propose a novel multi-site network (MS-Net) for improving prostate segmentation by learning robust representations, leveraging multiple sources of data. To compensate for the inter-site heterogeneity of different MRI datasets, we develop Domain-Specific Batch Normalization layers in the network backbone, enabling the network to estimate statistics and perform feature normalization for each site separately. Considering the difficulty of capturing the shared knowledge from multiple datasets, a novel learning paradigm, i.e., Multi-site-guided Knowledge Transfer, is proposed to enhance the kernels to extract more generic representations from multi-site data. Extensive experiments on three heterogeneous prostate MRI datasets demonstrate that our MS-Net improves the performance across all datasets consistently, and outperforms state-of-the-art methods for multi-site learning.
I. INTRODUCTION
Prostate MRI segmentation needs aggregated multi-site data because single-site samples are scarce, but inter-site heterogeneity makes joint training difficult. MS-Net addresses this challenge with site-specific normalization and multi-site knowledge transfer, improving segmentation consistently across three datasets.
- Single-site prostate MRI datasets are often too small for data-driven CNNs, motivating aggregation across hospitals for robust model training.
- Differences in imaging protocols, endorectal coil usage, and population demographics create inter-site heterogeneity that complicates multi-site training.
- Multi-site learning jointly trains one network on S site datasets to improve segmentation consistently across those sites, excluding external-site performance from its scope.
- MS-Net uses Domain-Specific Batch Normalization to normalize features separately for each site and Multi-site-guided Knowledge Transfer to learn shared robust representations.
- Experiments on three prostate MRI datasets show consistent accuracy gains over Joint, Separate, and state-of-the-art multi-site learning approaches.
II. RELATED WORKS
Prior prostate segmentation work includes atlas-based, deformable, handcrafted-feature, and deep-learning methods, while multi-domain learning motivates domain-specific modeling. MS-Net combines a universal network with site-specific branches to learn shared representations from heterogeneous datasets.
- Medical image analysis has used intensity normalization and handcrafted features or classifiers to address appearance differences across domains.
- Multi-domain learning studies single networks across diverse visual domains and finds that domain-specific parameters can handle domain-specific nuances.
- Earlier prostate segmentation methods include multi-atlas registration, deformable models, graph cuts, feature-based classifiers, and deep-learning networks.
- MS-Net contains a universal network and site-associated auxiliary branches, with DSBN layers addressing heterogeneity and transferred knowledge guiding shared kernels.
A. Domain Specific Batch Normalization Layer
Batch-normalization statistics vary across prostate MRI sites, making shared normalization problematic. MS-Net therefore assigns site-specific normalization parameters and statistics to handle inter-site differences during training and testing.
- Batch normalization normalizes channel-wise internal representations and applies trainable affine parameters to improve feature discrimination and learning speed.
- BN collects moving means and variances during training and uses them for feature normalization during testing.
- Single-site models show visible cross-site variations in BN statistics, especially in high-dimensional middle layers.
- Sharing global BN statistics across heterogeneous sites can hinder generic representation learning and degrade testing performance through inaccurate normalization.
- DSBN assigns each site individual BN variables and applies its collected moving statistics to corresponding-site test data.
B. Multi-site-guided Knowledge Transfer
MSKT trains site-specific auxiliary branches alongside a universal network, transferring their combined knowledge to promote robust, domain-agnostic representations across heterogeneous datasets.
- MSKT is designed to guide shared CNN kernels toward wide valleys among datasets rather than sharp dataset-specific minima.The motivation is that different auxiliary representations and predictions provide additional multi-site information during universal-network training.
- MSKT adds auxiliary branches that independently learn site-specific knowledge while a universal network learns shared representations.Each auxiliary branch matches the universal decoder architecture and is trained for one site.
- The universal network receives both ground-truth supervision and multi-site knowledge transferred from all auxiliary branches together.The transfer uses a dice-coefficient-like loss to align universal probability maps with auxiliary prediction masks.
- The knowledge-transfer loss operates on activation outputs represented by batch size, spatial dimensions, and channel number.The auxiliary masks are converted to one-hot format so their dimensions match the universal probability maps.
- Jointly training auxiliary branches and the universal network also acts as positive feature regularization for the universal encoder.
C. Network Architecture and Implementation Details
MS-Net uses an adapted 2D Residual-UNet backbone and trains universal and auxiliary parameters from site-specific mini-batches before deploying only the universal network.
- Network backbone: The backbone is an adapted 2D Residual-UNet with four down-sampling and four up-sampling blocks.A 2D architecture is used because through-plane resolution varies substantially across sites.
- Training procedure: Training begins with datasets from S sites and randomly initialized universal encoder, universal decoder, and auxiliary parameters.
- Training procedure: Each iteration uses S mini-batches, with one batch drawn from each site dataset.
- Training procedure: The procedure computes the universal-network loss, updates the universal encoder and decoder, and then retains only these parameters for deployment.Auxiliary parameters are not kept for testing.
2) Objective functions and training procedure:
The objective jointly optimizes auxiliary and universal segmentation losses, knowledge transfer, and L2 regularization, while training alternates between auxiliary branches and the universal network.
- Objective functions: The objective functions update auxiliary branches and the universal network using their respective segmentation losses plus L2 regularization.
- Objective functions: The universal objective balances segmentation loss and knowledge-transfer loss with α, set to 0.5 in the implementation.
- Objective functions: L2 regularization applies to the trainable universal encoder, universal decoder, and auxiliary-branch parameters, with η set to 1e−4.
- Training procedure: For each iteration, site-specific batches are fed to the network while auxiliary branches and the universal network are trained alternately.
- Training procedure: After training, all auxiliary branches are removed and only the universal network remains for inference.
3) Implementation details:
Implementation uses TensorFlow and multi-GPU training, with augmentation and post-processing applied to support prostate segmentation across three site datasets.
- Implementation details: The framework was implemented in Python with TensorFlow and trained on three NVIDIA TitanXp GPUs, assigning one GPU to each site.
- Implementation details: Training used Adam for 30000 iterations with batch size 5 and an initially 1.0 × 10−3 learning rate decayed by power 0.95 every 500 iterations.
- Implementation details: Random horizontal flipping and random shifting were used as data augmentation to mitigate overfitting.
- Implementation details: Post-processing selects the largest connected 3D volume as the final segmentation mask.
IV. EXPERIMENTS
The experiments evaluate three heterogeneous prostate MRI sites, preprocessing strategies, and baseline multi-site training behavior. Whitening normalization helps, but preprocessing alone does not fundamentally resolve inter-site heterogeneity.
- Datasets and evaluation: The evaluation uses 30, 30, and 19 T2-weighted MRI samples from Sites A, B, and C, respectively.The datasets come from three different institutions and are divided into training and testing subsets.
- Datasets and evaluation: Scanner, field-strength, coil, resolution, and population differences create substantial inter-site discrepancies across the three datasets.Most Site C images are from patients with prostate cancer, introducing a semantic difference in the prostate region.
- Preprocessing and metrics: Each sample is intensity-normalized to zero mean and unit variance before network input, with Dice coefficient and average symmetric distance used for evaluation.Performance is reported as mean±std to represent average performance and cross-subject variance.
- Normalization analysis: Whitening normalization improves Joint performance most among the tested intensity-normalization methods, raising Dice by 3.15% over Joint without intensity normalization.Bias-field correction and noise filtering provide limited and inconsistent additional gains across sites.
- Normalization analysis: Careful preprocessing remains insufficient to fundamentally solve heterogeneity in multi-site network training.The remaining discrepancy is attributed to factors such as coil usage and resolution differences that preprocessing cannot address.
C. Effectiveness of Our Multi-site Learning Method
MS-Net is compared with Separate and Joint baselines using heterogeneous-site segmentation, statistical testing, convergence, and inference-time analyses. It consistently improves segmentation across all three sites while retaining similar inference time to Joint.
- Comparison with baseline settings: DSBN outperforms Separate on all three sites in both Dice score and average symmetric distance.This result supports using domain-specific feature normalization to address inter-site discrepancy.
- Comparison with baseline settings: MS-Net further improves on DSBN, achieving average segmentation Dice of 91.66% and overall average symmetric distance of 0.77mm.The authors attribute these gains to more robust representations from the shared kernels.
- Comparison with baseline settings: MS-Net produces more accurate segmentation masks and clearer boundaries than Joint and Separate across all three sites.Qualitative examples compare ground truth, Joint, Separate, and the proposed approach.
- Statistical analysis: All paired t-tests yield p-value smaller than 0.05 against the Joint and Separate baselines, indicating statistically significant improvements.Tests use Dice score and average symmetric distance, with single-site and overall p-values reported.
- Computational analysis: The Joint, Separate, and Ours approaches take 5.5, 3.5, and 3 hours to converge, respectively.The proposed approach also converges faster and has a smoother overall Dice-loss curve than both baselines.
- Computational analysis: Inference takes 15.035s, 16.460s, and 15.224s for Joint, Separate, and Ours, respectively, making Ours similar to Joint.The comparison averages inference time over ten runs on three samples from each site.
2) Comparison with the state-of-the-art:
MS-Net is compared with multi-site and multi-domain learning baselines, including adapter, dual-stream, and feature-recalibration methods. It outperforms the state-of-the-art Parallel-Adapter in overall segmentation performance.
- Comparison with the state-of-the-art: The comparison includes USE-Net, Dual-Stream, Series-Adapter, and Parallel-Adapter methods for multi-site or multi-domain learning.The methods address heterogeneity through channel recalibration, shared bottleneck layers, or domain-adapter modules.
- Comparison with the state-of-the-art: MS-Net outperforms Parallel-Adapter by 0.78% in Dice score and 0.07mm in average symmetric distance overall.The authors present this comparison as evidence that MS-Net learns robust representations from multi-site data.
D. Ablation Analysis of Our Approach
Ablation experiments examine the knowledge-transfer loss ratio, its interaction with DSBN, and portability across network backbones. The results support moderate transfer weighting and show gains across several architectures.
- MSKT with different loss ratio: MSKT generally improves segmentation when the knowledge-transfer loss ratio α is around 0.5, whereas α = 1.0 does not guarantee an acceptable increase.Ground-truth supervision and transferred knowledge are complementary in learning robust representations.
- MSKT with different loss ratio: Applying MSKT to Joint improves Dice by 0.57% (90.25% vs. 90.82%), compared with a 0.75% gain on DSBN.The authors suggest DSBN gives shared kernels more capacity to learn generic multi-site representations.
- Different network backbones: MS-Net improves overall Dice by 1.41%, 1.27%, and 1.92% with Res-UNet, Dense-Unet, and Mobile-Unet backbones, respectively.These experiments demonstrate feasibility across different segmentation network designs.
- Different network backbones: Transferred multi-site knowledge provides a larger Dice improvement under the lightweight network-backbone setting.The authors connect the larger gain to the stronger role of transferred-knowledge regularization for lightweight models.
V. DISCUSSION
The paper addresses multi-site prostate segmentation by combining heterogeneous MRI data while targeting consistent performance across internal sites. MS-Net uses domain-specific normalization and transferred multi-site knowledge, but requires site information and does not establish external-site generalization.
- Problem setting: Multi-site learning jointly trains one network on datasets from different sites to improve segmentation consistently across those internal sites.
- Method: MS-Net compensates for inter-site heterogeneity with domain-specific batch normalization and enhances shared kernels through multi-site-guided knowledge transfer.
- Results: Experiments on three heterogeneous prostate MRI datasets show consistent improvement over baseline and state-of-the-art multi-site approaches.
- Limitations: Inference requires the testing sample's site prior so the network can select corresponding batch-normalization statistics.
- Limitations: For unseen sites, unknown domain statistics and trainable variables limit DSBN-based generalization, while external-site performance is outside the main study scope.
- Future work: Future work targets improving performance on both internal and external unseen sites and extending multi-site learning to other medical-imaging tasks.