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

arXiv: https://arxiv.org/abs/2506.15748 · HTML: https://mlsanity.com/paper/2506.15748/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 43 references that are in the map.
- First posted on arXiv 18 Jun 2025.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [A Survey Of Methods For Explaining Black Box Models](https://mlsanity.com/paper/1802.01933/) (2018)
- [Explaining Machine Learning Classifiers through Diverse Counterfactual Explanations](https://mlsanity.com/paper/1905.07697/) (2019)
- [Understanding disentangling in β-VAE](https://mlsanity.com/paper/1804.03599/) (2018)
- [Adversarially Robust Generalization Requires More Data](https://mlsanity.com/paper/1804.11285/) (2018)
- [GAN-based Synthetic Medical Image Augmentation for increased CNN Performance in Liver Lesion Classification](https://mlsanity.com/paper/1803.01229/) (2018)
- [Diffusion Models for Counterfactual Explanations](https://mlsanity.com/paper/2203.15636/) (2022)
- [GANterfactual - Counterfactual Explanations for Medical Non-Experts using Generative Adversarial Learning](https://mlsanity.com/paper/2012.11905/) (2020)
- [Automatic Knee Osteoarthritis Diagnosis from Plain Radiographs: A Deep Learning-Based Approach](https://mlsanity.com/paper/1710.10589/) (2017)
- [Automatic Grading of Individual Knee Osteoarthritis Features in Plain Radiographs using Deep Convolutional Neural Networks](https://mlsanity.com/paper/1907.08020/) (2019)
- [Quantifying Radiographic Knee Osteoarthritis Severity using Deep Convolutional Neural Networks](https://mlsanity.com/paper/1609.02469/) (2016)
- [Automatic Detection of Knee Joints and Quantification of Knee Osteoarthritis Severity using Convolutional Neural Networks](https://mlsanity.com/paper/1703.09856/) (2017)
- [Mitigating attribute amplification in counterfactual image generation](https://mlsanity.com/paper/2403.09422/) (2024)
- [Counterfactual Explanations for Medical Image Classification and Regression using Diffusion Autoencoder](https://mlsanity.com/paper/2408.01571/) (2024)
- [Understanding Deep Learning via Decision Boundary](https://mlsanity.com/paper/2206.01515/) (2022)
- [Transformer with Selective Shuffled Position Embedding and Key-Patch Exchange Strategy for Early Detection of Knee Osteoarthritis](https://mlsanity.com/paper/2304.08364/) (2023)
- [Temporal Evolution of Knee Osteoarthritis: A Diffusion-based Morphing Model for X-ray Medical Image Synthesis](https://mlsanity.com/paper/2408.00891/) (2024)

## Abstract (arXiv)

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.