Trace Reconstruction for DNA Data Storage using Language Models

The authors thank the reviewer for the feedback. In the following we address the concerns and questions in the order as pointed out by the reviewer.

• The reviewer noted that “unlike traditional dynamic programming-based algorithms, the proposed GPT model relies heavily on dataset-specific statistics, as seen in real-world experiments where fine-tuning was necessary’'. Interestingly, our proposed method generalizes relatively well under a distribution shift. Specifically, we see that a model pretrained on the IDS channel and evaluated on the Microsoft dataset, which is distributed differently than the training distribution (different error probabilities and the sequences exhibit long-range dependencies and are therefore not uniform at random as is pretraining data), still gives similar and often better performance than state-of-the-art reconstruction algorithms like ITR. Finally, in practice finetuning is usually possible since in DNA data storage the error distributions for a given system are either known or can be measured.

• Regarding “Could training GPT with various error probabilities enable broader applicability, including for real data?” We saw that training on the uniform distribution $\mathcal{U}[0.01;0.1]$ achieves good performance. Training directly on the estimated noise values of the Microsoft did not perform better than training on $\mathcal{U}[0.01;0.1]$. Therefore, we believe that our suggested training generation mechanism is already a good choice, it covers the error probabilities of most DNA data storage systems. Our method already outperforms existing methods on real data by a wide margin.

• The issue with the performance drop for larger clusters in Section 5.3.2 (Figure 9) could be resolved with additional hyperparameter tuning. We updated the plots accordingly.

• Regarding missing data and code. We uploaded the code and datasets, including all hyperparameters for each experiment as well as reproduction instructions.

We hope the new experiments and clarifications as well as making the code available address the reviewer’s concerns, and we would greatly appreciate reconsideration of the score. We are happy to provide further clarifications if needed.

We thank the reviewer for their comments and feedback.

Weaknesses: The reviewer's main concern is novelty. The novelty of our proposed methods lies in proposing next-token prediction as an approach to trace reconstruction and showing that it achieves state-of-the-art performance. Yes, next-token prediction is widely used in ML, however it is not obvious at all that this paradigm enables state-of-the-art performance for trace reconstruction. This is testified by the fact that there are several other prior works that propose deep learning methods for trace reconstruction (such as DNAformer and RobuSeqNet) that rely on other ideas and perform significantly worse as demonstrated in our paper. Regarding the naming of the proposed method, we updated our draft with a formal name (TReconLM which stands for Trace Reconstruction with a Language Model) for readability and revised throughout the paper to improve readability.

Questions:

1. We reformulated the sentence on Line 199 to increase clarity. We do in fact have experiments where we study the performance when the train and test distribution differ significantly in the finetuning experiments, and show that the corresponding performance drop can be overcome with finetuning.

2. The main body now contains comparisons to both RobuSeqNet and DNAformer.

3. It is beneficial to train separate models, specifically we saw that training separate models for each value of the cluster size N gives the best performance. Since this is prohibitively expensive in terms of training compute we chose to train only two models instead of nine, one for each N from two to ten. Regarding “In Line 420, the author(s) refer to "sequences of length two to five". The reviewer suspects that this might actually refer to the cardinality of clusters, with the range extending from N=2 to N=5.” Yes, thank you for pointing that out. We updated the draft accordingly.

4. A subcluster is a cluster split further into subclusters. Real-world data may contain clusters that have more than 10 sequences. We split these into subclusters to process these larger clusters with our method.

5. We believe that by finetuning we can adapt to the data-specific errors, e.g., high number of insertions towards the sequence end in the Noisy-DNA dataset, as well as learning dataset-specific patterns, e.g., the Microsoft dataset contains long-range dependencies. This explains the significant difference between the finetuned models and the baseline methods on real-world data. However, the success rates in Figure 6 and 8 for N = 2 are 0.174 and 0.402, which is far from perfect reconstruction.

We hope the new experiments and clarifications address the reviewer’s concerns, and we would greatly appreciate reconsideration of the score. We are happy to provide further clarifications if needed.

Many thanks for the feedback and the positive evaluation of our work!

Major weakness:

• We agree that the model might struggle with data that is far from the training data. From the finetuning experiments in Figure 7, we see that without finetuning, when the data is relatively far from the training data, ITR sometimes performs slightly better than our method, but with finetuning on the target distribution our method TReconLM significantly outperforms ITR and all other baselines. For DNA data storage this is not an issue, since there we can adapt to the distribution with finetuning.

• Regarding burst errors, they can be handled to some extent. The reads of the Noisy-DNA dataset show a high value of “C” and “T” towards the sequence end. The results show that our reconstruction method is able to handle this type of error.

Minor weaknesses:

• Thanks for the suggestion, we fixed that.
• We also considered other sequence-to-sequence models, specifically state-space models based on the Mamba architecture, however transformers worked better so we proceeded with them.

Questions:

• Regarding compatibility with nanopore, yes the method is compatible with nanopore reads. In Section 5.3.2 the data (Microsoft data) is from nanopore and with that data it works well. Regarding longer reads, at the moment the synthesis technologies that scale well do all not give sequences of length more than about 200, thus in the near and medium future the reads will remain relatively short. We think, however that the method also scales to larger read lengths and plan to add experiments on that. In our work, padding of the reads is not necessary, and we work directly with the nucleotides without one-hot-encoding, but one-hot-encoding with padding has been successfully applied by Bar-Lev et al., 2024.

• The method is to some extent robust to a shift in training and test distribution (see Finetuning experiments in Section 5.3). The Microsoft dataset contains sequences that exhibit long-range dependencies and are therefore not uniform at random. Still pretrained transformer models can perform similarly and often better compared to state-of-the-art reconstruction algorithms like ITR. When finetuning on real data is not possible but we have knowledge about the distribution it is typically possible to generate finetuning data.

• We are working on an experiment to find the minimum number of traces for a given set of error probabilities. As these experiments are still running, we are unfortunately not yet able to give a quantitative answer to this question. However, in the new Appendix F we study a variant of this question.

We thank the reviewer for the detailed feedback and insightful comments.

Regarding the Weaknesses:

• Regarding “The paper showcases a successful use of LLMs to solve a real problem. But the solution is not particularly surprising, especially as written.’’ We found it remarkable that the method works significantly better on a real-world problem of importance and found it not entirely expected given that there are several competing trace-reconstruction methods, including neural network-based methods (like DNAformer and RobuSeqNet) that perform significantly worse.

• In order to better understand how our GPT-based method, which we now termed TReconLM for readability, solves the trace reconstruction better than competing methods, we carried out the following additional experiments:

• In Appendix E, we plot the attention maps of the last layer and find that typically, when estimating a given position, the attention maps attend to the right position in the traces.

• To investigate why TReconLM does better than baselines like ITR we analyzed problem instances and found that for easy problem instances where the traces are close, ITR and TReconLM perform equally well, but on harder problem instances where the traces are further away, TReconLM performs better.

• Regarding important insights on what makes the method work: There are several details that one needs to get right for the method to work well, for example varying the error probabilities during the pretraining stage as we do through sampling the error probabilities $p_\mathrm{I}$, $p_\mathrm{D}$, and $p_\mathrm{S}$ from a uniform distribution is important. In particular, a model trained on the uniform distribution $\mathcal{U}[0.01;0.1]$ for the error probabilities and tested on 0.05 performs better than a model only trained on 0.05. As another example, we also tried chain-of-thought-like targets (see Appendix B), that did not perform well, so directly predicting the sequence estimate works well.

• Thanks for the suggestion to study the impact of model size further. We added a new section with preliminary scaling laws (Section 5.4) where we study the performance of transformers ranging from 3M to 300M parameters trained up to 32M problem instances. We see from the preliminary scaling plot that further increasing the number of training problem instances will likely give further benefits. The corresponding experiments are still running and will be completed in a week or two. From the current results, it can already be seen how the method improves as we invest more compute, and that moderately sized transformers like 80M ones can do very well.

• Regarding key technical details: We uploaded our code, including all hyperparameters for each experiment and run times of the GPUs during the rebuttal, and we included the training details in the paper.

Questions:

1. We added a comparison to GPT4oMini as a baseline in Appendix C.2 with zero, three, and five-shot evaluation on synthetic data. It can be seen that our method gives the best result in the considered range.

2. In order to provide some interpretability of our method we added visualizations of the attention matrices in the last layer of the transformer models of both pretrained and finetuned models.

3. As suggested, in order to push the limits of LLMs we increase the number of reads and evaluate on high noise values ($p_\mathrm{UB}=[0.1;0.2]$ see Section 5.2.1). However, the models for this experiment require further training.

4. Thanks for being open to changing your score. We think it is remarkable that TReconLM (a GPT) performs so well at trace reconstruction, which is not entirely expected given that several prior deep learning methods perform worse (RobuSeqNet and DNAformer). Moreover, it is critical to get several details right, like how to formulate the next-word prediction, and how exactly to generate training instances.

Minor comments:

We thank the reviewer for pointing out the typo. We added tables for the Hamming and Levenshtein distances as well as success rates in the Appendix G for an exact comparison of the results.