Visual Disentangled Diffusion Autoencoders: Scalable Counterfactual Generation for Foundation Models

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.

💡 Research Summary

The paper tackles the persistent problem that large foundation models (FMs) such as CLIP, despite impressive zero‑shot capabilities, still rely on spurious correlations (“Clever Hans” behavior). Existing mitigation techniques either need group annotations (e.g., Group‑DRO) or depend on costly gradient‑based adversarial optimization to produce visual counterfactual explanations (VCEs). To overcome these limitations, the authors introduce Visual Disentangled Diffusion Autoencoders (DiDAE), a two‑stage, gradient‑free framework that couples a frozen FM encoder with a disentangled dictionary and a conditional diffusion decoder.

In the first stage, the FM Φ maps an image x to a dense semantic embedding z_sem = Φ(x). Rather than manipulating z_sem directly, DiDAE learns an invertible dictionary Ω that decomposes the embedding into interpretable semantic components c = Ω(z_sem). Ω can be obtained either via supervised alignment (e.g., Orthogonal Procrustes using known attribute labels) or via unsupervised methods (e.g., sparse autoencoders, SVD). Each component corresponds to a human‑understandable concept such as “blond hair”, “male”, or “makeup”.

The second stage uses a diffusion autoencoder decoder D_θ. The decoder receives the edited embedding z’_sem = Ω⁻¹(c′) together with a stochastic spatial code x_T (preserving high‑frequency details) and produces a counterfactual image ˆx = D_θ(z’_sem, x_T) in a single forward pass. Because the FM encoder remains frozen, DiDAE inherits its robust semantic manifold while the decoder learns only to reconstruct images from (z_sem, x_T) pairs.

Two concrete algorithms are presented for counterfactual generation:

1. Component Reflection (Algorithm 1) – For a chosen component k, the coefficient c_k is simply negated (c′_k = –c_k) while all other coefficients stay unchanged. This creates a clean, disentangled edit that flips the presence/absence of a single attribute without affecting others.

2. Distilled Boundary Inversion (Algorithm 2) – A downstream linear classifier f is distilled into a linear probe wᵀz. For each component direction v_k, an analytic scalar α is solved such that wᵀ(z_sem + α v_k) = –wᵀz_sem, guaranteeing that the decision score is exactly inverted. This yields counterfactuals that cross the classifier’s decision boundary while altering only the targeted semantic direction.

Using these tools, the authors propose two model‑correction strategies:

• Projection – Once a spurious direction d_spur is identified, embeddings are orthogonally projected onto the complement of d_spur (z_robust = z – (z·d_spur) d_spur). This removes the bias without any fine‑tuning of the FM.

• Counterfactual Knowledge Distillation (CFKD) with DiDAE (DiDAE‑CFKD) – DiDAE automatically generates a large set of labeled counterfactual images for each discovered component. By clustering components once (the “pre‑clustered teacher”), the labeling effort scales with the number of components rather than with the product of downstream tasks, models, and counterfactuals. The student model (e.g., a ResNet‑18) is then retrained on the augmented dataset, forcing it to rely on causal features.

Experiments are conducted on two heavily poisoned benchmarks: a synthetic Square dataset (spurious background intensity) and CelebA‑Blond (spurious gender correlation). A 98 % poisoning ratio ensures that naive models learn the shortcut. DiDAE‑CFKD dramatically outperforms gradient‑based baselines (ACE, DIME, Diff‑ICE) in both speed (average generation time ≈0.12 s vs. >3 s) and downstream accuracy (12–18 percentage‑point gain on balanced test sets). Qualitatively, DiDAE produces disentangled edits—e.g., toggling “blond hair” while leaving “male” unchanged—whereas ACE often entangles hair and eyebrows or injects adversarial noise.

Key contributions are: (i) a gradient‑free, interpretable manipulation of frozen FM embeddings via a disentangled dictionary; (ii) a diffusion‑based decoder that yields high‑fidelity, diverse counterfactuals in a single pass; (iii) scalable integration with CFKD that reduces labeling overhead and effectively mitigates shortcut learning. The work bridges recent interpretability research on foundation model latent spaces (e.g., sparse autoencoders) with practical counterfactual generation, opening avenues for fast, large‑scale bias correction in vision models. Future directions include extending the dictionary to non‑linear manifolds, applying the framework to multimodal FM (e.g., CLIP‑like text‑image models), and building interactive tools for end‑users to explore and edit model behavior in real time.