Embed-Search-Align: DNA Sequence Alignment using Transformer Models

TL;DR

The paper tackles DNA sequence alignment by reframing it as an Embed-Search-Align problem using Reference-Free DNA Embedding (RDE) to generate embeddings for reads and reference fragments. Alignment is performed via near-neighbor search in a shared embedding space, enabled by a contrastive self-supervised training objective and a scalable vector-store index, achieving near-traditional-aligner performance (e.g., ~99% recall for 250 bp reads on the human genome). The work demonstrates cross-chromosome and cross-species generalization, suggesting embeddings capture structural DNA patterns beyond training data and enabling potential pan-genome and de novo assembly applications. Future directions include speedups, specialized training for shorter reads, broader species coverage, and extending the framework to assembly tasks and pan-genomic representations.

Abstract

DNA sequence alignment involves assigning short DNA reads to the most probable locations on an extensive reference genome. This process is crucial for various genomic analyses, including variant calling, transcriptomics, and epigenomics. Conventional methods, refined over decades, tackle this challenge in 2 steps: genome indexing followed by efficient search to locate likely positions for given reads. Building on the success of Large Language Models in encoding text into embeddings, where the distance metric captures semantic similarity, recent efforts have explored whether the same Transformer architecture can produce embeddings for DNA sequences. Such models have shown early promise in classifying short DNA sequences, such as detecting coding/non-coding regions, and enhancer, promoter sequences. However, performance at sequence classification tasks does not translate to sequence alignment, where it is necessary to search across the genome to align each read, a significantly longer-range task. We bridge this gap by framing the Sequence Alignment task for Transformer models as an "Embed-Search-Align" task. In this framework, a novel Reference-Free DNA Embedding model generates embeddings of reads and reference fragments, which are projected into a shared vector space where the read-fragment distance is used as a surrogate for alignment. Technical contributions include: (1) Contrastive loss for self-supervised training of DNA sequence representations, facilitating rich reference-free, sequence-level embeddings, and (2) a DNA vector store to enable search across fragments on a global scale. DNA-ESA is 99% accurate when aligning 250-length reads onto a human genome (3gb), rivaling conventional methods such as Bowtie and BWA-Mem. DNA-ESA exceeds the performance of 6 Transformer model baselines such as Nucleotide Transformer, Hyena-DNA, and shows task transfer across chromosomes and species.

Figures (2)

• Figure 5: RDE convergence plots. Both plots show the loss pr. step. For clarity, we smooth the loss using a moving average
• Figure 6: Motivating the task of de-novo assembly using the distance adjacency matrix of the UMAP-generated $k$-nearest neighbor ($k=10$) network of reads from Thermus Aquaticus: We assume that reads $r_1, r_2, \dots, r_C$ belonging to a species are embedded into the RDE embedding space $h(r_1), h(r_2), \dots, h(r_C)$ using a model trained on a known reference genome belonging to a different species. For this illustrative numerical example we ensure that reads are created with the following properties: first, no read is a substring of any other read (efficient algorithms for the case where reads can be substrings of others will be covered in our future work); second every read has a corresponding read that it overlaps with for each end, and finally, the union of the reads cover the ground-truth reference genome. Thus, for every read $r_i$ there exist reads $r_k$ and $r_j$ such that the ordered triplet $(r_j, r_i, r_k)$ constitutes a fragment of the reference. In a perfect embedding scenario, the closest neighbors of $r_i$ (i.e. $k=1$ in an $k$-NN search) in the embedding space should be $r_j$ ad $r_k$. For this ideal case, if a network is constructed where each read (node) is connected to its two nearest neighbors, and the nodes are arranged based on the order in which they appear in the reference chromosome, one would get a bi-diagonal adjacency matrix. When ground truth is not known, the adjacency matrix will be a permuted version of a bi-diagonal matrix, but a complete assembly can still be performed by a greedy algorithm that starts with any node and pieces together the neighboring reads in the embedded space in time that grows linearly in the number of reads. Simulated reads, each of length $500$ and with random overlaps of lengths drawn uniformly $\mathcal{U}([150,250])$ with its neighboring reads at each end, are sampled from $4$ of the chromosomes in Thermus Aquaticus, and are embedded into the RDE embedding space trained with a reference human genome. The sorted adjacency matrix of the UMAP umap-generated $k$-NN ($k=10$) network is presented above for each chromosome. The dark off-diagonal distance values in the $k$-NN adjacency graph with $k=10$show that the distance properties of actual embeddings of simulated reads from Thermus Aquaticus are almost coincident with that of the ideal case discussed above. Moreover, just as in the ideal case, by finding an approximate solution to the Traveling Salesman Problem in the unsorted network, one can find a walk (the assembly) that is close to the original reference. The results are presented in Table \ref{['tab:assembly']}.