Efficient Learning of Stationary Diffusions with Stein-type Discrepancies

TL;DR

This paper introduces SKDS, a Stein-type kernel deviation from stationarity, which efficiently ensures a learned diffusion's stationary distribution matches a target, with theoretical guarantees and practical improvements over existing methods.

Contribution

The paper proposes SKDS, a novel Stein-based formulation that guarantees stationarity alignment and offers computational advantages over previous kernel deviation methods.

Findings

SKDS guarantees stationary distribution alignment when it vanishes.

Empirical results show SKDS achieves comparable accuracy to KDS with lower computational cost.

SKDS outperforms most baselines in experiments.

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 $\epsilon$-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.