通过核苷酸级的监督与基因组压缩实现基因组语言模型的功能上下文学习.

paper：Functional In-Context Learning in Genomic Language Models with Nucleotide-Level Supervision and Genome Compression

0. TL; DR

这项工作从训练目标和数据构建的角度重新审视了长上下文基因组基础模型的设计，引入了两个关键技术：

因子化核苷酸监督（Factorized Nucleotide Supervision, FNS）：该技术通过概率边缘化，将高效的k-mer分词与单核苷酸似然联系起来，从而在不牺牲计算效率的情况下恢复了单核苷酸级别的分辨率。
基因组压缩预训练（Genome Compression Pretraining, GCP）：该策略通过关注以基因为中心和调控区域，重塑了训练数据的分布。

本文提出了GENERator-v2，一个支持长达98k碱基对上下文的自回归基因组基础模型家族。在零样本评估和下游微调基准测试中，GENERator-v2相比其前代模型持续改进，同时实现了更高的推理效率。

1. 背景介绍

高通量测序技术的飞速发展导致了基因组数据前所未有的积累，这激发了人们对基因组基础模型（GFM）的日益增长的兴趣。这些模型旨在直接从原始DNA序列中学习通用的表示，以支持基因组注释、变异解释和序列设计等多种下游任务。

现有的GFM在架构和训练目标上有所不同。掩码语言模型（Masked language models, MLM）如DNABERT-2和Nucleotide Transformer，在序列理解和判别性预测任务上表现出色，但其双向掩码目标限制了其生成能力。自回归模型（Autoregressive models）如HyenaDNA和Evo2，通过预测下一个词元来解决这个问题，但所有这些方法在扩展到基因组长度时都面临一个共同的挑战：如何在长程上下文、单核苷酸分辨率和计算效率之间取得平衡。

长上下文建模在基因组学中尤为重要，因为有意义的预测通常需要整合分布在广大基因组区间内的多个基因和调控单元的信息。然而，在单核苷酸分辨率下对长基因组上下文进行建模，对于标准的Transformer架构来说计算成本过高。

为了缓解这一矛盾，研究者提出了多种策略。一些模型通过引入卷积组件来扩展有效上下文长度，如Enformer和HyenaDNA。另一些则采用状态空间模型（SSMs），如Caduceus，其计算复杂度随序列长度线性增长。然而，这些架构上的创新在实践中将额外上下文转化为性能提升的能力有限。

另一种方法是通过k-mer分词来减少序列长度。但这种方法牺牲了单核苷酸分辨率，并且对分词的起始位置敏感。

在这项工作中，作者没有选择复杂的架构创新，而是回归本源，从训练目标和数据构建的角度重新思考问题。作者提出了GENERator-v2，它保留了标准的Transformer架构和高效的k-mer分词，但通过两个核心技术（因子化核苷酸监督（FNS）和基因组压缩预训练（GCP））解决了现有方法的局限性。

2. GENERator-v2

GENERator-v2的核心创新在于其训练策略，而非模型架构。它采用了与GENERator-v1相同的Transformer解码器骨干，但引入了FNS和GCP来提升性能。

2.1 因子化核苷酸监督 (Factorized Nucleotide Supervision, FNS)

标准的k-mer分词将一段DNA序列（如6个碱基）视为一个不可分割的原子单元。在训练时，模型的目标是预测正确的k-mer。这意味着，一个预测出的k-mer即使只与真实值相差一个核苷酸，其受到的惩罚也和一个完全错误的预测相同。这种“非黑即白”的监督方式与生物学现实不符，并可能导致训练不稳定。

FNS通过将每个k-mer的预测任务分解为$k$个独立的单核苷酸预测任务来解决这个问题。模型仍然输出一个关于所有可能k-mer的概率分布，但损失函数是在核苷酸级别上计算的。

对于一个k-mer词汇表$V = B^k$（其中$B = {A,C,G,T}$），模型输出每个k-mer $v$的概率$p(v)$。FNS首先计算在第$j$个位置出现某个碱基$b$的边缘概率：

\[p^{(j)}(b) = \sum_{v \in V: v(j)=b} p(v)\]

然后，对于一个真实的k-mer目标$y = (y^{(1)}, \dots, y^{(k)})$，FNS损失被定义为所有$k$个位置上负对数似然的总和：

\[L_{FNS} = -\sum_{j=1}^{k} \log p^{(j)}(y^{(j)})\]

FNS恢复了单核苷酸分辨率，使模型能够进行单核苷酸变异（SNP）级别的概率推理。通过提供更密集的梯度信号，使训练过程（尤其是在后期）更加稳定。FNS完全在损失函数层面操作，不改变模型架构或分词方式，因此保留了k-mer分词带来的计算优势。

2.2 基因组压缩预训练 (Genome Compression Pretraining, GCP)

真核生物基因组中，功能性信号（如基因和调控元件）是极其稀疏的，被大量的“背景”序列所分隔。在长上下文模型中，这意味着大部分计算能力被浪费在处理这些低信息量的区域上，而模型很少有机会学习跨越多个功能单元的长程依赖关系。

GCP是一种数据层面的预训练策略，它通过构建一个“压缩”的基因组来解决这个问题。

首先从参考基因组中识别出所有已注释的功能区域（如以基因为中心的区域和调控元件）。
然后按照它们在基因组上的原始顺序，将这些区域的DNA序列拼接在一起，并在区域之间插入一个特殊的分隔符（<s>）。
最后模型在这个压缩后的、信息密集的序列上进行训练。

GCP提供了高效的上下文利用，一个训练窗口现在可以包含数十甚至数百个功能单元，即使它们在原始基因组中相距数兆碱基。GCP在数据层面实现了类似于“稀疏注意力”的效果，鼓励模型学习功能单元之间的相互作用关系，即“下一个基因”的预测。与真正的稀疏注意力机制相比，GCP的方法更稳定，因为它允许窗口内的所有词元参与标准的密集自注意力计算。

3. 实验分析

作者在一系列零样本评估和微调基准测试中，系统地评估了GENERator-v2的性能。

3.1 序列恢复（Sequence Recovery）

这是一个零样本任务，用于评估模型的内在生成准确性和效率。模型被给予一个6144 bp的输入序列，并被要求生成接下来的30个核苷酸。

图A (真核生物)：GENERator-v2相比其前代v1版本有稳定和系统的提升。在相同参数规模下，GENERator-v2-1B优于Evo2-1B，而GENERator-v2-3B的性能与参数量大得多的Evo2-7B相当。
图B (原核生物)：在不同菌群中，GENERator-v2和Evo2各有胜负，但总体性能相当，并且都远超Evo-7B。
图C (生成效率)：GENERator系列模型的推理速度比Evo系列快了几个数量级。例如，完成一个基准测试，GENERator-v2-3B用时0.62小时，而Evo2-7B需要27.12小时。这凸显了k-mer分词和Transformer架构在效率上的巨大优势。

3.2 变异效应预测 (Variant Effect Prediction)

在ClinVar数据集上进行零样本评估，预测SNP的致病性。该任务直接检验模型预测的核苷酸级概率是否具有生物学意义。

（见上图D）GENERator-v2相比GENERator-v1在该任务上有明显提升，这直接证明了FNS的有效性。在相似的参数规模下，GENERator-v2-1B的性能优于Evo2-1B。在计算成本几乎相同的情况下，GENERator-v2-3B的性能与Evo2-7B持平。

这些结果表明，GENERator-v2能够在不依赖多序列比对的情况下，准确预测变异的功能影响，同时保持了高效的计算性能。

3.3 嵌入表示分析 (Embedding Representations)

通过UMAP可视化，分析不同模型学习到的序列嵌入的结构，以探究GCP对表示学习的影响。

内容词元嵌入（Content Token Embedding）：GENERator-v2的嵌入清晰地分离出了DNA链方向和密码子阅读框相位这两种局部序列特征，并且聚类结构高度对称。GENERator-v1也能捕捉到这些特征，但结构不那么规整。Evo2的嵌入结构最弱，分离不明显。
分隔符嵌入（Separator Embedding）：对于GENERator-v2，分隔符<s>的嵌入不再编码局部信息，而是根据物种的分类进行聚类。这表明，在GCP训练下，分隔符<s>学会了充当一个“摘要词元”，总结前一个功能区域的信息，以预测下一个功能区域。这正是“功能性上下文学习”在表示层面的体现。GENERator-v1（未使用GCP）的分隔符嵌入则不具备这种摘要功能。

这些结果表明，GCP成功地引导模型学习到了从局部句法特征到高级生物学语义的层次化表示。

3.4 下游任务微调

在修订版的Nucleotide Transformer (NT)基准任务上，对GENERator-v2进行微调，并与Enformer, DNABERT-2, HyenaDNA等一系列模型进行比较。

在绝大多数任务上，GENERator-v2的性能都优于其前代GENERator-v1，并达到或超过了所有其他基线模型，包括NT-v2 (500M)和Caduceus (8M)。这些结果表明，通过FNS和GCP带来的改进，能够有效地迁移到各种有监督的下游任务中。