causal-concept-graphs-latent-space

name	causal-concept-graphs-latent-space
title	Causal Concept Graphs in LLM Latent Space for Stepwise Reasoning
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2603.10377
keywords	["Interpretability","Causal Analysis","LLM","Mechanistic Interpretability","Reasoning"]
description	Extract sparse causal concept graphs from LLM activations using SAE and DAGMA, then validate through ablation to identify causally influential features. Bridges mechanistic interpretability with causal inference for understanding reasoning flow.

Technique: Learning Causal Structures in Monosemantic Concept Space

Understanding how LLMs perform reasoning requires identifying not just what features matter, but how they causally interact. This approach extracts sparse, interpretable concepts via Sparse Autoencoders (SAEs), learns a directed acyclic graph (DAG) over them using DAGMA, then validates graph edges through targeted interventions to ensure structural causality rather than mere correlation.

This bridges mechanistic interpretability with causal inference, enabling step-by-step analysis of reasoning processes.

Core Concept

The method operates in three stages:

1. Monosemantic Concept Extraction: SAE with TopK gating produces consistently sparse concept activations (e.g., 13 active per example)

2. DAG Learning via DAGMA: Linear structural equation model over top-64 concepts, enforcing acyclicity with matrix exponentials

3. Causal Fidelity Validation: Ablation experiments measuring how concept removal affects downstream behavior, distinguishing causal from spurious edges

Architecture Overview

• SAE backbone: Encoder-decoder architecture with TopK gating
• Concept bank: 1000s of learned concept vectors, sparsely activated
• DAG learner (DAGMA): Linear structural equation model with acyclicity constraint
• Intervention system: Ablation mechanism for causal fidelity scoring
• Attribution analysis: Identify high-centrality causal concepts per reasoning step

Implementation Steps

Step 1: Extract Sparse Monosemantic Concepts

Use SAE with TopK gating to maintain consistent sparsity without magnitude shrinkage.

import torch
import torch.nn as nn

class SAEWithTopKGating(nn.Module):
    def __init__(self, activation_dim=2048, num_concepts=1024, k_active=13):
        super().__init__()
        self.activation_dim = activation_dim
        .num_concepts = num_concepts
        .k_active = k_active

        
        .encoder = nn.Linear(activation_dim, num_concepts)

        
        .decoder = nn.Linear(num_concepts, activation_dim)

     ():
        
        batch_size, seq_len, _ = activations.shape

        
        concept_logits = .encoder(activations)  

        
        topk_values, topk_indices = torch.topk(
            concept_logits,
            k=.k_active,
            dim=-
        )

        
        sparse_concepts = torch.zeros_like(concept_logits)
        sparse_concepts.scatter_(-, topk_indices, topk_values)

        
        reconstructed = .decoder(sparse_concepts)

        
        sparsity = .k_active / .num_concepts

         reconstructed, sparse_concepts, sparsity