Learning Perturbation-Specific Representations through Conditional Energy Prior for Cellular Response Prediction

Yuchen Yan; Yatao Bian

cs.LGOct 13, 2026

Learning Perturbation-Specific Representations through Conditional Energy Prior for Cellular Response Prediction

Authors: Yuchen Yan, Yatao Bian

Organizations: Department of Computer Science, National University of Singapore

Abstract

AI-based cellular perturbation prediction is essential to accelerate drug discovery and therapeutic identification. However, perturbation-specific signals are generally obscured by cellular complexity, making it difficult for existing methods to correctly capture perturbation effects. To address this problem, we propose PerturbEnergy, which introduces a conditional energy-based prior over perturbation representations. It assigns lower energy (higher probability) to representations compatible with the given perturbation and higher energy to others, thereby explicitly tilting the latent space toward perturbation-preferred regions. Besides, PerturbEnergy can be integrated into autoencoder architectures as we find its Evidence Lower Bound (ELBO) is equivalent to the ELBO of the complete autoencoder. We further prove that optimizing PerturbEnergy increases the lower bound on the conditional mutual information between representation and perturbation. Extensive experiments on four benchmark datasets demonstrate the effectiveness of our PerturbEnergy: it improves perturbation recognition especially in latent space with improvement of up to 83.2%; It also achieves state-of-the-art performance for predicting top-50/100/200 differentially expressed gene(DEG) response and remaining competitive on highly variable gene(HVG) prediction. The code is provided: https://github.com/BlueWhaleLab/PerturbEnergy

Figures & tables

Figure 1: Perturbation classification accuracy. Existing methods show limited perturbation recognition rate both in representation and in expression.

Figure 2: Our PerturbEnergy framework. It reshapes the latent prior to perturbation preferred regions by assigning low energy to perturbation-compatible representations. Thus we can sample the representation with more perturbation-specific information for the decoder.

Figure 3: Performance on parts of the HVG group metrics.

\uparrow/\downarrow

means the higher/lower is better.

Method	Replogle			Tahoe100M
Method	Top-50	Top-100	Top-200	Top-50	Top-100	Top-200
Linear	0.2738 $\pm$ 0.0086	0.2390 $\pm$ 0.0072	0.2045 $\pm$ 0.0058	0.4536 $\pm$ 0.0072	0.3866 $\pm$ 0.0053	0.3032 $\pm$ 0.0040
Linear (scGPT)	0.2449 $\pm$ 0.0081	0.2081 $\pm$ 0.0066	0.1738 $\pm$ 0.0052	0.3748 $\pm$ 0.0060	0.3257 $\pm$ 0.0044	0.2588 $\pm$ 0.0034
CellFlow	0.2174 $\pm$ 0.0067	0.1850 $\pm$ 0.0055	0.1536 $\pm$ 0.0043	0.3860 $\pm$ 0.0062	0.3318 $\pm$ 0.0044	0.2617 $\pm$ 0.0034
Squidiff	0.2173 $\pm$ 0.0055	0.1907 $\pm$ 0.0044	0.1658 $\pm$ 0.0035	0.4192 $\pm$ 0.0071	0.3889 $\pm$ 0.0051	0.3231 $\pm$ 0.0039
PerturbDiff	0.2612 $\pm$ 0.0096	0.2198 $\pm$ 0.0079	0.1808 $\pm$ 0.0063	0.3558 $\pm$ 0.0070	0.3079 $\pm$ 0.0051	0.2462 $\pm$ 0.0040

Table 2: RMSE (

\downarrow

) over the true top-

K

DEGs. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	Latent Representation				Expression Delta
Method	Replogle	PBMC	Norman2019	Tahoe100M	Replogle	PBMC	Norman2019	Tahoe100M
scGPT	18.5 $\pm$ 1.1	33.1 $\pm$ 1.0	74.4 $\pm$ 2.9	38.2 $\pm$ 1.9	21.5 $\pm$ 2.5	24.8 $\pm$ 1.4	73.1 $\pm$ 2.7	39.9 $\pm$ 2.3
PerturbDiff	9.1 $\pm$ 3.3	14.9 $\pm$ 2.3	29.5 $\pm$ 4.1	19.5 $\pm$ 3.8	18.4 $\pm$ 1.2	18.3 $\pm$ 1.5	52.6 $\pm$ 3.0	31.4 $\pm$ 3.4
STATE	16.3 $\pm$ 0.9	12.6 $\pm$ 1.1	23.1 $\pm$ 2.0	23.6 $\pm$ 3.6	24.5 $\pm$ 1.8	26.9 $\pm$ 2.8	70.0 $\pm$ 1.2	41.8 $\pm$ 4.1
scGen	24.6 $\pm$ 1.6	48.4 $\pm$ 2.5	81.6 $\pm$ 1.7	46.1 $\pm$ 2.5	19.7 $\pm$ 2.3	24.6 $\pm$ 1.9	68.2 $\pm$ 2.3	35.2 $\pm$ 1.5
CPA	22.2 $\pm$ 0.4	43.0 $\pm$ 1.9	78.2 $\pm$ 1.9	33.3 $\pm$ 4.0	16.3 $\pm$ 1.4	18.5 $\pm$ 1.0	60.1 $\pm$ 3.3	24.7 $\pm$ 1.6

Table 3: Perturbation classification accuracy. Best in bold , second best underlined .

Figure 4: Ablation and other study on Replogle. (a) is the ablation; (b) is the hyperparameter study about Langevin steps/size; (c) is the influence of candidate number $M$ . (d) is the training efficiency.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

z_{a}^{(k+1)}=z_{a}^{(k)}-\eta\nabla_{z_{a}}E_{\alpha}\left(z_{a}^{(k)},z_{b},a\right)+\sqrt{2\eta}\,\epsilon^{(k)}

Appendix

Algorithm 1 Langevin Sampling

Dataset	Perturbation Type	#Perturbed Cells			#Control Cells	#Genes	#HVGs	#Perturbations	#Cell Types	#Batches
Dataset	Perturbation Type	Train	Valid	Test	#Control Cells	#Genes	#HVGs	#Perturbations	#Cell Types	#Batches
Replogle	gene	572,545	4,825	26,878	39,165	6,642	2,000	2,023	4	56
PBMC	cytokine	6,657,336	150,484	2,260,453	629,701	62,710	2,000	90	18	12
Norman	gene	85,568	1,754	12,268	11,855	33,694	2,000	236	1	8
Tahoe100M	drug	89,495,239	553,076	8,790,272	2,330,156	40,352	2,000	1,137	50	14

Appendix

Table 4: (a) Perturbation datasets. Statistics include counts of perturbed cell across train/valid/test splits and counts of control cells, total gene vocabulary size and HVGs, and dataset diversity in terms of perturbations, cell types, and experimental batches.

Model	Family	LR	EMA	Model-specific settings
Linear	linear	—	—	closed form; no gradient training
scGPT	foundation	$1\!\times\!10^{-3}$	—	encoder frozen; head $512{+}512\!\to\!1024\!\to\!G$
scGen	VAE	$1\!\times\!10^{-3}$	—	$z=512$ , hidden $800{\times}2$ , no conditioning
CPA	VAE	$1\!\times\!10^{-3}$	—	$z=512$ , adv. weight 50, 3 adversary steps
stVAE	VAE	$1\!\times\!10^{-3}$	—	$z=512$ , latent loss 1.0, VAE loss 0.1
PerturbedVAE	VAE	$1\!\times\!10^{-3}$	—	$z_{i}=z_{v}=512$ , contrastive weight 1.0

Appendix

Table 5: Model-specific training settings. All other settings including batch shape, step budget, optimiser family, gradient clipping, checkpoint selection and evaluation are shared across models and identical to those used for PerturbEnergy.

Method	Replogle			Tahoe100M
Method	Top-50	Top-100	Top-200	Top-50	Top-100	Top-200
scGen	0.2336 $\pm$ 0.0078	0.1992 $\pm$ 0.0065	0.1661 $\pm$ 0.0051	0.4462 $\pm$ 0.0081	0.3818 $\pm$ 0.0059	0.3018 $\pm$ 0.0046
+ PerturbEnergy (Lang)	0.1931 $\pm$ 0.0064	0.1631 $\pm$ 0.0052	0.1356 $\pm$ 0.0041	0.3457 $\pm$ 0.0061	0.3023 $\pm$ 0.0044	0.2421 $\pm$ 0.0034
+ PerturbEnergy (Mis)	0.1902 $\pm$ 0.0063	0.1607 $\pm$ 0.0051	0.1341 $\pm$ 0.0040	0.3387 $\pm$ 0.0059	0.2966 $\pm$ 0.0043	0.2379 $\pm$ 0.0033
PerturbedVAE	0.2130 $\pm$ 0.0071	0.1817 $\pm$ 0.0058	0.1514 $\pm$ 0.0047	0.3984 $\pm$ 0.0075	0.3476 $\pm$ 0.0055	0.2791 $\pm$ 0.0042
+ PerturbEnergy (Lang)	0.1966 $\pm$ 0.0058	0.1692 $\pm$ 0.0048	0.1424 $\pm$ 0.0038	0.3892 $\pm$ 0.0073	0.3324 $\pm$ 0.0053	0.2653 $\pm$ 0.0040

Appendix

Table 6: Backbone transfer: RMSE (

\downarrow

) over the true top-

K

DEGs when the PerturbEnergy objective is applied to two other backbones. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Within each backbone block, the best entry per column is in bold .

Figure 5: Influence of training steps and model dimension on Replogle and Tahoe100M.

Figure 6: Visualization of representation in Tahoe100M.

Method	Replogle			Tahoe100M
Method	Top-50	Top-100	Top-200	Top-50	Top-100	Top-200
Linear	0.2282 $\pm$ 0.0076	0.1963 $\pm$ 0.0062	0.1634 $\pm$ 0.0049	0.3788 $\pm$ 0.0069	0.3169 $\pm$ 0.0046	0.2246 $\pm$ 0.0034
Linear (scGPT)	0.2114 $\pm$ 0.0076	0.1757 $\pm$ 0.0060	0.1418 $\pm$ 0.0046	0.3092 $\pm$ 0.0057	0.2624 $\pm$ 0.0039	0.1884 $\pm$ 0.0029
CellFlow	0.1925 $\pm$ 0.0063	0.1615 $\pm$ 0.0051	0.1306 $\pm$ 0.0039	0.3280 $\pm$ 0.0059	0.2767 $\pm$ 0.0039	0.1975 $\pm$ 0.0029
Squidiff	0.1675 $\pm$ 0.0047	0.1433 $\pm$ 0.0037	0.1205 $\pm$ 0.0029	0.3077 $\pm$ 0.0061	0.2785 $\pm$ 0.0042	0.2116 $\pm$ 0.0030
PerturbDiff	0.2304 $\pm$ 0.0092	0.1904 $\pm$ 0.0075	0.1519 $\pm$ 0.0059	0.2878 $\pm$ 0.0062	0.2415 $\pm$ 0.0042	0.1748 $\pm$ 0.0031

Appendix

Table 7: MAE (

\downarrow

) over the true top-

K

DEGs. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	Replogle			Tahoe100M
Method	Top-50	Top-100	Top-200	Top-50	Top-100	Top-200
Linear	0.9611 $\pm$ 0.0024	0.9676 $\pm$ 0.0018	0.9712 $\pm$ 0.0015	0.9202 $\pm$ 0.0031	0.9323 $\pm$ 0.0022	0.9456 $\pm$ 0.0015
Linear (scGPT)	0.9692 $\pm$ 0.0020	0.9753 $\pm$ 0.0015	0.9789 $\pm$ 0.0012	0.9406 $\pm$ 0.0026	0.9480 $\pm$ 0.0019	0.9573 $\pm$ 0.0013
CellFlow	0.9759 $\pm$ 0.0017	0.9809 $\pm$ 0.0012	0.9840 $\pm$ 0.0009	0.9372 $\pm$ 0.0030	0.9460 $\pm$ 0.0021	0.9563 $\pm$ 0.0015
Squidiff	0.9760 $\pm$ 0.0013	0.9795 $\pm$ 0.0009	0.9813 $\pm$ 0.0008	0.9116 $\pm$ 0.0034	0.9145 $\pm$ 0.0025	0.9250 $\pm$ 0.0019
PerturbDiff	0.9644 $\pm$ 0.0025	0.9720 $\pm$ 0.0019	0.9768 $\pm$ 0.0015	0.9443 $\pm$ 0.0036	0.9521 $\pm$ 0.0026	0.9610 $\pm$ 0.0019

Appendix

Table 8: PCC (

\uparrow

) over the true top-

K

DEGs. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	Replogle			Tahoe100M
Method	Top-50	Top-100	Top-200	Top-50	Top-100	Top-200
Linear	0.9142 $\pm$ 0.0050	0.9297 $\pm$ 0.0037	0.9383 $\pm$ 0.0031	0.7456 $\pm$ 0.0108	0.8009 $\pm$ 0.0078	0.8497 $\pm$ 0.0051
Linear (scGPT)	0.9295 $\pm$ 0.0044	0.9456 $\pm$ 0.0031	0.9546 $\pm$ 0.0025	0.8303 $\pm$ 0.0069	0.8621 $\pm$ 0.0051	0.8930 $\pm$ 0.0034
CellFlow	0.9444 $\pm$ 0.0035	0.9569 $\pm$ 0.0025	0.9646 $\pm$ 0.0020	0.8185 $\pm$ 0.0074	0.8556 $\pm$ 0.0053	0.8896 $\pm$ 0.0035
Squidiff	0.9466 $\pm$ 0.0028	0.9560 $\pm$ 0.0019	0.9608 $\pm$ 0.0016	0.7862 $\pm$ 0.0090	0.8028 $\pm$ 0.0072	0.8323 $\pm$ 0.0049
PerturbDiff	0.9177 $\pm$ 0.0054	0.9374 $\pm$ 0.0039	0.9488 $\pm$ 0.0032	0.8482 $\pm$ 0.0071	0.8767 $\pm$ 0.0052	0.9032 $\pm$ 0.0036

Appendix

Table 9:

R^{2}

(

\uparrow

) over the true top-

K

DEGs. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	PCC $\uparrow$	$R^{2}$ $\uparrow$	PDCorr $\uparrow$	LFCSpear $\uparrow$	LogFC-cos $\uparrow$
Linear	0.9772 $\pm$ 0.0011	0.9529 $\pm$ 0.0022	0.3323 $\pm$ 0.0129	0.2519 $\pm$ 0.0117	0.3368 $\pm$ 0.0131
Linear (scGPT)	0.9868 $\pm$ 0.0006	0.9732 $\pm$ 0.0013	0.4584 $\pm$ 0.0175	0.3607 $\pm$ 0.0145	0.4688 $\pm$ 0.0174
CellFlow	0.9706 $\pm$ 0.0012	0.9393 $\pm$ 0.0024	0.5559 $\pm$ 0.0187	0.4623 $\pm$ 0.0170	0.5648 $\pm$ 0.0185
Squidiff	0.9829 $\pm$ 0.0005	0.9649 $\pm$ 0.0009	0.4945 $\pm$ 0.0170	0.4536 $\pm$ 0.0141	0.5065 $\pm$ 0.0169
PerturbDiff	0.9875 $\pm$ 0.0008	0.9743 $\pm$ 0.0015	0.5470 $\pm$ 0.0174	0.5186 $\pm$ 0.0137	0.5557 $\pm$ 0.0173
STATE	0.9905 $\pm$ 0.0005	0.9805 $\pm$ 0.0010	0.5091 $\pm$ 0.0206	0.4323 $\pm$ 0.0183	0.5177 $\pm$ 0.0205

Appendix

Table 10: HVG-group correlation-based measures on Replogle, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups; the perturbation-discrimination scores are computed once over the whole test set and carry no such spread. Best in bold , second best underlined .

Method	DEPrec $\uparrow$	DERec $\uparrow$	DEOverlap $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Linear	0.2047 $\pm$ 0.0075	0.2047 $\pm$ 0.0075	0.1198 $\pm$ 0.0049	0.7200 $\pm$ 0.0076	0.1662 $\pm$ 0.0073
Linear (scGPT)	0.2590 $\pm$ 0.0106	0.2590 $\pm$ 0.0106	0.1615 $\pm$ 0.0073	0.7698 $\pm$ 0.0092	0.2409 $\pm$ 0.0115
CellFlow	0.0967 $\pm$ 0.0033	0.0967 $\pm$ 0.0033	0.0518 $\pm$ 0.0019	0.6708 $\pm$ 0.0054	0.0806 $\pm$ 0.0023
Squidiff	0.2566 $\pm$ 0.0104	0.2566 $\pm$ 0.0104	0.1592 $\pm$ 0.0071	0.7981 $\pm$ 0.0089	0.2475 $\pm$ 0.0114
PerturbDiff	0.2982 $\pm$ 0.0132	0.2982 $\pm$ 0.0132	0.1973 $\pm$ 0.0100	0.8096 $\pm$ 0.0085	0.2938 $\pm$ 0.0147
STATE	0.3049 $\pm$ 0.0134	0.3049 $\pm$ 0.0134	0.2024 $\pm$ 0.0101	0.7946 $\pm$ 0.0101	0.2997 $\pm$ 0.0149

Appendix

Table 11: HVG-group DEG-related scores on Replogle, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	RMSE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	EDist $\downarrow$	Sinkhorn $\downarrow$
Linear	0.1122 $\pm$ 0.0025	0.0144 $\pm$ 0.0007	0.0764 $\pm$ 0.0016	1.4655 $\pm$ 0.0592	1.0269 $\pm$ 0.0047
Linear (scGPT)	0.0860 $\pm$ 0.0017	0.0082 $\pm$ 0.0004	0.0585 $\pm$ 0.0010	0.9934 $\pm$ 0.0352	0.9918 $\pm$ 0.0035
CellFlow	0.1303 $\pm$ 0.0024	0.0186 $\pm$ 0.0007	0.0866 $\pm$ 0.0015	4.4345 $\pm$ 0.0812	0.6073 $\pm$ 0.0052
Squidiff	0.1021 $\pm$ 0.0012	0.0108 $\pm$ 0.0003	0.0656 $\pm$ 0.0007	3.7791 $\pm$ 0.0379	0.5664 $\pm$ 0.0035
PerturbDiff	0.0817 $\pm$ 0.0020	0.0079 $\pm$ 0.0005	0.0527 $\pm$ 0.0012	6.3453 $\pm$ 0.0534	0.5637 $\pm$ 0.0041
STATE	0.0725 $\pm$ 0.0016	0.0060 $\pm$ 0.0003	0.0491 $\pm$ 0.0009	0.7007 $\pm$ 0.0274	0.9903 $\pm$ 0.0038

Appendix

Table 12: HVG-group distribution-level distances on Replogle, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	PCC $\uparrow$	$R^{2}$ $\uparrow$	PDCorr $\uparrow$	LFCSpear $\uparrow$	LogFC-cos $\uparrow$
Linear	0.9705 $\pm$ 0.0009	0.9372 $\pm$ 0.0020	0.2015 $\pm$ 0.0123	0.0708 $\pm$ 0.0060	0.1960 $\pm$ 0.0123
Linear (scGPT)	0.9826 $\pm$ 0.0007	0.9636 $\pm$ 0.0014	0.2922 $\pm$ 0.0119	0.0881 $\pm$ 0.0049	0.2934 $\pm$ 0.0119
CellFlow	0.9784 $\pm$ 0.0013	0.9555 $\pm$ 0.0025	0.2301 $\pm$ 0.0067	0.2610 $\pm$ 0.0051	0.2294 $\pm$ 0.0069
Squidiff	0.9605 $\pm$ 0.0009	0.8978 $\pm$ 0.0016	0.3061 $\pm$ 0.0102	0.2693 $\pm$ 0.0044	0.3003 $\pm$ 0.0103
PerturbDiff	0.9805 $\pm$ 0.0007	0.9521 $\pm$ 0.0021	0.4435 $\pm$ 0.0125	0.3566 $\pm$ 0.0057	0.4429 $\pm$ 0.0126
STATE	0.9868 $\pm$ 0.0006	0.9724 $\pm$ 0.0012	0.3074 $\pm$ 0.0140	0.0808 $\pm$ 0.0061	0.3078 $\pm$ 0.0140

Appendix

Table 13: HVG-group correlation-based measures on PBMC, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups; the perturbation-discrimination scores are computed once over the whole test set and carry no such spread. Best in bold , second best underlined .

Method	DEPrec $\uparrow$	DERec $\uparrow$	DEOverlap $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Linear	0.1404 $\pm$ 0.0055	0.1404 $\pm$ 0.0055	0.0786 $\pm$ 0.0035	0.7303 $\pm$ 0.0045	0.1169 $\pm$ 0.0059
Linear (scGPT)	0.1693 $\pm$ 0.0067	0.1693 $\pm$ 0.0067	0.0975 $\pm$ 0.0046	0.7147 $\pm$ 0.0052	0.1469 $\pm$ 0.0070
CellFlow	0.1349 $\pm$ 0.0035	0.1349 $\pm$ 0.0035	0.0735 $\pm$ 0.0021	0.7689 $\pm$ 0.0030	0.1062 $\pm$ 0.0024
Squidiff	0.1501 $\pm$ 0.0063	0.1501 $\pm$ 0.0063	0.0853 $\pm$ 0.0041	0.8073 $\pm$ 0.0036	0.1363 $\pm$ 0.0058
PerturbDiff	0.1995 $\pm$ 0.0076	0.1995 $\pm$ 0.0076	0.1176 $\pm$ 0.0054	0.8215 $\pm$ 0.0039	0.1701 $\pm$ 0.0080
STATE	0.1929 $\pm$ 0.0083	0.1929 $\pm$ 0.0083	0.1149 $\pm$ 0.0060	0.7200 $\pm$ 0.0060	0.1724 $\pm$ 0.0089

Appendix

Table 14: HVG-group DEG-related scores on PBMC, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	RMSE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	EDist $\downarrow$	Sinkhorn $\downarrow$
Linear	0.0376 $\pm$ 0.0009	0.0016 $\pm$ 0.0001	0.0195 $\pm$ 0.0004	0.2725 $\pm$ 0.0137	0.4927 $\pm$ 0.0113
Linear (scGPT)	0.0299 $\pm$ 0.0010	0.0012 $\pm$ 0.0001	0.0167 $\pm$ 0.0004	0.2035 $\pm$ 0.0137	0.4902 $\pm$ 0.0111
CellFlow	0.0338 $\pm$ 0.0014	0.0017 $\pm$ 0.0002	0.0159 $\pm$ 0.0006	1.0656 $\pm$ 0.0340	0.2766 $\pm$ 0.0072
Squidiff	0.0529 $\pm$ 0.0013	0.0033 $\pm$ 0.0002	0.0217 $\pm$ 0.0006	1.6522 $\pm$ 0.0339	0.2937 $\pm$ 0.0077
PerturbDiff	0.0340 $\pm$ 0.0008	0.0014 $\pm$ 0.0001	0.0146 $\pm$ 0.0004	2.4800 $\pm$ 0.0287	0.2905 $\pm$ 0.0071
STATE	0.0264 $\pm$ 0.0009	0.0009 $\pm$ 0.0001	0.0138 $\pm$ 0.0005	0.1658 $\pm$ 0.0112	0.4879 $\pm$ 0.0112

Appendix

Table 15: HVG-group distribution-level distances on PBMC, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	PCC $\uparrow$	$R^{2}$ $\uparrow$	PDCorr $\uparrow$	LFCSpear $\uparrow$	LogFC-cos $\uparrow$
Linear	0.9837 $\pm$ 0.0022	0.9674 $\pm$ 0.0043	0.4668 $\pm$ 0.0426	0.2613 $\pm$ 0.0231	0.4710 $\pm$ 0.0422
Linear (scGPT)	0.9971 $\pm$ 0.0003	0.9941 $\pm$ 0.0007	0.9051 $\pm$ 0.0160	0.3811 $\pm$ 0.0145	0.9055 $\pm$ 0.0160
CellFlow	0.9772 $\pm$ 0.0027	0.9545 $\pm$ 0.0053	0.9137 $\pm$ 0.0270	0.3883 $\pm$ 0.0141	0.9141 $\pm$ 0.0146
Squidiff	0.9921 $\pm$ 0.0006	0.9839 $\pm$ 0.0012	0.8296 $\pm$ 0.0220	0.4656 $\pm$ 0.0150	0.8307 $\pm$ 0.0219
PerturbDiff	0.9975 $\pm$ 0.0003	0.9944 $\pm$ 0.0006	0.8985 $\pm$ 0.0205	0.5092 $\pm$ 0.0132	0.8984 $\pm$ 0.0205
STATE	0.9969 $\pm$ 0.0004	0.9936 $\pm$ 0.0008	0.9078 $\pm$ 0.0153	0.3700 $\pm$ 0.0164	0.9078 $\pm$ 0.0153

Appendix

Table 16: HVG-group correlation-based measures on Norman2019, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups; the perturbation-discrimination scores are computed once over the whole test set and carry no such spread. Best in bold , second best underlined .

Method	DEPrec $\uparrow$	DERec $\uparrow$	DEOverlap $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Linear	0.3251 $\pm$ 0.0119	0.3251 $\pm$ 0.0119	0.1964 $\pm$ 0.0086	0.9423 $\pm$ 0.0032	0.3358 $\pm$ 0.0161
Linear (scGPT)	0.6528 $\pm$ 0.0203	0.6528 $\pm$ 0.0203	0.4960 $\pm$ 0.0201	0.9777 $\pm$ 0.0053	0.7284 $\pm$ 0.0251
CellFlow	0.2190 $\pm$ 0.0097	0.2190 $\pm$ 0.0097	0.1242 $\pm$ 0.0062	0.8502 $\pm$ 0.0071	0.2118 $\pm$ 0.0089
Squidiff	0.5682 $\pm$ 0.0213	0.5682 $\pm$ 0.0213	0.4079 $\pm$ 0.0197	0.9735 $\pm$ 0.0035	0.6335 $\pm$ 0.0252
PerturbDiff	0.6467 $\pm$ 0.0217	0.6467 $\pm$ 0.0217	0.4909 $\pm$ 0.0216	0.9809 $\pm$ 0.0037	0.7091 $\pm$ 0.0268
STATE	0.6523 $\pm$ 0.0210	0.6523 $\pm$ 0.0210	0.4963 $\pm$ 0.0209	0.9757 $\pm$ 0.0064	0.7227 $\pm$ 0.0260

Appendix

Table 17: HVG-group DEG-related scores on Norman2019, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	RMSE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	EDist $\downarrow$	Sinkhorn $\downarrow$
Linear	0.0762 $\pm$ 0.0048	0.0067 $\pm$ 0.0009	0.0203 $\pm$ 0.0014	1.2427 $\pm$ 0.1400	1.3307 $\pm$ 0.0529
Linear (scGPT)	0.0325 $\pm$ 0.0019	0.0012 $\pm$ 0.0001	0.0132 $\pm$ 0.0006	0.5972 $\pm$ 0.0622	1.1605 $\pm$ 0.0364
CellFlow	0.0352 $\pm$ 0.0029	0.0016 $\pm$ 0.0003	0.0113 $\pm$ 0.0010	0.2976 $\pm$ 0.0452	1.2613 $\pm$ 0.0464
Squidiff	0.0551 $\pm$ 0.0022	0.0032 $\pm$ 0.0002	0.0174 $\pm$ 0.0007	0.9530 $\pm$ 0.0515	1.0328 $\pm$ 0.0339
PerturbDiff	0.0318 $\pm$ 0.0017	0.0011 $\pm$ 0.0001	0.0105 $\pm$ 0.0006	2.8086 $\pm$ 0.1059	0.8198 $\pm$ 0.0293
STATE	0.0337 $\pm$ 0.0021	0.0013 $\pm$ 0.0002	0.0127 $\pm$ 0.0007	0.3919 $\pm$ 0.0490	1.2014 $\pm$ 0.0364

Appendix

Table 18: HVG-group distribution-level distances on Norman2019, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	PCC $\uparrow$	$R^{2}$ $\uparrow$	PDCorr $\uparrow$	LFCSpear $\uparrow$	LogFC-cos $\uparrow$
Linear	0.9761 $\pm$ 0.0006	0.9451 $\pm$ 0.0018	0.3223 $\pm$ 0.0057	0.1999 $\pm$ 0.0040	0.3344 $\pm$ 0.0065
Linear (scGPT)	0.9783 $\pm$ 0.0005	0.9532 $\pm$ 0.0012	0.4708 $\pm$ 0.0076	0.2459 $\pm$ 0.0038	0.4803 $\pm$ 0.0076
CellFlow	0.9804 $\pm$ 0.0005	0.9583 $\pm$ 0.0012	0.5560 $\pm$ 0.0055	0.2803 $\pm$ 0.0033	0.5630 $\pm$ 0.0056
Squidiff	0.9343 $\pm$ 0.0009	0.8561 $\pm$ 0.0023	0.3047 $\pm$ 0.0050	0.1878 $\pm$ 0.0027	0.3067 $\pm$ 0.0049
PerturbDiff	0.9798 $\pm$ 0.0007	0.9529 $\pm$ 0.0013	0.5307 $\pm$ 0.0067	0.3664 $\pm$ 0.0033	0.5192 $\pm$ 0.0074
STATE	0.9815 $\pm$ 0.0004	0.9613 $\pm$ 0.0008	0.5694 $\pm$ 0.0063	0.2884 $\pm$ 0.0037	0.5755 $\pm$ 0.0065

Appendix

Table 19: HVG-group correlation-based measures on Tahoe100M, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups; the perturbation-discrimination scores are computed once over the whole test set and carry no such spread. Best in bold , second best underlined .

Method	DEPrec $\uparrow$	DERec $\uparrow$	DEOverlap $\uparrow$	AUROC $\uparrow$	AUPRC $\uparrow$
Linear	0.1318 $\pm$ 0.0038	0.1318 $\pm$ 0.0038	0.0720 $\pm$ 0.0023	0.7426 $\pm$ 0.0034	0.1015 $\pm$ 0.0030
Linear (scGPT)	0.2070 $\pm$ 0.0059	0.2070 $\pm$ 0.0059	0.1192 $\pm$ 0.0039	0.7713 $\pm$ 0.0040	0.1730 $\pm$ 0.0057
CellFlow	0.2339 $\pm$ 0.0056	0.2339 $\pm$ 0.0056	0.1361 $\pm$ 0.0038	0.7950 $\pm$ 0.0036	0.1972 $\pm$ 0.0059
Squidiff	0.1341 $\pm$ 0.0035	0.1341 $\pm$ 0.0035	0.0731 $\pm$ 0.0021	0.7502 $\pm$ 0.0029	0.0954 $\pm$ 0.0026
PerturbDiff	0.2054 $\pm$ 0.0054	0.2054 $\pm$ 0.0054	0.1176 $\pm$ 0.0035	0.7751 $\pm$ 0.0039	0.1659 $\pm$ 0.0051
STATE	0.2407 $\pm$ 0.0061	0.2407 $\pm$ 0.0061	0.1411 $\pm$ 0.0042	0.7951 $\pm$ 0.0039	0.2042 $\pm$ 0.0065

Appendix

Table 20: HVG-group DEG-related scores on Tahoe100M, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Method	RMSE $\downarrow$	MSE $\downarrow$	MAE $\downarrow$	EDist $\downarrow$	Sinkhorn $\downarrow$
Linear	0.1243 $\pm$ 0.0012	0.0159 $\pm$ 0.0004	0.0632 $\pm$ 0.0005	0.9646 $\pm$ 0.0147	0.7265 $\pm$ 0.0008
Linear (scGPT)	0.1166 $\pm$ 0.0010	0.0139 $\pm$ 0.0003	0.0629 $\pm$ 0.0004	0.9021 $\pm$ 0.0115	0.7206 $\pm$ 0.0008
CellFlow	0.1093 $\pm$ 0.0010	0.0122 $\pm$ 0.0003	0.0560 $\pm$ 0.0004	0.8195 $\pm$ 0.0103	0.7319 $\pm$ 0.0013
Squidiff	0.2058 $\pm$ 0.0010	0.0427 $\pm$ 0.0004	0.1101 $\pm$ 0.0005	6.7723 $\pm$ 0.0579	0.4257 $\pm$ 0.0011
PerturbDiff	0.1166 $\pm$ 0.0012	0.0140 $\pm$ 0.0003	0.0581 $\pm$ 0.0005	21.6085 $\pm$ 0.0591	0.3951 $\pm$ 0.0008
STATE	0.1062 $\pm$ 0.0007	0.0114 $\pm$ 0.0002	0.0560 $\pm$ 0.0004	0.7920 $\pm$ 0.0073	0.7239 $\pm$ 0.0008

Appendix

Table 21: HVG-group distribution-level distances on Tahoe100M, over all 2,000 highly variable genes. Values are the mean over held-out perturbation groups

\pm

the s.e.m. across those groups. Best in bold , second best underlined .

Explore similar work

Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation
Mar 7, 2026Andrea Giuseppe Di Francesco, Andrea Rubbi, Rishabh Jain +1Single-Cell Perturbation Modeling Retrieval-Augmented Generation
University of Cambridge Wellcome Sanger Institute
What Makes a Representation Good for Single-Cell Perturbation Prediction?
May 19, 2026Wenkang Jiang, Yuhang Liu, Yichao Cai +5Variational Autoencoders Disentangled Representation Learning
Australian Institute for Machine Learning, Adelaide University, Australia · Responsible AI Research Centre, Australia · College of Computer Science and Artificial Intelligence, Fudan University, China +3
PerturbPFN: Probing the Limits of Synthetic Priors in Drug Perturbation Modelling
Jul 26, 2026Yuche Gao, José Miguel Hernández-Lobato, Siyuan GuoSynthetic Data Pretraining Structural Causal Models
University of Cambridge, Cambridge, United Kingdom · Prior Labs, Freiburg, Germany

Mar 7, 2026J/K move · Enter open · S save

Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

Andrea Giuseppe Di Francesco, Andrea Rubbi, Rishabh Jain, Pietro Liò

University of Cambridge Wellcome Sanger Institute

Single-Cell Perturbation Modeling Retrieval-Augmented Generation Adaptive Retrieval

Open paper AskPDF arXiv

Predicting transcriptional responses to genetic perturbations is fundamental to functional genomics and therapeutic discovery. Recent deep learning models have shown promise in single-cell perturbation response prediction, but they typically generate each response in isolation, without explicitly leveraging related perturbations. We introduce PT-RAG (Perturbation-aware Two-stage Retrieval-Augmented Generation), a plug-in retrieval-and-conditioning module for generative cellular perturbation response. PT-RAG augments an existing perturbation-response backbone with learned access to related perturbation contexts. The key challenge is that relevance is not fixed in this setting: functionally related genes may elicit different effects across cell types. PT-RAG addresses this with a two-stage retrieval mechanism: GenePT-based semantic retrieval first identifies K candidate perturbations, after which a differentiable Gumbel-Softmax selector adaptively selects retrieved contexts conditioned on the control cell state, the query perturbation, and each candidate perturbation. We evaluate PT-RAG on two backbones, a STATE-style generator used as a frozen random reservoir and a fully trained scGPT, across cross-cell-type and cross-perturbation generalization tasks. PT-RAG consistently improves distributional similarity and often overall predictive quality; for example, on scGPT cross-cell-type results, the 2-Wasserstein distance drops by 5.9% relative to scGPT alone. The code to reproduce our experiments is available at https://github.com/difra100/PT-RAG_NIPS.

Abstract

Explore similar work

Differentiable Retrieval-Augmented Generation for Predicting Cellular Responses to Gene Perturbation

What Makes a Representation Good for Single-Cell Perturbation Prediction?

PerturbPFN: Probing the Limits of Synthetic Priors in Drug Perturbation Modelling