Weighted Diversified Sampling for Efficient Data-Driven Single-Cell Gene-Gene Interaction Discovery

Single-cell transcriptomic is a technology that profiles gene expression at the individual cell level. The profiled results, namely single-cell transcriptomic data, provide a unique landscape of gene expressions. In contrast to traditional bulk RNA-seq analysis, single-cell transcriptomic data allows for cell-level sequencing, which captures the variability between individual cells (Ata et al., 2020). Leveraging this high-resolution data allows scientists to gain insights into developmental processes, disease mechanisms, and cellular responses to environmental changes.

The single-cell transcriptomic data can be formulated as a set of high-dimensional and sparse feature vectors. We denote a single-cell transcriptomic dataset at $X$, where each cell $x\in X$ is a sparse vector with dimensionality $V\in\mathbb{N}_{+}$. Here $V$ represents the total number of genes we can observe in $X$. Since cell $x\in\mathbb{R}^{V}$ is a sparse vector, we can represent $x$ as a set $\{(i_{1},v_{1}),(i_{2},v_{2}),\cdots,(i_{k},v_{k})\}$. In this set, every tuple $(i,v)$ represents the expression of gene $i\in[V]$ with expression level $v\in\mathbb{R}$. Besides we can also denote cell $x$ as $[x_{1},x_{2},\cdots,x_{V}]$, where most of the $x_{i}$s are zeros.

In this data formulation, single-cell transcriptomic data for each cell is represented as a set of gene expressions, with different cells expressing varying genes. Additionally, even when two cells express the same gene, their expression levels may differ. Our research objective is to identify gene-gene interactions within the vocabulary $V$ that drive complex biological processes and disease mechanisms.

Here, we propose our Transformer architecture, CelluFormer, to learn gene-gene interactions within single-cell transcriptomic data. Based on the set formulation of single-cell transcriptomic data, we believe that the order of genes is arbitrary and biologically meaningless. Similar to scGPT (Cui et al., 2024), and scFoundation (Hao et al., 2024), our method adopts a permutation-invariant design. We define our permutation-invariant condition as follows.

We see Condition 2.1 as a fundamental difference between the proposed Transformer and the sequence Transformers (Vaswani et al., 2017) widely used in natural language processing. For sequence Transformers, we have to ingest sequential masks during the training to ensure that the current token does not interact with the future token. Additionally, during the inference, the sequence Transformer should perform a step-by-step generation for each token. As a result, the sequence Transformer does not satisfy Condition 2.1. Moreover, the difference between CelluFormer and a vision Transformer (Dosovitskiy et al., 2020) is that the vision Transformer has a fixed sequence length for every input data sample. However, the number of genes expressed in each cell can vary a lot. For example, according to Figure 1, the number of genes expressed in a single cell can be up to 12,000 or more. Thus, we utilize a padding mask for the classification downstream task. Additional details regarding the implementation of CelluFormer are provided in Appendix C.1.

We observe that there is a significant performance difference between Transformer models if we feed them with different styles of single-cell transcriptomic data. It is known that cells can be categorized into different types based on their functionality. For instance, neuronal cells represent the cell types that fire electric signals called action potentials across a neural network (Levitan and Kaczmarek, 2015). Our study suggests that Transformers should be trained on single-cell transcriptomic data from various cell types to achieve better performance. We showcase an example in Table 1. We train a Transformer model to classify whether a cell is an Alzheimer’s disease-infected cell or not. According to our study, CelluFormer proposed in Section 2.2 trained on neuronal cells outperforms traditional multilayer perceptron (MLP) with downstream training on a single cell type. However, we do not see this gap when we perform training of CelluFormer on a single cell type. As a result, we see that the Transformers generally prefer massive exposure to the single-cell transcriptomic data.

In this paper, we would like to accomplish the following objective.

We see the self-attention mechanism of Transformers on a cell’s set style gene expressions as a pathway to model gene-gene interactions. CelluFormer takes a cell $x$’s gene expressions and produces an attention map $A_{i,j}\in\mathbb{R}^{m\times m}$ at encoder block $i$ and attention head $j$. Here $m$ represents the number of genes expressed in cell $x$. Since Transformer architecture uses the Softmax function to produce $A_{i,j}$, we can view the $p$th row of $A_{i,j}$ as the interaction between gene $p$ and all other genes in $x$. As a result, an attention map is a natural indicator of gene-gene interactions. Moreover, if we have a perfect Transformer that takes a cell $x$ gene expressions and correctly predicts if it is infected by a disease, we view the attention map of this cell as an indicator of disease-oriented gene-gene interactions. Following this path, we propose a gene-gene interaction modeling approach as illustrated in Figure 2. For each cell $x$, we represent it as a set and generate a bag of embeddings from the gene embedding table. Next, we use the expression levels of each gene as a scaling factor for each gene’s embedding. Next, we take the average attention maps of all layers and all heads to obtain a gene-gene interaction map in this cell.

In Objective 2.2, we would like to see not only the gene-gene interactions just for cell $x$ but also the statistical evidence of how two genes interact in the dataset $X$. As a result, we propose to accumulate multiple cells’ averaged attention maps as illustrated in Figure 3. For $X$, we initialize $Z_{0}\in 0^{V\times V}$ matrix as the overall attention map before aggregation and $M_{0}\in 0^{V\times V}$ as the overall frequency dictionary before aggregation. Next, for each cell $x$ in the dataset, we remove its diagonal value in its averaged attention map as it represents self-interaction. Next, we perform scatter addition operations that merge $x$’s averaged attention map back to $Z_{0}$. We let $Z_{ij}$ add the interaction value of gene $v_{i}$ and $v_{j}$ in the average attention map of cell $x$ obtained in the Transformer model. Simultaneously, to eliminate the dataset bias of expressed genes, we count the number of appearances for each gene pair in the dataset. Once again, we perform scatter addition to record the counts back to $M_{0}$. This is done by updating $M_{0}$ through scatter addition, where $M_{ij}=M_{ij}+1$ for every occurrence of the gene pair $(v_{i},v_{j})$ in the dataset. Finally, we rank the off-diagonal values in $Z$ where $Z_{ij}\leftarrow\frac{Z_{ij}}{M_{ij}}$ to retrieve the top gene-gene interaction.

As illustrated in Section 2.4, once we have a pre-trained CelluFormer that can successfully predict whether a cell is infected by a disease or not with its gene expressions, we can perform gene-gene interaction discovery by passing massive cells into this model and get the accumulated attention map as Figure 3. However, this process requires data-intensive computation. For every cell in the dataset, we first need to compute the average attention map as illustrated in Figure 2. Next, we perform aggregations as shown in Figure 3. It is known that CelluFormer uses plenty of trainable parameters to achieve good disease infection classification performance. As a result, the computation complexity for generating a cell’s averaged attention map is expensive. Moreover, since the attention map for cell $x$ is $m\times m$, where $m$ is the number of genes expressed in $x$. According to Figure 1, we see that $m$ can be 12,000 or more. These giant attention maps consume the limited high bandwidth memory (HBM) in the graphics processing unit. Therefore, we have to perform batch-wise computation on a massive cell dataset for computing gene-gene interactions. Moreover, given the scale of the dataset, any sampling algorithm with a runtime that grows exponentially with the dataset size is impractical.

In this work, we would like to address this data-efficiency challenge by raising and asking the following research question: Can we find a representative and small subset from the large cell dataset and still perform successful gene-gene interaction discovery? Moreover, we would like the procedure for finding this small subset as efficient as possible.

We would like to answer this question by proposing a diversity score of a cell in the dataset. To begin with, we introduce the $\mathsf{Min}\text{-}\mathsf{Max}$ similarity between two cell’s gene expressions.

According to the definition, $\mathsf{Min}\text{-}\mathsf{Max}(x,y)\in[0,1]$. Higher $\mathsf{Min}\text{-}\mathsf{Max}$ means that two cell’s gene expressions are closer to each other. $\mathsf{Min}\text{-}\mathsf{Max}$ is widely viewed as a kernel (Li, 2015b; Li et al., 2021; Li and Li, 2021) in statistical machine learning. In this paper, we would like to define a kernel density on top of the $\mathsf{Min}\text{-}\mathsf{Max}$ similarity.

We view $\mathsf{Min}\text{-}\mathsf{Max}(q)$ density as an indicator of how diverse $q$ is in $X$. Smaller $\mathsf{Min}\text{-}\mathsf{Max}(q)$ means that all other $x\in X$ may be less similar to $q$, making $q$ a unique cell. On the other hand, higher $\mathsf{Min}\text{-}\mathsf{Max}(q)$ means that $X$ has some cells that have similar gene expressions with $q$, making $q$ less unique. However, to compute $\mathsf{Min}\text{-}\mathsf{Max}(q)$ for every $q\in X$ following Definition 3.2, we have to compute all pairwise $\mathsf{Min}\text{-}\mathsf{Max}(x,y)$ for any $x,y\in X$, which results in an unaffordable $O(n^{2}\mathbb{NNZ}(X))$ time complexity, where $n$ is the size of $X$ and $\mathbb{NNZ}(X)$ is the maximum possible number of genes expressed in a cell $x\in X$. To reduce this $n^{2}$ computation, we propose a randomized algorithm that takes advantage of 0-bit consistent weighted sampling (CWS) (Li, 2015a) hash functions.

Here the $o(1)$ is a minor additive term with complex form. For simplicity, we refer the readers to (Li et al., 2021), Theorem 4.4 for more details.

This work presents an efficient randomized algorithm that estimates $\mathsf{Min}\text{-}\mathsf{Max}$ density $\mathcal{K}(q)$ (see Definition 3.2) for every $q\in X$. As showcased in Algorithm 1, we initialize an array $A$ with all values as zeros. Next, we conduct a pass over $X$. In this pass, for every $x\in X$, we compute its hash values after $R$ independent hash functions. Next, we increment $A_{r,h_{r}(x)}$ with $1$. After this pass, we take another pass at the dataset, for every $x\in X$, we take an average over the $A_{r,h_{r}(x)}$ and build a density score $w_{x}$. We would like to highlight that Algorithm 1 requires only two linear scans of the dataset. The time complexity for this algorithm is $O(n\mathbb{NNZ}(X))$, which is linear to the dataset. Moreover, we show that Algorithm 1 produces an estimator to $\mathsf{Min}\text{-}\mathsf{Max}$ density.

We provide the proof of Theorem 3.4 in the supplementary materials.

We propose to use the inverse form of $\mathsf{Min}\text{-}\mathsf{Max}$ density in Definition 3.2 as a score for diversity. We define it as normalized inverse $\mathsf{Min}\text{-}\mathsf{Max}$ density as below.

We view the IMD $\mathcal{I}(q)\in[0,1]$ as a monotonic increasing function for the diversity of $q$. Higher $\mathcal{I}(q)$ means that all other $x\in X$ may be less similar to $q$, making $q$ a unique cell. Moreover, IMD can be directly used as a sample probability to generate a representative subset of $X$ for Objective 2.2. Given $X$, we perform sampling without replacement to generate a subset $X_{\mathsf{sub}}\subset X$, where $x\in X$ has the sampling probability $\mathcal{I}(x)$. The advantages of sampling with IMD (see Definition 3.5) can be summarized as follows.

• •

  The IMD $\mathcal{I}(q)$ can be an effective indicator for how diverse $q$ is in dataset $X$.

• •

  Computing IMD is an efficient one-shot preprocessing process with just two linear scans of $X$ with time complexity $O(n\mathbb{NNZ}(X))$, where $n$ and $\mathbb{NNZ}(X)$ is defined in Section 3.2.

• •

  The memory complexity of computing IMD is $O(RB)$, which can be viewed as constant since it is independent of $n$ and $\mathbb{NNZ}(X)$.

In the following section, we would like to evaluate the empirical performance of IMD in selecting a representative subset out of a massive single-cell transcriptomic dataset while still maintaining effective performance in data-driven gene-gene interaction discovery powered by CelluFormer.

Dataset: For the training dataset, we employ the Seattle Alzheimer’s Disease Brain Cell Atlas (SEA-AD) (Gabitto et al., 2023), which includes single nucleus RNA sequencing data of 36,601 genes (as 36,601 features) from 84 senior brain donors exhibiting varying degrees of Alzheimer’s Disease (AD) neuropathological changes. By providing extensive cellular and genetic data, SEA-AD enables in-depth exploration of the cellular heterogeneity and gene expression profiles associated with AD. To facilitate a comparative analysis between AD-affected and non-AD brains, we select cells from 42 donors classified within the high-AD category and 9 donors from the non-AD category, based on their neuropathological profiles. This selection criterion ensures a robust comparison, aiding in the identification of gene-gene interactions linked to AD progression (Gabitto et al., 2023). The dataset is comprehensively annotated, covering 1,240,908 cells across 24 distinct cell types. We selected 18 neuronal cell types as our final training dataset since we believe neuronal cells are more likely to reveal explainable gene-gene interactions that are related to Alzheimer’s Disease compared to non-neuronal cells. To better detect expression relationships among genes, we apply the Seurat Transformation Function (Stuart et al., 2019) to eliminate the problem of sequence depth difference.

Our baseline includes NID (Tsang et al., 2017), a traditional feature interpretation technique that extracts learned interactions from trained MLPs. NID identifies interacting features by detecting strongly weighted connections to a standard hidden unit in MLPs after training. We evaluated our CelluFormer model against the MLP model, with performance results presented in Table 1.

Additionally, to comprehensively evaluate RQ1, we utilized two existing single-cell large foundation models to assess our algorithm. Specifically, we fine-tuned two foundation models, scFoundation (Hao et al., 2024) and scGPT (Cui et al., 2023), to classify whether a cell is AD or non-AD (performance results are provided in Table 4). We then applied our gene-gene interaction discovery pipeline using the attention maps of these foundation models.

In the sampling experiments, we compare WDS with uniform sampling since none of them requires preprocessing time exponential to the dataset size.

To evaluate the effectiveness of our Transformer-based framework for gene-gene interaction discovery, we performed feature selection across seven different cell types used as inference datasets. Additionally, we used a dataset encompassing all neuronal cell types to assess the overall performance of various models. As shown in Table 2, Transformer-based methods, including CelluFormer, scGPT and scFoundation, significantly outperformed other baselines.

This result indicates that our proposed Transformer-based framework is more effective and stable at extracting general and global gene-gene interaction information. In addition, the foundation models, scGPT and scFoundation, achieved comparable performances with other baselines across some of the datasets. We attribute this outcome to two main factors. Overfitting to Pretrained Knowledge: A foundation model, particularly a large one, might have learned very general or specific knowledge during its pretraining phase. When fine-tuning for a specific task, the model might overfit the preexisting knowledge, leading to poor generalization of the new task data. Mismatch Between Pretraining and Fine-Tuning Data: If the data distribution for fine-tuning is significantly different from the data on which the foundation model was trained, the model might struggle to adapt, resulting in worse performance. A model trained from scratch on the specific task data may perform better as it directly optimizes for that data distribution.

We addressed these questions by comparing our weighted diversified sampling (WDS) method with uniform sampling across various sample sizes, ranging from 1% to 10% of the original dataset. We generated data subsets for each cell type using WDS and uniform sampling. We then applied our Transformer-based framework for feature selection at each sample size. Since CelluFormer consistently outperformed other baselines, we selected it as our base model. We repeated Each experiment five times and recorded the NES scores as the results. To evaluate the sampling methods, we calculated the average NES score across the five experiments. We also computed the Mean Square Error (MSE) between the NES scores from the sampling experiments and the ground truth derived from the entire dataset, as shown in Table 2. The evaluation results are presented in Table 3. We note that WDS consistently produced higher NES scores compared to uniform sampling. As the sample size increased, the NES scores from uniform sampling began to converge with the ground truth. In contrast, the NES scores from WDS consistently remained close to the ground truth, even at smaller sample sizes. The result indicates that while WDS offers a significant advantage in small samples by enabling the Transformer to capture a broader range of genetic interactions, its benefits diminish as more data becomes available. Moreover, we find that for some cell types, smaller samples of data outperformed larger samples of data on NES. This suggests that: (1) single-cell transcriptomic data may contain noises that affect gene-gene interaction discovery, and, (2) some complex gene-gene interaction patterns in single-cell transcriptomic data cannot be interpreted directly through attention maps. We also provide a detailed study on the choice of parameter $R$ in Algorithm 1 in Appendix D.2.