scBiGNN: Bilevel Graph Representation Learning for Cell Type Classification from Single-cell RNA Sequencing Data

Problem Statement. In our work, two biological networks are constructed based on the pre-processed gene expression data $\tilde{\mathbf{X}}\in\mathbb{R}^{N\times T}$. One is the gene-gene interaction graph $\mathcal{G}^{\text{g}}=(\mathcal{V}^{\text{g}},\mathbf{A}^{\text{g}})$ where $\mathcal{V}^{\text{g}}$ is the set of $T$ genes $\{v_{1}^{\text{g}},\ldots,v_{T}^{\text{g}}\}$ and $\mathbf{A}^{\text{g}}\in\mathbb{R}^{T\times T}$ is the associated adjacency matrix representing the interactions between genes, the other is the cell-cell graph $\mathcal{G}^{\text{c}}=(\mathcal{V}^{\text{c}},\mathbf{A}^{\text{c}})$ where $\mathcal{V}^{\text{c}}=\{v_{1}^{\text{c}},\ldots,v_{N}^{\text{c}}\}$ indicates the set of $N$ cells in the scRNA-seq dataset and $\mathbf{A}^{\text{c}}\in\mathbb{R}^{N\times N}$ is its adjacency matrix encoding the relationships between cells. Given $\tilde{\mathbf{X}}$ and the labels $\mathbf{Y}^{L}\in\mathbb{R}^{|\mathcal{V}^{L}|\times C}$ for a subset of annotated cells $\mathcal{V}^{L}\subset\mathcal{V}^{\text{c}}$, the goal is to predict the labels $\mathbf{Y}^{U}\in\mathbb{R}^{|\mathcal{V}^{U}|\times C}$ for the the unlabeled cells $\mathcal{V}^{U}=\mathcal{V}^{\text{c}}\backslash\mathcal{V}^{L}$, where $C$ is the number of cell types.

Gene-level GNNs. The gene-level GNNs aim to predict the label of each cell $v_{i}^{\text{c}}$ using its own gene expression data $\tilde{\mathbf{x}}_{i}$ (the $i$-th row of $\tilde{\mathbf{X}}$) and the gene-gene interaction network. Formally, they model the following label distribution $q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})$ for each $v_{i}^{\text{c}}\in\mathcal{V}^{\text{c}}$ with parameters $\phi$:

	$\displaystyle\left\{\mathbf{f}_{i,t}\right\}_{t=1}^{T_{i}}$	$\displaystyle=\operatorname{GNN}^{\text{g}}\left(\tilde{\mathbf{x}}_{i},\mathbf{A}^{\text{g}}\right),$
	$\displaystyle\mathbf{h}_{i}$	$\displaystyle=\operatorname{Read-Out}\left(\left\{\mathbf{f}_{i,t}\right\}_{t=1}^{T_{i}}\right),$
	$\displaystyle q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})$	$\displaystyle=\operatorname{Cat}\left(\mathbf{y}_{i}\Big{|}\operatorname{Clf}^{\text{g}}\left(\mathbf{h}_{i}\right)\right),$		(1)

where $\operatorname{GNN}^{\text{g}}(\cdot)$ denotes the node representation learning function of the gene-level GNN, $\left\{\mathbf{f}_{i,t}\right\}_{t=1}^{T_{i}}$ is the set of gene representations learned by $\operatorname{GNN}^{\text{g}}(\cdot)$, $T_{i}$ is the number of genes that have non-zero expression values in $\tilde{\mathbf{x}}_{i}$, $\mathbf{h}_{i}$ is the representation of cell $v_{i}^{\text{c}}$, $\operatorname{Clf}^{\text{g}}(\cdot)$ denotes a classifier, and $\operatorname{Cat}(\cdot)$ stands for the categorical distribution. In this way, gene-gene structural information can be exploited from the local biological view of each individual cell to enhance cell representation learning.

Cell-level GNNs. The cell-level GNNs classify cell types by using the relationships among all the cells in the dataset. Specifically, the workflow can be characterized as below:

	$\displaystyle\left\{\mathbf{h}^{\prime}_{i}\right\}_{i=1}^{N}$	$\displaystyle=\operatorname{GNN}^{\text{c}}\left(\left\{\tilde{\mathbf{x}}_{i},\mathbf{h}_{i}\right\}_{i=1}^{N},\mathbf{A}^{\text{c}}\right),$
	$\displaystyle p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}})$	$\displaystyle=\operatorname{Cat}\left(\mathbf{y}_{i}\Big{|}\operatorname{Clf}^{\text{c}}\left(\mathbf{h}^{\prime}_{i}\right)\right),$		(2)

where $\operatorname{GNN}^{\text{c}}(\cdot)$ is the node representation learning function, $\left\{\mathbf{h}^{\prime}_{i}\right\}_{i=1}^{N}$ is the set of cell representations learned by $\operatorname{GNN}^{\text{c}}(\cdot)$, $\operatorname{Clf}^{\text{c}}(\cdot)$ is a classifier, and $\theta$ is the model parameters. By aggregating features from neighboring cells with similar characteristics, cell-level GNNs extract the cell-cell structural information from a global view of the scRNA-seq dataset to improve single-cell classification.

Both gene- and cell-level GNNs have proved effective in scRNA-seq analyses and achieved superior performance for automatic cell type identification. Therefore, we propose scBiGNN, a bilevel graph representation learning method that aims to integrate these two levels of structural information to have a more comprehensive view of scRNA-seq data.

Directly cascading the above gene- and cell-level GNNs and training them end-to-end would face the scalability issue when the size of the scRNA-seq dataset is large. Moreover, as these two types of GNNs are complementary, end-to-end learning does not allow them to interact and enhance each other. Therefore, inspired by several previous studies gmnn ; text ; emgcn , we instead employ the EM algorithm to train these two GNN modules alternately and make them gradually reinforce each other for more accurate cell type classification.

To be concrete, our proposed scBiGNN maximizes the evidence lower bound (ELBO) of the log-likelihood of the observed cell labels:

$\displaystyle\log p_{\theta}(\mathbf{Y}^{L}|\tilde{\mathbf{X}})\geq\mathbb{E}_{q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})}\left[\log p_{\theta}(\mathbf{Y}^{L},\mathbf{Y}^{U}|\tilde{\mathbf{X}})-\log q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})\right]\triangleq\mathcal{L}_{\text{ELBO}}(\theta,\phi;\mathbf{Y}^{L},\tilde{\mathbf{X}}),$

(3)

where $q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})$ can be arbitrary distribution over $\mathbf{Y}^{U}$ (s.t. $q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})>0$ if $p_{\theta}(\mathbf{Y}^{L},\mathbf{Y}^{U}|\tilde{\mathbf{X}})>0$), and the equality holds when $q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})$ equals to the true posterior distribution $p_{\theta}(\mathbf{Y}^{U}|\mathbf{Y}^{L},\tilde{\mathbf{X}})$. As can be seen, Eq. (3) formalizes the objective function with two distributions $q_{\phi}$ and $p_{\theta}$, which can be modeled by the abovementioned gene- and cell-level GNNs respectively. According to the EM framework, we alternately train $q_{\phi}$ in the E-step (with $p_{\theta}$ fixed) and $p_{\theta}$ in the M-step (with $q_{\phi}$ fixed) to maximize $\mathcal{L}_{\text{ELBO}}(\theta,\phi;\mathbf{Y}^{L},\tilde{\mathbf{X}})$. Thus, we can separately optimize these two GNN modules through several iterations, leading to better scalability and making them interact and enhance each other. In what follows, we introduce the structures of these two GNN modules and how they reinforce each other in the training procedure.

In the E-step, the cell-level GNN (i.e., $p_{\theta}$) is fixed and the gene-level GNN (i.e., $q_{\phi}$) is optimized to maximize $\mathcal{L}_{\text{ELBO}}(\theta,\phi;\mathbf{Y}^{L},\tilde{\mathbf{X}})$. In this section, we first present the detailed structure of our employed gene-level GNN, and then we introduce the optimization of $q_{\phi}$ in the E-step.

Structure of Gene-level GNN. The input of the gene-level GNN module includes two embeddings. The first is gene embedding, which assigns a unique and learnable vector $\mathbf{e}^{\text{gene}}_{j}\in\mathbb{R}^{d_{g}}$ to each gene $v_{j}^{\text{g}}$. The second is count embedding, which transforms every non-zero gene expression value $\tilde{X}_{ij}$ of cell $v_{i}^{\text{c}}$ to a $d_{g}$-dimensional vector through an MLP, i.e., $\mathbf{e}^{\text{count}}_{ij}=\operatorname{MLP}(\tilde{X}_{ij})\in\mathbb{R}^{d_{g}}$. The input embedding of gene $v_{j}^{\text{g}}$ in cell $v_{i}^{\text{c}}$ is then defined as $\mathbf{e}_{ij}=\mathbf{e}^{\text{gene}}_{j}+\mathbf{e}^{\text{count}}_{ij}$. By gathering and stacking all the input embeddings of genes that have non-zero expression values in cell $v_{i}^{\text{c}}$, we have the input embedding matrix $\mathbf{E}_{i}\in\mathbb{R}^{T_{i}\times d_{g}}$ for cell $v_{i}^{\text{c}}$.

Inspired by attention-based models transformer ; gat , the gene-level GNN of scBiGNN adaptively learns gene-gene interactions via the self-attention mechanism. Specifically, the feature aggregation layer of $\operatorname{GNN}^{\text{g}}(\cdot)$ is defined as

	$\displaystyle\operatorname{R}^{(l)}\left(\mathbf{A}^{\text{g}}\right)$	$\displaystyle=\operatorname{softmax}\left((\mathbf{F}_{i}^{(l-1)}\mathbf{W}_{\text{g}}^{(l)})(\mathbf{F}_{i}^{(l-1)}\mathbf{W}_{\text{g}}^{(l)})^{\top}\right),$
	$\displaystyle\mathbf{F}_{i}^{(l)}$	$\displaystyle=\operatorname{MLP}_{\text{g}}^{(l)}\left(\operatorname{R}^{(l)}\left(\mathbf{A}^{\text{g}}\right)\mathbf{F}_{i}^{(l-1)}\mathbf{W}_{\text{g}}^{(l)}\right),$		(4)

where $\operatorname{softmax}\left(\cdot\right)$ is the row-wise softmax operation, $\mathbf{W}_{\text{g}}^{(l)}$ is a learnable weight matrix, $\operatorname{MLP}_{\text{g}}^{(l)}(\cdot)$ is an MLP, $\mathbf{F}_{i}^{(l)}$ is the output feature matrix of the $l$-th layer, and $\mathbf{F}_{i}^{(0)}=\mathbf{E}_{i}$ for each cell $v_{i}^{\text{c}}$. Moreover, in each layer of $\operatorname{GNN}^{\text{g}}(\cdot)$, we can also employ the multi-head attention mechanism transformer , which performs several different attention functions of Eq. (3.3) and concatenates their outputs as the feature matrix. After $L$ layers, we can obtain the set of gene representations $\{\mathbf{f}_{i,t}^{j}\}_{t=1}^{T_{i}}$ in which $\mathbf{f}_{i,t}^{j}$ is the $t$-th row of $\mathbf{F}_{i}^{(L)}$ and the superscript $j$ indicates that it corresponds to the gene $v_{j}^{\text{g}}$. The read-out function in Eq. (3.1) can take the form of the weighted sum of all the gene representations to produce the cell representation for $v_{i}^{\text{c}}$:

$\displaystyle\mathbf{h}_{i}=\operatorname{Read-Out}\left(\{\mathbf{f}_{i,t}^{j}\}_{t=1}^{T_{i}}\right)=\sum\nolimits_{j}\alpha_{j}\mathbf{f}_{i,t}^{j},$

(5)

where $\alpha_{j}$ can be a learnable weight for gene $v_{j}^{\text{g}}$ ($j=1,\dots,T$) or set as $1/T_{i}$ which performs average pooling. Then, $q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})$ can be obtained by feeding $\mathbf{h}_{i}$ to an MLP classifier, as shown in Eq. (3.1).

Optimization of Gene-level GNN. Maximizing $\mathcal{L}_{\text{ELBO}}(\theta,\phi;\mathbf{Y}^{L},\tilde{\mathbf{X}})$ in the E-step is equivalent to making $q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})$ approximate the true posterior distribution $p_{\theta}(\mathbf{Y}^{U}|\mathbf{Y}^{L},\tilde{\mathbf{X}})$. According to the EM framework, the standard operation is to minimize the Kullback-Leibler (KL) divergence $\operatorname{KL}(q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})\|p_{\theta}(\mathbf{Y}^{U}|\mathbf{Y}^{L},\tilde{\mathbf{X}}))$. However, directly optimizing the KL divergence is nontrivial, as it depends on the entropy of $q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})$, which is hard to deal with text . Moreover, for $v_{i}^{\text{c}}\in\mathcal{V}^{U}$, when $p_{\theta}(\mathbf{y}_{i}|\mathbf{Y}^{L},\tilde{\mathbf{X}})$ is close to the one-hot categorical distribution, it would face the instability issue caused by $\log y$ with $y$ approaching zero. Therefore, following several EM-based optimization methods gmnn ; text ; emgcn , we instead treat the fixed $p_{\theta}(\mathbf{Y}^{U}|\mathbf{Y}^{L},\tilde{\mathbf{X}})$ as the target and minimize the reverse KL divergence $\operatorname{KL}(p_{\theta}(\mathbf{Y}^{U}|\mathbf{Y}^{L},\tilde{\mathbf{X}})\|q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}}))$ to make $q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})$ approximate the target. Assuming that the labels of different cells are independent, which is reasonable as the organization and function of each cell are not determined by other cells, we can derive the following loss function for unlabeled cells:

$\displaystyle\mathcal{L}_{E}^{U}=-\mathbb{E}_{p_{\theta}(\mathbf{Y}^{U}|\mathbf{Y}^{L},\tilde{\mathbf{X}})}\left[\log q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})\right]=-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{V}^{U}}\mathbb{E}_{p_{\theta}(\mathbf{y}_{i}|\mathbf{Y}^{L},\tilde{\mathbf{X}})}\left[\log q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})\right],$

(6)

where $p_{\theta}(\mathbf{y}_{i}|\mathbf{Y}^{L},\tilde{\mathbf{X}})=p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}})$ is the categorical distribution predicted by the cell-level GNN for $v_{i}^{\text{c}}\in\mathcal{V}^{U}$ in the previous M-step. That is, the pseudo-labels predicted by the cell-level GNN are used for training the gene-level GNN. Additionally, the gene-level GNN can also be trained by the ground-truth labels of cells in $\mathcal{V}^{L}$. Therefore, we have the loss function for labeled cells:

$\displaystyle\mathcal{L}_{E}^{L}=-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{V}^{L}}\log q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i}),$

(7)

where $\mathbf{y}_{i}$ is the ground-truth label for cell $v_{i}^{\text{c}}\in\mathcal{V}^{L}$. Combining Eqs. (6) and (7), the overall objective in the E-step for optimizing $\phi$ is to minimize $\mathcal{L}_{E}=\mathcal{L}_{E}^{U}+\beta\mathcal{L}_{E}^{L}$ with $\beta$ being a balancing hyper-parameter (we set $\beta=1$ in this work), which can be solved via stochastic gradient descent (SGD). To be more specific, in each update step of SGD, we sample a batch of labeled cells $\mathcal{B}^{L}$ and a batch of unlabeled cells $\mathcal{B}^{U}$ to perform gradient descent $\phi\leftarrow\phi-\gamma_{\text{g}}\nabla_{\phi}\tilde{\mathcal{L}}_{E}$, where $\tilde{\mathcal{L}}_{E}=-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{B}^{U}}\log q_{\phi}(\hat{\mathbf{y}}_{i}|\tilde{\mathbf{x}}_{i})-\beta\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{B}^{L}}\log q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})$ with $\hat{\mathbf{y}}_{i}$ sampled by $\hat{\mathbf{y}}_{i}\sim p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}})$ for unlabeled cells and $\gamma_{\text{g}}$ is the learning rate.

In the M-step, the gene-level GNN (i.e., $q_{\phi}$) is fixed and the cell-level GNN (i.e., $p_{\theta}$) is updated to maximize $\mathcal{L}_{\text{ELBO}}(\theta,\phi;\mathbf{Y}^{L},\tilde{\mathbf{X}})$. Also, before we introduce the optimization of $p_{\theta}$, we present the detailed workflow of our cell-level GNN.

Structure of Cell-level GNN. Existing cell-level GNNs predetermine the cell-cell graph based on $\tilde{\mathbf{X}}$ improving . By contrast, we construct the cell-cell graph by performing $k$-nearest neighbors analysis on the cell representations $\{\mathbf{h}_{i}\}_{i=1}^{N}$ generated by the gene-level GNN in each EM iteration, i.e., $\mathbf{A}^{\text{c}}=\operatorname{kNN-graph}(\{\mathbf{h}_{i}\}_{i=1}^{N})$, since $\{\mathbf{h}_{i}\}_{i=1}^{N}$ should have better clustering characteristics than $\tilde{\mathbf{X}}$ for identifying cell types and its quality is enhanced during the optimization of $q_{\phi}$.

Given the gene expression data $\tilde{\mathbf{x}}_{i}$ and the cell representation $\mathbf{h}_{i}$, the input feature for each cell $v_{i}^{\text{c}}$ is defined as $\mathbf{z}_{i}=\operatorname{Concat}(\mathbf{h}_{i},\operatorname{MLP}(\tilde{\mathbf{x}}_{i}))\in\mathbb{R}^{d_{c}}$, where $\operatorname{Concat}(\cdot)$ denotes the concatenation operation for the input vectors. By stacking the input features of all the cells, we have the input feature matrix $\mathbf{Z}\in\mathbb{R}^{N\times d_{c}}$. And the structure of $\operatorname{GNN}^{\text{c}}(\cdot)$ can be modeled as:

	$\displaystyle\operatorname{R}^{(l)}\left(\mathbf{A}^{\text{c}}\right)=\mathbf{D}^{-1}\mathbf{A}^{\text{c}},\quad\mathbf{D}=\operatorname{diag}\left(\mathbf{A}^{\text{c}}\mathbf{1}_{N}\right),$
	$\displaystyle\mathbf{Z}^{(l)}=\operatorname{MLP}_{\text{c}}^{(l)}\left(\operatorname{R}^{(l)}\left(\mathbf{A}^{\text{c}}\right)\mathbf{Z}^{(l-1)}\mathbf{W}_{\text{c}}^{(l)}\right),$		(8)

where $\mathbf{1}_{N}$ is the all-ones vector of length $N$, $\operatorname{diag}(\cdot)$ returns a diagonal matrix with the input vector as the diagonal, $\mathbf{W}_{\text{c}}^{(l)}$ is a learnable weight matrix, $\operatorname{MLP}_{\text{c}}^{(l)}(\cdot)$ is an MLP, $\mathbf{Z}^{(l)}$ is the output feature matrix of the $l$-th layer, and $\mathbf{Z}^{(0)}=\mathbf{Z}$. As for $\mathbf{h}^{\prime}_{i}$ that is fed to an MLP classifier in Eq. (3.1), we find that setting $\mathbf{h}^{\prime}_{i}=\operatorname{Concat}(\mathbf{z}_{i}^{(0)},\dots,\mathbf{z}_{i}^{(L)})$ with $\mathbf{z}_{i}^{(l)}$ being the $i$-th row of $\mathbf{Z}^{(l)}$ practically achieves better classification performance for $p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}})$, which is similar to the jumping knowledge networks jknet .

Optimization of Cell-level GNN. Maximizing $\mathcal{L}_{\text{ELBO}}(\theta,\phi;\mathbf{Y}^{L},\tilde{\mathbf{X}})$ in the M-step is equivalent to minimizing the following loss function:

	$\displaystyle\mathcal{L}_{M}$	$\displaystyle=-\mathbb{E}_{q_{\phi}(\mathbf{Y}^{U}|\tilde{\mathbf{X}})}\left[\log p_{\theta}(\mathbf{Y}^{L},\mathbf{Y}^{U}|\tilde{\mathbf{X}})\right]$
		$\displaystyle=-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{V}^{U}}\mathbb{E}_{q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})}\left[\log p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}})\right]-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{V}^{L}}\log p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}}).$		(9)

Here, the first term is the loss for unlabeled cells, which trains $p_{\theta}$ using the pseudo-labels generated by the gene-level GNN in the previous E-step. The second term is the loss for labeled cells, which is the standard supervised learning with ground-truth labels. Minimizing Eq. (3.4) w.r.t. $p_{\theta}$ can also be solved via SGD. Concretely, in each update step of SGD, we sample $\hat{\mathbf{y}}_{i}\sim q_{\phi}(\mathbf{y}_{i}|\tilde{\mathbf{x}}_{i})$ for unlabeled cells and perform gradient descent $\theta\leftarrow\theta-\gamma_{\text{c}}\nabla_{\theta}\tilde{\mathcal{L}}_{M}$, in which $\tilde{\mathcal{L}}_{M}=-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{V}^{U}}\log p_{\theta}(\hat{\mathbf{y}}_{i}|\tilde{\mathbf{X}})-\sum\nolimits_{v_{i}^{\text{c}}\in\mathcal{V}^{L}}\log p_{\theta}(\mathbf{y}_{i}|\tilde{\mathbf{X}})$ and $\gamma_{\text{c}}$ is the learning rate.

To optimize the overall scBiGNN framework, we first pre-train the gene-level GNN with labeled cells to obtain the initial $q_{\phi}$. Then, the EM algorithm alternately optimizes $p_{\theta}$ (i.e., the cell-level GNN) and $q_{\phi}$ (i.e., the gene-level GNN) to reinforce each other. In the M-step, the cell representations generated by the gene-level GNN are used to construct the cell-cell graph, and the pseudo-labels predicted by $q_{\phi}$ together with the given cell labels are utilized to train the cell-level GNN. In the E-step, the pseudo-labels produced by $p_{\theta}$ and the ground-truth cell labels are used to train the gene-level GNN, in which the gene-gene interactions are adaptively learned. After several EM iterations, the performance of both cell- and gene-level GNNs is enhanced in cell type classification.

In this work, we conduct cell type classification experiments on five benchmark datasets to evaluate the performance of our scBiGNN, i.e., Zheng68K, Zhengsorted, BaronHuman, BaronMouse and AMB, which are widely used in many published papers for single-cell classification. These datasets can be directly downloaded from Zenodo (https://doi.org/10.5281/zenodo.3357167), and the detailed information is depicted in Table 1.

As for the gene-level GNN module, we use a one-layer model with four attention heads for the BaronMouse dataset, while a two-layer model with a single head in each layer is employed for the other datasets. Additionally, for Zheng68K and Zhengsorted datasets, the read-out function is set as the simple average pooling operation, while for the others we find that using learnable $\alpha_{j}$ in Eq. (5) performs better. As for the cell-level GNN module, a three-layer model is employed for cell graph representation learning. The feature dimension of each layer’s output is set as 32 for both GNN modules. All the MLPs used in scBiGNN have one hidden layer with 32 neurons and ReLU activation. The maximum number of EM iterations is set as 3, which is commonly enough for scBiGNN to converge, as shown in Figure 1.

We compare with a variety of baseline models to demonstrate the superior performance of scBiGNN. Specifically, following the previous study scgraph , eight methods are considered: support vector machine (SVM), linear discriminant analysis (LDA), nearest mean classifier (NMC), random forest (RF), scID scid , CHETAH chetah , SingleR singler and ACTINN actinn . In addition, three recent works that use gene-level GNNs are included: scGraph scgraph , sigGCN siggcn and HNNVAT adversarial . We also compare with the cell-level GNN, termed GCN-C, in which the cell-cell graph is constructed by principal component analysis and $k$-nearest neighbors analysis on the raw gene expression data improving .

We benchmark scBiGNN against all the baselines by performing five-fold cross validation on each dataset. The results of cell type classification accuracy are summarized in Table 2, in which the best performance in each column is highlighted in bold. As can be seen, among all the baselines, graph-based deep learning methods are the most accurate classifiers, verifying that the biological structural information is beneficial to identifying cell types from scRNA-seq data. Out of all the graph-based models, our scBiGNN consistently achieves the highest accuracy on all the benchmark datasets. Noticeably, it outperforms all the baselines with more than 2.5% accuracy improvement on the largest scRNA-seq dataset (i.e., Zheng68K). Moreover, from Table 2 and Figure 1 we can observe that $p_{\theta}$ always performs better than $q_{\phi}$. The reason is that $p_{\theta}$ explicitly integrates the bilevel graph representations of $\mathbf{h}_{i}$ that summarizes gene-level structural information and $\mathbf{A}^{\text{c}}$ that contains cell-level structural information. Overall, these results validate that the gene- and cell-level relationships are effectively mined by $p_{\theta}$ of scBiGNN to facilitate scRNA-seq classification.

Gene-gene Graph. Our scBiGNN adaptively learns gene-gene interactions by the self-attention function $\operatorname{R}^{(l)}\left(\mathbf{A}^{\text{g}}\right)$ in Eq. (3.3). We integrate all the attention matrices into a gene-gene interaction matrix $\tilde{\mathbf{A}}$ by taking the average across all the corresponding elements in all attention matrices for every gene-gene pair, where each element $\tilde{A}_{ij}$ measures the attention of gene $v_{i}^{\text{g}}$ paid to gene $v_{j}^{\text{g}}$. Then, we can obtain the importance score $s_{j}$ for each gene $v_{j}^{\text{g}}$, which is defined as $s_{j}=\sum_{i}\tilde{A}_{ij}$. We sort the importance scores and get five lists of top 50 important genes from five-fold cross validation. Figure 2(A) shows that most genes appear more than once in the five lists, indicating that scBiGNN is trustworthy and effective in learning consistent important genes across different training procedures.

Cell-cell Graph. The cell-cell graph in scBiGNN is constructed as the $k$-nearest neighbor ($k$NN) graph based on the cell representations learned from $q_{\phi}$. In Figure 2(B), we can observe that our model is insensitive to $k$. A small value of $k$ (e.g., $k=5)$ is enough for $p_{\theta}$ to achieve good performance, verifying the high quality of cell representations generated by $q_{\phi}$ in finding close neighborhoods. In Figure 2(C), we measure the homophily ratio (i.e., the proportion of intra-class edges in all the edges) of cell-cell graphs constructed by scBiGNN and the raw data $\tilde{\mathbf{X}}$ improving respectively. As can be seen, cell-cell graphs built by scBiGNN link more cells of the same type, which is critical to learning separable cell representations across different classes to ease GNN-based node classification emgcn .