Domain Invariant Model with Graph Convolutional Network for Mammogram Classification

TL;DR

This paper introduces DIM-GCN, a framework that disentangles disease-related features from irrelevant biases in mammogram classification, improving robustness on out-of-distribution samples by leveraging graph convolutional networks and variational auto-encoders.

Contribution

The paper proposes a novel domain invariant model with graph convolutional networks that explicitly disentangles disease-related features from irrelevant biases for better OOD generalization.

Findings

Superior OOD generalization performance on mammogram diagnosis.

Effective disentanglement of disease-related and irrelevant features.

Utilization of clinical attributes via GCN enhances feature representation.

Abstract

Due to its safety-critical property, the image-based diagnosis is desired to achieve robustness on out-of-distribution (OOD) samples. A natural way towards this goal is capturing only clinically disease-related features, which is composed of macroscopic attributes (e.g., margins, shapes) and microscopic image-based features (e.g., textures) of lesion-related areas. However, such disease-related features are often interweaved with data-dependent (but disease irrelevant) biases during learning, disabling the OOD generalization. To resolve this problem, we propose a novel framework, namely Domain Invariant Model with Graph Convolutional Network (DIM-GCN), which only exploits invariant disease-related features from multiple domains. Specifically, we first propose a Bayesian network, which explicitly decomposes the latent variables into disease-related and other disease-irrelevant parts that are provable to be disentangled from each other. Guided by this, we reformulate the objective function based on Variational Auto-Encoder, in which the encoder in each domain has two branches: the domain-independent and -dependent ones, which respectively encode disease-related and -irrelevant features. To better capture the macroscopic features, we leverage the observed clinical attributes as a goal for reconstruction, via Graph Convolutional Network (GCN). Finally, we only implement the disease-related features for prediction. The effectiveness and utility of our method are demonstrated by the superior OOD generalization performance over others on mammogram benign/malignant diagnosis.