arXiv CSopen access

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

Non-parametric recovery of causal diffusion mechanisms from steady-state observations Richard Schwank1 and Mathias Drton1,2 1 School of Computation, Information and Technology, Technical University of Munich, e-mail:

arXiv:2606.30467v1 [stat.ML] 29 Jun 2026

[email protected]; [email protected] 2 Munich Center for Machine Learning

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 nonparametrically 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. Keywords and phrases: Stationary diffusion process, Causal dynamical system, FokkerPlanck inverse problem, Reproducing kernel Hilbert space, Score matching, Logistic population growth in a random environment, Graphical continuous Lyapunov model.

1. Introduction Let x1 , . . . , xn ∈ Rd be multivariate data collected by observing each one of n independent individuals at a single time point. Suppose that probing individuals a second time is difficult/impossible, as is the case, for instance, when single-cell gene sequencing requires sacrificing the considered cells. We are interested in inference on the mechanism governing the time evolution of the random vector x from which the n observations were generated. For example, we may wish to use expression data to infer an underlying gene regulation or protein signaling mechanism (Lorch et al., 2026; Zhao et al., 2026; Varando and Hansen, 2020). Inspired by related literature (Peters et al., 2022) and models for gene regulation (Pratapa et al., 2020), we assume the unobserved evolution of x follows a time-homogeneous diffusion process x(t) given by the stochastic differential equation

dx(t) = b x(t) dt + Σ x(t) dw(t) with b : Rd → Rd an unknown drift function and w(t) a d-dimensional standard Brownian motion. The scaling function Σ is assumed to be known; often the case with Σ a constant multiple of the identity has been considered. We assume that our observations are representative, in the sense of being collected after the system has reached its steady state (and is in equilibrium). In other words, we assume the data x1 , . . . , xn is generated as an i.i.d. sample from a distribution with a density p that is a stationary probability density of the diffusion process (Lorch et al., 2024; Varando and Hansen, 2020). 1

R. Schwank and M. Drton/Causal Diffusions

2

Even under the equilibrium assumption, there are multiple drifts b inducing the same stationary density p, see Section 5.4 and Nickl and Ray (2020). One setting that allows unique recovery of the mechanism b are suitably sparse causal systems, i.e., the i-th drift component bi depends only on coordinates xj with j belonging to a (typically small) index set. We depict this sparsity graphically via the parent set of i in a directed graph whose nodes are the vector indices. In this paper, we consider a non-parametric model for drifts b that are causally structured according to a known directed acyclic graph (DAG). This setup arises for example when modeling the effect of a covariate process on an outcome process, see Section 5.1 and De Gregorio et al. (2025). We first show that the drift b is non-parametrically identified by the stationary density p under mild assumptions, i.e., that we can recover b uniquely given infinite data. On a technical level, we link b and p via the Fokker-Planck equation and show uniqueness if all drift components bi are p-integrable, a weak condition to ensure the diffusion process does not explode (Lee and Trutnau, 2021). This result generalizes existing work making strong parametric assumptions (Dettling et al., 2023). Our second contribution is a non-parametric kernel estimator for the drift b given a userspecified DAG. The estimator builds on a result from the identifiability section, which we interpret as an identification formula for the image of b under an integral operator. This formula is reminiscent of an identity from nonparametric score matching, and following Zhou et al. (2020), we approximate both sides by plug-in estimators. All that then remains is to invert the integral operator, for which we study different regularization schemes. Building on De Vito et al. (2005), we theoretically analyze the Tikhonov-regularized estimator and derive a high-probability concentration inequality. We deduce that any drift with components bi ∈ L2 (p) can be consistently estimated when an L2 -universal kernel is used. Additionally, if the Hilbert space corresponding to the chosen kernel approximates the drift sufficiently well, we derive a quantitative bound for the components’ L2 risks depending on their depth in the structure graph. For practical use, we also provide a cross-validation procedure to tune hyper-parameters. We illustrate our method on simulated data from population growth in a random environment and from a sigmoidal drift in seven variables. Further, we note that our plug-in approach aligns well with low frequency samples from a single ergodic process and performs well in regimes with low to moderate autocorrelation. Finally, we switch perspective from recovering to creating b, asking when a given density p is the stationary density of a diffusion with DAG-structured drift. One motivation for this question is that DAG-structured drifts induce irreversible processes that converge faster to their stationary distribution than the score-based drifts used in generative diffusion models. We prove that any sufficiently smooth density p can be generated by the complete DAG and illustrate accelerated convergence in a simple setting. Notation We boldface vectors and matrices in lower and upper case, respectively. For a ∈ Rd , A ∈ Rd×d and index sets S, S ′ ⊂ {1, . . . , d}, we write aS := (ai )i∈S , AS,S ′ := (Aij )i∈S,j∈S ′ , d / . For Lebesgue densities p : R → R+ we define the marginal pS (xS ) := Rand a−S := (ai )i∈S p(x) dx−S . For a1 , a2 reals, a1 ∨ a2 := max{a1 , a2 } and a1 ∧ a2 := min{a1 , a2 }. The Rd−|S| vector q-norm, q ∈ [1, ∞], is written ∥a∥q , with ∥a∥ := ∥a∥2 . We use the Lebesgue spaces Lq (Rk , λ) and abbreviate to Lq (Rk ) for Lebesgue’s measure or to q L (λ) if the ambient space is clear from λ. We also use the local versions Lqloc . Let Lq (Rk1 , Rk2 , λ) denote functions a : Rk1 → Rk2 with component functions ai ∈ Lq (Rk1 , λ) for all i ≤ k2 . We use Pk the Sobolev spaces W n,q (Rk , λ) and write ∇ for the (weak) gradient, ∇ · a := i=1 ∂i ai for the divergence, ∇2 for the Hessian matrix, and ∆ for the Laplace operator. All derivative operators may only be relative to a subset of variables, e.g., ∇2S a := (∂ij a)i,j∈S . When a is a function of one variable, we write a′ := ∇a. We write C n (Rk ) (and C0n (Rk )) for n-times continuously differentiable (and compactly supported) functions where n ∈ N ∪ {∞}. Let C n,0 (Rd × Rd )

R. Schwank and M. Drton/Causal Diffusions

3

denote the space of continuous functions a(x, y) such that ∂xα a exists and is continuous for all multi-indices |α| ≤ n. We denote function composition by a ◦ b. Finally, for Hilbert spaces A, B, L(A, B) denotes bounded linear operators F : A → B, whose adjoints we denote by F ∗ . 2. Related literature The idea to explain observations by repeated application of a causal mechanism over a long unobserved time window has parallels in the causality literature. For instance, it has been considered to give meaning to cyclic structural equation models (SCMs) as surveyed by Peters et al. (2017, Sect. 2.3.3). If convergence can be established, Bongers and Mooij (2018) prove that equilibria of random differential equations also follow SCMs. Although conceptually similar, the quantitative situation in this paper is different since our diffusion process does not converge to a constant almost surely. Therefore, a sample from the stationary distribution reveals information on common time history between coordinates. As a result, conditional independence relations are significantly reducted in the stationary distribution (Boege et al., 2025), which generally is not Markov with respect to the structure graph underlying the drift function. A key theme of this paper is identifiability: characterizing what information on the causal mechanism is uniquely determined by the stationary distribution when no time history is available. When the structure graph is known, it was shown by Dettling et al. (2023) that linear diffusion drifts can be fully recovered when the graph is free of 2-cycles; examples show generic identifiability for additional graphs. When the graph is an unknown DAG, the stationary distribution determines a graph equivalence class, which for linear drifts are finer than the Markov equivalence classes encountered for SCMs (Améndola et al., 2025). Additional interventional data can narrow down the drift even further (Inglese, 2002; Rohbeck et al., 2024). Identifiability of linear drifts is often studied under the term graphical Lyapunov model in the literature, as the Lyapunov equation links drift matrix and covariance matrix of the Gaussian stationary distribution. We note that sparsely interacting diffusion processes also feature in literature from probability theory (Lacker and Zhang, 2023). For drift estimation, maximum likelihood is an option as the Fokker-Planck equation yields the stationary density for a proposal drift. This approach is considered for parametric models by Pedretscher et al. (2019). Similarly, one can optimize distributional metrics between the observed distribution and the solution to the Fokker-Planck equation (Botvinick-Greenhouse et al., 2023). Setting computational demands in high dimensional parameter spaces aside, the main challenge is the curse of dimensionality when solving the stationary Fokker-Planck equation in dimensions greater than three or four (Sun and Kumar, 2014, Fig. 10). The alternative, to sample from the SDE until reaching steady state, also suffers from the curse of dimensionality in general (Amorino and Gloter, 2023). We are only aware of the work by Lorch et al. (2024) that scales to higher dimensions by optimizing a kernelized surrogate loss. Common to the aforementioned work is the need to optimize a complicated nonlinear function in the proposal drift, usually via gradient descent. In contrast, the method we develop here targets the mean squared error between proposal and ground truth drift directly; see Section 4.3. This is very convenient for the theoretical analysis controlling the L2 (p) distance between estimated and ground truth drift. Our work also has connections with score matching. First, we find drift learning to be equivalent to score learning in one dimension when the diffusivity is constant. Second, our estimator naturally extends nonparametric score learning methods (Zhou et al., 2020). Third, the score solves a conceptually related problem in generative diffusion modeling: learning a drift inducing a given stationary distribution without observing any time history (Song and Ermon, 2019).

R. Schwank and M. Drton/Causal Diffusions

4

3. Model identifiability 3.1. Diffusion processes with structured drifts and the Fokker-Planck equation Consider an individual with attribute vector x ∈ Rd , which evolves over time according to the time-homogeneous stochastic differential equation (SDE)

dx(t) = b x(t) dt + Σ x(t) dw(t) (1) understood in Ito-sense, where b : Rd → Rd is the drift function, w(t) is a standard d-dimensional Wiener process, and Σ : Rd → Rd×d is a full rank √ matrix-valued function. Equation (1) can be interpreted as x(t + ∆t) = x(t) + b(x(t))∆t + ∆t · Σ(x(t)) Z(t) for arbitrarily small ∆t > 0, where Z(t) is standard Gaussian. We are interested in drifts adhering to a causal structure; more precisely to a directed acyclic graph (DAG) with self-loops (Dettling et al., 2023). That is, we define a DAG to be a collection of directed edges i → j between nodes i, j ∈ {1, . . . , d} with exactly one path from node i back to itself, namely the self-loop i → i. Let Pa(i) := {j : j → i} be the parents of node i and let pa(i) := {j ̸= i : j → i} denote the proper parents. In particular, Pa(i) = pa(i) ∪ {i}. Let an(i) denote the proper ancestors, i.e., nodes with a path to i but excluding i itself. Definition 3.1. A function f : Rd → R only depends on xS for some index set S ⊂ {1, . . . , d} if f = f ◦ πS , where πS denotes the projection πS (x) = xS . A function b : Rd → Rd is structured according to a DAG D on nodes {1, . . . , d}, if every component function bi only depends on xPa(i) . In particular, if the drift b is DAG-structured, each component function bi can depend on xi due to self-loops. The reasoning is that linear drifts b(x) = Bx given by matrix B ∈ Rd×d can only be guaranteed to admit a stationary distribution if all eigenvalues have a negative real part. However, DAG-structured matrices can only satisfy this condition with a non-zero diagonal, since they are lower triangular after a suitable permutation of the indices. For a causal interpretation, it would be natural to assume that (Σi,k )dk=1 , and thus the entire SDE equation for dxi (t), only depends on the xj with j → i ∈ D. However, as subsequent results are not affected, we keep the structure of the matrix function Σ(x) unrestricted. We consider processes x(t) that approach an equilibrium distribution as t → ∞. Under suitable assumptions on b and D (Khasminskii, 2012, Sect. 4), this equilibrium distribution admits a probability density p with respect to Lebesgue measure that satisfies the stationary Fokker-Planck equation d 1 X ∂i ∂j (p Dij ) = ∇ · (b p), (2) 2 i,j=1 where the diffusivity parameter D : Rd → Rd×d is defined by D(x) := Σ(x)Σ(x)T . Hereafter, we omit “stationary” and refer to (2) as the Fokker-Planck equation. Before discussing how to interpret this partial differential equation (PDE) in general, we consider a simple setting used as illustration throughout the paper:

∂11 p(x) + ∂22 p(x) = ∂1 (b1 p)(x) + ∂2 (b2 p)(x)

∀x ∈ R2 .

⟲

⟳

Example 3.2 (Running example: 2-path). Consider an object with d = 2 attributes x1 , x2 , e.g., concentrations of two mRNAs in a cell, and assume regulation according to 1 → 2 , meaning the drift is of the form b(x1 , x2 ) = (b1 (x1 ), b2 (x1 , x2 )) ∈ R2 . √ Let p ∈ C 2 (R2 ) be a bounded probability density. Consider Σ to be 2 times the constant identity, meaning the left hand side of (2) becomes the Laplacian ∆p. Assume b ∈ C 1 (R2 , R2 ) and that the Fokker-Planck relation ∆p = ∇ · (b p) holds, i.e., (3)

R. Schwank and M. Drton/Causal Diffusions

5

√ If b has bounded first derivatives, the SDE dx(t) = b(x(t))dt+ 2dw(t) has a solution (x(t))t≥0 for arbitrary initial x(0) ∈ R2 . If this solution is positively recurrent, the density of x(t) approaches p at every point as t → ∞ (Khasminskii, 2012, Lemma 4.16/4.17). To minimize assumptions on the drift, we consider the following weak formulation. Definition 3.3. We say b ∈ L1loc (Rd , Rd , p) solves the Fokker-Planck equation with (Lebesgue) density p and diffusivity D ∈ L1loc (Rd , Rd×d , p), when Z Z 1 ⟨D, ∇2 φ⟩F dp ∀φ ∈ C0∞ (Rd ), (4) bT ∇φ dp = 2 d d R R P where ⟨A, B⟩F = tr(AT B) = ij Aij Bij is the Frobenius product of matrices A, B. If b, p, and D are twice differentiable in the classical sense, integration by parts and the fundamental lemma of the calculus of variations recovers the classical equation (2). 3.2. The one dimensional case We first consider recovering a univariate drift function b : R → R from a univariate stationary density p. This part provides intuition and serves as a base case for general DAGs. Let p ∈ C 2 (R) be a positive density that is stationary for the univariate SDE dx(t) = b(x(t))dt + σdw(t)

(5)

with drift b ∈ C 1 (R) and a constant σ > 0. Then, the Fokker-Planck equation (2) reads 1 2 ′′ σ p = (bp)′ 2

⇐⇒

b=

σ 2 p′ c + 2 p p

for c ∈ R,

′

using the fundamental theorem of calculus. Note that pp = (log p)′ is the score function of p. Hence, the Fokker-Planck equation determines the drift up to a one-parameter family (bc )c∈R . It is intuitive that bc=0 is the only sound drift in the family: c/p grows strongly for c ̸= 0, especially when p is light-tailed. In fact, the growth is so strong that x(t) is forced to leave the state space R and attain one of the symbolic values ±∞ at a finite time, called an explosion. For univariate SDEs, Feller’s explosion test (Lemma A.1) yields: 2

′

Lemma 3.4. The solution of the SDE (5) with drift bc (x) = σ2 pp + pc explodes with positive probability if c ̸= 0. An explosion in the process x(t) strongly contradicts our goal to model objects in a stable equilibrium. Especially, since one can show that x(t) following (5) with bc for c ̸= 0 has a positive probability of having exploded before any time T > 0 (Karatzas and Ruf, 2016). Consequently, 2 for our purposes, b = σ2 log(p)′ is the only reasonable drift that yields p as a stationary density. When considering an x-dependent diffusivity σ or multivariate diffusions, exploding drifts are much harder to characterize sharply. In this paper, we consider drifts from the class of functions absolutely integrable with respect to the density p, i.e., b ∈ L1 (R, p). This is a mild restriction, yet strong enough to rule out explosions under suitable assumptions on the diffusivity D (Lee 1 and Trutnau, 2021; Hwang et al., 2005). The R assumption b ∈ L (R,R p) consistently also singles out the score from the class (bc )c∈R : when R |p′ (x)|dx < ∞, then R |bc |dp < ∞ if and only if c = 0. The following Lemma demonstrates that the assumption b ∈ L1 (R, p) is also sufficient to recover the drift for x-dependent diffusivity σ. Lemma 3.5. Let b∗ , b ∈ L1 (R, p) both solve a Fokker-Planck equation (4) with the same density p > 0 and diffusivity D ∈ L1loc (R, p). Then, b = b∗ almost everywhere.

R. Schwank and M. Drton/Causal Diffusions

6

3.3. General DAGs Let the drift b : Rd → Rd be structured according to a DAG D. Building on the framework developed for the univariate case in the previous section, we show that requiring b ∈ L1 (Rd , Rd , p) is enough to recover the drift from the stationary density p. This integrability assumption was motivated by its ability to effectively rule out explosions in the associated diffusion process. The main result at the end of this section is derived via induction over a topological ordering of D. By shuffling the nodes, we may assume the ordering to be 1, . . . , d without loss of generality. This means any edge i1 → i2 in the DAG satisfies i1 ≤ i2 . It is important for the k-th induction step to have an equation involving only the components bi with i ≤ k + 1. We state such an equation below and afterwards discuss how it simplifies in the running Example 3.2. Lemma 3.6. Let the D-structured drift b ∈ L1 (Rd , Rd , p) solve the Fokker-Planck equation (4) with density p and diffusivity D ∈ L1 (Rd , Rd×d , p). For any index i ≤ d and any φ = φ̃ ◦ πPa(i) with φ̃ ∈ C0∞ (R|Pa(i)| ) it holds that Z Z Z 1 ⟨DPa(i),Pa(i) , ∇2Pa(i) φ⟩F dp. (6) bi ∂i φ dpPa(i) = − bTpa(i) ∇pa(i) φ dp + R|Pa(i)| Rd Rd 2 In the i-th induction step, bpa(i) is identified by assumption. To see that the right hand side of (6) suffices to determine bi , consider i = 2 in the running Example 3.2, for which equation (6) simplifies to the Fokker-Planck equation ∂2 (b2 p) = ∆p − ∂1 (b1 p). Assume this equation had two solutions b2 , b∗2 ∈ C 1 (R2 ) ∩ L1 (R2 , p) for the same right hand side. Taking the difference of both equations implies ∂2 (δ p) = 0, where we defined δ := b2 − b∗2 . By the fundamental theorem of calculus, (δp)(x) = f (x1 ) for some function f . The assumptions on b2 , b∗2 imply δ ∈ L1 (R2 , p), however Z Z Z |δ| dp = |f (x1 )|dx1 dx2 = ∥f ∥L1 (R) dx2 R2

R2

R

is finite only if ∥f ∥L1 (R) = 0, implying δ = 0 if p > 0 and therefore b2 = b∗2 . The following main result generalizes this proof idea to the equations (6). Theorem 3.7. Assume b∗ , b ∈ L1 (Rd , Rd , p) both solve a Fokker-Planck equation (4) with the same density p > 0 and same diffusivity D ∈ L1 (Rd , Rd×d , p). When b, b∗ are structured according to a common DAG D, then b = b∗ Lebesgue-almost everywhere. To summarize: the stationary distribution, the structure DAG D, and the diffusivity function D together are sufficient to recover the drift function. Next, we construct a practical estimator for this task when a sample of the stationary distribution is available. 4. Estimation 4.1. Reproducing kernel Hilbert spaces (RKHS) Kernel methods, popularized by support vector machines, are applied under different philosophies in the literature. Here, we view kernels as generating nonparametric function spaces, from which we select elements to approximate the drift function. Their advantages are closed-form optima and the ability of the kernel to absorb partial derivatives from the drift in the FokkerPlanck equation. We call k : Rk ×P Rk → R a kernel if k is symmetric, i.e., k(x, y) = k(y, x), and if k is positive m definite, meaning i,j=1 αi αj k(xi , xj ) ≥ 0 for all α ∈ Rm and x1 , . . . , xm ∈ Rk . An example is the Gaussian kernel exp(− 2γ1 2 ∥x − y∥22 ) with bandwidth γ > 0.

R. Schwank and M. Drton/Causal Diffusions

7

The function k(x, ·) : Rk → R, y 7→ k(x, y) is called the feature map corresponding to x ∈ Rk . On the linear span of the feature maps, define an inner product via ⟨k(x, ·), k(y, ·)⟩ := k(x, y). The RKHS given by the kernel k is the completion of the linear span with respect to ⟨·, ·⟩ and denoted by H. The functions f : Rk → R it contains have the reproducing property ⟨f, k(x, ·)⟩H = f (x). The kernel choice determines smoothness and many other properties of functions in H. For k continuous and bounded, the inclusion J : H → L2 (p) is continuous for any probability measure p on Rk , and ∥J∥L(H,L2 (p)) ≤ ∥k∥∞ . The inclusion into L2 (p) in particular allows us to use mean squared error. The adjoint J ∗ : L2 (p) → H, often called integral operator, is given by: Z (J ∗ f )(y) = k(x, y)f (x)dp(x). Rk

We frequently use L := J ∗ J. Under the above assumptions, L is compact and self-adjoint, with nuclear norm bounded by ∥k∥∞ (Steinwart and Christmann, 2008). During our statistical analysis, we approximate J, J ∗ and LP by sample versions. Let x1 , . . . , xn ∈ n k R , and let E n be Rn with the inner product ⟨v, w⟩E n := n1 i=1 vi wi . Following De Vito et al. n (2005), we define Jx : H → E , f 7→ (f (x1 ), . . . , f (xn )). Then, n

Jx∗ v =

1X vi k(xi , ·), n i=1

n

L̂f := Jx∗ Jx f =

1X f (xi )k(xi , ·). n i=1

4.2. Idea and setup Let the probability density p on Rd solve the Fokker-Planck equation with unknown drift b∗ ∈ L2 (Rd , Rd , p) but known diffusivity D ∈ L2 (Rd , Rd×d , p). Assume b∗ is structured according to a known DAG D. As we show below, we can actually regress on the unknown b∗ using only suitable expectations under the data distribution p. This is analogous to score matching, where one implicitly regresses on the true score. We approximate each drift component b∗i by functions from an RKHS Hi ⊂ L2 (Rd , p). Assuming each node in D has a self-loop to avoid case distinctions, we take the corresponding kernel ki (x, y) only depending on (xPa(i) , yPa(i) ). Consider the infinite sample Tikhonov regression bi,λ := argminb∈Hi ∥b − b∗i ∥2L2 (p) + λ∥b∥2Hi

(7)

for some regularization strength λ > 0. We have that bi,λ = −(Li + λI)−1 ζi ,

Z ζi (y) := − Rd

ki (x, y)b∗i (x)dp(x) ∈ Hi ,

(8)

where I denotes the identity on Hi . We now free ζi from the unknown b∗i , meaning we can compute the regression function bi,λ without requiring b∗i . Noticing the similarity between ζi and the left hand side of Lemma 3.6, we consider functions Ki,y : Rd → R indexed by y ∈ Rd with the property that ∂i Ki,y = ki (·, y): Lemma 4.1. Assume p is bounded, let i ≤ d, and let y ∈ Rd . If Ki,y ∈ W 2,∞ (Rd ) only depends on xPa(i) and ∂i Ki,y = ki (·, y), then ζi (y) = Ex∼p [ zi (x, y) ], where 1 zi (x, y) := b∗pa(i) (x)T ∇pa(i) Ki,y (x) − ⟨DPa(i),Pa(i) (x), ∇2Pa(i) Ki,y (x)⟩F . 2 This means we can express ζi as an expectation under p only involving quantities strictly before i in the topological order induced by the DAG D. For pa(i) = ∅ and D the constant

R. Schwank and M. Drton/Causal Diffusions

8

Table 1 Example kernels with corresponding Ki,y . Name

k(x, y)

Linear

∥x−y∥2 exp(− 2γ 2 2 ) xT y + c

Sigmoid

tanh(α xT y + c)

Derivative

∂xi ∂yi k̃(x, y)

Gaussian

γ

√

Ki,y (x)

i 2π Φ( xi −y ) · k(x−i , y−i ) γ 1 2 x y + xi (xT −i y−i + c) 2 i i 1 log cosh(α xT y + c) αyi

∂yi k̃(x, y)

identity matrix, zi reduces to Zhou et al. (2020). Following this work, we replace expectations with sample averages and consider an iterative plug-in estimator, which we formally introduce in Section 4.3. We consider alternative regularization strategies to Tikhonov in Section 4.4. The core idea also works without any regularization, see Section 4.6. 4.3. Estimator derivation Continuing in the setup from Section 4.2, the first step is to choose kernels ki (x, y) depending on (xPa(i) , yPa(i) ) to approximate b∗i for all i ≤ d. Our method requires the anti-derivative of ki with respect to xi , i.e., Ki,y such that ∂i Ki,y = ki (·, y) for almost all y ∈ Rd . We can compute Ki,y analytically for some popular kernels, see Table 1. Here, Φ denotes the standard Gaussian cdf. For derivative kernels, see Steinwart and Christmann (2008, Def. 4.35). Let x1 , . . . , xn ∈ Rd be samples from the density p. We compute the drift estimate b̂ := (b̂1 , . . . , b̂d ) component-wise along the causal order of D. The central object in each step is ζ̂i : Rd → R, which approximates ζi . Inspired by Lemma 4.1, we use already learned drift components (b̂j )j∈pa(i) to approximate zi via 1 ẑi (x, y) := b̂pa(i) (x)T ∇pa(i) Ki,y (x) − ⟨DPa(i),Pa(i) (x), ∇2Pa(i) Ki,y (x)⟩F , (9) 2 Pn and set ζ̂i (y) := n1 j=1 ẑi (xj , y). All that remains is to transform ζ̂i to the drift estimate b̂i . We present a Tikhonov regularized transformation here and point to Section 4.4 for alternative regularizers. The idea is to plug ζ̂i and L̂i from Section 4.1 into equation (8), i.e., b̂i ≈ −(L̂i + λI)−1 ζ̂i . For reasons discussed below, we use the following explicit formula instead, which is based on an extended representer theorem (Zhou et al., 2020, C.4.1) Lemma 4.2. Let f ∈ Hi , y ∈ Rd , and let λ > 0. Define Ki ∈ Rn×n via Ki,jl := ki (xj , xl ), set ki,y ∈ Rn to (ki,y )j := ki (y, xj ), and f ∈ Rn to fj := f (xj ). Then, ((L̂i + λI)−1 f )(y) =

1 1 f (y) − kTi,y (Ki + nλI)−1 f , λ λ

where I ∈ Rn×n is the identity matrix. We can now state the entire Tikhonov regularized estimation procedure. Definition 4.3. For each i ranging through the topological order of D, set ζ̂i (y) := n1 with ẑi from (9), and define ζ̂ i ∈ Rn via (ζ̂ i )j := ζ̂i (xj ). We define b̂i (y) :=

1 T 1 k (Ki + nλi I)−1 ζ̂ i − ζ̂i (y), λi i,y λi

where ki,y , Ki , I are as in Lemma 4.2 and λi > 0 are the regularization strengths.

Pn

j=1 ẑi (xj , y)

R. Schwank and M. Drton/Causal Diffusions

9

The remainder of this chapter considers tuning λ, alternative regularization strategies, and theoretically analyses b̂i . Finally, we did not formally apply (L̂i + λI)−1 to ζ̂i , since it is difficult to guarantee ζ̂i ∈ Hi . As a start, y 7→ Ki,y (x) from Table 1 usually lies not in Hi due to boundary value mismatch, and the partial derivatives of Ki,y do not simplify the matter. For the Gaussian kernel, defining Ki,y via integration from 0 to xi solves the problem and ζ̂i ∈ Hi is guaranteed (Vasilev, 2026). We implement this rule in the package, however a proof for general kernels seems difficult. 4.4. Alternative regularization methods The Tikhonov estimator approximates bi,λ := −(Li + λI)−1 ζi , see Section 4.2. The purpose of adding λI, namely to stabilize the otherwise ill-conditioned inversion of Li , can be achieved with many alternative methods: the Spectral cut-off regularization discards small eigenvalues of Li and the Landweber method performs early-stopped gradient descent on Li f = −ζi . The ν-Method is an accelerated version of the Landweber method. Zhou et al. (2020) provide explicit formulas for each method when Li is replaced by L̂i , similar to Lemma 4.2. We implement all methods in the Python package by plugging our ζ̂i into these formulas. We compare the regularizers empirically in Section 5. Theoretically, an advantage of spectral cut-off regularization is that it provably adapts to regression functions admitting a simple expansion in the eigenbasis of Ji Ji∗ (Dicker et al., 2017), but note our discussion in Section 4.5. 4.5. Theoretical analysis Let the Lebesgue density p on Rd be bounded and solve the Fokker-Planck equation with unknown drift b∗ ∈ L∞ (Rd , Rd ) but known diffusivity D ∈ L∞ (Rd , Rd×d ). Denote 1 B := ∥b∗ ∥L∞ (Rd ,Rd ) ∨ ∥D∥L∞ (Rd ,Rd×d ) . 2 Assume b∗ is structured according to a known DAG D. For i ≤ d, let ki ∈ C 0 (Rd × Rd ) be bounded and only depend on (xPa(i) , yPa(i) ). Let Ki,y (x) ∈ C 2,0 (Rd × Rd ) only depend on (xPa(i) , yPa(i) ) and satisfy ∂xi Ki,y (x) = ki (x, y). Further, assume κ := max max ∥∂xα Ki,· (·)∥L∞ (Rd ×Rd ) < ∞. i≤d |α|≤2

Let x1 , . . . , xn be iid random variables following p. Let b̂ := (b̂1 , . . . , b̂d ) be the drift estimator from Definition 4.3 with Tikhonov regularization strengths λ1 , . . . , λd > 0. Define the training evaluation vectors b∗i ∈ Rn via (b∗i )j := b∗i (xj ), and analogously b̂i , ζ i , ζ̂, and bi,λi ∈ Rn , where bi,λi is the population Tikhonov approximation (8). We first analyze how ζ̂i , ζˆi , and L̂i , the main estimated quantities in the definition of b̂i , concentrate around their ground truth counterparts. Here, we use L̂i since it later allows us a unified treatment of ki,y and K. Let ai := |Pa(i)| + 1 be the number of arguments of b∗i plus one. A Hilbert-space valued Hoeffding inequality yields: Lemma 4.4. For any i ≤ d it holds with probability at least 1 − δ that √ X a2i κB 2 log1/2 2δ √ +κ ∥b̂j − b∗j ∥E n ∥ζ̂i − ζi ∥L2 (p) ≤ n j∈pa(i)

(10)

R. Schwank and M. Drton/Causal Diffusions

10

Note that ∥ζ̂i − ζi ∥L2 (p) still depends on the sample x1 , . . . , xn ; it is the expected squared distance from ζi with respect to a fresh sample from p. We also bound the training error with a U-statistics approach: Lemma 4.5. Let n ≥ 3. For any i ≤ d it holds with probability at least 1 − δ that √ 1/2 1/2 2∨c1 X a2i κB 3 a2i κB · (log n + log δ ) √ ∥ζ̂ i − ζ i ∥E n ≤ + +κ ∥b̂j − b∗j ∥E n √ √ n c2 n j∈pa(i)

for some global positive constants c1 , c2 . Finally, another Hoeffding bound allows us to bound ∥L̂i − Li ∥L(Hi ,Hi ) . Lemma 4.6. For any i ≤ d it holds with probability at least 1 − δ that √ 2κ 2 log1/2 2δ √ ∥L̂i − Li ∥L(Hi ,Hi ) ≤ . n Next, we bound the effect of deterministic deviations on the generalization error of b̂i . Since Lemmas 4.4 and 4.5 involve the training error, we also bound it here. Lemma 4.7. If ∥ζ̂ i − ζ i ∥E n ≤ ε1 , ∥ζ̂i − ζ∥L2 (p) ≤ ε2 and ∥L̂i − Li ∥L(Hi ,Hi ) ≤ ε3 , then ε3 ∥b∗i ∥L2 (p) 2ε1 + , λi λ ! i √ 3/2 ε1 ε3 ε ε1 + ε2 ε3 ∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) + + 33/2 ∥b∗i ∥L2 (p) + + 3/2 . λi λi λ λ ∥b̂i − b∗i ∥E n ≤ ∥bi,λi − b∗i ∥E n +

i

(11) (12)

i

Inserting Lemmas 4.4, 4.5, and 4.6 into Lemma 4.7 yields the main result. We are interested in choosing λi as n grows, hence we hide constants that do not depend on n. Theorem 4.8. Let i ≤ d and λi (n) ≥ √cn for some constant c > 0. Then,
O(log n) √ log 2δ λi n X
O(1) + log 2δ ∥b̂j − b∗j ∥E n , λi

∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) +

(13)

j∈pa(i)

∥b̂i − b∗i ∥E n ≤ ∥bi,λi − b∗i ∥E n + +

O(log n) √ · log( 2δ ) λi n

2κ X ∥b̂j − b∗j ∥E n λi

(14)

j∈pa(i)

with probability at least 1 − δ for any δ ∈ (0, 1). Here, O(·) is as n → ∞. Theorem 4.8 bounds the generalization and training error of b̂i with high probability. We find √ the same rate (λi n)−1 √as De Vito et al. (2005) although we do not require ζ̂i ∈ Hi . This rate would diverge if λi < 1/ n, hence the restriction in the theorem statement. For a more concise result, we consider the expected generalization error, for which we combine (13) and (14) and finally integrate out δ.

R. Schwank and M. Drton/Causal Diffusions

11

Corollary 4.9. Let dij be the length of longest path from i to j in D, excluding self-loops and therefore possibly zero. Let di := maxj dji and set ri := 6 · 3di . Use Tj (λ) := ∥bj,λ − b∗j ∥L2 (p) to abbreviate the Tikhonov population approximation error. If 2/3dji

λi (n) ≥ n−2/ri ∨j∈an(i) Tj

(λj (n))

(15)

for all i ≤ d, then the expected generalization error is bounded as follows: h i O(log n) + O(1) E ∥b̂i − b∗i ∥L2 (p) ≤ n1/ri

X

1/3dji

Tj

(λj ) ∨ Tj (λj ).

(16)

j∈an(i)∪{i}

In particular, if Hi is dense in L2 (pPa(i) ) for all i ≤ d, the condition (15) allows choices λi (n) with limn→∞ λi (n) = 0, for which then h i lim E ∥b̂i − b∗i ∥L2 (p) = 0. n→∞

√ Further, consider Tj (λ) = O( λ · logαj λ1 ) as λ → 0 for all j ≤ d and constants αj ≥ 0. For 1 any ε ∈ (0, 2d ), the choice λi := n−2(1−εi )/ri with εi := di · ε guarantees that i h E ∥b̂i − b∗i ∥L2 (p) = O(n−(1−2εi )/ri ). This Corollary guarantees consistency of our method if Hi is dense in L2 (pPa(i) ), i.e., if ki is L2 universal. An important example is the Gaussian kernel (Steinwart and Christmann, 2008, 4.36), which we implement in the software package. The condition (15) requires that the regularization strength λi (n) may only decrease as fast as the the ancestor drifts b∗j can be approximated in the space Hj with regularization strength λj (n). In practice, we recommend to choose√λi via cross-validation (Section 4.6). The final part of Corollary 4.9 is inspired by Tj (λ) ≤ O( λ) for b∗j ∈ Hj . We allow for additional log factors since drifts of practical interest, e.g., b(x) = −x, may possess RKHS smoothness locally, but b∗i ∈ / Hi due to tail growth mismatch. If p is subGaussian, the tail can sometimes be controlled by such log-factors, e.g., b∗1 (x) = ax + f (x) √ satisfies Tj (λ) = O( λ log2 λ1 ) for any f in the Gaussian RKHS and a ∈ R (Vasilev, 2026). The √ √ growth Tj (λ) ≈ λ influenced our choice of ri and the exponent 3dji , since λ + (nα λ)−1 is balanced by λ(n) = n−2α/3 for α > 0. A different assumption on Tj (λ) would imply a different balancing choice ri , however note our discussion on source conditions below and note that the consistency result is unaffected. In the remaining section, we discuss how to improve the doubly-exponential rate ri and other open problems. First, we are not aware of any lower bounds for learning structured drifts from the stationary distribution. For the related problem of learning the score function ∇ log p = ∇p/p, see Section 5.4, one can achieve minimax optimal L2 (p) generalization error by estimating p̂ and essentially returning ∇p̂/p̂, see Wibisono et al. (2024). The closest parallel to this viewpoint is equation (20) in Theorem 5.1, expressing b∗i in terms of conditionals pi|1,...,i−1 , their derivatives and partial integrals. This suggests learning b∗i may be harder than learning the score, however equation (20) could potentially simplify. A well-known method to guarantee faster rates are source conditions of the form b∗i = (Ji Ji∗ )r fi for some r > 0 and fi ∈ L2 (p) (Dicker et al., 2017; Zhou et al., 2020). The assumption b∗i ∈ Hi corresponds to r = 21 . Regularization methods with sufficient qualification can guarantee ∥bi,λ − b∗i ∥L2 (p) = O(λr ), however a source condition for r > 21 depends on p through the integral operator Ji∗ . Since b∗ is completely determined by p, the source condition becomes a condition on p, which we leave to future work.

R. Schwank and M. Drton/Causal Diffusions

12

P The term driving the rate ri is λ1i j∈pa(i) ∥b̂j −b∗j ∥E n . Could b̂i be de-coupled from (b∗j )j∈pa(i) , similar to nonparametric score learning? Unfortunately not, since the stationary density p and the parent set pa(i) alone do not determine b∗i , see Appendix B.3. However, we believe the factor 1/λi is too conservative in general. It arises from the mismatch that replacing b∗i by some other bi ∈ L2 (pPa(i) ) in zi structurally affects Ex∼p [ zi (x, ·) ], e.g., it may not lie in the image of the integral operator Ji∗ . As the proof of Theorem 4.8 illustrated, the integral operator Ji∗ controls the inverse power of λ. In Appendix B.3, we show how to potentially link Ex∼p [ zi (x, ·) ] with Ji∗ for certain kernels, although tighter theoretical guarantees may require replacing ζ̂i by an analytically better behaved estimator. In particular, a sufficiently strong bound √ on ∥ζ̂i − ζi ∥Hi would allow us to follow the argument by Zhou et al. (2020), leading to 1/ nλi . That being said, the current choice ζ̂i being a sample average has its own advantages, e.g., see Section 5.3. 4.6. Hyperparameter selection Let the estimate b̂i,λ depend on a hyperparameter λ, e.g., the regularization strength in Tikhonov regularization. We derive a cross-validation scheme to adaptively select λ, assuming all estimates b̂1 , . . . , b̂i−1 earlier in the causal order have already been tuned. Our procedure estimates the risk up to an additive constant, similar to cross-validation in nonparametric density estimation. If b̂i,λ ∈ Hi , which for example can be guaranteed for the Gaussian kernel (Section 4.3), we use that Z ∗ 2 ∥b̂i,λ − bi ∥L2 (p) ∝ b̂2i,λ dp − 2⟨b̂i,λ , J ∗ b∗i ⟩Hi (17) Rd

R up to the additive constant Rd (b∗i )2 dp, where we applied the adjoint relation ⟨J b̂i,λ , b∗i ⟩L2 (p) = ⟨b̂i,λ , J ∗ b∗i ⟩Hi for the inclusion J : Hi → L2 (p), see Section 4.1. Let b̂i,λ;I be learned on the training data (xi )i∈I . We estimate the right hand side of equation P (17) using test data (xj )j∈J , where I ∩ J = ∅. Then, |J1 | j∈J b̂i,λ;I (xj )2 is an unbiased R estimate of Rd b̂2i,λ;I dp. For the second term, recall that −J ∗ b∗i = ζi by (8). We estimate ζ̂i;J := P 1 j∈J ẑi (xj , ·). To summarize, we estimate (17) by |J | 2⟨b̂i,λ;I , ζ̂i;J ⟩Hi +

1 X b̂i,λ;I (xj )2 . |J | j∈J

For spectral regularization, b̂i,λ:I is a linear combination of kernel basis functions and we compute ⟨b̂i,λ;I , ζ̂i;J ⟩Hi using the kernel trick. Otherwise, we project ζ̂i;J onto the eigenspaces of L̂i;J correspondingP to the largest eigenvalues {γ : γ ≥ γmin } for some threshold γmin > 0, where L̂i;J f := |J1 | j∈J ki (xj , ·)f (xj ). This is formally justified if ζi lies in the image of Li , e.g., when b∗i ∈ Hi . We can then compute the inner product by the kernel trick, see Appendix B.4. We use standard k-fold crossvalidation, meaning {1, . . . , n} is split into k bins B1 , . . . , Bk of roughly equal size. Set Jl := Bl and Il := {1, . . . , n} \ Jl for l ≤ k. Apply the method described above on (Ii , Ji ) and average the resulting k estimates for a final risk estimate. Repeat over a λ-grid and choose λ with lowest risk. Some practical considerations are as follows. We mostly use kernels for which the eigenvalues of Li decay very rapidly and set γmin = 0.01 by default. How computationally expensive the cross-validation is depends on the regularizer. For example, the Landweber regularizer and its accelerated versions actually obtain bi,λ for all λ > λ0 when computing bi,λ0 . Due to double descent phenomena, we recommend testing very small λ0 . Finally, keep in mind that estimation

R. Schwank and M. Drton/Causal Diffusions

Learned drift vector field Rel. generalization error

Training on dependent samples

0.2 x2 0.4 0.25

x1

13

2

1 0.1

1.75

0.5 Sample autocorrelation

0.9

Fig 1: Learning the drift of SDE (18) from 1000 equilibrium samples. Left: good agreement between learned (blue) and ground truth drift (grey) up to low data density corners; Section 5.1. Right: Sampling from a single diffusion path with varying observation gap ∆t; Section 5.3. Alt text: Left panel: Arrow plot of learned versus ground truth drift. Right panel: line plot of empirical generalization error versus sample autocorrelation. The error remains roughly stable below a correlation of 0.75 and increases strongly afterwards.

errors compound along the causal order, which affects ζ̂i;J and therefore also the accuracy of cross-validation. 5. Simulations and extensions 5.1. Synthetic population dynamics in a random environment Let the population size x2 (t) of some species be influenced by an observable exogenous process x1 (t), e.g., precipitation. Consider logistic growth with capacity decaying to zero when the exogenous factor x1,t is large: !   0.5 1 0  · (1 − x1 (t))  dx(t) = dw(t). (18) dt + 1 0 21 x2 (t) − x2 (t) 2x2 (t) 1+exp(x 1 (t))

⟲

⟳

The numerical constants were chosen to guarantee stable persistence of the population (Li et al., 2011). From this stationary distribution, we non-parametrically recover the drift vector field. In particular, we discover that there is a logistic mechanism and that the carrying capacity is 1/(1 + exp(x1 )). The synthetic training data are n = 1000 vectors in R2 , which approximate an independent sample from the stationary distribution of (18). Practically, this could mean n isolated populations of the species, and we only require one measurement of the population number and the exogenous factor per location. See Section 5.3 for multiple samples. We apply our Tikhonov regularized estimator, requiring as input the diffusivity function 0 ( 10 .5x ) and the structure graph 1 → 2 ; see Example 3.2. We use the Gaussian kernel for 2 both drift components. The regularization parameters λ1 , λ2 are selected via cross-validation; see Section 4.6. The left panel of Figure 1 compares the estimated and ground truth drift vector fields on a rectangle encompassing the bulk of the training data. Both fields agree very well. Discrepancies

R. Schwank and M. Drton/Causal Diffusions

Sigmoid kernel

Node 1 Node 4 Node 7 Tikhonov Spectral Landweber

1

Rel. generalization error 10 10 10

0

Gaussian kernel

14

2

1

5 4

6

10

3

2

103

104 Training samples

103

104 Training samples

3

7

Fig 2: Relative generalization error ∥b̂i − b∗i ∥2L2 (p) /∥b∗i ∥2L2 (p) versus training sample size for selected i = 1, 4, 7. Bottom right: structure graph of the 7-variate distribution. Alt text: Two plots compare the generalization error of the Gaussian kernel and the Sigmoid kernel, with multiple lines representing various drift components and regularization methods.

increase towards the top right corner, a low-density region of the stationary distribution, which contribute negligibly to the L2 (p) generalization error, our metric of choice inspired by the score learning literature (Wibisono et al., 2024). Extending this regression model to multiple covariate processes requires information on the interplay between the covariates, see Appendix B.3. A simple model is to assume independent covariate processes (De Gregorio et al., 2025, eq. 33). More generally, if a causal hierarchy between the processes can be specified, our estimator can be applied directly. If the covariate processes are coupled, e.g., Varughese and Fatti (2008) study the effect of a bi-variate coupled exogenous diffusion on a birth-death process, the covariate dynamics may not be identified from the equilibrium alone. If the covariate drifts can be learned from a different data source, our method can still be applied to the outcome drift component, i.e., to learn how the covariate processes affect the outcome process from equilibrium observations only. 5.2. Synthetic data Consider a seven-dimensional drift structured by the graph in Figure 2, defined via X bi (x) := wj tanh(−qTj xPa(i) ), j∈Pa(i)

where (qj )j∈Pa(i) is a randomly generated orthogonal basis of R|Pa(i)| such that qj,i > 0, and where wj are random positive weights. This construction ensures the drift is reverting towards the origin. We take the diffusivity D to be the constant identity and generate between n = 103 and n = 2 · 104 samples from the stationary distribution by MCMC. We train with Tikhonov, Spectral, and accelerated Landweber regularization using the Gaussian and the Sigmoid kernel. Regularization strengths are tuned by cross-validation, see Section 4.6. For each training set size, we record the relative squared generalization errors ∥b̂i − b∗i ∥2L2 (p) /∥b∗i ∥2L2 (p) for all nodes i = 1, . . . , 7. To reduce the O(n3 ) computational complexity scaling of Tikhonov and Spectral regularization, we apply Nyström’s method, see Vasilev (2026),

R. Schwank and M. Drton/Causal Diffusions

15

√ and reduce K to O( n) rows and columns, covering all cases in Sutherland et al. (2018). This usually works well, however these two methods produced a handful of outliers. In practice, we recommend to keep as many rows and columns as possible, or apply the accelerated Landweber method with naturally only roughly O(n2 ) scaling, which had no performance outliers. The median relative generalization error over N = 100 experiments is shown in Figure 2. The first takeaway is that our method works as intended on a problem of this scale. Predictably, the sigmoid kernel performs better than the Gaussian kernel as it matches the ground truth drift structure. Spectral and Landweber regularization significantly outperform Tikhonov regularization for the sigmoid kernel, consistent with theory that higher qualification methods adapt to drifts which are low complexity in the RKHS. Spectral regularization outperforms for the Gaussian kernel on the root node, possibly because its cross-validation procedure requires one approximation step less. Finally, note that the accelerated Landweber regularization has a rougher error curve due to its discrete number of steps; for example the median step number transitions from two to three after 104 samples for the sigmoid kernel on node four. We did not include the method by Lorch et al. (2024) in this experiment since their implementation assumes known linear self-regulation and does not tune hyper-parameters. Instead, we apply it to a linear drift in Appendix C.1 and demonstrate its ability to approximate DAGstructured drifts under oracle stopping. 5.3. Dependent samples Our main motivation was to require only a single equilibrium sample per individual, however our estimator can also be meaningful when (xj )j=1,...,n are collected from one individual at times tj := j ·∆t for some ∆t > 0. Namely, if ∆t > 0 is independent of n, also termed Pnthe low-frequency regime (Gobet et al., 2004), and if the evolution is ergodic, the average n1 j=1 zi (xj , y) converges to ζi (y) in probability. This suggests our estimator may also be consistent for ergodic processes under low frequency sampling. The finite sample performance depends on how close to equilibrium the individual was at time t = 0. We repeat the simulation from Section 5.1 with low frequency samples, meaning we simulate the SDE (18) by the Euler-Maruyama method, and take n = 1000 samples with time distance ∆t. A generous burnin period ensures x1 approximately follows the stationary distribution. We consider 10 different choices for ∆t, such that the Lag 1 autocorrelation of (xj )j=1,...,n varies between 0.04 and 0.96, and fit the Tikhonov regularized estimator with cross-validation. We estimate the mean squared generalization error relative to ∥b∥2L2 (p) on a fresh sample and repeat this experiment N = 100 times. The right part of Figure 1 shows the average relative generalization error versus the average autocorrelation for each ∆t. As expected, high autocorrelation (∆t → 0) decreases performance. Still, performance is reasonable for low and medium autocorrelation. Although our estimator does not make full use of the available information, we include this example here since multivariate nonparametric implementations for low frequency sampled data are rare. 5.4. Irreversible diffusion models Given observations x1 , . . . , xn from an unknown density p, generative modeling aims to produce a fresh sample x̃1 , . . . x̃m from p as accurately as possible (Song and Ermon, 2019). Diffusion models learn the score function s(x) := ∇ log p(x) and sample via the Langevin SDE √ (19) dx̃(t) = s(x̃(t))dt + 2dw(t),

R. Schwank and M. Drton/Causal Diffusions

=0 =1 =5 = 10 = 20

15 Energy distance

16

10 5 0

0

Time

2

Fig 3: Consider score s and causal drift b from Theorem 5.1 for some 5-variate Gaussian density √ p; see Appendix C.2. We simulate dx(t) = (s − ω · (s − b))(x(t))dt + 2dw(t) and plot the energy distance between x(t) and the target p for various ω, t. As discussed in Section 5.4, ω ̸= 0 speeds up convergence. Alt text: Multiple lines show the decrease in energy distance over time. The larger the absolute value of omega, the faster the decay.

which has p as its stationary distribution as desired. If a causally structured drift inducing p exists, it could augment the score drift: assume the function b also induces p as a stationary density when used as a drift in (19). Since both then solve the Fokker-Planck equation, we have ∇ · (bp) = ∆p = ∇ · (sp), therefore ∇ · ((b − s) p) = 0. Applying the results by Hwang et al. (2005) to the drift functions s + ω · (b − s) for ω ∈ R, we find that any ω ̸= 0 usually leads to strictly faster convergence against p compared to the score function s alone, i.e., ω = 0. We illustrate this experimentally for a Gaussian distribution p in Figure 3; for simulation details see Appendix C.2. A causally structured drift b is guaranteed to exist for any given density p when the DAG is complete as we show below. This generalizes a result for Gaussian distributions (Dettling et al., 2023). Note that, even for the complete DAG, b ̸= s; for example by Schwarz’s theorem ∂i sj = ∂j si , which does not hold for b. Theorem 5.1. Let 0 < p ∈ C 2 (Rd ) be a probability density and σ be constant. For any j ≤ d, define 1 : j := 1, . . . , j, let fj+1 and Fj+1 denote the density and cumulative distribution function of xj+1 given x1:j , respectively, and assume Fj+1 ∈ C 3 (Rj ). 2 p′ (x ) Define b : Rd → Rd component-wise as b1 (x) := σ2 p11 (x11 ) and bj+1 recursively as bj+1 :=

σ 2 ∂j+1 fj+1 + ∆1:j Fj+1 + 2(∇1:j log p1:j )T ∇1:j Fj+1 (∇1:j Fj+1 )T b1:j + . 2 fj+1 fj+1

(20)

Then, b is structured according to the complete DAG with topological order 1, . . . , d (including 2 self loops), and satisfies the Fokker-Planck equation σ2 ∆p = ∇ · (bp). Note that a different topological ordering than 1, . . . , d could further affect the convergence speed. On a high level, perturbing the score by the causal drift makes the diffusion process irreversible, which has long been known to speed up convergence (Hwang et al., 2005; ReyBellet and Spiliopoulos, 2015). For linear drifts, the non-reversible perturbation leading to the fastest convergence has been characterized by Lelièvre et al. (2013). While the aforementioned literature takes a sampling perspective, i.e., assuming the target density p is (partially) known, generative modeling assumes only observations from p are available.

R. Schwank and M. Drton/Causal Diffusions

17

6. Discussion and limitations We established nonparametric identifiability under very weak conditions from which future estimators, also parametric ones, can profit. A canonical extension would be to consider cyclic models and models with unknown structure graphs, in parallel with recent advances for linear drifts. A further important research direction concerns the case of unknown diffusivity, which remains open even for linear systems. Our proposed estimator is computationally feasible, even when the Fokker-Planck forward problem becomes intractable due to the curse of dimensionality. We theoretically guarantee the estimator’s consistency in L2 (p), which practically can be achieved by the Gaussian kernel in our package. Although the consistency result assumes bounded drifts for technical convenience, the estimator shows strong empirical performance even when the drift is unbounded. Extending our method to new kernels requires access to its univariate partial integrals. We have so far only used kernels from Table 1 allowing analytical expressions, however numerical approximations to the univariate integral should also work. A notable theoretical limitation of our estimator is the recursive plug-in strategy, which potentially could accumulate estimation errors exponentially along the topological order. Experimentally, the more significant contributor to the estimation error seems to be the number of parent variables, also see Vasilev (2026). Our approach, similar to the referenced score literature, trades global optimality in a joint loss for computational efficiency and interpretability of the component functions, as these are optimized individually. Practically, our method could potentially be improved by taking biases of cross-validation into account (Arlot and Celisse, 2010). Finally, note that many popular kernels do not produce drift estimates which can be sampled from with constant diffusivity since they decay to zero for large input. One can instead fit kernels weighted with a linear component at the cost of an extra hyperparameter (Sriperumbudur et al., 2017). Acknowledgements We thank Viktor Vasilev for his contributions to this project through his master’s thesis, conducted under the supervision of R.S. and M.D., as referenced throughout the paper. We acknowledge support from the DAAD programme Konrad Zuse Schools of Excellence in Artificial Intelligence, sponsored by the Federal Ministry of Research, Technology and Space, from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement No 883818), and from the German Federal Ministry of Education and Research and the Bavarian State Ministry for Science and the Arts. The authors take full responsibility for the contents of this article. References Amorino, C. and Gloter, A. (2023). Estimation of the invariant density for discretely observed diffusion processes. Statistics 57 213–259. Améndola, C., Boege, T., Hollering, B. and Misra, P. (2025). Structural Identifiability of Graphical Continuous Lyapunov Models. Arcones, M. A. and Giné, E. (1993). Limit theorems for U -processes. Ann. Probab. 21 1494– 1542. Arlot, S. and Celisse, A. (2010). A survey of cross-validation procedures for model selection. Stat. Surv. 4 40–79. Bhattacharya, R. and Waymire, E. C. (2023). Continuous parameter Markov processes and stochastic differential equations. Graduate Texts in Mathematics 299. Springer, Cham.

R. Schwank and M. Drton/Causal Diffusions

18

Boege, T., Drton, M., Hollering, B., Lumpp, S., Misra, P. and Schkoda, D. (2025). Conditional independence in stationary distributions of diffusions. Stochastic Process. Appl. 184 Paper No. 104604, 16. Bongers, S. and Mooij, J. M. (2018). From random differential equations to structural causal models: The stochastic case. arXiv preprint arXiv:1803.08784 3. Botvinick-Greenhouse, J., Martin, R. and Yang, Y. (2023). Learning dynamics on invariant measures using PDE-constrained optimization. Chaos 33 Paper No. 063152, 22. De Gregorio, A., Frisardi, D., Iacus, S. and Iafrate, F. (2025). Adaptive elastic-net estimation for sparse diffusion processes. Stat. Inference Stoch. Process. 28 Paper No. 22, 35. De Vito, E., Rosasco, L., Caponnetto, A., De Giovannini, U. and Odone, F. (2005). Learning from examples as an inverse problem. J. Mach. Learn. Res. 6 883–904. Dettling, P., Homs, R., Améndola, C., Drton, M. and Hansen, N. R. (2023). Identifiability in continuous Lyapunov models. SIAM J. Matrix Anal. Appl. 44 1799–1821. Dicker, L. H., Foster, D. P. and Hsu, D. (2017). Kernel ridge vs. principal component regression: minimax bounds and the qualification of regularization operators. Electron. J. Stat. 11 1022–1047. Evans, L. C. (2010). Partial differential equations, second ed. Graduate Studies in Mathematics 19. American Mathematical Society, Providence, RI. Gobet, E., Hoffmann, M. and Reiß, M. (2004). Nonparametric estimation of scalar diffusions based on low frequency data. Ann. Statist. 32 2223–2253. Hwang, C.-R., Hwang-Ma, S.-Y. and Sheu, S.-J. (2005). Accelerating diffusions. Ann. Appl. Probab. 15 1433–1444. Inglese, G. (2002). Recovering a vector field with the aid of controlled noise. J. Inverse IllPosed Probl. 10 187–193. Karatzas, I. and Ruf, J. (2016). Distribution of the time to explosion for one-dimensional diffusions. Probab. Theory Related Fields 164 1027–1069. Khasminskii, R. (2012). Stochastic stability of differential equations, second ed. Stochastic Modelling and Applied Probability 66. Springer, Heidelberg. Lacker, D. and Zhang, J. (2023). Stationary solutions and local equations for interacting diffusions on regular trees. Electron. J. Probab. 28 Paper No. 4, 37. Lee, H. and Trutnau, G. (2021). Existence, uniqueness and ergodic properties for timehomogeneous Itô-SDEs with locally integrable drifts and Sobolev diffusion coefficients. Tohoku Math. J. (2) 73 159–198. Lelièvre, T., Nier, F. and Pavliotis, G. A. (2013). Optimal non-reversible linear drift for the convergence to equilibrium of a diffusion. J. Stat. Phys. 152 237–274. Li, X., Gray, A., Jiang, D. and Mao, X. (2011). Sufficient and necessary conditions of stochastic permanence and extinction for stochastic logistic populations under regime switching. J. Math. Anal. Appl. 376 11–28. Lorch, L., Krause, A. and Schölkopf, B. (2024). Causal Modeling with Stationary Diffusions. In AISTATS 2024. Proceedings of Machine Learning Research 238 1927–1935. PMLR. Lorch, L., Zhang, J., Bunne, C., Krause, A., Schölkopf, B. and Uhler, C. (2026). Latent Causal Diffusions for Single-Cell Perturbation Modeling. Nickl, R. and Ray, K. (2020). Nonparametric statistical inference for drift vector fields of multi-dimensional diffusions. Ann. Statist. 48 1383–1408. Pedretscher, B., Kaltenbacher, B. and Pfeiler, O. (2019). Parameter identification and uncertainty quantification in stochastic state space models and its application to texture analysis. Appl. Numer. Math. 146 38–54. Pereverzyev, S. (2022). An introduction to artificial intelligence based on reproducing kernel Hilbert spaces. Compact Textbooks in Mathematics. Birkhäuser/Springer, Cham.

R. Schwank and M. Drton/Causal Diffusions

19

Peters, J., Bauer, S. and Pfister, N. (2022). Causal Models for Dynamical Systems. In Probabilistic and Causal Inference: The Works of Judea Pearl. ACM Books 36 671–690. ACM. Peters, J., Janzing, D. and Schölkopf, B. (2017). Elements of causal inference. Adaptive Computation and Machine Learning. MIT Press, Cambridge, MA. Pratapa, A., Jalihal, A. P., Law, J. N., Bharadwaj, A. and Murali, T. M. (2020). Benchmarking algorithms for gene regulatory network inference from single-cell transcriptomic data. Nature Methods 17 147–154. Rey-Bellet, L. and Spiliopoulos, K. (2015). Irreversible Langevin samplers and variance reduction: a large deviations approach. Nonlinearity 28 2081–2103. Rohbeck, M., Clarke, B., Mikulik, K., Pettet, A., Stegle, O. and Ueltzhöffer, K. (2024). Bicycle: Intervention-Based Causal Discovery with Cycles. In Proceedings of the Third Conference on Causal Learning and Reasoning 236 209–242. PMLR. Rosasco, L., Belkin, M. and De Vito, E. (2010). On learning with integral operators. J. Mach. Learn. Res. 11 905–934. Song, Y. and Ermon, S. (2019). Generative Modeling by Estimating Gradients of the Data Distribution. In Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, Vancouver, BC, Canada 11895–11907. Sriperumbudur, B. K., Fukumizu, K., Gretton, A., Hyvärinen, A. and Kumar, R. (2017). Density Estimation in Infinite Dimensional Exponential Families. J. Mach. Learn. Res. 18 57:1–57:59. Steinwart, I. and Christmann, A. (2008). Support vector machines. Information Science and Statistics. Springer, New York. Sun, Y. and Kumar, M. (2014). Numerical solution of high dimensional stationary FokkerPlanck equations via tensor decomposition and Chebyshev spectral differentiation. Comput. Math. Appl. 67 1960–1977. Sutherland, D. J., Strathmann, H., Arbel, M. and Gretton, A. (2018). Efficient and principled score estimation with Nyström kernel exponential families. In AISTATS 2018. Proceedings of Machine Learning Research 652–660. PMLR. Varando, G. and Hansen, N. R. (2020). Graphical continuous Lyapunov models. In Proceedings of the 36th Conference on Uncertainty in Artificial Intelligence (UAI) 989–998. PMLR. Varughese, M. M. and Fatti, L. P. (2008). Incorporating environmental stochasticity within a biological population model. Theoretical Population Biology 74 115–129. Vasilev, V. (2026). Nonparametric Estimation of Causal Lyapunov Models, Master’s thesis, Technical University of Munich. Wibisono, A., Wu, Y. and Yang, K. Y. (2024). Optimal score estimation via empirical Bayes smoothing. In COLT. Proceedings of Machine Learning Research 247 4958–4991. PMLR. Zhao, W., Fertig, E. J. and Stein-O’Brien, G. (2026). CycleGRN: Inferring Gene Regulatory Networks from Cyclic Flow Dynamics in Single-Cell RNA-seq. bioRxiv. Zhou, Y., Shi, J. and Zhu, J. (2020). Nonparametric Score Estimators. In ICML 2020. Proceedings of Machine Learning Research 119 11513–11522. PMLR.

R. Schwank and M. Drton/Causal Diffusions

Appendix A

20

Proofs for the identifiability section

Theorem 3.7. Assume b∗ , b ∈ L1 (Rd , Rd , p) both solve a Fokker-Planck equation (4) with the same density p > 0 and same diffusivity D ∈ L1 (Rd , Rd×d , p). When b, b∗ are structured according to a common DAG D, then b = b∗ Lebesgue-almost everywhere. Proof. We show the claim by induction over the topological ordering of D, which we assume to be 1, . . . , d without loss of generality as discussed. Setting i = 1 in Lemma 3.6 and recalling that pa(1) = ∅ by assumption, we obtain Z Z 1 D11 (x)φ′′ (x1 )p(x)dx ∀φ ∈ C0∞ (R), βφ′ dp1 = 2 Rd R where β can be either of b1 and b∗1 . Taking R the difference of both versions of this equation, once for β := b1 and once for β := b∗1 , we find R δφ′ dp1 = 0 where δ := b1 − b∗1 . This means δ · p1 has a weak derivative which is zero almost everywhere, implying that δ · p1 = 0 almost everywhere (Lemma D.2). As p1 > 0 by p > 0, we conclude δ = 0 almost everywhere, i.e. b1 = b∗1 almost everywhere. Assume bS = b∗S with S = {1, . . . , i − 1} for some i ≥ 2. In particular, bpa(i) = b∗pa(i) . Taking R the difference of equation (6) with b and equation (6) with b∗ , we find R|Pa(i)| δ ∂i φ dpPa(i) = 0 with δ := bi − b∗i . Since φ ranges through C0∞ (R|Pa(i)| ), the function δ · pPa(i) has one zero weak derivative, which implies δ · pPa(i) = 0 almost everywhere (Lemma D.2). By pPa(i) > 0 we conclude bi = b∗i almost everywhere. Rz Lemma A.1 (Feller’s explosion test (Bhattacharya and Waymire, 2023)). Let I(z) := 0 σ22 b(v)dv. Consider Z ∞ Z y Z 0 Z 0 −I(y) I(v) −I(y) e e dvdy and e eI(v) dvdy. 0

0

−∞

y

If the left integral is finite, xt solving the SDE (5) explodes towards +∞ with positive probability. Similarly, the right integral indicates explosion towards −∞. 2

′

Lemma 3.4. The solution of the SDE (5) with drift bc (x) = σ2 pp + pc explodes with positive probability if c ̸= 0. Proof. We consider the case c > 0 and prove explosion towards +∞ by showing that the left integral in Feller’s explosion test (Lemma A.1) is finite. When c < 0, then the right integral is finite and explosion is towards −∞. First, compute Z z Z z Z z 2c 2c 2 ′ b (v)dv = (log p) (v) + dv = log p(z) − log p(0) + dv. I(z) = 2 c 2 p(v) 2 p(v) σ σ σ 0 0 0 R∞ Ry R∞R∞ We show that 0 v eI(v)−I(y) dydv < ∞. By Fubini’s theorem, 0 e−I(y) 0 eI(v) dvdy < ∞, which is the claim. For 0 ≤ v ≤ y, we compute Z ∞Z ∞ Z ∞Z ∞ p(v) − Rvy σ22c dw p(w) eI(v)−I(y) dydv = e dydv = p(y) 0 v 0 v Z ∞ Z  −σ 2 p(v) ∞ d  − Rvy σ22c dw p(w) e dydv = 2 dy 0 v Z ∞ 2 Z ∞ 2 R   2c σ σ − ∞ dw p(v) 1 − e v σ2 p(w) dv ≤ p(v) · 1 dv < ∞. 2 2 0 0

R. Schwank and M. Drton/Causal Diffusions

21

Lemma 3.5. Let b∗ , b ∈ L1 (R, p) both solve a Fokker-Planck equation (4) with the same density p > 0 and diffusivity D ∈ L1loc (R, p). Then, b = b∗ almost everywhere. Proof. Taking the difference of the Fokker-Planck equation (4) for b and b∗ respectively, we find Z ((b − b∗ )p)(x) φ′ (x)dx = 0 ∀φ ∈ C0∞ (R). R

From Lemma D.2 it follows that (b − b∗ )p = 0 almost everywhere. By p > 0 it follows that b = b∗ almost everywhere. Lemma 3.6. Let the D-structured drift b ∈ L1 (Rd , Rd , p) solve the Fokker-Planck equation (4) with density p and diffusivity D ∈ L1 (Rd , Rd×d , p). For any index i ≤ d and any φ = φ̃ ◦ πPa(i) with φ̃ ∈ C0∞ (R|Pa(i)| ) it holds that Z Z Z 1 (6) bi ∂i φ dpPa(i) = − bTpa(i) ∇pa(i) φ dp + ⟨DPa(i),Pa(i) , ∇2Pa(i) φ⟩F dp. 2 |Pa(i)| d d R R R Proof. Abbreviate A := Pa(i). When |A| = d, the equation to prove is simply the Fokker-Planck equation (4) which holds by assumption. In the following, we consider the case |A| < d. For n ∈ N, define the approximation φn (x) := φ(x) · γn (x−A ) with γn ∈ C0∞ (Rd−|A| ) from Lemma D.4. Particularly, C := supn∈N ∥γn ∥W 2,∞ < ∞ and φn ∈ C0∞ (Rd ), which implies Z Z 1 bT ∇φn dp = ⟨D, ∇2 φn ⟩F dp (21) 2 Rd Rd by the Fokker-Planck relation (4). We examine the limit as n → ∞. Let j ≤ d. When j ∈ A, we have |∂j φn (x)| = |∂j φ̃(xA ) · d γn (x−A )| ≤ ∥∂j φ̃∥L∞ · C < ∞ and ∂j φn (x) → R ∂j φ̃(x−A ) · 1 =R φ(x) as n → ∞ for all x ∈ R . The dominated convergence theorem implies Rd bj ∂j φn dp → Rd bj ∂j φ dp as n → ∞. When j ∈ / A, we have |∂j φn (x)| = |φ̃(xA ) · ∂j γn (x−A )| ≤ ∥φ̃∥L∞ · C < ∞ and ∂j φn (x) → φ̃(x ) · 0 = 0 as n → ∞ for all x ∈ Rd . The dominated convergence theorem implies −A R b ∂ φ dp → 0 as n → ∞. Rd j j n With a similar case distinction for ⟨D, ∇2 φn ⟩F , by letting n → ∞ in equation (21) we obtain that Z Z 1 T ⟨DA,A , ∇2A φ⟩F dp. bA ∇A φ dp = 2 d d R R All R that remains R for proving the Lemma’s claim is to solve for the bi ∂i φ-term and to recall that b ∂ φ dp = b ∂ φ dpA since both bi and ∂j φ only depend on xA . Rd i i R|A| i i Appendix B B.1

Proofs and details for the theoretical analysis

Setup

Lemma 4.1. Assume p is bounded, let i ≤ d, and let y ∈ Rd . If Ki,y ∈ W 2,∞ (Rd ) only depends on xPa(i) and ∂i Ki,y = ki (·, y), then ζi (y) = Ex∼p [ zi (x, y) ], where 1 zi (x, y) := b∗pa(i) (x)T ∇pa(i) Ki,y (x) − ⟨DPa(i),Pa(i) (x), ∇2Pa(i) Ki,y (x)⟩F . 2

R. Schwank and M. Drton/Causal Diffusions

22

Proof. Let y ∈ Rd , let i ≤ d, and abbreviate A := Pa(i). By assumption, there is K̃ ∈ W 2,∞ (R|A| ) such that K̃(xA ) = Ki,y (x) almost everywhere. As p was assumed to be bounded, there is a sequence (K̃n )n∈N ⊂ C0∞ (R|A| ) with ∥K̃ − K˜n ∥W 2,2 (R|A| ,pA ) → 0 as n → ∞ by Lemma D.3. Set Kn (x) := K̃n (xA ). By Lemma 3.6, we have Z Z Z 1 bi ∂i Kn dpA = − (b∗pa(i) )T ∇pa(i) Kn dp + ⟨DA,A , ∇2A Kn ⟩F dp. (22) 2 |A| d d R R R Using Hölder’s inequality, we have Z Rd

(b∗A )T ∇A Ki,y dp −

Z Rd

(b∗A )T ∇A Kn dp ≤

Z Rd

∥b∗A ∥2 ∥∇A Ki,y − ∇A Kn ∥2 dp n→∞

≤ ∥b∗ ∥L2 (Rd ,Rd ,p) · ∥K̃ − K̃n ∥W 2,2 (R|A| ,pA ) −−−−→ 0. Together with an analogous result for ⟨D, ∇2A Kn ⟩F , this implies the Lemma’s claim by letting n → ∞ in equation (22). B.2

Statistical analysis

Lemma 4.4. For any i ≤ d it holds with probability at least 1 − δ that √ X a2 κB 2 log1/2 2δ √ +κ ∥b̂j − b∗j ∥E n ∥ζ̂i − ζi ∥L2 (p) ≤ i n

(10)

j∈pa(i)

Proof. By the triangle inequality and ∥∂j Ki,· (·)∥L∞ ≤ κ for all j ∈ pa(i), we find that n

∥ζ̂i − ζi ∥L2 (p) ≤

1X zi (xj , ·) − ζi n j=1

+κ L2 (p)

X

∥b̂j − b∗j ∥E n .

j∈pa(i)

Note E[ zi (x1 , ·) ] = ζi (·) using Lemma 4.1 since ∂i Ki,y = ki (·, y) almost everywhere for all y ∈ Rd , and that X X ∥ζi − zi (x1 , ·)∥L2 (p) ≤ ∥ζi ∥L∞ (Rd ) + κB + κB ≤ a2i κB. j∈pa(i)

j,k∈Pa(i)

Then, by Hoeffding’s inequality for separable Hilbert spaces (Pereverzyev, 2022) it holds that q n a2i κB 2 log( 2δ ) 1X √ zi (xj , ·) − ζi ≤ n j=1 n L2 (p)

with probability at least 1 − δ. Lemma 4.5. Let n ≥ 3. For any i ≤ d it holds with probability at least 1 − δ that √ 1/2 1/2 2∨c1 X a2i κB 3 a2i κB · (log n + log δ ) n √ ∥ζ̂ i − ζ i ∥E ≤ + +κ ∥b̂j − b∗j ∥E n √ √ n c2 n j∈pa(i)

for some global positive constants c1 , c2 .

R. Schwank and M. Drton/Causal Diffusions

23

Pn P Pn Proof. By ζ̂i = n1 k=1 zi (xk , ·) + j∈pa(i) n1 l=1 (b̂j (xl ) − b∗j (xl )) · ∂xj Ki,· (xl ), the triangle inequality gives that v !2 u n n u1 X 1 X X t ∥ζ̂ i − ζ i ∥E n ≤ zi (xk , xj ) − ζi (xj ) + κ ∥b̂j − b∗j ∥E n . (23) n j=1 n k=1

j∈pa(i)

Define the following kernel hx indexed by the parameter x ∈ Rd : hx (a, ã) := (zi (a, x) − ζi (x)) · (zi (ã, x) − ζi (x)) . Further define the index sets S(j) := {(k, l) ∈ {1, . . . , n}2 | k ̸= l, k ̸= j, l ̸= j}. We multiply out the square under the root in (23). Using supa,ã∈Rd |hx (a, ã)| ≤ (a2i κB)2 and that {1, . . . , n}2 \ S(j) contains at most 3n elements, we find that n

1X n j=1

n

1X zi (xk , xj ) − ζi (xj ) n k=1

!2 ≤

n 1X 1 3a4i κ2 B 2 + n n j=1 n2

X

hxj (xk , xl ).

k,l∈S(j)

Next, we apply a U -statistic bound to the second term. Let x ∈ Rd be a deterministic constant. The kernel hx is symmetric and bounded by h∞ := (a2i κB)2 independently of x. For k ̸= l, we have E[ hx (xk , xl ) ] = 0 and a 7→ E[ hx (a, xl ) ] is the constant zero function, hence hx is p-canonical. By Arcones and Giné (1993), there are positive global constants c1 , c2 such that for any j ∈ {1, . . . , d} and t > 0:  P



  c2 t hx (xk , xl ) ≥ t  ≤ c1 exp − ⇒ h∞ k,l∈S(j)       X 1 c2 t · n c2 t · n2 hx (xk , xl ) ≥ t  ≤ c1 exp − ≤ c1 exp − P 2 . n h∞ · (n − 1) h∞

1 n−1

X

k,l∈S(j)

By independence of xj and (xk , xl )k,l∈S(j) the above inequality also holds for the conditional probability given xj = x. Taking the expectation with respect to xj , we have shown:     X 1 c2 t · n P 2 hxj (xk , xl ) ≥ t  ≤ c1 exp − . n h∞ k,l∈S(j)

Using the triangle inequality and a union bound, it follows that     n X X 1 c2 t · n 1 h log(n) ∞   ≤ c1 exp − . P hxj (xk , xl ) ≥ t + n j=1 n2 c2 n h∞ k,l∈S(j)

1 ,2)/δ) The claim follows by taking t := h∞ log(max(c and from sub-additivity of the square root. c2 n

Lemma 4.6. For any i ≤ d it holds with probability at least 1 − δ that √ 2κ 2 log1/2 2δ √ ∥L̂i − Li ∥L(Hi ,Hi ) ≤ . n

R. Schwank and M. Drton/Causal Diffusions

24

Proof. Apply Hoeffding’s inequality on the space of Hilbert-Schmidt operators using ∥Li ∥HS ≤ κ from Section 4.1 and ∥L̂i ∥HS ≤ κ. Also see (Pereverzyev, 2022). Lemma 4.7. If ∥ζ̂ i − ζ i ∥E n ≤ ε1 , ∥ζ̂i − ζ∥L2 (p) ≤ ε2 and ∥L̂i − Li ∥L(Hi ,Hi ) ≤ ε3 , then ε3 ∥b∗i ∥L2 (p) 2ε1 + , λi λ ! i √ 3/2 ε1 ε3 ε ε3 ε1 + ε2 + 33/2 ∥b∗i ∥L2 (p) + + 3/2 . ∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) + λi λi λ λ ∥b̂i − b∗i ∥E n ≤ ∥bi,λi − b∗i ∥E n +

i

(11) (12)

i

Proof. Fix i ∈ {1, . . . , d}; we write λ instead of λi to reduce indices. To apply functional analysis tools, we note Lemma 4.2 may equivalently be stated as 1 f (L̂i + λI)−1 f = − Jx∗ (Jx Jx∗ + λI)−1 Jx f + λ λ

(24)

for all f ∈ Hi . Further, note our definition of b̂i can be written as b̂i =

1 ∗ 1 J (Jx Jx∗ + λI)−1 ζ̂ i − ζ̂i . λ x λ

(25)

Our proof uses the following decomposition: b∗i − b̂i = b∗i − bi,λ + bi,λ − (L̂i + λI)−1 (−ζi )

(26)

+ (L̂i + λI)−1 (−ζi ) − b̂i , where bi,λ := (Li + λI)−1 (−ζi ) = (Li + λI)−1 J ∗ b∗i is the Tikhonov regularizer of b∗i . For the bound on b̂i − b∗i , we evaluate both sides of (26) at the random variables x1 , . . . , xn and take ∥ · ∥E n on both sides. As that Evaluation of RKHS elements can be written by applying the sampling operator Jx : Hi → E n , we have found that ∥b∗i − b̂i ∥E n ≤ ∥b∗i − bi,λ ∥E n + ∥Jx ((Li + λI)−1 − (L̂i + λI)−1 )ζi ∥E n + 1 1 ∥Jx Jx∗ (Jx Jx∗ + λI)−1 (ζ̂ i − ζ i )∥E n + ∥ζ̂ i − ζ i ∥E n , (27) λ λ where we denoted bi,λ = (bi,λ (x1 ), . . . , bi,λ (xn )) ∈ Rn . Using A−1 − B −1 = B −1 (B − A)A−1 for invertible operators A, B, we have that (Li + λI)−1 − (L̂i + λI)−1 = (L̂i + λI)−1 (L̂i − Li )(Li + λI)−1 .

(28)

Recall that L̂i = Jx∗ Jx , that Li = J ∗ J and finally that ζi = −J ∗ b∗i . We therefore bound: 1 ∥Jx (Jx∗ Jx + λI)−1 ∥L(Hi ,E n ) ≤ √ λ 1 ∗ −1 ∗ ∥(J J + λI) J ∥L(L2 (p),Hi ) ≤ √ . λ

(29) (30)

One way to see (31) is to multiply the operator with its adjoint and to subsequently apply µ 1 the spectral theorem for J ∗ J. Using supµ≥0 (µ+λ) 2 ≤ λ and taking the square root yields (31), also see De Vito et al. (2005). Equation (30) follows by replacing J with Jx . Applying an

R. Schwank and M. Drton/Causal Diffusions

25

operator norm factorization bound, the second term on the right hand side of (27) is bounded by ε3 ∥b∗i ∥L2 (p) · λ−1 . Using ∥Jx Jx∗ (Jx Jx∗ + λI)−1 ∥L(E n ,E n ) ≤ 1 and ∥ζ̂ i − ζ i ∥E n ≤ ε1 by assumption, the bottom line of (27) is bounded by 2ε1 · λ−1 . Thus, ∥b∗i − b̂i ∥E n ≤ ∥b∗i − bi,λ ∥E n +

ε3 ∥b∗i ∥L2 (p) 2ε1 + λ λ

as claimed. Next, we study b̂i − b∗i in the L2 (p) norm and again use the decomposition (26). Making the embedding J : Hi → L2 (p) explicit for all elements from Hi on the right side of (26) and applying the triangle inequality, we find that ∥b∗i − b̂i ∥L2 (p) ≤ ∥b∗i − Jbi,λ ∥L2 (p) + ∥J((Li + λI)−1 − (L̂i + λI)−1 )ζi ∥L2 (p) + 1 1 ∥JJx∗ (Jx Jx∗ + λI)−1 (ζ̂ i − ζ i )∥L2 (p) + ∥ζ̂i − Jζi ∥L2 (p) . (31) λ λ Treating the second term on the right hand side with (28) and (30), we find that 1 ∥J((Li + λI)−1 − (L̂i + λI)−1 )ζi ∥L2 (p) ≤ ∥J(L̂i + λI)−1 ∥L(Hi ,L2 (p)) · ε3 · √ · ∥b∗i ∥L2 (p) . λ Let Oλ := (L̂i + λI)−1 . To resolve the mismatch between population operator J and sample operator L̂i , we consider (JOλ )∗ JOλ = Oλ Jx∗ Jx Oλ + Oλ (J ∗ J − Jx∗ Jx )Oλ . Using ∥Oλ Jx∗ ∥L(E n ,Hi ) = ∥Jx Oλ ∥L(Hi ,E n ) ≤ √1λ by (29) and Oλ∗ = Oλ , and using ∥Oλ ∥L(Hi ,Hi ) ≤ 1 λ , we have shown r ∥J(L̂i + λI)−1 ∥L(Hi ,L2 (p)) = ∥JOλ ∥L(Hi ,L2 (p)) ≤

√ ε3 1 ε3 1 + 2 ≤√ + . λ λ λ λ

We treat the third term in (31) similarly. With Uλ := (Jx Jx∗ + λI)−1 , we find that (JJx∗ Uλ )∗ JJx∗ Uλ = Uλ Jx Jx∗ Jx Jx∗ Uλ + Uλ Jx (J ∗ J − Jx∗ Jx )Jx∗ Uλ . Using ∥Jx Jx∗ Uλ ∥L(Hi ,E n ) ≤ 1 and ∥Jx∗ Uλ ∥L(E n ,Hi ) ≤ √1λ , we have shown that ∥JJx∗ (Jx∗ Jx + λI)−1 (ζ̂ i − ζ i )∥L2 (p) ≤ ε1 ·

r 1+

ε3 . λ

Collecting all terms, we have shown that ∥b∗i − b̂i ∥L2 (p) ≤ ∥b∗i − Jbi,λ ∥L2 (p) + as claimed.



 √ ε1 ε3 ε1 ε3  ε3 3/2 ε2 + ∥b∗i ∥L2 (p) + + 3/2 + λ λ λ λ λ

R. Schwank and M. Drton/Causal Diffusions

26

Theorem 4.8. Let i ≤ d and λi (n) ≥ √cn for some constant c > 0. Then,
O(log n) √ log 2δ λi n
X O(1) + log 2δ ∥b̂j − b∗j ∥E n , λi

∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) +

(13)

j∈pa(i)

∥b̂i − b∗i ∥E n ≤ ∥bi,λi − b∗i ∥E n + +

O(log n) √ · log( 2δ ) λi n

2κ X ∥b̂j − b∗j ∥E n λi

(14)

j∈pa(i)

with probability at least 1 − δ for any δ ∈ (0, 1). Here, O(·) is as n → ∞. Proof. Any O(·)-notation we use is with respect to n → ∞. We define the shorthand PTE := P ∗ n j∈pa(i) ∥b̂j − bj ∥E for the cumulative parent training error. Further, note that for any any 0 ≤ α ≤ γ ≤ 1 and β ≥ 2, it holds that logα (β/2) logα (β/δ) ≤ + log(2)α−γ . γ logγ (2) δ∈(0,1) log (2/δ) sup

(32)

Let δ ∈ (0, 1) and insert δ/3 into Lemmas 4.6, 4.5, and 4.4. By equation (32), we may unify √ log1/2 6δ = log1/2 ( 2δ ) · O(1) and a2i κB 3 = log1/2 ( 2δ ) · O(1), to name two examples. Finally, applying a union bound over 3δ , we obtain an event of probability at least 1 − δ on which simultaneously √ O( log n) √ ∥ζ̂ i − ζ i ∥E n ≤ · log1/2 2δ + κ PTE, (33) n O(1) (34) ∥ζ̂i − ζi ∥L2 (p) ≤ √ · log1/2 2δ + κ PTE, n O(1) (35) ∥L̂i − Li ∥L(Hi ,Hi ) ≤ √ · log1/2 2δ . n Inserting these bounds into Lemma 4.7, we have with probability at least 1 − δ that √ O(1) O( log n) 2κ 1/2 2 ∗ ∗ √ ∥b̂i − bi ∥E n ≤ ∥bi,λi − bi ∥E n + √ · log · log1/2 2δ + PTE . δ + λi λi n λi n √ Combining O(1) + O( log n) = O(log n) and using log1/2 2δ = O(1) · log 2δ by equation (32), the claim (13) follows. Continuing, Lemma 4.7 shows that O(1) O(1) √ ∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) + √ · log1/2 ( 2δ ) + · log3/4 ( 2δ ) λi n (λi n)3/2 √  √  O( log n) 2κ 1 O( log n) O(1) 1/2 2 1/2 2 √ √ + log log log1/4 2δ δ + λ PTE + 3/2 δ + κ PTE λi n n n1/4 i λi ! √ √ O(1) · log1/4 2δ O( log n) O( log n) κ 1/2 2 3/4 2 ∗ √ √ = ∥bi,λi −bi ∥L2 (p) + log ( δ )+ √ 3/2 log ( δ )+ 2+ PTE . λi λi n (λi n) (λi n)1/2 √ By assumption, (λi n)−1/2 = O(1). Applying equation (32) to the log terms, in particular 2 + log1/4 2δ = O(1) · log 2δ , yields the claim (14).

R. Schwank and M. Drton/Causal Diffusions

27

Corollary 4.9. Let dij be the length of longest path from i to j in D, excluding self-loops and therefore possibly zero. Let di := maxj dji and set ri := 6 · 3di . Use Tj (λ) := ∥bj,λ − b∗j ∥L2 (p) to abbreviate the Tikhonov population approximation error. If 2/3dji

λi (n) ≥ n−2/ri ∨j∈an(i) Tj

(λj (n))

(15)

for all i ≤ d, then the expected generalization error is bounded as follows: h i O(log n) E ∥b̂i − b∗i ∥L2 (p) ≤ + O(1) n1/ri

X

1/3dji

Tj

(λj ) ∨ Tj (λj ).

(16)

j∈an(i)∪{i}

In particular, if Hi is dense in L2 (pPa(i) ) for all i ≤ d, the condition (15) allows choices λi (n) with limn→∞ λi (n) = 0, for which then h i lim E ∥b̂i − b∗i ∥L2 (p) = 0. n→∞

√ Further, consider Tj (λ) = O( λ · logαj λ1 ) as λ → 0 for all j ≤ d and constants αj ≥ 0. For 1 ), the choice λi := n−2(1−εi )/ri with εi := di · ε guarantees that any ε ∈ (0, 2d h i E ∥b̂i − b∗i ∥L2 (p) = O(n−(1−2εi )/ri ). Proof. Let δ ∈ (0, 1). By inserting δ/d in Theorem 4.8, absorbing d into the O(·) terms via (32), and applying a union bound, we have that ∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) +

X O(1) O(log n) √ ∥b̂j − b∗j ∥E n , log( 2δ ) + log 2δ λi λi n

(36)

2κ X O(log n) √ · log( 2δ ) + ∥b̂j − b∗j ∥E n λi λi n

(37)

j∈pa(i)

∥b̂i − b∗i ∥E n ≤ ∥bi,λi − b∗i ∥E n +

j∈pa(i)

for all i ∈ {1, . . . , d} with probability at least 1 − δ. We now recursively plug equation (37) into equation (36) and transform to something that’s amenable to taking expectation. Let (λi )i=1,...,d satisfy (15). Via induction over di ≥ 1, we show that log 2δ λi

X

∥b̂j − b∗j ∥E n ≤

j∈pa(i)

X O(log n) 2 2 1/3dmi log + O(1) (Tm (λm ) ∨ Tm (λm ))(Qm + log2 2δ ) δ n1/ri m∈an(i)

−b∗ ∥2

∥b

m,λm m En for all i ≤ d on this event with probability at least 1 − δ, where Qm := ∥bm,λ . First −b∗ ∥2 m

m L2 (p)

consider di = 1. Then, an(i) = pa(i), otherwise we had di ≥ 2. In particular, pa(j) = ∅ for all j ∈ an(i). Further, dmi = 1 for all m ∈ an(i). By (37) we have that log 2δ λi

X

∥b̂j − b∗j ∥E n ≤

j∈pa(i)

X ∥bj,λi − b∗j ∥L2 (p) X O(log n) p √ log2 ( 2δ ) + log( 2δ ) Qi . λi λi λj n

j∈pa(i)

j∈pa(i)

−2/3

For the second term, use λ−1 ≤ ∥bj,λj − b∗j ∥L2 (p) for all j ∈ pa(i) by construction, and apply i √ log( 2δ ) Qi ≤ log2 ( 2δ ) + Qi . Regarding the first term, consider that 1√ ≤ n2/ri · n2/rj · n−1/2 = n2/18 · n2/6 · n−1/2 = n−1/18 = n−1/ri , λi λj n

R. Schwank and M. Drton/Causal Diffusions

28

and absorb |pa(i)| · O(log n) = O(log n) to see the induction claim. Let i be a node such that the induction claim has been shown for di − 1. Note that dj ≤ di − 1 for all j ∈ an(i), otherwise there would exist a path of length di + 1 with i as endpoint via node j. By (37), log 2δ λi

X

∥b̂j − b∗j ∥E n ≤

j∈pa(i)

X ∥bj,λj − b∗j ∥L2 (p) X O(log n) p √ log2 ( 2δ ) + log( 2δ ) Qi λi λi λj n j∈pa(i)

j∈pa(i)

+

1 X 2κ log 2δ λi λj j∈pa(i)

X

∥b̂m − b∗m ∥E n . (38)

m∈pa(j)

By r1j = 6·31dj ≥ 6·3d1i −1 = r3i , we have 1 λi n1/rj

≤ n2/ri · n−3/ri = n−1/ri .

(39)

√ Since (λj n)−1 ≤ n2/rj · n−1/2 ≤ n−1/rj , the first term on the right hand side of (38) is n) bounded by O(log log2 2δ using (39). Treat the second term like in the induction start, using n1/ri λ−1 ≤ Tj (λj )−2/3 for j ∈ pa(i). Recalling the induction hypothesis, we find that the third term i on the right hand side of (38) is bounded by X O(log n) X log2 2δ + O(1) 1/r j λi n

j∈pa(i)

X

j∈pa(i) m∈an(j)

(39) O(log n)

≤

n1/ri

(Qm + log2 2δ ) 1/3dmj (λm ) ∨ Tm (λm )) (Tm λi 1/3dmj

log2 2δ + O(1)

X

Tm

X

j∈pa(i) m∈an(j)

(λm ) ∨ Tm (λm ) (Qm + log2 2δ ). 2/3dmi (λm ) Tm

If Tm (λm ) ≤ 1, use dmi ≥ dmj +1 for j ∈ pa(i), and therefore 3−dmj −2·3−dmi ≥ 13 3−dmj ≥ 3−dmi . −1/2dmi

Otherwise, bound Tm (λm ) ≤ 1. Absorbing |pa(j)| into O(1), we have shown the induction claim. Now, consider (36), and for di ≥ 1 plug in the induction claim to obtain: ∥b̂i − b∗i ∥L2 (p) ≤ ∥bi,λi − b∗i ∥L2 (p) + + O(1)

X

O(log n) √ log λi n

2 δ

+ 1di ≥1 ·

O(log n) log2 ( 2δ ) n1/ri

1/3dmi

∥bm,λm − b∗m ∥L2 (p) ∨ ∥bm,λm − b∗m ∥L2 (p) (Qm + log2 ( 2δ )).

m∈an(i)

for all i ≤√d with probability at least 1 − δ, where 1di ≥1 is 1 if di ≥ 1 and zero otherwise. Now, insert (λi n)−1 ≤ n−1/ri . Then, integrate the bound as demonstrated in Lemma D.5 and use E[ Qm ] = 1. Using 3dii = 30 = 1 yields the theorem claim. If Hi is dense in L2 (R|Pa(i)| , pPa(i) ) for all i ≤ d, it follows that inf b∈Hi ∥b − b∗i ∥2L2 (p) = 0. This implies limλ→0 Ti (λ) = 0 (Steinwart and Christmann, 2008, λi to the right h Sct. 5.4). Setting i

hand side of (15) yields limn→∞ λi (n) = 0 and limn→∞ E ∥b̂i − b∗i ∥L2 (p) = 0 by (16). √ 1 Let Tj (λ) = ∥bj,λ − b∗j ∥L2 (p) = O( λ · logαj λ1 ) as λ → 0 for all j ≤ d and let 0 < ε < 2d . Set −2(1−εj )/rj λj (n) := n for all j ≤ d. Let i ≤ d. For any j ∈ an(i) we have by λj → 0 that 1/3dji

dji

∥bj,λj − b∗j ∥L2 (p) = O(n−(1−εj )/(rj 3

)

dji

· logαj /3

(n2(1−εj )/rj )) = O(n−(1−εi )/ri ),

since rj 3dji ≤ ri and n(εj −εi )/ri decays faster to zero than the log-term as n → ∞. In particular, (λi )i≤d satisfies (15) for n large enough. The leading term in equation (16) is Ti (λi ) = O(n−(1−εi )/ri · logαi n2(1−εi )/ri ) = O(n−(1−2εi )/ri ), which proves the claim.

R. Schwank and M. Drton/Causal Diffusions

B.3

29

Theory discussion details

First, we demonstrate that the stationary distribution and parent set alone do not identify b∗3 in a three variable example. Consider the following matrices       0.5 0.125 0.1563 −1 0 0 −1.1111 0.44444 0 0.5625 0.21094  B1 := 0.5 −1 0  B2 :=  0 −0.88889 0  . Σ :=  0.125 0.1563 0.21094 0.683594 0.5 0.5 −1 0.3125 0.63889 −1 Then, both drift functions B1 x and B2 x induce a Gaussian with mean zero and covariance Σ as stationary distribution when the constant identity matrix is used as diffusivity, as seen by the continuous Lyapunov equation. Note that the third drift component, i.e., the third row, differs even though node 3 has the same parent set in both drifts functions. In other words, the stationary distribution and the parent set alone do not identify the third drift component here. It also depends on the parent drifts, which differ in both cases since 1 is a parent of 2 in B1 and 2 a parent of 1 in B2 . For the second part of the discussion, consider the setup from example 3.2. We learn the drift b∗2 (x1 , x2 ) using a shift-invariant product kernel k2 ((x1 , x2 ), (y1 , y2 )) = k(x1 − y1 ) · k(x2 − y2 ) for some univariate function k.R We assume k(−x) = k(x) and k ′ (−x) = −k ′ (x), and that k is x2 integrable. We use K2,y (x) := −∞ k2 ((x1 , z), y)dz. Under sufficient tail decay of k, Z x2 ∂1 K2,y (x) =

∂x1 k2 ((x1 , z), y)dz = k ′ (x1 − y1 )

−∞

Z x2 k(z − y2 )dz = −∞

− k ′ (y1 − x1 )

Z x2 −y2

Z ∞ k(z)dz = −∂y1 k(x1 − y1 )

−∞

k(x2 − z)dz. y2

Let ε(x1 ) := b∗1 (x1 ) − b(x1 ) for some proposal drift b. Then, Z   1 ∂1 K2,y (x)ε(x1 )dp(x) = ζ2 (y) − E ∂1 K2,y · b − 2 ∆K2,y = 2 Z ∞ Z R Z ∞ −∂y1 k2 (x, (y1 , z))ε(x1 )dp(x)dz = − ∂y1 (J2∗ ε)(y1 , z) dz. y2

R2

y2

In other words, using the incorrect root drift b for learning ζ2 leads to a bias, which in the infinite sample limit can be expressed in terms of J2∗ ε. The bias in ζ2 induces a bias on the second drift component via L2 (p) ∋ b2,λ2 = J2 (J2∗ J2 + λ2 I)−1 ζ2 . Therefore, the final bias induced by b ̸= b∗1 heuristically is:  Z  ∞

J2 (J2∗ J2 + λ2 I)−1 −

∂y1 (J2∗ ε)(y1 , z) dz .

y2

Since ∥J2 (J2∗ J2 + λI)−1 J2∗ ∥L(L2 (p)) ≤ 1, the bias could be up to order ∥ε∥L2 (p) = ∥b∗1 − b∥L2 (p) . This would improve the current bound λ−1 · ∥b∗1 − b∥L2 (p) , which is due to ∥(J2∗ J2 + R∞ λI)−1 ∥L(H2 ,H2 ) ≤ λ−1 . However, it is difficult to bound the effect of y2 ∂y1 · dz on J2∗ ε, in R∞ particular since y2 ·dz does not map into H due to boundary value mismatch, making the above bias heuristic actually ill-defined. Future research could derive alternative estimators of ζ2 beyond sample averages of K2,y (xj ), which may be better behaved analytically. B.4

Cross-validation

The following describes how the cross-validation algorithm approximates ⟨b̂i,λ;I , ζ̂i;J ⟩Hi . Let Ki = (ki (xj , xl ))j,l∈J be the kernel matrix on the test data and let γ be an eigenvalue

R. Schwank and M. Drton/Causal Diffusions

30

of |J1 | K with eigenvector u. Then, γ is an eigenvalue of L̂i,J with eigenfunction v(y) := P √1 j∈J uj ki (xj , y) (Rosasco et al., 2010). Let γmin > 0 and γ1 , . . . , γs be all eigenvalues of |J |·γ 1 |J | K with magnitude greater than γmin . The projection onto the corresponding eigenfunctions

(vj )j=1,...,s can be written as argminα∈Rs ∥

s X

αj vj − ζ̂i;J ∥2H = argminα∈Rs

j=1

−2

s X

X jklm

αj αk uj,l uk,m ki (xl , xm ) √ |J | γj γk

1 2 αj (UΓα)T h, uj,k ζ̂i;J (xk ) = argminα∈Rs (UΓα)T KUΓα − p |J | |J |γj k∈J |J | X

p j=1

√ where U ∈ R|J |×s has the u-vectors as columns, where we defined Γ := diag( γ −1 ) ∈ Rs×s and hk := ζ̂i;J (xk ), and where we ignored the additive constant ∥ζ̂i;J ∥2H in the optimization. Using |J1 | ΓUT KUΓ = I, the minimizer is given by 1 ΓUT h. α∗ = p |J | For an RKHS function f , we therefore approximate ⟨f, ζ̂i;J ⟩Hi by √1 (UΓα∗ )T f , where fj := |J |

f (xj ) for j ∈ J . All in all, our crossvalidation method optimizes 1 2 ∥bi ∥22 + p (UΓα∗ )T bi , |J | |J | where bi = (b̂i,λ;I (xj ))j∈J . Appendix C C.1

Simulation and extension details

KDS for DAG-structured drifts

The kernel deviation from stationarity (KDS) by Lorch et al. (2024) is an alternative approach to fit nonparametric drift functions to samples from the stationary distribution of an SDE. We extend the accompanying Python package to DAG-structured drifts by simply adding input masks in the existing Neural network class MLPSDE. Note that the current KDS package version assumes bi (x) = −xi + fΘ (x−i ) for a two-layer neural network f with trainable weights Θ. We tried activating some low-level functions in the package which generalize to arbitrary self-regulation, but obtained NaNs in the KDS training, even when using the MLPSDE class. Therefore, we train on data from a Gaussian distribution with drift b1 (x1 ) = −x1 , b2 (x1 , x2 ) = x1 − x2 , and b3 (x1 , x2 , x3 ) = − 21 x1 + 12 x2 − x3 matching the package assumptions. We train on n = 1000 samples using the default package hyperparameters and record the generalization error for varying numbers of gradient steps. After 40000 steps, the average generalization error noticeably increased and we did not observe double descent even when training 10 times longer. The optimal recorded generalization error averaged 0.07 ± 0.006 on N = 50 experiments. This confirms that KDS approximates the true drift when oracle-stopped. We are not aware of package functionality for adaptive early stopping; reducing the learning rate significantly could lead to convergence at the expense of longer training. For comparison, our Tikhonov regularized estimator achieves an average generalization error of 0.2 ± 0.01 on this problem while adaptively choosing the regularization strength. Additionally it learns the self-regulation −xi ; in particular it must learn b1 and also fit a three-variable function b3 instead of a two-variable function like KDS.

R. Schwank and M. Drton/Causal Diffusions

C.2

31

Irreversible diffusion models

We consider a 5-dimensional Gaussian density p with mean zero and score function s(x) rounded to three digits shown below. Also consider the alternative drift function b(x) defined by     −1.185 0.275 −0.032 0.147 0.282 −1 0 0 0 0  0.275 −1.016 0.241 −0.201 0.015   0.5 −1 0 0 0      x; s(x) = −0.032 0.241 −1.210 0.348 −0.243 x. −0.1 0.3 −1 0 0 b(x) =       0.147 −0.201 0.348 −0.756 −0.007  0.4 −0.6 1 −1 0  0.282 0.015 −0.243 −0.007 −0.832 0.7 −0.1 −0.6 0.1 −1 Both matrices are linked via the Lyapunov equation such that for all ω ∈ R the SDE √ dx(t) = (s − ω · (s − b))(x(t))dt + 2dw(t) has the density p as its stationary density. One also could obtain b via Theorem 5.1 from p directly, however the Lyapunov method is simpler to compute and must yield the same result by our identifiability theory. (ω) For ω ∈ {0, 1, 5, 10, −20} we simulate i = 1, . . . , 100 trajectories xi (t) from t = 0 to T = 2 using the Euler-Maruyama method with 5000 steps. The initialization points are drawn from a Gaussian distribution with mean vector (5, 5, 5, 5, 5) and identity covariance and the initialization is shared across different ω. At each time t we compute the energy distance between (ω) (ω) x1 (t), . . . , x100 (t) and an independent sample of size 1000 from p. We then plot this distance depending on t and ω. Theorem 5.1. Let 0 < p ∈ C 2 (Rd ) be a probability density and σ be constant. For any j ≤ d, define 1 : j := 1, . . . , j, let fj+1 and Fj+1 denote the density and cumulative distribution function of xj+1 given x1:j , respectively, and assume Fj+1 ∈ C 3 (Rj ). 2 p′ (x ) Define b : Rd → Rd component-wise as b1 (x) := σ2 p11 (x11 ) and bj+1 recursively as bj+1 :=

σ 2 ∂j+1 fj+1 + ∆1:j Fj+1 + 2(∇1:j log p1:j )T ∇1:j Fj+1 (∇1:j Fj+1 )T b1:j + . 2 fj+1 fj+1

(20)

Then, b is structured according to the complete DAG with topological order 1, . . . , d (including 2 self loops), and satisfies the Fokker-Planck equation σ2 ∆p = ∇ · (bp). Proof. For reference, note the following identities for j ≥ 1: fj+1 (x) =

p1:(j+1) (x1:(j+1) ) p1:j (x1:j )

Z xj+1 Fj+1 (x) =

fj+1 (x1:j , z) dz. −∞

We show via induction over j that bj only depends on x1:j (thus proving drift structuring according to the complete DAG) and that σ2 ∆1:j p1:j = ∇1:j · (b1:j p1:j ), 2 which proves the theorem claim when taking j = d. The case j = 1 holds by definition. We next show that if the induction claim holds for j, it also holds for j +1. From the above identities, we see that fj+1 and Fj+1 only depend on x1:(j+1) . Since b1:j only depends on x1:j by induction hypothesis, it follows that bj+1 only depends on

R. Schwank and M. Drton/Causal Diffusions

32

x1:(j+1) , as desired. To see that the PDE holds, multiply the constructed bj+1 with p1:(j+1) , obtaining p1:(j+1) bj+1 =

σ2 σ2 p1:j ∂j+1 fj+1 + p1:j ∆1:j Fj+1 + 2 2 σ2 · 2 · (∇1:j p1:j )T · ∇1:j Fj+1 − (p1:j b1:j )T · ∇1:j Fj+1 . 2 2

The first two terms can be simplified to σ2 p1:j ∆1:(j+1) Fj+1 by definition of Fj+1 . We now add zero to the right hand side, more specifically   2 Fj+1 · σ2 ∆1:j p1:j − ∇1:j · (p1:j b1:j ) , 2

which is zero by the induction hypothesis. All terms involving σ2 read σ2 σ2 σ2 p1:j ∆1:(j+1) Fj+1 + · 2 · (∇1:j p1:j )T · ∇1:j Fj+1 + Fj+1 · ∆1:j p1:j . 2 2 2 Since ∂j+1 p1:j = 0, we have ∆1:(j+1) p1:j = ∆1:j p1:j and also (∇1:j p1:j )T · ∇1:j Fj+1 = (∇1:(j+1) p1:j )T · ∇1:(j+1) Fj+1 . Recalling the Laplacian product rule ∆(f g) = f ∆g + 2∇f · ∇g + g∆f , we can simplify all terms 2 involving σ2 to σ2 ∆1:(j+1) (p1:j Fj+1 ). 2 2

The remaining terms not involving σ2 read −(p1:j b1:j )T · ∇1:j Fj+1 − Fj+1 · ∇1:j · (p1:j b1:j ) = −∇1:j · (b1:j p1:j Fj+1 ) by the divergence product rule. Collecting the results, we have shown p1:(j+1) bj+1 =

σ2 ∆1:(j+1) (p1:j Fj+1 ) − ∇1:j · (b1:j p1:j Fj+1 ). 2

Taking ∂j+1 on both sides, changing the derivative order by Schwarz’s theorem, and using ∂j+1 Fj+1 = fj+1 as well as that neither p1:j nor b1:j depend on xj+1 , we obtain ∂j+1 (p1:(j+1) bj+1 ) =

σ2 ∆1:(j+1) (p1:j fj+1 ) − ∇1:j · (b1:j p1:j fj+1 ). 2

Simplifying p1:j fj+1 = p1:(j+1) and rearranging yields the induction claim. Appendix D

Technical Lemmas

Definition D.1 (Standard Mollifier). Define η : Rd → R by   ( C exp ∥x∥12 −1 ∥x∥2 ≤ 1 2 η(x) := 0 ∥x∥2 > 1 R R with C such that Rd η(x)dx = 1. Finally, let ηε (x) := ε1d η( xε ) for ε > 0 and note Rd ηε (x)dx = 1.

R. Schwank and M. Drton/Causal Diffusions

33

R Lemma D.2. Let u ∈ L1 (Rk ), k ≥ 1, with Rk u ∂l φ dx = 0 for all φ ∈ C0∞ (Rk ) and one index l ∈ {1, . . . , k}. Then u = 0 almost everywhere on Rk . Proof. The assumptions imply that u has a weak derivative in xl direction, given by zero almost everywhere. We write ∂l u = 0. First, we study the case k = 1. Then, u ∈ W 1,1 (R) with u′ = 0. In this case, u is a.e. equal to an absolutely continuous function with (ordinary) derivative u′ = 0 on every finite interval (a, b); see problem 5.10.5 in (Evans, 2010). It follows that u is zero globally by u ∈ L1 (R). Next, consider k ≥ 2. By reordering input arguments, we can assume l = k without loss of generality. For ε > 0, set uε := u ∗ ηε , the convolution with the standard mollifier ηε from Definition D.1. Then, by the proof of theorem 1 in section 5.3.1 of (Evans, 2010), • uε ∈ C ∞ (Rk ) with • ∂k uε = ηε ∗ ∂k u = ηε ∗ 0 = 0, and • uε → u in L1 (V ) as ε → 0 for every V ⊂⊂ Rk . The first two bullet points imply that uε (x) = ũε (x−k ) for functions ũε ∈ C ∞ (Rk−1 ). Let B ⊂ Rk−1 be an Euclidean ball of arbitrary (finite) radius around the origin. For any ε > 0 and n ∈ N, by the triangle inequality ∥uε ∥L1 (B×[−n,n]) ≤ ∥uε − u∥L1 (B×[−n,n]) + ∥u∥L1 (Rk ) . Construct a sequence ε(n), such that ε(n) ≤ n1 and ∥uε − u∥L1 (B×[−n,n]) ≤ 1 for all ε ≤ ε(n). Noting that ∥uε ∥L1 (B×[−n,n]) = 2n∥ũε ∥L1 (B) , this implies 2n∥ũε(n) ∥L1 (B) ≤ 1 + ∥u∥L1 (Rk ) ⇐⇒ ∥ũε(n) ∥L1 (B) ≤

1 + ∥u∥L1 (Rk ) . 2n

By assumption, ∥u∥L1 (Rk ) < ∞, so the right hand side converges to zero and consequently ũε(n) → 0 in L1 (B). For arbitrary a ∈ R, it holds that ∥uε(n) ∥L1 (B×[a,a+1]) = ∥ũε(n) ∥L1 (B) → 0, implying uε(n) → 0 in L1 (B × [a, a + 1]). Recalling that ε(n) → 0 by construction, the third bullet point implies u = 0 almost everywhere on B × [a, a + 1]. Covering Rk by varying B and a yields u = 0 almost everywhere on Rk . Lemma D.3. When the Lebesgue density p is bounded, then for any φ ∈ W 2,∞ (Rk ) there exists (φn )n∈N ⊂ C0∞ (Rk ) with ∥φn − φ∥W 2,2 (Rk ,p) → 0 as n → ∞. Proof. Let n ∈ N, set Br := {x ∈ Rk | ∥x∥2 < r} for r ≥ 0 and let φ ∈ W 2,∞ (Rk ). It holds that φ ∈ W 2,∞ (Bn+1 ) ⊂ W 2,2 (Bn+1 ). By theorem 1 in Section 5.3.1 of (Evans, 2010), there is ε(n) ≤ 0.5 such that ∥(φ ∗ ηε(n) ) − φ∥W 2,2 (Bn ) ≤ n1 , where ∗ denotes convolution and ηε is the standard mollifier from Definition D.1. Further, for any multi-index α with |α| ≤ 2 we have Dα (φ ∗ ηε(n) ) = ηε(n) ∗ (Dα φ) on Bn+0.5 . Define φn := γn · (φ ∗ ηε(n) ) with γn from Lemma D.4. Recall γn ∈ C0∞ (Rk ) with support in Bn+0.25 . Since φ ∗ ηε(n) ∈ C ∞ (Bn+0.5 ), it follows φn ∈ C0∞ (Rk ). Further, for any multi-index α with |α| ≤ 2, by Leibnitz’ formula X α  α D φn = Dβ γn · (ηε(n) ∗ Dα−β φ). β β≤α

R. Schwank and M. Drton/Causal Diffusions

34

R First, notice |(ηε(n) ∗ Dα−β φ)(x)| ≤ ∥Dα−β φ∥∞ · Rd ηε(n) (y)dy = ∥Dα−β φ∥∞ for all x ∈ Rk . Second, ∥Dβ γn ∥L∞ (Rk ) ≤ supn ∥γn ∥W 2,∞ (Rk ) < ∞ by Lemma D.4. Combining these facts yields supn ∥φn ∥W 2,∞ (Rd ) < ∞. Thus, ∥φn − φ∥2W 2,2 (Rd ,p) = ∥φn − φ∥2W 2,2 (Bn , p) + ∥φn − φ∥2W 2,2 (Rd \Bn , p) ≤  Z ∥φn − φ∥2W 2,2 (Bn ) ∥p∥L∞ (Rk ) + 2 ∥φ∥2W 2,∞ (Rd ) + sup ∥φn ∥2W 2,∞ (Rd ) n

p(x)dx,

Rd \Bn

which converges to 0 when n → ∞ as desired.

Lemma D.4. There is a sequence (γn )n∈N ⊂ C0∞ (Rk ) satisfying γn (x) = 1 for ∥x∥2 ≤ n, γn (x) = 0 for ∥x∥2 ≥ n + 41 , and supn∈N ∥γn ∥W 2,∞ (Rk ) < ∞. Proof. Choose γ ∈ C ∞ (R) with γ = 1 for x ≤ 0 and γ(x) = 0 for x ≥ 41 . Define γn (x) := γ(∥x∥2 − n). By construction, γn (x) = 1 for ∥x∥2 ≤ n and γn (x) = 0 for ∥x∥2 ≥ n + 41 . Since ∥·∥ ∈ C ∞ (Rd \{0}), it holds that γn ∈ C0∞ (Rk ). Further, ∥γn ∥L∞ (Rk ) ≤ ∥γ∥L∞ ([0,1/4]) < ∞. For any i ≤ k, we have ∥∂i γn (x)∥L∞ (Rk ) = sup |γ ′ (∥x∥2 − n) · (∂i ∥ · ∥2 )(x)| ≤ ∥γ ′ ∥L∞ ([0,1/4]) · ∂i ∥ · ∥2 L∞ (∥x∥≥1) . x∈Rk

For any i, j ≤ k we have ∥∂ij γn (x)∥L∞ (Rk ) = sup |γ ′′ (∥x∥2 − n) · (∂i ∥ · ∥2 )(x) · (∂j ∥ · ∥2 )(x) + γ ′ (∥x∥2 − n) · (∂ij ∥ · ∥2 )(x)| x∈Rk ′′

≤ ∥γ ∥L∞ ([0,1/4]) · ∂i ∥·∥2 L∞ (∥x∥≥1) · ∂j ∥·∥2 L∞ (∥x∥≥1) +∥γ ′ ∥L∞ ([0,1/4]) · ∂ij ∥·∥2 L∞ (∥x∥≥1) . From ∥∂xβ ∥ · ∥2 L∞ (∥x∥≥1) < ∞ for any multi-index β with 1 ≤ |β| ≤ 2 the claim follows. Lemma D.5. Let A, B be nonnegative random variables and α, β > 0 be constants such that A ≤ B + β logα 2δ with probability at least 1 − δ for all δ ∈ (0, 1). Then, E[ A ] ≤ 2 E[ B ] + β · c(α) for some constant c only depending on α. Proof. Let t > 0 be arbitrary. We have that

P( A > t + β logα 2 ) = P( B + (A − B) − β logα 2 > t ) ≤ P B > 2t +P A > B + 2t + β logα 2 . t Further, letting δ := 2 exp(−( 2β + logα 2)1/α ) ∈ (0, 1), it holds that

t P A > B + 2t + β logα 2 ≤ 2 exp(−( 2β + logα 2)1/α ). We have that Z ∞ Z ∞ E[ A ] = P( A > t )dt ≤ 1 · β logα 2 + P( A > t + β logα 2 )dt 0 0 Z ∞ Z ∞
α t t ≤ β log 2 + P B > 2 dt + 2 exp(−( 2β + logα 2)1/α )dt. 0

0

t Transforming 2t 7→ u and 2β 7→ u in the integrals proves the theorem claim.