cs.CVSep 22, 2026
SaveDecoupling Disease, Covariates, and Individual Variability: A Unified Disentanglement Framework for Medical Image Classification
Organizations: Department of Radiology & Faculty of Medical Imaging Technology, Ruijin Hospital, College of Health Science and Technology, Shanghai Jiao Tong University School of Medicine, Shanghai 200025, China · Clinical Neuroscience Center, Ruijin Hospital, Shanghai Jiao Tong University School of Medicine, Shanghai 200025, China · School of Data Science, Fudan University, Shanghai 200433, China · Institute of Science and Technology for Brain-Inspired Intelligence, Fudan University, Shanghai 200433, China · Department of Radiology, Shandong Provincial Hospital Affiliated to Shandong First Medical University, Jinan 250031, China · Department of Clinical Psychology, School of Medical Psychology, Air Force Medical University, Xi’an 710032, Shaanxi, China
Abstract
Accurately isolating disease-related features from confounding covariates (e.g., age, gender, site) and individual variations remains a fundamental challenge in medical image classification. Traditional regression-based approaches may ignore non-linear relations between image features and true covariates. To overcome this issue, we present a generalized Medical Imaging Disentanglement Learning (MedIDL) framework. MedIDL maps image features into three mutually orthogonal latent spaces through specialized disentanglement heads: a disease classification head guided by a supervised loss, a covariate-alignment head constrained by cross-subject similarity matching, and a Gaussian head absorbing individual variations. We evaluated our framework across 7 datasets encompassing diverse imaging modalities. MedIDL outperforms state-of-the-art supervised and self-supervised classification methods in accuracy across all datasets. Association analyses demonstrate that MedIDL successfully isolates target-specific latent representations. Gradient-based interpretability mappings localize pathognomonic patterns aligning with established clinical literature.
Figures & tables
Fig. 1: Overview of the datasets, covariate association, and advantage of the proposed framework. (a) illustrates the dataset description, showing the sample sizes across diverse modalities. (b) visualizes the most significantly associated biomarker(s) for a covariate (e.g. age or scanner site) within the HCs of a dataset. The image biomarkers are functional/structural connectivity for fMRI/DTI, intensity for sMRI and X-ray. The -values were false discovery rate (FDR)-corrected in each dataset. (c) conceptually illustrates the latent representation spaces, highlighting that conventional methods typically enforce linear relationships between the extracted covariate features and true covariates, whereas our framework allows for non-linear manifold alignment.
Fig. 2: Overall architecture of the proposed disentangled representation learning framework. For any given input modality (e.g., 3D images, brain graphs, or 2D images), the data is processed by an encoder and subsequently projected into three separate latent spaces via dedicated heads. The disease head extracts disease-specific features ( ) optimized for separating patients from HCs. The individual head captures normalized, individual-specific latent variations ( ). The covariate head extracts demographic and site-related features ( ), which are supervised by aligning the channel-wise subject similarity matrices with the corresponding ground truth similarity matrices from the covariates (e.g., age, gender, and site). Finally, the three disentangled components are concatenated and fed into a decoder to reconstruct the original input.
| fMRI DTI ADHD SCZ Presbycusis PPMI SCZ Method Accuracy a AUC Accuracy AUC Accuracy OvR-AUC Accuracy AUC Accuracy AUC BrainGNN 62.22 3.87 63.04 4.42 64.86 3.57 64.48 2.73 83.09 4.41 84.02 4.51 67.54 3.85 69.02 4.17 68.61 3.26 69.22 3.40 NEGAT 62.08 2.94 62.95 3.35 63.33 4.56 64.09 3.91 82.79 3.85 82.95 3.36 69.31 3.52 70.53 3.92 68.25 3.82 68.72 3.91 GraphCL 62.86 4.71 62.21 4.27 65.31 4.72 65.19 4.40 83.45 4.49 82.50 3.92 70.65 4.58 72.36 4.20 68.29 3.41 68.88 3.62 JOAO 61.92 3.72 62.12 3.84 64.92 5.11 64.80 4.63 83.15 3.77 82.69 4.43 70.28 3.59 71.92 3.47 68.37 4.10 69.30 3.86 LaGraph 63.88 2.68 63.92 2.63 65.25 3.83 64.46 4.30 83.85 4.02 83.19 3.46 71.58 3.72 72.20 3.91 69.24 3.69 70.11 4.03 AGCL 63.25 4.18 62.63 4.29 67.25 4.18 65.35 3.64 84.15 4.17 83.55 3.86 72.48 4.25 72.85 3.74 69.92 3.88 70.16 4.05 GATE 62.25 3.84 62.71 4.03 66.28 3.71 66.05 4.12 83.97 3.47 83.27 2.49 72.18 3.62 72.55 4.10 68.33 3.92 69.81 3.60 BrainGCL 62.05 3.59 61.52 3.70 65.27 2.97 64.31 3.59 82.85 2.98 83.06 4.14 72.11 4.13 72.46 3.86 68.25 4.05 69.03 3.72 DMG 62.57 2.95 62.08 3.51 67.14 3.95 68.52 4.41 84.08 3.19 83.95 4.39 72.65 3.50 73.07 4.03 70.09 4.14 71.58 3.95 Ours 65.24 3.17 b 66.03 4.05 70.02 4.06 71.05 4.57 86.22 3.36 87.05 5.04 74.46 3.52 74.91 4.10 73.42 4.06 74.08 4.02 |
TABLE I: Classification performance on fMRI (ADHD-200, SCZ, Presbycusis) and DTI (PPMI, SCZ) connectomes.
| 3D MRI 2D X-ray ADNI Lung Method Accuracy a AUC Accuracy AUC v-CNN 80.53 4.15 81.14 3.58 82.35 2.30 83.05 1.96 ResNet 81.32 3.72 80.88 4.05 82.95 2.35 83.52 2.06 ViT 81.59 3.90 82.22 3.74 83.48 2.02 84.15 1.65 GF-Net 83.07 4.16 83.25 3.10 85.05 1.59 85.68 2.15 V-Mamba 82.29 3.90 84.02 4.07 85.28 1.94 85.75 2.23 DiffMed-v2 84.46 4.18 84.55 3.29 86.26 2.06 86.81 1.45 SimCLR 83.88 2.95 84.09 3.34 84.49 1.45 84.50 1.59 ADIOS 83.58 3.52 83.86 3.75 85.15 2.29 85.42 1.88 SGLA-Net 82.09 3.85 83.11 3.79 84.30 1.92 84.85 2.01 Ours 85.71 3.52 b 86.48 4.83 87.60 3.42 87.91 2.48 |
TABLE II: Classification performance on 2D/3D vision tasks.
Fig. 3: Quantitative validation of feature disentanglement using Kendall’s Tau rank correlation across seven datasets. Each row represents a ground-truth covariate (age, gender, or site) for evaluating the correlation. A dot within a subplot denotes the mean correlation over all the subjects in a test set, where the correlation for a subject measures the order consistency between similarities from this subject to all the others measured by a distinct feature (by color) and similarities measured by the true covariate.
Fig. 4: Comprehensive ablation and sensitivity analysis of the MedIDL framework across diverse imaging modalities. (Row 1) Impact of the unified latent embedding dimension ( ) on classification accuracy. Shaded regions indicate standard deviations. The dashed line denotes the top competing method. (Row 2) Robustness validation across various modality-aware backbones. The dashed line denotes the top competing method. (Row 3) Component-wise ablation of the tri-branch disentanglement module. (Row 4) Joint hyperparameter sensitivity heatmaps for the covariate alignment weight ( ) and latent reconstruction weight ( ), suggesting a stable, broad high-accuracy zone (darker red regions). (Row 5) Superiority of our implicit latent disentanglement (solid bars) over traditional input-level de-confounding via linear regression. Dark hatched bars represent explicit de-confounding (regressing out covariates prior to classification), while light hatched bars denote classification without covariate removal.
Fig. 5: Quantitative validation of biomarker importance via counterfactual perturbation analysis across seven datasets. A plot depicts the degradation trajectories of the model’s mean prediction probability as input features are progressively perturbed/masked (replaced by means from HCs) from the patients in the test sets. Solid red lines represent the targeted perturbation of the “top important features” identified by our framework. Dashed blue lines with shades serve as the baseline representing the mean and standard deviation from 5 “random perturbations”. To accommodate the distinct spatial topologies of different modalities, the -axis denotes the number of ROIs perturbed for fMRI and DTI, 3D voxels masked for sMRI, or 2D pixels masked for X-ray.
Fig. 6: Saliency maps identified by the proposed framework across seven datasets. (Top): For graph-based brain networks, the framework maps the most discriminative ROIs in the MNI space, where the marker size scales proportionally with the node’s importance score. (Bottom): For voxel/pixel-based datasets, gradient-based feature attribution maps localize structural and radiological anomalies. On ADNI (sMRI), the gradient is with respect to T1 intensity after affine registration. Hence, positive values in the cortical region indicate gray matter atrophy; negative values in the ventricle indicate enlarged ventricle. On Lung (X-ray), red rectangles are pulmonary lesions annotated by a clinical expert.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
| Dataset & Modality | Class | Sites | Subjects | Gender (F/M) | Age ( ) | Scanner Hardware | Acquisition Parameters |
| ADHD-200 ( fMRI ) | ADHD HC | 4 | 275 205 | 79/196 105/100 | 11.5 2.8 12.1 3.2 | Siemens/Philips 3T | TR = 2000/2500 ms, TE = 15–30 ms, FA = 75 ∘ /90 ∘ , voxel size 3 mm 3 |
| SCZ ( fMRI ) | SCZ HC | 3 | 137 190 | 62/75 92/98 | 24.8 7.5 24.9 8.2 | 3.0T Siemens Magnetom Trio Tim scanner | TR = 2 s, TE = 30 ms, flip angle = 90 ∘ , FOV = 220 220 mm 2 |
| Presbycusis ( fMRI ) | Presbycusis N-PTA HC | 1 | 130 154 112 | 55/75 70/84 48/64 | 63.2 3.3 63.6 4.6 63.1 3.2 | 3.0T Philips Achieva MR scanner | TR = 2 s, TE = 35 ms, FOV = 240 240 mm 2 , resolution = 3.75 3.75 mm 2 , 35 slices, thickness = 4 mm, 240 vols |
| PPMI ( DTI ) | PD HC | 15 | 152 41 | 103/49 28/13 | 61.3 9.1 59.5 10.8 | 3.0T Siemens Tim Trio scanner | TR/TE = 900/88 ms, FA = 90 ∘ , 72 slices, = 1000 s/mm 2 , 64 diffusion dirs, voxel size = 2 2 2 mm 3 , Cardiac-triggered |
| SCZ ( DTI ) | SCZ HC | 3 | 137 190 | 62/75 92/98 | 24.8 7.5 24.9 8.2 | 3.0T Siemens Magnetom Trio Tim scanner | = 1000 s/mm 2 , 32 dirs, resolution = 2 2 2 mm 3 |
| ADNI ( sMRI ) | MCI HC | 4 | 201 238 | 101/100 125/113 | 76.5 7.0 73.7 8.0 | Multi-site 1.5T/3T | T1-weighted MPRAGE, isotropic 1 mm 3 voxels |
TABLE I: Demographic, clinical, and detailed imaging acquisition characteristics across the seven datasets.