PlantCAD2: 被子植物跨物种功能注释的长上下文DNA语言模型.

0. TL; DR

PlantCAD2 是一个具有长上下文、植物特异性的DNA语言模型，它具有单核苷酸分辨率，并在65个被子植物基因组上进行了预训练。

在12项零样本测试任务中，参数量为6.76亿的PlantCAD2在其中10项任务中超过了参数量高达70亿的Evo2模型，展示了其高效捕捉进化保守性的能力。在7个跨物种的功能基因组学任务中（包括染色质可及区域、基因表达和蛋白质翻译预测），PlantCAD2的性能均优于参数量为10亿的AgroNT模型。PlantCAD2的8192 bp上下文窗口显著提高了在玉米等大基因组中预测染色质可及性的性能，强调了长程上下文对于建模远距离调控的重要性。

1. 背景介绍

破译DNA序列如何编码分子功能、表型和适应性是生物学的终极目标之一。随着测序成本的急剧下降，大规模的基因组测序项目正在产生海量的基因组数据。然而，功能注释却远远落后，尤其是在植物中，标记数据仅限于少数模式物种和作物。这一差距凸显了对能够仅从原始序列中学习并能将知识迁移到其他植物物种的计算模型的迫切需求。

最近，基于自监督策略预训练的基础模型为大规模解释基因组序列开辟了新的可能性。与需要大量标记数据的传统有监督学习方法不同，基础模型在海量未标记数据上进行预训练，然后通过在少量标记数据上进行微调，即可适应特定的下游任务。这种方法已在蛋白质科学领域取得巨大成功。

然而，相比蛋白质语言模型，DNA语言模型面临着更大的挑战。植物基因组中普遍存在的重复序列会干扰预训练；非编码区比编码区保守性更低；而且，对双链DNA的建模需要考虑反向互补（reverse complementary, RC）对称性和双向上下文。

虽然已有一些植物特异性的DNA语言模型，如GPN和AgroNT，但它们存在局限性。例如，PlantCAD的作者们之前开发的第一代模型PlantCAD1（PlantCaduceus），其上下文窗口只有512 bp，限制了其捕捉长程调控相互作用的能力。而AgroNT虽然上下文窗口更长，但其k-mer分词策略牺牲了单核苷酸分辨率。通用DNA模型Evo2虽然参数巨大，但其预训练数据跨越整个生命之树，可能无法很好地捕捉植物特有的调控模式。

为了解决这些问题，作者推出了PlantCAD2，一个专门为被子植物基因组量身定制的、具有长上下文窗口的改进版DNA语言模型。

2. PlantCAD2框架

PlantCAD2在第一代PlantCAD的基础上，从架构、上下文长度、参数规模和系统发育广度四个方面进行了重大改进。

2.1 模型架构

PlantCAD2的核心架构基于Caduceus，并用Mamba2模块替换了原来的Mamba1模块。

Mamba2是一种高效的状态空间模型（SSM），其计算复杂度随序列长度呈线性增长，而非Transformer的二次方增长。这使得PlantCAD2能够高效处理长序列。

模型保留了Caduceus的双向处理（同时考虑上游和下游序列）和反向互补（RC）等变性设计，这对于理解DNA的生物学特性至关重要。

作者训练了三种不同规模的模型：PlantCAD2-S (88M参数), PlantCAD2-M (311M参数), 和PlantCAD2-L (694M参数)。

2.2 预训练策略

作者精心挑选了65个被子植物的基因组，最大化了系统发育的多样性，以避免模型偏向于某些研究得较多的科属（如禾本科和十字花科）。

得益于Mamba2的效率，PlantCAD2的上下文窗口从PlantCAD1的512 bp扩展到了8,192 bp，这对于捕捉远距离调控至关重要。

模型采用单核苷酸分词，保留了最高的序列分辨率。采用掩码语言建模（Masked Language Modeling, MLM）作为目标。为了减轻基因组中大量重复序列对训练的干扰，作者在计算损失时，对重复区域的权重进行了降低。损失函数定义为：

\[\mathcal{L} = \frac{\sum_{i \in M} w_i \cdot \ell(y_i, \hat{y}_i)}{\sum_{j \in M} w_j}\]

其中，$w_i$是第$i$个核苷酸的权重，如果该位置在重复区域内，则$w_i = 0.1$，否则$w_i = 1.0$。

2.3 下游任务应用策略

3. 实验分析

3.1 模型性能与扩展性分析

作者评估了不同模型大小和上下文长度对预训练性能和推理速度的影响。

3.2 零样本预测进化保守性

作者设计了三个独立的任务来评估模型在零样本条件下预测进化保守性的能力。

3.3 零样本预测转录和翻译关键连接点

通过零样本方式，评估模型恢复被掩码的关键基序（如ATG起始密码子、GT剪接供体位点等）的能力。

在玉米和番茄（一个未参与预训练的物种）上，PlantCAD2在所有四种连接点（TIS, TTS, 剪接供体/受体）的预测准确率上，都一致地、大幅度地优于其前身PlantCAD1、GPN和Evo2。同样，Evo2在需要双向信息的任务上表现不佳，而PlantCAD2的双向和RC等变性设计使其能够稳健地处理所有类型的连接点。

3.4 零样本预测结构变异的功能影响

首次尝试使用DNA语言模型来预测小片段缺失（一种结构变异）的功能影响。

PlantCAD2预测的缺失影响分数（$Δ \log P$）与基于多序列比对的进化保守性分数（phyloP）高度相关。模型能够准确地区分位于高度保守区域和非保守区域的缺失，其AUROC性能优于参数量大得多的Evo2。这表明PlantCAD2学习到的表示能够推广到单核苷酸变化之外，捕捉与非编码区结构变异相关的更广泛的序列依赖性。

3.5 微调预测染色质可及性

通过LoRA参数高效微调，训练PlantCAD2来预测全基因组的染色质可及区域（ACRs）和细胞类型特异性的ACRs。

3.6 微调预测基因表达和蛋白质丰度

通过LoRA微调，训练PlantCAD2预测基因表达和蛋白质翻译水平（基于Ribo-seq数据）。