Technical Summary: DNAChunker: Learnable Tokenization for DNA Language Models

1. Problem Statement

DNA language models (LMs) have emerged as powerful tools for representing genomic sequences, yet their efficacy is critically dependent on how raw nucleotides are converted into model inputs. Unlike natural language, DNA lacks canonical boundaries (e.g., spaces or punctuation), making the choice of tokenization strategy a significant bottleneck.

Current approaches typically rely on fixed tokenization schemes:

• Single Nucleotides: While simple, this results in excessively long sequences, making long-range interaction modeling computationally expensive.
• Fixed-size k-mers: These are brittle; a single insertion, deletion, or shift can completely alter the tokenized output, even if the biological function remains unchanged.
• Byte-Pair Encoding (BPE): While frequency-driven, BPE often fails to capture functional importance. It tends to fragment biologically significant motifs (e.g., transcription factor binding sites) because the most frequent substrings are often non-functional repetitive elements.

These fixed schemes are inherently blind to genomic context, struggling to adapt to biological shifts, indels, and local repeats. The paper posits that a learnable, context-dependent segmentation strategy is required to group DNA sequences into meaningful, variable-length units that reflect biological organization.

2. Methodology: DNACHUNKER

The authors introduce DNACHUNKER, a bidirectional masked DNA language model (MLM) that integrates a learnable, adaptive segmentation module. The architecture is designed to compress sequences into context-dependent chunks while retaining nucleotide-level fidelity for masked prediction.

2.1. Architecture Overview

The model follows a hierarchical encoder-decoder structure with a central main network:

1. Encoder (Stage 1 & 2):
   • Input: Raw base-pair embeddings are processed by lightweight bidirectional Mamba layers (Schiff et al., 2024).
   • Segmentation: A routing network predicts boundaries between adjacent positions based on cosine dissimilarity between projected query and key vectors.
   • Compression: Adjacent positions are merged into chunks if they are highly similar, producing a shorter, coarser-grained sequence.
   • Mask Protection: A critical mechanism ensures that [MASK] tokens (functional artifacts of pretraining) are never merged with neighboring bases. Boundaries are enforced immediately before and after masked positions to prevent the model from learning segmentation patterns based on mask placement rather than genomic context.
2. Main Network:
   • The compressed segment sequence is processed by a standard Transformer (30 blocks) with Rotary Position Embeddings (RoPE). This network serves as the primary locus for long-range contextual reasoning.
3. Decoder (Hierarchical Dechunking):
   • The decoder reconstructs base-pair resolution in two stages, mirroring the encoder.
   • Bidirectional Smoothing: Unlike autoregressive decoders, this module leverages context from both directions. It applies probability-gated smoothing to the piecewise-constant representations generated by pasting segment embeddings.
   • Masked Residual Gating: Residual connections from the encoder are gated based on whether a position contains a mask. If a segment contains a mask, the residual is zeroed out, ensuring the reconstruction relies on the main network's learned representations rather than leaked encoder context.

2.2. Training Objectives

• Masked Language Modeling (MLM): The model is pretrained on the human reference genome (GRCh38/hg38).
• Repeat Down-weighting: To prevent the model from memorizing repetitive elements (which are easy to predict but biologically less informative), loss weights are applied: repetitive regions are down-weighted by 0.1, while non-repetitive regions retain a weight of 1.0.
• Ratio Loss: An auxiliary loss is used to control the compression ratio of the chunking layers, guiding the model toward a target retention ratio (0.33 per stage).

3. Key Contributions

1. Bidirectional Dynamic Tokenization: The paper adapts dynamic tokenization from autoregressive settings to the masked DNA pretraining setting. This enables the model to learn token boundaries bidirectionally, a capability previously unexplored in masked genomic models.
2. State-of-the-Art Performance: DNACHUNKER achieves competitive or state-of-the-art results across five diverse benchmarks (Nucleotide Transformer, Genomic Benchmarks, BEND, DNALongBench) using only the human reference genome for pretraining, despite having significantly fewer parameters (172M) than many baselines (e.g., Nucleotide Transformer's 2.5B).
3. Biologically Meaningful and Robust Tokens:
   • Motif Preservation: The learned tokenizer treats functional motifs (e.g., TF-binding sites) as single coherent chunks, whereas fixed tokenizers (BPE, k-mers) frequently fragment them.
   • Mutation Resilience: The segmentation remains stable under single-nucleotide variants (SNVs), insertions/deletions (InDels), and structural variants (SVs), unlike fixed schemes where a single shift cascades into a complete tokenization change.
   • Adaptive Granularity: The model automatically allocates finer granularity (shorter chunks) to functionally enriched regions (exons, promoters) and coarser granularity (longer chunks) to repetitive sequences (SINEs), aligning with biological information density.

4. Experimental Results

The model was evaluated on five benchmark suites:

• Nucleotide Transformer (NT) Benchmark: DNACHUNKER achieved the best total average Matthews Correlation Coefficient (MCC) of 0.772 and the best average rank (1.67) across 18 datasets, outperforming larger models like NT-multi (2.5B parameters) and GENERator (1.2B parameters).
• Revised NT Benchmark: It surpassed learnable tokenization baselines (MxDNA, PatchDNA), particularly in splice-site annotation, demonstrating that adaptive segmentation can maintain base-level sensitivity where needed.
• Genomic Benchmarks: It achieved the second-best average rank, performing on par with GENERator while using 7x fewer parameters and training only on the human genome.
• DNALongBench: DNACHUNKER surpassed the strong baseline Caduceus-PH on all five long-range tasks (up to 1 Mb context), including enhancer-target gene interactions and eQTL prediction, even when using a frozen backbone with linear probing.
• BEND Benchmark: It achieved the best average rank (1.9) across seven tasks, including chromatin accessibility and variant effect prediction.

Efficiency Analysis: Despite the added complexity of a learnable segmentation module, DNACHUNKER is more compute-efficient than fixed tokenization schemes for long contexts. At 16k-bp inputs, it requires significantly fewer FLOPs than BPE or single-nucleotide tokenization because the adaptive compression aggressively reduces the effective sequence length.

5. Significance and Claims

The paper claims that DNACHUNKER offers a more effective and efficient approach to genomic sequence modeling by treating tokenization as a learnable, data-driven process rather than a fixed heuristic.

• Biological Alignment: The learned segments reflect biological organization, preserving functional motifs as singular units and adapting chunk length to information density. This contrasts with fixed schemes that are "brittle" under shifts and mutations.
• Efficiency: By compressing repetitive regions and expanding functional regions, the model improves long-context efficiency without sacrificing nucleotide-level resolution.
• Generalization: The results suggest that end-to-end, data-driven segmentation can yield genomic language models that better reflect biological reality, achieving high performance without the need for massive multi-species pretraining or billion-parameter scales.

The authors conclude that their work demonstrates the viability of integrating fully learnable, dynamic chunking into masked DNA language models, jointly learning bidirectional representations and biologically meaningful token boundaries.