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

We thank the reviewer for their comments and kindly point them to the general comments in addition to specific comments below. Kindly also review the responses to other reviewers.

[Q1] Have you evaluated the method on real data/data with noise?

Before we discuss the additional evaluation results we have included in the revised version, we wanted to address why we used ART as a simulation toolkit. An important consideration was that for such reads, we know the ground truth (making evaluation much more precise) while ensuring that the reads have the same distribution of errors as one would expect from real world sequencers.

In fact, ART is an industry-standard simulator that is used to benchmark almost every new aligner that is introduced. ART mimics the variants, insertions/deletions (INDELs) and other error models found in reads generated by expensive Illumina machines. The ART simulator used in our experiments was fine-tuned on an Illumina MiSeqv3 machine model (https://www.illumina.com/systems/sequencing-platforms/miseq/specifications.html). BWA-Mem and several other aligners have been validated on such simulated data: Table 1 here: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2705234/ demonstrates performance on WGSIM from SAMTools, which is a more primitive version of ART (https://bmcbioinformatics.biomedcentral.com/articles/10.1186/1471-2105-14-184). ART was also used in the simulation study of the 1000 Genomes datasets to profile the quality of the resource: https://www.niehs.nih.gov/research/resources/software/biostatistics/art/index.cfm

We also want to note that Pure Reads serve an important purpose: They establish an upper bound on the performance of the aligner (as noisy reads would only have lower performance), allowing us to benchmark the maximum possible performance for Transformer-based LLM models as shown in Fig. 2. These pure-read results make it superfluous to benchmark the existing Transformer models on noisy data / real data.

In the revision: While we believe that ART reads with $Q_{PH} >= 30$ are sufficient to benchmark aligners -- Next-gen Illumina machines generate reads at Phred score $Q_{PH} >= 30$, we address the reviewers’ request by evaluating DNA-ESA on reads with $Q_{PH} \in [10,20]$, $I,D = 1$%; see Table 4 (left).

We are pleased to report that DNA-ESA performance at $Q_{PH} \in [10,20]$, $I,D = 1$% is only about $1$% less than the performance reported at $Q_{PH} \in [30,60]$, $I,D = 0.01$%.

Performance on real read data: In Table 4, we have also included the performance of DNA-ESA on real reads generated by PacBio CCS ($Q_{PH} = 20$): and read length $\in \{250, 350, 500\}$.

Here, too, DNA-ESA performs as well on real read data as on ART-generated reads reported in our original submission. This further supports our original evaluation platform.

[Q2] What are the runtime resource requirements?

The model requires a GPU with >10GB of vRAM to run inference on DNA-ESA. The model uses 20GB of CPU memory during inference. This is an initial version of the model, and as noted in Sec. Future Work, will continually become more performant using methods such as model distillation, compilation, and distributed processing techniques. Efficient transformers are still an area of active research which have seen notable improvements over the last few years (e.g. FastAttention, Long-range architectures such as HyenaDNA and similar).

[Q3] Does the model generalize to less related organisms? Presumably there will be some dependence on the degree of homology.

We thank Reviewer 3 for their curiosity in understanding whether the model generalizes to less related organisms. Accordingly, we have expanded the inter-species experiments in the transfer learning section (Table 3) to include Thermus Aquaticus -- an ancient extremophile bacterium -- and Acidobacteriota reference sequences: two species that are highly dissimilar to the rest. We are pleased to report that the model -- originally trained to align reads to fragments on Chr. 2 of humans -- continues to perform well (in fact even with top-K = 1, the recall is >93%)! This adds further support to the claim that the DNA-ESA architecture truly learns the alignment task rather than the sequences themselves seen during training.

We urge the reviewer to kindly read the general comments as well as our replies to other reviewers.

[Q1] Provide support to the hypothesis that a fragment and its subsequences are close in the embedding space.

We agree with the reviewer and we have updated Figure 1 to visualize both reads and fragments. We see convincing visual evidence that both the reads and the fragments align within the embedding space.

[Q2] Make the notations more readable, e.g. in Eq. 2, q is used in SA(r,R). Eq. 6, really difficult to guess the meanings.

We have resolved this discrepancy and now use $q$ in Eq. 2 as well to denote the start position. Furthermore, for Eq. 6, we realize that the asterisks on $F$ and $q$ are difficult to read. So we have added clarity to the equation in the new version.

[Q3] Comparison to existing classical methods (Bowtie-2, Minimap, BWA-Mem-2)

We have presented the side-by-side results of Bowtie-2 and DNA-ESA in Table 4 on two noisy datasets. We are happy to report that the performance of Bowtie-2 is only 3-4% above the recall performance of DNA-ESA. The quality of the recalled fragments (of those reads that are successfully aligned) is comparable among the two models.

[Q4] “Provide details of the transformer. Do you do paddings to short reads?”

Clarified section 3.3: “Reads and shorter fragments were padded to equal the length of the longest sequence sampled in a batch.”

[Q5] What happens to negative samples, i.e. a random read that is not similar to any reference fragment?

We thank the reviewer for the insight about including additional metrics such as Accuracy, Precision and F1 score in addition to Recall: the classic case of evaluating specificity vs. sensitivity. We are happy to report in Table 2 that these metrics are all high: negative samples were drawn from Chr. 3 with positive samples from Chr. 2.

[Q6] How do you decide the boundaries of a reference fragment in the embedding space?

We are unclear about what the reviewer means by boundary. Our guess is that the reviewer asks for clarification about how we go to the exact location of a read from the top-K fragments. Note that each fragment has meta-data corresponding to the reference location from which the fragment is drawn. Below we also briefly recap the alignment process. (Step 1) Identify the long fragments corresponding to the most likely match to a read. Top-K neighbor is computed in log complexity; (Step 2) Fine-alignment of the read to the fragments: Compute Smith-Waterman distance between the read and each of the candidate fragments. This distance metric is thresholded. We sweep across several of these thresholds in Table 3 and describe the process in Appendix C.1.

[Q7] Reorder the figures so it flows with the text.

We extend our gratitude to the reviewer for their valuable comments. We hope to convince the reviewer that the current order is justified.

Figure 1 plays a pivotal role by immediately illustrating the generated embedding space, which forms the foundation for all subsequent downstream applications. This depiction holds particular appeal for Machine Learning enthusiasts, as it showcases the training method's unique ability to generate trajectories within the embedding space.

Figure 2 provides a side-by-side performance comparison of various Transformer-based models on pure reads (reads without any noise). Despite the apparent simplicity of this task, it becomes evident that competing Transformer-based models struggle to address it. Finally, with the motivation firmly established, Figures 3 and 4 delve into the intricacies of the method, describing the finer points of how these representations are generated.

[Q8] Appendix A, add legend for the lines

We have added the following to the figure text: “DNA-ESA convergence plots. Each plots show the loss pr. step (grey). For clarity, we smooth the loss using a moving average (black).”

[Q9] Appendix B, B.3 the complexity is G, not log(G) as you have to compare with all reference fragments.

It's crucial to highlight that the search complexity across G is, indeed, logarithmic (log(G)).

This logarithmic scaling represents a fundamental advantage derived from utilizing vector stores in our task. Without the efficiency gained from vector comparisons, resorting to direct string representations for matching the read with every fragment might as well have worked.

Vector stores opt for a highly accurate approximate K-nearest neighbor method in addition to other proprietary optimizations to search that include Small-World approximations, hierarchical organization of graph data and concurrent search. This is the entire reason behind FAISS’ ubiquity and Pinecone’s rapid success as a vector store solution. This innovation facilitates the breakthrough in our representation learning solution to alignment.

We thank Reviewer 1 for their comments about our submission. We appreciate their recognizing the promise of the approach, specifically (a) the effectiveness of the sequence embeddings generated by DNA-ESA (as demonstrated in search); (b) the generalization across species and the flexibility of the approach in encoding reads and fragments of varying lengths.

We urge the reviewer to also review the general comments in addition to our responses to other reviewers. Below we address the specific concerns:

[Q1] “I felt that the paper is a very dense read for the general ML audience at ICLR for folks who do not have DNA sequencing background, and it will be great to make the paper more accessible.”

Acknowledging the interdisciplinary nature of our work, we concur with the reviewer's perspective that it is our responsibility, as the authors, to bridge the gap between cutting-edge techniques in NLP and the longstanding challenges in bioinformatics and genomics. To address this, we have taken the following steps:

In Eq. 1, we define a read in its simplest form: a sequence of bases. In Eq.s 2 and 3, we write the objective of sequence alignment through the application of the sharding rule [P1]. Furthermore, in Section 3.1, we present the rationale behind a straightforward inequality constraint Eq. 4 (modified now to incorporate homologs / repeats) within the representation space.

It is worth noting that our approach to motivating sequence alignment, as outlined, is itself innovative and designed to resonate with the machine learning community. We remain optimistic that the reviewer finds merit in this novel perspective.

[Q2] “The performance for short reads is worse than long reads, given that short reads are more commonly used, this may affect how this system can be actually used.”

The reviewer rightly notes that short-read aligners are more commonly used than long-read aligners. Our approach does not intend to replace existing short-range alignment methods. Rather, methods for sequencing have seen rapid advances in cost and their ability to produce increasingly longer reads (PacBio and Oxford Nanopore, up to 10s of kilobases long), and accordingly, aligners to map such long reads are also of recent interest in the bioinformatics community.

[Q3] “Can the authors comment on how this paper is a good fit for ICLR, and the steps the author may take to make this paper more accessible to the ML audience?”

Representation learning holds a pivotal role within the framework of ICLR. The treatment of DNA sequences demands a computational approach grounded in deep learning, as evidenced by numerous recent works in the field [1,2,3]. Notably, the guidelines for ICLR explicitly outline a track dedicated to "Applications in biology, and related fields."

In our endeavor to make our paper accessible to the broader ICLR community, which includes individuals with expertise in machine learning, we kindly direct the reviewer's attention to the first part of the response. We believe our work aligns with the conference's objectives, particularly within the specified track.

[1] Nguyen, Eric, et al. "Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution." arXiv preprint arXiv:2306.15794 (2023).

[2] Dalla-Torre, Hugo, et al. "The nucleotide transformer: Building and evaluating robust foundation models for human genomics." bioRxiv (2023): 2023-01.

[3] Ji, Yanrong, et al. "DNABERT: pre-trained Bidirectional Encoder Representations from Transformers model for DNA-language in genome." Bioinformatics 37.15 (2021): 2112-2120.