The GB.Tissue model is developed on spatial single-cell transcriptomic data. To incorporate spatial cell information, K nearest neighbor cells (K=8 in our case) are retrieved for each center cell. The center cell and neighbor cell expression vectors are concatenated as model input. Two-dimension positional rotary embedding are introduced to encode both the gene and cell information. The first dim is gene index and second dim is cell index. The overall training scheme is similar to scFoundation with an asymmetric encoder-decoder architecture.

We evelauted the model on two spatial data task, including predicting niche label and cell density. Please refer to GB.tissue-downstream-tasks for more details. Below shows the brief description and evaluation metrics.

niche label type prediction: the task is to predict niche type of each cell given spatial expression data (in total 6 types).

cell density prediction: the task is to predict neighbor cell number of a target cell given the expression profiles.

We introduce how to finetune and evaluate our pre-trained GB.Tissue foundation models for downstream tasks. These tasks can be classified into the following categories:

Note: All the following scripts should be run under ModelGenerator/.

The related data is deposited at https://huggingface.co/datasets/genbio-ai/tissue-downstream-tasks. Please download the data and put under ModelGenerator/downloads as cell_density or niche_type_classification. Under each sub-directory, there are three files denote different split (xx.train.h5ad, xx.val.h5ad, xx.test.h5ad).

For each .h5ad, several obs attributes should be included to reprezent the spatial (coordinate) information (like x, y), the label information (like niche_label). All the column fields will be specified in the following config.yaml file.

Note: the file scRNA_genename_and_index.tsv includes all the corresponding gene name and index in h5ad file.

We fully finetune GB.Tissue for niche label type prediction.

Note:

The filter_columns includes label column and spatial coordinate column. rename_columns keep unchanged and will be used for running.

Once finished run, there will be several ckpt file under the specified output directory default_root_dir. Then we can use the ckpt to evaluate on test dataset.

Note: ckpt_path is the finetuned checkpoint path.

The config file is like experiments/GB.Tissue/cell_density_regression.yaml, all the fintuning running and evaluation are similar as classification task.

We can dump embedding for a .h5ad file. The script is as:

The output file will be under specified output_dir like ./logs/emb.xenium/lightning_logs/pred_output. Each batch will be saved and a merged one will also be generated as predict_predictions.pt. The predict_predictions.pt file satcks all batches:

We can retrieve all the gene embedding and aggregate into cell embedding (like max pooling):