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

Siyuan Li,Zedong Wang,Zicheng Liu,Di Wu,Cheng Tan,Jiangbin Zheng,Yufei Huang,Stan Z. Li

2024-06-03

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.

Genomics,Artificial Intelligence

What problem does this paper attempt to address?

This paper presents a framework called VQDNA, which aims to improve the genome tokenization method in multi-species genome sequence modeling. Current genome language models often use manually designed tokenization strategies, but these strategies may fail to capture the most discriminative patterns in the limited vocabulary of genome data. VQDNA learns a vector quantization codebook as a trainable genome vocabulary to adaptively transform the genome into a pattern-aware embedding representation, achieving end-to-end tokenization. To further enhance its performance, the paper proposes hierarchical residual quantization (HRQ) to enrich the genome vocabulary in a coarse-to-fine manner by designing codebooks at different scales.

Experiments on multiple genome datasets demonstrate that VQDNA outperforms existing genome language models in terms of parameter efficiency and performance. Specifically, VQDNA performs well in various downstream tasks, including SARS-CoV-2 variant classification, and its HRQ vocabulary shows the ability to perceive subtle patterns with biological significance, revealing its potential in broader biological applications. Therefore, VQDNA is a versatile framework for multi-species genome sequence modeling, capable of handling large-scale and diverse genome analysis tasks.