Multi-Modal and Multi-Attribute Generation of Single Cells with CFGen

We thank Reviewer FqvM for investing the time in reviewing our work. We are glad to read positive comments about our presentation and contribution. We also appreciate the constructive criticism and feedback we received. We hope our clarifications and changes to the manuscript will contribute to a more positive assessment of our manuscript.

Handling discrete count data via negative binomial distribution is presented as a "contribution" of this paper. Why?

We apologize if our phrasing generated some lack of clarity. We will elaborate more on our contribution here. Our goal is not to claim that we are the first to introduce the concept of modeling scRNA-seq and ATAC-seq data with discrete distributions, this is standard practice in many settings. However, most existing models leveraging such a likelihood formulation are VAE-based approaches (like scVI) with a strong focus on learning biologically meaningful and batch-free representations rather than optimizing for synthetic data generation. In other words, scVI and its relatives do not aim at fully regularizing the latent space to standard Gaussian and allow the retaining of biological structure for interpretable representations. This, however, hinders the generation potential.

On the other side of the spectrum, there are models like scGAN and scDiffusion which, unlike scVI, optimize for pure generation. Such models do not factor the probabilistic properties of the data into their models and rely on powerful generative schemes to learn a continuous approximation of the data. The lack of consideration of such properties is a disadvantage when modeling complex domains like discrete biological data since they exhibit skewness, sparsity and overdispersion that do not resonate well with continuous models.

Our CFGen contributes to implementing the best of both worlds. The model is founded on rigorous considerations of the data characteristics (cells are modeled using a negative binomial likelihood) while optimizing for synthetic data generation (unlike common scVAE models) and still providing representation-learning-based tasks available in scVI, like batch correction. In summary, our contribution is not the choice of the likelihood itself, but rather its combination with a versatile flow-based generative framework encompassing both generation and manipulation of single-cell representations.

According to the paper, "... the proposed factorisation is novel". In the factorisation of Eq. 5, what is the rational behind conditioning the latent factor z on library size?

To answer this question, we start comparing how we handle the size factor to the scheme implemented by models like scVI. Single-cell VAEs provide an option to learn a log-normal distribution over the size factor that one can sample from. This is achieved similarly as we propose, namely by fitting a log-normal distribution over the library size in the data (it can also be conditional on the batch variable). The main difference between our approach and scVI is that in the latter the latent cellular state variable and the size factor are sampled independently, in CFGen we provide the option to bias the sampling by the size factor. In other words, we sample a latent state $\mathbf{z}$ that accounts for cell size, while scVI does not.

Why is this important? To generate from scVI, one would sample a batch, sample a size factor conditioned on the batch, sample a latent code independently and then generate a cell using the conditional decoder's output scaled by the sampled size factor (which is not involved in the decoding, only in the scaling). In CFGen, we would instead bias the sampling of the latent code on the library size. Now, if the size factor is somewhat uniform within the batch, the approach in scVI might work. However, if the batch key represents a coarse annotation (e.g., the source study of a dataset in modern atlases), the library size can vary significantly within the batch. This variability reflects differences inherent to the batch annotation (see Fig. A18 for examples using the Human Cell Atlas dataset). Sampling a cell state and library size independently in such cases could result in scaling a decoded cell state by an incompatible size factor, producing unrealistic outcomes. Instead, conditioning on the size factor is more appropriate, as it biases the sampling of the latent state toward regions exhibiting a specific size factor. This approach can be combined with coarsely annotated variables for a more targeted conditional generation.

In summary, conditioning on size factor steers the generated state towards regions of the cell space where scaling for such size factor makes sense.

Conditional independence assumption.

Thank you for the question. First, we would like to point out that the definition in Eq. 5 is the joint distribution for the generative model conditioned on a single attribute. The (conditional) independence assumption is a standard practice in multi-attribute generative modeling [1, 2, 3, 4], where the central mathematical formulation relies on rewriting the multi-attribute conditional probability distribution as a factorization over single-attribute conditional densities:

$p(\mathbf{z}|y_1, y_2,...,y_K) \propto p(\mathbf{z})\prod_{i=1}^{K}\frac{p(\mathbf{z}|y_i)}{p(\mathbf{z})}$.

To obtain the product, one must first assume conditional independence of the factors $y_i$ between each other given an observed $\mathbf{z}$. In the context of score-based models, this formulation allows to expression of the conditional score as a composition of an unconditional model and single-attribute conditional models (see Eq. 17 in the manuscript). In these regards, our assumption follows the line of established works in the field of generative models, where we extend the idea to Flow Matching to fit in our modeling framework (see Prop. 5 and the Appendix sections A.2 to A.4 for a detailed breakdown of the components).

To add more intuition to the above formulation, we assume that, given a state $\mathbf{z}$ and two conditioning variables $y_1$ and $y_2$, observing $\mathbf{z}$ sufficiently captures the interactions between $y_1$ and $y_2$. None of the cited works have raised concerns regarding the arising of V-structure dependencies violating the assumption. In our setting, we can intuitively think that, provided that we have a rich cell representation $\mathbf{z}$, if we can observe $\mathbf{z}$, then adding $y_1$ does not give more information about $y_2$, and vice versa. Hence, $p(y_1|\mathbf{z}, y_2)=p(y_1|\mathbf{z})$. For example, if $y_1$ is cell type and $y_2$ is batch, then we are assuming that if one can observe a cell state $\mathbf{z}$ and the batch $y_2$, the information from the variable $y_2$ does not give any additional information on $y_1$ that is not already contained in $\mathbf{z}$. Since we model a very informative $\mathbf{z}$, our results as well as the cited papers show evidence that the assumption is reasonable.

Guidance scheme.

Our method is an extension to classifier-free guidance using multiple attributes. From a theoretical perspective, the model differs from the standard classifier-free guidance approach in that it models the contribution of multiple single-attribute flows additively on the final generative model. More in detail, in Section A.4 we proved that learning and simulating the additive vector field in Eq. 8 allows us to generate observations compositionally on the chosen attributes.

Compositionality means that one can use a single model to generate observations from either all conditions, subsets thereof and unconditionally. From a training perspective, the difference from using all attributes at a time as in classifier-free guidance is indeed that each step only sees one attribute. But the differences also involve sampling, where we compose $K$ different vector fields weighted by their respective $\omega_i$. Unlike standard classifier-free guidance or simple joint conditioning by multiple variables, here we are able to tune up or down the effect of individual attributes without altering the others, which is useful in single-cell tasks like batch effect removal, where the extent of correction can be modulated by how strong the batch component is in the dataset.

Reasons for the performance of scDiffusion

Thank you for your suggestion. As we also described in our first answer, diffusion models such as scDiffusion indeed optimize for the generation task. However, in the context of scRNA-seq, some modeling choices may be responsible for hindering the performance of the generation of scRNA-seq data:

• scDiffusion does not take into account important properties of single-cell data, such as sparsity, overdispersion and discreteness. Although normalizing the data is an option to ensure continuity, most normalization methods preserve zeros in the data and a non-linear mean-variance trend. Continuous models usually benefit from centered and non-sparse input, making the structure of scDiffusion with a continuous decoder sub-optimal.
• The model relies on classifier-based guidance, therefore conditional sampling is heavily dependent on the performance of the classifier on individual labels. This structurally challenges the application of scDiffusion to rare cell type generation or any guiding schemes using hard-to-classify attributes.
• For very smaller datasets like PBMC3k, training SDE-based diffusion models is empirically complex and unstable.

In our framework, all these aspects are overcome by training a latent Flow Matching model with a discrete likelihood scheme and classifier-free guidance. We add a description of potential limitations in the scDiffusion model to Appendix D.5.

[1] Liu, Nan, et al. "Unsupervised compositional concepts discovery with text-to-image generative models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[2] Shi, Changhao, et al. "Exploring compositional visual generation with latent classifier guidance." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[3] Du, Yilun, Shuang Li, and Igor Mordatch. "Compositional visual generation with energy based models." Advances in Neural Information Processing Systems 33 (2020): 6637-6647.

[4] Liu, Nan, et al. "Compositional visual generation with composable diffusion models." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.

The authors addressed all of my major concerns. So I'm increasing my score from 5 to 6.

We are thankful for the time DU3L invested in reviewing our paper. We are glad to hear the reviewer found our contribution solid and the paper well written. We appreciate the constructive criticism and updated our paper in the direction of the suggested changes.

Fig 3. is not really clear to me. Firstly, I suggest adding contrasting colors for points representing generated and real data. Secondly, what are the red points representing?

We acknowledged the reviewer's point of view and modified the figure accordingly. We introduced a stronger red-blue contrast between real and generated cells and used the same color scheme consistently across the whole figure to represent real and generated cells. Now both in the zero-guidance case and in the guided generation scheme blue points represent real cells and red points generated cells. We also made the legend bigger.

We briefly summarize the plot in case it remains unclear. On the left, enclosed by the boxes, we show the qualitative generative performance of the model when setting guidance parameters to 0. Such an approach allows a faithful modeling of the whole single cell distribution since the flow is not guided by any attribute. Next to the boxes, we show the effect of increasing a single guidance parameter over the counterpart, hence generating cells with increased specificity from an intersection of attributes. The middle plots with points highlighted in red are the generated cells for a certain guidance scheme (hence, specific guidance weights). On the opposite sides of such plots, with cells colored in blue, we highlight the real points coming from the two categories used for guidance. We include the latter plots as a reference since we expect guidance parameters to become increasingly good at modeling intersections of such attributes.

I also suggest perhaps adding a quantitative metric (perhaps an oracle model that predicts the attributes) as well.

We added such a metric to Tab. 12 in the Appendix. We train as an oracle a 3-layer MLP classifier with a softmax head on the real data. Specifically, we derive one classifier per guidance attribute. Upon generation with different guiding schemes, we apply the two classifiers to the generated cells. We expect that, when the model is guided on a single attribute only (hence, the counterpart attribute has guidance strengths $\omega=0$), the oracle only assigns high probability to the class involved in guidance. This is indeed what occurs in Tab. 12, in the first row, only the cell type (NeurIPS) and tissue (C. Elegans) attributes are used for guidance. Therefore, the oracle assigns the cells generated under this scheme a high probability of being part of the guiding biological annotation class, while the probability for the mouse ID and donor classes we chose is low (the model did not use them to guide generation). Upon increasing the donor and mouse ID weights, the oracle predicts generated cells to be part of the guiding classes from both attributes with high probability.

I also suggest removing the bars from Fig. 2b as they make it hard to observe the overlapping density curves which are easier to infer from.

We modified the plot according to the suggestion. The new version of Fig.2 can be found in the updated manuscript.

For Sec 5.2, it might be worthwhile to also add a comparison with CFGen just trained on RNA-data in order to measure the effects of using multimodal data for training.

We added CFGen trained on RNA data only to Tab. 2. The model achieved top performance in terms of RBF-MMD and second-best performance (below its multimodal counterpart) in terms of Wasserstein-2 distance. Hence, multi-modal generation on PBMC10k performs on par with its RNA-only counterpart. From our results, we infer that modeling multi-modal data does not hamper the performance of our model despite the increased amount of information to synthesize.

We acknowledge the Reviewer's feedback and are thankful for the positive comments on the quality of the work. We particularly value the received criticism, as it pushed us to extend the scope of our contribution to the gene imputation task. We hope Reviewer eqfJ will find our new results insightful.

The authors do not discuss the computational complexity of the proposed method. A more detailed breakdown of computational requirements, including training and sampling times for the proposed method and the baselines, would improve the paper.

We included a new section to the Appendix (Section H.1) reporting a detailed breakdown of the runtime and computational complexity of our model. More specifically, we added information on CFGen's runtime in Fig. A1, Tab. 5 and Tab. 6. We describe the insights derived from our new results here:

• Fig A1 illustrates how the generation runtime changes as a function of hyperparameters. Since CFGen is a latent generative model, we consider different runtime curves for different latent space sizes as a function of the following hyperparameters. The size of the latent space significantly influences the runtime, since it establishes the size of the state space where the generative ODE is simulated. The sampling runtime increases approximately linearly with respect to the number of cells and genes, while neural-network-related hyperparameters appear not to impact the simulation time significantly.
• In Tab. 5 We compare training times between the different models across benchmarked datasets. Of course, smaller VAE-based models like scVI are faster to train than diffusion and Flow-Matching-based counterparts. Between CFGen and scDiffusion, the latter exhibits longer combined training times on larger datasets like Tabula Muris and HLCA, but also on PBMC10K. Importantly, scDiffusion's training requires optimizing an autoencoder, a denoising diffusion network and a cell type classifier used for guidance separately, while CFGen only trains an autoencoder and a Flow Matching model.
• Tab. 6 reports the sampling times of the trained generative models. Intuitively, VAE models are still faster, since they do not have to simulate generative ODEs and SDEs across time. However, our results also suggest that such VAEs are qualitatively and quantitatively limited on large datasets (Fig. A7) and downstream applications such as augmentation (Fig. A11) and multi-modal marker generation (Fig. A8-A9). CFGen can generate atlas-level datasets of >500k cells in around 8 seconds (see the HLCA column). Conversely, scDiffusion requires approximately half an hour for the task and is therefore not a sensible candidate for augmenting datasets with millions of cells. We highlight the aspects that make CFGen faster and more performing than scDiffusion in Appendix D.5.

We want to thank the reviewer for their feedback and suggestions, which will surely contribute to improving the quality of our experimental validation. We also thank the reviewer for their positive comments about our contribution.

scVI should be included as a baseline in Figure 2.

We include the comparison with scVI in Appendix H.2 of the revised manuscript. More specifically, we provide two new pieces of information:

• The plot comparing the distribution of the number of zeros per cell and mean-variance trend between real and generated cells by CFGen and scVI (Fig. A2).
• Quantitative metrics to evaluate how well different models approximate the sparsity and overdispersion of the real data (Tab. 7).

As you can see from Fig. A2, both scVI and CFGen model the sparsity and overdispersion characteristics solidly in the considered datasets. This is expected since both use a decoding mechanism optimized under a negative binomial likelihood model. However, note that capturing the sparsity and overdispersion trend does not necessarily boil down to modeling the whole transcription state properly. Indeed, in Fig. A7 (Appendix) we show that scVI struggles to retrieve a realistic single-cell representation on larger datasets, due to the lower flexibility of its generative model (see Tab. 1 for quantitative metrics confirming this aspect). On the contrary, CFGen closely samples from the data distribution.

We complement Fig.2 by evaluating the following quantitative metrics for each generative model:

• The 1D Wasserstein-2 distance between the vectors of per-cell number of zeroes from real and generated data.
• The 1D Wasserstein-2 distance between the empirical mean-variance ratio vectors of real and generated cells.

Note that we have to use a distributional distance because generated cells are not the same "items" as real cells but new objects, therefore we cannot evaluate a correlation between them. The lower such metrics get, the more closely a model's sparsity and mean-variance ratio resemble the one from the data. Results in Tab. 7 report that CFGen is the best model at approximating sparsity and overdispersion on three datasets out of four.

It is unclear how the classifier guidance strength is determined.

The only downstream task where guidance is used is batch correction on the NeurIPS and C.Elegans datasets in Section 5.5. In Appendix H.9 we provide an intuition of our selection process. In batch correction, cells are transported to noise and then back again to data guided by a biological and a target batch covariate. The stronger the guidance strength parameters $\omega_{\mathrm{bio}}$ and $\omega_{\mathrm{batch}}$ the more biological conservation and batch conversion will be emphasized.

If one merely observes the scIB metric in Tab. 13 computed over different guidance strength parameters, they will choose the highest possible guidance strengths, since they provide the best aggregation within cell types and batches. However, in Fig. A16a-b and A17a-b, we show that scIB metrics could be misleading and should be accompanied by qualitative evaluation. Indeed, increasing guidance parameters too much in the translation task leads to an unnatural collapse of the variability in the data beyond the one explained by batch and biological annotations. On another note, we found that guidance strength parameters surrounding values of 1 and 2 are sufficient at both preserving signal in the data and performing correction without over-squashing the cell representations. An example of unwanted effects is presented in Fig. A16, where biological preservation results in unnatural clustering for both datasets.

Finally, it is important to consider the extent of the batch effect present in the data. For example, in C. Elegans the batch effect is mild, therefore we select the parameters $\omega_{\mathrm{bio}}=2,\omega_{\mathrm{batch}}=1$ since they provide better performance than $\omega_{\mathrm{bio}}=1, \omega_{\mathrm{batch}}=2$ and $\omega_{\mathrm{bio}}=1, \omega_{\mathrm{batch}}=1$ (Tab. 13). In the NeurIPS dataset, we observe the opposite effect and therefore select $\omega_{\mathrm{bio}}=2, \omega_{\mathrm{batch}}=1$, since as soon as $\omega_{\mathrm{bio}}>1$ we obtain the unnatural biological structure in Fig A17c, which violates smooth temporal single-cell trajectories.

In conclusion, we recommend first evaluating the extent of batch effect in the data and then sweeping over combinations of guidance weights, selecting the configuration that achieves the best scIB metric values without inducing unrealistic single-cell representations.

For batch correction, is CFGen's performance (in terms of the Batch and Bio scores) sensitive to varying the guidance parameters? How does one tune the guidance parameters in practice?

Please, see the answer above.

Quantitative results are lacking when evaluating the compositional classifier guidance in Section 5.3. The change in MMD and WD with respect to the target distribution when increasing guidance strength can suffice.

We include the suggested results in Fig. A15 (Appendix H.8). In the experiment, we increase the guidance parameters for both considered attributes in parallel from 0.1 to 2.5 and evaluate how well the generated cells approximate the real cells at the intersection of the two labels. Notably, increasing the guidance weights improves modeling the combinations of attributes, with both lower MMD and Wasserstein-2 distance values.

Does data augmentation with CFGen improve performance for a logistic regression model?

We explored such a direction in the new Fig. A12. More in detail, we apply the setting described for scGPT in Section 5.4 to CellTypist [1], a famous cell type annotation model based on logistic regression. Interestingly, a similar trend is observed for the linear classifier as the one reported for scGPT, where the rare cell type classification accuracy on unseen patients increases upon data augmentation. This is particularly evident in the PBMC covid dataset where the majority of cell types exhibit a boost in predictive performance over held-out patients. In HLCA the overall trend still favors an improvement in rare cell type classification, though less pronounced and consistent than the PBMC covid dataset. Overall, our results suggest our augmentation strategy can apply to multiple classifiers.

We would like to thank again Reviewer kEhi for their consideration and remain available for further clarifications.

[1] Cippà, Pietro E., and Thomas F. Mueller. "A First Step Toward a Cross-tissue Atlas of Immune Cells in Humans." Transplantation 107.1 (2023): 8-9.

Dear Reviewer kEhi,

Thank you once again for taking the time to provide thoughtful feedback on our work. Your insightful comments have significantly contributed to enhancing our paper in the following ways:

• Providing a more comprehensive set of baselines.
• Clarifying model selection practices in the batch correction task.
• Demonstrating promising generative performance results in the multi-attribute generation task.
• Offering additional evidence that augmentation with CFGen improves rare cell type classification using a logistic regression model.

We hope these revisions address the key points the reviewer raised and would greatly appreciate hearing whether the reviewer believes their concerns have been fully addressed. As always, we are happy to discuss any further remaining questions or suggestions.

Thank you again for reviewing our work.

Best regards,

The Authors

Dear Reviewer kEhi,

Thank you for your thoughtful feedback and acknowledging our rebuttal and contribution to the field. We sincerely appreciate your decision to increase your score to acceptance and your recognition of our work's value to computational biology.

We are also grateful for your suggestion to update the conclusion to reflect our new batch correction results. We will incorporate the provided text into the revised version of our manuscript.

Thank you again for your support and contribution to improving our paper.

Best regards,

The Authors