CellFluxRL: Biologically-Constrained Virtual Cell Modeling via Reinforcement Learning¶

会议: ECCV2026
论文: ECCV 2026 Poster
领域: 计算生物学
关键词: 虚拟细胞, 流匹配, 强化学习后训练, 生物奖励模型, 细胞形态生成

一句话总结¶

CellFluxRL 不改动生成架构，而是把七项可自动计算的生物学评估器（MoA 一致性、核是否被胞质包含、核圆度、核与胞质的尺寸/数量统计）当作奖励，对已训练好的流匹配虚拟细胞模型 CellFlux 做 RL 后训练，把「看起来逼真」的细胞图像推向「生物学上可信」，七项奖励全部优于基线，再用 best-of-N 测试时扩展进一步放大收益。

研究背景与动机¶

湿实验一直是药物发现的主要瓶颈：做一次实验要买试剂耗材、要等数周甚至数月来合成药物并培养细胞。因此「虚拟细胞」——在计算机里模拟细胞对扰动的响应——被寄予厚望。随着扩散模型、流匹配这类生成模型与高通量成像筛选技术同时成熟，图像式虚拟细胞成为一个具体可行的形态：给定一张未受扰动的对照细胞图像和一个化合物扰动，直接生成扰动之后的细胞图像。CellFlux 是这条路线上的当前最优模型，它把问题重述为「分布到分布」的变换，用流匹配在同一批次内把对照分布输运到扰动分布，从而绕开批次效应。

但作者观察到一个系统性的失败：这些模型生成的图像往往局部看起来非常真实，整体上却违反基本的物理与生物学约束——最典型的例子是细胞核被生成到了细胞质之外。这类样本在像素层面无可指摘，在生物学上却完全不成立，直接摧毁了虚拟细胞在下游药物筛选里的可用性。问题的根源在于目标函数的层级错配：流匹配学的是速度场，它的损失作用在像素与分布层面，只保证生成的样本离真实扰动的分布「整体上近」，而「核必须在胞质内」「核的大小要符合该类药物的作用机制」这类约束是全局的、结构性的，像素级目标既看不见也没有理由去满足它们。

更麻烦的是，这类约束恰恰天然非可微：判断 MoA 是否一致要用一个分类器、判断核是否被胞质包含要做分割、判断尺寸是否合理要算群体统计量，全都无法直接反传。本文的切入角度正是利用这一点——这些约束虽然不可微，但每一条都已经有现成或易得的评估器可以给生成图打分，而「用不可微的评估器当奖励去优化生成模型」正是强化学习擅长的事。核心 idea：把领域先验写成一组可自动计算的生物奖励，用 RL 后训练流匹配虚拟细胞模型（正/负隐式策略对比更新 + KL 正则防奖励劫持），并让同一套奖励在推理时充当验证器做 best-of-N 选择。

关于名字：「Flux」继承自基础模型 CellFlux，指的是流匹配中学到的流（速度场）——把对照细胞分布输运成扰动后分布的连续变换，与物理学里的「通量」无关。

方法详解¶

整体框架¶

方法的输入是同一批次内的一张对照细胞图像 $x_0$（三通道荧光显微图，分别标记细胞核、细胞骨架、线粒体等结构）与一个化合物扰动 $c$，输出是预测的「处理后」图像 $\hat x_1$。基础模型 CellFlux 学一个速度场 $v_\theta$，在第 $t$ 步沿线性插值路径把对照分布 $p_0$ 连续输运到扰动分布 $p_1(\cdot\mid c)$；由于显微成像是破坏性的（细胞被固定染色），同一细胞处理前后的配对数据不存在，只能拿到非配对的对照组与扰动组图像集合，所以模型必须做分布级变换而非逐点映射，同时必须以同批次的对照图像为起点，才能只学扰动引起的变化而不学批次间的技术差异。

本文不训练新的生成架构，而是在这个预训练速度场之上叠加一层 RL 后训练，整个流程是一个「采样—评估—对比更新」的闭环：每一步固定一张 $x_0$ 和 $c$，从当前策略采样一组候选图；七个奖励函数逐个给候选打分并加权求和；分数在组内归一化成相对优势，决定每个候选是正例还是负例；再用同一份样本构造正、负两个隐式速度场做对比式更新，同时用 KL 项把新模型拴在预训练模型附近。推理阶段还可以复用同一套奖励，对同一条件采 $N$ 张图挑奖励最高的一张。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["对照细胞图像 x0 + 扰动条件 c"] --> B["生物奖励套件<br/>七项约束评估器"]
    A --> C["组内 rollout<br/>与相对优势归一化"]
    B --> C
    C --> D["对比式 RL 更新<br/>正负隐式策略 + KL"]
    D -->|滞后副本 EMA| C
    D --> E["奖励驱动的测试时扩展<br/>best-of-N 挑选"]
    E --> F["生物可信的扰动后细胞图像"]

关键设计¶

1. 生物奖励套件：把领域知识写成七项可自动计算的评估器

这是全文最核心的贡献。作者没有去训一个「偏好模型」，而是从三个尺度上把细胞生物学的既有知识直接翻译成可计算的奖励函数，共七项、分三类。生物功能层面只有一项但最重要：MoA（作用机制）一致性。药物按其靶向的细胞过程分类（微管解聚、DNA 复制抑制、肌动蛋白破坏……），不同 MoA 会产生截然不同的形态签名，所以「生成图能否被正确分类到真实 MoA 类」就是生物保真度的强指标。奖励直接取一个在真实扰动图上预训练好的 MoA 分类器对真值类的预测概率，$r_{\mathrm{MoA}}(\hat x_1,c)=p_{\mathrm{cls}}(y_c\mid \hat x_1)$。结构有效性层面有两项：其一是「核在胞质内」，用 Cellpose 从生成图里分割出核与胞质掩膜，奖励取所有核与胞质掩膜交集面积占胞质面积的比例，直接惩罚「各通道各画各的、空间上对不上」这类崩坏；其二是「核圆度」，因为核形状本身随 MoA 系统变化（微管解聚剂会产生碎裂、不规则的核，另一些扰动则保持光滑圆润），所以先按 $4\pi\cdot\text{area}/\text{perimeter}^2$ 算出生成图里每个核的圆度并取均值，再与该 MoA 类真实分布的均值 $\mu^{(y_c)}$ 比较、用其标准差 $[\sigma^{(y_c)}]^2$ 缩放，取负的归一化偏差。形态统计层面有四项：最大核尺寸、最大胞质尺寸、核数量、胞质数量，每一项都算成「生成图统计值与该 MoA 真实分布的均值之差，除以该类的标准差，再取负」的形式，即

\[r_{\text{morphological},s}(\hat x_1,c)=-\frac{\bigl[s(\hat x_1)-\mu_s^{(y_c)}\bigr]}{\sigma_s^{(y_c)}},\quad s\in\{\text{NucSize},\text{CytoSize},\text{NucCount},\text{CytoCount}\}\]

用最大尺寸而不是平均尺寸，是为了避免画面边缘被截断的细胞把统计量拉偏；核/胞质数量则约束细胞密度，防止模型凭空多造或漏掉细胞。

这样设计的好处是奖励覆盖了从高层生物语义到低层几何统计的三个尺度，全是像素级流匹配损失看不见、却由真实扰动分布必然满足的性质——换句话说，这些奖励表达的是「真值分布本来就满足什么」，因此优化它们与原始的分布匹配目标并不冲突，只是把模型容量集中到对物理正确性最关键的那些维度上。而且七项奖励全部复用已有工具（预训练分类器、现成分割模型、统计量），无需新增标注。

2. 组内 rollout 与相对优势归一化：把绝对分数变成「同组里谁更好」

七个奖励量纲和尺度完全不同：MoA 落在 $[0,1]$，形态统计是负的归一化偏差，正负号与量级都不可比。直接拿加权和当梯度信号既不稳定也无意义，本文的做法是在组内做相对比较。每一轮固定同一个 $(x_0,c)$，从当前策略采样 $m$ 张候选构成一个组；候选之间的多样性主要来自基础模型在输运前注入的高斯噪声以及 ODE 离散化误差，作者发现这点变异对区分正负样本已经足够。每张候选先算加权总奖励 $r(\hat x_1,c)=\sum_k w_k r_k(\hat x_1,c)$，再减去组内均值、按常数 $Z_c$ 归一化并裁剪到 $[0,1]$，得到「最优性奖励」$r_{\mathrm{adv}}$（原文 Eq. 11，⚠️ 缓存中该式的裁剪上下界与 $Z_c$ 的定义有排版损失，精确形式以原文为准）。

$r_{\mathrm{adv}}$ 同时承担两个角色：它既是一个 0—1 的连续权重，说明「这张样本离组内最优有多近」，也因此决定这张样本在接下来被当作正例还是负例、以及正例的力度有多大。用组内均值而非全局基线做参照，是因为不同条件下可达的奖励水平差别很大（细胞密度不同、MoA 不同），组内比较才能把「这个条件下什么算好」分离出来。

3. 对比式 RL 更新：用正/负隐式策略绕开不可解的似然

奖励全是非可微的，只能走 RL；但扩散与流匹配模型的样本似然不可解，无法像语言模型那样直接按奖励高低缩放某条轨迹的概率。本文采用 DiffusionNFT 的思路：不估计似然，而是在流的前向过程上构造一对隐式策略，让一个模型同时扮演正负两个角色——

\[v^{+}_{\theta}(x_t,c,t):=(1-\gamma)\,v^{\mathrm{old}}(x_t,c,t)+\gamma\,v_{\theta}(x_t,c,t)$$ $$v^{-}_{\theta}(x_t,c,t):=\gamma\,v^{\mathrm{old}}(x_t,c,t)-\gamma\,v_{\theta}(x_t,c,t)\]

其中 $v^{\mathrm{old}}$ 是数据采样策略（对 $v_\theta$ 做指数滑动平均得到的滞后副本），$\gamma$ 控制引导强度。训练时把 $r_{\mathrm{adv}}$ 当权重，让正速度项去拟合高奖励样本的速度目标、负速度项去拟合低奖励样本（等价于远离它们），再补一项 $\beta$ 倍 KL 散度把后训练模型拴在预训练策略附近。

这套构造的两个关键点值得强调。第一，正向策略是旧策略与当前策略的凸组合，因此隐含地实现了「既偏离预训练模型、又不过度偏离」的正则：$\gamma$ 越大引导越弱、模型越贴 $v^{\mathrm{old}}$，$\gamma$ 越小则允许改动越激进。第二，整套推导只用流的前向过程，不假设起点是标准高斯先验——这正好适配本文「同批次对照分布 → 扰动分布」的分布到分布设定，而 FlowGRPO、DanceGRPO 这类把生成过程 MDP 化的 GRPO 方法默认从高斯先验出发，在这里无法直接套用。KL 项在这里不只是防过拟合，更是防奖励劫持：没有它，模型完全可能学会生成极端但能骗过评估器的图像（比如把核画得巨大以博取某个统计奖励），而不是真的变正确。

4. 奖励驱动的测试时扩展：同一套评估器当验证器用

既然奖励函数已经能自动给生成图打分，它们就不必只用于训练。推理时给定同一个 $(x_0,c)$ 采样 $N$ 张候选，直接选加权总奖励最高的那张输出（best-of-N），不需要任何额外训练，只消耗推理算力。这与推理模型里用验证器做 best-of-N 的做法同源：作者观察到 RL 与测试时扩展是互补而非替代关系——RL 提升的是「候选池」本身的质量分布，所以在 $N$ 较小时收益最明显（基础分布已经不错，随手挑一张就够好），而 $N$ 增大时两条曲线都单调上升、RL 曲线始终位于基础模型上方，意味着在任何固定算力预算下 RL 后训练模型都更划算。

损失函数 / 训练策略¶

训练目标是最大化七项奖励的加权和，同时约束后训练模型与预训练模型的 KL 散度不要超过阈值，即 $\max_{\theta'}\mathbb{E}_{x_0\sim p_0,\,c\sim p(c)}\bigl[\sum_{k}w_k r_k(\hat x_1,c)\bigr]$ 受 KL 约束（原文 Eq. 2；⚠️ 约束的具体形式与阈值定义在缓存排版中不完整）。实现上采用 DiffusionNFT 在线 RL 算法，v_old 按指数滑动平均跟随 v_θ。

奖励权重为 $r=5.0\,r_{\mathrm{MoA}}+2.0\,r_{\mathrm{Nuc\text{-}in\text{-}Cyto}}+r_{\mathrm{Roundness}}+r_{\mathrm{NucSize}}+r_{\mathrm{CytoSize}}+r_{\mathrm{NucCount}}+r_{\mathrm{CytoCount}}$：MoA 与「核在胞质内」经验上更难优化，所以给了更高权重。作者也试过对所有归一化奖励统一取权重 1.0（不需要任何先验知识来调权），同样优于基础模型。KL 权重 $\beta=1$，其余超参沿用 DiffusionNFT。训练 1200 步，单张 H100 跑 32 小时；推理时 best-of-N，$N=4$。

实验关键数据¶

主实验¶

数据集为 BBBC021——化学扰动高通量显微成像数据集，含 98K 张三通道 96×96 图像，覆盖 26 种化合物、归为 12 类 MoA，训练/测试划分沿用 CellFlux。对比对象包括生成式基线 PhenDiff、IMPA 与基础模型 CellFlux。奖励类指标越高越好，FID/KID 越低越好。

指标	PhenDiff	IMPA	CellFlux	CellFluxRL	+TTS (N=4)
MoA	0.18	0.12	0.26	0.34	0.56
Nuc-in-Cyto	0.91	0.79	0.88	0.96	0.97
Roundness	-0.24	-0.32	-0.34	-0.26	-0.19
NucSize	-0.88	-1.02	-2.21	-1.04	-0.38
CytoSize	-0.97	-1.59	-1.09	-0.65	-0.41
NucCount	-2.61	-1.05	-0.83	-0.53	-0.28
CytoCount	-3.22	-1.39	-1.03	-0.68	-0.33
Overall	-5.20	-3.19	-2.44	0.46	3.15
FID ↓	41.94	35.70	20.36	24.01	23.19
KID ↓	0.028	0.029	0.015	0.014	0.016

消融实验¶

表下为「只优化单个奖励」与「七项奖励联合优化」（CellFluxRL）的对比，每列是一个只针对该奖励做 RL 后训练的模型。

评测指标	MoA-only	Roundness-only	NucSize-only	CellFluxRL（全部）
MoA	0.412	0.263	0.301	0.337
Nuc-in-Cyto	0.852	0.907	0.904	0.956
Roundness	-0.342	-0.198	-0.298	-0.264
NucSize	-3.319	-2.328	-0.702	-1.038
CytoSize	-1.198	-0.919	-0.917	-0.652
NucCount	-0.760	-0.685	-0.818	-0.528
CytoCount	-0.922	-0.905	-0.946	-0.677

关键发现¶

RL 后训练在所有七项奖励上全面超过基础模型，总奖励从 -2.44 提升到 0.46，增益横跨生物功能、结构约束、形态统计三类，说明「像素级目标看不见全局约束」这一诊断是对的，且 RL 确实能补上。
MoA 是增益最显著的一项：0.26 → 0.34，对应的分类准确率从 0.61 升到 0.66（PhenDiff 0.56、IMPA 0.53）。换句话说，RL 让生成图更容易被正确识别为「该类药应有的形态」。
单奖励消融揭示了组合奖励的价值：只优化 MoA 能得到更高的 MoA 分数（0.412 > 0.337），但 NucSize 崩到 -3.319；只优化 NucSize 能把 NucSize 做到 -0.702，其他项普遍变差。而联合优化的 CellFluxRL 在每一项上都排到第二或第三（除在个别项上与专门优化的模型有差距），说明多个奖励叠加时能实现跨维度均衡而不是此消彼长。
图像质量基本保持但不是主要目标：CellFluxRL 的 KID 0.014 为所有方法最优，但 FID 24.01 反而比基础模型 CellFlux 的 20.36 差；加上 TTS 后 FID 回到 23.19。作者也明确表示 FID/KID 不是优化目标，本文要的是生物正确性。
Roundness 一项上 PhenDiff（-0.24）略优于 CellFluxRL（-0.26），但 +TTS（-0.19）实现反超——这是表中唯一一处 CellFluxRL 未取得领先（含并列）的单项，提醒读者「全面超越」的表述对基础模型成立，对全部基线则有个别例外。
KL 权重 $\beta$ 的敏感性呈权衡：把 $\beta$ 从 1.0 调到 1.3，结构与形态类指标偏好更小的 KL 权重，而生物功能类奖励随 $\beta$ 略升。作者据此推断——要提高 MoA 层面的生物准确度需要更实质的模型改动，而结构/形态的修正更容易实现。实际使用中 $\beta$ 就是这个权衡的旋钮，默认取 1.0 偏重结构与形态，若以 MoA 一致性为优先则应调大。
测试时扩展随 $N$ 单调改善，且 RL 后训练显著提升了扩展效率：RL 曲线在任何 $N$ 下都位于基础模型之上，小 $N$ 处优势最大——此时基础分布的质量最能决定结果。

亮点与洞察¶

一套评估器三用：同一组生物奖励同时充当评测指标、训练奖励和推理时的验证器。这个「一物三用」的抽象是本文最漂亮的地方，它把 RLHF / 推理模型里已经成熟的「奖励模型 + best-of-N」范式干净地迁移到了科学生成任务上，而且不需要人类偏好标注。
用领域知识替代人类偏好：奖励不是训出来的，而是由现成工具（预训练 MoA 分类器、Cellpose 分割）与统计量组成。这规避了奖励模型训练成本与偏好数据采集，也让奖励本身可解释、可审计——每一项奖励对应一条明确的生物学陈述。
把「评估器的不可微性」从障碍转成了方法选择的理由：因为约束不可微，所以只能 RL，而 RL 恰好能利用这些现成评估器。这条逻辑反过来也解释了为什么不做架构修改（如注入物理约束层）：那样要为每条约束设计可微代理，成本高且容易失真。
算法与设定匹配的论证很清楚：作者没有盲目套用 GRPO 系列，而是指出 MDP 化方法默认高斯先验、与「从同批次对照图出发」的分布到分布设定不兼容，因此选前向过程式的 DiffusionNFT。这种「先讲清设定约束、再选算法」的论证方式值得借鉴。
可迁移思路：任何「有非可微规则/仿真器可当裁判」的科学生成任务（材料结构生成、医学影像合成、分子构象采样）都能照搬这套「奖励做指标 + 训练 + 验证器」的三段式，尤其是那些真值分布天然满足某些统计不变量（质量守恒、几何包含关系、群体统计）的场景——把这些不变量写成负的归一化偏差，就是现成的奖励。

局限与展望¶

作者承认的核心局限是奖励全靠人工设计，依赖领域专家把生物学知识翻译成可计算形式，新领域门槛高；他们提出的方向是用 LLM 从科学文献自动生成可执行的奖励函数。
只在一个数据集上验证：BBBC021 规模不大（98K 图、96×96、26 个化合物、12 类 MoA），MoA 类别少且类别间形态差异本就明显，扩展到 JUMP / RxRx 这类百万级、上千扰动的数据集时奖励是否仍然可分、RL 是否仍然稳定，论文没有回答。
奖励模型本身可能引入偏差：MoA 分类器就是在 BBBC021 上预训练的，用它的预测概率当奖励相当于让生成模型去迎合该分类器的决策边界；分类器的错误或捷径会被 RL 放大（奖励劫持的另一种形式），而论文只用 KL 项做了间接防护，没有做「换一个分类器奖励是否仍有效」的检验。
成本不低：1200 步训练需要单张 H100 跑 32 小时，推理时 best-of-N 还要 $N$ 倍的采样开销；论文只给了 $N=4$ 的定量结果，$N$ 更大时的曲线只在图里给了趋势，没有报告收益何时饱和。
建模范围止于形态图像，不涉及基因表达、转录组或信号通路——与「虚拟细胞」的完整愿景（多组学、可解释机制）仍有距离，「biology-constrained」在这个意义上指的是形态学层面的约束。
改进方向：把部分奖励做可微松弛（如可微分割 / 软 IoU）以降低 RL 的方差与采样成本；用多分类器集成或对抗式奖励审计来抑制奖励劫持；把奖励扩展到时间序列（活细胞成像）而非固定染色终点。

评分¶

新颖性: ⭐⭐⭐⭐ 首次系统性地把「可计算生物奖励 + RL 后训练 + 奖励做验证器」这套范式落到图像式虚拟细胞上，奖励设计与三类七项的划分有实质领域洞察；但核心 RL 算法沿用 DiffusionNFT，贡献主要在问题诊断与领域落地。
实验充分度: ⭐⭐⭐ 主表、单奖励消融、KL 敏感性、TTS 曲线都做了，且对比了三个生成基线；但只在 BBBC021 一个数据集上验证，缺少跨数据集与跨基础模型的泛化实验，TTS 也只报告了 $N=4$。
写作质量: ⭐⭐⭐⭐ 从失败案例（核跑出胞质）到「像素级目标看不见全局约束」的诊断再到奖励设计，逻辑链条非常清晰，算法与设定的匹配也讲得明白；部分公式在排版上有损失。
价值: ⭐⭐⭐⭐ 提出的「同一套领域评估器既做指标、又做训练奖励、又做推理验证器」对科学生成模型有普适参考价值，且给社区提供了 FID/KID 之外的生物学评测维度。