SCALE:Scalable Conditional Atlas-Level Endpoint transport for virtual cell perturbation prediction

Abstract

Virtual cell models aim to enable in silico experimentation by predicting how cells respond to genetic, chemical, or cytokine perturbations from single-cell measurements. In practice, however, large-scale perturbation prediction remains constrained by three coupled bottlenecks: inefficient training and inference pipelines, unstable modeling in high-dimensional sparse expression space, and evaluation protocols that overemphasize reconstruction-like accuracy while underestimating biological fidelity. In this work we present a specialized large-scale foundation model SCALE for virtual cell perturbation prediction that addresses the above limitations jointly. First, we build a BioNeMo-based training and inference framework that substantially improves data throughput, distributed scalability, and deployment efficiency, yielding 12.51* speedup on pretrain and 1.29* on inference over the prior SOTA pipeline under matched system settings. Second, we formulate perturbation prediction as conditional transport and implement it with a set-aware flow architecture that couples LLaMA-based cellular encoding with endpoint-oriented supervision. This design yields more stable training and stronger recovery of perturbation effects. Third, we evaluate the model on Tahoe-100M using a rigorous cell-level protocol centered on biologically meaningful metrics rather than reconstruction alone. On this benchmark, our model improves PDCorr by 12.02% and DE Overlap by 10.66% over STATE. Together, these results suggest that advancing virtual cells requires not only better generative objectives, but also the co-design of scalable infrastructure, stable transport modeling, and biologically faithful evaluation.

Figures (3)

• Figure 1: Overview of SCALE.(a) Conceptual formulation of virtual cell perturbation prediction. Given control and perturbed single-cell populations measured under matched biological context, the task is framed as conditional transport from control to perturbed cell states. Rather than assuming one-to-one correspondence between cells, the model learns population-level perturbation-induced transitions over unordered cellular sets. (b) SCALE for scalable and biologically faithful perturbation prediction. SCALE encodes cellular populations into a set-aware latent space and learns stable endpoint-aligned transport under perturbation conditions. This design is coupled with scalable training infrastructure and biologically grounded evaluation to support large-scale virtual cell modeling.
• Figure 2: SCALE module detail. (a) Hierarchical set-aware encoder. This module firstly encodes gene-level cell inputs into latent cell embeddings through stacked set-transformer blocks and aggregate the cell-wise information with a set equivariant fusion layer.(b) Conditional JiT-based velocity field network. This module integrates condition features, time encoding, and cell latent embeddings to model perturbation latent transportation and predict target embeddings through velocity fitting.
• Figure 3: Ablation study. We analyze three key design choices in SCALE: condition pooling/fusion, JiT training formulation, and the design of prior distribution. (a) Adaptive condition aggregation is critical, as mean pooling causes a clear performance drop, while basic and seed attention yield substantially better results. (b) JiT parameterization consistently improves over the base flow formulation, with endpoint-oriented training (x-pred/x-loss) achieving the strongest overall performance. (c) A Gaussian-control mixed prior distribution further improves transport learning, whereas masking-based variants severely degrade performance, suggesting that increasing path difficulty can be beneficial but overly destructive corruption is harmful in the latent space.