Flow Maps for Trajectory Inference to Improve Learned Conditional Laws

Trajectory inference seeks to reconstruct a stochastic process from unpaired samples of its temporal marginals. Since the true process is inherently unidentifiable from the limited observed data, encoding inductive biases into trajectory inference methods is essential to select a plausible approximate process. Many recent methods construct an interpolant process between observed marginals and try to match its probability path with a neural ODE or SDE, but little attention is paid to the learned joint distribution law across time. In many applications, there exists information on the characteristics of this law. In single-cell differentiation, for example, biological knowledge specifies plausible transitions between cell types, and we seek an efficient way to incorporate such knowledge. This would require constraining model-generated trajectories directly, which conventionally involves numerical integration and differentiation through the solver, making training computationally expensive. To address this, we propose CLIFT, a method using (stochastic) flow maps to efficiently impose differentiable constraints on learned conditional laws. A flow map provides parallelisable access to finite-time transitions, creating a natural interface for conditional-law losses whose gradients flow directly to the model parameters. Across single-cell datasets, CLIFT improves adherence to known biological differentiation rules while maintaining competitive marginal distributional fidelity and being significantly more efficient than imposing constraints through numerical rollouts.