VQDNA: Unleashing the Power of Vector Quantization for Multi-Species Genomic Sequence Modeling

TL;DR

VQDNA tackles the limitation of hand-crafted genome tokenization by learning a discriminative genome vocabulary with a vector-quantized codebook (VQ-VAE) and augments it with Hierarchical Residual Quantization (HRQ) to create coarse-to-fine token representations. The approach follows a three-stage pipeline: learn the vocabulary (Stage 1), pre-train a Transformer encoder with masked modeling on tokenized genomes (Stage 2), and fine-tune on diverse downstream tasks (Stage 3). Across 32 genome datasets, VQDNA and especially HRQ achieve state-of-the-art or competitive results with fewer parameters, and analyses on SARS-CoV-2 demonstrate biologically meaningful, fine-grained pattern awareness in the learned vocabulary. The work advances cross-species genomic sequence modeling by providing a data-driven, scalable genome vocabulary that can generalize to broad genomics tasks and potentially enable generation-style applications in the future.

Abstract

Similar to natural language models, pre-trained genome language models are proposed to capture the underlying intricacies within genomes with unsupervised sequence modeling. They have become essential tools for researchers and practitioners in biology. However, the hand-crafted tokenization policies used in these models may not encode the most discriminative patterns from the limited vocabulary of genomic data. In this paper, we introduce VQDNA, a general-purpose framework that renovates genome tokenization from the perspective of genome vocabulary learning. By leveraging vector-quantized codebooks as learnable vocabulary, VQDNA can adaptively tokenize genomes into pattern-aware embeddings in an end-to-end manner. To further push its limits, we propose Hierarchical Residual Quantization (HRQ), where varying scales of codebooks are designed in a hierarchy to enrich the genome vocabulary in a coarse-to-fine manner. Extensive experiments on 32 genome datasets demonstrate VQDNA's superiority and favorable parameter efficiency compared to existing genome language models. Notably, empirical analysis of SARS-CoV-2 mutations reveals the fine-grained pattern awareness and biological significance of learned HRQ vocabulary, highlighting its untapped potential for broader applications in genomics.

Figures (6)

• Figure 1: Performance of fine-tuned VQDNA and other genome language models across downstream tasks on 32 datasets, including Epigenetic Mark Prediction (EMP) for Yeast, Transcription Factor Prediction on mouse and human genome (TFP-M and TFP-H), Covid Variants Classification (CVC), Promoter Detection (PD), Core Promoter Detection (CPD), Splice Site Prediction (SSP), and Editing Efficiency Prediction (EEP). The circle size indicates the parameter scale of each model. Notably, NT-2500M-1000g is with 2537M model parameters, while our VQDNA has only 110M.
• Figure 2: An overview of our three-stage training pipeline of VQDNA. (a) VQ genome vocabulary learning with large-scale multi-species genome sequences. (b) Masked modeling pre-training of the Transformer encoder with frozen genome vocabulary. (c) Fine-tuning the pre-trained encoder with an MLP head for various downstream genome analysis tasks.
• Figure 3: Illustration of our Hierarchical Residual Quantization (HRQ) as genome word embedding for VQDNA framework. We instantiate HRQ with a 6-layer encoder and decoder with two hierarchical codebooks after the output of $3$-th and $6$-th layers in practice.
• Figure 4: Illustration of RQ and our HRQ in a two-dimensional space with a two-layer quantization case. We use purple for the current hierarchical input $H^{(n)}$(or the residual in RQ),green for the second layer encoder output $Z^{(2)}$,orange for the input $Z^{(1)}$ and output hierarchical embeddings $\hat{H}^{(n)}$ per layer, and pale orchid for the ultimate embeddings $\hat{Z}^{(n)}$ after $n$-layer quantization.
• Figure 5: Visualization of the HRQ codebooks on CVC dataset by UMAP 2018UMAP. The label of each code is obtained by calculating the most relevant class with Grad-CAM cvpr2017gradcam of the linear classifier learned upon HRQ-tokenized sequences. The pentagon dots stand for codes of the layer-3 codebook, while the pale circle is the layer-6 ones. The result shows great intra- & inter-lineage pattern-awareness of HRQ vocabulary.