TIDAL: A Temporal Causal Diffusion Framework for Visualizing Knee Osteoarthritis Treatment Outcomes

Anthony Cuturrufo, Fred Xu, Arvind Murari Vepa, David R. Bakal, Veena K. Ranganath, Yizhou Sun, Wei Wang

Keywords: computer vision, medical imaging

TL;DR: TIDAL is a novel causal diffusion framework that generates counterfactual medical images to mitigate confounding bias in longitudinal knee osteoarthritis treatment outcome prediction.

Abstract: Generating realistic patient-specific counterfactual images of treatment outcomes from longitudinal medical imaging is a challenging task, often complicated by confounding and selection bias in observational datasets. To address this challenge, we propose TIDAL (Temporal IPW Diffusion Adversarial Learning), a novel longitudinal causal diffusion framework that integrates causal inference techniques directly into diffusion model training. TIDAL utilizes a Stable Diffusion backbone conditioned on patient history and incorporates two key causal adaptations: (1) Temporal Inverse Propensity Weighting (IPW) that reweights the diffusion loss based on treatment propensity scores; and (2) Domain Adversarial Training that encourages treatment-invariant representations. We demonstrate TIDAL's effectiveness by simulating knee osteoarthritis (OA) progression with longitudinal X-Rays from the Osteoarthritis Initiative (OAI). Performance is assessed using image fidelity metrics and causally-relevant Individual Treatment Effect (ITE) metrics for OA features like Kellgren-Lawrence grade. Our experiments show that TIDAL outperforms baseline approaches, achieving 21.52% reduction in image generation error and 18.43% improvement in causal validity, demonstrating significant improvements for longitudinal medical counterfactual generation.

Submission Number: 112