Backward Kolmogorov Transport: Sampling Invariant Laws of Reversible Diffusions from Trajectory Data

Abstract

We consider a reversible diffusion with unknown drift, observed through trajectories, and construct samplers of its invariant law. The trajectories contain samples of this law, but a sampler built from them by density or score estimation inherits the errors of these estimates, in particular on the barriers between metastable states. The trajectories also contain the dynamics, and for a reversible diffusion the ratio between the law of the process and its invariant law solves the backward Kolmogorov equation, so that the eigenpairs of the generator propagate it in closed form. Backward Kolmogorov transport (BKT) evaluates this ratio with eigenpairs estimated from the trajectories and coefficients computed once from the initial particles, and moves the particles along the Wasserstein gradient flow of the Kullback-Leibler divergence. The particles do not interact, BKT defines a transport map, and no density, score or drift is estimated. With exact eigenpairs BKT transports the spectral projection of the initial law exactly when this projection is positive. For estimated eigenpairs we prove a transport identity in which truncation enters only through the initial law and estimation only through the residual of the eigenpairs, and a local stability bound in the 2-Wasserstein distance. With coefficients computed from the transported particles, the sampling error of the retained modes cancels, and a central limit theorem identifies the asymptotic variance. Experiments on Ornstein-Uhlenbeck processes, multi-well potentials, separable products and alanine dipeptide, including comparisons with particle methods based on density estimates, agree with these results, and with an accurate spectrum the errors of BKT reach or fall below those of independent samples.