Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning Causal Alignment for Reliable Disease Diagnosis

Mingzhou Liu, Ching-Wen Lee, Xinwei Sun, Yu Qiao, Yizhou Wang

TL;DR

The paper tackles the problem of aligning machine learning–based medical diagnoses with radiologists’ causal decision pathways to avoid spurious correlations. It introduces a causal alignment framework that uses counterfactual generation to identify decision-relevant regions and a causal alignment loss to enforce alignment with radiologist explanations, optimized with the implicit function theorem and conjugate gradient to handle the implicit counterfactual solver. A hierarchical extension leverages attribute annotations via CCCE-based causal attribution, enabling alignment at both attribute and image-region levels. Experiments on lung nodule and breast mass datasets show improved CAM precision and malignancy accuracy, indicating faithful, transferable alignment to expert reasoning. The work advances trustworthy medical AI by grounding model decisions in clinician-driven causality rather than superficial correlations.

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.

Learning Causal Alignment for Reliable Disease Diagnosis

TL;DR

The paper tackles the problem of aligning machine learning–based medical diagnoses with radiologists’ causal decision pathways to avoid spurious correlations. It introduces a causal alignment framework that uses counterfactual generation to identify decision-relevant regions and a causal alignment loss to enforce alignment with radiologist explanations, optimized with the implicit function theorem and conjugate gradient to handle the implicit counterfactual solver. A hierarchical extension leverages attribute annotations via CCCE-based causal attribution, enabling alignment at both attribute and image-region levels. Experiments on lung nodule and breast mass datasets show improved CAM precision and malignancy accuracy, indicating faithful, transferable alignment to expert reasoning. The work advances trustworthy medical AI by grounding model decisions in clinician-driven causality rather than superficial correlations.

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.
Paper Structure (20 sections, 3 theorems, 28 equations, 9 figures, 4 tables, 1 algorithm)

This paper contains 20 sections, 3 theorems, 28 equations, 9 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related works
  3. Problem setup & Background
  4. Methodology
  5. Casual alignment loss
  6. Optimization
  7. Hierarchical alignment
  8. Experiment
  9. Experimental setup
  10. Comparison with baselines
  11. Ablation study
  12. Visualization
  13. Conclusion and discussion
  14. Causal alignment theory
  15. Network architectures
  16. ...and 5 more sections

Key Result

Theorem 4.1

Consider two vectors $x, \theta$, and a differentiable function $T(x,\theta)$. Let $x^* := \arg\min_x T(x,\theta)$. Suppose that: i) the argmin is unique for each $\theta$, and ii) the Hessian matrix $H$ is invertible. Then $x^*(\theta)$ is a continuous function of $\theta$. Further, the Jacobian ma

Figures (9)

  • Figure 1: (a) Two decision chains with different causal structures but present similar correlation patterns. The left chain "$\text{mass } (X) \!\to\! \text{attributes } (A) \!\to\! \text{label } (Y)$" aligns with radiologists, while the right chain "$\text{shortcut} \!\to\! \text{label } (Y)$" is misaligned. However, both $(X,Y)$ and $(\text{shortcut}, Y)$ are correlated, due to the confounding bias between the shortcut and $Y$. (b) Our approach for learning causally aligned models. We first identify features and attributes that causally influence the model's decision, then align them to that of radiologists in a hierarchical fashion. (c) Class Activation Mapping (CAM) visualization and comparison of CAM precision on lung cancer diagnosis.
  • Figure 2: The schematic overview of our method. (a) We adopt a hierarchical structure that first provides attribute descriptions for the image, and then shows the diagnosis result. (b) Training with the proposed alignment loss. In the forward pass, a counterfactual image $x^*$ is generated and used to compute the alignment loss $\mathcal{L}_{\text{align}}$ relative to the expert's annotation $m$. In the backward pass, we use an implicit gradient solver to obtain the gradient $\nabla_\theta \mathcal{L}_{\text{align}}$ and use it to update the parameter $\theta$.
  • Figure 3: Causal diagram of radiologists' decision process. $A$ and $Y$ denote the expert's annotations of the attributes and the decision label, respectively.
  • Figure 4: CAM visualization. Each row denotes different cases. The first column is the input images, where nodules and masses are marked by red bounding boxes. The second to seventh columns are CAMs of compared baselines and our method, respectively. See Appx. \ref{['appx.more-visu']} for more results.
  • Figure 5: Network architecture used in lung nodule classification.
  • ...and 4 more figures

Theorems & Definitions (6)

  • Theorem 4.1: Implicit Function Theorem (IFT), krantz2002implicit
  • Remark A.3
  • Proposition A.4
  • proof
  • Proposition A.5
  • proof