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

TL;DR

VQDNA introduces a novel vector quantization-based framework for genome tokenization, enabling adaptive, pattern-aware embeddings that improve genome modeling and reveal biologically significant mutation patterns.

Contribution

It proposes VQDNA with hierarchical residual quantization for improved genome vocabulary learning and demonstrates its effectiveness across multiple datasets.

Findings

Outperforms existing genome language models in accuracy and efficiency

Reveals biologically meaningful mutation patterns in SARS-CoV-2

Enriches genome vocabulary with hierarchical codebooks

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.