Single-cell-level perturbation-induced and condition-related signal estimation with batch effect removal using NDreamer

Abstract

Advances in sequencing technologies and the growing volume of single-cell data have created unprecedented opportunities for uncovering gene expression patterns causally induced by experimental perturbations or statistically associated, but not necessarily causal, with disease conditions. However, current analytical methods inadequately account for batch effects and data sparsity or fail to capture the inherent non-linearity in single-cell data, leading to biased estimation. To address these limitations, we developed NDreamer that combines neural discrete representation learning and matching to remove batch effects and estimate perturbation-induced or condition-associated signals at single-cell resolution. NDreamer outperformed existing methods by using mutual information loss on discrete latent variables to disentangle cells’ intrinsic features from conditions or batch effects, while preserving both global and local variance within batches and conditions via triplet and local neighborhood loss. We applied NDreamer to multiple datasets across platforms, organs, and species and validated and benchmarked its performance in removing batch effects and estimating perturbation-induced or condition-associated signals. In particular, we applied NDreamer to an Alzheimer’s disease cohort, revealing biologically relevant gene expression patterns that distinguish dementia patients from controls.

Graphical Abstract

Graphical Abstract.

Introduction

Single-cell RNA sequencing (scRNA-seq) has become a cornerstone technology for studying cellular heterogeneity and detecting gene expression patterns at cell-type level. Research has been conducted to identify gene expression changes induced by perturbations, including genetic modifications [1], chemical treatments [2], and environmental stimuli [3]. Advanced technologies such as perturb-seq [1, 4] have further enabled us to study cellular responses to modifications in disease-related risk genes across diverse cell types.

In single-cell studies, accurately quantifying the influence of perturbations on individual cells and studying how these responses vary across cellular populations remain challenging, since resolving this heterogeneity is essential for uncovering causal mechanisms of cellular regulation and for enabling context-specific therapeutic interventions. It is well known that cells with distinct intrinsic features, such as cell type, subtype, cell cycle stage, microenvironment, chromatin accessibility, and other pre-existing states, respond differently to perturbations. Here, we define the perturbation-induced effect in each cell as the individual treatment effect (ITE) and its intrinsic features as effect modifiers in the causal inference framework [5]. Since we only observe the perturbation label and gene expression of each cell, it is difficult to dissect intrinsic variability from perturbation-induced effects without direct information on intrinsic features. Estimating ITE is further complicated by the fact that we only observe a cell’s transcriptome under one perturbation condition and do not observe its counterfactual state. Counterfactual causal inference approaches are helpful for estimating ITE in perturbation studies.

Early approaches rely on aggregating single-cell data to pseudo-bulk data and compare the expression of cells in the perturbed condition with cells in the control condition. They average the expression of cells of the same type in each sample, failing to capture fine-grained heterogeneity and can produce different results depending on the chosen aggregation resolutions. For example, CD4+ and CD8+ T cells respond differently to CD28 costimulation, where CD4+ T cells sustain proliferation, whereas CD8+ T cells show a limited response [6]. Thus, aggregating all T cells could mask these subtype-specific responses, leading to insignificant results in the analysis of T cell responses. Later methods estimate single-cell-level ITE using single-cell digital twins and generative models [7, 8] or counterfactual matching based on perturbation-independent factors [3, 9]. While these models have advanced the single-cell ITE estimation by identifying and disentangling perturbation-independent factors and isolating perturbation-related ITE responses from the gene expression matrix, they have several limitations. First, the unknown or subtle bias [10, 11] in the expression profile generated by the generative model [7, 8] can impact the ITE estimation and downstream analysis. Second, the inherent non-linearity in single-cell data makes linear decomposition [12] not suitable for analysis. Third, independence measured by correlation coefficients does not guarantee independence, and the discriminator in adversarial training to establish independence may not rigorously lead to independence under various hyperparameter settings. Lastly, it is unknown if the distance in the perturbation-independent factor space is distorted, thus, the causal weighting or matching in the downstream analysis may not be valid.

In single-cell studies with multiple batches, the estimation of ITE can also be biased due to batch effects which previous ITE estimation methods did not account for. A large number of methods have been developed to remove batch effects [12, 13]. Although many existing methods effectively learn cell embeddings that are largely insensitive to batch effects, most of them do not output denoised gene expression profiles. Here, denoised expression refers to reconstructed or corrected expression values in which technical noise and batch-associated variation are reduced while condition-related biological signals are preserved. This distinction is important because batch-invariant embeddings can support downstream integration, but they do not provide directly interpretable gene-level expression estimates. Moreover, some methods may overcorrect the data, removing condition-related signals together with unwanted batch effects [12].

Moreover, large-scale scRNA-seq studies collected on human samples comparing disease patients and healthy controls hold great promise for identifying disease (condition)-associated gene programs and biological pathways across cell types, shedding light on disease mechanisms and potential therapeutic targets. However, estimating single-cell-level condition-related signals (CRS) still remains challenging. Here, CRS can be viewed as ITE when comparing between disease and control conditions, but can only be interpreted as correlations without causal relationships.

To address these challenges, we propose NDreamer (Neural Discrete learning for decomposing condition-Related or perturbation-induced signals, Effect modifiers, And Measurement ERror in scRNA-seq), to remove batch effects and estimate perturbation-induced or condition-associated signals at single-cell resolution. NDreamer can be applied to single-cell experiments, including multiple perturbations or condition labels with or without batches, and estimate the ITE/CRS for each cell and the conditional average treatment effect (CATE) or conditional average condition-related signals (CACRS) for specific cell populations (Fig. 1a). We benchmarked NDreamer against existing single-cell-level condition-related signal estimation methods as well as batch effect-denoising methods and showed that NDreamer outperformed all current methods. Applied to datasets across sequencing platforms, organs, diseases, and species, NDreamer effectively identified heterogeneous responses of cells to genetic editing, chemical treatments, and disease-associated gene expression patterns across conditions. We also applied NDreamer to an Alzheimer’s disease snRNA-seq dataset and uncovered disease-related expression patterns across cell types in the middle temporal gyrus (MTG).

Figure 1.

The general form of gene expression in single-cell sequencing studies and the pipeline of NDreamer. (a) Data structure in the single-cell studies. Given a cell-by-gene expression matrix and a cell-by-metadata annotation matrix. The aim is to estimate the perturbation-induced or condition-related expression patterns. The perturbation or condition can be a chemical perturbation, a genetic perturbation, or a disease status. (b) The pipeline of NDreamer. Given a scRNA-seq dataset with multiple batches, conditions/perturbations, and cell types, NDreamer first generates an effect-modifier embedding in which cells are mixed across conditions and batches while intrinsic cellular variation is preserved, then uses matching based on the effect modifier embedding to generate the denoised expression, where batch effects are removed, and cells’ intrinsic features are preserved. It then estimates the ITE/CRS based on the denoised expression using matching. Dot colors represent different batches, and dot shapes represent different cell types.

Methods

Unified modeling of gene expression in single-cell studies

We assume that the observed gene expression of a single cell can be viewed as the combination of its true expression, batch effects (if present), and independent and identically distributed (i.i.d.) random noise. The true expression of a cell depends on the perturbation (e.g. chemical perturbations, environmental stimuli, genetic modifications) or condition (e.g. disease or control) label and its intrinsic characteristics, referred to as effect modifiers, such as cell type and cell state. Batch effects, arising from technical bias in measurement, are determined by the interaction between batch-specific attributes, the cell’s effect modifier features, and its experimental condition. Formally, the expression level of gene j in cell i, denoted as , can be modeled as

where is the true expression of gene j in cell i, represents the intrinsic features of cell i, is its condition label, is the batch effect on gene j in cell i, represents the batch to which cell i belongs, and is the random noise where . Here, and can be non-linear functions. Since batch effects are related to measurement error, we assume . This means among cells with the same intrinsic features and condition label , batch effects may shift the observed expression in different directions across batches, but these technical deviations are centered around zero. Therefore, averaging the observed expression of gene jacross cells with the same x and cfrom multiple batches provides a point estimate of the true expression component , up to random noise. This is reasonable because represents the systematic expression determined by cell-intrinsic characteristics and condition, whereas batch effects arise from technical factors such as sample processing, sequencing depth, or reagent variability. These technical factors may introduce batch-specific deviations, but after conditioning on the same biological state and condition, they are not expected to contribute a consistent biological signal in one direction across batches.

To dissect the condition-related responses and the condition-independent effect modifiers from the observed gene expression, we applied a neural network to extract the effect modifier that is independent of batch and condition while preserving the biological variation within each condition and batch. Then, for each cell in the non-control condition, we estimated its denoised expression by finding its nearest neighbor in batches under the same condition in the effect modifier space and averaging their corresponding gene expression. Since and , the average can be considered as an estimate of . We then estimate the counterfactual expression of cells in the control condition using the same operation, but for batches under the control condition. by: subtracting the counterfactual expression from the estimated denoised expression, we obtain the ITE/CRS for cell i. For each cell in the control condition, its denoised expression is calculated in the same way as calculating the counterfactual expression of cells in the non-control condition. by: aggregating the ITE/CRS of certain cell populations, we can estimate the CATE and apply the corresponding nonparametric test to test its statistical significance (Fig. 1b). For convenience, we only discuss the estimation of ITE since the estimation methods of ITE and CRS are the same under the proposed framework. The notations used in the methods description is tabulated in Supplementary Notes S1.

Overview of the NDreamer deep learning model

Effectively extracting batch- and condition-independent effect modifiers [3] while preserving within-batch and within-condition biological variation, without introducing distortions, remains challenging. Moreover, to make sure that the transformation of the condition- and batch-independent effect modifier embedding is not an arbitrary transformation, we assume that the effect modifier space should contain the information of within-batch and within-condition between-cell variance, where the within-batch and within-condition local and global distance structures should be retained in the effect modifier embeddings. To address this, we employed neural discrete representation learning [14, 15], which generates categorical latent variables from the input expression profiles. Following discretization, a subsequent neural network transformation was applied to produce the final effect modifier embedding as a continuous vector. To constrain the embedding space and reduce distance distortions, we applied the variational inference techniques [16] to the final effect modifier embedding and used a neural network decoder to reconstruct the expression (Fig. 2a).

Figure 2.

Overview of the NDreamer deep learning model for effect modifier extraction. (a) The structure of the deep learning model. Input expression is first transformed to a set of categorical latent variables according to the neural discrete representation learning, then it is transformed to a continuous embedding, which is treated as the final effect modifier embedding, and then re-parametrized similar to VAE. Then a decoder is used to reconstruct the expression. For the categorical latent variables, we apply the independent and dependent loss using normalized mutual information. For the effect modifier embedding, we applied the triplet loss and the local neighborhood loss to preserve biological conservation. (b) Graphical illustration of the triplet loss function. The distance between cells within the same high-resolution unsupervised cluster should remain consistently smaller than the distance between cells from different low-resolution unsupervised clusters. Points with different colors now represent cells from different high-resolution clusters, while low-resolution clusters are indicated by dashed circles. (c) Graphical illustration of the local neighborhood loss. Local neighborhood structures calculated through the Gaussian kernel in the effect modifier space should mirror those in the raw expression space.

To ensure that the effect modifier embedding is independent of both batch and condition so that it can be used for matching and estimating ITE or CRS [3], we first minimize the normalized mutual information between the categorical latent variable and the batch and condition (Fig. 2a). By decoupling the categorical latent variable from the batch and condition, the transformation of the categorical latent variable, the final effect modifier embedding, also becomes independent of these variables. Compared to the correlation-coefficient-based independence measurements as well as discriminator-based independence constraints [17], the mutual information-based approach can more robustly capture complex independence patterns (Supplementary Notes S3).

After establishing independence through mutual information loss, we applied three loss functions to preserve within-batch and within-condition biological variation while preventing distortions. Firstly, the selection of the categorical latent variable was designed to depend on the unsupervised clustering patterns observed in the raw expression space within each condition and batch. Secondly, we employed a triplet loss to ensure that the distance between cells within the same high-resolution unsupervised cluster remained consistently smaller than the distance between cells from different low-resolution unsupervised clusters [18]. Together, these two loss functions safeguard the global biological structure. Beyond preserving global structure, we implemented a local neighborhood loss to maintain local structure, ensuring that the local neighborhood structures in the effect modifier space mirrored those in the raw expression space (Fig. 2c).

Estimation of the effect modifier embedding

Encoder with neural discrete representation learning and variational inference

For cell i with expression , batch , and condition , we utilize a set of codebooks, each containing embeddings of dimensionality , to discretize the input into a latent discrete representation. The discretization operation introduced in VQ-VAE [14, 15] has been used in computer vision to improve the stability of image generation. In our study, it enables more reliable control of independence by allowing direct measurement via mutual information, which is both necessary and sufficient to indicate independence. It is also more time efficient than adversarial training since we have fewer backward propagations, and supports an extremely large and expressive space of possible combinations—for example, with 1024 embeddings per codebook and 32 codebooks, there are 1024³² possible combinations, which is expressive enough (Supplementary Notes S3).

Specifically, first, for each codebook, is transformed by a codebook-specific neural network into an intermediate latent space . Each embedding in the codebook is then probabilistically selected based on the softmax of the negative Euclidean distance between and the embeddings. The probability of selecting the j-th embedding in a given codebook is computed as:

To make the random embedding selection operation differentiable, we employ the Gumbel-Softmax trick [19], enabling backpropagation through the probabilistic codebook selection.

Putting together the results from each codebook, we generate a codebook choice matrix P of shape , recording which embedding is mapped for each codebook. After the embeddings are selected, we concatenate the chosen embeddings from each codebook to form a new latent space , which serves as the discrete representation of .

To ensure robust alignment between and the selected code and make the update of embeddings differentiable, a commitment loss with two components is employed similar to the VQ-VAE [14]. The first term penalizes the difference between and , while the second term enforces the embeddings to remain close to , preventing instability in the codebook updates. The full commitment loss is expressed as:

where is the stop-gradient operator, and β (default to 0.25) controls the balance between the two terms.

To improve numerical continuity and reduce distortion in the estimated effect modifier embedding, we applied the variational inference technique similar to the standard Variational Autoencoder (VAE) framework. Based on the new latent space , we apply a multi-layer perceptron (MLP) to further transform into two outputs: a mean vector and a standard deviation vector , where refers to the number of dimensions of the effect modifier embedding. The standard deviation vector is then exponentially transformed to ensure positivity and the mean vector is treated as the effect modifier embedding . Since is independent of the condition and batch as constrained by the mutual information loss mentioned later, , as a transformation of , is also independent of the condition and batch. At last, we resample the latent effect modifier representation by:

where reparameterization is used to enable backpropagation through the sampling operation, expressed as

Following the Evidence Lower Bound (ELBO) loss in the training of traditional VAE [16], the Kullback–Leibler (KL) divergence loss is introduced. The KL-divergence between the approximate posterior and the prior is given by

where j indexes the dimensions of the latent effect modifier embedding space. The resulting resampled effect modifier embedding serves as the final encoder output derived from the input , integrating the benefits of both probabilistic embedding discretization and variational inference.

In addition, we found that some embeddings in the codebook are almost never used in model training, which reduces the model’s efficiency. To solve this problem, we reset the embeddings that are never selected every 30 steps when training, where we randomly replace the unused embeddings with the intermediate latent space embedding produced by the randomly chosen real data.

Decoder with conditional generation

We then reconstruct the gene expression to let the model learn from the data using the resampled effect modifier embedding along with the batch embedding and the condition embedding. Written in formula:

where and are the batch and condition embedding, and refers to the decoder neural network.

As part of the VAE ELBO, the reconstruction loss is defined as the mean squared error (MSE) between the measured normalized expression (Details of data normalization are provided in Supplementary Notes S5) and the reconstructed expression :

where j indexes the dimensions of the expression vector, and is the total number of highly variable genes input to the model.

To further emphasize the interaction between the effect modifier , batch embedding , and condition embedding , we reinforce the reconstruction loss with an additional cosine distance loss

where ·p denotes the dot product, and is the vector norm.

Normalized mutual information loss

To encourage independence between the choice of the embeddings in the codebooks and the batch assignments, as well as the condition assignments, while preserving meaningful biological cluster information, we introduce the normalized mutual information loss. The choice of the embeddings in the codebooks is inherently categorical, allowing us to compute entropy and mutual information directly from the probability distributions. We start by introducing the calculation of mutual information between the choice of the embeddings in the codebooks and the batch.

Consider B cells in a mini-batch (note that the mini-batch in deep learning training is different from the batch in the scRNA-seq study), let the batch assignment vector represent the experimental batch assignments. The assignment vector is first one-hot encoded into a matrix , where is the number of unique batch categories. For one codebook, the joint probability , indicating the probability of a cell being simultaneously assigned to one of the batches and one of the embeddings in this codebook, can be calculated by

where is the probability matrix of codebook embedding selection for this codebook. This matrix-form operation enables acceleration by parallel computing.

The marginal probabilities of the batch and embedding assignment P(b) and P(e) are obtained by summing over the respective dimensions of :

The marginal entropy of the batch and embedding assignment can then be calculated by:

The mutual information is then calculated by:

To stabilize the training and avoid NaN values, we use the arithmetic normalized mutual information (ANMI) as the loss function.

We then compute the ANMI of other codebook assignments with the batch and average them as the mutual information between the codebook assignments and the batch . By minimizing the mutual information, the loss encourages the latent variables to remain independent of the batch assignments.

Similarly, we calculate and minimize the ANMI between the codebook assignments and the condition assignment . In addition, within each unique batch-condition combination, we maximize the ANMI between the codebook assignments and the unsupervised cluster label with low and high cluster resolution and to retain the global structure as in the original gene expression space.

Triplet loss

The triplet loss is designed to preserve the relative distances in the effect modifier embedding, ensuring that the structure of the embedding space is not distorted [18]. The key idea is that cells belonging to the same small high-resolution cluster in the original space should be embedded closer to the anchor cell than cells from different low-resolution clusters. For a given cell i, the triplet consists of the anchor cell i, a positive cell with effect modifier embedding from the same high-resolution cluster, and a negative cell with effect modifier embedding from a different low-resolution cluster. The triplet loss for a single triplet is defined as:

where is the embedding of the anchor cell, and are the embeddings of the positive and negative cells, respectively, denotes the squared Euclidean distance, and α >0 is a margin that enforces separation between the anchor-negative and anchor-positive distances.

For each cell i, such triplets are sampled, and the triplet losses are summed to compute the total triplet loss for the cell. This approach ensures that the effect modifier embeddings reflect meaningful relationships between cells while maintaining their relative distances in the embedding space.

Local neighborhood loss

Besides using triplet loss and ANMI between codebook embedding selection and cluster information to preserve global biological conservation, we also designed the local neighborhood loss for retaining the local neighborhood structure in the effect modifier space within each unique batch-condition combination. For a reference cell i, we first identify its nearest neighbors in the original expression space and the effect modifier space, resulting in a cell subset m (). We then use the Gaussian kernel to calculate the relative distance between the reference cell and the m cells in both spaces:

where refers to the gene expression vector or the effect modifier embedding in cell i. This calculation results in the relative distance in the gene expression space and effect modifier space and . The local neighborhood loss is calculated as the mean squared error between and to encourage consistency between the neighborhood structure in the gene expression space and the effect modifier space. The hyperparameter settings for the above-mentioned loss functions are provided in Supplementary Notes S2.

ITE/CRS estimation

Since the computation of ITE and CRS, as well as CACRS and CATE are the same, we will only discuss the calculation of ITE and CATE below. To estimate the ITE for each cell, we mitigate confounding effects arising from variations in intrinsic cell features by searching for the nearest neighbors in the effect modifier space in other batches under the same condition and the reference control condition. The nearest neighbors are searched based on the Euclidean distance. These nearest cells are then averaged across batches within the same condition to obtain the cell’s denoised expression.

Let and represent the sets of batches that have cells in the treatment condition and the control condition, respectively. For each batch , the nearest cell to cell i is defined as:

where is the Euclidean distance between and . The denoised expression under the treatment condition, , is computed as:

Similarly, the denoised expression under the control condition is:

The individual treatment effect for cell i is then given by:

This framework effectively adjusts for variations in intrinsic cell features, providing a robust estimation of the ITE.

CATE/CACRS estimation and statistical test

To calculate the CATE for a subgroup S (e.g. cells of the same cell type), the average ITE across all cells in S for gene j is computed as:

where denotes the number of cells in subgroup S. To evaluate whether the distribution of ITEs for gene j, is systematically shifted toward positive or negative values, we use a sign-based test. We begin by formulating the null hypothesis that the observed individual treatment effect () values are equally likely to be positive or negative since the true value is 0 and there is continuous random noise. To test this hypothesis, each is transformed into a binary variable where

Under the null hypothesis, these binary variables follow a Bernoulli distribution with P=0.5. Next, the binary outcomes are aggregated within a subgroup S, where we calculate by adding the count of positive outcomes () with 0.5 times the count of . For large , the central limit theorem enables the approximation of by a normal distribution with mean and variance . The z-score is then computed to quantify the deviation of from its expected value, defined as

Finally, P-value is calculated using the standard normal distribution’s cumulative distribution function,

to assess the statistical significance of the observed deviation under the null hypothesis.

Results

NDreamer demonstrates robustly improved performance across experimental perturbation datasets

To systematically evaluate the performance of NDreamer with existing single-cell-level methods for perturbation-effect analysis, we conducted comprehensive benchmarking across diverse datasets across different sequencing platforms, species, and organs. Our comparison involves a set of currently state-of-the-art models, including CINEMA-OT [3], Mixscape [9], scCAPE [17] (code was downloaded from pip on 17^th November 2024), and scGen [7] (implemented by pertpy [2], code was downloaded on 13^th December 2024). We applied the default hyperparameter settings provided by their GitHub tutorial page for these methods across all datasets. We treated the perturbation- and batch-free embeddings for previous methods as their estimated effect modifier embeddings. Since the ground truth of ITE is absent in real datasets, our comparison is based on two categories of metrics calculated from the effect modifier embedding, which are also used in the evaluation in previous studies [3, 9]:

Condition label-mixing performance in the effect modifier space. Given that the intrinsic characteristics of cells prior to any perturbation are independent of the perturbation label, the extracted effect modifier embeddings for cells in the two conditions (perturbation versus control), if plotted together, should be well-mixed. Of note that here we refer the perturbation and control as conditions. Borrowing the ideas from the systematic batch-effect removal benchmarking papers [13, 20], this can be evaluated by batch-mixing metrics, where different perturbations are considered as different batches here. Specifically, we use the widely used metrics of average silhouette width (ASW_batch) [20], k-nearest-neighbor batch-effect test (kBET) [21], batch local inverse Simpson’s index (bLISI) [22], and True Positive Rate [23] for batch-mixing evaluation.

Conservation of biological information for the effect modifier space. The extracted effect modifier embeddings should preserve the intrinsic features and within-perturbation variations of cells, which can be characterized by cell type annotation. Therefore, we employed the classical label conservation metrics to evaluate local neighborhoods (1 minus cell-type LISI, cLISI) [22], global cluster alignment (Adjusted Rand Index (ARI), geometric normalized mutual information (NMI)), and relative distances (cell-type ASW) [13].

In addition to these metrics, we computed an overall score by linearly scaling the values of each metric to a range of 0 to 1 and then averaging the scaled values across all metrics. We used five publicly available datasets for the benchmarking: (1) The peripheral blood mononuclear cell (PBMC) dataset, including 16, 893 cells from IFN-beta-stimulated and control conditions [24]. (2) The rhinovirus infection (virus) dataset where we compared the 12, 255 cells between the mock condition and the exposure to 2% cigarette-smoke extract perturbation [3]. (3) A small PBMC stimulation (PBMC1) dataset, which contains 5027 cells from the stimulated and control conditions [3]. (4) The ECCITE dataset, containing 20 729 cells of three annotated cell cycle phases (G1, G2M, S) obtained from three technical replicates measured using ECCITE-seq with 25 target genes generated from the stimulated THP-1 cell line [9]. (5) The autism spectrum disorder (ASD) dataset containing brain cells with 35 neurodevelopmental delay (ND) risk gene mutations from 18 batches measured by Perturb-Seq [1].

We begin with the PBMC dataset. As illustrated in Fig. 3a, NDreamer achieves the highest overall score, far ahead of other methods, and provides the most satisfactory uniform manifold approximation and projection [25] (UMAP) visualization with clear cell type separation and condition-mixing (Fig. S1). It demonstrates the clearest cell type separation and the most effective mixing of cells from two different conditions. Notably, NDreamer is the only method that successfully captures the variation between dendritic cells and other cell types. Among the other methods, CINEMA-OT achieves the second-best performance in mixing cells across conditions; however, it incorrectly splits a single cluster of CD4 T cells into two subgroups and fails to clearly distinguish between CD14+ Mono, FCGR3A+ Mono, and dendritic cells. The remaining three methods, while achieving relatively good cell type separation, fail to integrate cells from the two conditions effectively.

Figure 3.

Benchmarking of NDreamer and other single-cell experimental perturbation analysis methods. (a–c) Quantitative evaluation of NDreamer and other methods on the PBMC (a), virus (b), and PBMC1 (c) datasets. (d) Scatterplot of average biological conservation score against the average condition-mixing score for each method. (e) Quantitative evaluation of NDreamer and other methods on the ECCITE dataset from the aspects of batch mixing, biological conservation, and condition mixing. (f) Quantitative evaluation of NDreamer and other methods on the ASD dataset from the aspects of batch mixing, biological conservation, and condition mixing. (g) UMAP of the effect modifier embedding of NDreamer and CINEMA-OT, colored by condition, batch, and cell cycle phase.

We then compared NDreamer to other methods using the virus and PBMC1 datasets. As shown in Fig. 3b and c, Supplementary  Figures  S2 and S3, NDreamer consistently outperformed the other methods. When averaging biological conservation and condition-mixing performance across all benchmarked methods on the three datasets, NDreamer emerged as the only method that achieved strong performance in both aspects (Fig. 3d). In contrast, Mixscape demonstrated limited condition-mixing capability, while the other three methods showed inadequate performance in preserving biological conservation.

However, in some single-cell experimental perturbation studies, multiple batches are present, and batch effects can significantly influence analysis results. This challenge is further amplified in genetic perturbation studies, where datasets often contain tens of genetic mutations. Previous methods, which only support comparisons between two conditions and lack batch-effect removal functionality, are less applicable in such scenarios. We illustrate this using the ECCITE and ASD datasets.

As shown in Fig. S4a, the raw data of the ECCITE dataset exhibit pronounced batch effects, with clear separation between replication 1 and the other replications. We compared NDreamer (with batch-effect removal enabled) against four other methods in comparing genetically perturbed and unperturbed cells. In Fig. 3g, CINEMA-OT fails to mitigate batch effects, displaying a clear batch-related separation in the effect modifier space.

To systematically evaluate these methods when batch effects exist, we computed condition mixing, biological conservation, and batch-effect removal performance with the same metrics for condition mixing but applied to batches instead. As shown in Fig. 3e and f, NDreamer outperformed all other methods when these three metrics were combined, delivering embeddings with excellent condition-mixing and batch-mixing performance (Fig. S4 and S5). In addition, we also benchmarked these methods by treating each genetic mutation as a condition and enabling batch-effect removal, if applicable, where we found that NDreamer still outperformed all other methods (Fig. S6). The detailed results of the experimental perturbation study benchmarking are available in Supplementary Table S1.

Systematic evaluation of NDreamer: Does it produce condition-related signal-preserved expression with batch effects removed?

Evaluation metrics

In addition to extracting effect modifiers and estimating ITE in experimental perturbation studies, NDreamer can also be applied to observational studies to denoise condition-related expressions from raw expressions influenced by batch effects and estimate CRS. To systematically evaluate its performance, we benchmarked NDreamer against current state-of-the-art models, including CellANOVA [12], scDisInFact [26], AIF [27], and scDML [28], using the default hyperparameter settings for these methods. To systematically evaluate which method can produce denoised expressions that are both free from batch effects and preserve condition-related signals, drawing inspiration from previous studies [12], we conducted benchmarking based on four evaluation aspects (Fig. 4a):

Figure 4.

Benchmarking of NDreamer and other integration or batch effect correction methods on producing condition-related signal-preserved denoised expression. (a) Illustration of the criteria of a good batch-effect-removal method. For the denoised expression, the batches should be well-mixed within one condition, cells in different conditions should be well-separated, and the variation between cell types should be retained. (b) Evaluation of condition-related signal preservation using the neighboring cell condition versus central cell condition matrix for NDreamer, CellANOVA, scDisInFact, and AIF. (c) Evaluation of data distortion. The distributions of the PCC between raw and denoised expression of each gene’s logit (adjusted P-value) obtained from DEG analysis across 4 methods. (d) Radar plot of method comparison using five evaluation perspectives: Data distortions in the denoised expression; condition-related signal preservation for denoised expression (signal preservation); batch-condition-mixing performance for the effect modifier embedding (EM-batch condition mix); biological conservation for the effect modifier embedding (EM-bio conservation); batch-mixing performance within each condition for the denoised expression (denoised-batch mixing) on the T1D dataset. (e) The overall score comparison between the methods of NDreamer, CellANOVA, scDisInFact, and AIF across the T1D, kidney, mouse, ECCITE, and ASD datasets. (f) Radar plot of methods comparison using five evaluation perspectives (upper) and the overall score comparison (lower) on the simulation dataset. (g) Evaluation of MSE between the predicted and ground truth ITE in the simulation dataset (left) and mean PCC across genes between the predicted and ground truth denoised expression (right).

1. Condition-related signal preservation for denoised expression. In single-cell studies, cells from different conditions (e.g. disease versus healthy control) are expected to exhibit distinct differences, even when they belong to the same cell type or share similar intrinsic features. Capturing these cross-condition differences is crucial for understanding condition-specific biological variation. To assess whether these differences are preserved after integration, same as CellANOVA, we first identify each cell’s 30 out-of-batch nearest neighbors (oobNNs) within the integrated embedding. Out-of-batch nearest neighbors are defined as cells from other batches (or samples) that are closest to a given cell in the integrated space. This approach ensures that comparisons are not confounded by batch-specific biases. For each cell, we then calculate the proportion of its oobNNs that belong to the same condition as the cell itself. If a cell represents a state with condition-specific differences, a higher proportion of its oobNNs is expected to come from the same condition. For quantitative comparison, we aggregate these proportions across cells and stratify them by condition, visualizing their density distributions. Effective integration is indicated by a higher proportion of neighbors belonging to the same condition as the central cell, reflecting the successful recovery of condition-specific signals. Conversely, failure to capture these signals results in uniform distributions, lacking enrichment for same-condition neighbors. To formalize this, we construct a neighboring cell condition versus central cell condition matrix. The value at the i-th row and j-th column in this matrix represents the average proportion of oobNNs from the j-th condition for cells in the i-th condition. The trace of this matrix, normalized by the number of rows, serves as a performance measure: a value closer to 1 indicates that the matrix closely resembles an identity matrix, signifying better recovery of condition-specific signals.

2. Data distortion in the denoised expression. While denoised data inevitably differ from raw expression data, it is critical that denoising does not distort major biological differences already present within each condition and batch. Because cell-type-specific expression patterns are expected to be largely preserved after appropriate denoising, we used differentially expressed gene (DEG) concordance between raw and denoised data as a data-distortion check. For example, a marker gene for a given cell type in the raw data should remain a marker for the same cell type after denoising, rather than being attenuated, lost, or spuriously assigned to another cell type. Specifically, DEGs between one cell type and all other cells within each condition-batch combination (CBC) should remain as consistent as possible between raw and denoised expressions. To ensure a fair comparison, raw and denoised expressions were analyzed using the same DEG pipeline, including the same cell labels, CBC stratification, gene set, statistical test, and Benjamini-Hochberg (BH) adjustment. We computed the BH-adjusted P-values for each gene in each cell-type-by-CBC comparison for both raw and denoised data. The Pearson correlation coefficient (PCC) between the logits of adjusted P-values from raw and denoised data was calculated for each comparison, with the mean PCC across all CBCs serving as the first metric. In addition, because a method could preserve DEG ranking while changing the scale of statistical significance, we evaluated DEG classification consistency using an adjusted P-value cutoff of 0.05. Using DEGs identified from raw expression data as the reference and those from denoised data as predictions, we calculated prediction accuracy and F1 score as the second metric.

3. Batch-mixing performance within each condition for denoised expression. After batch-effect removal, the cell’s expression should be determined only by its intrinsic features and the condition it belongs to. As a result, cells from different batches within the same condition should be well-mixed across batches. To evaluate this, we average the batch-mixing metrics ASW_batch, bLISI, kBET, and True positive rate that are applied to the batch across conditions.

4. Batch-condition-mixing performance and biological conservation for the effect modifier embeddings. NDreamer, CellANOVA, and scDisInFact all first identified a batch- and condition-free embedding for each cell, then generate a denoised expression based on these embeddings. As a result, the prerequisite of generating a reliable denoised expression is that the embeddings must be well-mixed across conditions and batches, and preserve the within-batch or within-condition variations between cell types. To evaluate this, we applied the four batch-mixing metrics: ASW_batch, bLISI, kBET, and True positive rate for batch-condition-mixing evaluation, and ARI, ASW_label, cLISI, and NMI for biological conservation evaluation.

In addition to these metrics, we computed an overall score by first averaging the metrics within the same aspect, then linearly scaling these values of each aspect to a range of 0.2 to 1, and then averaging the scaled values across all metrics. Since scDML only generates embeddings for cells instead of denoised expression, it is only evaluated on the two effect modifier space-related evaluation aspects, and the scores for other aspects are set to 0.

Real datasets

Besides the above-mentioned ECCITE and ASD datasets, we also used three additional large datasets with obvious batch effects (Fig. S6) to evaluate the batch-effect-removal and condition-related-signal-preservation performance: (1) The type 1 diabetes (T1D) dataset that contains 69,645 cells from 24 batches under three conditions (11 healthy individuals, 5 individuals with T1D, and 8 individuals with no clinical presentation of T1D but positive for β-cell autoantibodies (AAB)) [29]. (2) The human kidney multiomics atlas (kidney) dataset containing 282,610 cells from 36 donors and 47 batches with 24 ‘Control’ batches and 23 ‘Disease’ batches, where 29 of them are scRNA-seq data, and 18 of them are snRNA-seq data. The analysis of this dataset includes multi-modal data integration [30]. (3) The mouse radiation therapy (mouse) dataset containing 91,752 cells from the epithelial and lamina propria layers of the intestinal segments sequenced at 4 distinct times by two separate technical groups from mice under the control and radiation-treated conditions [12]. We will focus on illustrating the T1D dataset in the following paragraph (Fig. S7).

We first evaluated the condition-related signal preservation. As shown in Fig. 4a and b, the oobNNs for NDreamer’s denoised expression are predominantly from the same condition, resulting in a neighboring cell condition versus central cell condition matrix that closely resembles an identity matrix, with a trace similarity of 0.973. In contrast, the other two methods exhibit less effective separation between conditions, with CellANOVA achieving a trace similarity of 0.442 and scDisInFact achieving a trace similarity of 0.624.

Next, we evaluated the extent of data distortion. We began by visualizing the distribution of PCCs between the adjusted P-values of genes in the raw and denoised expressions for each cell type and CBC. As shown in Fig. 4c, NDreamer achieved the highest mean and median PCC values, outperforming both CellANOVA and scDisInFact.

We further benchmarked NDreamer against previous methods using the metrics of batch-mixing performance within each condition for denoised expression, batch-condition-mixing performance for the effect modifier embeddings, and biological conservation for the effect modifier embeddings. The results are presented in Fig. 4d. NDreamer outperforms all prior methods across all evaluation metrics, with the sole exception of a marginally lower biological conservation score than CellANOVA for the effect-modifier embeddings.

We further benchmarked NDreamer against CellANOVA, scDisInFact, and AIF using additional datasets, including the ECCITE, ASD, kidney, and mouse datasets, and calculated the overall scores. Due to the absence of fine-grained cell-type annotations for the mouse dataset, biological conservation metrics were not evaluated for this dataset. Across all datasets, NDreamer consistently outperformed the other methods with the highest overall scores (Fig. 4e, Supplementary Table S2).

Simulation dataset

Besides real datasets, we also performed a simulation study where we know the ground truth ITE/CRS and the ground truth denoised expression to benchmark NDreamer with previous studies. To generate the simulation data, we first used Symsim [31] to generate the baseline gene expression matrix, where cells are designed to come from different cell types. We then randomly assign cells to six batches: three assigned to the control condition and three assigned to the perturbed condition. Then, conditioning on the cell type information, perturbation-induced effects are added, where 60% of the genes are randomly set to be up- or down-regulated with three levels of regulation. The perturbation-induced effects are then added to the baseline expression for cells in the perturbation group. After that, batch effects are added to the expression matrix using Symsim, and the raw counts are generated through negative-binomial sampling.

We compare NDreamer with CellANOVA, AIF, scDML, scDisInFact, scCAPE, Mixscape, scGen, and CINEMA-OT on the simulation dataset using the five evaluation metrics illustrated above, if applicable. As shown in Fig. 4f, NDreamer outperforms all previous methods across nearly all metrics, with only moderately strong performance on biological conservation in the effect modifier space, and achieves the highest overall score. We then compared the predicted denoised expression with the ground truth using the mean gene-wise PCC. For each gene, we computed the PCC between the predicted and ground-truth denoised expression vectors across all cells, and then averaged the PCCs across genes. A higher average PCC indicates a more accurate estimation. As shown in Fig. 4g, NDreamer achieved the highest average PCC compared with CellANOVA, scDisInFact, and AIF. For the estimated single-cell-level perturbation-induced signal, we compared the estimated values with the ground truth using mean squared error, where a lower MSE indicates better estimation. NDreamer also achieved the lowest MSE compared with CINEMA-OT and Mixscape.

NDreamer uncovers differential gene expression patterns in a large Alzheimer's disease cohort

In addition to benchmarking NDreamer against other methods, we applied it to the SEA-AD dataset [32]. This dataset includes 1 378 211 single-nucleus RNA-sequenced (snRNA-seq) cells derived from the MTG region of 47 donors labeled as 'normal' and 42 donors labeled as 'dementia' in the 'disease' annotation (Fig. 5a). Identifying expression pattern shifts between cells in Alzheimer’s disease patients and healthy controls has long been a challenge. Our objective here is to estimate the CRS for each cell and perform statistical tests for the CRSs in certain cell populations of interest.

Figure 5.

NDreamer uncovers differential gene expression patterns in the SEA-AD dataset. (a, b) Overview of the SEA-AD MTG snRNA-seq dataset. There are 47 donors labeled as 'normal' and 42 donors labeled as 'dementia' in the 'disease' annotation. UMAP visualization of the raw expression colored by cell type (a) and condition (b). (c) The GO enrichment analysis results for the common top upregulated genes for all cell types (left) and the UMAP visualization of these genes on the UMAP of raw expressions. (d–g) The GO enrichment analysis of the top down-regulated genes (left) and the UMAP visualization of the most significant DEGs except for the commonly up-regulated genes based on the raw expression for (d) astrocytes, (e) oligodendrocytes, (f) Lamp5 inhibitory neurons, and (g) excitatory neurons.

We began by visualizing the estimated effect modifier embedding, denoised expression, and CRSs of each cell using UMAP. As shown in Fig. S8, the effect modifier embedding demonstrates good mixing of conditions and batches (Fig. S9), while the denoised expression effectively separates cells from different conditions. Additionally, distinct clusters corresponding to different cell types (the 'Subclass' annotation) are evident across the effect modifier embeddings, denoised expressions, and CRSs. These results indicate that NDreamer successfully extracts the condition- and batch-independent effect modifier, preserves condition-related signals and intrinsic cell feature variations, and effectively removes batch effects.

After estimating the CRS and aggregating it by cell type to calculate and statistically test the CACRS, we observed that most cell types share the same top 100 upregulated genes with the smallest adjusted P-values (Supplementary Table S3), including ATP6, ND2, ND3, ND4, COX2, and CYTB, while showing downregulation of the FTH1 gene (p.adj < ). These estimates align with the UMAP visualization of the raw expression data (Fig. 5b and c), where regions with higher dementia cell density exhibit altered expression levels of these genes compared to regions with higher normal cell density. Furthermore, our findings corroborate previous studies, where increased expression of mitochondrial genes [33, 34] and disruption in iron metabolism are identified as the hallmarks of oxidative damage in AD brains [35]. To further investigate, we identified the intersection of the top 100 genes with the lowest adjusted P-values for the CACRS in each cell type and conducted a gene ontology (GO) enrichment analysis. This analysis revealed that these genes are strongly associated with cellular respiration and energy metabolism, consistent with prior research highlighting abnormalities in these processes as key features of AD pathology [36, 37].

We then conducted a statistical analysis of the CRSs for astrocytes. In addition to the previously mentioned upregulated genes common across most cell types, the most significantly downregulated genes included PCDH9-AS2, ALDH1A1, and PPP1R9A-AS1 (p.adj < , Fig. 5d), which have also been reported in previous studies [38–40]. Next, we examined the GO terms enriched among the top 100 downregulated genes. These genes were primarily associated with the positive regulation of immune cell proliferation, which is also consistent with previous research [41].

Next, we analyzed the CRS and CACRS for oligodendrocytes (oligo). In addition to the commonly upregulated genes, we observed increased expression of PLCG2 for oligo under the dementia condition, while CARNS1 and LDB3 were identified as the most significantly downregulated genes (p.adj < ). These findings are consistent with previous studies [42–44]. GO enrichment analysis of the top 100 downregulated genes further highlights the vulnerability of oligodendrocytes and disruptions in myelin integrity in dementia conditions compared to normal conditions [42]. For the Lamp5 inhibitory neuron cell type, CXCL14, TRPC3, and KIT genes are the most significantly down-regulated genes, which agree with previous studies [45], and indicate the abnormality of the neuroinflammation [46], which are also reflected by the enriched GO terms.

Lastly, we observed that most excitatory neurons shared the same top downregulated genes, including TIMP3, CTNNA3, and CNTN5 (p.adj < ), which have been previously reported [47, 48] and are known to be associated with amyloid precursor protein processing [49] and tau pathology [50]. GO enrichment analysis of the top 100 downregulated genes further revealed the enriched pathways related to chemotaxis, inflammation, and ion transport.

In addition, we performed donor-level testing to avoid inflation of statistical significance from treating individual cells as independent biological replicates. For each cell type, CRS estimates were averaged within each donor, and statistical testing was then performed across dementia donors. The genes highlighted above remained among the top 10 genes with the lowest donor-level adjusted P-values (Supplementary Table S4, p.adj < ) in the corresponding cell type, indicating that these disease-associated CRS patterns are reproducible across donors rather than being driven solely by large cell-level sample size. For comparison, NDreamer identified an average of 897 significant genes across cell types, whereas Seurat identified an average of 431 significant genes across cell types under the same cutoff. The top AD-associated genes identified by NDreamer discussed above were also among the top 10 most upregulated or downregulated genes identified by Seurat in their corresponding cell types.

Exploration of batch effects, individual treatment effect, and sensitivity test

Exploration of batch effects

The batch effect for each cell can be directly estimated by subtracting the denoised expression from the raw expression, enabling visualization of the composition of batch effects. As shown in Fig. 6a, batch effects were found to be cell-type-specific, with each batch exhibiting unique batch effects, similar to the estimated batch effects by CellANOVA. This suggests that the true batch effects are complex functions of cells’ intrinsic features and batch-specific characteristics. To further explore this, we plotted batch effects against factors contributing to batch composition, such as scRNA-seq versus snRNA-seq in the kidney dataset and sequencing time in the mouse dataset. Consistent with previous studies [12], we found that the batch effect is closely associated with these factors, highlighting the multifaceted nature of batch effects.

Figure 6.

Exploration of batch effects, individual treatment effect, and sensitivity test. (a) UMAP visualizations of the estimated batch effects in the ECCITE dataset, colored by cell cycle phase and batch; the kidney dataset, colored by cell type and data modality; and the mouse datasets, colored by compartment and sequencing date. (b) The number of DEGs with BH-adjusted P-values below 0.05, tested on CD8T cells in the PBMC dataset, for the NDreamer and CINEMA-OT methods. (c) UMAP visualization of the top 2 most significant DEGs for the CD8T cells in the PBMC dataset estimated by NDreamer and CINEMA-OT. (d) The number of DEGs with BH-adjusted P-values below 0.05, tested on CD8T cells in the randomly permuted PBMC dataset, for the NDreamer and CINEMA-OT methods. (e) The ablation study of NDreamer on the PBMC dataset. The condition-mixing metrics and the biological conservation metrics are used. We include the models where we remove the dependent loss (No_Dependent), the local neighborhood loss (No_Neighbor), the triplet loss (No_Triplet), all three mentioned biological conservation-related losses (No_Structure-Related), the independent loss (No_Independent), and all the previously mentioned losses (VQ-VAE).

Testing for skewness

A critical requirement for the estimated ITE is that it should not be skewed, where the ITE values for one gene for cells in a certain group should be symmetric around 0 when there is no systematic biological difference in the expression level of this gene in this cell group between two conditions. To test this, we compared the ITEs estimated by NDreamer to those of CINEMA-OT, an innovative and powerful method for estimating single-cell-level ITE causal effect, calculating BH-adjusted P-values for the ITEs of CD8T cells using a nonparametric method to test if ITEs are significantly larger or smaller than 0 using the above-mentioned method. As shown in Fig. 6b, NDreamer identified only 25 genes with adjusted P-values less than 0.05. However, CINEMA-OT reported 1971 out of 2000 highly variable genes as significant. This implausibly high number likely reflects skewness in CINEMA-OT's ITE estimation. We then visualized the two most significant genes with the lowest adjusted P-values (Fig. 6c) by plotting the UMAP colored by the expression level of these genes and the distribution of ITEs for these genes in CD8T cells. Genes identified by NDreamer exhibited clear differences between the two conditions, while those identified by CINEMA-OT appeared to lack meaningful significance, likely arising from skewness. To further investigate, we randomly permuted condition labels for each cell and re-ran both methods. In this scenario, where significant P-values should be rare and driven only by random chance, NDreamer identified only two genes with adjusted P-values below 0.05 out of 2000 HVGs. In contrast, CINEMA-OT still reported 1977 genes as significant (Fig. 6d). This finding strongly suggests that the ITEs estimated by CINEMA-OT may be skewed towards negative.

Nearest neighbor matching is all you need

Nearest neighbor matching, optimal transport matching, and k-nearest neighbor weighting are three widely used methods for covariate adjustment in estimating ITEs in causal inference. However, optimal transport and k-nearest neighbor weighting methods may introduce additional skewness due to the sparsity and high proportion of zeros in single-cell data. Calculating counterfactual cells using (weighted) averages of multiple cells can result in a loss of sparsity in the counterfactual cells, potentially introducing skewness when subtracting the counterfactual expression from the reference cell expression. As a result, theoretically, any weighting method, if not one-to-one matching, would introduce the skewness towards the negative direction in the ITE estimation in scRNA-seq data, as shown in the ITEs of CINEMA-OT, where most ITEs are negative. A detailed proof of why weighting-based methods are skewed is provided in Supplementary Notes S4. In addition, since cells from different conditions and batches are well-mixed in the NDreamer-produced effect modifier space, there is no need to use optimal transport to adjust for the imperfect condition mixing performance as in previous methods.

Ablation study

To identify the key components of NDreamer’s deep learning model, we conducted an ablation study using the PBMC dataset (Fig. 6e, Supplementary  Figure  S10). The study shows that each loss function contributes to either condition-mixing performance or biological conservation. Specifically, removing the independent loss function (No_Independent) significantly decreases condition-mixing performance, demonstrating that the mutual information-based independent loss effectively identifies condition-independent effect modifiers. Similarly, eliminating the triplet loss (No_Triplet), local neighborhood loss (No_Neighbor), dependent loss (No_Dependent), or all three structure-related loss functions (No_Structure-Related) results in a notable reduction in biological conservation metrics. This highlights the critical role of these losses in preserving local and global biological structure. Moreover, comparing NDreamer to the original VQ-VAE model reveals a dramatic improvement across all metrics. This demonstrates that the inclusion of these loss functions successfully adapts NDreamer for causal inference in single-cell data. We also tested the sensitivity of NDreamer to the number of neighbors in the local neighborhood loss. As shown in Supplementary Table S5, the model’s performance is not sensitive to the change of the number of neighbors, while when it is set to 20, it achieves the highest average score of NMI and kBET.

Discussion

Advances in single-cell sequencing techniques and the growing volume of single-cell data have created unprecedented opportunities for uncovering statistically significant gene expression patterns induced by perturbations or associated with specific disease conditions. However, existing analytical approaches face significant limitations. Many methods for estimating experimental effects at the single-cell level are either skewed or ineffective at extracting condition-independent effect modifiers while preserving biological variation. Moreover, they struggle to scale for large experimental studies involving multiple batches with batch effects. Similarly, current batch-effect removal methods often overlook the non-linear characteristics of single-cell data, fail to achieve adequate batch- and condition-mixing performance, or remove the condition-related signals while correcting for batch.

To address these challenges, we developed NDreamer, a robust framework designed to denoise expression profiles while preserving condition-related signals and effectively estimating individual treatment effects at the single-cell level. Through the combination of loss functions that ensure batch- and condition-independent while preserving global and local biological structure, NDreamer outperformed existing methods across all datasets. Additionally, with optimized tensor computation designs for the mutual information loss and the triplet loss calculations, NDreamer achieves exceptional computational efficiency, finishing model training on the PBMC dataset (2 conditions, 16,893 cells) in under one minute and on the large SEA-AD dataset (1 378 211 cells, 89 batches) in 3837 seconds (Supplementary Table S6) with GPU memory less than 8GB when the batch size is set to 2048. The detailed summary of the running time and GPU usage for all datasets analyzed in the manuscript is provided in Supplementary Table S6. The GPU memory increases linearly with batch size, and the computation time linearly increases with the number of batches, the number of conditions, and the number of cells in the dataset.

However, it is essential to carefully examine the data before applying NDreamer. For causal interpretation in experimental perturbation studies, NDreamer relies on an effect-modifier representation that is minimally dependent on condition and batch. In observational disease-control studies, however, the estimated CRS should be interpreted as condition-associated rather than causal. In some cases, however, the true effect modifier may not be entirely independent of these factors, where NDreamer may not be ideal in such cases. For instance, in the ASD dataset, all cells of the "Vascular" cell type are unperturbed, making the condition inherently linked to the intrinsic features of the cell. Similarly, in the kidney dataset, certain cell types are common in the scRNA-seq modality but are rare in the snRNA-seq modality. Without additional information beyond the expression matrix, condition annotations, and batch annotations, these dependencies cannot be fully resolved based on information theory principles. We also acknowledge that, in the comparison of scRNA-seq data between patients and healthy controls, the effect modifier may not be fully independent of the condition or batch (e.g. there might be cell type composition differences between patients and healthy controls). In those cases, the three loss functions in NDreamer are designed to retain biological conservation, which will help correct this unwanted correlation. The model can still effectively extract the effect modifiers while preserving biological conservation, but the biological conservation metrics may inevitably experience a slight decrease. A practical suggestion in these situations is to exclude cells that inherently carry batch or condition information. For example, in the ASD dataset, if a cell is identified as belonging to the "Vascular" cell type, it can be excluded since it directly indicates the negative control condition. This preprocessing step can help mitigate potential confounding and enhance the reliability of downstream analyses.

To our knowledge, this is the first study to integrate neural discrete representation learning with mutual information loss to establish independence, demonstrating the practical applicability of mutual information loss beyond its theoretical justification. Unlike previous methods that rely on adversarial training and discriminators to enforce independence, our approach is significantly more efficient, especially when the encoder is large (e.g. Transformers), as it reduces the need for multiple backpropagation steps required for discriminator training. Additionally, neural discrete representation learning offers distinct advantages over conventional VAEs, where the latent representations are continuous vectors, particularly when auxiliary discrete-related objective functions, such as the mutual information loss, are involved.

Finally, our theoretical analysis and empirical results suggest that one-to-one matching can avoid the negative skew introduced by weighting-based approaches under the sparse single-cell settings considered here and in experimental perturbation studies; otherwise, weighting may introduce the skewness towards the negative direction for the estimated ITE due to the sparsity of single-cell data. Although there is no formal theoretical guarantee for the nearest neighbor matching, it still performed the best among the approaches considered here.

Key Points

• We developed NDreamer, a computational framework for single-cell-level perturbation-induced or condition-related signal estimation. NDreamer explicitly accounts for batch effects in single-cell perturbation-induced signal estimation and effectively removes batch effects using causal matching.

• NDreamer demonstrates superior performance in comparative benchmarks spanning batch mixing, condition mixing, biological conservation, data distortion, and signal preservation.

• NDreamer uncovers Alzheimer’s disease-related signals at single-cell resolution, providing new insights into immune and oxidative damage.

Author contributions

Xiao Xiao: Conceptualization, Methodology, Software, Formal analysis, Writing—original draft. Hongyu Zhao: Methodology, Formal analysis, Writing— review & editing. Zuoheng Wang: Supervision, Funding acquisition, Methodology, Formal analysis, Writing—review & editing. All authors read and approved the final manuscript.

Conflict of interest

None declared.

Funding

This study was supported by the National Institutes of Health (NIH) grant R01LM014087.

Data availability

The PBMC dataset is available from the GEO under accession number GSE96583 and can be downloaded from the scGen tutorial https://drive.google.com/uc?id=1r87vhoLLq6PXAYdmyyd89zG90eJOFYLk. The rhinovirus infection dataset and another PBMC dataset are available at https://datadryad.org/stash/dataset/doi:10.5061/dryad.4xgxd25g1. The ECCITE-seq dataset is available from the GEO under accession number GSE153056 and can be downloaded from pertpy [2] using the code ‘pertpy.data.papalexi_2021()’. The ASD dataset is available from the GEO under accession number GSE157977. The T1D dataset is available from the GEO under accession number GSE148073 and can be downloaded from https://cellxgene.cziscience.com/collections/51544e44-293b-4c2b-8c26-560678423380. The kidney dataset is available from the GEO under accession number GSE211785, and we use the version ‘GSE211785_Susztak_SC_SN_ATAC_merged_PreSCVI_final.h5ad.gz’. The mouse radiation therapy dataset is available from the GEO under accession number GSE280883. The SEA-AD MTG snRNA-seq dataset is available at https://cellxgene.cziscience.com/collections/1ca90a2d-2943-483d-b678-b809bf464c30. The SEA-AD dataset is available at https://cellxgene.cziscience.com/collections/1ca90a2d-2943-483d-b678-b809bf464c30.

Code availability

The NDreamer package is available at https://github.com/lugia-xiao/NDreamer with detailed tutorials. All code used to reproduce the results shown in the article are available at https://github.com/lugia-xiao/NDreamer_reproducible.