Visual Disentangled Diffusion Autoencoders: Scalable Counterfactual Generation for Foundation Models

TL;DR

The paper addresses the vulnerability of foundation models to spurious correlations by introducing Visual Disentangled Diffusion Autoencoders (DiDAE), a gradient-free framework that decomposes frozen model embeddings into interpretable semantic directions via a disentangled dictionary and decodes edits with a diffusion autoencoder. It enables scalable, disentangled counterfactual generation and, when combined with Counterfactual Knowledge Distillation (CFKD), achieves state-of-the-art mitigation of shortcut learning on synthetic and real-world benchmarks. Through two gradient-free analysis methods and two correction strategies (Projection and CFKD), DiDAE delivers fast counterfactuals (up to 64 per second) and substantial downstream gains, outperforming gradient-based baselines and several explainer approaches. The approach generalizes across domains and offers potential extensions to discrete modalities and more powerful generative backbones like latent diffusion models.

Abstract

Foundation models, despite their robust zero-shot capabilities, remain vulnerable to spurious correlations and 'Clever Hans' strategies. Existing mitigation methods often rely on unavailable group labels or computationally expensive gradient-based adversarial optimization. To address these limitations, we propose Visual Disentangled Diffusion Autoencoders (DiDAE), a novel framework integrating frozen foundation models with disentangled dictionary learning for efficient, gradient-free counterfactual generation directly for the foundation model. DiDAE first edits foundation model embeddings in interpretable disentangled directions of the disentangled dictionary and then decodes them via a diffusion autoencoder. This allows the generation of multiple diverse, disentangled counterfactuals for each factual, much faster than existing baselines, which generate single entangled counterfactuals. When paired with Counterfactual Knowledge Distillation, DiDAE-CFKD achieves state-of-the-art performance in mitigating shortcut learning, improving downstream performance on unbalanced datasets.

Figures (3)

• Figure 1: Comparison of traditional gradient-based counterfactuals (Left) versus the proposed DiDAE approach (Right) on a CelebA classifier trained on the "Blond Hair" label. The label is spuriously correlated with "Heavy Makeup" and "Attractive" and anti-correlated with "Male." Traditional methods require slow, iterative gradient updates through the diffusion process, often resulting in adversarial noise or entangled changes (e.g., changing hair color and eyebrows simultaneously). In contrast, DiDAE utilizes a frozen foundation model $\Phi$ to decompose embeddings into disentangled semantic components $\mu$. Counterfactuals are generated via simple linear reflection in this semantic space, followed by decoding via a diffusion decoder. This gradient-free approach is significantly faster and produces disentangled edits (e.g., toggling "Blond Hair" without affecting other attributes). On the left one can see an example for a counterfactual generated with ACE, which actually flipped the Blond Hair confounding Male classifiers decision, but the counterfactual seems to be primarily adversarial in nature. Furthermore, one can see in pixel-wise intensity difference map right to the counterfactual that the systematic changes made where distributed over the hair and the eyebrows. One should note that ACE would produce better results for features with a smaller pixel-foodprint (like smiling) and if the explained classifier would not rely on spuriously entangled features.
• Figure 2: Visualizations of 4 example dimensions of our foundation models with DiDAE. For Square, the components correspond to the 4 latent dimensions (foreground, background, X position, Y position). For CelebA, the components reveal different attribute dimensions correlated with the "Male" attribute.
• Figure 3: Visualizations of counterfactuals for the first 4 SVD dimensions found for the Square dataset in its foundation model space. One can see that Comp1 clearly corresponds to the foreground color and Comp4 clearly corresponds to the background color. Comp2 and Comp3 appear to be related to the x- and y-position, which were not disentangled perfectly. However, there is also no unique solution how the x- and y-axis could be disentangled, e.g. $\mathcal{B}_1 = \{ (0,1), (1,0) \}$ and $\mathcal{B}_2 = \{ (1,1), (-1,1) \}$ are both correct orthgonal bases.