Gumbel-Softmax Discretization Constraint, Differentiable IDS Channel, and an IDS-Correcting Code for DNA Storage

We would like to thank the reviewer for his/her valuable comments, which will undoubtedly help improve the quality of our work. Below, we provide detailed responses to each of the raised concerns. We hope the reviewer will consider revising the rating to a more positive score based on these clarifications and improvements.

Q1: Reword the experimental results and have more comparative experiments.

We apologize for the confusion caused by the presentation of the experimental results. To better clarify the results and highlight the advancements made by the proposed method, we will revise the experimental section and include additional comparative experiments.

The reviewer may refer to the separate rebuttal Section Comparative Experiments below, as well as the revised PDF, which will be provided later, for further details.

Q2: Can the differentiable IDS channel be extended to realistic DNA storage channel?

This is an insightful question. It was widely believed that the biochemical procedures involved in DNA storage introduce errors that follow specific patterns or distributions, such as energy constraints or the avoidance of long homopolymer runs, etc. However, there were few detailed, standardized descriptions of the exact principles governing errors. This may be due to the complexity of the underlying processes, which are often equipment-dependent and may vary across different experimental setups.

The answer is yes to this question that, the differentiable IDS can be extended to any IDS channel. It only requires a customized generation method for the error profile $p$ that corresponds to the specific IDS channel. The proposed differentiable IDS method accurately introduces errors to the sequence $c$ based on the predefined error profile $p$. Once the error patterns of your IDS channel are well-defined, you can generate the profile $p$ according to these rules, and the differentiable IDS will faithfully replicate the behavior of your IDS channel.

However, accurately defining the error rules is challenging. Can we apply the proposed differentiable IDS to undefined realistic channels? As briefly mentioned in the manuscript, generative deep learning models have been used to simulate how DNA sequences are transformed during DNA storage, such as in

• S. Kang, et.al, Generative adversarial networks for DNA storage channel simulator 2023.

By incorporating the differentiable IDS as a module within such generative models, the generative model may learn to generate error profiles, rather than directly generating retrieved DNA sequences. This may alleviate the burden of the model. While this is a straightforward application of differentiable IDS, it falls outside the scope of this current work.

Comment 1. The proposed contributions are somewhat loosely connected.

We agree with the reviewer that, from the perspective of bioinformatics or DNA storage, the contributions may seem loosely connected. However, we believe that, from the standpoint of machine learning and learning representation, these contributions are essential and closely aligned with overcoming the challenges of developing an autoencoder-based IDS correcting code.

The proposed Gumbel-Softmax Discretization Constraint induces a hard representation from the deep learning features. We believe that this approach helps bridge the gap between continuous deep models and discrete applications, and it could be extended to other tasks that face similar discretization challenges.

The proposed Differentiable IDS channel introduces a model that mimics IDS operations while preserving the ability for gradient back-propagation. First, it demonstrates that the Transformer can efficiently mimic IDS operations. Second, it serves as a general module for neural networks, applicable to various tasks where discrete IDS operations play a role.

We kindly remind the reviewer that, since this is a submission to a machine learning conference, the novelties should perhaps be re-evaluated from the perspectives of both the machine learning community and the DNA storage community.

Comment 4. Novelty of applying Gumbel-Softmax Discretization Constraint.

We apologize that our manuscript led to a misunderstanding by the reviewer.

The Gumbel-Softmax was originally introduced to enable sampling from categorical distributions in VAEs, thereby enhancing their generative capabilities. In such contexts, researchers typically encourage the distribution to maintain high entropy (i.e., avoid a one-hot style distribution) to prevent model degeneration.

In our case, we found that applying the Gumbel-Softmax effectively constrains its input vector to approximate a one-hot vector during training, and we provided a brief mathematical analysis to support this observation. This use diverges from the original sampling task, with the only similarity being the shared formula. To the best of our knowledge, no existing research has explored this specific approach.

It is true that Figure 1 in the original paper on applying Gumbel-Softmax in deep learning, “Categorical reparameterization with Gumbel-Softmax”, appears to illustrate quantization. However, the Figure 1 is actually focused on how the hyperparameter temperature $\tau$ influences the sampling process. Specifically, if $y=GS(x;\tau)$ represents a sample $y$ drawn from the Gumbel-Softmax distribution, Figure 1 demonstrates how the temperature $\tau$ controls the 'hardness' of the sample $y$ —from a hard sample of 'one-hot' vector ($\tau=0$) to a soft sample of 'uniform' distribution ($\tau=\infty$).

However, in our work, the Gumbel-Softmax is applied to induce the input $x$ of $y=GS(x;\tau)$ to resemble a one-hot vector. When applied as a constraint in a non-generative network, Gumbel-Softmax drives the input $x$ to a harder form. For instance, by applying the constraint, $x$ is more like the form (0,1,0,0) rather than (0.2,0.3,0.1,0,4), where the former corresponds directly to a nucleotide base T in ATGC, while the latter is difficult to interpret as a single base.

In summary, the role of Gumbel-Softmax in this work differs significantly from its original purpose. In fact, framing Gumbel-Softmax as a sampling method does not align with the proposed Gumbel-Softmax Discretization Constraint. We also intuitively suspect that introducing other sources of noise, aside from Gumbel-Softmax, would produce a similar effect in inducing hard features.

We thank the reviewers for their valuable comments. Below, we provide detailed responses to address each concern, and we hope the reviewer will consider raising the ratings to a positive score.

Q1. Error profile has length larger than the codeword.

This is due to the sync error, particularly the insertion, occurs without moving the index pointer to the next position. For example, inserting CCC at position 1 in a length-4 sequence ‘ATGC’ to obtain ‘A CCC TGC’ requires a length-7 profile of (0, insert C, insert C, insert C, 0, 0, 0).

Q2. Why using transformer for DIDS?

We did not explore diverse architectures for DIDS because the transformer-based model is currently the most suitable choice for our data structure. As mentioned in our response to Q1, the error profile often differs in length from the codeword, creating asynchrony that many neural network architectures, such as CNNs, struggle to handle. In contrast, the transformer model naturally captures global relationships within the input and is better suited for handling multi-input scenarios, such as the codeword and profile, through individual embeddings. Furthermore, as the reviewer noted, this DIDS model is effective. Therefore, we opted to use this reasonable and effective model without exploring alternative architectures, as doing so would deviate from the primary focus of this work.

We will revise the manuscript to provide a clearer explanation of how the transformer model works.

Q3. Why using the auxiliary reconstruction?

As mentioned in Line 287 and Section 5.2, the autoencoder (AE) does not naturally align with error-correcting codes (ECC). While conventional AEs focus on feature compression, ECCs are designed to generate parity checks. As demonstrated in the ablation study, directly optimizing Equation (15) does not enable the AE to behave as an ECC in our experiments.

Inspired by systematic codes, which embed the input message within the codeword, we introduced an auxiliary reconstruction task to enable the encoder to replicate the input message. As for whether the resulting code is systematic or not, it is determined by the model during training and is challenging to confirm. Variations in base mappings, permuted messages, and positional choices make it impractical to verify all possible systematic cases exhaustively.

Q4. Why transformer-based encoder/decoder. & Time Complexity.

The reasoning behind using a transformer-based encoder/decoder is similar to and partially addressed in our response to Q2. First, the model needs to handle inputs of variable lengths, as sync errors can alter the codeword length. Second, we aim to leverage the transformer’s logical capabilities, which may be suited to mimic some of the combinatorial rules underlying established codes, such as the VT code.

It is well known that the transformer has $O(n^2)$ complexity. We acknowledge the limitation of not exploring the numerous efficient transformer variants. However, we wish to emphasize the novelties of our contributions, which make AE-based IDS-correcting codes possible, as they are: the Gumbel-Softmax Discretization Constraint and the Differentiable IDS channel. As this is an emerging topic, there are many potential directions to explore, and it may not be feasible to investigate all of them within a single work.

Given the parallelization capabilities of GPUs and the codeword length limited by current biotechnology, the transformer's $O(n^2)$ complexity results in acceptable time consumption. We decoded 1,280,000 codewords on an RTX 3090, with the results listed below (the encoder, which shares the same structure, exhibits similar performance):

This table will be included in the revised PDF.

Q5. The one-hot vector sequence and the letter sequence.

We apologize for the misunderstanding and will revise the manuscript for clarity.

The output of the encoder and the input of the decoder are represented as probability vectors or one-hot vectors. This setup allows for gradient backpropagation, enabling joint optimization of the encoder and decoder. If letter representations were used between encoder and decoder in training phase, the letter embedding in the decoder would block the gradient.

During testing, the encoder output $c$ is processed by $\hat{c}=onehot(argmax(c))$ before being input into the decoder, where $argmax(c)$ represents the codeword in familiar letter-sequence form. In Line 404, the Gumbel-Softmax constraint is applied to force the encoder's output from a probability vector to a more one-hot-like vector. This constraint eliminates the discrepancy between $c$ and $\hat{c}=onehot(argmax(c))$ during testing (if $c=(1,0,0)$ in one-hot form, then $\hat{c}=c$, if $c=(0.5,0.25,0.25)$ in non-one-hot form, then $\hat{c}=(1,0,0)\neq c$. ). Thus, this constraint is a key contribution of our work.

Thank you for providing your response. Although the authors offer adequate answers, I still lack confidence that this work is suitable for presentation at a top-tier venue like ICLR. The use of the transformer architecture feels unnatural and seems disconnected from prior works rather than sequentially building upon them. I believe the paper would have been stronger if it started with a fundamental auto-encoder architecture and progressively demonstrated how their proposed methods design a code suitable for the DNA channel, eventually applying the transformer architecture later in the narrative. Additionally, I still have doubts about whether the comparisons with existing methods are fair. For these reasons, I will maintain my original score.

We are thankful for the reviewer’s valuable comments. Below, we provide detailed responses to address each concern. We hope the reviewer will consider raising the score to a positive rating.

Q1: Comparison experiments. & What will change for larger error probabilities?

We agree that providing results for a broader range of experimental settings is essential. As suggested by the reviewer, we have extended the experiments in two aspects.

First, we conducted comparison experiments with the combinatorial code from Cai, the segmented code DNA-LM, and the well-known heuristic code HEDGES, despite the challenges in aligning these methods under the same settings. We have a separate rebuttal page on this, the reviewer may refer to the rebuttal Section Comparison Experiments in the following.

Second, we extended our experiments to include channels with error probabilities in (0.5%, 1%, 2%, 4%, 8%, 16%) at a typical code rate of 0.67 (100/150). The NERs are presented in the following table, where each cell represents the NER of a model trained on the channel specified by the row header and tested on the channel specified by the column header.

The results demonstrate that the proposed method is compatible across different IDS error probabilities. As expected, models trained on channels with higher error probabilities exhibit compatibility with channels with lower error probabilities. In most cases, models trained and tested on similar channels achieve better performance. We will revise the manuscript adding such results.

Q2&3: Is applying the combinatorial single IDS correcting code directly good enough?

We apologize for not providing a clear comparison in the original manuscript.

When applying a single IDS-correcting code and treating the received sequence as the decoded sequence for failed decoding, the performance is directly related to the probability of multiple errors occurring. For example, consider 1% deletion on a 150-length sequence. The probability that more than one error occurs is approximately $1-0.99^{150}-C_{150}^{1}\cdot 0.01\cdot 0.99^{149} = 44.3$%. Moreover, treating the received sequence as the decoded sequence keeps the deletion shift, which introduces nearly half-length bits as errors on average. The final NER will be on the order of 10%.

We conducted experiments on both directly applying a combinatorial code and a segmented code. The results are presented in the following rebuttal Section Comparison Experiment. The numbers suggest that the performance of directly applying a combinatorial code is suboptimal.

Regarding the conventional VT code, we anticipate similar results to those observed with Cai’s code. This expectation arises because Cai’s code is an extension of the VT code, designed to correct single IDS errors on a 4-ary alphabet. Cai code and VT code behave the same in failing to correct multiple IDS errors and indel errors, respectively.

We sincerely thank the reviewer for their valuable comments. Below, we provide detailed responses to address each concern. We hope these clarifications and additional insights will encourage the reviewer to consider raising the rating score.

Q1. Comparisons with existing methods.

We acknowledge our oversight in not including sufficient comparison analysis in the original manuscript. Although aligning experimental settings for established codes was challenging, we conducted comparative experiments involving Cai’s combinatorial code, the segmented code DNA-LM, and the well-known heuristic code HEDGES. The results demonstrate the commendable performance of the proposed method, which we believe further validates its effectiveness.

The reviewer may refer to the separately provided rebuttal Section, Comparisons Experiment, or the revised PDF, which will be submitted subsequently.

Q2. Generalization to broader IDS channel settings. & Real-world DNA sequencing data.

As demonstrated in the manuscript, the proposed method is designed to benefit from customizing codes for specific IDS channel settings. While it is true that a single trained model may not be optimal for all IDS channel settings—indeed, no code can achieve that—the proposed approach offers a method for the rapid development of customized codes tailored to different settings.

Prompted by the reviewer's concern, we realized that limiting the experiments to the 1% error channel and the HOM, ASC, and DES channels was insufficient to show the model's generalization to diverse channel settings. To address this, we have extended our experiments to include channels with error probabilities of 0.5%, 1%, 2%, 4%, 8%, and 16% at a typical code rate of 0.67 (100/150). The NERs are presented in the following table, where each cell represents the NER of a model trained on the channel specified by the row header and tested on the channel specified by the column header.

The results demonstrate that the proposed method is compatible across different IDS error probabilities. As expected, models trained on channels with higher error probabilities exhibit compatibility with channels with lower error probabilities. In most cases, models trained and tested on similar channels achieve better performance. We will revise the manuscript adding such results.

In the proposed customized code, the channels require a well-defined setting. For real-world DNA sequencing data, achieving a clear and unified depiction of such IDS channels remains an open research topic. Moreover, the realistic IDS channels are highly dependent on various wet-lab experimental conditions, the specific biochemical equipment employed, and other experimental factors. Unfortunately, conducting two rounds of wet-lab experiments—one for investigating channel characteristics and another for validating code performance—is beyond our current resources. Therefore, we hope that the reviewer will consider our work to be more aligned with machine learning and coding theory.

Q3. Ablation Studies and hyperparameter optimization.

We apologize for not including comprehensive ablation studies and hyperparameter optimization in the original manuscript.

This work involves two main hyperparameters: the weight $\mu$ of the auxiliary reconstruction loss and the temperature $\tau$ in the Gumbel-Softmax Discretization Constraint formula. We have included the ablation study for the auxiliary loss, weight $\mu$, and the Gumbel-Softmax constraint in the original manuscript. However, the hyperparameter optimization for the temperature $\tau$ was not provided.

To address this gap, we conducted experiments with different choices of the temperature $\tau \in$ {$0.125, 0.25, 0.5, 1, 2, 4, 8$}, while keeping other settings at their default values. The results are listed below.

This table illustrates that the lowest NER is reported at the default setting $\tau=1$. When the value of $\tau$ is increased or decreased from this default, the NER worsens.

We will include a subsection or appendix in the revised manuscript displaying the training curves for different values of $\tau$, as this will provide more informative insights than the table alone.