预训练基因组变异语言模型更好地建模功能基因组.

paper：Pre-training Genomic Language Model with Variants for Better Modeling Functional Genomics

0. TL; DR

UKBioBERT是一个利用英国生物银行（UK BioBank）约30万欧洲裔个体的遗传变异数据进行预训练的DNA语言模型。UKBioBERT生成的嵌入信息丰富，能够有效地识别基因功能。

在与一系列SOTA基因组语言模型的基准比较中，UKBioBERT在区分不同功能基因的聚类任务上表现最佳。在细胞系层面，将UKBioBERT的嵌入信息整合到现有的序列到功能模型中，能够显著提高基因表达预测的准确性。

1. 背景介绍

将DNA视为一种“生命语言”来解读，是理解遗传学与基因组学复杂性的核心挑战。基因组中的功能元件，其活性受到附近基因组区域上下文的深刻影响，这暗示了DNA序列中存在着类似于自然语言中的“一词多义”和“远距离语义关系”。因此，利用自然语言处理中的先进技术，特别是基于Transformer架构的基础模型来建模DNA序列，已成为一个富有前景的研究方向。

目前，基因组语言模型主要分为两大流派：

掩码语言模型 (Masked Language Model, MLM): 以BERT为代表，通过预测序列中被随机“掩码”掉的部分，来学习双向的上下文信息。DNABERT和GENA-LM等模型均采用此策略。这类模型擅长生成高质量的序列表示，用于下游的分类或回归任务。
因果语言模型 (Causal Language Model, CLM): 以GPT为代表，通过自回归地预测下一个token，来学习序列的生成规律。DNAGPT, HyenaDNA, 和Evo等模型属于此类，它们更侧重于生成新的DNA序列。

尽管gLMs取得了显著进展，但仍存在一些局限性：

模型适用性问题： CLM类模型在生成新序列时可能存在“幻觉”问题，而MLM类模型生成的嵌入质量受多种因素影响，其可靠性仍需检验。
变异信息利用不足：虽然已有模型（如GPN, GENA-LM）尝试在预训练中加入单核苷酸多态性（SNPs）信息，但它们大多将此作为一种数据增强手段，且下游应用仍主要基于参考基因组。这种设计未能充分挖掘遗传变异在真实生物学场景（如个性化基因表达调控）中的核心作用。

为了更有效地建模DNA语言，并解决现有方法的局限性，作者提出了UKBioBERT。该模型旨在通过在预训练阶段，系统性地整合来自大规模人群队列（UK BioBank）的真实遗传变异信息，来学习一个更深刻、更具生物学意义的基因组表示，并将其应用于序列到功能的分析任务中，特别是个性化的基因表达预测。

2. UKBioBERT模型

作者的研究分为两个核心部分：首先，预训练一个名为UKBioBERT的基因组语言模型；然后，将其与现有的序列到功能模型融合，构建一个名为UKBioFormer的应用模型。

2.1 UKBioBERT: 基于变异的持续预训练

UKBioBERT的基础架构是DNABERT2，一个在人类参考基因组上预训练好的、强大的BERT模型。作者采用了持续预训练（continued pre-training）的策略，即在DNABERT2的检查点基础上，利用包含遗传变异的新数据进行进一步的训练。训练数据来源于英国生物银行（UK BioBank）中约30万欧洲裔个体的基因组数据，共包含超过1300万个变异。

作者首先将UK BioBank中的变异信息，通过替换、插入或删除等操作，“编辑”到人类参考基因组序列上，从而生成大量包含真实遗传多样性的“个性化”DNA序列。

预训练任务采用标准的掩码语言模型（MLM）：随机地“掩码”掉序列中的一部分核苷酸，然后让模型去预测这些被掩码的核苷酸是什么。模型通过最小化预测与真实核苷酸之间的交叉熵损失来进行优化。

预训练完成后，UKBioBERT可以作为一个强大的编码器。对于任意一段DNA序列（如一个基因的序列），可以将其输入到UKBioBERT中，并提取模型最后一层输出的[CLS] token的表示，或者对所有token的表示进行池化，从而得到一个代表该序列语义信息的嵌入向量。

2.2 UKBioFormer: 融合gLM与序列到功能模型

为了将UKBioBERT学习到的、蕴含丰富变异信息的基因组知识，应用于具体的序列到功能任务，作者进一步构建了UKBioFormer。

UKBioFormer巧妙地融合了两个模型：

Enformer: 一个强大的、基于Transformer的序列到功能模型，它能够从DNA序列中直接预测基因表达等功能基因组学信号。
UKBioBERT: 预训练好的基因组语言模型。

对于一个给定的DNA序列，作者同时使用Enformer和UKBioBERT对其进行编码，分别得到一个“功能”嵌入和一个“语言”嵌入。然后，将这两个嵌入拼接起来，形成一个更丰富的、融合了两种信息的联合表示。最后，将这个联合表示送入一个预测头，来预测该基因在特定个体中的表达水平。

在训练UKBioFormer时，为了提高效率和防止过拟合，作者采用了PEFT策略。UKBioBERT的权重被冻结，只作为特征提取器。Enformer的权重则进行微调。作者通过实验发现，一种有效的PEFT策略是只微调Enformer中Transformer的少数几层，而非全部。通过这种方式，UKBioFormer能够以较低的计算成本，将UKBioBERT的先验知识高效地注入到Enformer的预测任务中。

UKBioZoi是UKBioFormer的一个变体，它使用Borzoi（另一个序列到功能模型）代替Enformer作为基础架构，其设计思想是相同的。

3. 实验分析

作者通过一系列精心设计的实验，从多个角度验证了UKBioBERT和UKBioFormer的性能和价值。

3.1 实验一：UKBioBERT学习到的基因表示质量评估

作者提出了一个新的评估标准：一个好的基因组语言模型，其生成的嵌入应该能够有效地区分具有不同生物学功能的基因。

图c (与SOTA模型比较): 在与DNABERT, HyenaDNA, GPN, Llama3.1等超过15种gLM和LLM的比较中，UKBioBERT在区分不同功能基因的聚类任务中，取得了最高的综合得分。这证明了其生成的基因嵌入具有最佳的生物学区分度。
图d (UMAP可视化): UMAP降维图直观地显示，UKBioBERT的嵌入空间中，不同功能的基因（如蛋白编码基因）形成了最清晰、最独立的簇。
图b (数据规模影响): 随着预训练中使用的变异数据量的增加，模型的聚类性能也随之提升。这直接证明了整合大规模真实遗传变异信息对于学习高质量基因表示的重要性。

3.2 实验二：UKBioBERT增强细胞系基因表达预测

作者将UKBioBERT的嵌入，作为额外信息加入到EPInformer（一个预测细胞系基因表达的模型）中。

图a： 在K562和GM12878两种细胞系、CAGE-seq和RNA-seq两种测量技术上，加入了UKBioBERT嵌入的EPInformer（模式E和G），其预测准确性（PCC）均优于原始的EPInformer（模式C）和其他变体。
图b： 散点图显示，预测值与真实值高度相关（PCC高达0.902），证明了整合UKBioBERT嵌入的有效性。

3.3 实验三：UKBioBERT解释基因表达的可预测性

作者探究了为什么有些基因的表达可以被序列很好地预测，而另一些则不能。

图a： 在对41个代表性基因的个性化表达预测任务中，UKBioBERT（结合ElasticNet）的性能与SOTA的Performer模型相当。
图b,c (可预测性的解释): 作者发现，一个基因的可预测性（predictability）与其遗传力（heritability）之间没有显著相关性。相反，它与基于UKBioBERT嵌入的个体可聚类性（clustering performance）之间存在极强的正相关（PCC=0.81）。这提出了一个全新的视角：只有当具有相似表达水平的个体，其对应的基因组序列在嵌入空间中也能形成清晰的簇时，这个基因的表达才是可预测的。
图d (跨种族预测): 作者进一步展示了模型在跨种族预测中的局限性。用欧洲裔数据训练的模型，在预测非裔美国人数据时性能显著下降。

3.4 实验四：UKBioFormer提升个性化基因表达预测

作者进一步评估了融合模型UKBioFormer的性能。

图a,b： 对于那些本身可预测性较好的基因，UKBioFormer的性能一致地优于其基础模型Performer和UKBioBERT。
图c,d (联合训练): 作者发现，对于共享增强子（enhancer）或通路（pathway）的基因进行联合训练，并不总能带来性能提升。这表明，对于个性化表达预测，目前单基因模型仍然是更优的选择。

3.5 实验五：UKBioFormer的可解释性与eQTL发现

作者利用计算机模拟突变（In Silico Mutagenesis, ISM）和梯度归因（gradient attribution）等方法，探究了UKBioFormer解释变异效应的能力。

图b (eQTL效应预测): 与Enformer相比，UKBioFormer预测的eQTL效应值与GTEx数据库中的真实值对齐得更好。对于基因JUP，UKBioFormer正确预测了71%的eQTL效应方向，显著高于Enformer（53%）。
图c,e (梯度归因): 通过可视化碱基的重要性，作者发现，在eQTL位点（如rs9910080）引入变异，会引起梯度的显著变化，并且这种变化方向与已知的效应一致。
图d,f (ISM分析): ISM分析也证实，等位基因的改变会导致预测表达水平的相应变化，与eQTL效应吻合。模型还能发现与JUN等转录因子相关的motif。

这些实验全面地证明，通过整合大规模真实变异数据，UKBioBERT和UKBioFormer为从基因组序列预测和解释个性化基因表达，提供了强大、高效且可解释的新工具。