Learning Strand-Equivariant Tokenizations of DNA with MirrorDNA

DNA comprises two strands, and regulatory label prediction tasks generally expect a sequence and its reverse complement to yield the same result. However, DNA language models are typically trained and evaluated on a single strand at a time, leading to inconsistency. BPE and $k$-mer tokenizers split between 92-100% of transcription-factor motifs differently on the two strands, while NT-v2, DNABERT-2, and HyenaDNA change 25-31% of their predictions when the test sequences are reverse-complemented.

To address this issue, we introduce MirrorDNA, a compact DNA language model equipped with a learned tokenizer that ensures mirror-symmetric token boundaries across reverse-complement sequences, irrespective of its learned parameters.

We demonstrate that this symmetry is maintained throughout the entire model. Consequently, the same model can yield identical answers on both strands when the label is independent of the strand, and it can still distinguish between strands when necessary – a capability that no strand-blind model can achieve on its own. By considering only predictions that are correct on both strands, MirrorDNA more than doubles the Matthews Correlation Coefficient (MCC) of NT-v2, DNABERT-2, and HyenaDNA.

With only 2.3 million parameters, MirrorDNA outperforms Caduceus models trained on the same data and is not significantly different from NT-v2 (98 million parameters) on Genomic Benchmarks or from DNABERT-2 (117 million parameters) on Nucleotide Transformer tasks. Additionally, the model remains accurate on inputs that are sixteen times longer than those it was trained on.

If trained with reverse-complement augmentation, MirrorDNA changes approximately one prediction in ten, unless its outputs are averaged over both strands. Remarkably, the tokenizer learns to treat the splice-donor signal AGGT as a single token without any additional annotation. The MirrorDNA code is available at https://github.com/NIRJHOR-DATTA/MirrorDNA.