Species-aware DNA language models capture regulatory elements and their evolution

Abstract

The rise of large-scale multi-species genome sequencing projects promises to shed new light on how genomes encode gene regulatory instructions. To this end, new algorithms are needed that can leverage conservation to capture regulatory elements while accounting for their evolution. Here we introduce species-aware DNA language models (LMs), which we trained on more than 800 species spanning over 500 million years of evolution. Investigating their ability to predict masked nucleotides from context, we show that DNA LMs distinguish transcription factor and RNA-binding protein motifs from background non-coding sequence. Owing to their flexibility, DNA LMs capture conserved regulatory elements over much further evolutionary distances than sequence alignment would allow. Remarkably, DNA LMs reconstruct motif instances bound in vivo better than unbound ones and account for the evolution of motif sequences and their positional constraints, showing that these models capture functional high-order sequence and evolutionary context. We further show that species-aware training yields improved sequence representations for endogenous and MPRA-based gene expression prediction, as well as motif discovery. Collectively, these results demonstrate that species-aware DNA language models are a powerful, flexible, and scalable tool to integrate information from large compendia of highly diverged genomes.