UNISELF: Unified MS Lesion Segmentation

• UNISELF is a unified deep learning framework for automated MS lesion segmentation that handles multi-contrast MR images with high accuracy.
• It integrates a 2.5D U-Net backbone with conditional and test-time instance normalization to mitigate domain shifts and missing modalities.
• The self-ensembled lesion fusion approach and extensive augmentation deliver robust performance validated across multiple challenging datasets.

UNISELF is a unified deep learning framework for automated multiple sclerosis (MS) lesion segmentation from multi-contrast magnetic resonance (MR) images, designed to optimize both in-domain accuracy and out-of-domain generalization, even when trained on a single source dataset with limited annotated data. The method integrates a 2.5D U-Net backbone, conditional instance normalization (CondIN), test-time instance normalization (TTIN), and a novel self-ensembled lesion fusion algorithm to achieve robust and accurate lesion mapping across heterogeneous clinical imaging protocols and missing input modalities (Zhang et al., 6 Aug 2025).

1. Architecture and Feature Normalization

The core of UNISELF is a five-level 2.5D U-Net operating on volumetric MR data, with an encoder–decoder structure and a standard sequence of convolutions, nonlinearity, and normalization. At each level:

• Two blocks are applied, each consisting of a $3 \times 3$ convolution, ReLU activation, and feature normalization.
• Down-sampling is achieved with $2 \times 2$ max pooling, while up-sampling employs nearest-neighbor interpolation and a $3 \times 3$ convolution.
• The final layer is a $1 \times 1$ convolution followed by a sigmoid activation to yield a lesion probability map.

Feature normalization during training is selectable among BatchNorm (BN), InstanceNorm (IN), or conditional instance normalization (CondIN). All variants maintain learnable affine parameters $(\gamma, \beta)$. InstanceNorm operates per-sample, per-channel with the following statistics for a feature map $x \in \mathbb{R}^{n \times c \times h \times w}$:

$$\mu_{i,j} = \frac{1}{h\,w}\sum_{p=1}^h\sum_{q=1}^w x_{i,j,p,q}, \quad \sigma^2_{i,j} = \frac{1}{h\,w}\sum_{p,q} (x_{i,j,p,q} - \mu_{i,j})^2,$$

$$\hat{x}_{i,j,p,q} = \frac{x_{i,j,p,q} - \mu_{i,j}}{\sqrt{\sigma^2_{i,j}+\epsilon}}, \quad y_{i,j,p,q} = \gamma_j\,\hat{x}_{i,j,p,q}+\beta_j.$$

Conditional IN adapts the affine parameters based on the modality present.

2. Test-Time Instance Normalization (TTIN)

UNISELF addresses domain shifts and missing input contrasts at inference via TTIN. Instead of applying training-domain batch statistics (as in BN), TTIN recomputes the mean and variance for each channel of each test input (with batch size = 1):

$$\mu^{\text{test}}_j = \frac{1}{H\,W}\sum_{h,w} x_{j,h,w}, \quad (\sigma^2)^{\text{test}}_j = \frac{1}{H\,W}\sum_{h,w} (x_{j,h,w} - \mu^{\text{test}}_j)^2,$$

$$\hat{x}_{j,h,w} = \frac{x_{j,h,w} - \mu^{\text{test}}_j}{\sqrt{(\sigma^2)^{\text{test}}_j+\epsilon}}, \quad y_{j,h,w} = \gamma_j\,\hat{x}_{j,h,w}+\beta_j.$$

This normalization removes domain-specific bias in latent features at test time. t-SNE visualization demonstrates that TTIN collapses feature clusters from disparate sites into a single manifold, indicating effective mitigation of domain shift. Contrast dropout (CD) in training, when combined with TTIN, ensures robustness to missing imaging contrasts and prevents feature distributions from being biased toward particular modalities, such as FLAIR (Zhang et al., 6 Aug 2025).

3. Self-Ensembled Lesion Fusion

UNISELF employs a test-time self-ensemble framework combining extensive multi-orientation test-time augmentation (TTA) and a two-threshold connected lesion fusion process:

• Multi-orientation TTA: Eight spatial augmentations (four in-plane rotations and horizontal flips) are applied to each of the three anatomical planes (axial, sagittal, coronal), resulting in $2 \times 2$0 augmented inputs. The U-Net processes each input slice-wise to produce binary lesion masks $2 \times 2$1 after thresholding at 0.5, which are then mapped back to the original orientation.
• Confidence Map and Fusion: A voxel-level confidence map is computed,

$2 \times 2$2

Two confidence thresholds $2 \times 2$3 (empirically $2 \times 2$4, $2 \times 2$5) define high-confidence seeds ($2 \times 2$6) and a low-confidence region ($2 \times 2$7):

$2 \times 2$8

The final segmentation starts from $2 \times 2$9 and grows to include any 26-connected voxels in $3 \times 3$0, iteratively until convergence. This two-step procedure mimics Canny's hysteresis thresholding, improving the delineation of lesion extent and balancing false positives and missed detections better than ordinary majority voting.

4. Training Data, Preprocessing, and Augmentation

UNISELF is trained on the ISBI 2015 longitudinal MS lesion challenge dataset, which consists of 5 subjects with 21 annotated timepoints and four MR contrasts (T1w, T2w, PDw, FLAIR). Each scan is rated by two experts. Preprocessing includes rigid registration to $3 \times 3$1 MNI space, brain extraction, dura removal, and N4 bias correction.

Data augmentation includes:

• 3D random elastic or affine transforms (TorchIO, $3 \times 3$2),
• Randomly dropping out subsets of modalities via contrast dropout,
• Multi-orientation slicing, flips, and rotations,
• Random selection of rater labels for each training iteration.

The model receives as input a "2.5D" stack of three consecutive slices across all available modalities, generating predictions slice-wise. The loss function is voxel-wise $3 \times 3$3:

$3 \times 3$4

where $3 \times 3$5 is the predicted probability and $3 \times 3$6 is the ground-truth label. Optimization is performed using Adam with a learning rate of $3 \times 3$7, batch size of 12, for 150 epochs (300 iterations per epoch) (Zhang et al., 6 Aug 2025).

5. Benchmarking and Quantitative Performance

UNISELF is evaluated across multiple domains with significant domain shift phenomena, including different acquisition protocols, scanner types, and missing imaging contrasts. Metrics include Dice similarity (DSC), precision (PPV), sensitivity (TPR), lesion-wise TPR (LTPR), lesion-wise FPR (LFPR), volume correlation (VC), and the Carass et al. overall challenge score:

$3 \times 3$8

Summary of Test Results

Dataset & Method & Score (mean ± SD)	DSC	PPV	TPR	LFPR	LTPR	VC
ISBI 2015	UNISELF	93.28 ± 7.40	0.664 ± 0.141	0.882 ± 0.099	0.569 ± 0.199	0.100 ± 0.102
MICCAI 2016	UNISELF (orig)	0.770 ± 0.057	–	–	–	–
MICCAI 2016	UNISELF (no FLAIR)	0.715 ± 0.066	–	–	–	–
UMCL	UNISELF (orig)	0.745 ± 0.085	–	–	–	–
UMCL	UNISELF (no FLAIR)	0.668 ± 0.100	–	–	–	–
Private multisite	UNISELF	0.783 ± 0.074	–	–	–	–
Private multisite (no FLAIR)	UNISELF	0.722 ± 0.086	–	–	–	–
Private + artifacts	UNISELF	0.758–0.775 (mean ≈ 0.77)	–	–	–	–

UNISELF matches or exceeds the performance of state-of-the-art methods on in-domain tasks and substantially outperforms all benchmarks for out-of-domain datasets and challenging artifact scenarios. Competing methods, including ALL-Net and Tiramisu, show markedly inferior robustness in the presence of domain shift, missing contrast, or imaging artifacts (Zhang et al., 6 Aug 2025).

6. Generalization and Robustness Under Domain Shifts

TTIN enhances generalizability by dynamically aligning feature statistics to the input domain, removing the bias introduced by training-distribution batch statistics. This is particularly critical when domain shifts are induced by scanner variation or unavailable modalities. Contrast dropout during training only achieves distributional invariance in combination with TTIN. For this reason, UNISELF consistently avoids degradation under missing FLAIR or substantial acquisition artefacts, while baseline models fail.

The self-ensembled lesion fusion approach complements TTIN by effectively aggregating information across orientations, augmentations, and confidence levels. The two-stage detection and growth strategy mimics Canny’s hysteresis thresholding, facilitating more accurate lesion extent recovery and balancing detection sensitivity and specificity. Connected component growth from high-confidence seeds further restores lesion contiguity that is often lost with simple majority voting.

In empirical comparison, UNISELF retains performance on challenging perturbations—such as motion, noise, ghosting, and anisotropy—where artifact-based augmentations in baseline models do not confer adequate robustness. This suggests an inherent advantage of combining test-time adaptation (TTIN), extensive multi-orientation ensembling, and contrast-invariant training protocols for segmentation tasks within highly variable clinical imaging environments (Zhang et al., 6 Aug 2025).