Structure Learning on Clustered Data

Traditional causal discovery methods typically assume a homogeneous population, where a single directed acyclic graph (DAG) describes the causal relationships for all observations. This assumption fails in clustered data—such as multi-patient proteomic studies—where causal effects may share a common global structure but vary in strength or polarity across different clusters. The authors aim to bridge this gap by extending the classical mixed-effects modeling framework to the domain of DAG structure learning.

The authors propose a structural equation model that decomposes the weighted adjacency matrix $W^{(i)}$ for each cluster $i$ into a fixed-effects matrix $B$ (common to all clusters) and a random-effects matrix $U^{(i)}$ (cluster-specific deviations). To ensure the resulting causal structure is valid, they enforce a constraint that the union of the fixed-effects graph and the random-effects graph remains acyclic.

To make this computationally tractable, they utilize a differentiable log-determinant characterization of acyclicity. This transforms the combinatorial problem into a continuous optimization task, allowing the use of a path-following proximal gradient algorithm. This method leverages batched updates across clusters, making it scalable to hundreds of nodes.

The proposed framework successfully recovers both population-level causal structures and cluster-specific variations that standard iid-based estimators miss. The authors provide theoretical guarantees for model identifiability and asymptotic recovery of the true structure. Experiments on both synthetic data and real-world proteomic datasets demonstrate that the method effectively detects dependencies that are otherwise obscured by the assumption of population homogeneity.

This work provides a principled way to perform causal discovery in hierarchical or clustered data without sacrificing the scalability of modern differentiable structure learning. By allowing for cluster-specific deviations, researchers can gain more nuanced insights into causal mechanisms in fields like personalized medicine, economics, and psychology, where individual or group-level heterogeneity is the norm rather than the exception.