PertMix: Learnable Distribution Mixing for Unseen Perturbation Prediction

Predicting single-cell perturbation responses requires inferring changes in cellular states from perturbation conditions, yet existing deep learning models generalize poorly to unseen perturbations. By examining the local properties of the learned condition–response map, we identify directional derivative mismatch and anomalous curvature under discrete perturbation supervision as factors associated with generalization failure. Motivated by these findings, we introduce mixup to provide intermediate conditions and response targets between observed perturbations, regularizing the model along condition interpolation paths. However, standard mixup linearly interpolates perturbation conditions and cellular responses using the same coefficient, making it difficult to construct reliable distributional supervision aligned with the mixed conditions. We propose PertMix, a learnable distribution-mixing method for unseen perturbation prediction. PertMix adaptively constructs response distributions from perturbation representations and endpoint cell populations by learning mixing coefficients and nonlinear distribution transformations. We develop two optimization variants based on bilevel gradient optimization and policy-gradient optimization, respectively, using predictive feedback on observed perturbation responses to learn the mixing rules. Experiments on single-cell perturbation prediction show that PertMix improves prediction of unseen perturbations, demonstrating the value of learnable distribution mixing for perturbation generalization.