Identifying Drift, Diffusion, and Causal Structure from Temporal Snapshots

In plain words

Many real systems drift and jitter at the same time: cells change type, and pollutants spread through groundwater. Often you cannot follow one cell or one particle over time. You only get snapshots, a crowd of measurements at a few moments, with no way to link a measurement to the one before it.

The paper asks whether the snapshots alone can reveal the rules behind the motion: how strongly each variable pushes the others (the drift), how noisy each direction is (the diffusion), and which variables influence which. It finds that the answer is usually yes. The rules can only be confused when the starting crowd is shaped in a special symmetric way, and even then such a confusion is rare among possible rule sets.

The method, APPEX, works in two alternating steps. It first guesses which measurement at one time could have come from which measurement at the next time. Then it fits the rules that best explain those guessed paths. Repeating both steps on simulated systems gave more accurate rules than the common single-cell baseline. Its results are on simulated data and are reported by the authors.

TL;DR

Linear additive-noise SDEs are identifiable from temporal marginals whenever the initial distribution is not weakly auto-rotationally invariant. If it is, identifiability still holds for almost every parameter choice §3.2. APPEX estimates drift and diffusion by alternating anisotropic entropic optimal transport (Sinkhorn) with closed-form maximum likelihood §5.3. On simulated systems with d = 3 to 10, it beat Waddington-OT on drift and diffusion error and correlation tab. 1, tab. 2. The official code is MIT-licensed, but several settings in the paper differ from the code (see Paper vs code).

Problem

Trajectories are often unobservable: destructive sequencing and sensors that see only plume migration give only marginal snapshots §1. Such marginals often do not identify the SDE. Two classical sources are a shared stationary law and rotational drift §2. Earlier methods avoid this by assuming an irrotational drift and isotropic diffusion with known scale §1, §2. Those restrictions exclude feedback loops and repressilator-type dynamics, and they exclude variable or correlated noise §2. The paper also notes that a wrong diffusion prior tends to distort the drift estimate and the reverse §2.

Method

The paper works with the time-homogeneous linear additive-noise model dX_t = A X_t dt + G dW_t, with observational diffusion H = G Gᵀ §2. Only H is identifiable, not G §2.

Identifiability (§3).

1. Weak auto-rotational invariance (weak ARI) is defined through the cumulants of order n ≠ 2 of the initial law, which must be invariant under a common matrix B §3.2, Def. weak ARI.
2. Theorem (linear additive noise): every (A, H) produces a unique family of marginals if and only if X_0 is not weakly ARI §3.2.
3. Theorem (generic identifiability): for almost every (A, H), the covariance curve of the process identifies (A, H). This holds even when X_0 is weakly ARI. The exceptional set has Lebesgue measure zero §3.2.
4. Under isotropic diffusion, non-identifiability forces every marginal to be invariant under the same generalized rotation e^{Bθ}. Observing marginals that change over time therefore pins down the diffusivity §3.1.
5. Observing marginals at equilibrium is a separate source of non-identifiability; for linear SDEs, rotation is the other one §3.1.

Causal graph (§4). 6. A directed edge i → j exists if and only if drift component b_j depends on X^(i). For linear drift this means A_{j,i} ≠ 0 §4, lemma. 7. A pairwise latent confounder is present when H_{i,j} ≠ 0, because the two variables share a Brownian component §4, lemma. The paper notes that not all noise relations are recoverable from H §4.

APPEX (§5, alg. 1). Given marginals p_0, ..., p_{N-1} and initial guesses (A^(0), H^(0)), each iteration does the following §5.3. 8. Trajectory inference: for each consecutive pair of marginals, solve an anisotropic entropic OT problem whose transition kernel comes from the current reference SDE §5.1. 9. Sample paths by chaining the N-1 coupling plans from p_0 §5.1, alg. 1. 10. Fit the drift by closed-form MLE on the sampled paths, then fit the diffusion by closed-form MLE using the new drift §5.2. 11. After the last iteration, threshold the drift and diffusion entries to get the causal graph §5.3, remark on causal discovery.

The paper proves that, in the population limit, the KL divergence between the estimated path law and the reference SDE does not increase from one iteration to the next §5.3, lemma. It does not prove convergence of the nonconvex iteration itself §7.1.

# APPEX  [alg. 1]
# Input:  marginals p_0..p_{N-1}, each (M, d); dt; num_iters
# Output: drift (d, d), diffusion H (d, d)

def appex(marginals, dt, num_iters):
    # Start from zero drift and an isotropic reference diffusion  [§6]
    drift = zeros((d, d))
    diffusion = sigma_sq * identity(d)
    for k in range(num_iters):
        # Trajectory inference: one AEOT plan per consecutive pair  [§5.1]
        plans = [
            aeot(marginals[i], marginals[i + 1], drift, diffusion, dt)
            for i in range(N - 1)
        ]
        # Sample paths by chaining the plans from p_0  [§5.1; alg. 1]
        paths = sample_paths(marginals[0], plans, num_paths)
        # Drift first: the diffusion MLE uses the new drift  [§5.2]
        drift = mle_drift(paths, dt)
        diffusion = mle_diffusion(paths, drift, total_time)
    return drift, diffusion

Trajectory inference (§5.1, app. E). The reference kernel for each pair is a Gaussian whose mean and covariance come from the current (A, H). With a linearised mean, the covariance is H times dt. Sinkhorn then gives the plan with uniform marginals on both sides §5.1, §6.

# Anisotropic entropic OT for one pair of marginals  [§5.1]
# source, target: (M, d) samples; drift: (d, d); diffusion: (d, d)

def aeot(source, target, drift, diffusion, dt):
    # Mean of each source point after one linearised step  [§6]
    drifted = source + dt * source @ drift.T     # (M, d)
    # Residual of every target point against every drifted source  [§6]
    residuals = target[None] - drifted[:, None]  # (M, M, d)
    # Gaussian kernel with covariance (H + jitter) * dt  [§6]
    kernel = gaussian_pdf(residuals, cov=(diffusion + jitter) * dt)
    # Sinkhorn with uniform marginals on both sides  [§5.1]
    return sinkhorn(kernel, uniform(M), uniform(M))

MLE (§5.2, app. F). With paths on the N measurement times, the drift and diffusion estimates are closed-form. The displayed sums below are the finite-sample form given in the paper.

# Closed-form MLE from sampled paths  [§5.2]
# paths: (P, N, d); total_time = T; dt = T / (N - 1)

def mle_drift(paths, dt):
    # Increments and starting states over the N-1 steps  [§5.2]
    states = paths[:, :-1]                       # (P, N-1, d)
    increments = paths[:, 1:] - paths[:, :-1]    # (P, N-1, d)
    # Cross and Gram sums, pooled over paths and time  [§5.2]
    cross = einsum("pti,ptj->ij", increments, states)
    gram = einsum("pti,ptj->ij", states, states)
    # Solve for drift = cross @ inv(gram) / dt  [§5.2]
    return solve(gram.T, cross.T).T / dt


def mle_diffusion(paths, drift, total_time):
    # Increments minus the deterministic drift part  [§5.2]
    dt = total_time / (paths.shape[1] - 1)
    states = paths[:, :-1]
    residuals = (paths[:, 1:] - paths[:, :-1]) - dt * states @ drift.T
    # Average of residual outer products over paths, divided by T  [§5.2]
    sums = einsum("pti,ptj->ij", residuals, residuals)
    return sums / (paths.shape[0] * total_time)

Two details are easy to miss. The finite-sample drift formula sums over N-2 steps in both numerator and denominator; the population limit in the paper sums the denominator over N-1 marginals §5.2. The code follows the finite-sample form APPEX.py. Second, the trajectories are sampled from the coupling plans, so the MLE is run on a sampled set of paths rather than on the full law §5.1.

Hyperparameters

Values for the simulated experiments in §6 unless marked. "Code" means the official repository at the commit in the front matter, and the file is relative to its root.

Results

Author-reported, on simulated linear additive-noise SDEs. Setup: M = 500 trajectories, N = 20 marginals, Euler–Maruyama dt_EM = 0.01, and 30 APPEX iterations §6. Baseline (WOT) is the first APPEX iteration, which the paper states is equivalent to MLE after Waddington-OT trajectory inference §6. Hardware and runtime are not stated.

Revisited non-identifiable pairs (§6.1, fig. 4). Three classical non-identifiable SDE pairs were redone with an X_0 that is not weakly ARI (10 replicates each). Mean absolute error for A and H falls with the iteration in all cases. The first iteration (WOT-like) is the worst. The hardest case is the high-noise pair with drift −10 and noise √10 fig. 4. Values are read from the figure and not given in the text.

Random systems (§6.2, tab. 1 and tab. 2). Mean ± standard error over 10 systems per dimension.

Table: Mean absolute error (lower is better) tab. 1

Table: Correlation with the true parameters (higher is better) tab. 2

The paper also shows per-iteration correlation curves for d = 3 and d = 10 fig. 5. It reports occasional dips in accuracy, which it attributes to finite data, a Sinkhorn iteration cap, and building on earlier wrong estimates §6.2.

Causal graph (§6.3, tab. 3 and tab. 4). Structural Hamming distance over random graphs. Lower is better. Values are mean ± standard error.

Table: directed-edge SHD, no latent confounders tab. 3

Table: with correlated diffusion, p = 0.25 tab. 4

Example graphs for four systems are in fig. 6. Consistency as the number of samples per marginal grows is checked in app. G; that section's results are not summarised here.

Paper vs code

Repository: https://github.com/guanton/APPEX at commit 685c0d90c16a75708e69e439295e6bbf41147e87, MIT licence naming all seven authors. The paper links it in its Code Availability section. The R repository linked there (HydroML/X0isAllYouNeed) was not read.

• Initial diffusion. Paper: H^(0) = σ² I with σ² ∼ tr(H) · 10^Unif(−1,1) §6. Code: H^(0) = I, so σ² = 1, and the experiment driver passes no initial values APPEX.py experiments.py.
• Drift MLE denominator. Paper's population limit sums N-1 marginals in the denominator §5.2. Paper's finite-sample formula and the code sum only N-2 steps in both the numerator and the denominator APPEX.py.
• Kernel covariance. Paper's reference transition covariance is the integral of e^{As} H e^{Aᵀs} §5.1, app. E. Paper's experiments use a linearised kernel for d ≥ 2 §6. Code uses H dt with a 1e-8 jitter APPEX.py.
• Kernel sign. Paper's text writes the kernel with exp of a positive quadratic form §5.1. The code uses a Gaussian density, which has the negative sign APPEX.py.
• Number of marginals. Paper: N = 20 §6. Code: int(T/dt) + 1 = 21 time points, t = 0 to 1 data_generation.py experiments.py.
• Initial support scale. Paper: "each entry" magnitude uniform in [2, 10] §6. Code: the Euclidean norm of each point is uniform in [2, 10] utils.py.
• Rejection rule. Paper's Algorithm 1 has no rejection alg. 1. Code has an optional rule that rejects a candidate if NLL, W1, W2 or MMD rises by more than 0.01 APPEX.py APPEX.py. It is off in the experiments experiments.py experiments.py.
• Which iteration is reported. Paper does not say which iteration the tables use. Code takes WOT as iteration 0 and APPEX as the last iteration plot_experiment_results.py plot_experiment_results.py.
• Number of sampled paths. Paper does not state it. Code draws 1000 paths per iteration APPEX.py APPEX.py.
• Sinkhorn caps. Paper mentions capping Sinkhorn iterations §6.2 but gives no number. Code caps at 1000 (linear) and 500 (log) with relative tolerances 1e-2 and 1e-5 APPEX.py APPEX.py.
• Killed marginals. The paper marks the killed-marginal experiments as to-do §6. The code has a killed option, which is off by default experiments.py data_generation.py.
• Experiments run. Experiments 1 to 3 are defined, but the driver runs only experiment 3, version 1 experiments.py.
• Least squares. Paper's closed-form drift uses a pseudo-inverse form §5.2. Code solves the least-squares problem instead, which matches the inverse when the Gram matrix is invertible APPEX.py utils.py.

Pitfalls

• Identifiability depends on X_0. The guarantee holds only when the initial law is not weakly ARI. The paper says the default d-point X_0 meets this §6, but this card did not re-check that proof. Verify it for any new initial distribution.
• Continuous observation is assumed. The identifiability theorems assume marginals at all times. The paper shows that sampling at discrete times can alias rotations §3.2. The experiments use dt = 0.05 and rely on the covariance evolution instead.
• Generic is not uniform. "Almost every parameter" means outside a measure-zero set. The paper gives no finite-sample bound on how often estimates fail §3.2, §7.1.
• Only H is identifiable. G, and so the sign and path of each latent source, cannot be recovered §2, §4. Pairwise confounders are detected from nonzero H entries. Shared sources among three or more variables are not resolved by this lemma §4.
• Thresholds are tuned to the simulator. The causal edge threshold 0.5 and confounder threshold 1 match the simulated minimum weights §6.3. Real data needs a separate choice.
• Cost. Each consecutive pair needs an M×M kernel and a Sinkhorn solve, repeated for every iteration. The kernel is built in a Python loop over trajectories APPEX.py. The paper calls this computationally intensive and not suited to very large datasets §7.1.
• Linear-domain underflow. Gaussian kernel values can underflow in higher dimensions. The code switches to the log domain only for the random experiments experiments.py. Applying the linear version to larger d is risky.
• Loose linear Sinkhorn stop. The linear version stops when the relative change in u is below 1e-2 APPEX.py. That is loose, so the plan may not be fully converged.
• Exact d = 1 path. The exact MLE takes the log of a ratio of sums, so a non-positive ratio gives NaN. The code does not check for this APPEX.py.
• Sampling noise. Each iteration resamples paths, and the code sets no seed inside APPEX. The seed is set only in the replicate driver experiments.py. Repeat runs will differ.
• Iterations can destabilise. The paper reports occasional dips in accuracy §6.2. The default code has no safeguard against them.
• KL monotonicity is a population result. The decreasing-KL statement §5.3 holds for the exact population quantities. The code does not compute KL. It computes NLL, W1, W2 and MMD proxies only when check_convergence is set experiments.py.
• Rate of convergence. Asymptotic rates and sample complexity are left open §7.1.

Lineage

• Builds on: Sinkhorn-based trajectory inference with a Schrödinger bridge formulation §2, §5.1. The paper cites Chizat et al. (2022) and Shen et al. (2024) for this; neither has a card here.
• Compared against: Waddington-OT (Schiebinger et al., 2019), which assumes isotropic diffusion with a known scale. The paper's baseline is its first iteration §6. No card here.
• Related identifiability work: Wang et al. (2024), Guan et al. (2025), Terpin et al. (2024) and Brogat et al. (2024) are cited in the paper for related identifiability or estimation results §2, §7. None has a card here.
• Superseded by: not stated.