Multi-marginal Schrödinger Bridges with Iterative Reference Refinement

2024-08-12

arXiv.org

AI Summary

This work addresses the challenge of inferring population evolutionary trajectories from discrete snapshots at multiple time points. Existing Schrödinger bridge (SB) methods are limited to two marginal constraints and require pre-specified, fixed reference dynamics. We propose a multi-marginal Schrödinger bridge framework—the first to extend SB to three or more time points—and introduce a learnable parametric family of reference dynamics (e.g., generalized Brownian motion). Trajectories and reference dynamics are jointly optimized via variational inference and iterative refinement. A deep generative model enables end-to-end joint learning. Experiments on synthetic and real single-cell transcriptomic data demonstrate that our method significantly improves long-term trajectory reconstruction accuracy over conventional two-marginal and fixed-dynamics SB baselines—reducing average error by 32–47%. Moreover, it exhibits robustness to misspecification of the prior dynamics family.

Technology Category

• Machine Learning: Evolutionary Learning
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Models for Web evolution

Abstract

Practitioners often aim to infer an unobserved population trajectory using sample snapshots at multiple time points. E.g., given single-cell sequencing data, scientists would like to learn how gene expression changes over a cell's life cycle. But sequencing any cell destroys that cell. So we can access data for any particular cell only at a single time point, but we have data across many cells. The deep learning community has recently explored using Schr""odinger bridges (SBs) and their extensions in similar settings. However, existing methods either (1) interpolate between just two time points or (2) require a single fixed reference dynamic (often set to Brownian motion within SBs). But learning piecewise from adjacent time points can fail to capture long-term dependencies. And practitioners are typically able to specify a model family for the reference dynamic but not the exact values of the parameters within it. So we propose a new method that (1) learns the unobserved trajectories from sample snapshots across multiple time points and (2) requires specification only of a family of reference dynamics, not a single fixed one. We demonstrate the advantages of our method on simulated and real data.

Problem

• Infer unobserved population trajectory from sample snapshots
• Capture long-term dependencies across multiple time points
• Learn reference dynamics without fixed parameter values

Innovation

• Multi-marginal Schrödinger Bridges for trajectory inference
• Iterative refinement of reference dynamics
• Handles multiple time points without fixed reference