PaperlayerView exact revision on arXiv

Source-linked AI summary

Joint Causal Structure and Cluster Discovery Using Variational Inference

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

arXiv:2608.22212v1cs.LGcs.AIstat.ML

TL;DR

Causal discovery often needs to model interactions among variable groups, but the clusters and their causal relationships may both be unknown. This paper uses variational inference to infer clusters and causal structure jointly, with experiments showing effective recovery on synthetic and real datasets.

  • Problem

    Cluster-DAG discovery requires inferring both the partition of variables into clusters and the causal connections among those clusters from observational data.

    Source: §2.1 · p. 2

  • Method

    The paper uses variational inference to learn approximate posteriors over latent clusters and cluster-DAG structures for linear and nonlinear vector-valued structural causal models.

    Source: §1 · p. 2, §5 · p. 14

  • Results

    Across synthetic and real datasets, the proposed approaches show effective cluster and graph recovery compared with baselines in the reported experiments.

    Source: §4.3 · p. 13, §4.4 · p. 13, §5 · p. 14

  • Takeaways & Limitations

    The experiments support joint cluster and causal-structure discovery for linear and nonlinear vector-valued models on the tested synthetic and real datasets.

    Source: §5 · p. 14

    Exact posterior computation requires summing over all C-DAGs, and the search space grows with the number of variables and clusters.

    Source: §3.3.2 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Group-level causal questions motivate moving beyond scalar-variable models, yet cluster structure is often unknown. This paper uses variational inference to jointly infer latent clusters and causal structure from observational data.

1 Introduction

  • Causal discovery learns DAG structure from observational data, but scalar-output SCMs may not suit vector-valued mechanisms or group-level questions such as dependencies between brain regions.The paper motivates this setting with brain-region rather than individual-neuron dependencies.Source: §1 · p. 1
  • Most C-DAG methods assume known groupings, while Niu et al. (2022) proposed score-based joint cluster/DAG learning, particularly for distributions unfaithful to the full DAG.Related work also considers inference over a known pair of clusters (Wahl et al., 2023) and grouping variables in the presence of exogenous confounding (Kawahara et al., 2010).Source: §1 · pp. 1–2
  • The paper jointly infers latent clusters and causal structures from observational data in linear and nonlinear Gaussian additive-noise vSCMs using variational inference.It learns approximate posterior distributions over both latent components; the Bayesian treatment can support uncertainty quantification and active interventions, and experiments on synthetic and real data demonstrate effectiveness.Source: §1 · p. 2

2 Background

The section defines cluster-DAG discovery as jointly inferring a partition of variables and the directed acyclic graph over its clusters from observational samples. It introduces vector-valued causal models and linear-Gaussian, nonlinear additive-noise, and marginal-likelihood formulations.

2.1 Problem Statement

  • Cluster-DAG discovery infers a partition of variables and the directed acyclic graph over its clusters from observational samples.Each cluster is a disjoint subset of variable indices.Source: §2.1 · p. 2

2.2 Vector-Valued Structural Causal Models

  • A vector-valued structural causal model assigns each cluster from parent variables with potentially dependent noise, inducing a joint distribution that factorizes according to the clusters and graph.The model's functions map parent variables to cluster-valued variables.Source: §2.2 · p. 2

2.2.1 Linear Gaussian Vector Valued SCMs

  • In a linear-Gaussian vSCM, each cluster is a linear mapping of its parents with zero-mean multivariate Gaussian noise, yielding a product-of-Gaussians distribution.The noise covariance is cluster-specific.Source: §2.2.1 · pp. 2–3

2.2.2 Non-Linear Additive Noise Model

  • The nonlinear additive-noise model uses a neural network for each cluster's parent-to-cluster function, implemented as one network with inputs and outputs masked by cluster and parents (Lorch et al., 2021; Zheng et al., 2020).The network mean for each cluster is given by its parent variables and cluster-specific parameters.Source: §2.2.2 · p. 3

2.3 Marginal Likelihood of DAG

  • For linear-Gaussian DAGs, Normal-Wishart parameter priors allow the Gaussian model parameters to be marginalized out, yielding a closed-form marginal likelihood.The cited formulation applies to observations restricted to a subset of variables and their parents.Source: §2.3 · p. 3

3 Methodology

The method uses variational inference to jointly discover latent clusters and their causal structures from observational data. Its posterior over C-DAGs can represent multiple structures and quantify uncertainty.

3.1 Generative Model

  • Samples are modeled by a linear Gaussian vSCM or nonlinear additive model, with latent cluster assignments C and an upper-triangular edge matrix EC defining an acyclic cluster graph.The model infers C and EC from the observed data using Bayesian inference.Source: §3.1 · p. 4

3.2 Prior Distributions

  • The method uses a uniform prior over cluster assignments and a sparsity-enforcing prior over graph edges; the upper-triangular parameterization also ensures acyclicity.The graph prior includes a tunable sparsity hyperparameter λs.Source: §3.2 · pp. 4–5

3.3.1 Marginal Likelihood

  • For linear Gaussian models, integrating out causal parameters yields a closed-form marginal likelihood, but computing the posterior requires summing over all C-DAGs and is infeasible.This motivates approximating the posterior over clusters and structures with variational inference.Source: §3.3.1 · p. 5, §3.3.2 · p. 5

3.3.3 Variational Inference for C-DAGs

  • The method approximates the joint posterior over clusters and structures with a tractable variational family, maximizing the ELBO instead of directly minimizing KL divergence to the intractable posterior.The resulting distribution over C-DAGs can capture multiple structures and quantify uncertainty.Source: §3 · p. 4, §3.3.3 · pp. 5–6

3.4.1 Marginal Likelihood

  • For nonlinear vSCMs, marginalizing causal parameters is intractable, so the method point-estimates them and uses a variational posterior whose ELBO depends on those parameters.The posterior approximation follows the same variational approach, with a different ELBO.Source: §3.4.1 · p. 6, §3.4.2 · p. 6

3.5.1 Distribution over Clusters

  • For cluster assignments, the method considers a factorized categorical variational model and a linear autoregressive model that can capture dependence among variables.The autoregressive alternative addresses the limitation that independent assignments may not capture correlated cluster structure.Source: §3.5.1 · pp. 6–7

3.5.2 Distribution over graphs

  • For graph edges, the method considers independent binary, linear autoregressive, conditional, and autoregressive conditional variational models to capture edge dependencies and dependence on cluster assignments.The autoregressive conditional model combines graph topology with global cluster context.Source: §3.5.2 · pp. 7–8

3.6 Gradient Approximation of ELBO

  • The method optimizes the ELBO with gradient-based updates, estimating gradients for discrete C and EC using REINFORCE with an exponential moving average baseline and Monte Carlo sampling.The baseline is used for variance reduction.Source: §3.6 · pp. 8–9

4 Experimental Results

Experiments compare linear and nonlinear variational models with Niu et al. (2022) on synthetic and real-world data, using graph, cluster, and posterior-approximation metrics.

4.1 Evaluation Metrics

  • Evaluation uses HD for posterior distance when exact enumeration is feasible, E-SHD for EC accuracy, ExG-SHD for joint C and EC approximation error, AUROC for expanded-graph edge beliefs, and E[RI] for clustering similarity.The expanded graph is defined from C and EC; C-DAG enumeration is feasible up to d≤7 when k=4.Source: §4.1 · pp. 10–11

4.3 Synthetic Dataset

  • Our approach consistently outperforms baseline on tested 4- and 7-variable synthetic data; AR-C-Fact-E scores ExG-SHD 0.02 and E-RI 0.99 on the linear 4-variable fork, while several nonlinear variants score ExG-SHD 0.00 on 4-variable forks and v-structures.Source: §4.3 · pp. 12–13

4.2 Experimental Setup

  • The experiments compare linear and nonlinear models with the baseline, using 20 bootstrapped baseline datasets and 1000 iterations per run.Variants combine factorized or autoregressive cluster distributions with factorized, autoregressive, or conditional edge distributions.Source: §4.2 · pp. 11–12

4.4 Real World Datasets

  • On real-world datasets, our approach performs better than the baseline on most metrics; the climate Neural Conditional Edge model scores ExG-SHD 0.04, E-SHD 0.0, AUC 1.0, and E-RI 1.0 versus baseline E-RI 0.55.For protein data, the linear autoregressive cluster model achieves ExG-SHD and E-SHD of 0, while the Neural Conditional Edge model reaches E-RI 0.80 versus baseline 0.55.Source: Table 3 · p. 13, §4.4 · pp. 13–14

5 Conclusion

The paper develops variational inference for jointly learning cluster assignments and C-DAG structures, including alternative models for assignment and graph dependencies. It extends the framework to nonlinear vector-valued SCMs, reports effectiveness against baselines on synthetic and real data, and identifies scaling to high-dimensional datasets as future work.

  • The method learns posteriors over cluster assignments and C-DAG structures using categorical cluster distributions and Bernoulli edge distributions.It considers factorized and autoregressive linear assignment models, alongside independent binary, linear autoregressive, conditional, and autoregressive conditional graph models.Source: §5 · p. 14
  • The framework extends to nonlinear vector-valued SCMs by parameterizing causal mechanisms with neural networks, with inference techniques for both linear and nonlinear models.Source: §5 · p. 14
  • Experiments on synthetic and real datasets show that the proposed approaches are effective compared with baselines.Source: §5 · p. 14
  • Scaling the framework to high-dimensional datasets is identified as future work.Source: §5 · p. 14

A Proof of Proposition 1

The proof derives a lower bound on the marginal log likelihood of the data from an expression involving variational and model probabilities.

  • The expression log QZ (EC) + log QT (C) −log P (D, EC, C) + log P (D) gives a lower bound on the data’s marginal log likelihood.Source: A Proof of Proposition 1 · p. 16

B Proof of Proposition 2 · C Data Generating C-DAGs for Real and Synthetic Datasets · C.1 Data Generating C-DAG for 4 variable 3 cluster dataset

The proof section frames Proposition 2 through the data’s marginal log likelihood, while the data-generation section describes synthetic and real-world C-DAG settings and a four-variable, three-cluster example.

B Proof of Proposition 2

  • The proof considers the data’s marginal log likelihood and then invokes expectation in its derivation.The supplied proof passages provide only these steps, not the full derivation.Source: B Proof of Proposition 2 · p. 16

C Data Generating C-DAGs for Real and Synthetic Datasets

C.1 Data Generating C-DAG for 4 variable 3 cluster dataset

  • Figure 4 shows the data-generating C-DAGs for the dataset with d = 4 variables and k = 3 clusters.Source: §C.1 · p. 16

C.2 Data Generating C-DAG for 7 variable 3 cluster dataset · C.3 Non-Linearities for synthetic Non-Linear data

The appendix specifies data-generating C-DAGs for a 7-variable, 3-cluster dataset and describes how synthetic nonlinear data are generated from C-DAGs using element-wise nonlinearities after parent combinations.

C.2 Data Generating C-DAG for 7 variable 3 cluster dataset

C.3 Non-Linearities for synthetic Non-Linear data

  • Synthetic nonlinear data are generated from the C-DAGs in Figures 4 and 5.Source: §C.3 · p. 17
  • An element-wise nonlinear operation follows the linear combination of parents to ensure nonlinearity.Source: §C.3 · p. 17
  • Table 4 describes the nonlinearities applied after parent combinations.Source: Table 4 · p. 17, §C.3 · p. 17

C.4 Details for Protein Dataset

The protein-signalling evaluation uses 853 observations over 11 features and an expert-provided ground-truth graph. The graph indicates two disconnected clusters, and Figure 6b is used as the ground-truth C-DAG for metrics despite no true C-DAG being known.

C.4 Details for Protein Dataset

  • The Sachs protein-signalling dataset has 853 observations and 11 features, with an expert-provided ground-truth graph showing two disconnected clusters.One cluster contains Plcg, PIP3, and PIP2; all remaining nodes form the other cluster.Source: §C.4 · p. 17
  • Because the dataset has no known true C-DAG, Figure 6b is used as the ground-truth C-DAG for reporting metrics.Source: §C.4 · p. 17

C.5 Details for Climate dataset

The climate evaluation uses surface-temperature data from ENSO and British Columbia (BCT), with two latent clusters and an established causal influence from ENSO to BCT. The passages describe its preprocessing and evaluation setup but report no result for the proposed method.

C.5 Details for Climate dataset

  • The climate dataset covers ENSO and British Columbia (BCT), where ENSO’s causal influence on BCT is established; Wahl et al. (2023) recover this relationship with around 59% accuracy.The data contain two clusters, whose dimensionalities can be controlled by adjusting grid-box size.Source: §C.5 · p. 18
  • The dataset yields 73 annual temperature samples after deseasonalizing each sample by monthly means, smoothing with Gaussian kernel bandwidth σ = 120 months, and averaging October–December ENSO and January–March BCT temperatures; BCT is offset by one time lag.Source: §C.5 · p. 18
  • The evaluation considered surface temperature from four randomly chosen grid points in each region (d = 8, k = 4); Figure 7 shows the climate dataset’s ground-truth C-DAG used to calculate metrics.Source: §C.5 · p. 18, Figure 7 · p. 18

D Experimental Details

The experiments used tuned optimization settings and model-specific training protocols across datasets. Autoregressive models were trained with repeated restarts, while the neural formulation used a two-layer softplus MLP and longer dataset-dependent training runs.

Experimental Details

  • RMSprop at learning rate 0.01 with exponential decay (gamma 0.9 every 500 epochs) was optimal for autoregressive models; λs was 100 for the 7-variable, 3-cluster chain and 1 for other datasets.The search compared learning rates 0.1, 0.01, and 0.001; SGD, RMSprop, and Adam; and exponential and cosine-annealed schedulers.Source: D Experimental Details · p. 18
  • Autoregressive experiments used 10 random restarts and 10,000 epochs, Xavier-uniform initialization, and Tesla V100-SXM2-32GB GPUs; linear-model hidden dimensions were searched from 4 to 128.The selected hidden-state dimensions and variation across restarts are referenced in Tables 5 and 6, respectively.Source: D Experimental Details · p. 19
  • The neural formulation used a two-hidden-layer, 64-unit softplus MLP with d-dimensional input and output, AdamW weight decay, and dataset-specific λs values.λs was 1000 for protein signalling, 30 for climate, and 500 for synthetic nonlinear datasets; 64 was selected from hidden dimensions 16, 32, 64, and 128.Source: D Experimental Details · p. 19
  • Neural experiments ran 10 times for 75,000 epochs on real-world data and 50,000 on synthetic data, using cosine decay for the first half of training.The learning rate decayed from 5 × 10−3 to 2 × 10−5; ELBO used 30 samples, and performance metrics used 1,000 predictive-distribution samples.Source: D Experimental Details · p. 19
Loading 2608.22212v1…