LOCALE: Local-Alignment Embeddings for Noise-Robust DNA Search at Scale

Abstract

Searching petabase-scale repositories of raw sequencing data such as the NIH Sequence Read Archive (SRA) could unlock biological insights that remain inaccessible. Existing sequence-search methods either do not scale well or rely on exact $k$-mer matching that is brittle to genomic variation, which is high both within and across species and is further compounded by sequencing error. We recast sequence search as dense retrieval: we learn sequence embeddings whose inner-product similarity ranks locally aligned sequences above unaligned ones. We introduce LOCALE, an encoder trained with an InfoNCE objective on biologically informed augmentations: overlapping crops of parent sequences corrupted with substitutions, insertions, and deletions. This training recipe, rather than a specialized architecture, consistently induces variation-robust representations across four encoders spanning three tokenization schemes, improving Recall@$R_q$ by 7--39 points at 10% genomic variation. On a 50-accession SRA benchmark LOCALE retains 63.2% Recall@$R_q$ at 10% variation while existing k-mer indexes and DNA embedding methods fall below 35%, and on the full 4,571-accession MetaGraph benchmark (296 Gbp, 2.06 billion vectors) it achieves 4.5x the AUPRC of MetaGraph, the state-of-the-art k-mer index, with a single-node GPU IVF-PQ index retaining a 3.6x advantage at MetaGraph's query latency. Our encoder and training strategy generalize to real biological variance: on a cross-genotype Hepatitis B task, querying reads from one genotype against accessions of another with >8% natural divergence, LOCALE reaches 0.813 Recall@R_q versus 0.518 for MetaGraph.