Bipartite Graph Attention-based Clustering for Large-scale scRNA-seq Data

scRNA-seq clustering is a critical task for analyzing single-cell RNA sequencing (scRNA-seq) data, as it groups cells with similar gene expression profiles. Transformers, as powerful foundational models, have been applied to scRNA-seq clustering. Their self-attention mechanism automatically assigns higher attention weights to cells within the same cluster, enhancing the distinction between clusters. Existing methods for scRNA-seq clustering, such as graph transformer-based models, treat each cell as a token in a sequence. Their computational and space complexities are $\mathcal{O}(n^2)$ with respect to the number of cells, limiting their applicability to large-scale scRNA-seq datasets. To address this challenge, we propose a Bipartite Graph Transformer-based clustering model (BGFormer) for scRNA-seq data. We introduce a set of learnable anchor tokens as shared reference points to represent the entire dataset. A bipartite graph attention mechanism is introduced to learn the similarity between cells and anchor tokens, bringing cells of the same class closer together in the embedding space. BGFormer achieves linear computational complexity with respect to the number of cells, making it scalable to large datasets. Experimental results on multiple large-scale scRNA-seq datasets demonstrate the effectiveness and scalability of BGFormer.

Single-cell RNA sequencing measures gene expression in cells, with clustering as a key step in analyzing these data to identify different cell types. However, clustering on datasets with a large number of cells is very challenging because it typically requires considering relationships between every pair of cells, leading to extremely high computational cost. To address this, we propose a fast clustering method that introduces a set of learnable anchors to represent all cells. Instead of directly computing cell-to-cell relationships, we compute relationships between cells and these anchors. Our approach improves clustering performance while significantly reducing both computational cost and runtime on large-scale datasets.