Causal Representation Learning with Observational Grouping for CXR Classification

2025-06-25

AI Summary

This study addresses the unidentifiability of causal representations and the limited generalization and robustness of models in medical imaging (e.g., chest X-rays). We propose an observation-grouping–based identifiable causal representation learning framework. Methodologically, we construct observation groups stratified by sensitive attributes—such as race, sex, and imaging view—and enforce invariance of causal features across these groups within an end-to-end unsupervised learning paradigm, thereby enabling latent disentanglement and causal structure discovery. Our key contributions are: (i) the first integration of observation grouping into causal representation learning for medical images, ensuring identifiability of causal variables; and (ii) explicit modeling of invariance under multiple confounding sources, enhancing interpretability and fairness. Experiments demonstrate substantial improvements in cross-domain generalization across multiple disease classification tasks and effective mitigation of bias induced by sensitive attributes, establishing a novel paradigm for trustworthy AI-assisted diagnosis.

Technology Category

• Machine Learning: Causal Learning
• Reasoning under Uncertainty: Causality
• Computer Vision: Representation Learning for Vision

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

Identifiable causal representation learning seeks to uncover the true causal relationships underlying a data generation process. In medical imaging, this presents opportunities to improve the generalisability and robustness of task-specific latent features. This work introduces the concept of grouping observations to learn identifiable representations for disease classification in chest X-rays via an end-to-end framework. Our experiments demonstrate that these causal representations improve generalisability and robustness across multiple classification tasks when grouping is used to enforce invariance w.r.t race, sex, and imaging views.

Problem

• Identifies causal relationships in chest X-ray data generation
• Improves generalisability of disease classification features
• Enhances robustness using grouping by race, sex, views

Innovation

• Grouping observations for identifiable causal learning
• End-to-end framework for CXR disease classification
• Enforcing invariance across race, sex, and views