# Causal Concept Graphs in LLM Latent Space for Stepwise Reasoning
**Authors**: Md Muntaqim Meherab, Noor Islam S. Mohammad, Faiza Feroz
> Corresponding author Noor Islam S. Mohammad:
institutetext: Daffodil International University, Dhaka, Bangladesh email: {meherab2305101354, faiza.cse}@diu.edu.bd institutetext: New York University, Brooklyn, NY, USA email: noor.islam.s.m@nyu.edu
## Abstract
Sparse autoencoders can localise where concepts live in language models, but not how they interact during multi-step reasoning. We propose Causal Concept Graphs (CCG): a directed acyclic graph over sparse, interpretable latent features, where edges capture learned causal dependencies between concepts. We combine task-conditioned sparse autoencoders for concept discovery with DAGMA-style differentiable structure learning for graph recovery, and introduce the Causal Fidelity Score (CFS) to evaluate whether graph-guided interventions induce larger downstream effects than random ones. On ARC-Challenge, StrategyQA, and LogiQA with GPT-2 Medium, across five seeds ( $n{=}15$ paired runs), CCG achieves $\mathrm{CFS}=5.654\pm 0.625$ , outperforming ROME-style tracing ( $3.382\pm 0.233$ ), SAE-only ranking ( $2.479\pm 0.196$ ), and a random baseline ( $1.032\pm 0.034$ ), with $p<0.0001$ after Bonferroni correction. Learned graphs are sparse (5–6% edge density), domain-specific, and stable across seeds.
## 1 Introduction
Mechanistic interpretability has made rapid progress: we can localise semantic features and circuits in transformers and extract sparse, monosemantic dictionaries from residual streams [6, 21, 2, 5, 26]. What remains hard is the dynamic question: for multi-step reasoning, which internal features interact, and in what order, as computation unfolds. This gap matters for reliability and safety: without tracing internal reasoning, we cannot robustly diagnose failures or distinguish genuine reasoning from shortcut strategies [24]. Existing tools only partially address this. Model editing methods (e.g., ROME/MEMIT) precisely localise single factual associations [17, 18, 9] but are not designed for distributed, compositional reasoning. Concept Bottleneck Models provide interpretability via an explicit concept layer [13, 31] but require a human-specified vocabulary and supervision. We combine sparse feature discovery with causal structure learning: we first extract task-conditioned concept features from activations, then learn a DAG over those features, yielding Causal Concept Graphs (CCG) with no manual concept annotation.
#### Contributions.
We contribute (i) a task-conditioned sparse autoencoder with TopK gating and neuron resampling that achieves a stable 5.1% L0 activation rate on reasoning inputs (Section 3.1); (ii) a DAGMA-based causal graph learner over concept activation matrices that recovers sparse DAGs with 5–6% edge density (Section 3.2); (iii) the Causal Fidelity Score (CFS), a numerically stable intervention-based metric for evaluating whether the learned graph identifies concepts with large causal reach (Section 3.3); and (iv) multi-seed experiments on three reasoning benchmarks showing consistent, statistically significant improvements over strong baselines (Section 4). *[Error downloading image: 2603.10377v1/ccg_pipeline_springer.png]*
Figure 1: CCG pipeline. Stage 1: task-conditioned SAE on GPT-2 Medium residual activations (Layer 12) with TopK gating ( $K{=}256$ , $k{=}13$ ; 5.1% L0). Stage 2: DAGMA learns a sparse DAG over the top-64 concepts per domain. Stage 3: CFS evaluates intervention faithfulness ( $\mathrm{CFS}{=}5.654$ ; $p{<}0.0001$ vs. baselines).
## 2 Related Work
#### Mechanistic interpretability.
Transformers exhibit identifiable circuits and algorithms [20, 21, 27, 19]. Superposition explains widespread polysemanticity [6], motivating sparse autoencoders for feature dictionaries [2, 5] that remain interpretable at scale [26]. We focus on feature-to-feature interaction structure during computation. Causal tracing and model editing ROME and successors localise and edit factual associations in mid-layer computations [17, 18, 9]. These methods target single associations; our goal is multi-feature, multi-step causal structure. Our intervention-based evaluation is inspired by this line but operates on learned feature graphs.
Causal structure learning DAG learning spans classic constraint/score-based methods [25, 3, 22] and continuous relaxations [32], with DAGMA improving numerical behaviour near optima [1]. We adapt DAGMA to task-structured concept activations rather than i.i.d. tabular variables. Concept-based explanations. TCAV links user-defined concepts to predictions [11]; CBMs and variants enforce/approximate a concept layer but rely on a predefined vocabulary or supervision [13, 30, 31]. CCG instead discovers concepts from activations and learns their dependencies. LLM reasoning Chain-of-thought prompting elicits explicit intermediate steps [28, 14], and mechanistic studies analyse specific reasoning circuits [27, 15]. Closest in spirit are causal intervention analyses of internal computation [7, 29], though typically at component level; we target sparse concept features and their learned DAG structure.
## 3 Methodology
### 3.1 Stage 1: Task-Conditioned Sparse Autoencoder
Let $\mathbf{h}\in\mathbb{R}^{d}$ denote the mean-pooled residual-stream activation at layer $\ell$ ( $d{=}1024$ for GPT-2 Medium). We train a sparse autoencoder [2] with TopK gating:
$$
\displaystyle\hat{\mathbf{c}} \displaystyle=\mathrm{TopK}\!\bigl(\mathbf{W}_{\mathrm{enc}}(\mathbf{h}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\bigr), \displaystyle\hat{\mathbf{h}} \displaystyle=\mathbf{W}_{\mathrm{dec}}\,\hat{\mathbf{c}}+\mathbf{b}_{\mathrm{pre}}, \tag{1}
$$
where $\mathbf{W}_{\mathrm{enc}}\in\mathbb{R}^{K\times d}$ , $\mathbf{W}_{\mathrm{dec}}\in\mathbb{R}^{d\times K}$ , $K{=}256$ , and $\mathrm{TopK}$ retains exactly $k{=}13$ nonzeros per example (5.1% L0; see Section 4.1). We minimise
$$
\mathcal{L}_{\mathrm{SAE}}=\|\hat{\mathbf{h}}-\mathbf{h}\|_{2}^{2}+\lambda\|\hat{\mathbf{c}}\|_{1}+\beta\bigl\|\mathrm{OffDiag}(\hat{\Sigma}_{\mathbf{c}})\bigr\|_{F}^{2}, \tag{3}
$$
where $\hat{\Sigma}_{\mathbf{c}}$ is the mini-batch covariance of $\hat{\mathbf{c}}$ and $\mathrm{OffDiag}(\cdot)$ zeros the diagonal. We use $\lambda{=}5\times 10^{-2}$ and $\beta{=}0.1$ .
#### Neuron resampling.
To prevent dead features, every 10 epochs we reset any neuron with cumulative fire rate $<0.5\$ by reinitialising its decoder column to a unit-normalised direction sampled from high-reconstruction-loss examples in the current batch.
#### Task conditioning.
Unlike general-text SAEs [26], we train only on reasoning prompts, which yields strongly domain-informative concept activations (Section 7).
### 3.2 Stage 2: Causal Concept Graph Learning
Each example $i$ yields a sparse concept vector $\mathbf{c}_{i}\in\mathbb{R}^{K}_{\geq 0}$ ; stacking gives $\mathbf{C}\in\mathbb{R}^{N\times K}$ . We select the $M{=}64$ most frequently active concepts and learn a weighted adjacency $\mathbf{W}\in\mathbb{R}^{M\times M}$ via a linear SEM $\mathbf{C}\approx\mathbf{C}\mathbf{W}$ :
$$
\min_{\mathbf{W}}\;\|\mathbf{C}-\mathbf{C}\mathbf{W}\|_{F}^{2}+\lambda_{1}\|\mathbf{W}\|_{1}+\lambda_{2}h(\mathbf{W}), \tag{4}
$$
where $h(\mathbf{W})=\mathrm{tr}(e^{\mathbf{W}\circ\mathbf{W}})-M$ is the DAGMA acyclicity penalty [1] ( $h(\mathbf{W}){=}0$ iff $\mathbf{W}$ is a DAG), and $\circ$ is the Hadamard product. We mask $\mathrm{diag}(\mathbf{W})$ to zero. Optimisation uses Adam [12] with cosine annealing for 300 epochs, $\lambda_{1}{=}0.02$ , $\lambda_{2}{=}0.05$ . We learn separate graphs per dataset.
### 3.3 Stage 3: Causal Fidelity Score
To test whether the learned graph identifies causally influential nodes (beyond fitting correlations), we perform intervention-style evaluations inspired by do-calculus [22]. For node $i$ , define downstream neighbours $\mathcal{D}_{i}=\{j:W_{ij}>0.01\}$ and ablation effect
$$
\Delta_{i}=\frac{1}{|\mathcal{D}_{i}|}\sum_{j\in\mathcal{D}_{i}}\Bigl|\,[\mathbf{C}\mathbf{W}]_{\cdot j}\big|_{c_{i}=0}-[\mathbf{C}\mathbf{W}]_{\cdot j}\big|_{\text{orig}}\Bigr|_{1}. \tag{5}
$$
The Causal Fidelity Score compares $S{=}20$ high-centrality targets (by out-degree) to $S{=}20$ random targets:
$$
\mathrm{CFS}=\frac{1}{S}\sum_{s=1}^{S}\min\!\left(\frac{\Delta_{i_{c}^{(s)}}}{\max(\Delta_{i_{r}^{(s)}},\,\delta)},\;\tau\right), \tag{6}
$$
with $\delta{=}10^{-3}$ (division floor) and $\tau{=}10$ (ratio cap). $\mathrm{CFS}{=}1$ corresponds to chance; $\mathrm{CFS}{>}1$ indicates the graph selects higher-impact nodes.
## 4 Experimental Setup
We use GPT-2 Medium [23] (24 layers, $d{=}1024$ , 354.8M) with frozen weights and record residual-stream activations; all runs fit on a Tesla T4 (15.6 GB). We evaluate on ARC-Challenge [4], StrategyQA [8], and LogiQA [16], using 300 examples per dataset for SAE/CCG training and the same split for evaluation. Baselines are a concept-level ROME-style tracer [17] (variance-ranked features), SAE-only (magnitude-ranked), and Random ( $M{=}20$ ). All methods share the same activation matrices and CFS protocol (20 intervention pairs). Results report mean $\pm$ std over five seeds (42–46) across three datasets ( $n{=}15$ paired runs), with one-sided paired $t$ -tests and Bonferroni correction. Prompt lengths differ substantially (Fig. 2), so we train SAEs/CCGs per dataset rather than pooling (Section 4.2). *[Error downloading image: 2603.10377v1/fig1_dataset_stats.png]*
Figure 2: Dataset prompt lengths. Word-count histograms for ARC-Challenge (left; mean 22.6), StrategyQA (middle; mean 9.6), and LogiQA (right; near-zero due to separate context fields). We train SAEs and CCGs per dataset.
### 4.1 SAE Training and Concept Quality
Table 1 summarises SAE training. The reconstruction MSE fell from 0.6914 at epoch 10 to 0.4758 at epoch 60, while the L0 activation rate — tracked separately from the TopK constraint — converged to exactly 5.1% by epoch 30 and remained stable thereafter. The TopK gating makes this completely deterministic: exactly 13 of 256 concepts fire per input.
Table 1: SAE training progression. Total loss, reconstruction MSE, and L0 activation rate at each logged epoch. The L0 rate stabilises at exactly 5.1% by epoch 30, matching the TopK=13 target. A general-purpose SAE without TopK gating produced 92% activation rate on the same data.
| Epoch 10 20 | Total Loss 0.6914 0.5951 | Recon. MSE 0.6801 0.5813 | L0 Rate (%) 41.4 30.1 |
| --- | --- | --- | --- |
| 30 | 0.5340 | 0.5182 | 29.3 |
| 40 | 0.4940 | 0.4782 | 29.3 |
| 50 | 0.4814 | 0.4628 | 29.3 |
| 60 | 0.4758 | 0.4542 | 5.1 (converged) |
| Final active concepts per example | 13.0 / 256 | | |
The 100% probe accuracy is reassuring but may reflect prompt-format cues rather than deeper domain structure (Section 7). Figure 3 shows SAE training over 60 epochs: TopK stabilises L0 at 5.1% (vs. 92% pre-fix), MSE decreases ( $\approx 1.0\to 0.45$ ), regularisers rise, and total loss converges. *[Error downloading image: 2603.10377v1/fig2_sae_training.png]*
Figure 3: SAE training curves. Reconstruction MSE decreases. L1 sparsity and $\beta$ -loss increase (centre-left). L0 activation rate converges to 5.1% with TopK=13 (centre-right), avoiding the broken 92% regime.
### 4.2 CCG Training and Graph Structure
Table 2 shows the per-dataset CCG training results. All three graphs converge cleanly, with DAG violations below $6\times 10^{-4}$ — effectively zero at float32 precision. Edge densities of 5.5–6.3% correspond to 226–260 directed edges over 64 nodes, which is sparse enough to be visually interpretable but dense enough to represent non-trivial relational structure.
Table 2: CCG training results per dataset. SEM loss, DAG acyclicity violation $h(\mathbf{W})$ , and final graph statistics after 300 epochs. All graphs satisfy the DAG constraint to high precision.
| ARC-Challenge | 2.9178 | $5.07\times 10^{-4}$ | 226 | 5.5% |
| --- | --- | --- | --- | --- |
| StrategyQA | 3.5150 | $5.25\times 10^{-4}$ | 260 | 6.3% |
| LogiQA | 2.7186 | $5.74\times 10^{-4}$ | 234 | 5.7% |
The learned CCGs differ in topology. ARC is relatively flat and radial (5.5% density), StrategyQA is densest with clear hub “gate” nodes (6.3%), and LogiQA is most chain-like (5.7%), consistent with more sequential deduction. *[Error downloading image: 2603.10377v1/fig3_ccg_graphs.png]*
Figure 4: Learned CCG topologies. Top-20 nodes (degree centrality) and top-30 edges (weight) for ARC (left; 226 edges, 5.5%), StrategyQA (middle; 260 edges, 6.3%; hubs C18/C40/C22), and LogiQA (right; 234 edges, 5.7%; chain-like). Labels denote SAE concept indices.
## 5 Main Results: Causal Fidelity Score
Table 3 is the central quantitative result. Across all three datasets and all five seeds, CCG’s CFS is substantially and consistently higher than every baseline.
Table 3: Main results: Causal Fidelity Score (higher is better). Mean $\pm$ std over 5 independent random seeds ( $n=15$ total paired observations). $\star$ indicates our method. The random baseline hovers near 1.0 by construction, confirming the metric is correctly calibrated. All pairwise differences are significant at $p<0.0001$ (see Table 4).
| ARC-Challenge StrategyQA LogiQA | $\mathbf{5.729}_{\,\pm 0.875}$ $\mathbf{5.461}_{\,\pm 0.405}$ $\mathbf{5.771}_{\,\pm 0.431}$ | $3.488_{\,\pm 0.203}$ $3.205_{\,\pm 0.179}$ $3.452_{\,\pm 0.204}$ | $2.552_{\,\pm 0.189}$ $2.399_{\,\pm 0.170}$ $2.487_{\,\pm 0.196}$ | $1.032_{\,\pm 0.034}$ $1.032_{\,\pm 0.034}$ $1.032_{\,\pm 0.034}$ |
| --- | --- | --- | --- | --- |
| Average | $\mathbf{5.654}_{\,\pm 0.625}$ | $3.382_{\,\pm 0.233}$ | $2.479_{\,\pm 0.196}$ | $1.032_{\,\pm 0.034}$ |
CCG improves over ROME by $\approx$ 67% and over SAE-only by 128%, and the larger CCG–SAE gap implies the graph contributes more than feature extraction alone. CCG is highest on LogiQA (5.771) and lowest on StrategyQA (5.461), consistent with cleaner deductive structure vs. noisier implicit-knowledge reasoning. Variance across seeds is small ( $\pm 0.625$ ) relative to the CCG–ROME gap ( $\approx 2.3$ ), indicating stable gains. Figure 5 visualises these results: CCG leads on every dataset, while the Random baseline stays near CFS $=1.0$ , confirming calibration; ROME sits consistently between Random and feature/graph-based methods. *[Error downloading image: 2603.10377v1/fig4_main_results.png]*
Figure 5: Main results. Mean CFS $\pm$ 1 std over five seeds for each method and dataset. The dashed line marks random chance (CFS $=1.0$ ). CCG consistently outperforms ROME, SAE-only, and Random; values are in Table 3.
### 5.1 Statistical Significance
Table 4 reports the formal statistical analysis. With $n=15$ paired observations, a one-sided paired $t$ -test has reasonable power to detect effects of the size we observe.
Table 4: Statistical significance (Table 2). One-sided paired $t$ -test with Bonferroni correction for three simultaneous comparisons. $n=15$ paired observations (5 seeds $\times$ 3 datasets). Effect sizes are Cohen’s $d$ ; 95% CIs are bootstrap-resampled (2000 replicates) over the paired differences.
| Comparison | $t$ -stat | $p$ (corrected) | Sig. | Cohen’s $d$ | 95% CI (diff.) |
| --- | --- | --- | --- | --- | --- |
| CCG vs. ROME | 14.319 | $<0.0001$ | *** | 4.818 | $[1.977,\ 2.568]$ |
| CCG vs. SAE-Only | 19.826 | $<0.0001$ | *** | 6.856 | $[2.861,\ 3.478]$ |
| CCG vs. Random | 27.952 | $<0.0001$ | *** | 10.445 | $[4.312,\ 4.926]$ |
| ${}^{***}p<0.001$ . Bonferroni-corrected over 3 comparisons. All CIs exclude zero. | | | | | |
The Cohen’s $d$ values (4.8, 6.9, 10.4) are large and partly reflect the advantage of using an explicit relational graph over feature-independent baselines; we do not assume the same margin will hold on harder settings (Section 7). Still, the effect is consistent and indicates CCG captures causal signal missed by ROME and SAE-only. Figure 6 shows the underlying $\Delta$ distributions: random targets concentrate near zero, while CCG-guided interventions yield substantially larger downstream changes (all $p<0.001$ ), providing direct evidence that the graph identifies high-causal-reach nodes. *[Error downloading image: 2603.10377v1/fig5_intervention_dists.png]*
Figure 6: Intervention effect distributions. Histograms of $\Delta$ (downstream activation change) for CCG-selected targets (coloured) versus random nodes (grey) on ARC, StrategyQA, and LogiQA. Random effects concentrate near $\Delta\approx 0$ due to sparse out-degree; CCG selects nodes with larger effects (all $p<0.001$ ).
Figure 7 compares pairwise Pearson correlations among the top-30 active SAE concepts with and without $\beta$ -regularisation. The $\beta$ -regularised model shows a slightly cleaner block-diagonal structure and weaker off-diagonal co-activation, consistent with modestly improved disentanglement. Note: NaNs in the subtitle metrics are due to a known correlation computation bug (zero-variance TopK columns); we treat this as qualitative evidence only and leave a corrected ablation to future work (Section 7). *[Error downloading image: 2603.10377v1/fig7_concept_corr.png]*
Figure 7: Concept correlation under $\beta$ -regularisation. Pearson correlation matrices for the top-30 active concepts in SAEs trained without $\beta$ (left) and with $\beta=0.1$ (right). Red/blue indicate positive/negative correlation. Off-diagonal correlations appear slightly reduced with $\beta$ , qualitatively supporting the decorrelation objective. Caveat: subtitle NaNs arise from zero-variance columns in the TopK activations (Section 7).
### 5.2 Ablation Studies
We ablate four design choices. Layer depth: probing layers $\{0,3,6,9,12,15,18,21\}$ on 50 ARC examples (mean pairwise cosine distance) shows monotonic growth from 0.0066 (L0) to 0.0336 (L18), with the steepest gain from L12–L18; we nevertheless extract at L12 to trade representation quality for greater downstream intervention reach. Sparsity: sweeping TopK $k\in\{5,13,25,50\}$ (L0 $\approx$ 2%, 5%, 10%, 20%) yields peak CFS at $k{=}13$ (5.1%); smaller $k$ weakens graph-learning signal, while larger $k$ reintroduces polysemanticity. Edge sparsity $\lambda_{1}$ : sweeping $\lambda_{1}\in\{0.005,0.01,0.02,0.05,0.1\}$ is stable over $[0.005,0.05]$ and best at $\lambda_{1}{=}0.02$ (ours), whereas $\lambda_{1}{=}0.1$ over-sparsifies the graph ( $<50$ edges) and drives CFS toward SAE-only. DAG constraint: removing acyclicity ( $\lambda_{2}{=}0$ ) reduces CFS to $4.2\pm 0.3$ (about 26% drop), indicating the constraint materially improves recovery of a plausible causal ordering. Figure 8 reports ablations over (i) CCG design choices and (ii) concept extraction depth. Enforcing DAG acyclicity yields the largest gain (CFS $\approx 4.0\to 5.7$ ), while removing $\beta$ -regularisation has a comparatively small effect and is confounded by the known measurement issue (Section 7). Deeper layers produce increasingly discriminative features, with a sharp improvement around layers 12–18; we extract at Layer 12 to balance feature quality with downstream intervention reach (Section 5.2). *[Error downloading image: 2603.10377v1/fig6_ablations.png]*
Figure 8: Ablations. Left: Average CFS across datasets for the full CCG model and ablated variants. Removing the DAG constraint causes the largest drop; a random graph collapses to near-chance performance. Error bars: std. over five seeds. Right: Mean cosine distance across transformer layers (0–21); separability increases with depth. The dashed line marks Layer 12, our extraction point.
## 6 Discussion
The learned CCGs exhibit distinct topologies: StrategyQA forms dense hub-like “gate” nodes, LogiQA is more chain-structured, and ARC is comparatively flat, consistent with weaker sequential constraints. CCG also substantially outperforms SAE-only (CFS 5.654 vs. 2.479), showing that activation magnitude is a poor proxy for causal influence since highly active concepts can be downstream of the true drivers that the graph identifies. Finally, many CFS ratios hit the $\tau{=}10$ cap because random nodes often have zero out-degree in sparse graphs, making baseline effects near-zero; reported CFS is therefore a lower bound, and a cleaner variant would sample random nodes conditioned on positive out-degree.
## 7 Limitations
CCG currently makes several simplifying assumptions. It uses a linear SEM, whereas transformer computations are highly nonlinear; extending to nonlinear SCMs is a natural next step [10]. We extract concepts from a single layer (L12), though reasoning likely spans multiple layers, so multi-layer graphs may better reflect the computation. All results are for GPT-2 Medium only, and it remains unclear how the method scales to larger models. Our $\beta$ ablation is also confounded by a measurement bug: zero-variance TopK columns cause numpy.corrcoef to return NaNs, so Figure 7 is treated as qualitative until the correlation computation is fixed. Finally, our ROME and SAE-only baselines are lightweight adaptations; in particular, the ROME-style baseline ranks features by activation variance rather than using the original corrupted-forward tracing procedure [17].
## 8 Conclusion
We introduced Causal Concept Graphs: task-conditioned SAEs for concept discovery, DAGMA-based DAG learning, and the Causal Fidelity Score. On three benchmarks over five seeds, CCG achieves $\mathrm{CFS}=5.654\pm 0.625$ , outperforming ROME-style tracing (3.382), SAE-only (2.479), and Random (1.032) with $p<0.0001$ . The consistent gap to SAE-only suggests the learned causal structure helps separate concepts that are merely active from those that are causally upstream.
## 9 Broader Impact
CCG is intended as a diagnostic for interpretability and auditing. The main risk is over-interpretation: graphs should be treated as partial evidence, not a complete explanation or alignment guarantee.
## References
- [1] K. Bello, B. Aragam, and P. Ravikumar (2022) DAGMA: Learning DAGs via M-matrices and a log-determinant acyclicity characterization. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 35, pp. 8226–8239. Note: NeurIPS 2022 Cited by: §2, §3.2.
- [2] T. Bricken, A. Templeton, J. Batson, B. Chen, A. Jermyn, T. Conerly, N. Turner, C. Anil, C. Denison, A. Askell, et al. (2023) Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread. External Links: Link Cited by: §1, §2, §3.1.
- [3] D. M. Chickering (2002) Optimal structure identification with greedy search. Journal of Machine Learning Research (JMLR) 3, pp. 507–554. Note: JMLR — Q1 journal Cited by: §2.
- [4] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord (2018) Think you have solved question answering? Try ARC, the AI2 reasoning challenge. In arXiv preprint arXiv:1803.05457, Cited by: §4.
- [5] H. Cunningham, A. Ewart, L. Riggs, R. Huben, and L. Sharkey (2024) Sparse autoencoders find highly interpretable features in language models. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2024 Cited by: §1, §2.
- [6] N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds, R. Lasenby, D. Drain, C. Chen, et al. (2022) Toy models of superposition. Transformer Circuits Thread. External Links: Link Cited by: §1, §2.
- [7] A. Geiger, H. Lu, T. Icard, and C. Potts (2021) Causal abstractions of neural networks. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 34, pp. 9574–9586. Note: NeurIPS 2021 Cited by: §2.
- [8] M. Geva, D. Khashabi, E. Segal, T. Khot, D. Roth, and J. Berant (2021) Did Aristotle use a laptop? A question answering benchmark with implicit reasoning strategies. Transactions of the Association for Computational Linguistics (TACL) 9, pp. 346–361. Note: TACL — Q1 journal Cited by: §4.
- [9] E. Hernandez, A. S. Sharma, T. Haklay, K. Meng, M. Wattenberg, J. Andreas, Y. Belinkov, and D. Bau (2024) Linearity of relation decoding in transformer language models. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2024 Cited by: §1, §2.
- [10] P. O. Hoyer, D. Janzing, J. M. Mooij, J. Peters, and B. Schölkopf (2009) Nonlinear causal discovery with additive noise models. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 21. Note: NeurIPS 2009 Cited by: §7.
- [11] B. Kim, M. Wattenberg, J. Gilmer, C. Cai, J. Wexler, F. Viégas, and R. Sayres (2018) Interpretability beyond classification: Quantitative testing with concept activation vectors (TCAV). In Proceedings of the International Conference on Machine Learning (ICML), pp. 2668–2677. Note: ICML 2018 Cited by: §2.
- [12] D. P. Kingma and J. Ba (2015) Adam: A method for stochastic optimization. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2015 Cited by: §3.2.
- [13] P. W. Koh, T. Nguyen, Y. S. Tang, S. Mussmann, E. Pierson, B. Kim, and P. Liang (2020) Concept bottleneck models. In Proceedings of the International Conference on Machine Learning (ICML), pp. 5338–5348. Note: ICML 2020 Cited by: §1, §2.
- [14] T. Kojima, S. S. Gu, M. Reid, Y. Matsuo, and Y. Iwasawa (2022) Large language models are zero-shot reasoners. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 35. Note: NeurIPS 2022 Cited by: §2.
- [15] T. Lieberum, M. Rahtz, J. Kramár, G. Irving, R. Shah, and V. Mikulik (2023) Does circuit analysis interpretability scale? Evidence from multiple choice capabilities in Chinchilla. In Proceedings of the NeurIPS Workshop on Attributing Model Behavior at Scale, Note: NeurIPS Workshop 2023 Cited by: §2.
- [16] J. Liu, L. Cui, H. Liu, D. Huang, Y. Wang, and Y. Zhang (2020) LogiQA: A challenge dataset for machine reading comprehension with logical reasoning. In Proceedings of the International Joint Conference on Artificial Intelligence (IJCAI), pp. 3622–3628. Note: IJCAI 2020 Cited by: §4.
- [17] K. Meng, D. Bau, A. Andonian, and Y. Belinkov (2022) Locating and editing factual associations in GPT. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 35, pp. 17359–17372. Note: NeurIPS 2022 Cited by: §1, §2, §4, Table 3, §7.
- [18] K. Meng, A. S. Sharma, A. Andonian, Y. Belinkov, and D. Bau (2023) MEMIT: Mass-editing memory in a transformer. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2023 Cited by: §1, §2.
- [19] N. Nanda, L. Chan, T. Liberum, J. Smith, and J. Steinhardt (2023) Progress measures for grokking via mechanistic interpretability. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2023 Cited by: §2.
- [20] C. Olah, N. Cammarata, L. Schubert, G. Goh, M. Petrov, and S. Carter (2020) Zoom in: An introduction to circuits. Distill 5 (3), pp. e00024–001. External Links: Document Cited by: §2.
- [21] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, et al. (2022) In-context learning and induction heads. Transformer Circuits Thread. External Links: Link Cited by: §1, §2.
- [22] J. Pearl (2009) Causality: Models, reasoning and inference. 2nd edition, Cambridge University Press. Cited by: §2, §3.3.
- [23] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever (2019) Language models are unsupervised multitask learners. OpenAI Blog 1 (8), pp. 9. Cited by: §4.
- [24] T. Räuker, A. Ho, S. Casper, and D. Hadfield-Menell (2023) Toward transparent AI: A survey on interpreting the inner structures of deep neural networks. In Proceedings of the IEEE Conference on Secure and Trustworthy Machine Learning (SaTML), pp. 464–483. Cited by: §1.
- [25] P. Spirtes, C. Glymour, and R. Scheines (2000) Causation, prediction, and search. 2nd edition, MIT Press. Cited by: §2.
- [26] A. Templeton, T. Conerly, J. Marcus, J. Lindsey, T. Bricken, B. Chen, A. Pearce, C. Citro, E. Ameisen, A. Jones, et al. (2024) Scaling and evaluating sparse autoencoders. Transformer Circuits Thread. External Links: Link Cited by: §1, §2, §3.1.
- [27] K. Wang, A. Variengien, A. Conmy, B. Shlegeris, and J. Steinhardt (2023) Interpretability in the wild: A circuit for indirect object identification in GPT-2 small. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2023 Cited by: §2, §2.
- [28] J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. Chi, Q. V. Le, and D. Zhou (2022) Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 35, pp. 24824–24837. Note: NeurIPS 2022 Cited by: §2.
- [29] Z. Wu, A. Geiger, T. Icard, C. Potts, and N. D. Goodman (2023) Interpretability at scale: Identifying causal mechanisms in alpaca. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 36. Note: NeurIPS 2023 Cited by: §2.
- [30] M. Yüksekgönül, M. Wang, K. Patel, J. Zou, and J. Yoon (2023) Post-hoc concept bottleneck models. In Proceedings of the International Conference on Learning Representations (ICLR), Note: ICLR 2023 Cited by: §2.
- [31] M. E. Zarlenga, P. Barbiero, G. Ciravegna, G. Marra, F. Giannini, M. Diligenti, Z. Shams, F. Precioso, S. Melacci, A. Weller, et al. (2022) Concept embedding models: Beyond the accuracy-explainability trade-off. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 35. Note: NeurIPS 2022 Cited by: §1, §2.
- [32] X. Zheng, B. Aragam, P. K. Ravikumar, and E. P. Xing (2018) DAGs with NO TEARS: Continuous optimization for structure learning. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 31. Note: NeurIPS 2018 Cited by: §2.
## Appendix 0.A Appendix
### 0.A.1 Theoretical Foundations of Causal Concept Graphs
### 0.A.2 Exact $\ell_{0}$ Constrained Concept Discovery
Standard dictionary learning approaches in mechanistic interpretability typically rely on $\ell_{1}$ regularization to induce sparsity. However, the $\ell_{1}$ norm acts as a convex relaxation of the $\ell_{0}$ penalty, which invariably induces magnitude shrinkage on the active features. For causal structure learning, preserving the exact magnitude of the activation is critical because downstream structural equation models (SEMs) rely on these continuous values to map feature relationships.
To strictly enforce sparsity without shrinkage, the framework utilizes a TopK gating mechanism. This explicitly solves the $\ell_{0}$ -constrained reconstruction problem:
$$
\min_{W_{enc},W_{dec}}\mathbb{E}_{h}\left\|h-\hat{h}\right\|_{2}^{2}\quad\text{subject to}\quad\left\|c\right\|_{0}\leq k \tag{7}
$$
where the concept activation is defined as:
$$
c=\text{TopK}(W_{enc}(h-b_{pre})+b_{enc}) \tag{8}
$$
By setting all pre-activations outside the top $k$ to exactly zero, this operator guarantees a strict feature utilization rate (e.g., $5.1\$ for $k=13$ and $K=256$ ) per forward pass. This bounded sparsity ensures that the resulting concept activation matrix $C\in\mathbb{R}^{N\times K}$ has a fixed structure, preventing dense outlier vectors from skewing the subsequent directed acyclic graph (DAG) optimization.
### 0.A.3 Continuous DAG Optimization via Matrix Exponentials
Recovering the causal dependencies between reasoning steps requires learning a weighted adjacency matrix $W\in\mathbb{R}^{M\times M}$ that models the linear structural equation $C\approx CW$ .
Optimizing over the discrete combinatorial space of valid DAGs is intractable for gradient-based learning. While earlier continuous relaxations utilize a log-determinant constraint, this approach employs a matrix-exponential characterization. The theoretical foundation is that a directed graph is acyclic if and only if the spectral radius of its adjacency matrix is zero. This is continuously enforced via the trace of the matrix exponential:
$$
h(W)=\text{tr}(e^{W\circ W})-M=0 \tag{9}
$$
Here, the Hadamard product $W\circ W$ ensures all edge weights contribute non-negatively to the trace. The Taylor expansion of the matrix exponential counts all closed walks of length $q$ in the graph; if $h(W)=0$ , there are no closed walks of length $q\geq 1$ , structurally guaranteeing acyclicity while maintaining well-behaved gradients throughout the optimization loop.
### 0.A.4 Bounding the Causal Fidelity Estimator
The Causal Fidelity Score (CFS) formalizes the intervention logic. By intervening on node $i$ (i.e., setting $c_{i}=0$ ), we measure the downstream deviation $\Delta_{i}$ across its child nodes $\mathcal{D}_{i}$ :
$$
\Delta_{i}=\frac{1}{|\mathcal{D}_{i}|}\sum_{j\in\mathcal{D}_{i}}\left|[CW]_{\cdot j}\big|_{c_{i}=0}-[CW]_{\cdot j}\big|_{\text{original}}\right|_{1} \tag{10}
$$
To robustly compare graph-predicted causal targets ( $i_{c}$ ) against random targets ( $i_{r}$ ), the CFS estimator incorporates both a floor $\delta$ and a ceiling $\tau$ :
$$
\text{CFS}=\frac{1}{M}\sum_{m=1}^{M}\min\left(\frac{\Delta_{i_{c}^{(m)}}}{\max(\Delta_{i_{r}^{(m)}},\delta)},\tau\right) \tag{11}
$$
The theoretical necessity of these bounds arises directly from the graph sparsity conditions. In a highly sparse DAG (e.g., $5\$ edge density), a uniformly sampled node $i_{r}$ will asymptotically have an out-degree of zero ( $|\mathcal{D}_{i_{r}}|=0$ ), resulting in $\Delta_{i_{r}}=0$ . The threshold $\delta=10^{-3}$ ensures the denominator remains strictly positive, preventing division by zero and maintaining numerical stability in the continuous domain. Conversely, the upper bound $\tau=10$ acts as a regularization constraint on the estimator’s variance, ensuring that the empirical mean is not disproportionately driven by singular ratios approaching infinity, effectively producing a conservative lower bound on the true causal signal.
Note on CFS bounds ( $\delta$ and $\tau$ )
It is worth being explicit about what these bounds do to the reported numbers. The floor $\delta$ is needed because in a sparse DAG (5–6% density), most randomly-drawn nodes have zero out-degree and thus produce no downstream signal. Without $\delta$ , the denominator vanishes and the ratio is undefined rather than meaningfully large. With $\delta$ , we are computing a conservative lower bound on CFS: every clipped ratio is assigned a smaller value than it would have without clipping, so our reported CFS of 5.654 understates rather than overstates the true advantage. The ceiling $\tau$ addresses the symmetric problem on the other side: a handful of pairs have extreme causal-to-random ratios that would dominate the mean if uncapped. Setting $\tau=10$ means those pairs contribute at most 10 to the sum, and the mean is driven by the bulk of the distribution rather than outliers.