Non-parametric recovery of causal diffusion mechanisms from steady-state observations

Abstract: We consider sparse multivariate stochastic systems that evolve in continuous time according to a causal mechanism and present methodology to recover the system's time-infinitesimal transition mechanism from mere cross-sectional data. This observational paradigm is motivated by applications such as gene expression analysis, where destructive experimental techniques may only allow recording data once over a cell's lifetime. Precisely, we assume the system follows a time-homogeneous diffusion process that has reached an equilibrium distribution at observation time. Further, we assume the causal mechanism is fully described by the diffusion drift, is acyclic, and its causal structure graph is known. In this setting, we prove that the full causal mechanism, i.e., the drift function, can be non-parametrically identified under a weak non-explosion criterion. We derive a non-parametric kernel estimator for this challenging inverse problem and prove its consistency. Moreover, we propose a cross-validation scheme for hyperparameter tuning, illustrate the behavior of our estimator in simulations, and we discuss connections with irreversible generative diffusion models and low-frequency sampled data.

• The paper introduces a novel kernel estimator that non-parametrically recovers causal drift functions from stationary distributions.
• It establishes theoretical identifiability conditions with finite-sample error bounds for DAG-structured, high-dimensional stochastic systems.
• Synthetic experiments validate the method’s efficiency, offering insights on kernel selection and regularization in equilibrium inference.

Non-parametric Recovery of Causal Diffusion Mechanisms from Steady-State Observations

Problem Setting and Motivation

This paper addresses a fundamentally challenging inverse problem: non-parametrically reconstructing the causal, time-infinitesimal transition mechanism (the drift function of an SDE) governing high-dimensional stochastic systems, from equilibrium (steady-state) cross-sectional data alone. The motivation stems from practical domains, such as single-cell genomics, where destructive measurement protocols preclude repeated sampling from the same evolving entity. The focus is on sparse, causally structured multivariate diffusions with known directed acyclic graph (DAG) structure and fully known (possibly state-dependent) diffusivity.

Theoretical Identifiability from Equilibrium Data

Central to the work is establishing under what conditions the drift function $\mathbf{b}$ of a diffusion process can be uniquely and non-parametrically inferred from its stationary density $p$ and DAG structure. The system dynamics are defined by the SDE:

$$d{\mathbf{x}(t)} = \mathbf{b}\big(\mathbf{x}(t)\big)\,d{t} + \Sigma\big(\mathbf{x}(t)\big)\,d{\mathbf{w}(t)},$$

where $\mathbf{b}$ is the unknown drift, $\Sigma$ is known, and the process has achieved equilibrium.

The main theoretical result proves that, provided the drift is $L^1$-integrable against the stationary density and its causal structure matches the known DAG (with self-loops allowed), $\mathbf{b}$ can be uniquely identified almost everywhere, given $p$ and $\Sigma$. This holds under a very mild non-explosion condition that subsumes both constant and state-dependent diffusivity—greatly generalizing prior linear or parameterized identifiability results.

Kernel-based Non-parametric Estimation of Drift

To operationalize identification, the authors introduce a non-parametric kernel estimator for each DAG-structured component $b_i$ of the drift. Each component is modeled in an RKHS matching the variable's parental separation in the DAG. The approach:

• Uses expectations derived from the Fokker-Planck equation to express the population least-squares estimator for $p$0 in an integral-equation form, naturally regularized via Tikhonov or alternative spectral regularization.
• Avoids explicit solution of the stationary Fokker-Planck PDE, thus sidestepping the curse-of-dimensionality in high-dimensional settings.
• Leverages closed-form kernel solutions where possible, supporting both Gaussian and more structure-adaptive kernels (e.g., sigmoid kernels) along with analytic kernel integrals to efficiently evaluate required quantities.

The plug-in estimation algorithm advances sequentially along a topological ordering of the DAG, recursively applying the estimator to parent components. Hyperparameters (e.g., regularization strength) are adaptively selected via a cross-validation risk estimator tailored to the L² drift estimation framework.

Statistical Guarantees and Numerical Results

Rigorous finite-sample analysis demonstrates that the estimator is consistent and provides explicit high-probability and expected $p$1 error bounds for each drift component. The generalization error exhibits polynomial decay with sample size, with the rate explicitly depending on the topological depth of each node in the DAG and the richness of the kernel RKHS. The recursive plug-in structure, while computationally efficient, induces error propagation along causal depth; the theoretical rates are doubly exponential in DAG height, although the analysis suggests these bounds may be pessimistic compared to practical observations.

Extensive synthetic experiments validate the estimator, examining both:

• Recovery of structured nonlinear drifts in low and moderately high dimensions,
• Empirical generalization error behavior as a function of sample size, kernel choices, and regularization schemes.

Across scenarios, estimate accuracy is highly competitive, especially for kernels well-matched to underlying drift structure and for shallow DAGs.

Figure 1: Learning the drift of a two-dimensional SDE from equilibrium observations; the estimated and true drift vector fields align closely within the high-density region of the stationary distribution.

Figure 2: Relative generalization error for selected drift components versus sample size, comparing Gaussian and sigmoid kernels as well as different regularization strategies.

Extensions, Connections, and Theoretical Implications

The framework is broadly extensible, particularly to:

• Recovery of drifts in systems with low-frequency temporal sampling (ergodic trajectories),
• Connections with non-reversible generative diffusion models, whereby DAG-structured irreversible drifts yield provably faster mixing than classical Langevin (score-function-based) samplers,
• Cases where the structure graph is unknown, or the diffusion coefficient itself is to be estimated — although these remain open challenges.

The theoretical contribution fundamentally clarifies what aspects of dynamical causality are recoverable from steady-state distributions, especially under minimal modeling assumptions, and provides a computationally tractable, theoretically justified methodology.

Figure 3: Accelerated mixing demonstrated for a Gaussian target under drift interpolation between the score function and a causal DAG-structured drift, showing faster convergence with increased irreversibility.

Conclusion

This work rigorously establishes non-parametric identifiability and recovery of sparse, DAG-structured causal drift mechanisms of stochastic dynamical systems from steady-state samples, leveraging a novel kernel-based methodology with demonstrable statistical and computational practicality. The results suggest concrete avenues for future research, including extensions to unknown structure and diffusivity, rate refinements, and integration with modern generative modeling and causal inference paradigms. The presented approach highlights the important intersection of equilibrium statistical inference, stochastic process theory, and nonparametric machine learning.