AntigenLM: 结构感知的流感DNA语言建模.

paper：AntigenLM: Structure-Aware DNA Language Modeling for Influenza

0. TL; DR

AntigenLM 是一个在流感病毒的八个基因片段上进行预训练的生成式DNA语言模型。AntigenLM能够捕捉进化约束并在不同任务间实现泛化。在对时间序列的血凝素（HA）和神经氨酸酶（NA）序列进行微调后，AntigenLM能够准确预测未来在不同地区和亚型中出现的抗原变体，包括那些在训练中未见过的变体，其性能优于传统的系统发育和基于进化的模型。此外，它还能实现近乎完美的亚型分类。

1. 背景介绍

流感病毒进化迅速，以逃避宿主免疫，导致季节性流行病并需要频繁更新疫苗。因此，准确预测病毒进化对于疫苗株的选择和减轻全球公共卫生负担至关重要。目前，世界卫生组织（WHO）的疫苗推荐整合了大规模的基因组监测、抗原表征和流行病学监测。

传统的预测方法依赖于系统发育树的动态和基于突变的预测模型。最近，深度学习模型，如beth-1，已被证明能够准确预测进化分支的突变轨迹。与此同时，基因组基础模型（如DNABERT, HyenaDNA）和蛋白质基础模型（如ESM）也为抗原序列表示提供了另一种强大的方法。

然而，病毒的进化是由整个基因组范围内的协同相互作用所塑造的。这包括RNA-RNA相互作用、区段重配的限制，以及聚合酶区段（PB1/PB2/PA）与抗原蛋白HA和NA之间的协同适应。忽视这种结构性上下文的模型，如仅关注单个位点突变的beth-1，会割裂生物学信号，限制其泛化能力和可解释性。而通用的基础模型，由于在异构的多物种基因组上进行训练，往往难以捕捉到这些基因组范围内的、高阶的约束。

此外，尽管基于蛋白质水平的预测能捕捉到抗原进化的某些方面，但许多对病毒适应性至关重要的因素（如同义突变、非编码调控元件、RNA二级结构、包装信号等）对仅蛋白质模型是不可见的。这些观察共同促使作者开发一个针对流感病毒的、全基因组、核苷酸水平的语言模型，以实现准确的抗原序列预测。

为此，作者引入了AntigenLM，一个在预训练期间明确保留基因区段顺序和功能单元完整性的生成式DNA语言模型。通过强制模型学习基因区段的方向性和抗原边界，AntigenLM学习到的表示整合了局部序列依赖性和全局基因组结构，使其能够跨进化分支、亚型和地理区域进行泛化。

2. AntigenLM框架

AntigenLM是一个基于Transformer的框架，专为模拟甲型流感病毒基因组而设计。该模型源自GPT-2架构，但为适应生物序列学习的独特挑战进行了重新设计。

2.1 模型架构

AntigenLM的核心是一个解码器Transformer骨干，其参数相对紧凑（6层，384个隐藏维度，6个注意力头），但支持长达13,000个位置的嵌入范围，足以在不进行截断的情况下对完整的流感基因组进行建模。

在骨干网络之上，作者实现了一个双头设计，以支持多任务学习：

语言建模头：以自回归的方式预测下一个核苷酸词元，用于捕捉进化动态。
分类头：将特定位置的隐藏状态投射到亚型logits，用于监督式的序列级判别任务。

2.2 功能单元编码 (Functional-Unit Encoding)

AntigenLM采用了一个两阶段的功能单元编码策略，以同时保留全基因组上下文和区段级结构。

预训练阶段：将流感病毒的所有八个基因区段（PB2, PB1, PA, HA, NP, NA, MP, NS）按照固定的顺序拼接成一个连续的全基因组序列。这使得Transformer能够捕捉跨区段的长程共进化依赖关系，例如HA和NA之间的补偿性突变。
微调阶段：使用哨兵词元（sentinel tokens），如<HA>和<NA>，来明确界定功能区域。这不仅能引导注意力机制，还能在自回归生成时将解码限制在每个抗原内部，防止跨区段生成。

2.3 微调任务

作者针对两个下游任务对AntigenLM进行了微调：

任务一：抗原进化预测（生成任务）

将来自同一地区、连续三个过去时间点（如月份）的HA/NA序列对拼接成一个长序列作为prompt。每个序列块由<subtype>, <HA>, <NA>等哨兵词元分隔。

\[\text{block}^{(1)}_{\text{past}} \text{block}^{(2)}_{\text{past}} \text{block}^{(3)}_{\text{past}} \text{block}^{(\star)}_{\text{future}}\\ \text{block}^{(i)} = \text{<subtype><HA>HA}^{(i)}\text{<NA>NA}^{(i)}\text{<sep>}\]

优化标准的因果语言建模损失（causal language modeling loss），即预测未来时间点（$t+1$）的HA/NA序列。

\[L_{CLM} = -\sum_{t=1}^{T-1} \log p(x_{t+1} | x_{\leq t})\]

在推理时，模型仅接收三个历史序列块作为输入，并从<HA>词元开始自回归地生成未来的抗原序列。

任务二：亚型分类（判别任务）

输入包含一个HA/NA对的简单序列：<HA>HA<NA>NA<sep>。

将最后一个词元位置的隐藏状态作为整个序列的紧凑表示，通过一个线性层将其投射到亚型logits，并使用交叉熵损失进行训练。

3. 实验分析

3.1 预训练消融研究

为了剖析跨区段依赖、核苷酸级变异、基因组完整性和区段顺序对模型性能的贡献，作者设计了五种不同的预训练策略进行消融研究：

全基因组（Full-genome）：标准的AntigenLM策略。
不完整基因组（Incomplete-genome）：随机裁剪拼接后的基因组，破坏了部分基因边界。
分段式（Segment-wise）：每个训练样本只包含一个基因区段，移除了所有跨区段上下文。
仅抗原（核苷酸）：只使用HA和NA的核苷酸序列。
仅抗原（蛋白质）：只使用HA和NA的氨基酸序列。

在训练动态上，全基因组模型的损失下降最快，收敛到最低点。分段式模型收敛较慢，且最终损失更高。不完整基因组模型的损失几乎没有下降。这表明，完整的跨区段上下文就像一个辅助的监督信号，提高了样本效率；而破坏功能单元的完整性或结构则会严重削弱学习信号。

3.2 下个月HA/NA序列预测

评估模型捕捉短期进化动态的能力。

（图A）AntigenLM（全基因组预训练）能够平滑地进行逐月预测，其预测序列与真实序列的平均氨基酸不匹配数（Mean AA Mismatch）非常低（HA为3-4个，NA为1-2个）。通用的语言模型HyenaDNA和ProtGPT2生成的序列在长度上偏差很大，无法计算不匹配数。不完整基因组基线模型经常生成无效序列，且不匹配率极高。分段式和仅抗原模型的性能与AntigenLM相当但略有下降。

这些结果表明，保持至少HA和NA区段的完整性对于生成有效的预测至关重要，而整合全基因组的上下文能进一步提高预测的准确性。

3.3 预测下一季度的流行毒株

评估模型在疫苗株选择相关任务上的表现，预测下一季度（season T+1）的流行毒株序列。

（图B）AntigenLM在H1N1和H3N2两个亚型上都一致地取得了最低的氨基酸不匹配数。相比当前的WHO疫苗推荐系统，AntigenLM将不匹配数减少了超过70%；相比基于位点的模型beth-1，减少了50%。这凸显了AntigenLM在预测准确性上的显著优势，及其在指导疫苗设计方面的潜力。

3.4 跨亚型和地理区域的泛化能力

实验设置：

跨亚型迁移：在低资源亚型H7N9上进行测试，该亚型在微调数据中占比极低（0.3%）。
地理泛化：将美国的数据作为未见过的区域进行留出，仅在欧洲和亚洲的序列上微调模型，然后在美国数据上进行评估。

结论：

图C (跨亚型)：尽管训练数据极度稀缺，AntigenLM仍然能够为H7N9生成与主要亚型相当的预测。这证明了其有效的跨亚型泛化能力。
图D (跨地理区域)：在未见过的美国数据上，AntigenLM在HA的预测上仍然显著优于beth-1。模型在100个测试案例中正确预测了90个进化分支的变化。

这些结果表明，AntigenLM学习到的是全球性的进化约束，而不是特定区域的突变模式，使其能够在数据采样不足的地区做出可靠的预测。

3.5 亚型分类

在预训练的骨干网络上微调一个简单的分类头，以评估模型学习到的表示的质量。

AntigenLM（全基因组）和仅抗原（核苷酸）模型都取得了近乎完美的分类准确率（F1 score=99.81%），其混淆矩阵呈现出清晰的对角线主导。相比之下，破坏了基因组完整性的不完整基因组和分段式变体则表现出明显的亚型混淆，特别是在稀有亚型上。

这表明AntigenLM学习到了一个结构良好、亚型可分的潜在空间，为自动化流感监测和实时毒株追踪等应用提供了支持。