Diffusion-based Counterfactual Augmentation: Towards Robust and Interpretable Knee Osteoarthritis Grading

2025-06-18

AI Summary

Clinical grading of knee osteoarthritis (KOA) on radiographs suffers from substantial inter-observer variability, poor robustness of deep learning models near decision boundaries, and limited interpretability. To address these bottlenecks, we propose the first counterfactual-augmented framework that explicitly couples diffusion models with clinically defined decision boundaries. Our method leverages SDE-driven latent-space navigation to generate controllable counterfactual samples in pathologically meaningful neighborhoods adjacent to boundaries. We introduce a self-correcting learning paradigm that transforms model uncertainty into structured supervision signals and impose topological constraints on the latent space to align with established KOA progression patterns. Evaluated on the OAI and MOST datasets, our approach consistently improves classification accuracy across multiple architectures, visualizes minimally detectable pathological changes, and yields clinically interpretable decision rationales. The implementation is publicly available.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Machine Learning: Calibration & Uncertainty Quantification

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Automated grading of Knee Osteoarthritis (KOA) from radiographs is challenged by significant inter-observer variability and the limited robustness of deep learning models, particularly near critical decision boundaries. To address these limitations, this paper proposes a novel framework, Diffusion-based Counterfactual Augmentation (DCA), which enhances model robustness and interpretability by generating targeted counterfactual examples. The method navigates the latent space of a diffusion model using a Stochastic Differential Equation (SDE), governed by balancing a classifier-informed boundary drive with a manifold constraint. The resulting counterfactuals are then used within a self-corrective learning strategy to improve the classifier by focusing on its specific areas of uncertainty. Extensive experiments on the public Osteoarthritis Initiative (OAI) and Multicenter Osteoarthritis Study (MOST) datasets demonstrate that this approach significantly improves classification accuracy across multiple model architectures. Furthermore, the method provides interpretability by visualizing minimal pathological changes and revealing that the learned latent space topology aligns with clinical knowledge of KOA progression. The DCA framework effectively converts model uncertainty into a robust training signal, offering a promising pathway to developing more accurate and trustworthy automated diagnostic systems. Our code is available at https://github.com/ZWang78/DCA.

Problem

• Addressing inter-observer variability in KOA grading
• Enhancing model robustness with counterfactual examples
• Improving interpretability of automated diagnostic systems

Innovation

• Uses diffusion model for counterfactual augmentation
• Employs SDE for latent space navigation
• Self-corrective learning improves classifier robustness