Disentangled diffusion autoencoders speed counterfactual edits in visual models
Towards Fast and Disentangled Counterfactuals for Visual Foundation Models
Machine LearningComputer Vision and Pattern Recognition
Summary
Computer vision models sometimes make decisions based on tricks or errors hidden inside their reasoning. The authors introduce a new method called Disentangled Diffusion Autoencoders (DiDAE) that helps find simple, clear changes to images which reveal these hidden tricks. This method is much faster than previous techniques because it doesn't rely on slow gradient calculations. It can also improve other models by teaching them to avoid these errors, making AI decisions more reliable.
What this means in practice
- •For computer vision engineers: Quickly generate meaningful image edits that expose and fix decision flaws in pretrained vision models.
- •For machine learning operations teams: Improve deployed classifiers by automatically identifying and removing spurious features without needing metadata.
Authors
Sidney Bender, Benedikt Kunz, Ahmed Zeid, Shinichi Nakajima, Klaus-Robert Müller, Marco Morik
Abstract
Foundation models remain vulnerable to spurious correlations and ``Clever Hans'' strategies. Explainable machine learning can find and remove such strategies for classifiers without metadata. For foundation models, no such option exists yet. We propose Disentangled Diffusion Autoencoders (DiDAE). DiDAE wraps a frozen foundation model in a conditional diffusion decoder. A counterfactual is one closed-form edit along a direction of a disentangled dictionary, followed by decoding. The dictionary can be supervised (Procrustes) or unsupervised (Singular Value Decomposition, Sparse Autoencoders). No gradients are needed, so DiDAE is up to 2000 times faster than the state of the art. We evaluate on six datasets, two synthetic and four real-world. In a desiderata-driven benchmark on three of them, its counterfactuals are on par with or better than the state of the art, and they repair downstream classifiers through Counterfactual Knowledge Distillation (CFKD), where they beat metadata-based correction. The same machinery can rank a pretrained dictionary against a trained classifier. It returns the few directions the classifier actually reads, each causally verified by a counterfactual that flips the decision, and repairs the classifier along those a teacher marks spurious. The workflow is plug-and-play in our open-source Peal library we publish alongside the paper. With a public dictionary and a pretrained decoder, all that remains is a cheap linear distillation of the classifier and its own fine-tuning.