Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Learning of Stationary Diffusions with Stein-type Discrepancies

Fabian Bleile, Sarah Lumpp, Mathias Drton

TL;DR

This work introduces the Stein-type Kernel Deviation from Stationarity (SKDS) to efficiently learn stationary diffusions whose stationary distribution matches a target $\mu$, establishing that SKDS vanishes exactly for reversible diffusions with density $\mu$ and providing convexity properties for linear parametrizations. Building on kernel methods and Stein discrepancies, SKDS offers a closed-form, computationally cheaper alternative to Kernel Deviation from Stationarity (KDS) while retaining theoretical guarantees. The authors develop a linear SDE parametrization that preserves linearity in the learning objective, derive an empirical SKDS estimator with provable quasiconvexity, and demonstrate substantial speedups and competitive accuracy across synthetic causal settings and interventions. Overall, SKDS enables scalable, principled causal learning with cyclic dependencies by leveraging stationary diffusion models and RKHS-based discrepancies, with practical impact for interventional prediction and inference.

Abstract

Learning a stationary diffusion amounts to estimating the parameters of a stochastic differential equation whose stationary distribution matches a target distribution. We build on the recently introduced kernel deviation from stationarity (KDS), which enforces stationarity by evaluating expectations of the diffusion's generator in a reproducing kernel Hilbert space. Leveraging the connection between KDS and Stein discrepancies, we introduce the Stein-type KDS (SKDS) as an alternative formulation. We prove that a vanishing SKDS guarantees alignment of the learned diffusion's stationary distribution with the target. Furthermore, under broad parametrizations, SKDS is convex with an empirical version that is $ε$-quasiconvex with high probability. Empirically, learning with SKDS attains comparable accuracy to KDS while substantially reducing computational cost and yields improvements over the majority of competitive baselines.

Efficient Learning of Stationary Diffusions with Stein-type Discrepancies

TL;DR

This work introduces the Stein-type Kernel Deviation from Stationarity (SKDS) to efficiently learn stationary diffusions whose stationary distribution matches a target , establishing that SKDS vanishes exactly for reversible diffusions with density and providing convexity properties for linear parametrizations. Building on kernel methods and Stein discrepancies, SKDS offers a closed-form, computationally cheaper alternative to Kernel Deviation from Stationarity (KDS) while retaining theoretical guarantees. The authors develop a linear SDE parametrization that preserves linearity in the learning objective, derive an empirical SKDS estimator with provable quasiconvexity, and demonstrate substantial speedups and competitive accuracy across synthetic causal settings and interventions. Overall, SKDS enables scalable, principled causal learning with cyclic dependencies by leveraging stationary diffusion models and RKHS-based discrepancies, with practical impact for interventional prediction and inference.

Abstract

Learning a stationary diffusion amounts to estimating the parameters of a stochastic differential equation whose stationary distribution matches a target distribution. We build on the recently introduced kernel deviation from stationarity (KDS), which enforces stationarity by evaluating expectations of the diffusion's generator in a reproducing kernel Hilbert space. Leveraging the connection between KDS and Stein discrepancies, we introduce the Stein-type KDS (SKDS) as an alternative formulation. We prove that a vanishing SKDS guarantees alignment of the learned diffusion's stationary distribution with the target. Furthermore, under broad parametrizations, SKDS is convex with an empirical version that is -quasiconvex with high probability. Empirically, learning with SKDS attains comparable accuracy to KDS while substantially reducing computational cost and yields improvements over the majority of competitive baselines.
Paper Structure (43 sections, 11 theorems, 117 equations, 4 figures, 4 tables, 1 algorithm)

This paper contains 43 sections, 11 theorems, 117 equations, 4 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Kernel Theory.
  4. Characterization of Stationarity.
  5. Identifiability.
  6. Kernel Deviation from Stationarity (KDS; Lorch:2024).
  7. Stein Discrepancy.
  8. Stein-Type Kernel Deviation from Stationarity
  9. Closed Form Representation
  10. Properties
  11. Learning Stationary Diffusions
  12. Parametrization
  13. Stein-Type KDS as a Learning Objective
  14. Comparison with KDS
  15. Synthetic Experiments
  16. ...and 28 more sections

Key Result

Lemma 4

Let $\mu$ be a probability density. Let $b$ and $\sigma$ be $\mu$-square-integrable, i.e., in $L^2_\mu$. Let $k(x,x)$ be continuously differentiable with $k$ and $\partial_{x_j} k(x,x)$$\mu$-integrable. Then there exists a unique representer function $g_{\mathcal{S}, \mu} \in \mathcal{H}^d$ such tha with explicit form $g_{\mathcal{S}, \mu}(\cdot) = \mathbb{E}_{X \sim \mu}[\mathcal{S}_1 K(X, \cdot)

Figures (4)

  • Figure 1: Stein-Type KDS for a stationary linear SDE: We applied SKDS to $n=5000$ samples from the target distribution $\mu$, using a one-dimensional Gaussian kernel with bandwidth $0.5$. The models correspond to the SDE $dX_t = -4(X_t - \alpha)dt + \sigma dB_t$ for different values of $\alpha$ and $\sigma$. 1) PDFs of the stationary distributions. 2) An SKDS contour plot in $\alpha$ and $\sigma$. 3) and 4) Partial derivatives of the SKDS objective function for both model alternatives, showing that the gradient vanishes at the parameter values of the ground-truth model.
  • Figure 2: Benchmarking results for $d = 20$ variables with an Erdős--Rényi causal structure. The Wasserstein-$2$ distance ($\mathcal{W}_2$) was computed from $10$ test interventions on unseen target variables across $50$ randomly generated systems. Box plots depict the medians and interquartile ranges (IQR), with whiskers extending to the largest value within $1.5$ times the IQR from the boxes.
  • Figure 3: Benchmarking results for $d = 20$ variables. The MSE (left column) and the Wasserstein-$2$ distance (right column) distance MSE were computed from $10$ test interventions on unseen target variables across $50$ randomly generated systems and $6$ different data generating methods (rows). Box plots depict the medians and interquartile ranges (IQR), with whiskers extending to the largest value within $1.5$ times the IQR from the boxes.
  • Figure 4: Nested Stein-type KDS operator defining the custom objective function integrated into the stadion framework stadion2024.

Theorems & Definitions (34)

  • Definition 1: SKDS Operator
  • Definition 2: Stein-type KDS
  • Example 3: SKDS as a learning objective
  • Lemma 4: Representer function $g_{\mathcal{S}, \mu}$
  • Lemma 5: SKDS Closed Form
  • Proposition 6
  • Theorem 7: SKDS characterizes reversible diffusions
  • Corollary 8: Convexity of SKDS
  • Example 9: Linear SDE parametrization
  • Proposition 10: SKDS empirical estimate convex
  • ...and 24 more