scKGBERT

scKGBERT: A Knowledge-Enhanced Foundation Model for Single-Cell Transcriptomics

This repository is the official implementation of scKGBERT.

<hr>

Brief introduction

scKGBERT is the first single-cell foundation model to integrate genetic knowledge graph pre-training and RNA sequence pre-training into a unified framework. The model utilizes a Gaussian attention mechanism within the transformer architecture to align key knowledge with RNA sequence representations, enhancing its interpretability in learning gene and cell profiles. This approach allows scKGBERT to understand and learn both the biological semantics and genetic contexts of genes, capturing extensive interactions between independent genes and dynamic network information within cells.

<div align=center><img src="./fig/overall.png" style="zoom:50%;" /> </div>

Model

Genetic Knowledge Pre-Training:

• We apply an unsupervised graph pretraining method on a genetic knowledge graph to capture foundational gene knowledge representation.
• The STRING database is used as the external knowledge graph, providing critical protein-protein interaction data, supporting pathway analysis, and aiding in identifying disease biomarkers.

scKGBERT Pre-Training:

• scKGBERT takes a significant volume of unsupervised single-cell RNA sequences as input.
• During the pre-training phase, scKGBERT employs multiple pre-training units, each consisting of an RNA sequence encoder module (S-encoder), a knowledge graph encoder module (K-encoder), and a Gaussian cross-attention layer. Using a masking strategy similar to the BERT model, scKGBERT masks 15% of the input single-cell transcriptomic data and uses the unmasked data to predict the hidden positions.

Fine-Tuning in Downstream Tasks:

• After unsupervised pre-training on a large volume of unlabeled single-cell transcriptomic sequencing data, scKGBERT is fine-tuned for various gene and cell-level tasks.
• This fine-tuning enables the successful application of the pre-trained model to complex downstream tasks, enhancing the performance of large pre-trained models by incorporating biological knowledge.

Requirements

To run our code, please install dependency packages.

huggingface_hub           0.15.1           
networkx                  2.8.4            
numpy                     1.24.3         
python                    3.8.17             
scikit-learn              1.2.2           
tokenizers                0.13.2          
torch                     1.13.0+cu117           
transformers              4.31.0               

Overview

This project mainly contains the following parts.

├── models                        # pre-trained scKGBERT model.
├── utils                         # Utility classes folder
│   ├── KG_TransR.py                    # TransR training code.
│   ├── pretrainer.py                   # The Pretrainer code is used to train a model on large-scale datasets to obtain general feature representations.
│   ├── scKGBERT_model.py                # scKGBERT model.
│   ├── tokenizer.py                    # The Tokenizer class is used to segment data into units and convert them into numerical forms that a model can process.
│   └── utils_pkg.py                    # Other utility classes used.
├── pretrainscKGBERT.py             #  The scKGBERT pre-training code.

Quick start

If you want to use our pre-trained model directly for molecular property prediction, please run the following command:

>> deepspeed --num_gpus=2  pretrainscKGBERT.py

Other functions

The processed single-cell transcriptome data can be used with our model to accomplish the following tasks:

• Gene annotation

• Cell and cell subtype annotation

• Cancer drug response

• biomarker analysis