FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

技术摘要：FreeBridge —— 用于细胞转变动态的变分薛定谔桥

问题陈述
高内涵成像分析对于量化细胞对化学和遗传扰动的反应至关重要。然而，存在一个根本性的限制：细胞在图像采集时处于化学固定状态，导致单个时间轨迹是不可观测的。因此，扰动建模被简化为推断控制组与处理组之间的随机传输，而此时仅已知边缘终点分布。

虽然近期的生成模型（如扩散模型、流匹配）在控制端点与扰动端点之间实现了强大的对齐，但仅靠端点一致性不足以确定中间演化过程。无数种随机过程都可以连接相同的边缘分布；若缺乏进一步约束，这些模型可能会遍历由观察到的单细胞形态所不支持的潜空间区域，从而导致生物学上不合理的中间状态。现有方法通常优化分布层面的目标函数，这使得中间细胞结构处于欠定状态，可能满足 FID 或 KID 等指标，却产生伪影。

方法论：FreeBridge 框架
作者提出了 FreeBridge，这是一个受约束的随机传输框架，被构造成一个“细胞引擎”（Cell Engine），它明确地将有效细胞状态的定义与转换机制分离。该方法在仅有端点监督的情况下运行，由三个核心组件组成：

1. 状态单元（显式单细胞几何结构）：

   • FreeBridge 不将图像视为群体层面的混合物，而是将原子状态定义为实例分割的单细胞表示（使用 Cellpose）。
   • 一个固定的预训练编码器 $\phi$ 将这些裁剪图像嵌入到潜空间 $z = \phi(x)$ 中。该编码器在传输学习期间保持冻结，从而建立了一个在学习转换规则之前就已确定的固定潜几何结构。
   • 通过观察到的细胞的潜嵌入构建一个经验支持库 $Z_{emp}$。该库定义了生物学有效形态的“流形”。

2. 转换规则（薛定谔桥传输）：

   • 控制组分布 $\mu_z$ 与扰动分布 $\nu_{c,z}$ 之间的转换被建模为一个受控随机微分方程 (SDE)：$dZ_t = u_t(Z_t) dt + \sigma dW_t$。
   • 目标是寻找一个漂移场 $u_t$，将初始分布传输到 $t=1$ 时的目标分布。
   • 与标准的最佳传输不同，该公式引入了特定任务的状态代价来正则化路径。

3. 几何正则化（经验支持约束）：

   • 为了防止随机过程遍历潜空间中的“真空区”，FreeBridge 引入了一个支持代价 $V(z) = \lambda \cdot d_{emp}(z)^2$，其中 $d_{emp}(z)$ 是到经验支持库 $Z_{emp}$ 最近邻的距离。
   • 总目标函数最小化动能（漂移幅度）与支持代价之和：
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • 这种正则化将中间状态锚定在经验观察到的形态流形上，而无需配对的细胞级轨迹。

核心贡献

• 公式化： 引入了一种用于单细胞转换建模的薛定谔桥公式，明确地将状态规范与随机传输解耦。
• 几何落地： 提出了一种基于非参数密度估计（到经验支持库的最近邻距离）的新型正则化项，用以将中间轨迹约束在生物学合理的区域内。
• 统一协议： 通过使用实例分割的单细胞裁剪图像，在多个数据集（BBBC021, RxRx1, JUMP）上建立了统一的评估协议，确保了方法间的公平比较。
• 性能： 证明了通过约束几何结构，不仅提高了端点保真度，还提升了语义保留能力（作用机制，Mechanism-of-Action）。

结果
在 BBBC021、RxRx1 和 JUMP 数据集上的评估显示，FreeBridge 表现出以下特性：

• 端点保真度： 与最先进的基准模型（PhenDiff, IMPA, CellFlux）相比，实现了具有竞争力的或更优的 Fréchet Inception Distance (FID) 和 Kernel Inception Distance (KID) 分数。在 BBBC021 上，它实现了最低的 FID (15.35) 和 KID (1.10)。
• 中间一致性： 显著降低了支持违规率 ($R_{viol}$)，该指标衡量了落在经验形态支持之外的中间状态比例。在 BBBC021 上，Free-Bridge 将 $R_{viol}$ 降低至 0.11，而 PhenDiff 为 0.38，CellFlux 为 0.31。
• 语义保留： 消融实验表明，移除几何正则化会降低作用机制 (MoA) 的分类准确率（在 BBBC021 上从 72.08% 下降到 65.08%），这表明约束中间演化有助于保留机械信号。
• 定性分析： 可视化证实，与经常出现纹理平滑或模式不完整的基准模型相比，FreeBridge 能更好地保留化合物特有的形态特征（例如：核紧缩、胞质收缩）。

意义与主张
论文声称，建模生物学转变不仅需要边界对齐；它还需要中间演化的显式几何落地。通过将随机传输锚定在经验支持的生物流形上，FreeBridge 提供了一种原则性的方法来生成具有生物学解释性的扰动动态。作者断言，这种公式化为“虚拟细胞”模拟提供了稳健的基础，解决了端点对齐模型无法捕捉真实中间细胞状态的关键缺陷。这项工作强调，如果没有几何约束，生成模型可能会产生统计上有效的端点，但其遍历的却是生物学上不支持的区域。

局限性
作者承认 FreeBridge 依赖于冻结的编码器；编码器预训练中缺失的形态学因素无法仅通过传输学习来恢复。此外，支持距离的计算随小批量大小呈二次方 ($O(B^2)$) 缩放，对于较大的支持库，可能需要近似最近邻搜索。