EntangledSBM for Interacting Particle Systems

• EntangledSBM is a framework that learns the stochastic dynamics of interacting systems by coupling each particle’s evolution with the joint configuration and velocities of all others.
• It generalizes the classical Schrödinger bridge problem to account for interaction-dependent dynamics, making it applicable to biomolecular dynamics and heterogeneous cell populations.
• The method minimizes a KL divergence and an entanglement loss using neural networks, Transformers, and Sinkhorn-based techniques to ensure coordinated and realistic trajectory simulations.

Entangled Schrödinger Bridge Matching (EntangledSBM) is a framework for learning the stochastic dynamics of interacting, multi-particle systems where each particle's evolution is coupled—entangled—with the joint configuration and velocities of all others. The formulation generalizes the classical Schrödinger bridge problem—originally developed for systems of single or independent particles—to settings such as biomolecular dynamics and heterogeneous cell population models, where interaction-dependent dynamics are fundamental. EntangledSBM brings together recent unifying perspectives on the bridge problem from stochastic control, generative modeling, and path-space statistical inference, enabling tractable and theoretically grounded simulation of interacting systems under complex energy landscapes (Tang et al., 10 Nov 2025, Kim, 27 Mar 2025).

1. Theoretical Underpinnings and Motivation

The classical Schrödinger bridge problem is to find the most likely stochastic evolution (path law) connecting two prescribed endpoint distributions, under a reference dynamics—often Brownian motion or an SDE. Given marginals $\pi_\mathcal{A}$ and $\pi_\mathcal{B}$ at $t{=}0,T$, the bridge $\mathbb{P}^\star$ solves

$$\min_b\; D_{\mathrm{KL}}(\mathbb{P}^b\|\mathbb{P}^0) \quad \text{s.t.}\quad X_0\sim\pi_\mathcal{A},\; X_T\sim\pi_\mathcal{B}$$

where $\mathbb{P}^0$ is the uncontrolled path law and $\mathbb{P}^b$ corresponds to the drift-controlled SDE.

Modern deep learning instantiations, such as Deep Schrödinger Bridge Matching (DSBM), utilize neural networks to parameterize $b_t^\theta$ and estimate the optimal bridge by data-driven regression on marginal-matching losses. These models, however, generally treat system components as independent or mean-field coupled—insufficient for systems where coordinated dynamics (as in proteins or communicating cells) are key.

EntangledSBM resolves this by introducing an "entanglement" in both the parametrization and inference: each particle's drift depends on the full system's configuration and velocities, capturing heterogeneous and time-evolving interdependencies.

2. Mathematical Formulation

The state at time $t$ is $$X_t = (R_t,V_t) \in \mathbb{R}^{n\times d} \times \mathbb{R}^{n\times d}$$, where $R_t$ are positions and $V_t$ velocities for $n$ particles in $d$ dimensions.

Controlled SDE dynamics (underdamped Langevin with bias $b^i$ for particle $i$): $$\begin{cases} d\,r^i_t = v^i_t\,dt \ d\,v^i_t = \frac{-\nabla_{r^i} U(R_t) + b^i(R_t,V_t)}{m_i}\,dt - \gamma v^i_t\,dt + \sqrt{\frac{2\gamma k_B T}{m_i}}\,dW^i_t \end{cases}$$ where $U$ is the potential, $\gamma$ the friction, $k_B$ Boltzmann constant, $T$ temperature, and $m_i$ the mass.

Entangled bias force: Each $b^i(R_t, V_t)$ is a function of the whole $(R_t, V_t)$. To guarantee progress toward the target marginal while allowing coordinated movement, the force is decomposed as: $$b^i(R_t, V_t) = \alpha^i(R_t,V_t)\,\hat s^i + (I - \hat s^i (\hat s^i)^\top) h^i(R_t,V_t)$$ where $$\hat s^i = \nabla_{r^i} \log\pi_{\mathcal{B}}(r^i_t)/||\nabla_{r^i}\log\pi_{\mathcal{B}}(r^i_t)||$$, $\alpha^i \geq 0$, and $h^i$ parameterizes lateral interaction. This structuring ensures that each particle does not increase its distance from the target while permitting coordinated detours.

Learning objective: The parameterized path law $\mathbb{P}^{b_\theta}$ is trained to minimize

$$\min_\theta~ D_{\mathrm{KL}}(\mathbb{P}^{b_\theta}\|\mathbb{P}^\star)$$

with $$\mathbb{P}^\star(X_{0:T}) \propto \mathbb{P}^0(X_{0:T})\pi_{\mathcal{B}}(R_T)$$. By stochastic optimal control theory, this is equivalent to minimizing

$$\mathbb{E}_{\mathbb{P}^{b_\theta}}\left[ \frac12 \int_0^T \|b_\theta(R_t, V_t)\|^2\,dt - \log\pi_{\mathcal{B}}(R_T)\right]$$

3. Connection to Unified Bridge Algorithms and Entanglement Regularization

The Unified Bridge Algorithm (UBA) (Kim, 27 Mar 2025) reinterprets bridge problems as conditional regression on bridge-pinned marginals and couplings, unifying both ODE and SDE perspectives. In DSBM, the forward and backward drifts are typically learned independently. EntangledSBM introduces a coupling—entanglement—between forward and backward drift functions.

This entanglement is imposed by:

• Maintaining a joint coupling $Q_\mathrm{ent}(x_0, x_1; y_1, y_0)$ that links forward and backward path endpoints.
• Introducing an explicit entanglement loss:

$$\mathcal{L}_\mathrm{ent} = \lambda\,\mathbb{E}_{t, x_t, y_t} \|v_f(t, x_t) - v_b(1-t, y_t)\|^2$$

where $v_f$ and $v_b$ are the forward and backward drifts, and $x_t$ and $y_{1-t}$ are paired intermediate states.

• Optionally including mutual-information regularization to ensure the coupled endpoint distributions remain well-posed yet entangled.

This design forces the forward and backward flows to align at intermediate states, removes asymmetries, and improves the fidelity of simulated trajectories.

4. Algorithmic Implementation

The training process comprises simulating controlled SDE trajectories, computing cross-entropy losses on generated rollouts, and updating the entangled drift networks.

High-Level Pseudocode

for rollout in range(N_rollouts): # (1) Simulate M trajectories under SDE with current b_theta rollouts = simulate_SDE(alpha_theta, h_theta, U, M) R.append(rollouts) for step in range(N_steps): # (2) Sample minibatch of trajectories from R batch = R.sample_batch() # (3) Compute importance weights w^* using endpoint likelihood ratios w_star = compute_weights(batch) # (4) Compute cross-entropy loss L_CE L_CE = loss_cross_entropy(batch, w_star, alpha_theta, h_theta) # (5) Gradient update optimizer.step(L_CE)

Key algorithmic elements include:

• Trajectories are simulated forward with the current entangled bias.
• A replay buffer $\mathcal{R}$ enables stabilized sampling.
• Cross-entropy loss is preferred over log-variance objectives, as it is convex in the path measure and empirically yields smooth, manifold-respecting interpolants.
• Neural parameterization employs Transformers (for input permutation invariance across particles) and separate MLP heads for $\alpha^i$, $h^i$. Kabsch alignment enforces rotational invariance in molecular tasks.
• Entanglement losses are applied at each step, and endpoint couplings are re-estimated between forward- and backward-generated samples using entropic OT/Sinkhorn methods.

5. Empirical Results and Comparative Evaluation

Cell-Cluster Dynamics under Drug Perturbation

• Dataset: Tahoe-100M single-cell atlas for A549 under Clonidine/Trametinib ($5\mu$M); PCA to $d \in\{50,100,150\}$.
• Metrics: RBF-MMD, 1-Wasserstein, 2-Wasserstein, both on seen and unseen (hold-out) clusters.
• Findings: EntangledSBM with cross-entropy loss and velocity conditioning achieves the lowest RBF-MMD and Wasserstein distances. Omitting velocity conditioning or using log-variance (LV) loss degrades performance, with the latter producing sharp off-manifold transitions.

Transition-Path Sampling for Molecular Systems

• Tasks: Alanine Dipeptide (dihedral angles), fast-folding proteins (Chignolin, Trp-cage, BBA; $n = 138-504$ atoms).
• Metrics: RMSD to target (after Kabsch alignment), target hit percentage (THP), energy of highest barrier crossed (ETS).
• Findings: EntangledSBM achieves state-of-the-art RMSD and THP on Alanine Dipeptide, superior hit rates across all proteins, and generates physically plausible transition paths that traverse realistic energy barriers and preserve the underlying manifold structure.

A summary of the main metrics and baseline comparisons:

Task	Metric	EntangledSBM	Baseline (best previous)
A549 Cell Clusters	RBF-MMD	Lowest	Degraded/LV loss
Alanine Dipeptide	RMSD, THP, ETS	SOTA	Slightly higher/lower
Chignolin/Proteins	Hit Rate, Barrier	Higher, Lower	Lower, Higher

6. Practical Considerations, Limitations, and Applications

• Scalability: Computational cost increases with number of particles $n$ due to all-to-all coupling (attention); approximations with random-batch or graph-sparse attention are plausible directions for large systems.
• Neural Architectures: U-Nets, time-conditioned MLPs, sinusoidal encodings, and FiLM layers are effective; Transformers are used for handling set-valued inputs.
• Stability: Entanglement loss ($L_\mathrm{ent}$) can destabilize training at high weights; annealing strategies for $\lambda$ are used.
• Efficiency: Minibatch Sinkhorn ($256{-}1024$ samples) is employed for endpoint coupling estimation; fast GPU-based OT solvers are recommended.
• Limitations: Demonstrated primarily on small/medium systems; scaling to thousands of particles requires further innovation.
• Applications: Enhanced sampling in biomolecular molecular dynamics, simulation of cell response to perturbation, and modeling any interacting particle system with data-driven endpoints.

A plausible implication is that EntangledSBM provides a blueprint for next-generation simulation methods in biological and molecular systems by systematically integrating interaction-aware path generation under tractable, data-calibrated objectives.

7. Conclusion and Future Directions

EntangledSBM extends the Schrödinger bridge framework to the domain of interacting, second-order stochastic systems, introducing theoretical guarantees (convexity of the objective, existence and uniqueness under regularity) while yielding empirical performance matching or surpassing specialized baselines in both single-cell and molecular trajectory modeling tasks.

Future directions include optimizing computational scaling, incorporating unbalanced and branching end marginals, and exploring application to large-scale agent-based or granular flow models. EntangledSBM serves as a foundation for modeling heterogeneous systems where snapshot data exist but the underlying interaction-driven dynamics are unknown (Tang et al., 10 Nov 2025, Kim, 27 Mar 2025).