中国人民大学团队开发SPACE:基于MoE的DNA基础模型新范式


中国人民大学团队开发SPACE:基于MoE的DNA基础模型新范式

仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接

将ScienceAI设为星标
第一时间掌握
新鲜的AIforScience资讯
作者|论文团队
编辑|ScienceAI
Sequence-to-function模型是一类强大的基因组模型,能够直接从DNA序列预测对应的功能特征,如染色质开放性、基因表达量、DNA剪切位点等。其中最具代表性的工作是2021年发表在《NatureMethods》的Enformer[ 1]。
近日,谷歌DeepMind发布的AlphaGenome[ 2]模型取得了重要突破。该模型将前身Enformer的输入序列长度扩展至1Mb,引入了剪切强度、DNA三维接触图等一系列新的训练任务,并实现了单碱基精度预测。
AlphaGenome在众多基因组预测任务上达到了当前SOTA性能,再次证明了数据扩展在基因组学建模中的巨大潜力,同时也证明了sequence-to-functionmodel这种有监督训练范式学习到的表征的巨大潜力。
然而,单纯的数据扩展并非提升模型能力的唯一路径。来自中国人民大学高瓴人工智能学院的研究团队提出了一个重要的替代思路:通过对模型架构进行符合生物学原理的系统性创新,可以更高效地提升DNA基础模型的性能与泛化能力。

论文链接:https ://arxiv. org/abs/2506. 01833
Github链接:https ://github. com/ZhuJiwei111/space
Huggingface链接:https ://huggingface. co/yangyz1230/space
SPACE:Species-ProfileAdaptiveCollaborativeExperts
Enformer和AlphaGenome设计范式存在两个内在的架构性制约:
1. 物种共享编码器的局限性:这种「一视同仁」的编码器难以有效区分并建模不同物种间基因调控的特有模式(species-specificcharacteristics)与进化上的保守机制。
2. 独立预测头的知识割裂:基因组的功能图谱,如染色质可及性、组蛋白修饰和转录因子结合,在生物学上是高度关联和协同作用的。独立的预测头无法捕捉这些图谱间的内在依赖关系,从而限制了模型对复杂基因调控网络的深层理解。
图1:SPACE模型架构总览。模型包含三个核心阶段:(1)基于CNN的局部上下文聚合模块;(2)引入物种感知混合专家(MoE)的Transformer编码器;(3)引入谱系分组和双层门控机制的增强型解码器。
为了应对上述挑战,我们设计的SPACE模型引入了两项关键的架构革新:
1. 物种感知编码器(Species-awareEncoder):我们在Transformer编码器中部署了稀疏混合专家(MoE)层,以替代标准的前馈网络。该模块由一个共享的「专家」网络池和多个物种特异性的「门控网络」组成。对于来自特定物种的DNA序列,其对应的门控网络会动态地选择并加权激活一部分专家网络进行计算。这种设计使得模型能自适应地为不同物种分配不同的计算资源和参数路径,从而有效解耦并学习物种间的特有及共享调控特征。此外,我们还引入了专家-物种互信息损失函数,以鼓励专家网络向物种特异化的方向分化。
2. 谱系分组增强解码器(Profile-groupedEnhancementDecoder):在解码阶段,我们依据生物学先验知识将功能相关的基因组图谱(如测量染色质可及性的DNase-seq和ATAC-seq)进行分组。在此基础上,一个新颖的双层门控专家加权聚合模块被用于增强初始预测。第一层门控根据物种和序列的全局信息,动态选择最相关的「专家组」,以捕捉进化上的保守调控模式。第二层门控则在选定的专家组内部,根据各图谱的初步预测特征,为每个图谱精确选择并加权最合适的专家,从而建模图谱间的功能依赖性。最终,增强后的预测结果通过残差连接与初始预测融合,确保了训练的稳定性和预测的精确性。
我们在多个公认的基因组学基准上,对SPACE模型的性能进行了全面和严格的评估。
下游基因组任务基准测试
我们在NucleotideTransformer(NT)论文提出的一套包含18项下游分类任务的基准上,将SPACE与多个主流DNA基础模型进行了比较。这些任务涵盖了染色质谱系预测、调控元件注释和剪接位点识别三大类别。
实验结果表明,SPACE模型在18项任务中的11项上取得了当前最佳(SOTA)性能,显著优于包括DNABERT-2、HyenaDNA以及NT系列在内的无监督预训练模型,也一致性地超越了其监督学习的基线模型Enformer。这一结果强有力地支持了我们的核心论点:一个架构设计优良的监督式预测模型,其学到的序列表示(representation)具有强大的泛化能力,可以作为一个高效的DNA基础模型。
图2:在NucleotideTransformer下游任务中的性能比较(MCC分数)。表格展示了SPACE与多个基准模型在18项任务上的表现。加粗字体表示该任务上的最优性能。SPACE在11项任务中达到SOTA。
跨物种泛化能力评估
为了直接评估SPACE架构改进对模型跨物种泛化能力的影响,我们在GUE基准上进行了测试。我们特意选择了与模型训练物种(人、鼠)进化距离较远的酵母(10项表观遗传标记预测任务)和病毒(COVID变种分类任务)进行评估。
如下图所示,与同样使用人鼠数据进行预训练的Enformer相比,SPACE在所有11项跨物种预测任务上均取得了显著的性能提升。尤其是在酵母的H3K4me3标记预测任务上,SPACE的性能相较于Enformer提升了27. 28个百分点,充分证明了我们提出的物种感知MoE编码器等架构在捕捉和迁移进化上保守的调控特征方面的有效性。
图3:在GUE基准上与Enformer的跨物种泛化能力比较。SPACE在所有酵母和病毒预测任务上的性能(MCC或F1分数)均显著高于Enformer,向上箭头(↑)后的数值表示SPACE的绝对性能提升点数。
总结
AlphaGenome展示了「数据规模为王」的强大潜力之际,我们的SPACE模型提供了另一条同样富有前景的技术路径:通过引入符合生物学原理的、动态自适应的先进架构,可以更深刻、更高效地捕捉DNA序列中的复杂调控语言。我们相信,这种将架构设计与生物学洞见深度融合的研究范式,将为未来构建更强大、更具解释性的基因组学AI模型铺平道路。我们已将所有代码和模型开源,以期与社区共同推动该领域的进步。
参考文献:

人工智能×[生物神经科学数学物理化学材料]
「ScienceAI」关注人工智能与其他前沿技术及基础科学的交叉研究与融合发展。
欢迎关注标星,并点击右下角点赞和在看。
点击阅读原文,加入专业从业者社区,以获得更多交流合作机会及服务。


文章作者: ZejunCao
版权声明: 本博客所有文章除特別声明外,均采用 CC BY 4.0 许可协议。转载请注明来源 ZejunCao !
  目录