Paper: A General Single-Cell Analysis Framework via Conditional Diffusion Generative Models (ICLR-2024)

Reviewer comment (Experiments):
The paper mentions using LLMs as the prior, but why not compare the proposed method with LLMs-based ones, such as those in [1]?

Author response:
A. We thank the reviewer for the suggestion. We first note that **reference [1] does not consider LLMs that handle natural languages**. Instead, the LLM it referred to are large-scale pre-trained models using single-cell omics data, such as scGPT [2], scBERT [3], and Geneformer [4].

Nevertheless, to provide more detailed comparisons, we ran scGPT, and Geneformer on our annotation benchmarking datasets. The annotation accuracies across three runs are summarized in the following table. We excluded scBERT due to its fine-tuning time exceeding 24 hours. Our results indicate that scDiff performs comparably against these large-scale pre-trained models, even without pre-training on hundreds of thousands of single-cell samples. 

| ACC | PBMC12K | Pancreas | HLCA | Immune | Brain | Liver |
| --- | --- | --- | --- | --- | --- | --- |
| Geneformer | 0.9787 +/- 0.0019 | -- (no ENSG) | 0.8328 +/- 0.0057 | 0.8556 +/- 0.0098 | 0.9338 +/- 0.0071 | 0.8711 +/- 0.0390 |
| scGPT | 0.9630 +/- 0.0055 | 0.9454 +/- 0.0374 | 0.8634 +/- 0.0200 | 0.9071 +/- 0.0108 | 0.9504 +/- 0.0029 | 0.8640 +/- 0.0379 |
| scDiff | 0.9670 +/- 0.0042 | 0.9680 +/- 0.0143 | 0.8931 +/- 0.0070 | 0.8442 +/- 0.0076 | 0.9473 +/- 0.0074 | 0.8439 +/- 0.0042 |

[2] Cui, Haotian, et al. "scgpt: Towards building a foundation model for single-cell multi-omics using generative ai." bioRxiv (2023): 2023-04.

[3] Yang, Fan, et al. "scBERT as a large-scale pretrained deep language model for cell type annotation of single-cell RNA-seq data." Nature Machine Intelligence 4.10 (2022): 852-866.

[4] Theodoris, Christina V., et al. "Transfer learning enables predictions in network biology." Nature (2023): 1-9.

Reviewer comment (Presentation):
From the part of the paper under Section 2.2 to equation (9), the author used a large amount of words to introduce basics of the diffusion model. I would suggest the author to move this part to appendix as this is not the contribution of this paper.

Author response:
A. We thank the reviewer for the suggestion. Indeed, the formulation of DDPM is not our contribution. However, this brief introduction provides the necessary background and the notations for our method, making it indispensable. We understand the reviewer’s concern about the clarity of our contribution. Thus, we have adjusted the manuscript to explicitly indicate that the DDPM introduction is preliminary to our method.

Reviewer comment (Reproducibility):
For the conditioner part, how does the cross-attention exactly used? Will it be able to automatically detect which task the model is focusing on?

Author response:
A. We thank the reviewer for the question and the opportunity for us to clarify this matter. Cross-attention combines **the cell-embedding with the condition information**, e.g., cell type, cell perturbation, and gene perturbation. Specifically, the query (Q) for each sample is the cell embedding, and the key (K) and value (V) are derived from the condition representations.

The cross-attention module is **downstream task-agnostic**, meaning it will function exactly the same regardless of the intended downstream task. In fact, the training objective of scDiff is already task-agnostic (the primary highlight of scDiff): it is trained to *recover clean gene expression patterns from noised ones, given the condition information associated with the cell*.

Reviewer comment (Theory):
There are quite a few confusing part in the paper. For instance, the paper claims that it forms the task into a posterior estimation, but how the prior is used in the model? I don't think the paper explicitly explains this. Also, for cell labeling and knowledge transfer, I don't thinks it's a generation task so that generative model is a good practice in this case. Intuitively cell lebeling and knowledge transfer is more like a prediction task to me.

Author response:
A. We thank the reviewer for the questions regarding the prior and the formulation of the annotation task as posterior generation. First, the **priors are the conditions used in the scDiff model**, e.g., cell types, perturbation, and gene perturbation. They are essential for realizing the downstream tasks. For example, (1) annotation is implemented by choosing the condition that results in the lowest reconstruction error for a given input cell, and (2) perturbation is implemented by applying the perturbed condition to the normal (unperturbed) cell.

Second, regarding the cell annotation task formulation as posterior estimation, we understand that the typical approach to classification is by using a “deterministic prediction head”. However, we point out that the goal of cell labeling and knowledge transfer can indeed be formulated as *estimating the posterior of gene expression under certain cell types or unknown conditions*, which is a generative task. We further note that **formulating classification problems into generative task is not uncommon** in recent machine learning communities. A notable example is using autoregressive generative language models to solve multiple choice (multi-class classification) or answer True/False (binary classification).

Reviewer comment (Novelty):
I don't feel the strong motivation of using a attention mechanism to help the model learn specific tasks. To me it's more like the proposed method assembles three tasks very hard to let the model to learn. Therefore, the proposed method lacks novelty as each individual task can still be solved in the same framework.

Author response:
A. We thank the reviewer for the series of questions regarding cross-attention usage in scDiff. First, we clarify that we *do not assemble tasks* using the cross-attention. Instead, as the highlight of our method, scDiff is trained in a *task-agnostic* manner by reconstructing clean gene expression from the noised one, given its corresponding conditions (e.g., cell type, perturbation, and gene perturbation). Similar to latent diffusion [5], the cross-attention here servers to fuse prior (condition) information into the model while not explicitly selecting one of them.

Second, we use cross-attention to combine the cell information with the condition information due to its **effectiveness and superior performance** in the downstream evaluations. A more naive implementation to combine cell embedding and condition information is directly adding up their embeddings (followed by a feed-forward network), rather than fusing through a cross-attention mechanism. To demonstrate the superiority of cross-attention over the naive adding approach, we included an additional ablation study comparing the two for the annotation tasks. We observe that cross-attention resulted in higher accuracy than the naive adding counterpart, with the only exceptions being Liver and Pancreas.

|  | Brain | HLCA | Immune | Liver | Pancreas | PBMC12K |
| --- | --- | --- | --- | --- | --- | --- |
| Default | 0.9473 +/- 0.0074 | 0.8931 +/- 0.0070 | 0.8442 +/- 0.0076 | 0.8439 +/- 0.0042 | 0.9680 +/- 0.0143 | 0.9670 +/- 0.0042 |
| No cross-attention | 0.9376 +/- 0.0019 | 0.8870 +/- 0.0012 | 0.8318 +/- 0.0028 | 0.8456 +/- 0.0054 | 0.9722 +/- 0.0014 | 0.9608 +/- 0.0042 |

[5] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*. 2022.

Reviewer comment (Novelty):
It is not clear what the novel aspect of scDiff model is from a machine learning perspective. Although the authors introduced posterior inference as a novel unified framework for several single-cell tasks, posterior inference through variational or generative processes is a well-explored area, even within the field of single-cell analysis.

Author response:
A. Thank you for raising the concern about novelty. There are some existing works utilizing variational or generative processes in single-cell analysis [1, 2], but they are limited to specific type of tasks. To the best of our knowledge, we are the first to present the idea of handling multiple types of tasks (e.g., cell type annotation, imputation and perturbation prediction) without modifying the model structure. Our approach is novel, utilizing conditional diffusion models with prior information. We empirically find out that with task-agnostic training process and task-specific generative process, *scDiff* is able to match or even outperform the current SOTA. It is important to note that the modularized conditioning strategy can be readily extended to further tasks, as long as their objectives can be formulated as or approximated by the posterior. 

[1] Lopez, Romain, et al. "Deep generative modeling for single-cell transcriptomics." *Nature methods* 15.12 (2018): 1053-1058.

[2] Lotfollahi, Mohammad, F. Alexander Wolf, and Fabian J. Theis. "scGen predicts single-cell perturbation responses." *Nature methods* 16.8 (2019): 715-721.

Reviewer comment (Experiments):
The reported results in the experimental sections do not convincingly demonstrate that scDiff significantly outperforms relevant models for the specific tasks at hand. Furthermore, the authors have not quantified the computational cost associated with having a unified model that covers multiple tasks, particularly in relation to the observed improvements. While the enhanced performance might not necessarily stem from the model's extension, it could be attributed to factors like the use of the diffusion process, attention units, or different implementations.

Author response:
A. We thank the reviewer for pointing out the concerns. We formulate our reponses in three aspects: (i) we justify the performance of *scDiff* and provide additional benchmark results, (ii) *scDiff* has a computational complexity that is nearly the same as a standard MLP, and (iii) we validate our design of *scDiff* through ablation studies. 

We first point out that scDiff does show **convincingly superior performance**: it achieved top performance for 11 out of the 15 primary benchmarks we have carried out. For the remaining four tasks, scDiff performed comparably to the best methods. More importantly, scDiff achieved top performances in diverse tasks with a single **unified framework**, unlike other baselines that are task-specific. Our wide variety of conditioners, spanning from the simple class embedding, to LLM, to graph neural networks, also highlight the **extendability** of our method as another significant novelty.

In addition, to benchmark *scDiff* with existing leaderboards, we present results follow the setting of OpenProblems in [denoising](https://openproblems.bio/results/denoising/) and [cell type annotation](https://openproblems.bio/results/label_projection/). For annotation task, we copied the accuracy of four top-performing baselines from the [leaderboard](https://openproblems.bio/results/label_projection/) with random splits. The results are as follows, where *scDiff* is evaluated across five runs. We observe that *scDiff* matches the best performance of the leaderboard.

| ACC | Pancreas | Tabula Muris | CeNGEN |
| --- | --- | --- | --- |
| Logistic regression (log CP10k) | 0.98 | 0.92 | 0.89 |
| Seurat reference mapping (SCTransform) | 0.98 | 0.90 | 0.83 |
| Multilayer perceptron (log scran) | 0.98 | 0.92 | 0.87 |
| XGBoost (log CP10k) | 0.97 | 0.86 | 0.84 |
| scDiff | 0.9800 +/- 0.0039 | 0.9257 +/- 0.0066 | 0.8889 +/- 0.0024 |

For denoising task, we selected three performant baselines from the [leaderboard](https://openproblems.bio/results/denoising/) regarding MSE. We report the mean and variance of *scDiff* and baselines across five runs. Note that the MSE results on [leaderboard](https://openproblems.bio/results/denoising/) are scaled for better illustration, we reproduced the benchmark for a fair comparison via the [public code repository](https://github.com/openproblems-bio/openproblems/tree/main/openproblems/tasks/denoising) without scaling. A conclusion can be drawn that *scDiff* outperforms the baselines across datasets.

| MSE | PBMC | Pancreas | Tabula Muris |
| --- | --- | --- | --- |
| MAGIC | 0.1887 +/- 0.0001 | 0.2319 +/- 0.0001 | 0.1842 +/- 0.0001 |
| DCA | 0.2184 +/- 0.0003 | 0.2683 +/- 0.0003 | 0.2160 +/- 0.0007 |
| ALRA | 0.2823 +/- 0.0007 | 0.3317 +/- 0.0001 | 0.2800 +/- 0.0006 |
| scDiff | 0.1810 +/- 0.0011 | 0.2079 +/- 0.0027 | 0.1680 +/- 0.0023 |


Regarding the computational complexity, we note that ***scDiff* has a computational complexity that is nearly the same as a standard MLP**. In the following table, we report the per-epoch computation time for scDiff against standard baselines for all our benchmarking tasks. It is crucial to acknowledge that diffusion models inherently demand a relatively longer time during the inference stage compared to other generative models. In our empirical observations, we find that the inference time of *scDiff* is negligible compared to the training time. Therefore, our emphasis is placed on training speed. 

From the table, we observe that in all cases, scDiff has comparable per-epoch runtime as the baselines, especially for MLP-based ones (scANVI, DCA, and scGEN). Referring to the model structure of *scDiff*, the central backbone of it comprises the cross-attention (CA) encoder blocks, where the tokens are the number of types of conditions (e.g., cell type, perturbation status). As mentioned in Appendix E.1 Table 3, the number of types of conditions will not exceed $5$. Thus, unlike the traditional transformer model that poses $O(N^2)$ complexity, where the token number $N$ is typically much larger than the hidden dimensions, in *scDiff* the feedforward module dominates the computation complexity of our CA block.

For the space complexity, we note that the required amount of GPU memory is determined by both the batch size and the number of genes. Empirically, we set $\text{batch size} = 2048$ throughout all the experiments. On the dataset with the highest number of genes, namely HLCA with 28,024 genes, *scDiff* necessitates 4.5 GB of GPU memory. Given the above discussion on training efficiency, it is noteworthy that *scDiff* can be scaled effectively to large datasets.

|  | Dataset | scDiff per-epoch runtime (sec) | Baseline per-epoch runtime (sec) | Baseline name |
| --- | --- | --- | --- | --- |
| Annotation | Brain | 9.63 | 9.82 | scANVI |
|  | HLCA | 13.60 | 12.73 | scANVI |
|  | Immune | 5.28 | 8.86 | scANVI |
|  | Liver | 2.83 | 2.97 | scANVI |
|  | Pancreas | 5.05 | 14.37 | scANVI |
|  | PBMC12K | 2.04 | 3.31 | scANVI |
| Denoising | 293T | 0.8847 | 0.53 | DCA |
|  | Jurkat | 0.5736 | 0.60 | DCA |
|  | PBMC1K | 0.2634 | 0.16 | DCA |
| Perturbation | HPoly | 1.96 | 1.83 | scGEN |
|  | PBMC | 3.16 | 3.53 | scGEN |
|  | Salmonella | 1.58 | 1.81 | scGEN |
| Gene Perturbation | Adamson | 83.88 | 1189.66 | GEARS |
|  | Dixit | 44.67 | 760.70 | GEARS |
|  | Norman | 155.88 | 491.97 | GEARS |

Finally, we have carried out additional experiments to **systematically test the effectiveness of scDiff**, including the usage of the diffusion process, cross-attention, and embedder/decoder types. Our results highlight the importance of the diffusion process, while other minor architectural designs, such as embedder and decoder, show a less noticeable impact on the final performance.

|  |  | Brain | HLCA | Immune | Liver | Pancreas | PBMC12K |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | Default | 0.9473 +/- 0.0074 | 0.8931 +/- 0.0070 | 0.8442 +/- 0.0076 | 0.8439 +/- 0.0042 | 0.9680 +/- 0.0143 | 0.9670 +/- 0.0042 |
| Decoder | 2-layer | 0.9445 +/- 0.0041 | 0.8911 +/- 0.0042 | 0.8444 +/- 0.0070 | 0.8589 +/- 0.0238 | 0.9735 +/- 0.0023 | 0.9700 +/- 0.0045 |
|  | 3-layer | 0.9378 +/- 0.0017 | 0.8994 +/- 0.0038 | 0.8383 +/- 0.0112 | 0.8062 +/- 0.0216 | 0.9731 +/- 0.0014 | 0.9701 +/- 0.0029 |
|  | 4-layer | 0.9409 +/- 0.0064 | 0.8989 +/- 0.0037 | 0.8414 +/- 0.0051 | 0.8199 +/- 0.0185 | 0.9763 +/- 0.0011 | 0.9744 +/- 0.0002 |
| Embedder | 2-layer | 0.9394 +/- 0.0029 | 0.8857 +/- 0.0089 | 0.8424 +/- 0.0061 | 0.8434 +/- 0.0239 | 0.9720 +/- 0.0011 | 0.9667 +/- 0.0039 |
|  | 3-layer | 0.9512 +/- 0.0048 | 0.8847 +/- 0.0053 | 0.8358 +/- 0.0081 | 0.8447 +/- 0.0224 | 0.9745 +/- 0.0026 | 0.9672 +/- 0.0039 |
|  | 4-layer | 0.9434 +/- 0.0048 | 0.8927 +/- 0.0018 | 0.8461 +/- 0.0014 | 0.8423 +/- 0.0100 | 0.9742 +/- 0.0012 | 0.9606 +/- 0.0036 |
| General | No cross-attention | 0.9376 +/- 0.0019 | 0.8870 +/- 0.0012 | 0.8318 +/- 0.0028 | 0.8456 +/- 0.0054 | 0.9722 +/- 0.0014 | 0.9608 +/- 0.0042 |
|  | U-Net residual | 0.9456 +/- 0.0047 | 0.8959 +/- 0.0050 | 0.8502 +/- 0.0075 | 0.8399 +/- 0.0103 | 0.9693 +/- 0.0014 | 0.9644 +/- 0.0021 |
|  | Remove diffusion | 0.5513 +/- 0.1390 | 0.1487 +/- 0.0313 | 0.1233 +/- 0.1661 | 0.1525 +/- 0.0214 | 0.3749 +/- 0.0627 | 0.5939 +/- 0.1144 |

Reviewer comment (Experiments):
Some of the experimental settings descriptions, such as those outlined in section 4.2.1, are not clearly explained and can be rather confusing. Additional clarification is needed in this regard.

Author response:
A. Thank you for suggestion of improving the clarity of our experiments’ descriptions. We have made according changes in Appendix E. Here we summarize the experimental setting of Section 4.2.1: For a specific dataset, we refer to the cell types that contain more cells than a threshold as majority types and the remaining as minority types. We randomly sample one cell per type from the minority types as the one-shot set, and the rest of the cells in minority types are used for evaluation. *scDiff* is first pre-trained on all cells in majority types and then fine-tuned on the one-shot set for $50$ epochs.

Reviewer comment (Evaluation):
In Table 1, it would be valuable to include the number of clusters and chance level for each dataset to provide a more comprehensive understanding of the results.

Author response:
A. We have provided the number of clusters (cell types) in our appendix (see Appendix Table 4). We decided to defer the dataset statistics, such as the number of cell types to the appendix to make the main presentations more compact. Nevertheless, to better guide readers find this information, we have adjusted the main text to explicitly refer the readers to the appendix for dataset statistics.

Furthermore, we have added the chance level by predicting the majority class all the time to the annotation results table to provide better sense of the performance as suggested.

Reviewer comment (Experiments):
In the context of the missing value imputation task, the assumption that all zero-expressed genes are missing and no actual zero-expressed genes exist may not be entirely accurate. It might be more biologically relevant for the model to learn the mask matrix, $M$, instead of assuming that $x_g > 0$. In biology, we know that non or less-expressed genes can still play a marker role in some cell types.

Author response:
A. We thank the reviewer for the question regarding our benchmarking strategy for the denoising task. We clarify that we do not assume all zero-expressed genes are missing. Specifically, we randomly mask out 10% of the non-zero entries and turn them into zeros in the training data. The assumption here is that all zero-expressed genes are "true zeros," following the benchmarking setting from a well-established paper [1].

Meanwhile, we understand the reviewer's concern about the simplistic assumption about zero-expressed genes. Thus, we have carried out an additional set of denoising benchmarks using molecular cross-validation (MCV) [2] to comprehensively evaluate our proposed method's biological relevance. MCV is a more rigorous data-splitting strategy that partitions the molecules in a cell into two parts. According to the experimental results of the newly added OpenProblems benchmarks in the replies to W2, we showed that scDiff outperforms the top performance compared to the baseline approaches.

[1] Van Dijk, David, et al. "Recovering gene interactions from single-cell data using data diffusion." *Cell* 174.3 (2018): 716-729.

[2] Batson, Joshua, Loïc Royer, and James Webber. "Molecular cross-validation for single-cell RNA-seq." *BioRxiv* (2019): 786269.

Reviewer comment (Theory):
Can you provide further elaboration on Equation 10? It is not entirely clear why the single-cell data are encoded using the suggested *"TimeEmbed"* function.

Author response:
A. In a nut shell, the *TimeEmbed* function is used to encode the noise level of the diffusion process. Diffusion denoising probablistic models (DDPM) [1] works by gradually denoising a signal starting from a normal Gaussian noise. Informally, the TimeEmbed function serves to inform the model at which stage the denoising process is at, and whether the current input signal is expected to be more or less noisy. We follow this TimeEmbed strategy from DDPM by adding the time embedding to the raw expression embedding.

[1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." *Advances in neural information processing systems* 33 (2020): 6840-6851.

Reviewer comment (Theory):
Why does scDiff utilize a linear encoder / decoder?

Author response:
A. We use linear Embedder and Decoder for simplicity. Conversely, the Encoder is implemented as cross-attention blocks that integrate the encoded cell embedding with the condition information. To better illustrate the effectiveness of linear embedder and decoder, we have included a new ablation study that replaces linear embedder/decoder with multi-layer perceptron (MLP) counterparts (see ablation study results in the response to W2). Our results show that although MLP embedder/decoder improves the performance over the linear counterparts in some datasets, their performance is comparable overall, indicating the effectiveness of our original design choice.

Reviewer comment (Experiments):
It would be insightful to understand the computational efficiency of using a model like scDiff compared to an equivalent model designed solely for solving one or a few downstream tasks.

Author response:
A. We thank the reviewer for suggesting analyzing the computational cost of scDiff against standard baselines. We note that **scDiff has a computational complexity that is nearly the same as a standard MLP.** The results are summarized in the reply to W2. We briefly explain this claim in the following and show empirical results supporting it.

The central backbone of scDiff comprises the cross-attention (CA) encoder blocks, where the tokens are the number of conditions (e.g., cell type, cell perturbation, gene perturbation). Thus, unlike the traditional transformer model that poses $O(N^2)$ complexity, where the token number $N$ is typically much larger than the hidden dimensions, the feedforward module dominates the computation complexity of our CA block.

Reviewer comment (Evaluation):
In Table 2a, for the missing value imputation task, the authors reported correlation values. However, is not the primary goal to approximate the value of gene expression rather than capturing the overall expression pattern (correlation at the gene population level)? It might be more informative if the authors report the average (normalized) error.

Author response:
A. As we alluded to in our response to Q1, we follow the benchmarking strategy described by the original paper [1], which suggested Pearson correlation for evaluation. Nonetheless, in the newly added MCV experiments, we followed OpenProblems and reported MSE metrics. In either case, scDiff shows comparable or better performance than the baselines.

[1] Van Dijk, David, et al. "Recovering gene interactions from single-cell data using data diffusion." *Cell* 174.3 (2018): 716-729.

Reviewer comment (Experiments):
In Table 2, it would be beneficial to include the number of genes used for each task per dataset.

Author response:
A. We thank the reviewer for the suggestion. We have already provided basic dataset statistics in the Appendix (see Table 4), including the number of genes, batches, cell types, and conditions. We further note that we did not apply any feature selection except for filtering out all zero genes. To better guide readers to find this information, we have adjusted the main text under the experiment section to highlight the availability of this information.

Reviewer comment (Presentation):
In Figure 3, why the top and bottom subfigures do not reveal the same relative performance pattern? Could you provide further elaboration on how they are related to each other?

Author response:
A. In Figure 3, the upper panel shows the correlation between the true and predicted gene expression of the perturbation; on the other hand, the bottom panel shows the MSE between the true and predicted (only for the top 20 differentially expressed (DE) genes provided by GEARS). For correlation (upper panel), the higher the better. Conversely, for MSE (bottom panel), the lower the better. We first note that the only performance ranking discrepancy is observed in the Norman benchmark, where scDiff achieved a higher (better) correlation but higher (worse) MSE than GEARS. This discrepancy might occur due to scDiff better capturing the global trend of the perturbed gene expression pattern, while GEARS more accurately captures the exact gene expression value for the top 20 DE genes.

Reviewer comment (Reproducibility):
W1. The model description is more of a sketch (e.g. even Fig 1) than a self-contained and unambiguous description. More detail about how and what conditioning information is used for each task is lacking.

Author response:
A. Thank you for the suggestions. In short, regardless of the task, all types of available conditions in the data are used for training.  Specifically, for a given dataset, we gather all conditions, including cell types, cell batch labels, cell context, and, if available, perturbation status. Along with diffused expression, all the conditions are then fed into *scDiff* to recover the clean expression. We kindly refer the reviewer to Appendix E.1 Table 3 for the available conditions of each dataset. 

As detailed in Appendix D, we apply task-specific inference strategies. For annotation, the core idea of the method lies in evaluating the prediction error of the input at various time steps given different cell types. Regarding imputation, we generate samples from the distribution of training data conditioned on the visible values. In perturbation prediction tasks for novel cell types and novel gene perturbations, *scDiff* is trained on the existing conditions in the training set, and the task objective is to generate samples from query conditions.

Reviewer comment (Experiments):
W2. Given that this is an application to a new domain, rather than evaluating it on a suite of tasks with the same mode, it might make sense to characterize the model and the validity of those choices for scRNA-seq datasets.

Author response:
A. To characterize the model and the validity of those choices, we present the ablation studies in annotation task regarding number of layers in embedder and decoder, adding residual connection and removing cross-attention. The results are summarized in the following table, where we report the accuracy on all the annotation datasets across 5 runs. We will elaborate on our motivation for the current model structure according to the results. 

The autoencoder structure has shown great success in single-cell analysis [1, 2]. Following the existing works [3, 4], we build our model with a relatively heavy 6-layer encoder and a light 1-layer decoder. Before the encoder, we utilize an embedder to project the input expression into the embedding space. In the encoder, we wish to inject the prior information through different conditions into the expression embeddings, where we choose the cross-attention mechanism due to performance concerns. We extract the embeddings of conditions through various conditioners and feed them into the key and value of the cross-attention blocks. The embedder serves as a projection from expression space to embedding space, while the decoder reverses this process. We empirically observed that a 1-layer neural network is sufficient for such transformation. We note that the model structure has not been optimized for downstream tasks. Tricks like adding U-Net style residual connections can still improve the performance.
|  |  | Brain | HLCA | Immune | Liver | Pancreas | PBMC12K |
| --- | --- | --- | --- | --- | --- | --- |
|  | Default  | 0.9473 +/- 0.0074 | 0.8931 +/- 0.0070 | 0.8442 +/- 0.0076 | 0.8439 +/- 0.0042 | 0.9680 +/- 0.0143 | 0.9670 +/- 0.0042 |
| Decoder | 2-layer | 0.9445 +/- 0.0041 | 0.8911 +/- 0.0042 | 0.8444 +/- 0.0070 | 0.8589 +/- 0.0238 | 0.9735 +/- 0.0023 | 0.9700 +/- 0.0045 |
|  | 3-layer | 0.9378 +/- 0.0017 | 0.8994 +/- 0.0038 | 0.8383 +/- 0.0112 | 0.8062 +/- 0.0216 | 0.9731 +/- 0.0014 | 0.9701 +/- 0.0029 |
|  | 4-layer | 0.9409 +/- 0.0064 | 0.8989 +/- 0.0037 | 0.8414 +/- 0.0051 | 0.8199 +/- 0.0185 | 0.9763 +/- 0.0011 | 0.9744 +/- 0.0002 |
| Embedder | 2-layer | 0.9394 +/- 0.0029 | 0.8857 +/- 0.0089 | 0.8424 +/- 0.0061 | 0.8434 +/- 0.0239 | 0.9720 +/- 0.0011 | 0.9667 +/- 0.0039 |
|  | 3-layer | 0.9512 +/- 0.0048 | 0.8847 +/- 0.0053 | 0.8358 +/- 0.0081 | 0.8447 +/- 0.0224 | 0.9745 +/- 0.0026 | 0.9672 +/- 0.0039 |
|  | 4-layer | 0.9434 +/- 0.0048 | 0.8927 +/- 0.0018 | 0.8461 +/- 0.0014 | 0.8423 +/- 0.0100 | 0.9742 +/- 0.0012 | 0.9606 +/- 0.0036 |
| General | No cross-attention | 0.9376 +/- 0.0019 | 0.8870 +/- 0.0012 | 0.8318 +/- 0.0028 | 0.8456 +/- 0.0054 | 0.9722 +/- 0.0014 | 0.9608 +/- 0.0042 |
|  | U-Net residual | 0.9456 +/- 0.0047 | 0.8959 +/- 0.0050 | 0.8502 +/- 0.0075 | 0.8399 +/- 0.0103 | 0.9693 +/- 0.0014 | 0.9644 +/- 0.0021 |
|  | Remove diffusion | 0.5513 +/- 0.1390 | 0.1487 +/- 0.0313 | 0.1233 +/- 0.1661 | 0.1525 +/- 0.0214 | 0.3749 +/- 0.0627 | 0.5939 +/- 0.1144 |

[1] Lopez, Romain, et al. "Deep generative modeling for single-cell transcriptomics." *Nature methods* 15.12 (2018): 1053-1058.

[2] Eraslan, Gökcen, et al. "Single-cell RNA-seq denoising using a deep count autoencoder." *Nature communications* 10.1 (2019): 390.

[3] Gong, Jing, et al. "xTrimoGene: An Efficient and Scalable Representation Learner for Single-Cell RNA-Seq Data." *NeurIPS* (2023).

[4] He, Kaiming, et al. "Masked autoencoders are scalable vision learners." *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*. 2022.

Reviewer comment (Experiments):
W3. The authors tune their own model, and report performance with default values for competing methods. I appreciate the transparency in communication, but it makes it harder to interpret results. Perhaps using benchmarks where train/test splits are fixed, and other methods are reported with tuned performance would be more informative. [Openproblems](https://openproblems.bio/) curates such benchmarks on a wide variety of such tasks for many easy-to-access datasets. Performing some evaluations with those datasets and comparing the pre-defined metrics on those tasks against the leaderboard there would certainly boost my confidence in results.

Author response:
A. Thank you for the valuable suggestions. To benchmark *scDiff* with existing leaderboards, we present results follow the setting of OpenProblems in [denoising](https://openproblems.bio/results/denoising/) and [cell type annotation](https://openproblems.bio/results/label_projection/). For annotation task, we copied the accuracy of four top-performing baselines from the [leaderboard](https://openproblems.bio/results/label_projection/) with random splits. The results are as follows, where *scDiff* is evaluated across five runs. We observe that *scDiff* matches the best performance of the leaderboard.

| ACC | Pancreas | Tabula Muris | CeNGEN |
| --- | --- | --- | --- |
| Logistic regression (log CP10k) | 0.98 | 0.92 | 0.89 |
| Seurat reference mapping (SCTransform) | 0.98 | 0.90 | 0.83 |
| Multilayer perceptron (log scran) | 0.98 | 0.92 | 0.87 |
| XGBoost (log CP10k) | 0.97 | 0.86 | 0.84 |
| scDiff | 0.9800 +/- 0.0039 | 0.9257 +/- 0.0066 | 0.8889 +/- 0.0024 |

For denoising task, we selected three performant baselines from the [leaderboard](https://openproblems.bio/results/denoising/) regarding MSE. We report the mean and variance of *scDiff* and baselines across five runs. Note that the MSE results on [leaderboard](https://openproblems.bio/results/denoising/) are scaled for better illustration, we reproduced the benchmark for a fair comparison via the [public code repository](https://github.com/openproblems-bio/openproblems/tree/main/openproblems/tasks/denoising) without scaling. A conclusion can be drawn that *scDiff* outperforms the baselines across datasets.

| MSE | PBMC | Pancreas | Tabula Muris |
| --- | --- | --- | --- |
| MAGIC | 0.1887 +/- 0.0001 | 0.2319 +/- 0.0001 | 0.1842 +/- 0.0001 |
| DCA | 0.2184 +/- 0.0003 | 0.2683 +/- 0.0003 | 0.2160 +/- 0.0007 |
| ALRA | 0.2823 +/- 0.0007 | 0.3317 +/- 0.0001 | 0.2800 +/- 0.0006 |
| scDiff | 0.1810 +/- 0.0011 | 0.2079 +/- 0.0027 | 0.1680 +/- 0.0023 |

Reviewer comment (Experiments):
Q1. Some of the datasets have a very small number of cells (~3000 for Jurkat). Can the authors clarify in which cases was the model trained from scratch (i.e. with randomly initialized weights)?

Author response:
A. We apologize for the obscurity in the manuscript. To clarify, *scDiff* is trained from scratch in each task on each dataset except few-shot cell type annotation. We further refer to the questions posted by reviewer Cfbr, where we concluded that our proposed *scDiff* is well optimized even for small datasets like the three datasets in imputation task. One promising future work is to verify the transferability of *scDiff* pre-trained on large-scale and heterogenous datasets.

Reviewer comment (Theory):
Q2. Eq. 1 uniform prior assumption seems to be a strong over-simplification. Would this hurt on unbalanced datasets (e.g. when considering cell type classification as in appendix C)?

Author response:
A. We thank the reviewer for the question. We note that the cell type distributions of the current annotation datasets are already quite skewed and unbalanced. In the following table, we calculate the proportion of the number of cells of the largest three and smallest three cell types. The superior cell type annotation performance of *scDiff* highlights the feasibility of the seemingly simplistic uniform prior. Nevertheless, we agree that more careful handling to take into account of the non-uniformity could be an interesting avenue to explore in the future.

| Proportion | PBMC12K | Pancreas | HLCA | Immune | Brain | Liver |
| --- | --- | --- | --- | --- | --- | --- |
| Majority 1 | 0.4167 | 0.3353 | 0.1445 | 0.1125 | 0.4081 | 0.4143 |
| Majority 2 | 0.1857 | 0.2545 | 0.1348 | 0.0899 | 0.1047 | 0.1320 |
| Majority 3 | 0.1352 | 0.1308 | 0.1061 | 0.0801 | 0.0810 | 0.0973 |
| Minority 3 | 0.0293 | 0.0020 | 0.0005 | 0.0015 | 0.0058 | 0.0143 |
| Minority 2 | 0.0283 | 0.0015 | 0.0004 | 0.0011 | 0.0035 | 0.0063 |
| Minority 1 | 0.0073 | 0.0004 | 0.0003 | 0.0010 | 0.0023 | 0.0024 |

Reviewer comment (Writing):
Q3. What is $L$ here? Is it simply the number of distinct conditions?  
> The goal of each conditioner is to extract a set of L numerical representations of an input condition c.

Author response:
A. We thank the reviewer for pointing out the vague definition of the number of conditions $L$. We confirm that $L$ is indeed the number of conditions, and we have updated the manuscript to clarify this further.

Reviewer comment (Writing):
Q4. In the same section about the conditioner, can you provide examples for:  
> The mapping here can be designed to suit the specific needs of different input types.

Author response:
A. Thank you for the suggestion regarding the examples of condition mappings. Here we provide some examples. The matrix-valued observed counts in imputation will go through a linear transformation; The class labels of cells are represented by learnable embeddings; The cell ontology textual descriptions will be processed by LLMs; The gene ontology graph will be fed into graph neural networks. We have added a few examples following the aforementioned sentence in the revised manuscript.

Reviewer comment (Experiments):
Q5. If there are only ~30 cell types in the dataset (e.g. Liver in Fig. 5), how is any high (e.g. 64 or 128 dim) dimensional embedding from LLM's helpful for the model for one-shot classification?

Author response:
A. We thank the reviewer for the question regarding why LLM helps in the few-shot learning setting. In short, LLM helps by providing prior knowledge about the relationships between different cell types, potentially bridging the training and few-shot testing cell types. BioLinkBERT (the LLM we used in our paper) is trained on a vast biomedical corpus, which contains many medical terms (such as cell types) and descriptions of how different terms relate. The trained BioLinkBERT model can thus provide inductive bias that reflects how similar two cell types are given their biomedical definitions (extracted from the cell ontology). We argue that this inductive bias is the key to its success in enhancing the few-shot annotation setting. For example, during training, the model may have access to training examples of B-cells. Suppose during the few-shot training, the two target cell types are plasma cells (a specific type of B-cell) and neuron cells (utterly distinct from the B-cell lineage). Then, BioLinkBERT can generate the cell type embedding of plasma cells that are more similar to the B-cell but less similar to the Neuron cell. This distinction is helpful for the model to classify the two different cell types correctly.

Reviewer comment (Novelty):
The application here is straightforward, and no core ML innovation was necessary to execute this project. While realizing how diffusion models map to this important modeling domain and executing the basics well does justify the paper, it should be noted that it is not innovating dramatically.

Author response:
A. We appreciate the reviewer's recognition of our efforts in adapting diffusion models to single-cell analysis. While we concur that our work may not lie in innovating core ML, we emphasize its novelty in applying these models to a new domain. Our approach not only introduces a fresh perspective to single-cell analysis but also demonstrates the untapped potential of diffusion models in this field.
Our work serves as a bridge, connecting advanced diffusion model research with pressing challenges in single-cell analysis. We believe this interdisciplinary fusion is an innovation in its own right, offering new methodologies and insights that enrich both fields. The practical implications and potential for future research stemming from our unified framework underscore its significance as a meaningful contribution to the scientific community.

Reviewer comment (Writing):
The authors repeatedly use the term "casting single cell tasks as posterior inference" throughout the paper, but in truth do not perform much inference beyond the classical training scenario of diffusion models. Language such as "we cast these tasks as probabilistic modeling using a shared model" would be more accurate, since inference is not the main focus and is relatively generically solved without inspecting posterior distributions per layer over specific conditions. This is a minor point, but important to consider for careful language use.

Author response:
A. We thank the reviewer for the insightful feedback regarding our terminology. In this work, we refer to our framework as *posterior estimation* to highlight that we are modeling the *posterior* of gene expression. This emphasizes the discrepancy between the proposed framework and the deterministic formulation of the downstream tasks. We agree that the inference stage is not the core of the story and the exact likelihood computation is not necessary regarding the downstream tasks. To be consistent, we have updated the manuscript by replacing *posterior estimation* with *posterior modeling* and *probabilistic modeling*.

Reviewer comment (Theory):
The imputation task is not entirely correct in how it is phrased as a conditioning setting. The predictive distribution likely contains outputs beyond the visible ones. Although the conditioning mechanism injected per layer increases affinity towards correctly imputing, the approach is modeling not exactly p(x|m*x) but rather p(x*|m*x), where x* is sampled from some distribution conditioned on embeddings of x*m. This is a fine approximation and not a reason to reject, but it would be preferable if the authors acknowledged this and left space for future improvements.

Author response:
A. Thank you for carefully inspecting our formulation of the downstream tasks. We agree that there is a gap between the generated samples and the missing values. We incorporate cell context as a condition during the training stage to model the posterior of the training data. In the generation stage, we approximate the missing values by sampling from the distribution of training data conditioned on the visible values. We add a footnote in the updated manuscript for clarification.

Reviewer comment (Experiments):
The scaling behavior of the studied systems with respect to different dataset sizes is not explored.

Author response:
A. Thank you for the valuable question about scaling of dataset size. We conducted a preliminary experiment to examine whether *scDiff* follows the scaling law regarding the number of cells in training data. Starting from the three large datasets we used in annotation task, we subsampled the training sets by randomly keeping a proportion of training cells and evaluate *scDiff* on the same testing sets. We summarize the annotation accuracy across five runs in the following table, where the subsample ratios range from 0.01 to 1.0. 

| Subsample ratio | HLCA | Immune | Brain |
| --- | --- | --- | --- |
| 0.01 | 0.2839 +/- 0.0135 | 0.3275 +/- 0.0238 | 0.7797 +/- 0.0035 |
| 0.05 | 0.5924 +/- 0.0077 | 0.5679 +/- 0.0239 | 0.7673 +/- 0.0048 |
| 0.1 | 0.8253 +/- 0.0189 | 0.6995 +/- 0.0128 | 0.8951 +/- 0.0403 |
| 0.2 | 0.8916 +/- 0.0022 | 0.8020 +/- 0.0096 | 0.9510 +/- 0.0056 |
| 0.3 | 0.8951 +/- 0.0064 | 0.8362 +/- 0.0036 | 0.9659 +/- 0.0044 |
| 0.4 | 0.8807 +/- 0.0079 | 0.8598 +/- 0.0013 | 0.9599 +/- 0.0019 |
| 0.5 | 0.8799 +/- 0.0059 | 0.8625 +/- 0.0069 | 0.9279 +/- 0.0058 |
| Full | 0.8931 +/- 0.0070 | 0.8442 +/- 0.0076 | 0.9473 +/- 0.0074 |

For all three datasets, we observe that the accuracy has already matched or even outperformed that of full dataset before the subsample ratio reaches 0.5. A possible explanation is that the scRNA-seq data is not as diverse as images in vision. This suggests that *scDiff* trained from scratch does not require large amount of training dataset for the annotation task. A possible next step is to check how cross-dataset pre-training will affect the results.

Reviewer comment (Experiments):
Diffusion models tend to be data and compute hungry, so it remains unclear how they behave in this context. Questions include whether they can be applied on one screen and whether they require broad pre-retraining.

Author response:
A. As depicted in the above table, *scDiff* achieves comparable or even higher accuracy in annotation before the subsample ratio reaches 0.5. In addition, it is worth noting that our *scDiff* is trained from scratch in each task on each dataset. It outperforms the baselines even in the small datasets (e.g., three imputation datasets have at most ~3000 cells). We can have a preliminary conclusion that diffusion model is not as data hungry as they are in vision domain.

We note that *scDiff* has a computational complexity that is nearly the same as a standard MLP. In the following table, we report the per-epoch computation time for scDiff against standard baselines for all our benchmarking tasks.  This supports that *scDiff* can be scaled to large datasets for possible pre-training. The experiment in Q1 indicates that increasing the dataset size has limited improvement for annotation task. However, it is still unclear whether cross-dataset pre-training will help to boost the performance in annotation and other downstream tasks. We leave it as a future work to explore the possibilities of large-scale pre-training for *scDiff.*

|  | Dataset | scDiff per-epoch runtime (sec) | Baseline per-epoch runtime (sec) | Baseline name |
| --- | --- | --- | --- | --- |
| Annotation | Brain | 9.63 | 9.82 | scANVI |
|  | HLCA | 13.60 | 12.73 | scANVI |
|  | Immune | 5.28 | 8.86 | scANVI |
|  | Liver | 2.83 | 2.97 | scANVI |
|  | Pancreas | 5.05 | 14.37 | scANVI |
|  | PBMC12K | 2.04 | 3.31 | scANVI |
| Denoising | 293T | 0.8847 | 0.53 | DCA |
|  | Jurkat | 0.5736 | 0.60 | DCA |
|  | PBMC1K | 0.2634 | 0.16 | DCA |
| Perturbation | HPoly | 1.96 | 1.83 | scGEN |
|  | PBMC | 3.16 | 3.53 | scGEN |
|  | Salmonella | 1.58 | 1.81 | scGEN |
| Gene Perturbation | Adamson | 83.88 | 1189.66 | GEARS |
|  | Dixit | 44.67 | 760.70 | GEARS |
|  | Norman | 155.88 | 491.97 | GEARS |

Reviewer comment (Experiments):
While the modeling framework is convincing, the field would benefit from investigating questions of data efficiency to increase the value of the paper’s contributions.

Author response:
A. As depicted in the above table, *scDiff* achieves comparable or even higher accuracy in annotation before the subsample ratio reaches 0.5. In addition, it is worth noting that our *scDiff* is trained from scratch in each task on each dataset. It outperforms the baselines even in the small datasets (e.g., three imputation datasets have at most ~3000 cells). We can have a preliminary conclusion that diffusion model is not as data hungry as they are in vision domain.

We note that *scDiff* has a computational complexity that is nearly the same as a standard MLP. In the following table, we report the per-epoch computation time for scDiff against standard baselines for all our benchmarking tasks.  This supports that *scDiff* can be scaled to large datasets for possible pre-training. The experiment in Q1 indicates that increasing the dataset size has limited improvement for annotation task. However, it is still unclear whether cross-dataset pre-training will help to boost the performance in annotation and other downstream tasks. We leave it as a future work to explore the possibilities of large-scale pre-training for *scDiff.*

|  | Dataset | scDiff per-epoch runtime (sec) | Baseline per-epoch runtime (sec) | Baseline name |
| --- | --- | --- | --- | --- |
| Annotation | Brain | 9.63 | 9.82 | scANVI |
|  | HLCA | 13.60 | 12.73 | scANVI |
|  | Immune | 5.28 | 8.86 | scANVI |
|  | Liver | 2.83 | 2.97 | scANVI |
|  | Pancreas | 5.05 | 14.37 | scANVI |
|  | PBMC12K | 2.04 | 3.31 | scANVI |
| Denoising | 293T | 0.8847 | 0.53 | DCA |
|  | Jurkat | 0.5736 | 0.60 | DCA |
|  | PBMC1K | 0.2634 | 0.16 | DCA |
| Perturbation | HPoly | 1.96 | 1.83 | scGEN |
|  | PBMC | 3.16 | 3.53 | scGEN |
|  | Salmonella | 1.58 | 1.81 | scGEN |
| Gene Perturbation | Adamson | 83.88 | 1189.66 | GEARS |
|  | Dixit | 44.67 | 760.70 | GEARS |
|  | Norman | 155.88 | 491.97 | GEARS |