SATURN：利用蛋白质大模型嵌入进行跨物种整合

伴随着众多大模型新方法的涌现，跨物种整合领域出现了许多利用蛋白质序列信息进行嵌入的整合方法，有一些具备相当的准确性。

这里记录一下用 Rosen, Y., Brbić, M., Roohani, Y. et al. Toward universal cell embeddings: integrating single-cell RNA-seq datasets across species with SATURN. Nat Methods (2024). https://doi.org/10.1038/s41592-024-02191-z 方法文章中数据进行测试的全流程。

原始代码和数据链接见https://github.com/snap-stanford/SATURN

1.环境准备

https://github.com/snap-stanford/SATURN/blob/main/requirements.txt

建议在配置GPU的环境下训练模型，否则会相当耗时

2.数据准备

物种单细胞数据

SATURN/Vignettes/frog_zebrafish_embryogenesis/dataloader.ipynb at main · snap-stanford/SATURN · GitHub

这里准备的是物种带注释的计数矩阵，并且都需要转成h5ad的格式

物种蛋白质嵌入文件

还需要准备每个物种的蛋白质序列，可以从NCBI-DATASET中下载

提取之前可以进行一定的预处理和数据清洗，见clean_fasta.py(SATURN/protein_embeddings at main · snap-stanford/SATURN · GitHub)

现有的蛋白质大模型很多，这里选择ESM2，提取蛋白质嵌入信息的torch文件详见extract.py

(esm/scripts/extract.py at main · facebookresearch/esm · GitHub)

具体的一些参数解释如下

#Extraction of the final-layer embedding for a FASTA file from the ESM-2 model
python scripts/extract.py esm2_t33_650M_UR50D \
       examples/data/some_proteins.fasta \
       examples/data/some_proteins_emb_esm2 \
       --repr_layers 0 32 33 \   #(default: final only) selects which layers to include embeddings from
       --include mean per_tok
#per_tok includes the full sequence, with an embedding per amino