Fast Diffusion-Based Counterfactuals for Shortcut Removal and Generation

Nina Weng*1, Paraskevas Pegios*1,2, Eike Petersen1, Aasa Feragen1,2, Siavash Bigdeli1
1 Technical University of Denmark, 2 Pioneer Centre for AI

ECCV 2024 (ORAL)
*Equal Contribution
arXiv ECCV 2024 Paper (coming soon) ECCV 2024 Supplementary (coming soon) Code CelebA Code Medical
MY ALT TEXT

Our proposed method FastDiME generates counterfactuals by synthetically removing or adding shortcut features, ensuring high image quality for accurate model assessment. Unlike standard XAI approaches which require visual inspection, our approach allows us to automatically validate the role of the shortcut features in the model's decision.

Abstract

Shortcut learning is when a model -- e.g. a cardiac disease classifier -- exploits correlations between the target label and a spurious shortcut feature, e.g. a pacemaker, to predict the target label based on the shortcut rather than real discriminative features. This is common in medical imaging, where treatment and clinical annotations correlate with disease labels, making them easy shortcuts to predict disease. We propose a novel detection and quantification of the impact of potential shortcut features via a fast diffusion-based counterfactual image generation that can synthetically remove or add shortcuts. Via a novel self-optimized masking scheme we spatially limit the changes made with no extra inference step, encouraging the removal of spatially constrained shortcut features while ensuring that the shortcut-free counterfactuals preserve their remaining image features to a high degree. Using these, we assess how shortcut features influence model predictions. This is enabled by our second contribution: An efficient diffusion-based counterfactual explanation method with significant inference speed-up at comparable image quality as state-of-the-art. We confirm this on two large chest X-ray datasets, a skin lesion dataset, and CelebA.

Efficient Gradient Estimation for Classifier Guided Diffusion

MY ALT TEXT

Convergence of counterfactual generation using DiME, GMD, and our method FastDiME in a 2D toy example. DiME uses an expensive nested loop of guided DDPM to get a clean input with quadratic complexity, and GMD uses one-step denoised input but requires backpropagation through the denoiser, making it computationally and memory intensive. Our method FastDiME approximates gradients efficiently with linear complexity, avoiding heavy backpropagation.



Proposed Method FastDiME

MY ALT TEXT

FastDiME serves dual purposes: (i) as a counterfactual explanation method for explaining classifiers' decisions and (ii) as a generation tool for our shortcut detection pipeline. At each step of the diffusion process, a noised image is sampled with the guidance of the counterfactual loss, leveraging information derived from the denoised image. This also enables the automatic extraction and application of a self-optimized mask, encouraging localized changes of shortcut features and preventing changes in regions less relevant to the task.



Proposed Shortcut Learning Detection Pipeline

MY ALT TEXT

We propose detecting shortcut learning using counterfactuals. Datasets with varying shortcut-label correlations are constructed, which are used to train a task classifier suspected of shortcut learning. The degree of shortcut learning is assessed by comparing confidence levels between original images and their FastDiME shortcut counterfactuals. Significant prediction differences, with the target task-label remaining consistent, indicate shortcut learning.



Counterfactual Explanations









Shortcut Detection with Counterfactuals

MY ALT TEXT

Our proposed pipeline allows us to detect and quantify shortcut learning in practice as the severity of shortcut learning correlates with the strength of the association in the training set.



Acknowledgements

Work on this project was partially funded by the Independent Research Fund Denmark (DFF, grant number 9131-00097B), the Pioneer Centre for AI (DNRF grant nr P1), the DIREC project EXPLAIN-ME (9142-00001B), and the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF20OC0062606). The funding agencies had no influence on the writing of the manuscript nor on the decision to submit it for publication.

arXiv Preprint BibTeX

@article{weng2023fast,
        title={Fast diffusion-based counterfactuals for shortcut removal and generation},
        author={Weng, Nina and Pegios, Paraskevas and Feragen, Aasa and Petersen, Eike and Bigdeli, Siavash},
        journal={arXiv preprint arXiv:2312.14223},
        year={2023}
      }