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

• Cellular response prediction could support therapeutic target identification, drug discovery, and personalized medicine, while experimental screening remains costly and labor-intensive.Source: §1 · p. 1
• Existing approaches encode cells and decode responses or model transitions between control and perturbed representations.Source: §1 · p. 1
• Existing methods show limited perturbation distinctiveness in latent representations and predicted expression, motivating direct enrichment of perturbation-specific information.Source: §1 · p. 2, Figure 1 · p. 2
• PerturbEnergy improves perturbation-specific information by up to 83.2%, achieves state-of-the-art top-50/100/200 DEG prediction, and remains competitive on HVG prediction.Source: §1 · pp. 2–3
• The conditional energy-based prior assigns lower energy to compatible representation–perturbation pairs, tilting the latent prior toward perturbation-preferred regions.Source: §1 · p. 2

• Linear models predict responses through expression-space regression or representations from foundation models such as scGPT.Source: §2 · p. 3
• VAE-based methods represent perturbation effects as additive latent factors or sparse latent representations, as in scGen, Biolord, and SAMS-VAE (Bereket & Karaletsos, 2023).Source: §2 · p. 3
• Other methods use GNNs, Transformers, diffusion, or flow matching to model perturbations or transitions between control and perturbed representations.Examples include GEARS, STATE, and diffusion- or flow-based approaches (Yuan et al., 2026).Source: §2 · p. 3
• PerturbedVAE (Jiang et al., 2026), stVAE, and CPA use contrastive or adversarial objectives, but these provide implicit constraints without a theoretical guarantee of preserved perturbation-specific information.Source: §2 · p. 3

3.1 PRELIMINARY

• The model separates basal state from perturbation-specific variation, addressing whether latent representations retain enough perturbation information for response prediction.Preliminary comparisons associate richer perturbation-specific information with better prediction performance.Source: §3.1 · pp. 3–4

3.2 CONDITIONAL ENERGY-BASED PERTURBATION PRIOR AND ENERGY-ELBO

• A conditional energy-based prior assigns lower energy to perturbation-compatible representations, and its Energy-ELBO is equivalent to optimizing the complete model ELBO.The prior explicitly tilts the latent space toward perturbation-preferred regions while retaining regularization to the conventional prior.Source: §3.2 · pp. 4–5

3.3 OPTIMIZATION FOR PERTURBENERGY

• Because the partition function is generally intractable, the method uses mismatched perturbations or Langevin dynamics to obtain negative-phase samples.At inference, it samples perturbation representations and selects the one with the lowest energy before decoding the response.Source: §3.3 · pp. 5–6

3.4 DISCUSSION

• The authors prove that optimizing PerturbEnergy increases a variational lower bound on conditional mutual information between perturbation and its latent representation.The bound concerns I(A; Za|Zb, C).Source: §3.4 · p. 6

4.1 EXPERIMENTAL SETTINGS

• Experiments use Replogle, PBMC, Norman2019, and Tahoe100M, evaluate HVG expression and DEG responses, and implement PerturbEnergy with a standard VAE backbone.The default Langevin settings are K = 20 and η = 10−4, and inference uses M = 8 candidates.Source: §4.1 · p. 7

4.2 EXPERIMENTAL RESULTS AND DISCUSSIONS

• Across four datasets, PerturbEnergy is competitive on HVG correlation and distance metrics and consistently outperforms baselines on DE metrics.The reported DE metrics include DEPrec, AUROC, AUPRC, and DEOverlap.Source: §4.2 · p. 7, Figure 3 · p. 8
• PerturbEnergy achieves the best DEG-group performance on Replogle, PBMC, and Tahoe100M and the second-best performance on Norman2019.The reported comparison uses RMSE over the true top-K DEGs.Source: §4.2 · p. 8, Table 2 · p. 9
• PerturbEnergy is best in latent-representation and expression-difference classification, with latent-space improvements of 83.2% on Replogle and 28.2% on PBMC.Source: §4.2 · p. 8, Table 3 · p. 9
• On Replogle, removing energy degrades performance; more Langevin steps and smaller step size generally help top-50 DEG RMSE but larger K slows training, while M ≤4 hurts and 16 or 32 add little.The study compares ablation, Langevin steps and step size, candidate number, and training efficiency.Source: §4.2 · pp. 8–9, Figure 4 · p. 9

• PerturbEnergy reshapes the latent prior toward perturbation-preferred regions and is reported to increase mutual information between representations and perturbations.Its effectiveness is validated experimentally on four datasets.Source: §5 · p. 10

• Generative AI tools were used to polish writing, fix grammar, paraphrase for clarity, and suggest experimental table and figure placement.Source: §6 · p. 10
• The authors state that AI was not used to develop the research idea or theory, review related work, collect data, or implement the method code.Source: §6 · p. 10

• The work applies AI-based prediction to biological cellular perturbation responses, aiming to improve accuracy for common perturbations.Source: §7 · p. 10
• The authors report no ethical issues, risk data release, harmful content, or conflicts of interest.Source: §7 · p. 10
• They state that the method should be properly applied and not used to assist harmful biomedical research.Source: §7 · p. 10

• The authors state that they provide rigorous theoretical proofs for each theorem and lemma.Source: §8 · p. 10
• The work's code is publicly available via an anonymous link.Source: §8 · p. 10
• The experiments use publicly available datasets, data-processing methods, and evaluation metrics, which the authors say supports reproducibility of the experimental and theoretical results.Source: §8 · p. 10

• The proof starts from the ELBO of log pΘ(x|c, a) and applies the joint-prior assumption and variational-inference formulation.Source: §A.1.1 · p. 14
• The derivation invokes the KL-divergence chain rule and rewrites the energy-prior expression as a KL term, expected energy, and log normalizer.Source: §A.1.1 · p. 14
• The derivation gives the complete ELBO in Equation 25 and identifies the resulting expression as the Energy-ELBO, proving Theorem 1.Source: §A.1.1 · p. 14

A.1.3 DERIVATION OF THE GRADIENT TO ENERGY PARAMETERS

• The ELBO gradient for energy parameters α is the expected difference between energy gradients under the energy posterior and the encoder distribution.The derivation removes the α-independent KL term and combines the energy expectation with the partition-function gradient.Source: §A.1.3 · p. 15

A.1.4 PROOF OF THEOREM 2: CONDITIONAL MUTUAL-INFORMATION LOWER BOUND

• Energy optimization increases a lower bound on I(A; Za|Zb, C) by assigning lower energy to compatible perturbation representations and higher energy to incompatible ones.The bound is H(A|Zb, C) + E[log rα(A|Za, Zb, C)] and is tight when rα correctly approximates the real perturbation posterior.Source: §A.1.4 · p. 16

A.2.2 OPTIMIZATION ALGORITHM

• Langevin sampling initializes latent perturbation representations from the encoder and iteratively updates them with Gaussian noise; optimization uses these or random-perturbation samples.The optimization algorithm updates energy parameters using the energy objective and updates the remaining model parameters using the ELBO gradient.Source: §A.2.2 · pp. 16–17

A.2.3 INFERENCE ALGORITHM

• At inference, the model scores candidate latent perturbation representations by energy, selects the lowest-energy candidate, and decodes it with the basal representation.The inference algorithm takes an unperturbed cell, its context, and a target perturbation, and returns a predicted cellular response.Source: §A.2.2 · p. 17

A.3.1 DATASET DESCRIPTIONS AND SETTINGS

• The benchmarks cover genetic, cytokine, drug, and combinatorial perturbations across Replogle, PBMC, Tahoe100M, and Norman2019.Norman2019 evaluates combinatorial generalization: 39 of 131 gene pairs are held out while all 105 single-gene perturbations remain in training, so each test pair's constituent genes are known.Source: §A.3.1 · pp. 17–19
• For Replogle, PBMC, and Tahoe100M, held-out contexts test adaptation to new contexts; Norman2019 instead holds out perturbation combinations.All datasets use 2,000 highly variable genes; PBMC and Tahoe100M use library-size normalization and log1p transformation, and Replogle uses processed data. Validation uses 2% of training cells, while controls remain in the training pool.Source: §A.3.1 · pp. 18–19

Baseline models

• The comparison includes linear and pretrained-transformer baselines, VAE, diffusion, flow, and transformer methods, trained under shared data splits and condition sampling.Models share per-dataset step budgets and Adam optimization; checkpoints are selected by validation MSE every 2,500 steps, using averaged weights for EMA models.Source: Baseline models · pp. 19–20, Training settings · p. 20, Validation and checkpoint selection · p. 21

Metric Definitions

• Metrics are computed per test condition and averaged, with DEG sets defined by Wilcoxon rank-sum tests between real perturbed and same-context control cells.The evaluation uses 2,000 highly variable genes; genes are ordered by ascending p-value with absolute expression change, and Sk denotes the top-k genes.Source: Metric Definitions · p. 21

HVG Group Metrics; DEG Group Metrics

• The metrics assess expression-level and perturbation-effect correlations, direction-of-change agreement, and prediction-to-expression distances; DEG evaluation computes RMSE, MAE, PCC, and R2 on true top-k DE genes (k = 50, 100, 200).Because most of the transcriptome remains nearly untouched, genome-wide averages can miss the response; with the top-k set fixed from real data, DEG evaluation no longer asks whether responding genes were found.Source: HVG Group Metrics · pp. 21–23, DEG Group Metrics · p. 23

A.3.4 OTHER VAE BACKBONE INTEGRATION

• Adding PerturbEnergy reduces scGen top-50 DEG RMSE by 18.6% on Replogle and 24.1% on Tahoe100M, with smaller gains for PerturbedVAE.For scGen, the reported RMSE changes are 0.2336 to 0.1902 on Replogle and 0.4462 to 0.3387 on Tahoe100M; the integrations add only the energy-based optimization term while preserving the original architectures and objectives.Source: §A.3.4 · pp. 23–24

A.3.5 INFLUENCE OF TRAINING SCALE AND MODEL DIMENSION

• Insufficient training steps degrade top-50 DEG RMSE across datasets and latent dimensions, while dimensions from 256 to 1024 generally perform best.The passage notes that d = 512 generally gives the best Replogle performance and that larger dimensions do not necessarily improve perturbation-effect capture.Source: §A.3.5 · pp. 24–25

A.3.6 VISUALIZATION

• PerturbEnergy representations separate perturbations more clearly than original expression, scGen, or PerturbedVAE representations in Tahoe100M visualizations.Some perturbations remain mixed in PerturbedVAE representations, including Betamethasone dipropionate and Triamcinolone.Source: §A.3.6 · p. 25, Figure 6 · p. 25

Results on Replogle Dataset; Results on Tahoe100M Dataset

• On Replogle, PerturbEnergy generally leads on HVG correlation and DEG-related scores and remains competitive on distribution distances; Tahoe100M shows the same general pattern.These results concern the 2,000 highly variable genes on Replogle and Tahoe100M, respectively.Source: Results on Replogle Dataset · pp. 25–26, Results on Tahoe100M Dataset · p. 26, Table 10 · p. 29

A.3.7 OTHER METRICS ON DEG GROUP

• PerturbEnergy ranks in the top two for DEG prediction on every dataset except Norman2019; its top-50 MAE improves 16.8% on Replogle and 11.4% on PBMC.On Replogle, PerturbEnergy-Mis lowers MAE from 0.1675 (Squidiff) to 0.1434; on PBMC, PerturbEnergy-Lang lowers it from 0.0511 (stVAE) to 0.0453. The passage also reports variable baseline performance across datasets and top-50 MAE values of 0.1812 for PerturbedVAE on Replogle and 0.0703 on PBMC.Source: §A.3.7 · p. 25, Table 7 · p. 26, Table 8 · p. 27

Results on PBMC Dataset

• On PBMC, PerturbEnergy outperforms or matches other baselines on correlation and distance measures, and its DEG-related score exceeds VAE baselines (0.2473 vs 0.2208).The passage attributes this benefit to modeling perturbation-specific information.Source: Results on PBMC Dataset · p. 26, Table 13 · p. 31, Table 14 · p. 32

Results on Norman2019 Dataset

• On Norman2019, diffusion methods perform best on HVG correlation or distance measures, while PerturbEnergy remains best on DEG-related measures.The passage notes that Norman2019 generally contains dual-genetic perturbations and proposes that diffusion captures their complex co-influenced patterns.Source: Results on Norman2019 Dataset · p. 26, Table 16 · p. 33, Table 17 · p. 34