Concept-SAE: Causal Probing in Visual Models

• Concept-SAE is a framework that leverages dual-supervised sparse autoencoders to enable direct causal intervention on visual model behavior.
• It partitions the latent space into concept tokens and free tokens to ensure semantic fidelity and disentanglement, supporting precision in failure diagnosis.
• The method’s staged training and counterfactual interventions reliably correct errors and localize adversarial vulnerabilities for enhanced model robustness.

Concept-SAE is a framework for active causal probing of visual model behavior that repurposes sparse autoencoders from an observational tool into a semantically grounded intervention interface. In the formulation introduced in "Concept-SAE: Active Causal Probing of Visual Model Behavior" (Ding et al., 26 Sep 2025), the latent space is partitioned into explicitly aligned concept tokens and unconstrained free tokens, and the concept tokens are trained with dual supervision on concept existence and spatial localization. The resulting representation is intended to support direct counterfactual intervention on internal concepts, as well as layerwise analysis of failure modes such as adversarial vulnerability.

1. Problem formulation and conceptual motivation

Concept-SAE starts from a limitation of standard SAEs: they can provide interpretable features, but those features are often ambiguous and ungrounded, which makes them unreliable for active, causal probing of model behavior. The framework is designed to address the gap between observational interpretability and mechanistic intervention. In the underlying paper, standard SAEs are described as providing an observational lens, whereas Concept-SAE is introduced to make semantically grounded, disentangled handles available for direct manipulation of internal states (Ding et al., 26 Sep 2025).

The framework also positions itself relative to concept bottleneck approaches. Concept Bottleneck Models and concept embeddings can enforce concepts at the output layer, but they do not scale well to intermediate representations and may not ensure disentanglement or semantic fidelity. Concept-SAE responds by anchoring only part of the latent space to human-defined concepts while leaving a residual channel for information not covered by the concept set. This division is central to its claim that concept-level intervention can be performed without collapsing the rest of the model’s internal representation.

Two design objectives organize the method. The first is semantic fidelity, meaning correspondence between a learned token and a human-understandable concept with ground-truth object-level or attribute-level meaning. The second is disentanglement, meaning statistical and functional separation between tokens so that interventions on one token do not simply co-manipulate unrelated factors. The method’s training strategy is built to enforce both properties jointly rather than treating them as post hoc evaluation criteria.

2. Architecture and staged training

The architecture splits the SAE latent space into two modules: concept tokens and free tokens. The concept side is implemented by a Concept Tokenizer and a Concept Aggregator. The free side is implemented by a Free Tokenizer and a Free Aggregator. The concept modules are intended to capture human-defined, semantically aligned features, while the free modules capture anything not explained by the concept set (Ding et al., 26 Sep 2025).

Training proceeds in three stages. First, the concept tokenizer is trained with dual supervision so that concept tokens capture as much target-concept information as possible. Second, the concept aggregator is trained to reconstruct the corresponding internal feature representation from the concept tokens. Third, the free modules are trained on the residual so that the remaining information can still be reconstructed without leaking concept information back into the free channel. The paper describes this as a staged training routine in which concept modules are trained first and free tokens are trained only afterward, specifically to prevent leakage or overlap between concept and free features.

For an internal feature representation $h$, the concept tokenizer predicts concept-specific scores and masks, while the concept aggregator reconstructs concept-conditioned internal features. The training objectives are given as:

$$\mathcal{L}_{tokenizer} = \lambda_1||\mathcal{S} - s||_2^2 + \lambda_2||\mathcal{M} - m||_2^2 + \lambda_3||W_{merge}||_1$$

$$\mathcal{L}_{aggr} = \lambda_1||\hat{h}_{concept} - h||_2^2 + \lambda_2\text{KL}(\text{sm}(W_{merge}) ||\text{sm}(W_{aggr}^\top) ) + \lambda_3||W_{aggr}||_1$$

$$\mathcal{L}_{free} = \lambda_1||\mathcal{A}_{free}(\mathcal{T}_{free}(h)) + \hat{h}_{concept} - h||_2^2 + \lambda_2||\mathcal{T}_{free}(h)||_1$$

These losses encode the main structural commitments of the framework: supervised concept recovery, reconstructive consistency, sparsity, and residual completion. The paper further states that $L_1$ penalties are applied on token or channel weights so that each concept representation is built from a minimal span of channels or patches.

3. Dual supervision and semantic grounding

The defining supervision signal in Concept-SAE is dual. For each human-defined concept, supervision is provided along two axes: an existence score indicating whether the concept is present, and a spatial mask indicating where the concept occurs. The existence score and mask are computed as

$$s_i = \text{Sigmoid}\left(z_i^{(d_t)} \cdot W_{score, i}^{(d_t)} + b_{score, i}\right)$$

and

$$m^{(d_s)}_i = z_i^{(d_t)} \cdot W^{(d_t \times d_s)}_{seg, i} + b^{(d_s)}_{seg, i}.$$

In the reported implementation, existence supervision is extracted automatically using a vision-LLM, while segmentation supervision is generated by a segmentation model and then refined using the existence judgments. If a concept is absent, its mask is forced to zero (Ding et al., 26 Sep 2025).

The purpose of this dual supervision is not only to indicate that a concept is present, but to constrain where the corresponding token should localize. The paper describes this as grounding semantically meaningful tokens. Tokens are tied to explicit, human-meaningful concepts such as blond hair, beard, or zebra stripes, and the concept tokens are evaluated by whether they reconstruct only their assigned regions rather than background or unintended factors. This grounding is further reinforced by the staged training procedure: free tokens are introduced only after the concept modules have already captured the concept-aligned portion of the representation.

A further reported result is that irrelevant concept filtering works cleanly in this setup. When irrelevant concepts that never appear in the data are introduced, the corresponding concept tokens remain reliably inactive. The paper also reports that free tokens remain semantically rich after the concept channel is carved out, so the residual module is not merely a reconstruction dump. This suggests a decomposition in which semantically named factors and unlabeled residual factors can coexist without the concept channel losing its interpretability.

4. Causal intervention and active probing

Once concept tokens are grounded, Concept-SAE uses them for direct intervention. The paper describes two forms of causal probing. The first is counterfactual intervention: concept scores are manually modified, internal features are regenerated through the concept aggregator, and downstream computation is rerun to test whether the model’s prediction changes in the expected direction. The second is failure diagnosis, where differences in concept-score patterns across correct and incorrect predictions are used to analyze model errors (Ding et al., 26 Sep 2025).

The intervention mechanism is explicitly token-level and internal. For example, the reported procedure includes manually setting scores for concepts such as beard or adam’s apple from $0$ to $1$ for a male image misclassified as female, or suppressing those concepts for a female image misclassified as male. The reconstructed internal representation is then passed to the downstream classifier block, and the observed prediction change is used as evidence of causal influence. The paper states that such interventions can reliably correct predictions, with deeper layers in ViT models being especially effective for causal correction.

This operationalizes a stronger notion of explanation than passive feature inspection. Standard SAE workflows often identify features correlated with a model’s behavior; Concept-SAE instead asks whether altering a feature-aligned concept token changes the behavior itself. The method therefore treats concept tokens as experimentally manipulable variables rather than merely descriptive coordinates in latent space.

The same design also enables layer-specific intervention analysis. Because concept tokens are learned over intermediate representations, the method can test where in the network a concept is both encoded and causally effective. The paper reports that intervention efficacy varies by layer, which is used as evidence that concept representations are not uniformly actionable across a model’s depth.

5. Quantitative evaluation and failure-mode localization

The paper evaluates Concept-SAE with several metrics intended to capture disentanglement, localization, ambiguity, and robustness. A central localization metric is the Localization Ratio (LocR),

$$\text{LocR} = \frac{\text{MSE}\left((h - \hat{h}_{concept}) \odot M_{background}\right)}{\text{MSE}\left((h - \hat{h}_{concept}) \odot M_{face}\right)}.$$

Higher LocR indicates that the concept reconstruction is concentrated on the intended region rather than the background. In the reported comparison against a Concept-Embedding Model, Concept-SAE achieves a substantially higher LocR; the paper gives the specific example of ResNet-18 layer 5, where Concept-SAE obtains $$\mathcal{L}_{tokenizer} = \lambda_1||\mathcal{S} - s||_2^2 + \lambda_2||\mathcal{M} - m||_2^2 + \lambda_3||W_{merge}||_1$$0 and the comparison model obtains $$\mathcal{L}_{tokenizer} = \lambda_1||\mathcal{S} - s||_2^2 + \lambda_2||\mathcal{M} - m||_2^2 + \lambda_3||W_{merge}||_1$$1 (Ding et al., 26 Sep 2025).

A second diagnostic is concept score entropy. The paper reports that entropy is markedly higher for incorrectly classified or adversarial samples, interpreting this as a signature of ambiguous or poorly disentangled internal concepts. This entropy-based view is extended to adversarial analysis by comparing clean and adversarial concept-score distributions across layers. The method computes Jensen-Shannon distance between clean and adversarial distributions for each layer; layers with larger divergence are treated as more vulnerable.

This layerwise diagnostic is then functionally validated. The procedure freezes all layers except the one with the highest Jensen-Shannon distance and finetunes that layer on a mixture of clean and adversarial data. The paper reports that the layers identified by Concept-SAE’s vulnerability metric yield the strongest adversarial gains, including examples with up to $$\mathcal{L}_{tokenizer} = \lambda_1||\mathcal{S} - s||_2^2 + \lambda_2||\mathcal{M} - m||_2^2 + \lambda_3||W_{merge}||_1$$2 robust accuracy. In this sense, the framework is არა merely descriptive of adversarial effects; it is used to localize where robustness interventions should be applied.

The paper also states that the dual-supervision strategy outperforms alternative methods in disentanglement, and that free tokens remain semantically rich after the concept-specific pathway is trained. Taken together, these results support the paper’s central claim that semantically grounded concept tokens can function simultaneously as explanatory variables, intervention targets, and diagnostic probes for layerwise failure analysis.

6. Relation to broader concept-aligned SAE research

Concept-SAE belongs to a broader line of work that augments unsupervised sparse coding with explicit semantic alignment. Several later or parallel frameworks make this tendency more explicit. AlignSAE introduces a "pre-train, then post-train" curriculum in which ontology concepts are bound to dedicated latent slots after an initial unsupervised training phase, with supervised concept slots separated from a free feature bank for reconstruction (Yang et al., 1 Dec 2025). CB-SAE prunes low-utility SAE neurons and augments the latent space with a lightweight concept bottleneck aligned to a user-defined concept set, reporting improvements of $$\mathcal{L}_{tokenizer} = \lambda_1||\mathcal{S} - s||_2^2 + \lambda_2||\mathcal{M} - m||_2^2 + \lambda_3||W_{merge}||_1$$3 in interpretability and $$\mathcal{L}_{tokenizer} = \lambda_1||\mathcal{S} - s||_2^2 + \lambda_2||\mathcal{M} - m||_2^2 + \lambda_3||W_{merge}||_1$$4 in steerability across LVLMs and image generation tasks (Kulkarni et al., 11 Dec 2025). DACO uses a curated dictionary of 15,000 multimodal concepts to initialize SAE decoders and automatically annotate atoms for safety steering in MLLMs (Luo et al., 10 Apr 2026).

These developments also connect to criticisms of a purely unsupervised SAE workflow. One paper argues that SAEs are powerful tools for discovering unknown concepts but are less effective for acting on known concepts, distinguishing hypothesis generation from direct manipulation or detection of prespecified concepts (Peng et al., 30 Jun 2025). Concept-SAE can be read as an explicit attempt to overcome that limitation by introducing direct supervision and spatial grounding into the SAE latent space. This suggests that the method is not a rejection of sparse discovery, but a hybridization of discovery and concept alignment for cases where causal action on named concepts is required.

Broader theoretical work reinforces the same point from a geometric perspective. "Projecting Assumptions: The Duality Between Sparse Autoencoders and Concept Geometry" argues that SAE architectures impose structural assumptions about what concepts are detectable, and that no universal SAE exists (Hindupur et al., 3 Mar 2025). A plausible implication is that Concept-SAE’s token partitioning and dual supervision do not merely improve usability; they alter the recoverable concept geometry by construction. At the same time, "Dense SAE Latents Are Features, Not Bugs" shows that many dense latents correspond to meaningful model representations such as position tracking, context binding, entropy regulation, alphabet signals, part-of-speech, and principal component reconstruction (Sun et al., 18 Jun 2025). This complicates any simple equation between interpretability and extreme sparsity, and helps explain why Concept-SAE preserves a free-token channel instead of forcing every useful factor into a narrowly supervised bottleneck.

Within this broader literature, Concept-SAE is most precisely characterized as a causally oriented, semantically grounded SAE variant for vision models. Its distinctive contribution is not sparse coding alone, nor concept supervision alone, but the use of dual-supervised concept tokens as manipulable internal variables for counterfactual probing, error diagnosis, and adversarial vulnerability localization (Ding et al., 26 Sep 2025).