Learning Causal Alignment for Reliable Disease Diagnosis

TL;DR

This paper introduces a causality-based framework for aligning machine learning decision processes with radiologists' reasoning in disease diagnosis, improving reliability and transferability of medical AI models.

Contribution

It proposes a novel causal alignment loss using counterfactuals and implicit function theorem to align model decision chains with expert reasoning.

Findings

Effective alignment with radiologists demonstrated on medical diagnosis tasks

Improved robustness and interpretability of diagnosis models

Causal alignment outperforms associational methods in transferability

Abstract

Aligning the decision-making process of machine learning algorithms with that of experienced radiologists is crucial for reliable diagnosis. While existing methods have attempted to align their diagnosis behaviors to those of radiologists reflected in the training data, this alignment is primarily associational rather than causal, resulting in pseudo-correlations that may not transfer well. In this paper, we propose a causality-based alignment framework towards aligning the model's decision process with that of experts. Specifically, we first employ counterfactual generation to identify the causal chain of model decisions. To align this causal chain with that of experts, we propose a causal alignment loss that enforces the model to focus on causal factors underlying each decision step in the whole causal chain. To optimize this loss that involves the counterfactual generator as an implicit function of the model's parameters, we employ the implicit function theorem equipped with the conjugate gradient method for efficient estimation. We demonstrate the effectiveness of our method on two medical diagnosis applications, showcasing faithful alignment to radiologists.