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

Conference: ECCV2026
Paper: ECCV 2026 Poster
Area: Computational Biology
Keywords: virtual cell, flow matching, reinforcement learning post-training, biological reward model, cell morphology generation

TL;DR¶

Without touching the generative architecture, CellFluxRL treats seven automatically computable biological evaluators (MoA consistency, nucleus-in-cytoplasm containment, nuclear roundness, and size/count statistics of nuclei and cytoplasm) as rewards, and post-trains the flow-matching virtual cell model CellFlux with RL to push "visually realistic" cell images toward "biologically meaningful" ones; all seven rewards improve over the baselines, and best-of-N test-time scaling amplifies the gains further.

Background & Motivation¶

Wet-lab experiments remain the primary bottleneck in drug discovery: a single experiment requires expensive reagents and consumables, and weeks or even months to synthesize compounds and culture cells. "Virtual cells" — simulating cellular responses to perturbations in silico — are therefore highly desirable. With diffusion models and flow matching maturing alongside high-throughput imaging screens, an image-based virtual cell became concrete: given an unperturbed control cell image and a chemical perturbation, generate the image of the cell after treatment. CellFlux is the state of the art on this route; it reformulates the problem as a distribution-to-distribution transformation and uses flow matching to transport the control distribution to the perturbed distribution within each batch, sidestepping batch effects.

Yet the authors observe a systematic failure: these models generate images that look locally realistic but globally violate basic physical and biological constraints — most strikingly, a cell nucleus generated outside the cytoplasm. Such samples are blameless at the pixel level and biologically impossible, which destroys the usefulness of virtual cells for downstream drug screening. The root cause is a mismatch of levels: flow matching learns a velocity field, its loss acts at the pixel and distribution level, and it only guarantees that generated samples are close to the true perturbed distribution on the whole. Constraints such as "the nucleus must lie inside the cytoplasm" or "nuclear size must match this compound's mode of action" are global and structural; a pixel-level objective neither sees them nor has any reason to satisfy them.

What makes this harder is that such constraints are inherently non-differentiable: checking MoA consistency needs a classifier, checking nuclear containment needs segmentation, and checking size plausibility needs population statistics — none of them can be back-propagated through. The paper's angle turns exactly this into an opportunity. Every such constraint already has an existing or readily obtained evaluator that can score a generated image, and "optimizing a generative model with a non-differentiable evaluator as the reward" is precisely what reinforcement learning is good at. Core idea: express domain priors as a set of automatically computable biological rewards, post-train the flow-matching virtual cell model with RL (contrastive positive/negative implicit-policy updates plus a KL regularizer against reward hacking), and reuse the very same rewards as verifiers for best-of-N selection at inference time.

A note on the name: "Flux" is inherited from the base model CellFlux and refers to the flow (velocity field) learned by flow matching — the continuous transformation that transports the control cell distribution into the perturbed one — not to physical flux.

Method¶

Overall Architecture¶

The input is a control cell image $x_0$ from a given batch (three-channel fluorescence microscopy, labelling structures such as the nucleus, cytoskeleton, and mitochondria) together with a chemical perturbation $c$; the output is the predicted "post-treatment" image $\hat x_1$. The base model CellFlux learns a velocity field $v_\theta$ that continuously transports the control distribution $p_0$ to the perturbed distribution $p_1(\cdot\mid c)$ along a linear interpolation path. Because microscopy is destructive (cells are fixed and stained), paired before/after observations of the same cell do not exist; only unpaired sets of control and perturbed images are available. The model must therefore learn a distributional transformation rather than a pointwise mapping, and must start from same-batch control images so that it learns only the perturbation-induced change rather than technical batch variation.

The paper does not train a new generative architecture; it adds an RL post-training layer on top of this pretrained velocity field. The whole procedure is a "sample — evaluate — contrastive update" loop: at each step a control image $x_0$ and a condition $c$ are fixed and a group of candidates is drawn from the current policy; seven reward functions score each candidate and are aggregated into a weighted sum; the scores are normalized within the group into a relative advantage that decides whether a candidate acts as a positive or a negative example; the same samples are then used to build positive and negative implicit velocity fields for a contrastive update, with a KL term tethering the new model to the pretrained one. At inference, the same rewards can be reused to draw $N$ images for a condition and keep the highest-reward one.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["control cell image x0 + perturbation c"] --> B["biological reward suite<br/>seven constraint evaluators"]
    A --> C["group rollout and<br/>relative advantage normalization"]
    B --> C
    C --> D["contrastive RL update<br/>positive/negative implicit policies + KL"]
    D -->|lagged EMA copy| C
    D --> E["reward-driven test-time scaling<br/>best-of-N selection"]
    E --> F["biologically credible perturbed cell image"]

Key Designs¶

1. Biological reward suite: turning domain knowledge into seven automatically computable evaluators

This is the paper's central contribution. Rather than training a "preference model", the authors translate existing cell-biology knowledge into computable reward functions at three scales, seven terms in three categories. Biological function has a single but most important term: MoA (mode of action) consistency. Drugs are grouped by the cellular process they target (microtubule destabilization, DNA replication inhibition, actin disruption, …), and different MoAs leave distinct morphological signatures, so whether a generated image is correctly classified into its true MoA is a strong indicator of biological fidelity. The reward is simply the probability that a classifier pretrained on real perturbed images assigns to the ground-truth class, $r_{\mathrm{MoA}}(\hat x_1,c)=p_{\mathrm{cls}}(y_c\mid \hat x_1)$. Structural validity has two terms. The first is nucleus-in-cytoplasm: Cellpose segments nuclei and cytoplasm from the generated image, and the reward is the fraction of cytoplasm-mask area covered by the intersection with nucleus masks, directly penalizing the failure mode where channels are generated independently and become spatially inconsistent. The second is nuclear roundness, since nuclear shape varies systematically across MoAs (microtubule destabilizers produce fragmented, irregular nuclei while other perturbations keep nuclei smooth and round): the roundness $4\pi\cdot\text{area}/\text{perimeter}^2$ of every nucleus in the generated image is averaged, compared against the mean $\mu^{(y_c)}$ of the real MoA-conditioned distribution, scaled by its variance $[\sigma^{(y_c)}]^2$, and turned into a negative normalized deviation. Morphological statistics contribute four terms — maximum nucleus size, maximum cytoplasm size, nucleus count, and cytoplasm count — each computed as the generated statistic minus the MoA-conditioned ground-truth mean, divided by that class's standard deviation, and negated:

\[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}\}\]

The maximum rather than the average size is used so that cells truncated at the image border do not bias the statistic, while the two count terms constrain cell density and stop the model from hallucinating or dropping cells.

The benefit of this design is that the rewards span three scales, from high-level biological semantics down to low-level geometric statistics, and every one of them is a property that the true perturbed distribution necessarily satisfies but that the pixel-level flow-matching loss cannot see. In other words, optimizing these rewards does not conflict with the original distributional objective; it concentrates model capacity on the axes that matter most for physical correctness. All seven rewards reuse existing tools (a pretrained classifier, an off-the-shelf segmentation model, simple statistics), so no new annotation is needed.

2. Group rollout and relative advantage normalization: turning absolute scores into "who is better within this group"

The seven rewards live on completely different scales: MoA lies in $[0,1]$ while the morphological statistics are negative normalized deviations, so neither sign nor magnitude is comparable. Using their weighted sum directly as a gradient signal would be unstable and meaningless; instead the paper compares within a group. Each iteration fixes the same $(x_0,c)$ and samples $m$ candidates from the current policy to form a group. Diversity within the group comes mainly from the Gaussian noise the base model injects before transport plus ODE discretization error, and the authors find this variation sufficient for distinguishing positive from negative generations. Each candidate first receives the weighted total reward $r(\hat x_1,c)=\sum_k w_k r_k(\hat x_1,c)$; the group mean is subtracted, and the result is normalized by a constant $Z_c$ and clipped to $[0,1]$ to yield the optimality reward $r_{\mathrm{adv}}$ (Eq. 11 in the paper; ⚠️ the clipping bounds and the definition of $Z_c$ are corrupted by typesetting in the cached text — refer to the original paper for the exact form).

$r_{\mathrm{adv}}$ plays two roles at once: it is a continuous 0–1 weight expressing "how close this sample is to the best in its group", and it therefore decides whether the sample is treated as a positive or a negative example and how strongly. Using the group mean rather than a global baseline as the reference matters because attainable reward levels differ greatly across conditions (different cell densities, different MoAs); only a within-group comparison isolates "what counts as good under this condition".

3. Contrastive RL update: positive and negative implicit policies to bypass intractable likelihoods

All rewards are non-differentiable, so RL is the only option; but sample likelihoods in diffusion and flow-matching models are intractable, so one cannot rescale a trajectory's probability by its reward the way language-model RL does. The paper follows DiffusionNFT: instead of estimating likelihoods, it constructs a pair of implicit policies on the flow's forward process, letting a single model play both roles —

\[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)\]

where $v^{\mathrm{old}}$ is the data-collection policy (an exponential-moving-average lag of $v_\theta$) and $\gamma$ controls guidance strength. During training, $r_{\mathrm{adv}}$ weights a contrastive objective in which the positive velocity is fitted to the velocity targets of high-reward samples and the negative velocity to those of low-reward samples (equivalently, pushed away from them), with a $\beta$-scaled KL divergence added to keep the post-trained model near the pretrained policy.

Two aspects of this construction deserve emphasis. First, the positive policy is a convex combination of the old and current policies, which implicitly regularizes "deviate from the pretrained model, but not too far": larger $\gamma$ means weaker guidance and a model closer to $v^{\mathrm{old}}$, while smaller $\gamma$ permits more aggressive change. Second, the derivation relies only on the flow's forward process and assumes no standard Gaussian prior — exactly what the paper's "same-batch control distribution → perturbed distribution" setting requires, whereas MDP-based GRPO variants such as FlowGRPO and DanceGRPO start from a Gaussian prior and cannot be applied here as is. The KL term is not merely a guard against overfitting; it is the defence against reward hacking, without which the model could learn to produce extreme images that fool the evaluators (for instance, inflating nuclei to harvest a size reward) rather than becoming genuinely more correct.

4. Reward-driven test-time scaling: the same evaluators double as verifiers

Since the reward functions can already score generated images automatically, they need not be confined to training. At inference, for a given $(x_0,c)$ the model samples $N$ candidates and simply outputs the one with the highest weighted total reward (best-of-N) — no extra training, only inference compute. This mirrors best-of-N verification in reasoning models: the authors observe that RL and test-time scaling are complementary rather than substitutes. RL improves the quality of the candidate pool itself, so its benefit is largest at small $N$ (the base distribution is already good enough that a random draw is fine), while as $N$ grows both curves rise monotonically with the RL curve always above the base model's — meaning that at any fixed compute budget the RL-post-trained model is the better deal.

Loss & Training¶

The training objective maximizes the weighted sum of the seven rewards subject to a KL constraint between the post-trained and pretrained models, i.e. $\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]$ under a KL constraint (Eq. 2 in the paper; ⚠️ the exact constraint form and threshold are incomplete in the cached text). Implementation uses the online RL algorithm DiffusionNFT, with v_old following v_θ by exponential moving average.

The reward weights are $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 and nucleus-in-cytoplasm are empirically harder to optimize and are therefore weighted higher. The authors also tried a uniform weight of 1.0 on all normalized rewards (requiring no prior knowledge for weight tuning), which likewise improves over the base model. The KL weight is $\beta=1$ and all other hyper-parameters follow DiffusionNFT. Training runs for 1200 steps on a single H100 for 32 hours; inference uses best-of-N with $N=4$.

Key Experimental Results¶

Main Results¶

The dataset is BBBC021, a high-content microscopy chemical perturbation dataset with 98K three-channel 96×96 images covering 26 compounds grouped into 12 MoAs; the train/test split follows CellFlux. Baselines are the generative methods PhenDiff and IMPA plus the base model CellFlux. Reward metrics are higher-is-better; FID/KID are lower-is-better.

Metric	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

Ablation Study¶

Below: "optimizing a single reward" versus "jointly optimizing all seven" (CellFluxRL). Each column is a model RL-post-trained for that one reward only.

Metric	MoA-only	Roundness-only	NucSize-only	CellFluxRL (all)
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

Key Findings¶

RL post-training beats the base model on all seven rewards, raising the overall reward from -2.44 to 0.46, with gains spanning biological function, structural constraints, and morphological statistics. This validates the diagnosis that pixel-level objectives cannot see global constraints and that RL can supply them.
MoA shows the largest gain: 0.26 → 0.34, corresponding to classification accuracy rising from 0.61 to 0.66 (PhenDiff 0.56, IMPA 0.53). RL makes generated images far more likely to be recognized as exhibiting the morphology expected of that drug class.
The single-reward ablation shows the value of combining rewards. Optimizing MoA alone yields a higher MoA score (0.412 > 0.337) but collapses NucSize to -3.319; optimizing NucSize alone reaches -0.702 on that metric while degrading others. Jointly trained CellFluxRL ranks second or third on every individual metric (falling short only of the specialist models on their own targets), so stacking rewards buys balance across dimensions rather than a trade-off.
Image quality is roughly preserved but is not the goal: CellFluxRL achieves the best KID (0.014) of all methods, yet its FID of 24.01 is worse than the base CellFlux's 20.36 (TTS brings it back to 23.19). The authors state explicitly that FID/KID are not the optimization target — biological correctness is.
Roundness is the one metric where a baseline leads: PhenDiff (-0.24) is slightly better than CellFluxRL (-0.26), though +TTS (-0.19) overtakes it. This is the only cell in the table where CellFluxRL does not lead (or tie), so the "improves on everything" claim holds against the base model but has this exception among all baselines.
The KL weight $\beta$ reveals a trade-off: raising it from 1.0 to 1.3, structural and morphological metrics prefer a smaller KL weight while biological function rewards improve slightly with a larger one. The authors infer that higher biological accuracy requires more substantial model change than structural or morphological fidelity does. In practice $\beta$ is the knob for this trade-off; the default 1.0 favours structural and morphological constraints, and a larger value is recommended when MoA consistency is the priority.
Test-time scaling improves monotonically with $N$, and RL post-training raises scaling efficiency: the RL curve sits above the base model's at any $N$, with the largest advantage at small $N$, where the quality of the base distribution matters most.

Highlights & Insights¶

One suite of evaluators, three uses: the same biological rewards serve simultaneously as evaluation metrics, training rewards, and inference-time verifiers. This abstraction is the paper's most elegant element; it cleanly transfers the mature "reward model + best-of-N" paradigm from RLHF and reasoning models to scientific generation, without any human preference annotation.
Domain knowledge instead of human preference: rewards are not trained but composed from existing tools (a pretrained MoA classifier, Cellpose segmentation) and simple statistics. This avoids the cost of reward-model training and preference collection, and keeps every reward interpretable and auditable — each corresponds to an explicit biological statement.
Non-differentiability is turned from an obstacle into the reason for the method: because the constraints cannot be differentiated, RL is the only route, and RL happens to exploit exactly these off-the-shelf evaluators. The same logic explains why no architectural change (e.g. a physics-constraint layer) is made — that would require designing a differentiable surrogate per constraint, which is costly and easily distorted.
The algorithm is matched to the setting with an explicit argument: rather than applying the GRPO family by default, the authors point out that MDP-based methods assume a Gaussian prior and are incompatible with a "start from same-batch control images" distribution-to-distribution setting, hence the forward-process DiffusionNFT. This habit of first stating the setting's constraints and only then choosing the algorithm is worth copying.
Transferable idea: any scientific generation task with non-differentiable rules or simulators that can act as a judge (material structure generation, medical image synthesis, molecular conformation sampling) can adopt the same "reward as metric + training signal + verifier" triad — especially where the true distribution necessarily satisfies statistical invariants (conservation laws, geometric containment, population statistics). Writing such invariants as negative normalized deviations produces a ready-made reward.

Limitations & Future Work¶

The limitation the authors acknowledge is that all rewards are hand-engineered and require domain experts to translate biological knowledge into computable form, which is a high barrier for new domains; their proposed direction is to use LLMs to automatically translate scientific literature into executable reward functions.
Validation is confined to a single dataset: BBBC021 is modest (98K images, 96×96, 26 compounds, 12 MoAs), and its MoA classes are few with visually distinct morphologies. Whether the rewards remain discriminative and RL stays stable when scaled to million-image, thousand-perturbation datasets such as JUMP or RxRx is not addressed.
The reward models themselves may inject bias: the MoA classifier is pretrained on BBBC021, so using its predicted probability as a reward amounts to making the generative model cater to that classifier's decision boundary. Classifier errors or shortcuts can be amplified by RL (another form of reward hacking), and the paper only guards against this indirectly with the KL term — it never tests whether swapping in a different classifier keeps the rewards effective.
The cost is non-trivial: 1200 training steps take 32 hours on one H100, and best-of-N inference multiplies sampling cost by $N$. Only $N=4$ is reported quantitatively; larger $N$ appears only as a trend in a figure, with no report of when returns saturate.
The modelling scope stops at morphological images and does not touch gene expression, transcriptomics, or signalling pathways, so it remains distant from the full vision of a virtual cell (multi-omics, interpretable mechanisms). "Biology-constrained" here means constraints at the morphological level.
Improvement directions: make some rewards differentiable relaxations (e.g. soft IoU / differentiable segmentation) to cut RL variance and sampling cost; use classifier ensembles or adversarial reward auditing to suppress reward hacking; extend rewards to time series (live-cell imaging) rather than fixed stained endpoints.

vs CellFlux: CellFlux uses flow matching for distribution-to-distribution perturbation modelling and starts from same-batch control images to remove batch effects, but optimizes only a pixel-level objective; this paper fully reuses its velocity field and adds a biological-constraint alignment layer through RL at the output. The difference is the level of optimization — one optimizes "does the distribution look alike", the other "is the physics right".
vs PhenDiff / IMPA: two generative approaches to predicting morphological perturbation responses (diffusion-based and generative-model-based respectively), neither with an explicit constraint mechanism, and both weak on MoA consistency (0.18 / 0.12 reward, 0.56 / 0.53 accuracy). This paper lifts the MoA reward to 0.34 (0.66 accuracy) at almost no cost in generation quality, showing that explicit reward alignment buys more biological usability than piling on generation quality.
vs FlowGRPO / DanceGRPO: these cast diffusion/flow-matching generation as an MDP and optimize it with GRPO, requiring a standard Gaussian prior. This paper uses DiffusionNFT's forward-process estimation instead, because a virtual cell must start from a same-batch control distribution, not from noise — a good example of an algorithm being chosen to fit the data setting.
vs physics-aware generation (PhysGDPO, RDPO, VisionReward, etc.): those works align image/video generation with human preferences or physical common sense. This paper belongs to the same "modify the objective to inject physical constraints" line, but swaps in computable biological evaluators from the domain, emphasizes that no architectural change is needed, and gains the extra benefit that the rewards can be reused for test-time scaling.
Insight: when the metrics for a task are expensive or non-differentiable, it is worth asking whether those metrics can double as training signals and as inference-time selectors. This paper answers yes — and using one suite of metrics three ways also removes the need for a separate reward model.

Rating¶

Novelty: ⭐⭐⭐⭐ The first systematic application of "computable biological rewards + RL post-training + rewards as verifiers" to image-based virtual cells, with genuine domain insight in the reward design and its three-category, seven-term split; the core RL algorithm is inherited from DiffusionNFT, so the contribution lies mainly in problem diagnosis and domain grounding.
Experimental Thoroughness: ⭐⭐⭐ Main table, single-reward ablation, KL sensitivity, and TTS curves are all present, with three generative baselines; however, everything is validated on BBBC021 alone, with no cross-dataset or cross-base-model generalization study, and TTS is reported only at $N=4$.
Writing Quality: ⭐⭐⭐⭐ The chain from the failure case (nucleus outside the cytoplasm) to the diagnosis ("a pixel-level objective cannot see global constraints") to the reward design is very clear, and the algorithm–setting match is well argued; some equations suffer from typesetting corruption.
Value: ⭐⭐⭐⭐ The "one suite of domain evaluators as metric, training reward, and inference verifier" recipe has general reference value for scientific generative models, and it gives the community biological evaluation axes beyond FID/KID.