Causal Differentiating Concepts: Interpreting LM Behavior via Causal Representation Learning

Thank you for your positive feedback on our work. We really appreciate your questions and comments on this work. To address your questions about the sensitivity of our method to specific modeling choices, we ran some additional experiments that evaluate misspecification of the latent dimension, violation of the 1-sparse assumption, and assess our ability to extract insights from additional layers of LMs.

---

Applying the method to intermediate layers (rather than the final LLM layer)

This is a great point. Indeed, intermediate layers of LLMs can often capture rich semantic information (possibly even richer than the last layer). Even though our experiments assess the last layer of LLMs, implementationally, our method can be directly applied to intermediate LLM layers too. Although our theoretical results do not directly translate to intermediate layers, a host of recent literature finds that latent concepts are linearly encoded in LM representations (linear representation hypothesis; Roeder et al. 2021; Jiang et al., 2024; Park et al., 2024). Since the contrastive learning constraints we develop in Eqn 4 are meant to disentangle linearly mixed concepts, we hypothesize that we can still use the method proposed in this paper for unsupervised discovery of unknown latent concepts present in the model’s internal representations.

To empirically test this, we extend our refusal case study to intermediate representations to test whether our method can extract relevant insights from intermediate LM layers as well. For simplicity, we consider the representation at the last token position in each layer of the Llama-3.1-8B model. We find that in the earlier layers, our approach does not yield disentangled representations. This is to be expected, as it has been noted before that initial LLM layers capture surface-level encodings of the inputs. However, past layer 9, the structure presented in Figure 2 starts to emerge. That is, we observe that the harmfulness and topicality dimensions identified in our experiments can also be disentangled in intermediate model representations. For a quantitative assessment, we calculate the average unsupervised disentanglement score between the model's internal representations and the last layer. We observe an unsupervised disentanglement score of $0.78 \pm 0.02$ for layers 10-30 compared to $0.39 \pm 0.20$ for the first 9 layers. Interestingly, these findings align with a recent paper (Zhao et al., 2025), released earlier this month, that finds that harmfulness and refusal directions can be extracted from LLM representations for Llama2-Chat-7B around layer 10.

It would be interesting for future exploration to study whether there are indeed interesting concepts that can be extracted from intermediate layers that are abstracted away or lost in later layers. Our method, hopefully, provides a useful tool for such discoveries and insights.

---

Violations of the 1-sparse assumption

This is a great question. Even though we restrict our discussion to 1-sparsity, we expect that our results can be extended to 2-sparse condition as well (or generally speaking, a p-sparse condition) following the proof strategy as Ahuja et al. (2022a). That is, if the causal differentiating concepts are p-sparse and non-overlapping, we can expect identifiability up to permutations and block-diagonal transforms, rather than our stronger results on identifiability up to permutations and scaling. Further, if causal differentiating concepts are p-sparse and overlapping, we expect that the latents at the intersection of these blocks to be identified up to permutations and scaling.

To validate this second setting in our case, we construct synthetic data with 2-sparse causal differentiating concepts for latent dimension d=3. Specifically, imagine three causal variables: occupation, gender, and experience that determine the income prediction behavior of the model, such that the $p(y|x)$ can be clustered into four classes: {(Doctors/Lawyers, Male, More experience), (Doctors/Lawyers, Female, Less experience); (Teachers/Nurses, Female, More experience), and (Teachers/Nurses, Male, Less experience). We see that in this data, groups $C_k$ and $C_l$ have 2 causal differentiating concepts. For example, for behavior classes (Doctors/Lawyers, Male, More experience) and (Teachers/Nurses, Female, More experience), the causal differentiating concepts are occupation and gender.

We observe that our method can still disentangle the three causal factors. The average disentanglement measures are comparable to our 1-sparse setting—DCI-D ($0.95 \pm 0.02$), DCI-C ($0.98 \pm 0.01$), and DCI-I ($1.0 \pm 0.0$). The reported numbers are across 5 runs with different initializations. Note that we get this disentanglement because the 2-sparse causal differentiating concepts are overlapping for different pairs of groups. This indicates that our method can be extended beyond the 1-sparse setting following the intuition from Ahuja et al. (2022a). However, we leave more detailed explorations of these cases to future work.

Can you please clarify what you mean by “distributed causal shifts”? Does this mean that for any two behavioral classes, all (or multiple) latent concepts are causally differentiating? In this case, we would expect that our method would not yield disentangled representations since some form of inductive biases might be necessary for disentanglement. Please let us know if this does not address your question.

---

Choice of latent dimensionality

This is an interesting question. We conducted experiments to test variation in the learned number of latents in synthetic experiment settings and to assess factors that can help in choosing the latent dimension in practice.

To this end, we consider two cases: (1) when the dimensionality of learned latents is higher than that of true mediating concepts, and (2) when the dimensionality of learned latents is lower than that of true mediating concepts.

In experiment 1, we consider synthetic experiments with the number of true mediating concepts set to 2 and the number of learned latent concepts set to 3 and 4. We find that if we relax the span constraints (line 147), that is, the learned perturbations do not need to span all latent concepts, the learned perturbations correspond to the true causal differentiating concepts. That is, the model disregards any extra latent(s) in learning the causal differentiating concepts, and the extra latents are solely some function of one of the other latent variables. Thus, the model achieves a high disentanglement score, but a lower completeness score (since more than one latent captures a true causal factor): DCI-D $0.95 \pm 0.03$ and DCI-C $0.81 \pm 0.14$ for $n=3$, and DCI-D $0.94 \pm 0.04$ and DCI-C $0.76 \pm 0.19$ for $n=4$. In sum, when the model's latent dimension is misspecified to be larger than the true latent dimension, we observe a degree of robustness in the performance.

With the span constraint, the learned perturbations no longer correspond to the true causal differentiating concepts. This leads to degradation in disentanglement scores, alongside lower completeness scores: DCI-D $0.64 \pm 0.11$ and DCI-C $0.81 \pm 0.14$ for $n=3$ and DCI-D $0.89 \pm 0.08$ and DCI-C $0.57 \pm 0.13$ for $n=4$. Interestingly, the disentanglement scores are still higher than the autoencoding (DCI-D $0.19 \pm 0.12$, DCI-C $0.24 \pm 0.15$) and autoencoding + prediction (DCI-D $0.08 \pm 0.07$, DCI-C $0.09 \pm 0.08$) baselines. But more importantly, overspecifying latent dimensionality consistently leads to a drop in completeness scores, which can be used as a signal for reducing the latent dimensionality as a hyperparameter.

For experiment 2, we consider synthetic experiments with the number of true mediating concepts set to 3 and the number of learned latent concepts set to 2. We find that underspecifying the latent dimensionality results in a significant drop in the disentanglement scores ($0.51 \pm 0.20$). This is to be expected since the learned latents cannot separate out the mediating concepts since there are not enough latents. In this case, we conclude that underspecifying latent dimensionality consistently leads to a lower disentanglement score, which can be used as a signal for increasing the latent dimensionality as a hyperparameter.

We restrict ourselves to synthetic experiments for sensitivity analysis because in semi-synthetic and non-synthetic experiments, a parallel comparison is not feasible since the dimensionality of true mediating concepts is integral to the dataset and is not dictated externally. We believe that findings from synthetic experiments would naturally extend to general cases. Even though non-synthetic settings lack ground-truth mediating concepts, not allowing direct evaluation of disentanglement and completeness scores, the disentanglement measures can be measured across multiple model runs as a signal for guiding the latent dimensionality, following Duan et al. (2021).

---

References

[1] Kartik Ahuja, Jason S Hartford, and Yoshua Bengio. Weakly supervised representation learning with sparse perturbations. NeurIPS 2022.
[2] Sunny Duan, Loic Matthey, Andre Saraiva, Nick Watters, Chris Burgess, Alexander Lerchner, Irina Higgins. Unsupervised Model Selection for Variational Disentangled Representation Learning. ICLR 2020.
[3] Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, Bryon Aragam, Victor Veitch. On the Origins of Linear Representations in Large Language Models. ICML 2024.
[4] Kiho Park, Yo Joong Choe, Victor Veitch. The linear representation hypothesis and the geometry of large language models. ICML 2024.
[5] Geoffrey Roeder, Luke Metz, Diederik P. Kingma. On Linear Identifiability of Learned Representations. ICML 2021.
[6] Jiachen Zhao, Jing Huang, Zhengxuan Wu, David Bau, Weiyan Shi. LLMs Encode Harmfulness and Refusal Separately. 2025.

Thank you for your feedback on our work. We share some clarifications on the validity of the assumptions and the applicability of our method, which will hopefully help address your concerns.

---

Would the proof or technique break if we have 2-sparsity instead of 1-sparsity?

Even though we restrict our discussion to 1-sparsity, we expect that our results can be extended to 2-sparse condition as well (or generally speaking, a p-sparse condition) following the proof strategy as Ahuja et al. (2022a). That is, if the causal differentiating concepts are p-sparse and non-overlapping, we can expect identifiability up to permutations and block-diagonal transforms, rather than our stronger results on identifiability up to permutations and scaling. Further, if causal differentiating concepts are p-sparse and overlapping, we expect that the latents at the intersection of these blocks to be identified up to permutations and scaling.

To validate this second setting in our case, we construct synthetic data with 2-sparse causal differentiating concepts for d=3. Specifically, imagine three causal variables: occupation, gender, and experience that determine the income prediction behavior of the model, such that the $p(y|x)$ can be clustered into four classes: {(Doctors/Lawyers, Male, More experience), (Doctors/Lawyers, Female, Less experience); (Teachers/Nurses, Female, More experience), and (Teachers/Nurses, Male, Less experience). We see that in this data, groups $C_k$ and $C_l$ have 2 causal differentiating concepts. For example, for behavior classes (Doctors/Lawyers, Male, More experience) and (Teachers/Nurses, Female, More experience), the causal differentiating concepts are occupation and gender.

We observe that our method can still disentangle the three causal factors. The average disentanglement measures (across 5 initializations) are comparable to our 1-sparse setting—DCI-D ($0.95 \pm 0.02$), DCI-C ($0.98 \pm 0.01$), and DCI-I ($1.0 \pm 0.0$). This indicates that our method can be extended beyond the 1-sparse setting as well.

With respect to the feasibility of 1-sparse assumption, we want to note that our 1-sparse assumption is more relaxed than that of Ahuja et al. (2022a) since we do not require that every pair of clusters have only 1 causal differentiating concept or that the pair of clusters that have 1 causal differentiating concepts are known beforehand, rather that every $z^r$ is 1-sparse for some pair of clusters. This assumption is easier to meet in practice, as unless some points in input space are not observed altogether, 1-sparsity would hold naturally. For instance, in the above example with 2-sparsity, we do not observe data points that are {Doctor/Lawyer, Female, More experience} at all. If we did, we can have one of following cases: (a) this is a separate behavior cluster, in which case this cluster and (Doctors/Lawyers, Male, More experience) will have gender as the causal differentiating concept and this cluster and (Teachers/Nurses, Female, More experience) will have occupation as the causal differentiating concept; (b) we have {Doctor/Lawyer, Female} cluster (including both more and less experience data points), in which case this and (Teachers/Nurses, Female, More experience) will still have occupation as the only causal differentiating concept; and so on. We hope this example helps convey why 1-sparsity is not a particularly strong assumption in practice, as long as sufficient variability in the input space is observed.

---

Linearity assumption

Log-linearity prediction is ubiquitous not only in causal representation learning literature (e.g., Hyvärinen and Morioka, 2016; Ahuja et al., 2022b) but also in machine learning. The idea is that the non-linearity in data is pushed into the feature extraction, and the outcome (e.g., next-token, label, etc.) is predicted linearly from the extracted features. Thus, the expressivity of such a setup is not akin to linear regression.

While log-linearity is certainly an assumption, previous work in CRL theoretically proves that unsupervised learning of disentangled representations is fundamentally impossible without inductive biases (Locatello et al., 2019). Thus, certain assumptions are inherent to the goal of disentanglement. In the context of LLMs, the log-linearity assumption may not be as limiting since a host of recent work has found that concepts are linearly encoded internally (linear representation hypothesis; Roeder et al. 2021; Jiang et al., 2024; Park et al., 2024), lending themselves to be discovered by our method without the need of supervision.

---

Theoretical novelty

The judgment of novelty is inherently subjective; however, we would like to add that our paper does not claim or promise novelty in proof technique (our proof sketch clearly lays out that the proof naturally follows using techniques from Ahuja et al. (2022b), Hyvärinen and Morioka (2016), and Ahuja et al. (2022a)). A lot of mathematical work, generally, and especially in CRL, focuses on using existing proof techniques to prove theorems in new settings for new insights. The goal of our paper is to develop novel assumptions that take LLM behavior into account, enabling us to use the CRL framework to interpret LLMs without the need for labeled examples, a setting which is majorly under-explored in causal representation learning, with the exception of few recent works (Rajendran et al., 2024; Ahuja et al. 2022b; Moran and Aragam, 2025).

Another thing to note is that the assumptions that we develop are novel. Reiterating an earlier point, our 1-sparsity assumption is more relaxed than that of Ahuja et al. (2022a) since we do not require that every pair of clusters have only 1 causal differentiating concept or that the pair of clusters that have 1 causal differentiating concept are known beforehand. Thus, the assumptions and proofs from previous work do not out-of-the-box translate to our setting. Rather, we identify more relaxed assumptions that can still use existing proof techniques for identifiable discovery of latent concepts.

---

Number of labels

Can you clarify why you believe that increasing the number of classes (presumably for the outcome y) could change the conclusion of our experiments, where we already consider behaviors that are non-binary? As we note in the paper, behaviors can be finer-grained than the number of classes, e.g., P(refuse|prompt) can be very high, high, medium, etc. Indeed, in the real-world experiments, we consider behaviors that are continuous valued and bucketed into 3 behavioral classes.

In terms of comparison with the scope of other interpretability literature, a host of interpretability work spanning SAEs, mechanistic interpretability, mediation analysis, and more study binarized setting, for example, refusal behavior in language models (whether the model refuses a query or not; e.g., Zhao et al., 2025, Yeo et al. 2025), truthfulness in LLMs (whether model output is truthful or not; e.g., Zou et al., 2025), subject-verb agreement (whether subject and verb agree or not; e.g., Marks et al. 2025) and so on. In any of these works and ours, it is not that we restrict the output space of LLMs—they can still produce free-form text—but we interpret a property of this free-form text (refusal, truthfulness, etc.). Thus, our method, similar to existing approaches, can analyze LLMs out-of-the-box, be it a pretrained model, an instruction-tuned model, or a finetuned model.

---

Replacing 1-sparsity with some “full rank” condition on concepts <-> label structure

Just to clarify, by “full rank condition on concepts <-> label structure”, did you mean that if the perturbation matrix $\Delta$, which is comprised of the latent concepts that differentiate some labels $c_k$ and $c_l$, is full rank? Based on that interpretation, we do not think full rank condition guarantees identifiability. Consider $\Delta$ as the set of true perturbations and $\Delta’$ as the set of guessed perturbations. Theorem 1 yields $\Delta’ = A \Delta$, where $A$ is an affine transformation matrix. Now, if we assume that $\Delta$ is full rank, we can conclude that $A$ would be invertible (since both $\Delta$ and $\Delta’$ have rank n, so would the matrix $A$). But, it does not say anything about $A$ being an identity matrix (or a permutation of an identity matrix). A similar argument appears in Ahuja et al. (2022a) to prove that full rankness of $\Delta$ implies that true latents can be identified up to affine transformation, but for identifiability up to permutation and scaling, they use additional sparsity conditions as well.

---

References

[1] Kartik Ahuja, Jason S Hartford, and Yoshua Bengio. Weakly supervised representation learning with sparse perturbations. NeurIPS 2022.
[2] Kartik Ahuja, Divyat Mahajan, Vasilis Syrgkanis, Ioannis Mitliagkas. Towards efficient representation identification in supervised learning. CLeaR 2022.
[3] Aapo Hyvärinen, Hiroshi Morioka. Unsupervised Feature Extraction by Time-Contrastive Learning and Nonlinear ICA. NeurIPS 2016.
[4] Samuel Marks, Can Rager, Eric J. Michaud, Yonatan Belinkov, David Bau, Aaron Mueller. Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models. ICLR 2025.
[5] Gemma E. Moran, Bryon Aragam. Towards Interpretable Deep Generative Models via Causal Representation Learning. 2025.
[6] Goutham Rajendran, Simon Buchholz, Bryon Aragam, Bernhard Schölkopf, Pradeep Ravikumar. From Causal to Concept-Based Representation Learning. NeurIPS 2024.
[7] Wei Jie Yeo, Nirmalendu Prakash, Clement Neo, Roy Ka-Wei Lee, Erik Cambria, Ranjan Satapathy. Understanding Refusal in Language Models with Sparse Autoencoders. 2025.
[8] Jiachen Zhao, Jing Huang, Zhengxuan Wu, David Bau, Weiyan Shi. LLMs Encode Harmfulness and Refusal Separately. 2025. [9] Andy Zou et al. Representation Engineering: A Top-Down Approach to AI Transparency. 2025.

Thank you for your positive feedback on the theory and experiments presented in our work. We appreciate your recognition of the value of our work in scaling causal abstractions through unsupervised discovery of high-level concepts that affect model behavior. We also appreciate your positive comments regarding our experiments and evaluation. To address the questions raised about the sensitivity to variation in some modeling choices, we ran 3 additional experiments that evaluate misspecification of the latent dimension and assess our ability to extract insights from intermediate layers of LMs.

---

Applying the method to intermediate language model layers

This is a great point. In our experiments, we restrict ourselves to interpreting the last layer of the language models; however, experimentally, our method can be directly applied to intermediate LM layers as well. Although our theoretical results do not directly translate to intermediate representations, a host of recent literature identifies that latent concepts are linearly encoded in LM representations (linear representation hypothesis; Roeder et al., 2021; Jiang et al., 2024; Park et al., 2024). Since the contrastive learning constraints we develop in Eqn 4 are meant to disentangle linearly mixed concepts, we hypothesize that we can still use the method proposed in this paper for unsupervised discovery of unknown latent concepts present in the model’s internal representations.

To empirically test this, we extend our refusal case study to intermediate representations to test whether our method can extract relevant insights from intermediate LM layers as well. For simplicity, we consider the representation at the last token position in each layer of the Llama-3.1-8B model, which is a common practice in other work studying internal LLM representations (Zou et al., 2022; Arditi et al., 2024). We find that in the earlier layers, our approach does not yield disentangled representations. This is to be expected, as it has been noted before that initial LLM layers capture surface-level encodings of the inputs. However, past layer 9, the structure presented in Figure 2 starts to emerge. That is, we observe that the harmfulness and topicality dimensions identified in our experiments can also be disentangled in intermediate model representations. For a quantitative assessment, we calculate the average unsupervised disentanglement score between the model's internal representations and the last layer. We observe an unsupervised disentanglement score of $0.78 \pm 0.02$ for layers 10-30 compared to $0.39 \pm 0.20$ for the first 9 layers. Interestingly, these findings align with a recent paper (Zhao et al., 2025), released earlier this month, that finds that harmfulness and refusal directions can be extracted from LLM representations for Llama2-Chat-7B around layer 10.

---

If the dimensionality of the bottleneck doesn't correspond well to the true number of factors, what happens?

This is an interesting question. We conducted experiments to test variation in the learned number of latents in synthetic experiment settings. We consider two cases: (1) when the dimensionality of learned latents is higher than that of true mediating concepts, and (2) when the dimensionality of learned latents is lower than that of true mediating concepts.

In experiment 1, we consider synthetic data with the number of true mediating concepts set to 2 and the number of learned latent concepts set to 3 and 4. We find that if we relax the span constraints (line 147), that is, the learned perturbations do not need to span all latent concepts, the learned perturbations correspond to the true causal differentiating concepts. That is, the model disregards any extra latent(s) in learning the causal differentiating concepts, and the extra latents are solely some function of one of the other latent variables. Thus, the model achieves a high disentanglement score, but a lower completeness score (since more than one latent captures a true causal factor): DCI-D $0.95 \pm 0.03$ and DCI-C $0.81 \pm 0.14$ for $n=3$, and DCI-D $0.94 \pm 0.04$ and DCI-C $0.76 \pm 0.19$ for $n=4$. In sum, when the model's latent dimension is misspecified to be larger than the true latent dimension, we observe a degree of robustness in the performance.

With the span constraint, the learned perturbations no longer correspond to the true causal differentiating concepts. This leads to degradation in disentanglement scores, alongside lower completeness scores: DCI-D $0.64 \pm 0.11$ and DCI-C $0.81 \pm 0.14$ for $n=3$ and DCI-D $0.89 \pm 0.08$ and DCI-C $0.57 \pm 0.13$ for $n=4$. Interestingly, the disentanglement scores are still higher than the autoencoding (DCI-D $0.19 \pm 0.12$, DCI-C $0.24 \pm 0.15$) and autoencoding + prediction (DCI-D $0.08 \pm 0.07$, DCI-C $0.09 \pm 0.08$) baselines. But more importantly, overspecifying latent dimensionality consistently leads to a drop in completeness scores, which can be used as a signal for reducing the latent dimensionality as a hyperparameter.

For experiment 2, we consider synthetic data with the number of true mediating concepts set to 3 and the number of learned latent concepts set to 2. We find that underspecifying the latent dimensionality results in a significant drop in the disentanglement scores ($0.51 \pm 0.20$). This is to be expected since the learned latents cannot separate out the mediating concepts since there are not enough latents. In this case, we conclude that underspecifying latent dimensionality consistently leads to a lower disentanglement score, which can be used as a signal for increasing the latent dimensionality as a hyperparameter.

We restrict ourselves to synthetic experiments for sensitivity analysis because in semi-synthetic and non-synthetic experiments, a parallel comparison is not feasible since the dimensionality of true mediating concepts is integral to the dataset and is not dictated externally. We believe that findings from synthetic experiments would naturally extend to general cases. Even though non-synthetic settings lack ground-truth mediating concepts, not allowing direct evaluation of disentanglement and completeness scores, the disentanglement measures can be measured across multiple model runs as a signal for guiding the latent dimensionality, following Duan et al. (2021).

---

References

[1] Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, Neel Nanda. Refusal in Language Models Is Mediated by a Single Direction. NeurIPS 2025.
[2] Sunny Duan, Loic Matthey, Andre Saraiva, Nick Watters, Chris Burgess, Alexander Lerchner, Irina Higgins. Unsupervised Model Selection for Variational Disentangled Representation Learning. ICLR 2020.
[3] Yibo Jiang, Goutham Rajendran, Pradeep Ravikumar, Bryon Aragam, Victor Veitch. On the Origins of Linear Representations in Large Language Models. ICML 2024.
[4] Kiho Park, Yo Joong Choe, Victor Veitch. The linear representation hypothesis and the geometry of large language models. ICML 2024.
[5] Geoffrey Roeder, Luke Metz, Diederik P. Kingma. On Linear Identifiability of Learned Representations. ICML 2021.
[6] Jiachen Zhao, Jing Huang, Zhengxuan Wu, David Bau, Weiyan Shi. LLMs Encode Harmfulness and Refusal Separately. 2025.
[7] Andy Zou et al. Representation Engineering: A Top-Down Approach to AI Transparency. 2025.

Thank you for your positive feedback on our work. We appreciate that you found our work to be novel, with sound theory and experiments. Below, we answer the questions posed in your review.

---

Expressiveness of log-linear functions

Good question. Hyvarinen and Morioka (2016) point out that log-linear prediction on top of flexible, highly non-linear feature extraction still has universal expressiveness. Indeed, this style of modeling is akin to nonlinear prediction via kernel regression or basis expansion, which remain flexible, nonlinear predictors.

---

Does Assumption 2 require that every pair must be differentiated by some $z^r$?

Even though this is not explicitly stated in the assumption, it is implicit in the problem definition that a difference in behavior would correspond to some mediating concept(s) $z^r$.

---

Does it require that each $z^r$ must differentiate exactly one pair of concepts?

No, it is not necessary that each $z^r$ must differentiate exactly one pair of groups. In the refusal case study, for instance, latent dimension 1 (interpreted as harmfulness) is a causal differentiating concept between harmful ($p(refusal|x)=high$) and pseudo-harmful queries ($p(refusal|x)=high$), and also between harmful and harmless queries ($p(refusal|x)=low$).

---

References

[1] Aapo Hyvärinen, Hiroshi Morioka. Unsupervised Feature Extraction by Time-Contrastive Learning and Nonlinear ICA. NeurIPS 2016.