Invariant Risk Minimization in Medical Imaging with Modular Data Representation

Abstract

Despite the effectiveness of deep neural networks trained with Empirical Risk Minimization (ERM) in medical imaging tasks, these models often exhibit performance degradation when faced with Out-of-Distribution (OoD) data, owing to potential biases in their predictive accuracy. Invariant Risk Minimization (IRM) seeks to rectify this issue by identifying invariant or causal correlations across various environments. However, its practical application does not consistently deliver the expected generalization performance in real-world scenarios. This paper addresses a potential limitation of the IRM framework, positing that the constraints enforced by IRM might not sufficiently guide the model in learning all causal features. In response, we propose a novel methodology leveraging modular neural networks within the IRM framework. Our approach aims to generate more diverse data representations, thereby enhancing the generalization performance of models trained with IRM. Experimental validation on three tasks — two medical image classification tasks, namely, Camelyon17-wilds and CheXpert, and a synthetic task, Colored MNIST — demonstrates significant improvements in generalization performance in both OoD settings and subpopulation shift cases.

Overview

IRM이 지배적인 invariant feature만 학습하는 한계를 modular neural network로 극복하여, 의료 영상에서의 OoD 일반화를 개선한다.

1. Modular encoder — 데이터 표현 모델을 \(N\) 개 모듈로 분할하여 각각이 서로 다른 invariant feature를 학습하도록 유도한다.
2. Competitive selection — Multi-head dot product attention으로 입력에 가장 관련 있는 \(k\) 개 모듈을 선택한다.
3. IRM optimization — 선택된 모듈의 가중 표현으로 IRM 목표를 최적화하여 다양한 invariant feature를 활용한 OoD 일반화를 달성한다.

제안 방법의 구조도. Modular data representation을 IRM 프레임워크 내에 통합한다.

Method

IRM은 invariant/causal 상관관계를 학습하려 하지만, 실제로는 가장 지배적인(dominant) invariant feature만 인코딩하는 한계가 있다. 우리는 modular neural network를 IRM 프레임워크에 통합하여, 각 모듈이 서로 다른 invariant feature를 학습하도록 유도한다.

• 데이터 표현 모델 \(\Phi\) 를 \(N\) 개의 모듈 \(\{f_n\}_{n=1}^N\) 으로 분할
• Multi-head dot product attention으로 모듈 간 competitive learning 수행
• 입력 자체가 query, 모듈 출력이 key/value로 작동
• Top-\(k\) 모듈을 선택하되, 비선택 모듈도 soft selection으로 유지 (module collapse 방지)

Camelyon17-wilds와 CheXpert 데이터셋의 환경별 예시 이미지.

Results

Colored MNIST

Camelyon17-wilds (OoD Medical Imaging)

CheXpert (Subpopulation Shift)

Camelyon17-wilds에서 ERM/IRM 대비 OoD 테스트 정확도 10% 향상. CheXpert에서 worst-case 정확도 9.4% 향상.

BibTeX

@inproceedings{bae2024invariant,
  author    = {Bae, Jun-Hyun and Kim, Chanwoo and Chang, Tae-Young},
  title     = {Invariant Risk Minimization in Medical Imaging with Modular Data Representation},
  booktitle = {International Conference on Electronics, Information, and Communication (ICEIC)},
  year      = {2024},
  doi       = {10.1109/ICEIC61013.2024.10457174}
}