MIMM-X: Disentangling Spurious Correlations for Medical Image Analysis

AI summary

Overview

• Research area: Causal representation learning for medical image analysis (computer vision), specifically mitigating shortcut learning from multiple spurious correlations.
• Technical level: Intermediate. Readers need some familiarity with deep learning training objectives, mutual information estimation, and classification benchmarks.
• Scope: The paper introduces MIMM-X, an extension of the Mutual Information Minimization Model (MIMM) that disentangles a primary clinical prediction task from multiple spurious factors at once, and validates it on three datasets (UK Biobank, NAKO, CheXpert) across MRI and X-ray.

What This Paper Is About

Deep learning models trained on medical images often latch onto dataset-specific quirks, such as a patient's sex, age, or which scanner a scan came from, rather than the actual anatomical evidence of disease. This "shortcut learning" produces models that look accurate in the training cohort but collapse when the mix of patients changes, which is dangerous in clinical use. The authors build a framework that suppresses these shortcuts by forcing the model's task-specific features to carry no information about several spurious factors simultaneously.

Key Contributions

1. A multi-confounder extension of MIMM. MIMM-X generalizes the earlier MIMM method from a single known spurious correlation to N spurious factors, estimating one mutual information value between the primary-task feature subvector and the stacked spurious-feature subvectors to keep computation efficient.
2. Two new architectural mechanisms. A Confounder Attention Weighter (CAW) that learns softmax-normalized attention weights over the N spurious feature subvectors, and a GradNorm-inspired dynamic loss scaling (DLS) scheme that adjusts each task's loss contribution, with the primary task's emphasis increasing across epochs.
3. Validation across three datasets and two modalities. Experiments use brain MRI from NAKO and UK Biobank and chest X-ray from CheXpert, covering induced spurious correlations (sex, age, dataset origin) and naturally occurring ones (age, co-occurring lung diseases).
4. A disentanglement-specific evaluation protocol. Beyond classification accuracy on validation, inverted, and balanced test sets, the authors measure cross-prediction accuracy, testing whether the primary-task features can predict the confounders and whether the confounder features can predict the primary task, where ideal disentanglement means near-random-guess performance. Code is publicly released at https://github.com/lab-midas/MIMM-X.

Main Findings

• Experiment 1 (brain MRI, primary task: age group young <51 / old >57 years; confounders: sex, NAKO/UKB dataset): MIMM-X with both CAW and DLS achieved 84.2% on the validation set, 82.8% on the inverted set, and 82.6% on the balanced set, a 1.4% drop across distributions and the highest inverted-set accuracy (82.8%).
• Baseline and dependence-correlation methods collapsed under distribution shift. Baseline dropped from 96.3% to 27.2% (a -61.6% drop) and the distance correlation (dCor) variant dropped from 96.2% to 41.7% (-54.5%). The original single-confounder MIMM fell below random guessing on the inverted set (74.4% validation to 45.7% inverted).
• Dataset rebalancing was the strongest competitor in Experiment 1 with only a -9.9% drop (91.8% to 81.9%), but it required a 224.3% increase in training sample size.
• Ablation of CAW and DLS: MIMM-X with DLS only reached 85.6% validation / 74.2% inverted; MIMM-X with CAW only reached 46.3% validation / 54.0% inverted (the paper states the CAW-without-DLS configuration gave the highest validation accuracy but dropped by over -60% on the inverted distribution); MIMM-X with both components gave 84.2% / 82.8% / 82.6%. The authors attribute DLS's role to preventing easily learned spurious tasks from dominating the primary task signal.
• Experiment 2 (chest X-ray, primary task: pleural effusion; confounders: sex, age young <50 / old >60): MIMM-X reached 86.3% validation, 70.4% inverted, and 76.5% balanced, the highest inverted accuracy among all compared methods. All competitors except rebalancing sat near 50% on the inverted set.
• Experiment 3 (chest X-ray, one induced confounder plus two natural ones): With sex induced as the spurious factor and age and presence of any of 12 co-occurring lung diseases (CoD) left at their natural distributions, MIMM-X achieved 81.7% validation, 70.2% inverted, and 72.9% balanced. Rebalancing reached a similar 69.4% inverted, while other methods fell below 60% on the primary task.
• Disentanglement measured by cross-prediction: In Experiment 1, predicting sex and dataset from the primary feature subvector gave 51.8% and 50.0% for MIMM-X and 47.6% / 48.1% for rebalancing, versus 70.3% / 70.0% for the baseline and 70.1% / 67.4% for dCor. In Experiment 2, MIMM-X gave 54.8% / 54.3%, and in Experiment 3 it gave 52.6% / 53.1% / 50.9% for sex, age, and CoD, the closest to random guessing among the compared methods.
• t-SNE visualizations (Figure 3): Primary-task features from MIMM-X showed no clear separation when colored by sex or by dataset origin, whereas reference methods formed visible clusters, indicating residual shortcut information.
• Confounder prediction stays easy for every method. Accuracy for predicting the spurious factors themselves remained high across methods and datasets, which the authors note is expected because those tasks are inherently less complex to learn than the primary clinical task.

Methodology in Plain English

The image is passed through a feature encoder that splits its output into N+1 equal-sized feature subvectors: one reserved for the primary task and one for each of N spurious factors. Separate classification heads read these subvectors, with a multi-task linear head handling the spurious factors. Before those spurious features reach their head, the Confounder Attention Weighter learns a weight per spurious factor so the model can emphasize whichever one is currently more relevant. A mutual information estimator (a MINE model) then measures how much information the primary-task features share with the stacked spurious features, and training penalizes that shared information.

Training alternates between updating the encoder, the classification heads, and the attention weighter on one batch, and running N_B−1 updates of the mutual information estimator. The overall loss combines cross-entropy terms for the primary task and each spurious task with the mutual information penalty weighted by lambda, where the cross-entropy terms are dynamically rescaled based on how large each task's loss is relative to the mean loss, and the primary task's scaling exponent grows across epochs.

To create testable shortcuts, the authors sub-sample their training data so that 90% of samples in each primary-task class come from a specific class of each spurious factor. Models are then evaluated on a validation set carrying the same correlations, an inverted set where the correlations are reversed, and a balanced set with no correlation. A model that has genuinely learned the causal signal should barely lose accuracy when moving from validation to inverted data.

Why This Matters

The work targets a failure mode with direct clinical consequences: a model that uses patient sex, age, or scanner type as a proxy for disease will misjudge patients who do not fit the training cohort's demographics. MIMM-X addresses this without altering the underlying data distribution, which the authors contrast with rebalancing approaches that require larger training sets and careful dataset design.

Real-world applications:

• Cross-site deployment of diagnostic models. Performance is tested explicitly on data combining two cohorts (NAKO and UK Biobank), mimicking the shift a model faces when moved between hospitals or scanners.
• Disease detection from chest X-ray. Pleural effusion prediction from CheXpert serves as a proxy for routine radiology triage where sex and age are strong historical correlates of disease labels.
• Age prediction from brain MRI. Age-group classification from structural MRI is a common neuroimaging benchmark where cohort identity and sex are known confounds.
• Auditing models for demographic bias. The cross-prediction disentanglement protocol offers a concrete test of whether a model's internal representations still encode sensitive attributes.

Industry relevance: the method requires no change to data collection or resampling infrastructure, adding only an auxiliary classification head and a mutual information estimator, which matters for teams that cannot easily rebalance large proprietary clinical datasets. The authors report all experiments ran on a single NVIDIA GeForce RTX 3090 GPU.

Future Directions

• Discovering unknown confounders automatically. The authors state that MIMM-X currently relies on specifying confounding factors in advance and identify automatic discovery as future work.
• Scaling beyond three confounders. The paper explicitly lists handling more than three spurious factors and more complex confounding structures as unexplored.
• Better mitigation of rebalancing's cost. The authors suggest more advanced sampling strategies could reduce the training-set inflation (224.3% in Experiment 1) that rebalancing incurs, though such strategies would still depend on knowing the spurious attributes.
• Applying this to strongly non-homogeneous cohorts. The paper notes that careful dataset design for rebalancing may not always be achievable in such cohorts, positioning distribution-preserving methods as an alternative worth pursuing.

Target Audience

Researchers and practitioners working on robust or causal deep learning for medical imaging, fairness and bias mitigation in clinical machine learning, and domain generalization. It is also relevant to clinical ML engineers evaluating whether a model's apparent accuracy reflects genuine disease signal or dataset artifacts, and to readers already familiar with the single-confounder MIMM method who want to see how it extends to realistic multi-confounder settings.

Authors’ abstract

Deep learning models can excel on medical tasks, yet often experience spurious correlations, known as shortcut learning, leading to poor generalization in new environments. Particularly in medical imaging, where multiple spurious correlations can coexist, misclassifications can have severe consequences. We propose MIMM-X, a framework that disentangles causal features from multiple spurious correlations by minimizing their mutual information. It enables predictions based on true underlying causal relationships rather than dataset-specific shortcuts. We evaluate MIMM-X on three datasets (UK Biobank, NAKO, CheXpert) across two imaging modalities (MRI and X-ray). Results demonstrate that MIMM-X effectively mitigates shortcut learning of multiple spurious correlations.