SAVE: A Generalizable Framework for Multi-Condition Single-Cell Generation with Gene Block Attention

TL;DR

SAVE is a Transformer-based framework that models multi-condition single-cell gene expression by grouping genes into blocks, enabling better generalization and simulation of unseen biological scenarios.

Contribution

It introduces gene block attention and a flow matching mechanism to improve multi-condition single-cell modeling and generalization to new conditions.

Findings

Outperforms state-of-the-art methods in generation fidelity.

Enables extrapolative generalization to unseen conditions.

Effective in low-resource and combinatorially held-out settings.

Abstract

Modeling single-cell gene expression across diverse biological and technical conditions is crucial for characterizing cellular states and simulating unseen scenarios. Existing methods often treat genes as independent tokens, overlooking their high-level biological relationships and leading to poor performance. We introduce SAVE, a unified generative framework based on conditional Transformers for multi-condition single-cell modeling. SAVE leverages a coarse-grained representation by grouping semantically related genes into blocks, capturing higher-order dependencies among gene modules. A Flow Matching mechanism and condition-masking strategy further enhance flexible simulation and enable generalization to unseen condition combinations. We evaluate SAVE on a range of benchmarks, including conditional generation, batch effect correction, and perturbation prediction. SAVE consistently outperforms state-of-the-art methods in generation fidelity and extrapolative generalization, especially in low-resource or combinatorially held-out settings. Overall, SAVE offers a scalable and generalizable solution for modeling complex single-cell data, with broad utility in virtual cell synthesis and biological interpretation. Our code is publicly available at https://github.com/fdu-wangfeilab/sc-save