Towards Fast and Disentangled Counterfactuals for Visual Foundation Models

1 INTRODUCTION

• Foundation models encode non-causal artifacts, but existing mitigation often needs confounder labels and counterfactual explainers are a speed and quality bottleneck.Gradient-based methods can be slow, noisy, and entangled; gradient-free alternatives often lack explicit semantic sparsification and diversification.Source: §1 · p. 1
• DiDAE edits one learned semantic component at a time, producing diverse, disentangled counterfactuals without gradient updates.Its dictionary can be supervised with Procrustes or unsupervised with SVD or sparse autoencoders.Source: §1 · p. 1
• DiDAE counterfactuals support CFKD, where a pre-clustered teacher makes one decision per direction rather than per counterfactual.The same framework ranks pretrained dictionary directions by verified classifier flips and supports correction along directions marked spurious.Source: §1 · p. 1
• In the ImageNet freight-car example, verified counterfactual ranking surfaces rail tracks as a shortcut and one CFKD iteration raises average group accuracy from 96.4% to 97.2%.The example uses a public dictionary and pretrained decoder; neither is trained for the task.Source: Figure 1 · p. 2
• The evaluation spans synthetic, natural, and medical imaging datasets and reports counterfactual-based downstream correction at up to 2000 times the speed.The workflow is available in the open-source PEAL library, which fetches pretrained decoders and public dictionaries so users need not train a generator.Source: §1 · pp. 1–2

Spurious Correlations and Model Correction

• Existing correction methods often depend on confounder annotations, while CFKD uses counterfactual augmentation but has been bottlenecked by counterfactual speed and quality.DiDAE-CFKD reduces teacher interaction to one decision per direction rather than one per counterfactual.Source: §2 · pp. 2–3

Visual Counterfactual Explainers (VCEs)

• Gradient-based visual explainers can be slow and constrained in feature sparsification, while gradient-free methods are faster but typically lack semantic sparsification and diversification.DiDAE adds dictionary-based edits, a trust region, classifier linkage, and verification to the diffusion-decoder building block.Source: §2 · pp. 3–4, Figure 2 · p. 3

Interpretability of Foundation Models

• Sparse-autoencoder concept galleries leave unclear what correlated or entangled concepts mean and which dictionary atoms a particular classifier actually uses.DiDAE renders controlled component edits and ranks atoms by counterfactuals that flip the classifier under study.Source: §2 · p. 4

3 METHODS

• The pipeline uses a frozen encoder, a dictionary and component encoder, a decoder, and an inversion; decoder and dictionary can be reused across classifiers or supplied by the domain.For natural images, public CLIP dictionaries and a matching decoder can avoid user-side generative-model training.Source: §3 · pp. 4–5

3.1 Conditional Decoder

• The conditional decoder is trained on image–representation pairs while the foundation encoder stays frozen, so it renders the representation the foundation model itself computes.Any conditional diffusion or flow-matching model qualifies if it conditions on the representation and admits inversion.Source: §3.1 · p. 5

3.2 Disentangled Dictionaries

• The dictionary may align annotated concepts by closed-form Procrustes, use SVD directions without annotations, or use an overcomplete sparse autoencoder.For sparse autoencoders, an atom-wise edit changes one coefficient while leaving the other coefficients and reconstruction residual untouched.Source: §3.2 · p. 5

3.3 From Direction to Edit

• Edits move a component to its empirical range bound, making the coefficient step data-bounded but not guaranteeing that the full representation remains on the data manifold.For classifier-aware edits, a distilled linear probe selects the direction toward the decision boundary; pair edits can move two atoms when one cannot cross within its trust region.Source: §3.3 · pp. 5–6

3.4 Inversion and Counterfactual Decoding

• Inversion supplies a code that reconstructs the input, and decoding it under edited conditioning transports the representation edit into the image while retaining image-specific information.Edit-friendly DDPM inversion reconstructs exactly by construction and is the default for quantitative results; DDIM inversion is deterministic but lossy.Source: §3.4 · p. 6

3.5 Disentangled Diffusion Autoencoders and Direction Ranking

• Ranking uses sparse regression along dictionary components to identify candidate atoms, then orders directions by decoded counterfactuals that flip the real classifier and move the intended component.The output is classifier-specific and interventional, and teacher-marked directions can feed the CFKD correction step.Source: §3.5 · p. 6

3.6 DiDAE-CFKD

• DiDAE-CFKD filters counterfactuals for student flips and in-distribution status, then uses a teacher’s labels to augment training and fine-tune the classifier.Automatic metadata labeling or a pre-clustered teacher can transfer a direction-level label across its counterfactuals and downstream models sharing the encoder.Source: §3.6 · p. 7

4 EXPERIMENTAL SETUP

• Table 1 maps each dataset to its confounder, the instances used in the four experimental slots, and the appendix sections with details.The setup section summarizes datasets and protocols; full splits, model instances, thresholds, metrics, and hyperparameters are in the appendices and tables.Source: §4 · p. 7

4.1 Datasets

• Six datasets cover counterfactual quality and correction on Square, CelebA-Blond, and Camelyon17, plus dictionary ranking on Sparse Numbers, NICO++, and two ImageNet class pairs.Correction training splits are poisoned so the confounder predicts the label, while the test split is balanced over class and confounder.Source: §4.1 · p. 7

4.2 Model Instances

• Synthetic datasets use a ResNet-18 encoder, while real images use CLIP or PLIP; dictionaries are fitted or reused without fitting them to a student.Quantitative results use edit-friendly DDPM inversion and an empirical trust region; correction targets include scratch-trained ResNet-18 students and frozen-feature probes.Source: §4.2 · p. 7

4.3 Protocols

• Counterfactual quality is assessed first by Gain on held-out balanced data, then by sample and trajectory inspection, and lastly by desiderata scores.Correction uses CFKD and compares with metadata-based baselines; ranking reports latent, ambient, and verified flips, with an oracle or human teacher marking directions.Source: §4.3 · pp. 7–8

5.1 Desiderata and Efficiency

• 35.5 ± 2.3 Gain on CelebA-Blond exceeds SCE’s 23.4, and 35.0 ± 10.0 on Camelyon17 leads within seed variance, while Square trails SCE, 77.0 to 90.6.Gain measures downstream accuracy on held-out, balanced data.Source: Table 2 · p. 9, §5.1 · p. 11
• 47.8 counterfactuals per second on Square contrasts with 0.02–0.03 for gradient-based explainers; each DiDAE counterfactual uses a closed-form edit, inversion, and decoder rollout.The paper attributes the speed difference to avoiding hundreds of backward passes.Source: §5.1 · p. 9, Table 2 · p. 9

Does a counterfactual change one concept at a time?

• DiDAE edits follow individual dictionary axes, unlike DAE’s diagonal or coincident attempts; the separated directions support per-cluster teacher feedback.On Square, DAE does not cross the boundary; on CelebA, it mixes factors in one direction.Source: Figure 4 · p. 10, Does a counterfactual change one… · p. 11

Do I need concept annotations to build the dictionary?

• Without annotations, SVD can recover foreground and background color on Square but may rotate spatial axes; a sparse dictionary separates superposed digits on Sparse Numbers.On the ResNet-18 student, Procrustes and SVD are within a point of Gain.Source: Do I need concept annotations to build… · p. 12

5.4 Downstream Model Correction

• DiDAE-CFKD beats all four metadata-based baselines, reaching 80.1 Gain on Square and 16.3 on NICO++ without group labels.The Square result uses a pre-clustered teacher; the NICO++ result uses Ranking DiDAE repair.Source: Table 4 · p. 10, §5.4 · p. 12

Could I just project the spurious direction out instead?

• Projection reaches 53.4 Gain on Square and 22.6 on CelebA-Blond, below counterfactual-based probe correction at 84.7 and 70.6, respectively.The paper attributes the difference to concepts being spread across redundant foundation-model directions.Source: Could I just project the spurious… · p. 12

5.5 Ranking DiDAE

• Ranking surfaces shortcuts across four settings, including tracks and graffiti for freight car and context directions for NICO++; counterfactuals show what those directions change.On the fireboat probe, the water jet ranks first, while the fireworks flips are borderline.Source: Figure 6 · p. 11, §5.5 · pp. 12–13
• CFKD on teacher-labeled spurious directions raises average group accuracy from 74% to 87% on Sparse Numbers, 57% to 64% on NICO++, 96.4% to 97.2% on freight car, and 95.5% to 98.1% on fireboat.On ImageNet, average group accuracy covers four class-by-confounder groups.Source: §5.5 · p. 13

6 CONCLUSION

• DiDAE wraps a frozen foundation model in a conditional diffusion decoder and edits along dictionary directions in closed form rather than using iterative gradient search.Its dictionary may be supervised, unsupervised, or overcomplete and sparse without changing the edit operator.Source: §6 · p. 13
• Ranking mode surfaced rail tracks as a previously unreported shortcut on a natural ImageNet probe, and one CFKD iteration reduced it.The mode returns directions a classifier causally reads, verified by counterfactual flips.Source: §6 · p. 13
• Up to three orders of magnitude faster, DiDAE enables CFKD-scale correction, with 77.0 ± 2.2 Gain on Square and 35.5 ± 2.3 Gain on CelebA-Blond; SCE remains ahead on Square and some desiderata columns.The paper says its case rests first on Gain, then on random samples showing disentangled, on-manifold edits, and last on desiderata, which are necessary but not sufficient.Source: §6 · p. 13
• Decoding and reclassifying counterfactual edits reorders dictionary directions, especially for the largest dictionary used in SAE interpretability.Source: §6 · p. 13

7 LIMITATIONS

• DiDAE inherits the quality of its dictionary and generator; reproducible edits are not interpretable when atoms lack human-recognizable concepts.Verification confirms that the intended coefficient moved, but does not establish that the atom's concept is recognizable.Source: §7 · p. 13
• On Square, SVD components mix spatial axes because variance, rather than semantics, defines them.Source: §7 · p. 14
• ImageNet yields fewer latent flips that survive rendering than CelebA or Square, and many surviving edits are not minimal; the proposed manifold explanation is untested.The authors attribute the difficulty to ImageNet spanning multiple visual modes, but state that they have not tested this hypothesis.Source: §7 · p. 14

8 FUTURE WORK

• Future directions include embedding-domain extensions with invertible decoders, foundation-model repair from ranked labeled pairs, and blinded human evaluation of unfiltered counterfactuals.Source: §8 · p. 14

APPENDIX A DESIDERATA OF COUNTERFACTUALS

• The evaluation measures four explanation desiderata and downstream Gain; construction overhead is difficult to quantify, and common realism and proximity scores are omitted.The authors note that distributional scores assess image sets, while proximity can reward diffuse edits that are less informative.Source: Appendix A · p. 16

A.1 Sufficiency

• NAFR requires both a target-model flip and a transferred flip on an independently trained surrogate, while diversity assesses differences between two counterfactuals for one factual.NAFR counts factuals rather than attempts; the evaluation metrics are marginal and can miss whether most inputs receive two orthogonal, on-manifold edits.Source: §A.1 · p. 16, Appendix A · p. 16

A.2 Understandability

• Understandability is measured by Hoyer sparsity of the oracle-representation change, with 1 for a perfectly sparse edit and 0 for a dense transformation.In the unfiltered CelebA-Blond samples, DiDAE usually realizes the intended semantic edit photographically, but a single edit may not flip a classifier that relies on both dimensions.Source: §A.2 · pp. 16, 18, Figure 9 · p. 17

A.3 Fidelity

• Fidelity proxies include surrogate transfer as Non-Adversarial Rate and an unbiasedness score that compares edits to dominant versus non-dominant features.The fidelity desideratum includes manifold realism, neighborhood stability, and alignment with the local decision-boundary normal, which cannot be checked directly in high dimensions with an unknown manifold.Source: §A.3 · p. 18

A.4 Efficiency

• Efficiency is reported as counterfactuals per second, measured using the largest batch that fits on one NVIDIA A100 with 80 GB.This is per-counterfactual cost; one-time construction overhead is described as comparable across the methods discussed.Source: §A.4 · p. 18

A.5 Gain

• Gain measures the percentage of the performance gap closed between the original poisoned-data model and the corrected student, using average group accuracy.The appendix cautions that Gain concerns a specific downstream task, where a few reasonable CFKD samples may sometimes suffice.Source: §A.5 · p. 18

APPENDIX B COUNTERFACTUAL KNOWLEDGE DISTILLATION (CFKD)

• CFKD has a teacher label generated counterfactuals as true or false; false examples are added with their original labels, then the student is fine-tuned on the augmented data.The teacher may be a human, an oracle, or a scalable pre-clustered approach.Source: Appendix B · pp. 18–19

• The appendix expands Section 4 with the datasets and their splits used in the experiments.Source: Appendix C · p. 19
• It describes the instances filling the method’s four slots and the students being corrected.Source: Appendix C · p. 19
• It gives the four protocols underlying the reported results in Section 5.Source: Appendix C · p. 19

C.1 Datasets

• For the ImageNet fireboat–lifeboat task, training examples are filtered by whether MSAE water-jet atom #5717 is active, while validation and test retain natural-split images.The natural fireboat–lifeboat probe reaches 99.6–100% held-out accuracy, so the shortcut is made load-bearing for the experiment.Source: §C.1 · p. 19
• Square and CelebA-Blond use 98% poisoning, Camelyon17 uses 100%, and these three are tested on 1,000 samples balanced by class and spurious attribute.Sparse Numbers, NICO++, and ImageNet use their own balanced splits.Source: §C.1 · p. 19

C.2 Model Instances

• The encoders are custom ResNet-18 representations for synthetic data, CLIP ViT-L/14 for CelebA and natural images, and pathology-aware PLIP for Camelyon17.The decoders are domain-specific diffusion autoencoders, PathLDM, or a representation autoencoder operating on frozen CLIP patch tokens.Source: §C.2 · pp. 19–20
• Dictionaries include supervised Procrustes, Square SVD, trained batch top-K SAEs, and a public 6,144-atom MSAE; none is fitted to a student.Quantitative experiments use edit-friendly DDPM inversion, with the stated DDIM exceptions and empirical trust region settings.Source: §C.2 · p. 20
• The corrected students are poisoned-data ResNet-18 classifiers and linear probes; a distilled probe approximates a student’s signed margin in encoder space.For ranking, the margin is regressed by lasso on sparse-dictionary coefficients; a probe using the same encoder needs no distillation.Source: §C.2 · p. 20

C.3 Protocols

• Counterfactual quality compares DiDAE with five explainers on Square, CelebA-Blond, and Camelyon17 using desiderata, downstream Gain, oracle-axis projections, and unselected samples.DAE represents the global gradient-free family, while SCE is retained across datasets to expose its spatial-separability fit.Source: §C.3 · p. 20
• Dictionary inspection visualizes Procrustes and SVD components on Square and CelebA, tests single-dimension deactivation on Sparse Numbers, and compares correction Gain for Procrustes versus SVD.Source: §C.3 · p. 20
• Downstream correction uses 800 training and 200 validation counterfactuals per student, one fine-tuning iteration, and teacher review after specified misclassification, confidence, and outlier filters.Gain in average group accuracy is compared with metadata-based baselines and with projecting spurious components out of a frozen embedding.Source: §C.3 · p. 20
• Ranking tests four classifiers and reports latent flips, ambient flips, and verified flips, then uses teacher-marked spurious directions for CFKD repair.The ImageNet teacher is human; Sparse Numbers and NICO++ use an oracle, and ranking tests whether a planted water-jet atom is recovered without being named.Source: §C.3 · pp. 20–21

APPENDIX D HYPERPARAMETERS

• All explainers receive two attempts per factual, 800 training and 200 validation counterfactuals, and one CFKD iteration, without an added CFKD hyperparameter search.The quantitative CelebA dictionary spans only Blond_Hair and Male; its generator and student both operate at 128 × 128.Source: Appendix D · p. 21

APPENDIX E DECODER ARCHITECTURES AND TRAINING

• All three decoder architectures condition on a frozen encoder representation and differ in what they denoise and how conditioning enters the network.The decoder is trained with the encoder frozen.Source: Appendix E · p. 21

E.1 Diffusion Autoencoder

• For Square, Sparse Numbers, and CelebA, a pixel-space diffusion autoencoder uses a U-Net with timestep and frozen representation conditioning injected through AdaGN.Resolutions are 64 × 64 for synthetic datasets and 128 × 128 on CelebA, matching the student classifier resolution.Source: §E.1 · p. 21

E.2 PathLDM

• Camelyon17 uses PathLDM, a pathology-pretrained latent diffusion model conditioned on frozen PLIP embeddings and run with 50 edit-friendly DDPM steps.The model is trained on Camelyon17 training patches; inversion and decoding operate in its latent space.Source: §E.2 · p. 21

E.3 Representation Autoencoder

• ImageNet and NICO++ use a two-stage representation autoencoder that decodes frozen CLIP patch tokens and models the resulting latent conditioned on the dictionary-decomposed CLIP embedding.The foundation encoder remains frozen while the decoder and transformer are trained; edited representations are decoded by the same procedure used for the pixel-space autoencoder.Source: §E.3 · pp. 21–22

APPENDIX F INVERSION SCHEMES

• The inversion interface maps an image and its representation to a latent code, while decoding maps that code and a possibly edited representation back to an image.Decoding with the unchanged representation reconstructs the input.Source: Appendix F · p. 22

F.1 DDIM Inversion

• DDIM inversion deterministically maps the input to a code under its original conditioning, then decodes from that code under the edited representation.The integrations cancel up to discretization error, which causes visible smoothing; the method uses 20 uniformly spaced steps.Source: §F.1 · p. 22

F.2 Edit-Friendly DDPM Inversion

• Edit-friendly DDPM inversion stores noise maps from independently sampled forward marginals and reuses them for reverse decoding with edited conditioning.This reconstructs the input up to numerical error and retains its noise identity; the diffusion autoencoder uses 20 steps and PathLDM uses 50.Source: §F.2 · pp. 22–23

F.3 Stochastic Inversion on the Flow-Matching Schedule

• Flow-matching inversion adapts stochastic inversion to linear interpolation; η spans deterministic-to-stochastic sampling, while tstart trades input fidelity against conditioning-driven edits.At tstart = 1, the inverted state nearly suffices on its own and edits are largely ignored; guidance at γ = 1 preserves reconstruction, while γ > 1 amplifies the conditioning difference.Source: §F.3 · p. 23

Inversion ablation

• The RAE is competitive on desiderata, with best sparsity, NA, and unbiasedness, but decodes at a third of the speed and reaches about half the Gain.The RAE comparison changes the generator as well as the inversion, so it is not a pure inversion ablation.Source: Appendix G · p. 23, Table 7 · p. 24
• Five of seven CelebA columns favor edit-friendly DDPM over DDIM, while Square gives DDIM a 9.4-point Gain advantage despite DDPM leading diversity, sparsity, and unbiasedness.On CelebA, the clearest DDPM advantages are diversity (+23.3) and Gain (+11.4).Source: Appendix G · p. 23, Table 7 · p. 24

Additional figures

• Figure 10 shows Procrustes-direction edits for four Square generative factors and four CelebA attributes, while Figure 11 shows deactivation of one SAE dimension with neighboring numbers retained.The CelebA attributes Blond_Hair, Male, Heavy_Makeup, and Attractive are correlated.Source: Appendix G · p. 23, Figure 10 · p. 24, Figure 11 · p. 24
• Figure 12 shows that the first four Square SVD directions include mixed spatial axes that variance alone cannot separate.Comp1 corresponds to foreground color and Comp4 to background color.Source: Appendix G · p. 23, Figure 12 · p. 25
• Figure 13 shows the student boundary shifting toward causal-feature reliance after one CFKD iteration on Square and CelebA.The Square boundary is exact, whereas the CelebA boundary is approximated by projecting onto oracle Male and Blond_Hair predictions.Source: Appendix G · p. 23, Figure 13 · p. 26