0. TL; DR

SNPBag是一个专注于单核苷酸多态性（single nucleotide polymorphism, SNP）的基础模型。SNPBag拥有8亿参数，在一个包含100万合成人类基因组的数据集上进行了预训练，总共处理了6万亿个SNP token。

在多项任务的基准测试中，SNPBag表现出卓越的性能。SNPBag能够将每个基因组的600万个SNP编码成一个仅0.75 MB的嵌入。这种紧凑的表示极大地便利了数据的存储、传输和下游应用。基于基因组嵌入，SNPBag能够快速地进行全球人群的祖源推断（ancestry inference），并能检测出远至12代的亲缘关系（genetic relationships）。

1. 背景介绍

人工智能（AI），特别是基于Transformer的基础模型（foundation models），已经彻底改变了生物医学研究。在蛋白质组学领域，AlphaFold系列模型精确预测蛋白质结构，加速了药物发现。在转录组学领域，Geneformer和scGPT等模型支持了细胞类型注释和基因调控网络推断。在基因组学领域，Nucleotide Transformer和EVO等模型通过处理跨物种的全基因组序列，实现了基因功能注释。

然而，在针对单一物种（如人类）内部的基因组变异（genomic variations）进行建模方面，仍然缺乏一个能够处理海量个体数据的AI基础模型。考虑到人类基因组数据的巨大体量和关键意义，这一空白尤为突出。从国际人类基因组单体型图计划（International HapMap Project）到千人基因组计划（1000 Genomes Project），再到英国生物银行（UK Biobank）等数百万级别的大型生物银行，人类基因组数据正以前所未有的速度累积。

但是，为人类基因组数据设计AI模型面临着独特的挑战：

数据规模的挑战：人类基因组极其庞大（约30亿个核苷酸），而生物银行已包含数百万个体。如果直接采用现有DNA语言模型的框架（如EVO），需要处理的token数量将达到10¹⁵级别，这远远超出了当前最先进LLM的训练规模，在未来几年内都是不可行的。
数据冗余的挑战：人类个体间的基因组序列高度相似（差异仅为0.1~0.6%），对全序列进行建模是极度冗余且低效的。

为了应对这些挑战，作者提出了一个专为生物银行级别的人类基因组数据设计的AI基础模型框架，其核心特征包括：

以变异为中心的表征 (Variant-centered tokenization): 关注多态性位点（如SNP），而非每个核苷酸，从而大幅减少计算量。
全基因组建模 (Whole-genome modeling): 分析整个基因组，以捕捉长距离的连锁不平衡（linkage disequilibrium）和全基因组范围内的上位性相互作用（epistatic interactions）。
紧凑的基因组嵌入 (Compact genome embeddings): 生成低维表示，以支持高效的下游应用。
多任务训练 (Multitask training): 在一个统一的框架内，解决多种人类基因组变异分析任务。

为了探索该框架的可行性，作者开发了SNPBag，一个专注于全基因组单核苷酸多态性（SNP）数据的基础模型。SNP是人类最常见的遗传变异形式，也是全基因组关联分析（GWAS）的基础。现有的SNP分析工具虽然有效，但在可扩展性、计算效率和任务特异性方面存在局限。SNPBag旨在通过一个统一的、可扩展的AI框架，克服这些限制。

2. SNPBag 方法

SNPBag的核心是一个基于BERT (Bidirectional Encoder Representations from Transformers)架构的模型，它通过自注意力机制（attention mechanism）来捕捉SNP位点之间的复杂依赖关系。

2.1 模型架构与预训练

SNPBag的输入不是原始的DNA序列，而是经过以变异为中心处理后的SNP序列。每个token代表一个SNP位点。输入序列由两部分信息组成：SNP的标识符（ID）和其对应的基因型（genotype）（编码为0, 1, 2）。

SNPBag的预训练任务是一个掩码基因型预测任务，类似于BERT的掩码语言模型。在每个输入序列中，作者随机地掩码（mask）掉85%~99%的SNP基因型，并让模型去预测这些被掩码位置的原始基因型。

输入序列首先通过一个线性嵌入层，被投影到一个512维的向量空间。然后，这些嵌入向量被送入一个由16个Transformer层组成的BERT编码器。每个Transformer层包含16个注意力头。通过多层的自注意力计算，模型能够学习到基因组上、甚至是长距离SNP之间的连锁不平衡和复杂的依赖关系。

BERT编码器的输出，通过一个多层感知机解码器，来预测每个被掩码位置的基因型概率。模型通过最小化预测基因型和真实基因型之间的交叉熵损失来进行优化。

2.2 多任务微调 (Multi-task Fine-tuning)

预训练好的SNPBag基础模型，可以通过更换解码器头的方式，被微调以适应多种不同的下游任务。

基因型填充 (Genotype Imputation): 使用与预训练相同的MLP解码器，但微调时采用特定SNP芯片（SNP array）的掩码模式。
单倍型定相 (Haplotype Phasing): 更换一个新的MLP解码器，使其预测相邻杂合位点之间的等位基因转换状态（switch states）。
基因组嵌入 (Genome Embedding): 训练一个额外的、由多个MLP组成的压缩头，将编码器的输出压缩成一个低维的基因组嵌入向量。

通过这种预训练-微调的范式，SNPBag能够将从海量数据中学到的通用基因组结构知识，高效地迁移到各种具体的分析任务中。

3. 实验分析

作者通过一系列精心设计的基准测试，全面地评估了SNPBag在多种基因组分析任务中的性能，并与现有的SOTA方法进行了比较。

3.1 实验一：基因型填充 (Genotype Imputation)

作者在HGDP数据集上，模拟了使用Illumina Omni2.5芯片数据，来填充（impute）整个22号染色体上所有SNP的任务。

图a, d (总体性能): 经过微调的SNPBag omni模型，其填充准确率达到了97.16%，显著优于所有传统的、包括使用了参考面板的SOTA方法（如BEAGLE5.2+IMPUTE5）。同时，其推理速度比传统方法快了13到20倍。
图b (按MAF分层): 在不同的次要等位基因频率（minor allele frequency, MAF）区间，SNPBag的性能都稳定地优于传统方法。
图c (按人群分层): SNPBag在所有大陆人群中都取得了最高的准确率。尽管在遗传多样性最高的非洲人群中性能有所下降，但这符合所有填充方法的普遍规律。

3.2 实验二：单倍型定相 (Haplotype Phasing)

作者在26个经过物理定相的HGDP个体上，评估了SNPBag的定相性能。

图a, b (准确率): SNPBag的转换错误率（switch error rate）为3.07%，显著优于所有不使用参考面板的传统方法，并与使用参考面板的SOTA方法（2.54~2.70%）的性能非常接近。
图c (速度): SNPBag的推理速度极快，比传统方法快了11到91倍。

这些结果表明，SNPBag在不依赖外部参考面板的情况下，实现了极具竞争力的定相精度和无与伦比的效率。

3.3 实验三：基因组嵌入与人群分类

作者使用预训练好的SNPBag模型，为1kG项目中的3202个个体生成了全基因组嵌入，并评估了这些嵌入在人群结构分析中的应用。

图a (UMAP可视化): UMAP降维图清晰地显示，SNPBag生成的基因组嵌入，能够完美地将不同大陆的人群（如欧洲人、东亚人、非洲人）分离开，并且能够进一步区分出次大陆级别的精细结构（如日本人与汉族人、芬兰人与南欧人）。
图b (分类准确率): 基于这些嵌入训练的分类器，在对19个亚人群进行分类时，平均F1 score达到了0.861，优于之前报道的其他机器学习模型。

这证明了SNPBag的基因组嵌入，在极大地压缩了数据（每个基因组仅0.75 MB）的同时，完整地保留了精细的群体遗传结构信息。

3.4 实验四：亲缘关系推断 (Relatedness Inference)

作者进一步探究了基因组嵌入是否能用于推断个体间的远亲关系。

图a (混淆矩阵): 作者在一个包含多达9代远亲的模拟家系数据上训练了一个MLP模型。结果显示，该模型能够准确地推断出远至12代的亲缘关系。
图b (准确率曲线): 对于近亲（1-2代），模型的准确率接近100%。对于远亲（7-9代），即使是精确匹配的准确率（D0）也达到了24~27%。如果允许1-3代的误差，那么在11代以内的准确率都保持在75%以上。
与Bonsai的比较： SNPBag的准确率曲线与工业界金标准23andMe的Bonsai软件高度一致。

这表明，SNPBag的嵌入保留了用于推断远亲关系所需的、长距离的单倍型信息，并且其推理效率极高（每对个体<0.02秒）。