Structure Learning on Clustered Data

2026-07-09

AI Summary

This study addresses a critical limitation in existing directed acyclic graph (DAG) structure learning methods, which assume data homogeneity and thus struggle with clustered data exhibiting cluster-specific variations, such as individual effects. To overcome this, the work introduces mixed-effects modeling into DAG learning for the first time, proposing a novel approach that jointly estimates a global causal structure alongside local cluster-level effects. Key innovations include a differentiable graph coupling mechanism that enforces acyclicity across the joint graph, an efficient first-order optimization scheme combined with cross-cluster batching, and theoretical guarantees of model identifiability and asymptotic consistency. Empirical evaluations on both synthetic and real-world clustered datasets demonstrate that the proposed method significantly outperforms current state-of-the-art approaches and effectively uncovers previously overlooked dependencies.

Abstract

Recent algorithmic advances have made directed acyclic graph (DAG) structure learning scalable for causal discovery. Yet, the currently available techniques assume a completely homogeneous population, precluding their application to clustered data where cluster-specific variations (e.g., patient-specific effects) are common. We address this issue by introducing a new approach that estimates a global structure while accounting for local cluster-level effects. The key idea is to extend the fixed- and random-effects framework of classical mixed models to the structure learning setting. Towards this end, we present a differentiable graph coupling mechanism that guarantees the union of the fixed- and random-effects graphs remains acyclic. Computationally, we provide a provably convergent first-order method and leverage efficient batched updates across clusters. Statistically, we establish identifiability of the model and show that our approach recovers the true structure asymptotically. In experiments on real and synthetic data, our proposal detects dependencies missed by alternative estimators, underscoring its value for structure learning in clustered settings.

Problem

• structure learning
• clustered data
• causal discovery
• heterogeneous populations
• DAG

Innovation

• structure learning
• clustered data
• mixed-effects models
• differentiable graph coupling
• causal discovery