eess.IVMar 4, 2026

Self-Supervised Multi-Stage Domain Unlearning for White-Matter Lesion Segmentation

Authors: Domen Preložnik, Žiga Špiclin

Organizations: University of Ljubljana, Faculty of Electrical Engineering, Slovenia

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Inter-scanner variability of magnetic resonance imaging has an adverse impact on the diagnostic and prognostic quality of the scans and necessitates the development of models robust to domain shift inflicted by the unseen scanner data. Review of recent advances in domain adaptation showed that efficacy of strategies involving modifications or constraints on the latent space appears to be contingent upon the level and/or depth of supervision during model training. In this paper, we therefore propose an unsupervised domain adaptation technique based on self-supervised multi-stage unlearning (SSMSU). Building upon the state-of-the-art segmentation framework nnU-Net, we employ deep supervision at deep encoder stages using domain classifier unlearning, applied sequentially across the deep stages to suppress domain-related latent features. Following self-configurable approach of the nnU-Net, the auxiliary feedback loop implements a self-supervised backpropagation schedule for the unlearning process, since continuous unlearning was found to have a detrimental effect on the main segmentation task. Experiments were carried out on four public datasets for benchmarking white-matter lesion segmentation methods. Five benchmark models and/or strategies, covering passive to active unsupervised domain adaptation, were tested. In comparison, the SSMSU demonstrated the advantage of unlearning by enhancing lesion sensitivity and limiting false detections, which resulted in higher overall segmentation quality in terms of segmentation overlap and relative lesion volume error. The proposed model inputs only the FLAIR modality, which simplifies preprocessing pipelines, eliminates the need for inter-modality registration errors and harmonization, which can introduce variability. Source code is available on https://github.com/Pubec/nnunetv2-unlearning.

Figures & tables

Figure 1. Comparison of MRI scans between different domains.
Source / SubsetVendorFSAMTrainTest
WMH / SeenPhilips3T2D2030
Siemens3T2D2030
WMH / UnseenPhilips3T3D 10
GE3T3D 30
GE1.5T3D 10
MSSEG / UnseenPhilips3T3D 15
Table 1. Dataset information and train / test splits as used in experimental evaluation.
Figure 2. (a) MRI preprocessing and (b) self-supervised multi-stage unlearning with nnU-Net.
Algorithm 1: Multi-Stage Self-Supervised Unlearning.
ee←\leftarrow Current epoch
ete_{t}←\leftarrow Total epochs
ewe_{w}←\leftarrow Warm-up epochs
ExE_{x}←\leftarrow Encoder block x={1,…,6}x=\{1,\ldots,6\}
CxC_{x}←\leftarrow Classifier x={1,…,6}x=\{1,\ldots,6\}
PatPat←\leftarrow Patience (on learning steps)
Table 4
Figure 3. Ablation study metric values with respect to learn-to-unlearn ratio (LUR=1-1, 2-1, 3-1, 4-1, 5-1), for various unlearning stages: A - from bottleneck, B - bottom three blocks, C - top three blocks, D - entire encoder .
Unlearning stages LURUnseen data - not used for model training: Philips 3T 3D, GE 3T 3D, GE 1.5T 3D
123456DSC ↑\uparrowTPR ↑\uparrowLTPR ↑\uparrowLFDR ↓\downarrowRVE ↓\downarrowRS ↓\downarrow
     ✓1 - 10.719 ± 0.1000.794 ± 0.1160.717 ± 0.0970.251 ± 0.1130.262 ± 0.2354.2
   ✓✓✓2 - 10.705 ± 0.1030.767 ± 0.1370.716 ± 0.1070.243 ± 0.1150.241 ± 0.2015.0
✓✓✓✓✓✓3 - 10.642 ± 0.1320.856 ± 0.0960.741 ± 0.1040.347 ± 0.1850.763 ± 0.5896.4
   ✓✓✓5 - 10.674 ± 0.1200.832 ± 0.1010.813 ± 0.0950.474 ± 0.1850.527 ± 0.4495.2
     ✓Self-superv.0.712 ± 0.0990.805 ± 0.1150.759 ± 0.1030.294 ± 0.1360.308 ± 0.2724.4
Table 2. Ablation study results on the 2017 WMH Segmentation Challenge’s unseen test split. Values are mean ±\pm standard deviation, while bold indicates best result for particular evaluation metric. LUR is the ratio of the number of learning vs. unlearning steps.
Model Classifier accuracy, AcciAcc^{i}
Unlearning stagesStepi=1i=2i=3i=4i=5i=6
nnU-Net* 100%100%100%100%100%100%
nnU-Net* SSMSUPost warm-up100%100%100%100%100%99%
Unlearning stages: 1 & 2 & 3 (Top 3)Post unlearning85%77%80%98%99%91%
nnU-Net* SSMSUPost warm-up100%99%99%99%99%99%
Unlearning stages: 4 & 5 & 6 (Bottom 3)Post unlearning100%100%100%67%48%50%
Table 3. Domain classifier accuracy, post warm-up and post-unlearning, per stage of the nnU-Net SSMSU on 2017 WMH dataset’s seen train split, comparing top- and bottom-three and unlearning of all stages. Values in bold indicate stages actively unlearned.
Figure 4. DSC with 95% confidence intervals for MSSEG 2016 unseen test set. For notes (*,**) see caption of Table 4 .
Figure 5. DSC with 95% confidence intervals per source domain of MSSEG unseen test set for nnU-Net SSMSU.
Figure 6. Axial cross-sections of MSSEG 2016 MR scans with superimposed manual ( top row ) and second best default nnU-Net ( middle row ) and the proposed SSMSU model segmentations ( bottom row ). Arrows highlight important differences.
WMH - seen test dataset
ModelDSC ↑\uparrowTPR ↑\uparrowLTPR ↑\uparrowLFDR ↓\downarrowRVE ↓\downarrowRS ↓\downarrow
nnU-Net No DA0.750 ± 0.120 †0.775 ± 0.132 †0.680 ± 0.160 †0.390 ± 0.160 †0.220 ± 0.3533.2
nnU-Net*0.773 ± 0.1080.785 ± 0.129 †0.679 ± 0.103 †0.210 ± 0.1310.205 ± 0.3882.0
nnU-Net* Preprocessing**0.732 ± 0.123 †0.748 ± 0.149 †0.630 ± 0.105 †0.219 ± 0.137 †0.232 ± 0.3713.8
Dinsdale et al. ( Dinsdale et al., 2021 )0.724 ± 0.149 †0.793 ± 0.123 †0.624 ± 0.108 †0.426 ± 0.187 †0.304 ± 0.4774.4
nnU-Net SSMSU (ours)0.772 ± 0.0970.810 ± 0.1240.718 ± 0.0970.262 ± 0.1600.197 ± 0.2281.6
Table 4. Segmentation evaluation on seen (WMH) and three unseen (MSSEG, MSLJ, ISBI) datasets, using respective seen and unseen test splits (Table 1 ). Values are mean ±\pm standard deviation, while bold indicates best results per each performance metric for methods in direct comparison.
Figure 7. Axial cross-sections of six MSSEG 2016 MR scans ( columns ) for the nnU-Net baseline ( top row ) and the proposed nnU-Net SSMSU ( bottom row ) with superimposed true positives ( green ), false positives ( blue ) and false negatives ( red ) ( bottom row ). Red frame indicates best result according to the DSC metric. Overall, nnU-Net baseline was best in 5, while the proposed nnU-Net SSMSU was best in 47 out of 51 test cases.
Figure 8. Lesion segmentation performance according to lesion size, i.e. small (SM), medium (MED) and large (LG); obtained on the MSSEG dataset.
More in this topic

Unsupervised Domain Adaptation

Unsupervised domain adaptation transfers models from labeled source data to unlabeled target domains with different data distributions.

See all 183 papers

Explore similar work

CardsList
Apr 1, 2026

Improving Generalization of Deep Learning for Brain Metastases Segmentation Across Institutions

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.

Yuchen Yang, Shuangyang Zhong, Haijun Yu +4

Institute of Medical Technology, Peking University Health Science Center, Beijing, China · Beijing Key Laboratory of Intelligent Neuromodulation and Brain Disorder Treatment, Beijing, China · Department of Radiation Oncology, University Hospital Erlangen, Erlangen, Germany

Medical Image SegmentationBrain Tumor SegmentationMulti-Source Domain AdaptationUnsupervised Domain AdaptationVariational Autoencoders
AskPDF
Apr 1, 2026

Automated Detection of Multiple Sclerosis Lesions on 7-tesla MRI Using U-net and Transformer-based Segmentation

Ultra-high field 7-tesla (7T) MRI improves visualization of multiple sclerosis (MS) white matter lesions (WML) but differs sufficiently in contrast and artifacts from 1.5-3T imaging - suggesting that widely used automated segmentation tools may not translate directly. We analyzed 7T FLAIR scans and generated reference WML masks from Lesion Segmentation Tool (LST) outputs followed by expert manual revision. As external comparators, we applied LST-LPA and the more recent LST-AI ensemble, both originally developed on lower-field data. We then trained 3D UNETR and SegFormer transformer-based models on 7T FLAIR at multiple resolutions (0.5x0.5x0.5^3, 1.0x1.0x1.0^3, and 1.5x1.5x2.0^3) and evaluated all methods using voxel-wise and lesion-wise metrics from the BraTS 2023 framework. On the held-out test set at native 0.5x0.5x0.5^3 resolution, 7T-trained transformers achieved competitive overlap with LST-AI while recovering additional small lesions that were missed by classical methods, at the cost of some boundary variability and occasional artifact-related false positives. On a held-out 7 T test set, our best transformer model (SegFormer) achieved a voxel-wise Dice of 0.61 and lesion-wise Dice of 0.20, improving on the classical LST-LPA tool (Dice 0.39, lesion-wise Dice 0.02). Performance decreased for models trained on downsampled images, underscoring the value of native 7T resolution for small-lesion detection. By releasing our 7T-trained models, we aim to provide a reproducible, ready-to-use resource for automated lesion quantification in ultra-high field MS research (https://github.com/maynord/7T-MS-lesion-segmentation).

Michael Maynord, Minghui Liu, Cornelia Fermüller +4

Department of Computer, Science University of Maryland College Park, College Park, 20742, MD, USA · Department of Neurology, University of Maryland School of Medicine, Baltimore, MD, 21201, USA · Baltimore VA Medical Center, VA Maryland Healthcare System, Baltimore, MD, 21201, USA

Medical Image Segmentation3D Medical Image SegmentationBrain MRISemantic SegmentationMagnetic Resonance Imaging
AskPDF
Apr 9, 2026

Component-Adaptive and Lesion-Level Supervision for Improved Small Structure Segmentation in Brain MRI

Small lesions in brain MRI are hard to segment because they occupy a tiny fraction of the volume and are dominated by background and larger lesions during voxel-wise optimization, so a model can reach a high Dice similarity coefficient (DSC) while missing many of them. We propose CATMIL, a training objective that adds two auxiliary terms to the standard nnU-Net Dice and cross-entropy loss without changing the architecture. The Component-Adaptive Tversky (CAT) term weights lesion voxels by the inverse size of their connected component, so each lesion contributes nearly equally regardless of volume. The lesion-level Multiple Instance Learning (MIL) term treats each lesion as a bag of voxels and penalizes lesions with no detected voxel. For multiple sclerosis lesion segmentation on MSLesSeg, CATMIL achieves the highest small-lesion recall (0.873 vs. 0.796 for Dice+CE; 95% CI of the difference +0.030 to +0.157, higher in all six test patients) and about 48% fewer missed lesions, with comparable DSC and HD95. The gain holds for lesions of at least 3 mm in diameter, the clinical reading size (recall 0.944 vs. 0.870). Standard losses produce no probability response to most small lesions they miss, so no threshold can recover them. The cost is more small false-positive components; a simple component-size filter removes most of them while keeping the sensitivity gain, and at matched lesion-wise precision CATMIL detects more small lesions with higher lesion-wise F1. An ablation attributes the detection gain to the MIL term. On a second dataset, 3D-MR-MS, CATMIL with the same loss weights again improves small-lesion recall, at a larger false-positive cost and slightly lower DSC. Code: https://github.com/luumsk/SmallLesionMRI

Minh Sao Khue Luu, Evgeniy N. Pavlovskiy, Bair N. Tuchinov

The Artificial Intelligence Research Center of Novosibirsk State University, Novosibirsk 630090, Russia

Medical Image SegmentationRobust Loss Functions
AskPDF