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

The paper addresses joint discovery of latent variable clusters and causal relationships among those clusters from observational data. It uses variational inference to learn approximate posteriors over cluster assignments and C-DAG structures, with experiments showing effectiveness on synthetic and real datasets.

  • Problem

    The paper studies how to infer both variable clusters and the causal structure among them when neither is known from observational data.

  • Method

    The approach uses variational inference with categorical cluster assignments and Bernoulli edge distributions, covering linear and neural-network parameterized nonlinear vector-valued SCMs.

  • Results

    Experiments on synthetic and real datasets showed the proposed approaches were effective compared with baselines, including variants achieving ExG-SHD of 0.00 and near-perfect real-world structural scores.

  • Takeaways & Limitations

    The framework provides a Bayesian variational treatment for jointly learning latent clusters and causal structures, with posterior distributions that can support uncertainty quantification and active interventions.

  • Takeaways & Limitations

    Exact Bayesian posterior computation requires summing over all C-DAGs, making closed-form inference computationally difficult as the cluster and edge space grows.

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

The paper targets causal discovery when variables are better modeled in groups, but both the groups and causal relationships among them are latent. It proposes variational inference to jointly learn clusters and cluster-level causal structures from observational data.

  • Motivation: Causal discovery traditionally models relationships among scalar variables, although some applications require mechanisms generating vector-valued groups.The motivation includes scientific applications such as brain imaging and climate modeling.
  • Related work: Group DAGs or C-DAGs represent directed acyclic relationships among clusters, groups, or partitions of variables.Existing work often assumes these groupings are known through domain knowledge.
  • Contribution: The paper jointly infers latent clusters and causal structures from observational data using variational inference.The approach learns an approximate posterior over both latent quantities and is formulated for linear and nonlinear Gaussian additive noise models based on vSCMs.
  • Contribution: Experiments on synthetic and real datasets demonstrate the effectiveness of jointly discovering clusters and causal structure over clusters.

2 Background

The background defines C-DAG discovery as recovering variable clusters and their directed connections, then describes vector-valued causal mechanisms and Gaussian DAG likelihoods used to model the data.

  • 2.1 Problem Statement: A C-DAG is a directed acyclic graph over disjoint variable clusters, with an adjacency matrix describing connections between clusters.The discovery task infers both the partition C and cluster-edge matrix EC from observational samples D.
  • 2.2 Vector-Valued Structural Causal Models: A vSCM assigns each cluster a vector-valued function of parent variables plus a potentially dependent noise vector.This formulation induces a joint distribution that factorizes according to the cluster partition and C-DAG edges.
  • 2.2.1 Linear Gaussian Vector Valued SCMs: The linear-Gaussian vSCM uses a matrix to map cluster parents to cluster variables and multivariate Gaussian noise with covariance Σ_Ci.
  • 2.2.2 Non-linear Vector Valued SCMs: The nonlinear additive-noise model parameterizes each cluster mechanism with a neural network whose inputs and outputs are masked by the cluster and its parents.A single fully connected network with parameters Θ models all cluster mechanisms.
  • 2.3 Marginal Likelihood of DAG: Gaussian DAG marginal likelihoods can be computed in closed form after integrating out parameters under Normal-Wishart priors.The likelihood is defined using parent sets and datasets restricted to selected variable subsets.

3 Methodology

The method jointly infers latent variable clusters and their causal structure from observational data using Bayesian modeling and variational inference. It approximates otherwise intractable posteriors with flexible distributions over cluster assignments and graph edges, supporting uncertainty quantification and active interventions.

  • Joint Bayesian discovery: The approach jointly discovers clusters and their corresponding causal structures by learning an approximate posterior over latent clusters and structures.The Bayesian formulation is intended to quantify uncertainty over multiple C-DAGs consistent with finite-sample data and support active interventions.
  • Generative model: The generative model treats cluster assignments C and graph edges EC as discrete latent variables, with EC parameterized as strictly upper triangular to enforce acyclicity.The model combines priors, likelihood, and parameters Θ in a joint distribution over the observed dataset and latent variables.
  • Linear and non-linear models: For linear vSCMs, marginal likelihoods can be computed by integrating out Θ, whereas non-linear vSCMs require point estimation of Θ and an ELBO dependent on those parameters.The non-linear case uses a joint distribution and variational approximation because parameter marginalization is not available in closed form.
  • Bayesian inference: Bayesian inference combines the likelihood with priors over clusters and graph structures to form the posterior distribution over C-DAGs.The framework can encode structural domain knowledge through priors, while the selected graph prior is designed to counter the greater number of dense DAGs.
  • Variational inference: Exact posterior computation is infeasible because evaluating the model evidence requires summing over the rapidly growing space of cluster assignments and C-DAGs.The method therefore approximates P(EC, C|D) with a tractable variational distribution rather than computing the posterior in closed form.
  • Variational inference: Variational inference learns a tractable approximation by minimizing posterior KL divergence indirectly through maximization of the Evidence Lower Bound.The variational family is factorized into distributions over cluster assignments and graph structure, with parameters learned through ELBO optimization.
  • Distribution over clusters: The method models cluster assignments with factorized or linear autoregressive categorical distributions, including autoregression to capture dependence among grouping variables.The autoregressive formulation is motivated by correlations among dimensions of C that independent categorical variables may not represent.
  • Distribution over graphs: For graph edges, the method considers independent binary, linear autoregressive, conditional, and autoregressive conditional variational families.These alternatives capture structural dependencies among edges and, in the conditional model, dependence of graph structure on cluster assignments.

4 Experimental Results

The experiments evaluate posterior and structure recovery across synthetic and real-world datasets using metrics for graph accuracy, clustering, and posterior approximation. The proposed approach generally outperforms the baseline, with performance depending on model formulation and dataset nonlinearities.

  • Experimental design: The evaluation compares the proposed Bayesian method with Niu et al. (2022) on synthetic and real-world datasets using multiple metrics.The baseline is bootstrapped to obtain multiple estimates for comparison with the Bayesian approach.
  • Evaluation metrics: Expected Expanded-Graph SHD evaluates full graph-recovery error by incorporating approximation errors from both cluster assignments and cluster-edge structures.The expanded graph is defined as G_expand = C E_C C^T.
  • Synthetic results: Across all tested metrics on 4-variable and 7-variable linear synthetic datasets, the proposed approach consistently outperforms the baseline.The autoregressive cluster model generally performs best among the linear variants.
  • Synthetic results: 0.00 ExG-SHD is achieved by several nonlinear variants on 4-variable fork and v-structure configurations.Autoregressive and conditional edge models outperform factorized edge models, while autoregressive cluster models perform less well in these settings.
  • Real-world results: On real-world datasets, the proposed approach performs better than the baseline on most metrics, and autoregressive cluster models approximate the true posterior better than factorized cluster models.Exact posterior enumeration was possible for both real-world datasets.
  • Real-world results: On the climate dataset, the Neural Conditional Edge model reaches ExG-SHD 0.04, E-SHD 0.0, AUC 1.0, and E-RI 1.0.The baseline scores ExG-SHD 15.75, E-SHD 1.1, and E-RI 0.55 on this dataset.

5 Conclusion

The paper develops variational inference for jointly learning cluster assignments and C-DAG structures, using probabilistic models for both components. It extends the framework to nonlinear vector-valued structural causal models and evaluates it on synthetic and real datasets.

  • 5 Conclusion: The method learns posterior distributions over clusters and structures using categorical cluster assignments and Bernoulli edge variables.It also considers factorized and autoregressive models for cluster assignments.
  • 5 Conclusion: The framework models cluster-graph structure with independent binary, linear autoregressive, conditional, and autoregressive conditional approaches.
  • 5 Conclusion: The framework extends to nonlinear vector-valued structural causal models by parameterizing causal mechanisms with neural networks.Inference techniques are developed for both linear and nonlinear models.
  • 5 Conclusion: Experiments on synthetic and real datasets show the effectiveness of the proposed approaches compared with baselines.The authors identify scaling to high-dimensional datasets as future work.

A Proof of Proposition 1

The proof introduces a KL divergence between variational distributions over graph edges and cluster assignments and their posterior. Expanding this divergence yields a variational lower bound on the data's marginal log likelihood.

  • A Proof of Proposition 1: The variational objective is expressed as KL(QZ(EC) QT(C) || P(EC, C|D)).
  • A Proof of Proposition 1: Expanding the objective produces log QZ(EC) + log QT(C) − log P(D, EC, C) + log P(D).
  • A Proof of Proposition 1: The resulting expression provides a lower bound on the marginal log likelihood of the data.

B Proof of Proposition 2

The supplied passages introduce the marginal likelihood of the data and describe the synthetic and real datasets used in the work. They do not provide the proof's intermediate derivation beyond the expectation step.

  • B Proof of Proposition 2: The proof considers the marginal log likelihood of the data.
  • B Proof of Proposition 2: The displayed expression includes the likelihood P(D|EC, C) and prior P(EC, C).
  • B Proof of Proposition 2: The derivation invokes the definition of expectation.
  • B Proof of Proposition 2: Synthetic datasets use linear vector structural causal models with d = 4, 7 variables and k = 3 clusters, focusing on chains, forks, and v-structures.
  • B Proof of Proposition 2: The real-data evaluation uses protein and climate datasets, with cluster ground truth explicitly identified for the protein data from an expert-generated DAG.

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

The seven-variable, three-cluster synthetic dataset is generated from C-DAG configurations, with nonlinear transformations applied to produce a nonlinear dataset. The associated table documents the nonlinearities used.

  • C.2 Data Generating C-DAG for 7 variable 3 cluster dataset: The data-generating C-DAG for this setting uses d = 7 variables and k = 3 clusters.
  • C.2 Data Generating C-DAG for 7 variable 3 cluster dataset: The figure configurations include fork, chain, and v-structure cases for synthetic C-DAG data generation.
  • C.2 Data Generating C-DAG for 7 variable 3 cluster dataset: The synthetic nonlinear dataset is generated from the C-DAGs shown in Figures 4 and 5.
  • C.2 Data Generating C-DAG for 7 variable 3 cluster dataset: An element-wise nonlinear operation follows the linear combination of parents to ensure nonlinearity.
  • C.2 Data Generating C-DAG for 7 variable 3 cluster dataset: Table 4 describes the element-wise nonlinearities applied after the linear combination of parents.

C.4 Details for Protein Dataset

The protein signalling evaluation uses 11 features and 853 observations, with an expert-provided ground-truth graph but no known true C-DAG. The ground-truth structure indicates two disconnected clusters.

  • The Sachs protein signalling dataset contains d = 11 features and N = 853 observations.Its expert-provided ground-truth graph is shown in Figure 6a.
  • No true C-DAG is known for the protein signalling dataset.
  • The ground-truth structure indicates two disconnected clusters: {Plcg, PIP3, PIP2} and all remaining nodes.

C.5 Details for Climate dataset

The climate evaluation uses surface-temperature samples from the Tropical Pacific and British Columbia, where the causal influence from ENSO to BCT is established. The experiment uses eight variables grouped into four points per region and evaluates against a ground-truth C-DAG.

  • The climate dataset contains surface-temperature samples from the Tropical Pacific (ENSO) and British Columbia (BCT).
  • The established causal influence of ENSO temperature variations on BCT is recovered by the referenced method with around 59% accuracy.
  • The experiment uses four randomly chosen grid points per region, giving d = 8 variables and k = 4 clusters.Figure 7 provides the ground-truth C-DAG used to calculate the metrics.

D Experimental Details

The experiments tune optimization and model settings across autoregressive and neural formulations, using repeated training runs and dataset-specific regularization. Training protocols include multiple restarts, long optimization schedules, and sampling-based ELBO and metric estimation.

  • Both factorized and linear autoregressive models undergo searches over learning rates, optimizers, schedulers, and λs values.RMSprop with learning rate 0.01 and exponential annealing gamma 0.9 was optimal across datasets.
  • Figure 8 illustrates autoregressive generation of cluster assignments using the linear model.
  • Experiments use 10 random restarts, 10,000 epochs, Xavier uniform initialization, and Tesla V100-SXM2-32GB GPUs.The linear autoregressive model searches hidden-state dimensions from 4 through 128.
  • The neural formulation uses a two-hidden-layer softplus MLP with 64 units per hidden layer and AdamW weight-decay regularization.The selected λs values are 1000 for protein signalling, 30 for climate, and 500 for synthetic nonlinear datasets.
  • Real-world datasets train for 75,000 epochs and synthetic datasets for 50,000 epochs, with cosine learning-rate decay during the first half.ELBO uses 30 samples, while performance metrics use 1000 predictive-distribution samples.
  • Tables 5 and 6 report hidden-state dimension choices and variation across random restarts for the autoregressive linear cluster model.
Loading 2608.22212v1…