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

Brief introduction

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

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

├── 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

Other functions

• Gene annotation
• Cell and cell subtype annotation
• Cancer drug response
• biomarker analysis