Technical Summary: DeMixPert

Problem Statement

Predicting transcriptome-wide responses to unseen genetic perturbations in single-cell data remains a significant computational challenge. Accurate prediction requires simultaneously recovering two distinct phenomena: perturbation-specific transcriptional shifts and heterogeneous cellular responses.

Existing methods often fail because they entangle the deterministic response structure with stochastic population-level variation. This entanglement causes dominant shared patterns (systematic variation driven by the basal cellular state) to mask weaker, perturbation-specific signals. Consequently, standard evaluation metrics can be dominated by shared variation rather than the specific effects of the perturbation, impairing Out-Of-Distribution (OOD) generalization. Furthermore, predicting responses under experimentally unobserved conditions requires disentangling three intertwined components:

1. Systematic Response: A basal-state-dependent structure shared across conditions.
2. Perturbation-Specific Response: Unique transcriptional effects induced by specific targets.
3. Population-Level Variation: Condition-specific subpopulations and their relative abundances.

Current approaches typically model the full perturbation response as a monolithic target, failing to explicitly separate these components or model the complex, condition-dependent distribution of cellular heterogeneity. While prior work has addressed individual aspects, most do not jointly disentangle basal-state-dependent systematic responses, perturbation-specific responses, and condition-dependent population-level variation.

Methodology: DeMixPert

The authors propose DeMixPert (Decomposed Response Modeling with Gaussian Mixtures), a framework designed to disentangle these components and model OOD single-cell perturbation responses. The architecture consists of three primary modules:

1. Deterministic Response Center Learning

To prevent systematic variation from obscuring perturbation-specific effects, DeMixPert constructs a deterministic response center ($\bar{r}_i$) by summing two separately learned components:

• Systematic Response ($g_i$): Derived from the basal state ($s_i$) of the control cell. This captures patterns shared across perturbation conditions that depend on the cellular context.
• Perturbation-Specific Response ($a_p$): Inferred from pretrained target embeddings ($e_p$) derived from scGPT. This component allows the model to generalize to unseen targets by leveraging biological relationships learned during pretraining.
• Formulation: $\bar{r}_i = g_i + a_p$. This center represents the expected deterministic shift in the latent response space.

2. Gaussian Prototype Invertible Network

To model the remaining population-level variation ($\epsilon_i$) around the deterministic center, DeMixPert employs a Gaussian prototype Invertible Network.

• Global Prototypes: A set of $K$ diagonal Gaussian prototypes ($\mathcal{N}(\mu_k, \text{diag}(\sigma_k^2))$) are learned globally from training residuals using an Expectation-Maximization (EM) algorithm. These represent reusable variation patterns.
• Adaptive Weighting: Instead of fixed weights, the model uses a Multi-Layer Perceptron (MLP) to predict condition-dependent mixture weights ($\alpha_i$) based on the concatenation of the basal state ($s_i$) and the perturbation embedding ($a_p$).
• Invertible Transformation: A conditional invertible network ($T_\theta$) maps a source variable sampled from the adaptive Gaussian mixture to the final population-level variation distribution. This allows for both forward generation of variations and inverse projection of observed variations for likelihood supervision.

3. Expression Reconstruction

The final predicted response latent is $\hat{r}_i = \bar{r}_i + \epsilon_i$. This is concatenated with the basal state representation ($s_i$) and decoded by an expression decoder ($D_\psi$) to reconstruct the post-perturbation gene expression profile ($\hat{x}^p_i$).

Training Objective

The model is trained end-to-end using a composite loss function:

• Response Alignment Loss ($L_{align}$): MSE between predicted and target response latents.
• Population-Variation Likelihood Loss ($L_{gm}$): Negative log-likelihood of target variations under the conditional Gaussian mixture, supervising both the mixture weights and the invertible transformation.
• Gene-Expression Reconstruction Loss ($L_{rec}$): MSE for reconstructing perturbed and control expressions.
• Distribution Loss ($L_{dist}$): Energy distance to align the generated and observed expression distributions at the population level.

Key Contributions

1. Decomposed Response Framework: DeMixPert proposes a response model for OOD single-cell perturbation prediction that jointly disentangles deterministic response structure from stochastic cellular heterogeneity and captures condition-specific response distributions using adaptive Gaussian mixtures.
2. Gaussian Prototype Invertible Network: The authors introduce a novel mechanism for modeling population-level variation using reusable Gaussian prototypes with context-adaptive weights. This allows the model to capture complex, condition-specific heterogeneity without assuming a unimodal distribution.
3. Superior OOD Performance: Quantitative analyses demonstrate that DeMixPert effectively recovers perturbation-specific responses while preserving population-level distributional fidelity, outperforming existing methods in unseen perturbation settings.

Experimental Results

DeMixPert was evaluated on four widely used single-cell genetic perturbation datasets: Adamson and Papalexi (single-gene perturbations) and Norman and Replogle (combinatorial perturbations). The evaluation utilized condition-disjoint splits to test generalization to unseen targets.

• Performance Metrics: The model was assessed using metrics for differential expression recovery (C-DEGs, DES), distributional agreement (Energy Distance, Wasserstein Distance), expression error (MSE), and perturbation discrimination (PDS).
• Key Findings:
  • On single-gene datasets (Papalexi, Adamson), DeMixPert achieved the highest C-DEGs scores (22.000 and 10.733) and the lowest Energy Distance values (0.171 and 0.448), significantly outperforming baselines like GEARS, scGPT, and STATE.
  • On combinatorial datasets (Norman, Replogle), DeMixPert achieved the highest Perturbation Discrimination Scores (0.981 and 0.938) and lowest Energy Distances (0.374 and 0.155).
  • Ablation Studies: Removing the decomposition mechanism (w/o Decomp) or the systematic component (w/o Sys) significantly degraded performance in Perturbation Discrimination Score (PDS) and Centroid Accuracy, confirming that separating these components is crucial for distinguishing perturbation-specific signals.
  • Interpretability: Visualizations confirmed that removing the systematic component ($g_i$) from the response representation increased cluster purity, indicating that the decomposition successfully isolates perturbation-specific structures from basal-state variation. Furthermore, a significant positive correlation was found between the divergence of learned Gaussian prototype weights and the energy distance between population distributions, validating the interpretability of the mixture composition mechanism.

Significance

The paper claims that DeMixPert addresses a critical limitation in current perturbation prediction methods: the entanglement of shared systematic variation with specific perturbation effects. By explicitly decomposing the response and modeling population heterogeneity through adaptive Gaussian mixtures, DeMixPert achieves superior generalization to unseen perturbations. The framework not only improves predictive accuracy but also offers a structured interpretation of heterogeneous cellular responses, separating transferable basal-state patterns from unique perturbation signatures. This approach provides a more robust foundation for in silico estimation of cellular responses in experimental settings where exhaustive measurement is infeasible.