Relative Entropy Estimation in Function Space: Theory and Applications to Trajectory Inference

TL;DR

This paper introduces a scalable method for estimating the KL divergence between probability measures on function space, improving trajectory inference evaluation in single-cell genomics.

Contribution

A novel framework for functional KL divergence estimation that enables coherent assessment of trajectory inference methods from snapshot data.

Findings

Functional KL closely matches analytic KL on benchmarks.

Path-space KL provides consistent evaluation of trajectory inference methods.

Discrepancies in inferred dynamics are revealed in regions with sparse data.

Abstract

Trajectory Inference (TI) seeks to recover latent dynamical processes from snapshot data, where only independent samples from time-indexed marginals are observed. In applications such as single-cell genomics, destructive measurements make path-space laws non-identifiable from finitely many marginals, leaving held-out marginal prediction as the dominant but limited evaluation protocol. We introduce a general framework for estimating the Kullback-Leibler divergence (KL) divergence between probability measures on function space, yielding a tractable, data-driven estimator that is scalable to realistic snapshot datasets. We validate the accuracy of our estimator on a benchmark suite, where the estimated functional KL closely matches the analytic KL. Applying this framework to synthetic and real scRNA-seq datasets, we show that current evaluation metrics often give inconsistent assessments, whereas path-space KL enables a coherent comparison of trajectory inference methods and exposes discrepancies in inferred dynamics, especially in regions with sparse or missing data. These results support functional KL as a principled criterion for evaluating trajectory inference under partial observability.