🤖 machine learning

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

This paper introduces a scalable, data-driven framework for estimating the Kullback-Leibler divergence between probability measures on function space, providing a principled metric to evaluate and compare trajectory inference methods under partial observability where traditional marginal-based assessments are insufficient.

Original authors: Chao Wang, Luca Nepote, Giulio Franzese, Pietro Michiardi

Published 2026-04-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chao Wang, Luca Nepote, Giulio Franzese, Pietro Michiardi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to reconstruct a movie from a few scattered, frozen frames.

In the world of single-cell biology (like studying how a stem cell turns into a heart cell), scientists can't film the process continuously. The "camera" is destructive: to see a cell's state, they have to kill it. So, all they have are snapshots of different cells at different times. They call this Trajectory Inference (TI). The goal is to guess the full "movie" (the path every cell took) based only on these still photos.

The Problem: The "Snapshots Lie"

The paper points out a major flaw in how scientists currently judge these movies. They usually check if the "frames" (snapshots) look correct.

The Analogy: Imagine two different movies.
- Movie A: A car drives smoothly down a highway.
- Movie B: A car drives in circles, then jumps, then drives smoothly.
- If you only look at the photo at the 1-minute mark, both cars might be in the exact same spot. Current evaluation methods say, "Great! Both movies are perfect!"
- But if you watched the whole movie, Movie B would be a disaster. The current methods miss the dynamics (the movement between the photos) because they only look at the marginals (the photos).

The Solution: The "Functional KL" (FKL)

The authors introduce a new way to measure how different two "movies" are. They call it Functional KL (Kullback-Leibler divergence in function space).

Instead of comparing just the photos, they compare the entire script and the choreography of the movement.

Here is how they do it, using a simple metaphor:

1. The "Flow" Metaphor

Imagine the cells are water droplets moving through a river.

The Old Way: You take a photo of the water at 10 AM and 2 PM. You check if the water looks the same in both photos.
The New Way (FKL): You look at the currents (the velocity fields) that push the water.
- The authors realized that if you know the "current" (the force pushing the water) at every single moment, you can mathematically calculate exactly how different two rivers are.
- They developed a formula that turns the complex problem of comparing infinite paths into a simple calculation: "How different are the currents pushing the water?"

2. The "Noise" Trick

To make this math work, they had to be clever about "noise."

The Analogy: Imagine trying to hear a whisper in a noisy room. If the room is too quiet, the whisper is hard to distinguish from silence. If the room is too loud, you can't hear anything.
They found that by adding a specific type of "rough" background noise (like static on a radio) to their math, they could make the "whisper" (the actual cell movement) stand out clearly. This allowed them to measure the difference between the real movie and the guessed movie with high precision.

What They Found

They tested their new "Movie Critic" (FKL) against the old "Photo Critic" (Marginal Metrics) using both computer simulations and real biological data.

The Old Critic is Confused: The old methods often gave conflicting results. Sometimes they said Method A was best; other times, they said Method B was best, depending on which specific photo you looked at.
The New Critic is Consistent: The FKL metric gave a clear, consistent ranking.
- Example: In one test, a method called TIGON looked great in the photos (low error). But the FKL metric revealed that TIGON was actually generating "smooth, boring" paths that didn't match the chaotic, real biological movement.
- Meanwhile, a method called SBIRR looked slightly worse in the photos but had the correct "choreography" (dynamics). The FKL metric correctly identified SBIRR as the better model.

Why This Matters

This paper provides a principled ruler for the future of biology.

Before, scientists were guessing which algorithm was best based on incomplete data (just the photos).
Now, they have a tool that checks if the story makes sense, not just the pictures.

In a nutshell: The authors built a new measuring tape that doesn't just check if the destination looks right; it checks if the journey was taken correctly. This helps scientists choose the right tools to understand how life develops, heals, and changes.

1. Problem Statement

Trajectory Inference (TI) aims to reconstruct latent dynamical processes from snapshot data (independent samples from time-indexed marginals), a common scenario in destructive measurement fields like single-cell genomics.

The Core Challenge: Because measurements are destructive, the full path-space law (the probability distribution over entire trajectories) is non-identifiable from finite marginals.
Limitations of Current Evaluation: Existing TI methods are primarily evaluated using marginal metrics (e.g., Earth Mover's Distance, Wasserstein distances) on held-out snapshots. The authors argue these metrics are fundamentally limited because:
1. They only compare distributions at discrete time points, ignoring temporal correlations and coupling between states.
2. Distinct trajectory models can produce identical marginals while inducing vastly different (and potentially incorrect) dynamics.
3. Rankings based on marginal metrics are often inconsistent, varying with the choice of validation time point, metric type, and hyperparameters.

The paper proposes a solution: a principled framework to estimate the Kullback–Leibler (KL) divergence between probability measures on function space, allowing for a direct comparison of full trajectory distributions rather than just snapshots.

2. Methodology

The authors introduce a tractable estimator for the Functional KL divergence (FKL) based on Functional Flow Matching (FFM) and infinite-dimensional optimal transport theory.

Theoretical Framework

Setting: The problem is formulated in a real separable Hilbert space $H$ (e.g., $L^2$ ). The goal is to estimate $KL(\nu_A \| \nu_B)$ , where $\nu_A$ and $\nu_B$ are probability measures over trajectories.
Assumptions:
1. Absolute Continuity: $\nu_A \ll \nu_B$ (Radon–Nikodym derivative exists).
2. Cameron–Martin Support: The measures are fully supported on the Cameron–Martin space of the noise covariance $C$ . This requires the noise to be "rougher" than the data distribution but still trace-class.
Derivation:
1. Velocity Fields: The authors utilize FFM, which models the evolution of measures via a velocity field $v_t(x)$ .
2. Continuity Equation: They apply the weak continuity equation using the Radon–Nikodym derivative and its logarithm as test functions.
3. Key Theorem: By linking the logarithmic gradients of the measures to the velocity fields, they derive a velocity-only representation of the KL divergence:
  $KL(\nu_A \| \nu_B) = \int_0^1 \int_H \frac{t}{1-t} \| v^A_t(x) - v^B_t(x) \|_{H_{\mu_0}}^2 d\mu^A_t(x) dt$
  Here, the divergence is expressed as the time-integrated squared discrepancy between the velocity fields of the two processes, weighted by the Cameron–Martin norm.

Implementation

Neural Operator: The velocity fields $v^A$ and $v^B$ are approximated jointly using a single Mesh-Informed Neural Operator (MINO) conditioned on a binary flag.
Stabilization: To handle the singularity at $t=1$ (due to the $\frac{t}{1-t}$ term), the network enforces the boundary condition $v(x, 1) = x$ via a subtraction-based parameterization.
Estimation: The FKL is estimated via Monte Carlo integration using samples from the reference measure, the noise distribution, and uniform time sampling, without requiring full ODE integration of the generative dynamics.

3. Key Contributions

General FKL Estimator: A theoretically grounded, tractable estimator for KL divergence between probability measures on function spaces, derived from infinite-dimensional flow matching theory.
Validation on Ground Truth: The estimator is validated on synthetic cases (Gaussian measures and linear SDEs) where analytic KL values are known, demonstrating high accuracy and numerical stability.
Comprehensive TI Benchmark: An extensive evaluation of state-of-the-art TI methods (including SBIRR, vSB, MSBM, MFL, AM, and TIGON) on both synthetic (Lotka-Volterra, Repressilator, Petal) and real-world scRNA-seq datasets.
Demonstration of Marginal Metric Failure: The paper provides empirical evidence that marginal-based rankings are inconsistent and often misaligned with visual inspection of dynamics, whereas FKL provides coherent, trajectory-level assessments.

4. Results

Accuracy: In special cases, the estimated FKL closely matches the analytic ground truth across varying dimensions and noise parameters.
Inconsistency of Marginal Metrics:
- Rankings of methods fluctuate wildly depending on the validation time point $\tau$ and the specific metric used (e.g., EMD vs. MMD).
- Methods can achieve low marginal distances (good snapshot fit) while generating completely wrong dynamics (e.g., TIGON on the Repressilator dataset produced smooth spirals instead of the correct oscillatory behavior, yet achieved competitive marginal scores).
FKL Superiority:
- FKL rankings align with visual inspection and biological plausibility. For instance, it correctly identifies SBIRR as superior on Lotka-Volterra and Repressilator, and MSBM on the Petal dataset, whereas marginal metrics often misrank these.
- Mode Seeking vs. Mode Coverage: The paper highlights the asymmetry of KL:
  - Forward KL ( $KL(\nu_A \| \nu_B)$ ) favors mode coverage (methods like MFL that cover broad support, even if off-data).
  - Reverse KL ( $KL(\nu_B \| \nu_A)$ ) favors mode seeking (methods like MSBM that stay strictly within high-density regions).
- This asymmetry reveals nuances in method behavior that marginal metrics cannot capture.

5. Significance

This work fundamentally shifts the evaluation paradigm for Trajectory Inference. By moving from finite-dimensional marginal comparisons to infinite-dimensional function-space divergences, the authors provide:

A Principled Criterion: FKL offers a mathematically rigorous way to evaluate how well a model captures the dynamics of a process, not just the static snapshots.
Resolution of Ambiguity: It resolves the "invisibility" of dynamic errors in current benchmarks, where models with incorrect lineage relationships or temporal couplings can appear optimal under marginal metrics.
Scalability: The proposed estimator is data-driven and scalable to realistic snapshot datasets, making it a practical tool for the single-cell genomics community and other fields relying on trajectory reconstruction from partial observations.

In conclusion, the paper argues that functional relative entropy should become the standard for evaluating trajectory inference under partial observability, as it directly measures the fidelity of the inferred dynamical laws.