Conditional Causal Representation Learning for Heterogeneous Single-cell RNA Data Integration and Prediction

1 Introduction

With the rapid advancement of single-cell RNA sequencing (scRNA-seq) technology, vast datasets are being generated across diverse laboratories and experimental conditions [Jovic et al., 2022]. These datasets often arise from different sequencing methods, perturbation conditions, or disease samples, resulting in significant heterogeneity. Efficiently modeling, integrating, and comparing such multidomain data has become a major challenge [Stuart et al., 2019; Argelaguet et al., 2021; Hao et al., 2021]. Addressing this issue is crucial for capturing cellular heterogeneity and dynamic changes, ultimately uncovering molecular mechanisms driving phenotypic diversity and disease susceptibility.

∗Corresponding author. Our code and supplementary information are available at https://github.com/fdu-wangfeilab/scConCRL.

Deep generative models have shown great potential in analyzing multi-domain single-cell data but struggle to distinguish biological signals within domains from non-biological signals between domains. Current integration methods typically preserve shared information but often introduce spurious correlations due to a limited identifiable foundation, leading to confusion or loss of biological variations [Aliee et al., 2024]. Consequently, generalization ability is restricted. Causal Representation Learning (CRL), an emerging technique, offers a promising solution by identifying latent factors and causal structures with causal explanations from observational data [Sch¨olkopf et al., 2021]. CRL accurately reflects causal relationships in biological signals and enhances model generalization in unseen scenarios.

In single-cell transcriptomics, gene expression levels represent cellular states but are often confounded by batch effects, drug perturbation, and other covariates. Existing integration methods focus on retaining invariant signals while removing potentially confusing signals [Luecken et al., 2022]. While effective for tasks like cell type annotation, this approach oversimplifies the dynamic cellular generation process and yields inaccurate reconstruction results. Methods like inVAE [Aliee et al., 2024] aim to only separate invariant and domain-specific signals to solve the issue but lack semantic understanding of latent variables and ignore dependencies between them, reducing model interpretability.

Many CRL methods assume independence among latent variables [Tejada-Lapuerta et al., 2023]. However, biological systems often exhibit strong dependencies, such as interactions in gene regulatory networks (GRNs). Ignoring these relationships hinders the model's ability to capture true causal mechanisms and limits its applicability in biological contexts. To effectively apply CRL to single-cell scenarios, it is essential to model these complex interactions and elucidate how gene relationships influence cellular state changes.

To address these challenges, we propose scConCRL, an interpretable conditional causal representation learning framework for integrating and predicting multi-domain scRNA-seq data. Our framework incorporates Structural Equation Modeling (SEM) to explicitly capture relationships between confounding domain variables and intrinsic molecular variables.

Genes are modeled as core molecular variables, while batch effects, perturbation conditions, and other domain information are encoded as latent domain variables. By leveraging a conditional disentanglement strategy, scConCRL performs batch effect correction and noise removal without requiring standard references. Furthermore, it models latent domainmolecular variable relationships, enabling robust perturbation predictions and enhanced analysis of complex biological systems.

Our contributions are as follows:

• We propose scConCRL, a causal representation learning framework that identifies both invariant and domainspecific signals, enabling reference-free integration in molecular space and improving generalization across datasets.
• We introduce a strategy that employs Structural Equation Modeling (SEM) to model domain variables (e.g., batch effects, perturbation) and molecular variables, enabling the framework to capture interactions between coarse-grained and fine-grained variables, thereby enhancing interpretability.
• We conduct comprehensive experimental evaluations comparing scConCRL with existing state-of-the-art methods and demonstrate scConCRL's superior performance in several tasks such as GRN inference, data integration, and perturbation prediction.

2 RelatedWork

2.1 Single-cell Integration

Single-cell multi-domain integration aims to combine data from various technical platforms or experimental conditions to uncover cell types, states, and functions. Traditional methods, such as Canonical Correlation Analysis (CCA) [Hao et al., 2023] and Mutual Nearest Neighbors (MNN) [Haghverdi et al., 2018], align data from two domains by minimizing discrepancies in their latent representations [Haghverdi et al., 2018; Xu et al., 2022]. These approaches are practical and interpretable but face challenges with multi-domain data. Deep learning-based methods [Lopez et al., 2018; Li et al., 2020b; Xiong et al., 2022; Yu et al., 2023] use unsupervised learning for data compression and align multi-domain distributions. However, this often achieves alignment only in the latent space, lacking interpretability and limiting their broader applicability in real-world scenarios.

For molecular-level integration tasks, some deep learning approaches adopt reference-based strategies and map data from different batches to a specific reference batch, but their performance depends heavily on the quality of the reference [Xiong et al., 2022]. In contrast, reference-free methods, such as Scanorama [Hie et al., 2019] and Beaconet [Xu et al., 2024], eliminate batch effects using global distribution information across datasets. While these methods avoid reference bias, their broad constraints can introduce artifacts that distort true biological variations. To address these issues, we propose a reference-free approach that extracts invariant signals and uses mechanism-driven constraints, enabling efficient integration of multi-domain data in the original molecular space while maintaining data quality.

2.2 Single-cell Condition Prediction

Single-cell condition prediction leverages sequencing data to forecast gene expression changes under various conditions. Variational Autoencoder (VAE)-based methods, such as scGen [Lotfollahi et al., 2019], achieve this by learning latent data distributions and identifying distribution shifts. Extensions like trVAE enhance cross-condition predictions by incorporating conditional variables, enabling the modeling of more complex conditions. scPreGAN [Wei et al., 2022] further improves prediction accuracy by capturing shared conditional signals and simulating the perturbation transfer process through adversarial training. Additionally, scShift [Dong and Kluger, 2023] incorporates advancements in causal representation learning to learn conditional and biological patterns, enabling more precise condition prediction. On this basis, scDisInFact [Zhang et al., 2024b] models two different types of domains separately, enabling combined predictions for multiple batches and conditions. Despite significant advances in single-cell condition prediction, challenges remain, particularly regarding the interpretability of models. In this context, scConCRL integrates latent condition and molecular variables, enabling the exploration of how condition influences molecular variables. This framework enhances interpretability and improves condition prediction accuracy.

2.3 Causal Representation Learning

Causal Representation Learning (CRL) aims to identify latent factors and uncover causal structures from observed data, providing interpretable representations. While VAEs excel at representation learning, they struggle to identify true latent factors in nonlinear contexts. Recent improvements, including the use of auxiliary variables and temporal structures, have enhanced VAEs' ability to capture causal mechanisms [Hyvarinen et al., 2019; Khemakhem et al., 2020; Yang et al., 2021]. In multi-domain settings, CRL extracts causal invariances across domains (e.g., different experimental conditions, environments, or datasets), revealing shared mechanisms [Zhang et al., 2024a]. For example, it has been validated on fMRI datasets to identify common causal factors affecting brain activity across conditions. In single-cell analysis, inVAE separated latent representations into invariant and domain-specific components, aiding the study of causal invariance [Aliee et al., 2024]. However, it did not deeply address the interpretation or relationships of latent variables. scConCRL advances this by explicitly modeling relationships between domain-dependent and molecular latent variables, enhancing interpretability and uncovering causal mechanisms in cellular processes.

3 Method

To establish some notation, let X ∈ Rn×m denote a singlecell gene expression matrix, where n represents the number of cells and m represents the number of genes. The matrix X mixes both biological signals and domain-specific information, such as technical noise or perturbation signals. We denote B as the variable encoding domain information (e.g., batch or condition). Our goal is twofold: first, to automatically disentangle the k-dimensional latent components, corresponding to batch effects, perturbation signals, and the basal component of cell state signals; and second, to uncover the underlying common GRN, represented by the weighted adjacency matrix Axx ∈ Rm×m, which is independent of the batch or condition to which each sample belongs, and the relationship between conditions and genes, captured by the matrix Axb ∈ Rm×k .

Confounding factors obscure the discovery of GRNs governing biological processes, leading to discrepancies across GRNs from different domains. scConCRL aims to derive common guidance across domains, inspiring GRN inference and accurately generating cellular information. By incorporating domain-specific factors that affect gene expression, scConCRL accounts for domain information, effectively disentangling true regulatory relationships from domain biases.

scConCRL comprises three main components: a Variational Autoencoder (VAE), a Structural Equation Model (SEM), and a disentanglement module. The encoder maps observed data X to both molecular and domain-specific latent variables. The SEM and decoder reconstruct the observed data using self-learned relationships. In addition, the disentanglement module encourages the latent representations to be partitioned into distinct variable sets. In this section, we provide a detailed description of each component and the optimization objectives.

3.1 Assumptions on the Generative Process

We aim to model the generative process of multi-domain scRNA-seq by exploring latent causal generative models. The observed data x is assumed to be generated by underlying latent causal variables, z, with the possibility of complex causal relationships among these variables. In addition, we introduce latent noise variables, n = {nb , nx }, that are exogenous to the system. Specifically, nx represents the basal biological factors, and nb represents domain-specific factors from batch effects or perturbations.

Furthermore, external domains b, such as technical bias or experimental perturbations, can lead to changes in the distribution of nb that influence the data distribution. The introduction of b allows us to model the effects of these domain-specific factors, offering greater flexibility in adapting to varying conditions. The generative model is defined as follows:

Here, the latent noise variables nx and nb follow Gaussian distributions, with nb being modulated by the observed variable b through two nonlinear mappings βj,1 and βj,2. The choice of a Gaussian distribution is based on the identifiability and expressiveness of latent variables. The latent cell embedding z is directly generated by the function g, which is defined within a SEM. Subsequently, a nonlinear mapping function f transforms z into the observed data x, with ϵ representing independent noise. For flexibility, we do not impose a discrete probability distribution on the observed data x.

To estimate the parameters of the scConCRL model, we employ variational inference, which allows us to approximate the posterior distribution of the latent variables. We factorize the posterior distribution as follows:

where δ(·) is the Dirac delta function. After this factorization, the evidence lower bound (ELBO) can be derived through straightforward computation:

According to the above equation, the encoder processes the observed expression data x and domain information b to derive the conditional distribution qϕ(n|x, b). Another network pψ uses the domain information to compute the prior distribution p(nb |b). Then, n is transformed by the layer defined by the SEM into causal representations z, which are subsequently fed into the function f to reconstruct observed data.

3.2 Structural EquationModel with Confounder

The Structural Equation Model (SEM) provides a structured approach to generate observed gene expression data. SEM is widely used in fields like economics and social science to model dependency relationships between variables within a system [Hair Jr et al., 2021]. In its traditional form, a standard SEM for a single domain is defined by the following equations:

Here, zs represents the basic cell embedding data in a single domain, where the matrix A is a learnable parameter that encodes the relationships between different variables in the observed data, such as gene-gene interactions. n denotes independent Gaussian variables, and the Gaussian distribution shows good performance in exponential family distributions [Liu et al., 2023]. Previous studies, such as DeepSEM [Shu et al., 2021], have utilized SEM to model gene relationships.

However, when handling multi-domain scRNA-seq data, the relationships modeled by SEM may be confounded by some covariates, particularly due to batch effects [Kaltenpoth and Vreeken, 2023]. Each domain, represented by b, has its own domain-specific network G b . A common network G = {V,A} is defined as the shared structure {G b }Bb=1 across all domains. This common structure excludes domain-specific edges, helping to reveal the true relationships [Huang et al., 2023].

To control the influence of confounding factors on common structure, we apply the backdoor criterion to enhance causal inference accuracy by controlling for confounding domain variables creating "backdoor paths" between variables, illustrated in Figure 1a. To integrate both true regulatory relationships and domain-specific interference, we use latent variables zb to represent domain-specific information associated with domain information b. We then expand the SEM model to include a larger set of variables, V = {vgi }mi=0 ∪ {vb }, where g denotes gene variables. This expanded model captures both direct gene-gene interactions and domain-specific effects, yielding a unified network structure consistent across domains [Kaltenpoth and Vreeken, 2023]. The SEM model for this extended set of variables is defined as:

In this formulation, ¯ A represents the edges between domain variables and genes, while Axx captures the relationships among the genes themselves. In accordance with the common GRN structure, we set the weight to 0 for the first k columns of ¯ A, corresponding to the edges from gene/domainrelated variables to domain-related variables.

Model identifiability refers to the extent to which model parameters can be uniquely determined from the observed data [Khemakhem et al., 2020]. Achieving model identifiability is crucial as it ensures that the SEM efficiently captures the relationships among genes. Based on the assumption of [Khemakhem et al., 2020; Kaltenpoth and Vreeken, 2023], we provide the foundation of identifiability in Appendix A. Further, we promote the adjacency matrix ¯ A to conform to a directed acyclic graph (DAG) structure. Instead of employing a conventional combinatorial DAG constraint, we adopt a continuous and differentiable constraint function [Zheng et al., 2018]. This function achieves a value of 0 solely when the adjacency matrix ¯ A represents a DAG. Specifically, we define the constraint function as follows:

Here tr(·) denotes the trace of a matrix. In addition, we use Ls = || ¯A||2F to ensure sparsity.

To facilitate the training process, we introduce rough gene relationships as guides on the latent space, incorporating the concept of gene injection into the latent variables [Yang et al., 2021]. Leveraging gene relationships, represented by the adjacency matrix Axx , we impose a specific constraint to minimize the distance between the gene expression data x and its regression form. This constraint, denoted as Lu, is formulated as follows:

where κ is a small constant.

3.3 Disentangling the Domain and Invariant Factors

To disentangle the latent factors nx and nb , we follow the principles outlined by [Chen et al., 2016; Higgins et al., 2017] promoting independence among the variables. Specifically, we make the posterior distribution of nb approximate the prior distribution conditioned on the domain information b. This allows us to replace any potential correlation between nx and nb with the correlation between nx and b.

We use adversarial learning to facilitate the disentangling process. The goal is to minimize the mutual information between nx and b, ensuring that domain-specific information contained in b does not influence the biological factors in nx . Adversarial learning has proven effective in domains such as domain generalization and factor disentanglement [Zhou et al., 2022].

The primary objective is to make nx indistinguishable by a batch discriminator. This discriminator is implemented as a multi-class classifier for categorical domain variables, which predicts the batch or domain label based on the input features and is parameterized by ϕd. The training process alternates between updating the feature encoder to learn meaningful features and training the discriminator to distinguish between domains. The adversarial loss, denoted as Ladv, is defined as follows:

To tackle the issue of imbalanced sample difficulty in singlecell analysis, we incorporate focal loss at the sample level, which helps the model focus more on harder-to-classify samples [Li et al., 2020a]. At the domain level, we dynamically adjust the loss weight for each domain based on its misclassification rate, using an Exponential Moving Average (EMA), as described by the following equations:

Nk is the number of domain k, p(k) i is the logit of sample i in domain k and I is the indicator function. In conclusion, adversarial loss drives the model to separate domain-specific information from biological information, resulting in a feature space invariant to domain variations. Minimizing this loss ensures that the learned variables nx remain independent of the domain information b.

3.4 Optimization and Inference of scConCRL

The training procedure of our model is formulated as the maximization of ELBO subject to certain constraints, expressed as equations (7), (6), and (8). To tackle this constrained optimization problem, we utilize the augmented Lagrangian algorithm, which leads to the formulation of a new

Here α, β, γ, λ, and ρ are hyperparameters that control the influence of each constraint term in the optimization process. Of note, in the inference process, we utilize the relationship represented by Axx to obtain integrated clean scRNA-seq data. For condition prediction, we set the condition domain to generate data under the specified conditions.

4 Experiment

4.1 Setup

Dataset. We simulated scRNA-seq datasets using the scMultiSim R package [Li et al., 2022] with default parameters and generated clean and noisy scRNA-seq data. For real-world datasets, we combined data from mouse Embryonic Stem Cells (mESC) from four protocols, alongside four different networks from the BEELINE benchmark [Pratapa et al., 2020] for GRN inference. To evaluate data integration, we used 4 human pancreas datasets (hPancreas), 4 mouse cortex datasets (mCortex), 8 human Peripheral Blood Mononuclear Cell datasets (hPBMC), and 14 human heart datasets (hHeart), all with gold standard cell type labels. For perturbation response prediction, we used a human PBMC dataset with seven cell types, including control and interferon-beta (IFN-β)-stimulated cells (PBMC). All datasets were preprocessed with normalization, logarithmic transformation, high-variable gene selection, and maxabs scaling. More information can be found in Appendix B.

Performance evaluation. For GRN inference evaluation, given the sparsity and incomplete nature of biological ground truths, we select EPR (Early Precision Rate), AUPRC ratio, and AUROC. To assess data integration performance, we use the scIB score [Luecken et al., 2022], which combines 8 metrics across two aspects: biological conservation and batch correction. The former indicates the representation quality while the latter assesses domain disentanglement. For perturbation response prediction tasks, we evaluate prediction accuracy using R2 scores [Lotfollahi et al., 2019]. Adjusted R2 evaluates the relationship between predicted data for a specific cell type and real perturbed data across all cell types. Further details on comparison methods, metrics, and implementation specifics can be found in Appendix C, D, and E.

4.2 GRN inference

GRN inference is one of the real application scenarios of causal structure inference and we evaluate scConCRL on multi-batch mESC dataset to assess its advantages, benchmarking its performance against state-of-the-art methods. We used two sets of highly variable genes (1000 and 2000 HVGs) from mESC datasets and compared scConCRL with two deep learning-based methods (DeepSEM and RegDiffusion [Zhu and Slonim, 2023]) and two traditional methods (PIDC [Chan et al., 2017] and GRNBOOST2 [Moerman et al., 2019]). The performance of scConCRL and the four baseline methods was assessed on true biological networks of varying scales. As shown in Table 1, scConCRL consistently outperforms the other methods, achieving higher accuracy in both non-cell-type-specific and cell-type-specific GRN inference. It excels with domain adaptation mitigating dataset domain shift issues, which affects the performance of deep learning baselines (DeepSEM and RegDiffusion). PIDC and GRNBOOST2 perform well when handling 1,000 HVGs but perform poorly when handling 2,000 HVGs, which highlights that scConCRL is a better choice in handling complex scenarios. Under the 2,000 HVGs setting, we further compared the performance of combining data from all four batches with using any single batch of data. We evaluated the four baseline methods on each single batch (labeled "Batch 0-3") and on the complete dataset (labeled "All"). For these baseline methods, combining the four batches does not improve performance compared to using individual batches, primarily due to variations in data quality and noise levels (Supplementary Tables 2-3). In contrast, scConCRL consistently outperforms using any single batch, demonstrating the importance of eliminating batch-related noise and bias for multi-batch data downstream analysis.

4.3 Integration in molecular space

Data integration was conducted to disentangle true biological variability from confounding variation. We have validated the model's ability to integrate data in the molecular space across multiple datasets, which requires not only achieving distributional alignment but also preserving highquality data. We compared scConCRL with other end-to-end molecular-level integration methods, as well as denoising + batch-correction two-step strategies which can recover data quality. As shown in the results of Table 2 and Supplementary Table 4, scConCRL significantly outperforms all other methods across all datasets. Specifically, scConCRL leads by 20-30% in batch correction on the hPBMC, mCortex, and hHeart datasets, and by 5% on the hPancreas dataset. In terms of biological conservation, scConCRL excels beyond all end-to-end molecular-level integration methods, preserving biological signals better than other end-to-end methods. Overall, two-step strategies can better preserve and restore biological signals, while their fusion results are less satisfactory. Seurat performs well on the hPancreas dataset but falls short on other datasets, indicating its limited applicability to large-scale datasets. The MNN and SCALEX methods may be constrained by the chosen reference batch, limiting their performance. Reference-free methods Beaconet and Scanorama fail to outperform reference-based methods due to weaker constraints. Supplementary Figures 1-4 confirmed scConCRL's dual strengths of batch correction and biological conservation. scConCRL demonstrates superior performance in data integration in the molecular space. Moreover, an ablation study based on the task (Appendix F and Supplementary Table 1) proves the effectiveness of our design.

Further, we tested the ability of our scConCRL algorithm for online integration analysis, where the model is trained on reference datasets and applied to infer data from a new query dataset. For each experiment, a single-batch subset was used as the query dataset, with the remaining datasets as the reference. We compared scConCRL to the SCALEX algorithm using the scIB score. In this experiment, batches were simplified into "reference" and "query" categories. As shown in Table 3, scConCRL outperforms SCALEX in both batch correction and biological feature preservation, especially achieving a 3% improvement in batch correction metrics. To assess the impact of data integration on data quality, we performed experiments on simulated datasets and visualized the integrated datasets alongside clean datasets, including both query and reference data. We compared scConCRL with two widely used denoising methods, MAGIC and DeepImpute, as well as SCALEX. scConCRL outperforms them, with seamless integration of the predicted reference (red) and query (blue) data, and the integrated dataset closely approximating the clean dataset. This demonstrates scConCRL's superior ability to denoise and integrate both observed and unseen data, restoring true expression levels (Figure 2).

4.4 Perturbation prediction

When it comes to the condition domain, scConCRL excels at disentangling the condition information and predicting perturbation responses, a crucial intervention task for analyzing drug effects. A typical perturbation dataset includes control data and real perturbed data, and models are used to predict the perturbed data. We trained models separately for each cell type using the rest and inferred unseen perturbed data based on the control data. Using the PBMC dataset, we compared the performance of scConCRL with baselines including scPreGAN [Wei et al., 2022], scDisInFact [Zhang et al., 2024b], and scShift [Dong and Kluger, 2023]. We visualized the distribution of control data, real perturbed data, and predicted perturbation data in Figure 3a and Supplementary Figure 5. scConCRL shows the best performance, with the predicted data closely matching the real perturbed data and clearly separated from the control data. In contrast, scPreGAN produces two separate prediction clusters, distanced from real data. scDisInFact was directionally correct but struggled to match real perturbed data. scShift's predictions are more scattered, further limiting its performance. We also calculated R2 and adjusted R2 to quantify performance. The average results per cell type are shown in Figure 3b, and the overall similarity between predicted and actual perturbed data is presented in Figure 3c. These results further demonstrate scConCRL's superior performance.

4.5 Key gene detection

In the scConCRL framework, the matrix Axb captures the relationship between conditions and genes, identifying condition-related key genes (CKGs) [Zhang et al., 2024b].

To evaluate the performance of scConCRL versus scDisInFact, we assessed CKG detection accuracy on the PBMC dataset. First, we used the Wilcoxon rank-sum test to identify the differentially expressed genes (DEGs) between real perturbed and control data, and then the ground truth was subsequently established based on p-values. For scConCRL, the absolute values of Axb are the CKG score for each gene, reflecting the likelihood of a gene being a CKG. As shown in Table 4, scConCRL consistently outperforms scDisInFact, demonstrating superior performance in CKG detection and highlighting the enhanced interpretability. To further validate scConCRL, we examined the overlap between predicted DEGs and ground truth. As shown in Figure 4a, scConCRL achieves nearly 70% overlap, indicating high accuracy. Pathway analysis of the top 100 predicted DEGs (Figure 4b) revealed strong associations with immune system functions and signaling pathways involved in immune modulation and viral defense.

5 Conclusion

In this paper, we introduce scConCRL, conditional causal representation learning that tackles the challenges of integrating and analyzing heterogeneous single-cell data. By modeling the relationships between latent domain variables and gene variables through SEM, scConCRL captures both coarse- and fine-grained variables' interactions. Evaluations on multi-domain scRNA-seq datasets demonstrate its superior performance in several tasks including GRN inference, data integration, perturbation prediction, and key gene detection. scConCRL reveals complex dependencies, offering an interpretable solution for understanding cellular diversity and dynamic gene regulatory mechanisms.