Entangled Schrödinger Bridge Matching

TL;DR

Entangled Schrödinger Bridge Matching (EntangledSBM) introduces a framework to learn interacting, second-order dynamics for multi-particle systems by using entangled bias forces that depend on all particle positions and velocities. The bias is parameterized with a Transformer and is conditioned on a target distribution to enable generalization to unseen endpoints, while training uses an off-policy cross-entropy objective to minimize the KL divergence to the optimal bridge path. The method is validated on two domains: simulating heterogeneous cell populations under perturbations and sampling transition paths in all-atom molecular dynamics, achieving accurate target distribution reconstruction and feasible transition paths across high-energy barriers. This approach advances trajectory learning beyond mean-field or snapshot-based methods and offers practical tools for drug discovery and biomolecular dynamics with improved generalization and trajectory realism.

Abstract

Simulating trajectories of multi-particle systems on complex energy landscapes is a central task in molecular dynamics (MD) and drug discovery, but remains challenging at scale due to computationally expensive and long simulations. Previous approaches leverage techniques such as flow or Schrödinger bridge matching to implicitly learn joint trajectories through data snapshots. However, many systems, including biomolecular systems and heterogeneous cell populations, undergo dynamic interactions that evolve over their trajectory and cannot be captured through static snapshots. To close this gap, we introduce Entangled Schrödinger Bridge Matching (EntangledSBM), a framework that learns the first- and second-order stochastic dynamics of interacting, multi-particle systems where the direction and magnitude of each particle's path depend dynamically on the paths of the other particles. We define the Entangled Schrödinger Bridge (EntangledSB) problem as solving a coupled system of bias forces that entangle particle velocities. We show that our framework accurately simulates heterogeneous cell populations under perturbations and rare transitions in high-dimensional biomolecular systems.

Figures (8)

• Figure 1: Entangled Schrödinger Bridge Matching. We consider the problem of simulating interacting multi-particle systems, where each particle's velocity depends dynamically on the velocities of the other particles in the system, and introduce EntangledSBM, a framework that parameterizes an entangled bias force capturing the dynamic interactions between particles.
• Figure 2: Simulated cell cluster dynamics with EntangledSBM under Clonidine perturbation with 150 PCs. Nearest neighbour cell clusters with $n=16$ cells are simulated over $100$ time steps to perturbed cells seen during training (Top) and an unseen perturbed population (Bottom). The gradient indicates the evolution of timesteps from the initial time $t=0$ (navy) to the final time $t=T$ (purple or turquoise). (A) Trajectories under base dynamics with no bias force. (B) Trajectories simulated with base and bias forces with (B) no velocity conditioning, (C) log-variance (LV) objective, and (D) cross-entropy (CE) objective.
• Figure 3: Comparison of EntangledSBM with and without velocity conditioning for cell cluster simulation under Trametinib perturbation. 50 PCs are simulated with the learned bias force trained with the CE objective without velocity conditioning (Top) and with velocity conditioning (Bottom) to (B) the perturbed population used for training and (C, D) the two unseen target populations.
• Figure 4: Transition paths generated with EntangledSBM. The energy landscape is colored from dark navy (low potential energy) to light purple (high potential energy) and plotted with the backbone dihedral angles $(\phi, \psi)$ for Alanine Dipeptide and the top two TICA components for the fast-folding proteins. The starting state $\boldsymbol{R}_0$ and target states $\boldsymbol{R}_\mathcal{B}$ are indicated with purple stars, and the final states of the simulated trajectories $\boldsymbol{R}_T$ are indicated with pink circles.
• Figure 5: Training and held-out cell clusters for the cell perturbation experiment. All cells are plotted with the top 2 PCs. (A) Clonidine ($5\mu L$) perturbation data containing the initial DMSO-treated control cells (dark blue), perturbed cluster for training (magenta), and held-out validation cluster (turquoise). (B) Trametinib ($5\mu L$) perturbation data containing the initial DMSO-treated control cells (navy blue), the perturbed cluster for training (purple), and two held-out perturbed clusters (magenta and turquoise).

Theorems & Definitions (15)

• Definition 3.1: Entangled Schrödinger Bridge Problem
• Proposition 3.1: Solving EntangledSB with Stochastic Optimal Control
• Proposition 4.1: Non-Increasing Distance from Target Distribution
• Proposition 4.2: Convexity and Uniqueness of Cross-Entropy Objective
• Proposition 4.3: Equivalence of Variational and Path Integral Objectives
• Proposition 4.4: Discretized Cross-Entropy
• Lemma B.1: Radon-Nikodym Derivative
• Theorem B.1
• Proposition B.1: Monotone Optimality of Entangled Bias Forces
• Lemma C.1