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, generative modeling

TL;DR: We propose TIDAL, a diffusion framework that generates patient-specific counterfactual knee X-rays using temporal IPW and adversarial invariance.

Abstract: Generating realistic patient-specific counterfactual images of treatment outcomes from longitudinal medical imaging is a challenging task, 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 observed treatment effects for OA features like Kellgren-Lawrence grade. Our experiments show that TIDAL significantly outperforms baseline approaches, achieving 21.52\% reduction in image generation error and 18.43\% improvement in observed treatment effects, demonstrating significant improvements for longitudinal medical counterfactual generation.

Primary Area: generative models

Submission Number: 23918