FreeBridge：细胞转变动力学的变分Schrödinger桥

# FreeBridge：面向细胞动态转换的变分薛定谔桥
来源：https://arxiv.org/html/2606.11286
11机构：石溪大学，纽约州石溪，美国
11邮箱：[email protected]机构：多伦多大学，安大略省多伦多，加拿大33机构：大学健康网络，安大略省多伦多，加拿大秦仁11脚注标志1君马海滨凌陈宇通讯作者。

###### 摘要

高内涵成像实验可量化细胞对化学和遗传扰动的反应，但由于细胞在采集时已被化学固定，单个细胞的连续轨迹无法直接观测。因此，扰动建模被归结为推断仅在控制组和处理组群体（作为独立的边缘分布）之间发生的随机输运。虽然最近的生成模型能够实现强端点对齐，但边界一致性并不能确定中间演化过程：多个随机过程可能在连接相同边缘分布的同时，穿过未被观测单细胞形态所支持的区域。我们提出 FreeBridge，一种在仅有端点监督条件下进行单细胞转换建模的薛定谔桥公式。FreeBridge 将原子状态定义为实例分割后的单细胞表征，从而建立一个固定的细胞流形，并通过经验潜在支持正则化学习在此几何约束内的随机输运。在 BBBC021、RxRx1 和 JUMP 数据集上，统一评估协议下，FreeBridge 保持了具有竞争力或更优的端点保真度和作用机制保持能力；在 BBBC021 上，它还进一步降低了中间支持违规。这些发现强调了几何基础对于生物学可解释的扰动动力学的重要性。项目页面：https://y-research-sbu.github.io/FreeBridge/。

## 1 引言

构建一个能够在硅胶中模拟细胞反应的虚拟细胞，长久以来一直是定量细胞生物学的核心愿望 [19 (https://arxiv.org/html/2606.11286#bib.bib19)]。基于图像的细胞表型分析使得在化学和遗传扰动下大规模测量形态反应成为可能 [5 (https://arxiv.org/html/2606.11286#bib.bib5),15 (https://arxiv.org/html/2606.11286#bib.bib15)]。RxRx1 和 JUMP [23 (https://arxiv.org/html/2606.11286#bib.bib23),6 (https://arxiv.org/html/2606.11286#bib.bib6)] 等大型扰动图谱的出现，以及近期将人工智能整合到虚拟细胞构建中的努力 [4 (https://arxiv.org/html/2606.11286#bib.bib4)]，为直接从成像数据建模扰动动力学创造了机会。

参考图例图 1：FreeBridge 概述。(a) 多细胞显微图像被分割成原子单细胞状态，形成状态单元。每个裁剪区域被嵌入为 z=φ(x)z=\phi(x)，定义控制端和扰动端点分布。在此固定潜在几何上，FreeBridge 学习一个时间条件漂移 uθ(z,t,c)u_\theta(z,t,c) 来建模薛定谔桥输运，并通过经验支持代价 V(z)V(z) 正则化中间状态。(b) 在 BBBC021 上，在统一单细胞协议下，FreeBridge 实现了比 PhenDiff、IMPA 和 CellFlux 更低的总体和条件级别 FID 和 KID。最近的生成模型学习控制分布和扰动分布之间的映射。扩散模型 [9 (https://arxiv.org/html/2606.11286#bib.bib9),20 (https://arxiv.org/html/2606.11286#bib.bib20),17 (https://arxiv.org/html/2606.11286#bib.bib17),22 (https://arxiv.org/html/2606.11286#bib.bib22)]、流匹配 [12 (https://arxiv.org/html/2606.11286#bib.bib12),14 (https://arxiv.org/html/2606.11286#bib.bib14),24 (https://arxiv.org/html/2606.11286#bib.bib24)] 以及包括 PhenDiff、IMPA 和 CellFlux [3 (https://arxiv.org/html/2606.11286#bib.bib3),16 (https://arxiv.org/html/2606.11286#bib.bib16),25 (https://arxiv.org/html/2606.11286#bib.bib25)] 在内的形态预测器通过将扰动预测表述为分布到分布的变换，实现了高端点保真度。然而，高内涵实验测量的是化学固定的细胞；个体轨迹无法观测，只有边缘端点分布可用。无穷多个随机过程可以连接相同的边缘分布 [7 (https://arxiv.org/html/2606.11286#bib.bib7),11 (https://arxiv.org/html/2606.11286#bib.bib11)]，因此端点对齐并不能确定中间演化过程。

现有的大多数方法优化分布级别目标以匹配控制群体和扰动群体 [9 (https://arxiv.org/html/2606.11286#bib.bib9),12 (https://arxiv.org/html/2606.11286#bib.bib12),3 (https://arxiv.org/html/2606.11286#bib.bib3),25 (https://arxiv.org/html/2606.11286#bib.bib25),16 (https://arxiv.org/html/2606.11286#bib.bib16)]。虽然这对端点真实性有效，但这种表述使得中间细胞结构欠定。原子细胞状态的概念通常是隐含的，并且与群体级统计量纠缠在一起，因此中间样本可能满足 FID [8 (https://arxiv.org/html/2606.11286#bib.bib8)] 或 KID [2 (https://arxiv.org/html/2606.11286#bib.bib2)]，同时却穿过了被观测单细胞形态弱支持的区域。

因此，一个扰动模型需要对允许的细胞状态进行明确定义，并需要一个在该状态空间内受约束的转换规则。我们将此结构描述为细胞引擎，包含一个状态单元和一个转换机制。原子状态被实例化为实例分割后的单细胞表征 [21 (https://arxiv.org/html/2606.11286#bib.bib21)]，从而建立一个独立于群体构成的几何基础流形。在该流形上，通过受薛定谔桥启发的随机输运过程 [7 (https://arxiv.org/html/2606.11286#bib.bib7),11 (https://arxiv.org/html/2606.11286#bib.bib11)] 对扰动动力学进行建模，并通过基于非参数密度估计 [18 (https://arxiv.org/html/2606.11286#bib.bib18),1 (https://arxiv.org/html/2606.11286#bib.bib1)] 的经验潜在支持约束进行正则化。

基于此公式，我们开发了 FreeBridge，一个在仅有端点监督下进行单细胞转换建模的约束随机输运框架。FreeBridge 的概览如图 1 (https://arxiv.org/html/2606.11286#S1.F1) 所示。通过明确地将状态规范与随机输运分离，并将随机演化锚定在一个经验支持的细胞流形上，FreeBridge 在保持端点保真度的同时约束了中间轨迹。在 BBBC021 [15 (https://arxiv.org/html/2606.11286#bib.bib15)]、RxRx1 [23 (https://arxiv.org/html/2606.11286#bib.bib23)] 和 JUMP [6 (https://arxiv.org/html/2606.11286#bib.bib6)] 上的评估显示，与先前的分布对齐方法 [3 (https://arxiv.org/html/2606.11286#bib.bib3),16 (https://arxiv.org/html/2606.11286#bib.bib16),25 (https://arxiv.org/html/2606.11286#bib.bib25)] 相比，它取得了具有竞争力或更优的生成性能；在 BBBC021 上，FreeBridge 还进一步减少了中间支持违规。

## 2 方法

### 2.1 概述：FreeBridge 作为细胞引擎

高内涵筛选提供了在控制和扰动条件下单细胞的端点群体。由于成像需要化学固定，每个细胞的时间轨迹无法观测。因此，扰动建模被表述为一个端点约束的、在经验单细胞分布之间的随机输运问题，而不是保持身份的演化问题。

如图 1 (https://arxiv.org/html/2606.11286#S1.F1)(a) 所示，我们在细胞引擎公式下开发了 FreeBridge，它将细胞状态规范与随机输运分开。状态单元定义了有效单细胞表征的允许几何结构。转换规则在该固定几何内模拟熵正则化的随机输运。这种分离确保了动力学学习不会隐式地重新定义生物学有效细胞状态的基本概念。

### 2.2 状态单元：显式单细胞几何

状态单元指定了细胞状态的原子表示。我们使用 Cellpose [21 (https://arxiv.org/html/2606.11286#bib.bib21)] 提取实例分割后的单细胞裁剪区域，确保每个观测对应一个物理上可识别的细胞，而不是一个场级别的混合。每个裁剪区域 xx 通过一个固定编码器 φ:X→Rd\phi:\mathcal{X}\rightarrow\mathbb{R}^d 进行嵌入，产生潜在表示：z=φ(x)z=\phi(x)。编码器在输运学习期间被冻结，以便在学习转换规则之前定义潜在几何。设 μ\mu 表示控制单细胞图像的分布，νc\nu_c 表示在扰动条件 cc 下的分布。它们在 φ\phi 下的推前分布为

μz=φ♯μ,νc,z=φ♯νc。\mu_z=\phi_\sharp\mu,\qquad\nu_{c,z}=\phi_\sharp\nu_c。(1)为了表征经验支持的细胞流形，我们构建了一个潜在支持库

Zemp={φ(xi)}i=1N。\mathcal{Z}_{\mathrm{emp}}=\{\phi(x_i)\}_{i=1}^N。(2)靠近 Zemp\mathcal{Z}_{\mathrm{emp}} 的区域对应于真实数据中观察到的形态。

### 2.3 转换规则：薛定谔桥输运

给定由状态单元定义的固定几何，转换规则指定了 μz\mu_z 和 νc,z\nu_{c,z} 之间的随机输运。我们通过受控的随机微分方程对潜在动力学进行建模：

dZt=ut(Zt)dt+σdWt,Z0∼μz。dZ_t=u_t(Z_t)\,dt+\sigma\,dW_t,\qquad Z_0\sim\mu_z。(3)目标是找到一个漂移场 utu_t，使得：

Z1∼νc,z。Z_1\sim\nu_{c,z}。(4)在指定参考过程或路径准则之前，仅端点匹配是欠定的。一旦固定了布朗参考过程和边界边缘分布，经典的薛定谔桥就变得明确：

minPu KL(Pu∥P0) s.t. Z0∼μz,Z1∼νc,z，\min_{\mathbb{P}^u}\mathrm{KL}(\mathbb{P}^u\|\mathbb{P}^0)\quad\text{s.t. }Z_0\sim\mu_z,\;Z_1\sim\nu_{c,z},(5)这等价于随机控制目标 [7 (https://arxiv.org/html/2606.11286#bib.bib7)]

minu E[12σ2∫01‖ut(Zt)‖2dt]。\min_u\mathbb{E}\left[\frac{1}{2\sigma^2}\int_0^1\|u_t(Z_t)\|^2dt\right]。(6)FreeBridge 遵循这种端点约束的桥观点，但通过一个数据驱动的状态代价对其进行增强，该代价将中间路径锚定到经验细胞形态流形上，产生了具有任务特定状态代价的广义薛定谔桥公式 [13 (https://arxiv.org/html/2606.11286#bib.bib13)]。随机项 σdWt\sigma dW_t 捕获了内在的细胞异质性。确定性输运为每个初始状态产生单一轨迹，而随机公式则允许条件一致的变异性。

### 2.4 中间动力学的几何正则化

一旦固定了边界边缘分布，桥的中间边缘分布由所选的参考过程决定；因此，偏离纯扩散参考需要一个明确的路径准则。在细胞建模中，中间状态对应于沿扰动反应的假定形态，并且应与经验观察到的单细胞结构保持一致。因此，我们将随机输运锚定到由状态单元定义的几何上。对于任何潜在状态 zz，定义经验支持距离：

demp(z)=minzi∈Zemp ‖z−zi‖。d_{\mathrm{emp}}(z)=\min_{z_i\in\mathcal{Z}_{\mathrm{emp}}}\|z-z_i\|。(7)这个最近邻统计量 [1 (https://arxiv.org/html/2606.11286#bib.bib1)] 作为密度支持的非参数代理 [18 (https://arxiv.org/html/2606.11286#bib.bib18)]。远离 Zemp\mathcal{Z}_{\mathrm{emp}} 的区域对应于未被观测细胞支持的潜在配置。我们惩罚沿轨迹的偏离：

Lsupport=E[∫01demp(Zt)2dt]。\mathcal{L}_{\mathrm{support}}=\mathbb{E}\left[\int_0^1 d_{\mathrm{emp}}(Z_t)^2 dt\right]。(8)这种正则化使中间边缘分布偏向经验支持的区域，而无需施加样本级别的配对约束（图 2 (https://arxiv.org/html/2606.11286#S2.F2)）。

实际估计。直接对完整支持库评估 demp(z)d_{\mathrm{emp}}(z) 需要每个查询进行 O(N)O(N) 次比较。训练期间，我们使用当前小批量中控制和扰动嵌入的并集来近似总体支持，ZB={z0(i)}i=1B∪{z1(i)}i=1B\mathcal{Z}_B=\{z_0^{(i)}\}_{i=1}^B\cup\{z_1^{(i)}\}_{i=1}^B。对于每个中间状态 Zt(i)Z_t^{(i)}，我们计算

d^emp(Zt(i))=minzj∈ZB ‖Zt(i)−zj‖2。\widehat{d}_{\mathrm{emp}}(Z_t^{(i)})=\min_{z_j\in\mathcal{Z}_B}\|Z_t^{(i)}-z_j\|_2。(9)这产生了在独立同分布采样下总体状态代价项的一个随机小批量近似。成对距离通过批量矩阵运算（torch.cdist）计算，导致 O(B2)O(B^2) 的批量复杂度。最近邻索引被视为不可微的，而梯度通过所选平方距离传播。

薛定谔桥的内部动力学。薛定谔桥约束了边界边缘分布；给定一个固定的参考过程，内部边缘分布就被确定了，因此改变中间状态的经验形态需要一个明确的状态代价，而不是仅靠端点监督 [7 (https://arxiv.org/html/2606.11286#bib.bib7),11 (https://arxiv.org/html/2606.11286#bib.bib11)]。因此，端点对齐本身并不能固定数据流形上所需的转换几何。图 2 (https://arxiv.org/html/2606.11286#S2.F2) 在合成流形上展示了这种效应。即使边界分布一致，在扩散参考下的输运也可能穿过潜在空间中稀疏支持的区域。这样的轨迹保持了端点统计量，但却改变了内部结构。支持正则化将随机演化限制在经验细胞流形上。

### 2.5 FreeBridge 目标

FreeBridge 目标被定义为具有任务特定状态代价 Vt(z)=λdemp(z)2V_t(z)=\lambda\,d_{\mathrm{emp}}(z)^2 的广义薛定谔桥：

minu E[12σ2∫01‖ut(Zt)‖2dt+λ∫01demp(Zt)2dt] s.t. Z0∼μz,Z1∼νc,z。\min_u\;\mathbb{E}\!\left[\frac{1}{2\sigma^2}\int_0^1\|u_t(Z_t)\|^2dt+\lambda\int_0^1 d_{\mathrm{emp}}(Z_t)^2 dt\right]\;\;\text{s.t.}\;\;Z_0\sim\mu_z,\;Z_1\sim\nu_{c,z}。(10)用于端点监督的批内配对。端点群体是未配对的，但训练需要批量级别的元组。对于每个条件 cc，我们独立采样 {z0(i)}i=1B∼μz\{z_0^{(i)}\}_{i=1}^B\sim\mu_z 和 {z1(i)}i=1B∼νc,z\{z_1^{(i)}\}_{i=1}^B\sim\nu_{c,z}，并在随机打乱后按索引配对。这种随机耦合仅用于优化，并不强加身份级别的对应关系或约束中间动力学。

### 2.6 数值近似与优化

我们将漂移参数化为 uθ(z,t,c)u_\theta(z,t,c)，使用一个时间条件神经网络。连续过程通过欧拉-丸山方法离散化：

Zt+Δt