Table of Contents
Fetching ...

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.

HAD: Hybrid Architecture Distillation Outperforms Teacher in Genomic Sequence Modeling

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 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 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 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.

Paper Structure

This paper contains 23 sections, 4 equations, 5 figures, 4 tables.

Figures (5)

  • Figure 1: Proposed Hybrid Architecture Distillation ($\textbf{HAD}$) pre-training pipeline. The upper branch performs feature alignment on visible nucleotides, distilling high-level knowledge from a pre-trained teacher model to the student model. The lower branch focuses on the low-level reconstruction of masked nucleotides, leveraging contextual information from the student's visible nucleotide representations.
  • Figure 2: Two-stage masking strategy in HAD. This strategy is designed to prevent information leakage and enhance feature learning for the student model during distillation.
  • Figure 3: Hybrid architecture of our student model, combining a bidirectional Gated Delta Net (GDN) backbone with a self-attention layer for efficient sequential processing and global context integration within a compact 1.1M parameter budget. Within the GDN, $\alpha$ serves as a data-dependent gate controlling memory erasure, while $\beta$ acts as the update strength from the delta rule.
  • Figure 4: Ablation on pretraining scheme with different model architectures (left) and teacher model size (right)
  • Figure 5: t-SNE visualization of pre-trained model representations on select downstream task data. The clustering visually demonstrates HAD's effective knowledge transfer from its teacher NTv2 through the distillation branch, particularly for distinguishing Enhancer-related features.