Joint Causal Structure and Cluster Discovery Using Variational Inference

Avni Rajpal; Anubhav Kumar; Rishabh Karnad; Mohammad Emtiyaz Khan; P. K. Srijith

cs.LGAug 23, 2026

Joint Causal Structure and Cluster Discovery Using Variational Inference

Authors: Avni Rajpal, Anubhav Kumar, Rishabh Karnad, Mohammad Emtiyaz Khan, P. K. Srijith

Organizations: Department of Artificial Intelligence, Indian Institute of Technology Hyderabad, India · RIKEN Center for AI Project, Tokyo, Japan

Abstract

Causal discovery aims to understand the relationships between individual random variables. In many applications, such as brain imaging and climate modeling, it is more meaningful to consider interactions among groups of variables. Existing methods assume that knowledge of such groups or clusters is explicitly available when modeling interactions. However, in practice, these clusters as well as the causal relationships among them, are latent. In this paper, we present a novel approach based on variational inference to simultaneously infer both the latent clusters and causal structures. We learn an approximate posterior over clusters and graph-structure by considering variational distributions based on categorical and Bernoulli models respectively. We derive variational lower bounds and estimation techniques to learn variational and model parameters. The effectiveness of our proposed methods for cluster and causal discovery are demonstrated on both synthetic and real data sets.

Figures & tables

Figure 1: Example of a cluster-DAG

Figure 2: Graphical model of the generative process

Figure 3: Example of a) C-DAG b) Expanded Graph

Table 1: Results on linear synthetic datasets. ExG-SHD, E-SHD, and HD: lower is better; E-RI and AUC: higher is better. NA: Not Applicable.

Table 2: Results on nonlinear synthetic datasets. ExG-SHD and E-SHD: lower is better; E-RI and AUC: higher is better.

Data- sets	Methods	ExG- SHD	E- RI	AUC	E- SHD	HD
Protein Dataset	Baseline	3.0	0.55	1.0	0.3	NA
	Lin Fact-C-E	27.4	0.57	1.0	1.0	1.0
	Lin AR-C-Fact-E	0.0	0.56	1.0	0.0	0.7
	Neur Fact-C-E	0.04	0.74	1.0	0.0	NA
	Neur AR-C-Fact-E	0.0	0.56	1.0	0.0	NA
	Neur Fact-C-Cond-E	0.0	0.80	1.0	0.0	NA

Table 3: Results on real-world datasets. NA means Not Applicable. ExG-SHD, E-SHD, HD lower is better. E-RI and AUC higher is better.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Data Generating CDAG for 4 variable 3 cluster synthetic dataset a) fork b) chain c) v-structure

Figure 5: Data Generating CDAG for 7 variable 3 cluster synthetic dataset a) fork b) chain c) v-structure

Dataset		C 1	C 2
4 variable 3 cluster	fork	Tanh	Tanh
	chain	Sin	ReLU
	v-structure	-	Cubic Polynomial
7 variable 3 cluster	fork	Sin	ReLU
	chain	ReLU	Sin
	v-structure	-	Tanh

Appendix

Table 4: Element-wise non-linearities applied after linear combination of parents

Panel of Figure 6 — Figure 6: Sachs Dataset

Figure 7: Ground truth C-DAG for climate dataset

Figure 8: Autoregressive generation of cluster assignments using the linear model

Dataset		Hidden-State Dimension
4 variable 3 cluster	fork	32
	chain	4
	v-structure	4
7 variable 3 cluster	fork	8
	chain	128
	v-structure	48

Appendix

Table 5: Linear Autoregressive Model hidden-state dimension choices

Dataset		E-RI±std	E-SHD±std
4 variable 3 cluster	fork	0.7±0.1	1.08±0.51
	chain	0.75±0.12	0.94±0.64
	v-structure	0.72±0.08	0.95±0.21
7 variable 3 cluster	fork	0.68±0.07	1.05±0.22
	chain	0.60±0.06	0.97±0.5
	v-structure	0.65±0.06	1.0±0.0

Appendix

Table 6: Variation across random restarts for autoregressive linear cluster model for selected hidden dimension size

Causal Discovery

Causal discovery infers cause-and-effect structures from observational and interventional data, distinguishing causal relationships from correlation.

See all 221 papers

From Local to Cluster: A Unified Framework for Causal Discovery with Latent Variables
Apr 24, 2026Zongyu LiCausal Effect Estimation Causal Discovery
Structure Learning on Clustered Data
Jul 9, 2026Ryan Thompson, Matt P. Wand, Veerabhadran BaladandayuthapaniCausal Structure Learning Causal Identifiability
School of Mathematical and Physical Sciences University of Technology Sydney Ultimo, NSW 2007, Australia · Department of Biostatistics University of Michigan Ann Arbor, MI 48105, USA
Interpretable Causal Discovery via Causal-Effect Constraints
Aug 12, 2026Cixuan Zhang, Guy Van den Broeck, Benjie WangBayesian Inference Causal Discovery
Computer Science Dept., Yale University, New Haven, Connecticut, USA · Computer Science Dept., University of California, Los Angeles, Los Angeles, California, USA

Apr 24, 2026J/K move · Enter open · S save

From Local to Cluster: A Unified Framework for Causal Discovery with Latent Variables

Zongyu Li

Causal Effect Estimation Causal Discovery Causal Abstraction Causal Discovery with Latent Confounding

Open paper AskPDF arXiv

Latent variables pose a fundamental challenge to causal discovery and inference. Conventional local methods focus on direct neighbors but fail to provide macro level insights. Cluster level methods enable macro causal reasoning but either assume clusters are known a priori or require causal sufficiency. Moreover, directly applying single variable causal discovery methods to cluster level problems violates causal sufficiency and leads to incorrect results. To overcome these limitations, this paper proposes L2C (Local to Cluster Causal Abstraction), a unified framework that bridges local structure learning and cluster level causal discovery. Unlike prior work that requires a complete manual assignment of micro variables to clusters, L2C discovers the partition automatically from local causal patterns. Our solution leverages a cluster reduction theorem to reduce any cluster to at most three nodes without loss of causal information, applies local causal discovery to identify direct causes, effects, and V structures in the presence of latent variables, and performs macro level causal inference via cluster level calculus on the learned cluster graph. L2C does not assume causal sufficiency, as latent variables are handled through local discovery. Theoretical analysis shows that L2C ensures soundness, atomic completeness, and computational efficiency. Extensive experiments on synthetic and real world data demonstrate that L2C accurately recovers ground truth clusters and achieves superior macro causal effect identification compared to existing baselines.

Abstract

Causal Discovery

Explore similar work

From Local to Cluster: A Unified Framework for Causal Discovery with Latent Variables

Structure Learning on Clustered Data

Interpretable Causal Discovery via Causal-Effect Constraints