FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

News

• Jun 2026 Camera-ready submitted; code & project page released.
• Jun 2026 FreeBridge accepted at MICCAI 2026.

Overview

Modeling how cells respond to perturbations is central to phenotypic drug discovery — yet high-content imaging fixes cells at acquisition, so a single cell's trajectory is never observed, only control and perturbed populations as separate marginals.

Matching endpoints alone leaves the interior undetermined: many stochastic processes share the same control and perturbed marginals while passing through morphologies that never occur in real cells.

FreeBridge is a Schrödinger Bridge for single-cell transition modeling under endpoint-only supervision — stochastic transport regularized to stay on the manifold of real single-cell morphologies.

• Endpoint-only supervision. Learns stochastic transport from control and perturbed marginals alone — no observed per-cell trajectories.
• Manifold-constrained transport. A Cell Engine separates a fixed single-cell state manifold from the learned bridge drift; an empirical support cost keeps intermediate states on observed morphology.
• State of the art. Cuts support violations 0.31 → 0.11 and overall FID 18.72 → 15.35 vs. CellFlux on BBBC021, with consistent gains on RxRx1 and JUMP.

Keywords: High-Content Imaging · Cellular Morphology Modeling · Schrödinger Bridges · Perturbation Dynamics

Method

FreeBridge casts endpoint-only supervision as a variational Schrödinger Bridge: among all stochastic processes matching the observed control and perturbed marginals, it seeks the one of minimal control energy whose intermediate states stay on the empirical morphology manifold, trading transport energy against a latent-support penalty $V(z_t)$:

The Cell Engine

The model is realized as a Cell Engine that separates state specification from stochastic transport. Microscopy fields are instance-segmented into single-cell crops, and a frozen encoder embeds each crop as $z=\phi(x)$; the collection of embeddings forms an empirical latent support set $\mathcal{B}$ that pins down the manifold of real morphologies. Transport is a time-conditioned drift driving the SDE $dz_t = f_{\theta}(z_t,t,c)\,dt + \sigma\,dW_t$ from control cells $z_0\sim\nu_{z,c}$ at $t=0$ to perturbed cells $z_1\sim\mu_z$ at $t=1$. The drift is trained variationally to minimize control energy while the support cost $V(z)=\min_{b\in\mathcal{B}}\|z-b\|^{2}$ penalizes excursions off the manifold; the weight $\lambda_{\text{bank}}$ trades transport energy against support feasibility. Intra-batch pairing within each plate supplies the coupling, and instance segmentation ensures the bridge operates on single cells rather than whole fields.

Why a Schrödinger Bridge?

GANs map noise directly to data and train unstably; diffusion models discretize generation into many denoising steps; both match the perturbed endpoint but say nothing about the path taken to reach it. A Schrödinger Bridge instead selects the minimal-energy stochastic process between two given marginals — exactly the control-to-perturbed setting. Adding the empirical support cost ties that process to observed morphology, so intermediate states remain biologically plausible rather than drifting through regions of latent space that no real cell occupies.

Datasets

FreeBridge is evaluated on three high-content imaging benchmarks under a single unified single-cell protocol.

Results

We evaluate FreeBridge on BBBC021, RxRx1, and JUMP under a unified protocol.

Quantitative comparison

FreeBridge improves overall and condition-level FID and KID over PhenDiff, IMPA, and CellFlux across all three benchmarks, and cuts the support-violation rate on BBBC021 by roughly 3×.

5k generated samples per condition, averaged over three seeds. R_viol (Support Violation Rate): fraction of intermediate states outside the empirical morphology support — a support-feasibility measure, not verified temporal ordering.

We compare endpoint morphology on BBBC021 against IMPA, PhenDiff, and CellFlux for two perturbations with distinct phenotypes.

Endpoint alignment leaves the interior underdetermined. We visualize trajectories across $t = 0 \to 1$ and sweep the support weight $\lambda_{\text{bank}}$.

Setting $\lambda_{\text{bank}}=0$ removes the support cost and recovers the CellFlux baseline — matching its R_viol of 0.31 in the main table.

We ablate the three ingredients of the Cell Engine, reporting endpoint fidelity and mechanism-of-action retention.

Full-model numbers are from an independent ablation run and differ marginally from the main table.

We sweep support weight $\lambda_{\text{bank}}$, base noise $\sigma$, and inference steps (NFE) on BBBC021.

Quick Start

The single-cell pipeline runs control/perturbed endpoints through a frozen latent encoder, then trains the bridge with Hydra experiment profiles.

# PyTorch-Lightning 1.x required (pinned to 1.8.5.post0) — 2.x is not supported
conda env create -f environment.yml
conda activate freebridge
pip install -e .

# crop single cells → train VAE → export latents → build endpoints
python data_prep/crop_bbbc021_singlecell.py --mask_root <masks>
python data_prep/train_vae_full.py
python data_prep/export_latents_from_vae.py
python data_prep/build_bbbc021_endpoints.py \
    --control_regex DMSO --target_regex taxol \
    --out_dir data/bbbc021

# default FreeBridge, support weight lambda_bank = 0.5
python train.py experiment=bbbc021

# support-weight sweep
python train.py experiment=bbbc021 state_cost.support_weight=0.0
python train.py experiment=bbbc021 state_cost.support_weight=1.0

Epoch 0 uses a straight-line warm start; manifold-constrained paths activate from epoch 1, so set max_epochs ≥ 2. Full setup and configs are in the GitHub repository.

@article{wang2026freebridge,
    title={FreeBridge: Variational Schr{\"o}dinger Bridges for Cellular Transition Dynamics},
    author={Wang, Xurui and Ren, Qin and Ma, Jun and Ling, Haibin and You, Chenyu},
    journal={arXiv preprint arXiv:2606.11286},
    year={2026},
}

Acknowledgements

Built on GSBM (Liu et al., ICLR 2024); the single-cell task and unified evaluation protocol follow CellFlux (Zhang et al., 2025). We thank the authors for releasing their code.