HAD: Hybrid Architecture Distillation Outperforms Teacher in Genomic Sequence Modeling

TL;DR

The paper introduces a Hybrid Architecture Distillation (HAD) method that uses distillation and reconstruction tasks to pre-train genomic sequence models more efficiently, surpassing larger teacher models in performance.

Contribution

The novel HAD approach combines distillation and reconstruction for genomic modeling, enabling smaller models to outperform larger teachers.

Findings

HAD achieves state-of-the-art results on genomic benchmarks.

HAD surpasses the performance of a 500x larger teacher model on some tasks.

Visualization shows HAD captures intrinsic genomic sequence patterns.

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.