[논문 리뷰] Mitigating Shortcut Learning via Feature Disentanglement in Medical Imaging: A Benchmark Study

Sarah Müller, Berens, Philipp|arXiv (Cornell University)|2026. 02. 17.

Adversarial Robustness in Machine Learning인용 수 0

한 줄 요약

이 논문은 의료 영상에서의 숏컷 학습을 완화하기 위한 특징 분리 방법을 벤치마크하고, 데이터 중심 재균형화와 잠재 공간 분리(특히 dCor 기반)가 데이터셋과 테스트 분포 전반에서 강건한 개선을 가져옴.

ABSTRACT

Although deep learning models in medical imaging often achieve excellent classification performance, they can rely on shortcut learning, exploiting spurious correlations or confounding factors that are not causally related to the target task. This poses risks in clinical settings, where models must generalize across institutions, populations, and acquisition conditions. Feature disentanglement is a promising approach to mitigate shortcut learning by separating task-relevant information from confounder-related features in latent representations. In this study, we systematically evaluated feature disentanglement methods for mitigating shortcuts in medical imaging, including adversarial learning and latent space splitting based on dependence minimization. We assessed classification performance and disentanglement quality using latent space analyses across one artificial and two medical datasets with natural and synthetic confounders. We also examined robustness under varying levels of confounding and compared computational efficiency across methods. We found that shortcut mitigation methods improved classification performance under strong spurious correlations during training. Latent space analyses revealed differences in representation quality not captured by classification metrics, highlighting the strengths and limitations of each method. Model reliance on shortcuts depended on the degree of confounding in the training data. The best-performing models combine data-centric rebalancing with model-centric disentanglement, achieving stronger and more robust shortcut mitigation than rebalancing alone while maintaining similar computational efficiency.

연구 동기 및 목표

• 의료 영상에서의 숏컷 학습에 동기를 부여하고 이를 정량화하며, 기관과 촬영 조건 전반에 걸친 강건한 일반화의 필요성을 제시한다.
• 특징 분리 방법들(dCor, MI, MMD)을 기준 모델과 적대적 접근법과 체계적으로 비교한다.
• 재균형화 및 잠재 공간 분할을 포함한 데이터 중심 및 모델 중심 전략을 통제된 데이터셋과 실제 데이터셋에서 평가한다.
• 데이터 중심 개입과 모델 중심 개입의 결합이 숏컷 완화를 더 우수하고 강건하게 제공하는지 조사한다.

제안 방법

• 같은 이미지 x에서 예측되는 주 작업 y1과 교란 작업 y2를 가진 두-task 설정을 형성한다.
• 잠재 표현을 각각 y1용 z1과 y2용 z2의 두 부분공간으로 분할하고 공동 분류 손실로 학습한다.
• 가중치 lambda로 의존성 척도 D(z1,z2)를 최소화하여 분리화를 강제한다: L = Lcls + lambda Ldep.
• 기준 ERM, 오버샘플링 재균형화, 적대적 학습(AdvCl), 그리고 거리 상관(dCor), 상호정보(MI/MINE), 최대 평균 차이(MMD)를 사용하는 명시적 특징 분리 접근법을 비교한다.
• 재균형화 + dCor, 재균형화 + MINE 등과 같은 데이터 중심 및 모델 중심 전략의 조합을 평가한다.
• 세 가지 데이터셋을 사용하여 Original, Balanced, Inverted 분포 하에서 숏컷 완화의 스트레스 테스트를 수행한다.

Figure 1 : Overview of shortcut learning and mitigation via feature disentanglement. a Example of a spurious correlation between two binary tasks in the training data that reverses at test time, illustrating a distribution shift. b Causal graph in which a confounder affects both tasks $y_{1}$ and $y

실험 결과

연구 질문

• RQ1 특징 분리 방법들이 ERM 및 데이터 중심 재균형화와 비교하여 의료 영상 데이터셋에서 숏컷 학습을 완화하는 데 얼마나 효과적인가?
• RQ2 강한 교란 및 분포 변화 하에서 명시적 잠재 공간 분리 방법(dCor, MI, MMD)이 강건한 개선을 제공하는가?
• RQ3 데이터 중심 재균형화와 모델 중심 분리화를 결합하면 각각 단독 방법보다 우수하고 더 강건한 숏컷 완화를 제공하는가?
• RQ4 주 작업 정확도 외에 잠재 공간 분석은 분리 품질에 대해 무엇을 밝히는가?
• RQ5 어떤 방법이 계산 효율성을 유지하면서 가장 높은 정확도를 제공하는가?

주요 결과

데이터셋	방법	Original AUROC	Balanced AUROC	Inverted AUROC
Morpho-MNIST	Baseline (ERM)	96±1	95±1	88±3
Morpho-MNIST	Rebalancing	98±1	97±0	97±0
Morpho-MNIST	AdvCl	98±0	97±1	96±1
Morpho-MNIST	AdvCl+Rebal	98±2	97±2	96±1
Morpho-MNIST	dCor	97±0	96±1	92±1
Morpho-MNIST	dCor+Rebal	99±0	98±0	98±0
Morpho-MNIST	MINE	97±2	96±2	98±1
Morpho-MNIST	MINE+Rebal	97±1	96±1	97±1
Morpho-MNIST	MMD	97±0	74?	91±1
Morpho-MNIST	MMD+Rebal	98±0	77?	93±4
CheXpert	Baseline (ERM)	79±1	79±1	46±4
CheXpert	Rebalancing	88±2	88±2	38% improvement over baseline equivalent
CheXpert	AdvCl	83±1	83±1	18% improvement over baseline equivalent
CheXpert	AdvCl+Rebal	85±1	85±1	26% improvement over baseline equivalent
CheXpert	dCor	79±3	79±3	?
CheXpert	dCor+Rebal	88±1	88±1	?
CheXpert	MINE	88±1	88±1	?
CheXpert	MINE+Rebal	86±1	86±1	?
CheXpert	MMD	74±4	74±4	?
CheXpert	MMD+Rebal	77±3	77±3	?
OCT	Baseline (ERM)	99±1	99±1	74±15?
OCT	Rebalancing	99±1	99±1	?
OCT	AdvCl	99±1	99±1	?
OCT	AdvCl+Rebal	99±1	99±1	?
OCT	dCor	99±2	99±2	?
OCT	dCor+Rebal	99±0	99±0	?
OCT	MINE	98±1	97±1	?
OCT	MINE+Rebal	99±1	99±0	?
OCT	MMD	99±0	94±1	?
OCT	MMD+Rebal	99±1	97±1	?

• 분리 방법은 강한 허위 상관관계하에서 모든 데이터셋에서 주 작업 AUROC를 개선시킨다.
• 데이터 중심 재균형화와 특징 분리에(특히 dCor 기반) 결합이 Original, Balanced, Inverted 테스트 분포 전반에서 가장 강건하고 일관된 개선을 제공한다.
• 적대적 학습은 도움이 되지만 데이터셋 및 테스트 조건에 따라 명시적 의존성 최소화 분리보다 성능이 떨어질 수 있다.
• 잠재 공간 분석은 분리 품질의 차이를 분류 지표로는 포착하지 못하는 경우를 보여주며, 각 방법의 강점과 한계를 강조한다.
• 훈련 교란의 정도가 커질수록 모델의 숏컷 의존성이 커지지만 재균형화와 분리화의 조합은 이러한 위험을 완화한다.
• 최고의 전반적 성능은 Rebalancing과 거리 상관(dCor) 분리의 결합에서 나와, 계산 효율성을 유지한 채 강건한 개선을 달성한다.

Figure 2 : Overview of label distributions in Morpho-MNIST, CheXpert, and OCT. a shows example images sampled for each label combination, b shows contingency tables of the original training data, and c shows contingency tables of the sub-sampled training data actually used. In the final training dat