HAD: Hybrid Architecture Distillation Outperforms Teacher in Genomic Sequence Modeling

📄 arXiv: 2505.20836v1 📥 PDF

作者: Hexiong Yang, Mingrui Chen, Huaibo Huang, Junxian Duan, Jie Cao, Zhen Zhou, Ran He

分类: cs.LG, q-bio.GN

发布日期: 2025-05-27

💡 一句话要点

提出混合架构蒸馏(HAD)，提升基因组序列建模中小模型性能，超越大模型教师。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 基因组序列建模 知识蒸馏 混合架构 自监督学习 掩码语言模型

📋 核心要点

  1. 现有基因组序列建模方法依赖大规模模型和数据，计算成本高昂，小模型性能不足。
  2. 提出混合架构蒸馏(HAD)，结合蒸馏和重构任务，提升小模型预训练效率和效果。
  3. 实验表明，HAD模型性能优异，甚至超越了参数量大500倍的教师模型，理解基因组序列能力更强。

📝 摘要（中文）

受自然语言领域掩码语言模型(MLM)成功的启发，自监督预训练和微调范式在DNA序列建模领域也取得了显著进展。然而，先前的方法通常依赖于海量的预训练数据或具有巨大参数的大规模基础模型，造成了巨大的计算负担。为了解决这个问题，许多工作试图使用更紧凑的模型来实现类似的结果，但仍然相差甚远。本文提出了一种混合架构蒸馏(HAD)方法，利用蒸馏和重构任务进行更有效和高效的预训练。具体来说，我们采用NTv2-500M作为教师模型，并设计了一种分组掩码策略，以对齐可见token的特征嵌入，同时在MLM预训练期间重构不可见token。为了验证我们提出的方法的有效性，我们在Nucleotide Transformer Benchmark和Genomic Benchmark上进行了全面的实验。与具有相似参数的模型相比，我们的模型取得了优异的性能。更令人惊讶的是，它甚至在某些子任务上超过了蒸馏上限——教师模型，该模型比学生模型大500多倍。最后，我们利用t-SNE进行更直观的可视化，这表明我们的模型可以获得对基因组序列中内在表示模式的深刻理解。

🔬 方法详解

问题定义：基因组序列建模领域，现有方法依赖于大规模预训练数据和大型模型，导致计算成本高昂。虽然一些工作尝试使用更紧凑的模型，但性能与大型模型相比仍有显著差距。因此，需要一种方法，能够在计算资源有限的情况下，训练出高性能的基因组序列建模模型。

核心思路：本文的核心思路是利用知识蒸馏技术，将大型教师模型的知识迁移到小型学生模型中。同时，结合掩码语言模型(MLM)的重构任务，使学生模型不仅学习教师模型的输出，还学习基因组序列的内在表示。通过这种混合的方式，可以更有效地训练小模型，使其在性能上接近甚至超越大型教师模型。

技术框架：HAD方法的整体框架包括以下几个主要步骤：1) 使用大型预训练模型（NTv2-500M）作为教师模型。2) 设计分组掩码策略，对输入基因组序列进行掩码。3) 学生模型同时进行两个任务：一是模仿教师模型对可见token的特征嵌入；二是重构被掩码的token。4) 使用混合损失函数，结合蒸馏损失和重构损失，优化学生模型。

关键创新：HAD方法的关键创新在于混合架构蒸馏，它同时利用了蒸馏和重构任务。传统的知识蒸馏主要关注输出层的知识迁移，而HAD方法通过对齐中间层的特征嵌入，使学生模型更好地学习教师模型的表示能力。同时，MLM的重构任务可以增强学生模型对基因组序列内在结构的理解。这种混合的方式可以更有效地利用教师模型的知识，并提高学生模型的泛化能力。

关键设计：HAD方法的关键设计包括：1) 分组掩码策略：将基因组序列分成多个组，并随机掩码其中的一些组，以增加模型的鲁棒性。2) 混合损失函数：结合蒸馏损失（例如，MSE损失）和重构损失（例如，交叉熵损失），以平衡知识迁移和序列重构。3) 学生模型的架构：可以使用Transformer或其他适合基因组序列建模的架构，但参数量要远小于教师模型。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，HAD方法在Nucleotide Transformer Benchmark和Genomic Benchmark上取得了优异的性能。与参数量相似的模型相比，HAD模型性能显著提升。更重要的是，在某些子任务上，HAD模型甚至超越了参数量大500倍的教师模型。t-SNE可视化结果表明，HAD模型能够学习到基因组序列中更复杂的表示模式。

🎯 应用场景

该研究成果可应用于基因组序列分析、疾病预测、药物发现等领域。通过训练更小、更高效的基因组序列模型，可以降低计算成本，加速相关研究的进展。此外，该方法还可以推广到其他序列建模任务中，例如蛋白质序列分析、RNA序列分析等，具有广泛的应用前景。

📄 摘要（原文）

Inspired by the great success of Masked Language Modeling (MLM) in the natural language domain, the paradigm of self-supervised pre-training and fine-tuning has also achieved remarkable progress in the field of DNA sequence modeling. However, previous methods often relied on massive pre-training data or large-scale base models with huge parameters, imposing a significant computational burden. To address this, many works attempted to use more compact models to achieve similar outcomes but still fell short by a considerable margin. In this work, we propose a Hybrid Architecture Distillation (HAD) approach, leveraging both distillation and reconstruction tasks for more efficient and effective pre-training. Specifically, we employ the NTv2-500M as the teacher model and devise a grouping masking strategy to align the feature embeddings of visible tokens while concurrently reconstructing the invisible tokens during MLM pre-training. To validate the effectiveness of our proposed method, we conducted comprehensive experiments on the Nucleotide Transformer Benchmark and Genomic Benchmark. Compared to models with similar parameters, our model achieved excellent performance. More surprisingly, it even surpassed the distillation ceiling-teacher model on some sub-tasks, which is more than 500 $\times$ larger. Lastly, we utilize t-SNE for more intuitive visualization, which shows that our model can gain a sophisticated understanding of the intrinsic representation pattern in genomic sequences.