ZIO: Order-Based Structure Learning for Zero-Inflated Count Data under Data Heterogeneity

Abstract

We propose a structure learning method for directed acyclic graph (DAG) models of zero-inflated multivariate count data. Although causal structure learning with zero-inflated Poisson models is identifiable, existing algorithms rely on local search over the DAG space, leading to poor scalability and a tendency to converge to suboptimal solutions. To address these limitations, we develop a novel order-based scoring framework that searches over variable orderings rather than directly over DAGs. Building on this framework, we further extend the method to multiple heterogeneous datasets by identifying a shared variable ordering while allowing each dataset to have its own edge set. We implement a stochastic hill climbing algorithm with a random-to-random (R2R) proposal operator to efficiently explore the ordering space. Simulation studies show competitive performance against existing methods and improved graph estimation accuracy as more heterogeneous datasets are jointly analyzed under a shared causal ordering. We apply the method to single-nucleus RNA sequencing data from a major depressive disorder study.

arXiv commentAccepted at NeurIPS 2026