HAD: Hybrid Architecture Distillation Outperforms Teacher in Genomic Sequence Modeling
Hexiong Yang, Mingrui Chen, Huaibo Huang, Junxian Duan, Jie Cao, Zhen Zhou, Ran He
TL;DR
The paper tackles the high computational cost of large pretrained models for genomic sequence modeling by introducing Hybrid Architecture Distillation (HAD), a dual-branch pretraining framework that blends high-level feature distillation with low-level masked nucleotide reconstruction. A compact $1.1$M-parameter student leverages a bidirectional Gated Delta Net (GDN) backbone augmented with a self-attention layer, and learns from a large NTv2-500M teacher through a structured two-stage masking strategy and a cross-attention reconstruction pathway. Empirical results on the Nucleotide Transformer Benchmark and Genomic Benchmark show HAD not only matches but often exceeds the performance of its substantially larger teacher, with an average Genomic Benchmark score of $0.875$ and leading performance on many NT tasks, underscoring the effectiveness of the hybrid distillation approach. The work demonstrates that intelligent integration of distillation and reconstruction, together with a carefully designed tokenizer-consistent masking scheme, can yield highly capable, efficient models for complex genomic sequence analysis, with practical implications for scalable genomic interpretation.
Abstract
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.
