TaxDistill：通过蒸馏基因组基础模型改进宏基因组分类注释

# TaxDistill: 通过蒸馏基因组基础模型改进宏基因组分类注释
来源: https://arxiv.org/html/2605.28868
叶荣野1,3,4,†, 李伦1,2,3,†, 罗铮1,3,4, 詹逸然1,3,4, 宋述慧1,2,3,4 
1 国家基因组科学数据中心，中国国家生物信息中心，北京 100101，中国
2 北京生物大数据智能治理与应用重点实验室，中国国家生物信息中心，北京 100049，中国
3 中国科学院北京基因组研究所，北京 100101，中国
4 中国科学院大学，北京 100049，中国
† 这些作者对这项工作贡献相同。

###### 摘要

宏基因组分类注释旨在识别环境样本中DNA片段的微生物起源。传统依赖序列相似性的方法常受限于高微生物多样性和参考数据库的不完整性，这促使了诸如Taxometer等学习方法的发展，这些方法通过事后校正来学习更具信息量的宏基因组序列表示。然而，这些方法在训练过程中通常依赖于通过相似性搜索工具生成的标签，这不可避免地引入了噪声，从而损害表示学习并降低分类性能。为解决这一问题，我们提出TaxDistill，一种用于宏基因组分类的知识蒸馏框架。我们引入GenomeOcean，一个拥有5亿参数的基因组基础模型，作为教师网络，用于提取深层语义特征并基于置信度生成软标签。通过将这些软标签信息蒸馏到一个轻量级学生网络中，TaxDistill有效减少了初始检索工具引入的标签噪声。在七个不同的CAMI2数据集上的全面实验表明，TaxDistill在大多数场景下优于现有基线。例如，在胃肠道数据集上，它将MMseqs2的F1得分从0.763提升至0.941，优于Taxometer基线。总体而言，TaxDistill为复杂宏基因组分析中的标签校正提供了一种可靠方法。

TaxDistill: 通过蒸馏基因组基础模型改进宏基因组分类注释

叶荣野1,3,4,†, 李伦1,2,3,†, 罗铮1,3,4, 詹逸然1,3,4, 宋述慧1,2,3,4††感谢:通讯作者。邮箱:[email protected] (https://arxiv.org/html/2605.28868v1/mailto:[email protected])
1 国家基因组科学数据中心，中国国家生物信息中心，北京 100101，中国
2 北京生物大数据智能治理与应用重点实验室，中国国家生物信息中心，北京 100049，中国
3 中国科学院北京基因组研究所，北京 100101，中国
4 中国科学院大学，北京 100049，中国
† 这些作者对这项工作贡献相同。

参见图题图1:宏基因组分析流程与研究动机

## 1 引言

宏基因组测序已成为分析复杂微生物群落的关键技术，本质上是在破译环境样本中隐藏的“生命语言”(Handelsman,2004 (https://arxiv.org/html/2605.28868#bib.bib6); Prabakaran and Bromberg,2025 (https://arxiv.org/html/2605.28868#bib.bib20); Levy Karin and Steinegger,2025 (https://arxiv.org/html/2605.28868#bib.bib12))。在临床病原体检测和疾病微生物组表征中，序列的分类注释是一个关键步骤，其目标是将测序读段或组装的contig精确映射到特定的分类节点上(Simon et al.,2019 (https://arxiv.org/html/2605.28868#bib.bib21); Chiu and Miller,2019 (https://arxiv.org/html/2605.28868#bib.bib4))。目前，主流方法主要依赖序列相似性搜索算法。(Wood et al.,2019 (https://arxiv.org/html/2605.28868#bib.bib27); Kim et al.,2016 (https://arxiv.org/html/2605.28868#bib.bib9); Kallenborn et al.,2025 (https://arxiv.org/html/2605.28868#bib.bib8); Kim and Steinegger,2024 (https://arxiv.org/html/2605.28868#bib.bib10))。这些方法在已充分表征的微生物上表现良好，但在参考数据库中代表性不足的稀有或新物种上，其性能往往会大幅下降(Meyer et al.,2022 (https://arxiv.org/html/2605.28868#bib.bib17))。

近年来，随着深度学习在序列建模方面的突破(Vaswani et al.,2017 (https://arxiv.org/html/2605.28868#bib.bib24); Ye et al.,2025 (https://arxiv.org/html/2605.28868#bib.bib28))，研究人员开始探索用于分类表示的神经网络方法。例如，Taxometer(Kutuzova et al.,2024 (https://arxiv.org/html/2605.28868#bib.bib11))是一种用于宏基因组序列分类的特征聚合方法，它结合了四核苷酸频率(TNFs)和丰度信息。该方法引入深度层次损失(Valmadre,2022 (https://arxiv.org/html/2605.28868#bib.bib23))以对齐分类树，从而将分类信号从标记序列平滑传播到未标记序列。然而，该模型的有效性面临一个关键瓶颈：其训练过程高度依赖于传统序列相似性工具生成的初始伪标签。在处理高度复杂的微生物场景时，此类检索工具常常产生大量错误分类和未分配节点(Meyer et al.,2022 (https://arxiv.org/html/2605.28868#bib.bib17))，不可避免地给后续训练引入严重的标签噪声。由于Taxometer仅采用轻量级的多层感知机(MLP)作为其特征编码器，当面对这些高噪声的硬标签时，受限于其容量和序列建模能力，模型容易过拟合错误信号并陷入表示坍塌(Zhang et al.,2016 (https://arxiv.org/html/2605.28868#bib.bib31); Liu et al.,2020 (https://arxiv.org/html/2605.28868#bib.bib16); Vishwakarma et al.,2025 (https://arxiv.org/html/2605.28868#bib.bib25))，削弱了其维持对噪声伪标签鲁棒性并进行自我校正的能力。基因组语言模型已在生命科学领域展现出广泛的应用潜力(Lin et al.,2023 (https://arxiv.org/html/2605.28868#bib.bib15); Brixi et al.,2026 (https://arxiv.org/html/2605.28868#bib.bib2); Cheng et al.,2024 (https://arxiv.org/html/2605.28868#bib.bib3); Ye et al.,2026 (https://arxiv.org/html/2605.28868#bib.bib29))。其中，最近提出的GenomeOcean是一个拥有40亿参数的生成式语言模型，在超过600 Gbp的大规模、复杂宏基因组组装序列上进行了预训练。与自然语言处理中的大语言模型类似，GenomeOcean(Zhou et al.,2025 (https://arxiv.org/html/2605.28868#bib.bib32))采用高效的字节对编码(BPE)分词策略来构建其基因组词汇表；它不仅能捕捉DNA序列中隐含的语法约束，还能建模复杂的远程依赖关系。

受图1 (https://arxiv.org/html/2605.28868#S0.F1)所示研究目标的驱动，我们引入了TaxDistill，一个新颖的宏基因组分类框架，专门关注contig级别的分类注释。与Taxometer类似，TaxDistill的核心定位是一个事后标签去噪模块，旨在校正初始基于检索的分类器结果。在保留Taxometer高效、轻量级架构作为学生网络的同时，我们引入了具有强大表达能力的GenomeOcean作为教师网络。通过利用从GenomeOcean蒸馏出的高维连续特征和富含暗知识的软标签置信度得分，我们有效地中和了传统序列检索工具引入的硬标签噪声。实验表明，这种知识蒸馏框架赋予了轻量级网络深度语义理解能力，其分类性能在复杂微生物环境的一系列基准测试中持续超越Taxometer基线。总之，我们的主要贡献如下：

1. 我们提出了TaxDistill，一个用于宏基因组分类注释的知识蒸馏框架。TaxDistill采用即插即用设计，可以直接与任何序列比对算法集成。据我们所知，本研究首次将宏基因组语言模型作为教师引入知识蒸馏框架，有效缓解了学生网络过拟合噪声标签的问题。
2. 我们的实验表明，软标签蒸馏有效地赋予了学生网络在宏基因组分类注释中的不确定性感知能力。通过选择性地将模糊边界上的高风险预测转换为未分类标签，TaxDistill实现了严格的假阳性控制，确保了对复杂现实应用的高可靠性。
3. 我们在CAMI2的七个多样化微生物环境数据集上进行了全面的基准测试，评估了多种主流序列分类器(MMseqs2(Kallenborn et al.,2025 (https://arxiv.org/html/2605.28868#bib.bib8)), Metabuli(Kim and Steinegger,2024 (https://arxiv.org/html/2605.28868#bib.bib10)), Kraken2(Wood et al.,2019 (https://arxiv.org/html/2605.28868#bib.bib27)))以及现有校正模型Taxometer(Kutuzova et al.,2024 (https://arxiv.org/html/2605.28868#bib.bib11))。实验结果表明，TaxDistill在大多数场景下表现优于基线模型。

## 2 相关工作

### 2.1 宏基因组序列分类

传统的宏基因组分类注释方法主要基于序列相似性和启发式匹配，广泛使用的工具包括Kraken2(Wood et al.,2019 (https://arxiv.org/html/2605.28868#bib.bib27)), Centrifuge(Kim et al.,2016 (https://arxiv.org/html/2605.28868#bib.bib9)), MMseqs2(Kallenborn et al.,2025 (https://arxiv.org/html/2605.28868#bib.bib8))以及最近提出的Metabuli(Kim and Steinegger,2024 (https://arxiv.org/html/2605.28868#bib.bib10))。尽管这些字典式检索方法计算效率极高，但在面对高度复杂的环境宏基因组样本或参考数据库中不存在的新型微生物时，往往会产生大量错误标签或模糊的分类预测。

为了克服序列比对方法的固有局限性，近年来研究人员开始引入深度学习架构来从基因序列中提取连续空间模式。值得注意的例子包括基于双向长短期记忆(Bi-LSTM)单元的DeepMicrobes(Liang et al.,2020 (https://arxiv.org/html/2605.28868#bib.bib14))，以及利用自注意力机制的MetaTransformer(Wichmann et al.,2023 (https://arxiv.org/html/2605.28868#bib.bib26))。尽管这些端到端序列分类模型在标准基准上表现出色，但其训练严重依赖于人工模拟的短读段和真实参考标签。这种植根于固定标签集的训练范式本质上是一种闭集学习。然而，现实世界的宏基因组环境极其复杂，充满了未被探索的微生物群落(Nayfach et al.,2021 (https://arxiv.org/html/2605.28868#bib.bib19); Thompson et al.,2017 (https://arxiv.org/html/2605.28868#bib.bib22))。理想化的固定标签集与现实环境中的微生物多样性之间的差异，常常导致这些模型在应用于真实环境数据时经历严重的域偏移，泛化性能显著下降。

为了解决传统固定标签学习范式的局限性，最近提出的Taxometer建立了一种新颖的事后标签校正方法。该模型根据初始分类器(例如MMseqs2)的检索结果，动态构建一个针对当前数据集的本地标签树。然而，其轻量级架构仍然极易过拟合初始检索工具引入的标签噪声。

### 2.2 知识蒸馏

知识蒸馏(KD)旨在将复杂教师模型的丰富表示能力转移给具有更紧凑结构的学生模型(Hinton et al.,2015 (https://arxiv.org/html/2605.28868#bib.bib7); Gou et al.,2021 (https://arxiv.org/html/2605.28868#bib.bib5))。近年来，大量研究表明KD在缓解从噪声标签学习的挑战方面具有显著优势(Li et al.,2017 (https://arxiv.org/html/2605.28868#bib.bib13); Müller et al.,2019 (https://arxiv.org/html/2605.28868#bib.bib18))。当面对包含大量错误的伪标签时，传统的硬标签常常导致轻量级网络过拟合。相比之下，教师网络生成的软标签捕捉到的信息揭示了不同类别之间的潜在相似性(Yuan et al.,2020 (https://arxiv.org/html/2605.28868#bib.bib30))。这种连续概率分布作为一种天然的正则化器(Yuan et al.,2020 (https://arxiv.org/html/2605.28868#bib.bib30); Ben-Baruch et al.,2024 (https://arxiv.org/html/2605.28868#bib.bib1))，有效地防止学生网络盲目地过拟合错误标注。

尽管参数化基因组基础模型近期已经出现，但如何将其深层语义知识高效蒸馏到轻量级网络中以用于宏基因组序列分类，仍未得到探索。提出的TaxDistill框架正是为了填补这一空白。

## 3 方法

参见图题图2:所提出的TaxDistill框架的整体架构。它由三个核心模块组成：多模态数据输入公式化、教师模型分支和学生模型分支。

在本节中，我们正式介绍所提出的TaxDistill框架。如图2 (https://arxiv.org/html/2605.28868#S3.F2)所示，该框架设计用于通过软标签蒸馏实现可靠的宏基因组分类。本研究创新性地提出了一个基于宏基因组语言模型的知识蒸馏框架，并将其应用于宏基因组分类注释任务。

教师分支采用预训练的GenomeOcean基础模型，并冻结其骨干网络。它从原始序列中提取深层语义特征，并通过一个可学习的分类头进行投影，输出类别概率分布。该分支通过深度层次损失独立优化。

相反，学生分支保持轻量级MLP架构，以确保低延迟推理。它处理一个由手工设计的TNF特征、跨K个环境的丰度以及总丰度组成的(103+K+1)维特征向量。

在联合优化阶段，引入KD损失来衡量教师分布与学生分布之间的差异。学生模型的参数由其自身的层次分类损失和KD损失共同更新。同时，教师模型仅由其自身的分类损失继续更新。详细的数学公式在3.1节 (https://arxiv.org/html/2605.28868#S3.SS1)中阐述。

### 3.1 问题形式化与符号说明