PACO: Partial-order-Augmented Continuous Optimization for Differentiable Causal Discovery

Differentiable structure learning transforms DAG discovery from a combinatorial search into a continuous optimization problem, enabling gradient-based recovery of causal structures. In many applications, researchers possess partial-order priors over the variables (e.g., from biological signaling cascades or temporal ordering), which substantially narrow the space of plausible causal graphs. Existing approaches integrate such priors by decomposing the prior order graph into a collection of paths and evaluating prior-aware acyclicity constraints path by path; this is correct but its constraint-evaluation cost scales linearly with the number of induced maximal paths, a bottleneck for multi-chain priors. We propose PACO, Partial-order-Augmented Continuous Optimization, which encodes the same partial-order compatibility through a single masked-reachability matrix, that masks the matrix exponential of the learned graph by the strict transitive closure of the prior. The constraint and its Frechet-adjoint gradient are computed in a single $2d{\times}2d$ block matrix exponential whose dominant cost is independent of the number of induced prior paths. We prove that PACO has the same exact zero-level feasible set as the path-decomposition formulation. Across synthetic and real-data experiments under matched priors and dataset seeds, PACO matches path-decomposition on structural recovery and substantially reduces wall-clock cost in multi-chain prior settings.