DNACHUNKER: Learnable Tokenization for DNA Language Models

TL;DR

This work tackles the tokenization bottleneck in DNA language modeling by introducing DNACHUNKER, a bidirectional masked DNA language model that learns dynamic, variable-length chunks through a two-stage hierarchical encoder and cross-attention-based dechunking. The approach yields state-of-the-art performance on two public benchmarks (Nucleotide Transformer and Genomic) with only 156M parameters, substantially fewer than prior baselines, while demonstrating robustness to mutations and interpretable tokenization that preserves functional regions at high resolution. Ablation studies show that the cross-attention dechunking and special-token handling are key architectural factors driving improvements, and the learned tokenizer biologically reallocates granularity to promoters, exons, and other functional elements while compressing repetitive regions. Overall, the method advances genomic language modeling by integrating learnable tokenization with scalable, bidirectional context, enabling more efficient and biologically faithful modeling of long-range DNA dependencies.

Abstract

DNA language models have emerged as powerful tools for decoding the complex language of DNA sequences. However, the performance of these models is heavily affected by their tokenization strategy, i.e., a method used to parse DNA sequences into a shorter sequence of chunks. In this work, we propose DNACHUNKER, which integrates a learnable dynamic DNA tokenization mechanism and is trained as a masked language model. Adopting the dynamic chunking procedure proposed by H-Net, our model learns to segment sequences into variable-length chunks. This dynamic chunking offers two key advantages: it's resilient to shifts and mutations in the DNA, and it allocates more detail to important functional areas. We demonstrate the performance of DNACHUNKER by training it on the human reference genome (HG38) and testing it on the Nucleotide Transformer and Genomic benchmarks. Further ablative experiments reveal that DNACHUNKER learns tokenization that grasps biological grammar and uses smaller chunks to preserve detail in important functional elements such as promoters and exons, while using larger chunks for repetitive, redundant regions.

Figures (4)

• Figure 1: Architecture, tokenizer robustness, and distribution of chunk size. (a) The architecture of DNAChunker. (b) Our tokenizer is robust against nucleotide-wise shifts or mutations, where we color the tokens to indicate that they are preserved despite the mutations. (c) Our DNAChunker dynamically represents functional elements (promoter, intron, exon) with high-resolution using smaller chunks, while compressing the non-functional repeatitive elements with larger chunks.
• Figure 2: Comparison with H-Net upon NT Benchmark.DNAChunker compared with H-Net trained with masked language modeling. Results are averaged over 10-fold cross validations with error bars.
• Figure 3: Token size distributions of BPE and DNAChunker. The BPE tokenizer (right) is compared against the two-stage DNAChunker tokenizer (left) on the H3C13, H3C14, and H3C15 genes. The plots visualize the average token size of BPE and the Stage 1 and Stage 2 token sizes of DNAChunker. Key genomic features are included as a reference, like gene bodies (shaded regions), conserved elements (PHYLOP), and SINE repeats.
• Figure 4: Token size distributions of BPE and DNAChunker. A comparison of the BPE tokenizer against our two-stage DNAChunker tokenizer on human chromosomes 1, 3, 4, 5, and 7. The plots visualize the distribution of chunk sizes for BPE and both the stage 1 and stage 2 outputs of DNAChunker. The distributions are categorized by key genomic features including Promoter, Intron, Exon, and Repeat.