Absorb & Escape: Overcoming Single Model Limitations in Generating Heterogeneous Genomic Sequences

We appreciated for your feedback. We have provided additional experiment results in response to your questions, this includes applying A&E algorithm on the Dirichlet Flow Matching in the new task and answers to common questions from other reviewers. In the following, we respond to your questions about this paper.

Q1. A&E performance on the DDSM task

We ran A&E on this task with a default threshold $T_{\text{absorb}}=0.85$. The same dataloader and random seed as the DFM repository were used. The A&E model comprises AR and DM components, specifically the language model and distilled DFM checkpoints provided in the DFM repository. As shown in the table below, A&E achieved state-of-the-art results of 0.0262 on the test split, despite using the second best DM, distilled DFM.

Q2. Validate the proposed approach in real world long sequences

The evaluation in Section 5 is based on real-world DNA sequences from the EPD database. For details on the dataset construction process, please refer to the global response. The only simulated data used is in Section 3, which illustrates the limitations of AR models and DM in handling heterogeneous data. Additionally, we have provided results on the performance of A&E in the transcription profile conditioned promoter sequence design task from the DDSM paper (see global response and Answer to Q3 below).

Q3. How did you train the model for DDSM in the case of multiple species?

We did not train DDSM for multi-species conditioning. The DDSM is evaluated in an unconditional setting in Section 5.1, where the training data are promoter sequences from {human, rat, Macaca mulatta, mouse}.

Q4. The benefits of the proposed DMs and why can’t one use previous pre-trained DMs, such as DDSM?

The Latent Diffusion Model (LDM) proposed in this paper serves as a baseline model. We found that this simple baseline outperformed existing models such as DDSM in the unconditional generation scenario (shown in Table 3 of Section 5.2). Our aim is not to prove that the proposed baseline LDM outperforms other discrete DMs. Instead, the focus of this paper is to demonstrate the effectiveness of the proposed sampling algorithm A&E by showing that the composed model can outperform single models.

Q5. What is the effect of absorption and escape in unconditional generation?

We have included additional results below. As you can see, the A&E algorithm achieves the best performance in unconditional generation.

Q6 How did you choose the length of segments in real-world and synthetic cases?

Algorithm 1 requires segment annotations (start and end positions of each segment) to be available. However, this information is not available in most use cases. To overcome this issue, Algorithm 2 is proposed to run without knowing the true segment annotations, as detailed in line 193, Section 4.2. In both cases, the user does not need to choose the segment lengths.

For synthetic data used in the toy example, we already know the segment annotations by construction.

Q7 What is the effect of the segment selection.

Similar to Q6, both algorithm 1 and algorithm 2 don't need to select segment. For line 5 in alg1, the segments could be sampled randomly.

Q8 Why one cannot use evolutionary based models for each segment?

In practical DNA generation, where the user only has DNA sequences without annotations, Algorithm 2 handles this scenario. In addition, conventional promoter design methods based on existing promoter libraries span only a short section of the promoter.

Q9 Examples of how the segments look like in real and generated data

The examples of generated sequences are available from the following Anonymous repo. Note that no segment annotations are available for the training data; we apply Algorithm 2 for generation.

We appreciate your feedback. Following your suggestions, we have provided additional results comparing A&E with DDSM and DFM in the global response, along with clarifications to address common questions from other reviewers. Below, we provide detailed responses to each of your questions.

Q1. If trained on the original DDSM dataset, how would DiscDiff perform?

We believe the DDSM task is easier than unconditional and class-conditioning generation. The DDSM task uses a transcription profile, a real-value array of the same length as the DNA sequence, as the condition, allowing the generative model to learn a direct mapping from condition to sequence.

This is evident when applying the DDSM model to the unconditional generation task proposed in this paper during the exploration stage. Initially, using the original score net design from DDSM for unconditional generation, it could not learn the training set distribution. We scaled up the original score net by 10 times, which allowed DDSM to perform reasonably in unconditional sequence generation.

Due to time constraints, we could not adapt the DiscDiff architecture to the DDSM task. However, we provide additional experiments with the A&E algorithm on the DDSM task by using the pretrained language model as the AR component and the DFM distilled model as the DM component. As shown below, A&E achieved a new state-of-the-art on this task.

Q2. Why deep AR models do not learn the heterogeneous structure

As detailed in line 120 section 3 of the original paper, given the assumption that a sequence consists of independent segments $seg_1$ and $seg_2$, AR models struggles to learn the independence between two elements $x_k \in seg_1$ and $x_{2k} \in seg_2$. This issue would be more challenging when the training data is insufficient. As AR models have intrinsic assumptions about token dependencies, overcoming this challenge could be difficult.

Q3. Benefit of multi-species promoter pretraining and Additional Experiment on DDSM dataset

We can train a single model for each species for promoter generation. However, for certain species (e.g., P. falciparum in EPD), we have limited data, which is insufficient for training a generative model. By training on cross-species data, our model leverages genomic similarities across different species, enabling promoter sequence generation for species with limited data. We expect future foundation models for genomics to make use of all available data, just like how LLMs are making use of the web-scale data, and being able to perform downstream tasks with contextual prompts.

See the response to Q1 for the additional exp. results.

Q4. DNADiffusion is cited but not explored in Table 3

We thank the reviewer for pointing it out. This is a typo. We mislabeled DNADiffusion as BitDiffusion in Table 3. We ran the training code from the DNADiffusion repository and benchmarked it on the EPD datasets. The details of the experiments are shown in Appendix D.

Q5 a) How motifs distribution were calculated

We use EPDnew analysis tools - motif distribution for plotting the motif distribution and get the motif distribution data.

Q5 b) How MSE is calculated in conditional generation

We clarify the definition of metric in the final section of global response. Note that Sei model is used for evaluating unconditional generation as Sei is for human genome, For conditional generation, the MSE is the MSE between the motif distribution of the generated sequence and nature sequence. As detailed in line 250 section 5.2.

Q6. Benchmarking against DFM

See the response to Q1.

Q7. Fig 3 doesn’t show any comparisons with other diffusion models (e.g., DDSM)

One reason is that training on the whole datasets is expensive. We first select the best-performing DM and AR components with unconditional generation in section 5.1 of the original paper. Secondly, section 5.2 demonstrates the effectiveness of the A&E algorithm by showing A&E outperforms the single model.

Q8. Distribution of enformers predictions

We are happy to include the single track expression level prediction results in the appendix of final draft, but due to the limitation of one page PDF, we don’t have additional space to include it here.

We appreciate your feedback on the paper. In our general response, we included additional results comparing A&E with other methods on additional datasets, a sensitivity analysis of the hyperparameter $T_{absorb}$, an explanation of how to define token-level probability from the diffusion model, and other common questions proposed by reviewers. We suggest reviewing the global response first. Below, we provide detailed responses to each of your questions.

Q1. Sensitivity Analysis over $T_{absorb}$

Figure 23 of the global response provides a sensitivity analysis of the A&E algorithm over $T_{absorb}$. As $T_{absorb}$ increases, the motif plot correlation between generated DNA and natural DNA first increases and then flattens. A larger $T_{absorb}$ encourages more frequent corrections by the AR model, improving the quality of generated sequences. However, a smaller $T_{absorb}$ is more computationally efficient as it requires fewer function evaluations. A&E is robust over a wide range of $T_{absorb}$. While it is best to use the validation dataset to choose the optimal value, we found that a value of 0.85 is generally appropriate for different tasks and scenarios.

Q2. Additional Task and Dataset

Following your suggestion, we have compared A&E with other SoTA discrete generation algorithm on the transcription profile conditioned promoter sequence design task from DDSM paper. As shown in Table 6 of the uploaded PDF, A&E with Language Model as the AR component and DFM distilled as the DM component achieves the smallest MSE of 0.0262, outperforming DFM and DDSM. This confirms the effectiveness of A&E across various tasks. Please see the global response for more details about the additional experiment.

Q3. Motivation for Promoter Generation Task

Similar to prior work (ExpGAN [1]), conventional promoter design methods based on existing promoter libraries span only a short section of the promoter. In contrast, deep generative models for promoter design explore the distribution properties of promoter sequences, enabling the generation of new promoters that classical methods cannot achieve. This motivation is shared by the DDSM and DFM papers and other DNA generation tasks.

Q4. How Enformer predictions are used for Evaluation

The Enformer outputs are used to show how the promoter sequence influences the downstream target gene instead of the bins corresponding to the promoter sequence itself. We inserted the generated promoter before the target gene, e.g., TP53, and checked how the promoter changed the expression level of the target gene. While the promoter sequence is only 128 bp, accounting for 1 bin in the prediction, the target gene is more than 10,000 bp long. For example, TP53 is about 25,000 bp, accounting for 195 bins in the Enformer output. Using SSE, we aggregate the change in expression levels across all the cell lines.

Q5. Token-Level Probabilities from the Diffusion Model ($p^{DM}$)

$p^{DM}$ at line 3 of Algorithm 2 should be $p^{DM}(\mathbf{x}_i) \in \mathbb{R}^{L\times4}$, representing the token level emission probability from the Diffusion Model. $p^{DM}(\mathbf{x}_i)$ can be retrieved for most of the discrete diffusion algorithms such as DDSM, DFM, and DiscDiff. For DNA generation, $p^{DM}(\mathbf{x})$ are usually stored as a variable logits $\in \mathbb{R}^{L\times4}$, where each row $p^{DM}(\mathbf{x}_i)$ represents token level probability distribution over {A,T,G,C}. $p^{DM}(\mathbf{x})$ are produced after the last sampling steps. In the latent diffusion model used in this paper, it is produced by the second-stage decoder (Appendix A), while in DDSM and DFM, it is produced by a 1D-CNN (output_channel=4) in the score net.

Q6. Description of EPD Dataset

We detailed the construction process of the DPE dataset in the global response. Briefly, we downloaded all the promoter records (30 million) from EPD datasets. Each record is a triple of (sequence, species, cell types, expression levels). Sequences could be duplicated in these records, so we aggregated these records by sequence, producing 160K unique sequences. Each sequence is a promoter-downstream sequence centered around the TSS of length 256 or 2048.

Q7. How the Model Learns to Differentiate Between Different Genomic Regions

Following the suggestion from reviewer crFN, we performed BLASTN on the two halves of sampled sequences against promoter-like and protein-like segments from the training set. A larger score indicates higher similarity. The results confirm that A&E better handles heterogeneity than the AR model.

Q8. Convergence criteria for Algorithm 1

The convergence criteria of Alg 1 is $\tilde{\mathbf{x}}^t$ converge to a fixed distribution, which can be achieved through $|\tilde{\mathbf{x}}^{t+1} - \tilde{\mathbf{x}}^t\| < \delta$.

[1] Zrimec, J., et al., 2022. Controlling gene expression with deep generative design of regulatory DNA

We sincerely appreciate your insightful reviews. Following your suggestions, we have presented additional experiments on DDSM dataset and more detailed description on the dataset in the general response and uploaded PDF. Hopefully our general reply addresses your concerns regarding the selection of $T_{absorb}$ and the extraction of $p^{DM}(\mathbf{x}_i)$. Below, we provide detailed responses to each of your questions.

Q1. Using the same evaluations as previous papers

We have hereby presented additional experiment results on the transcription profile conditioned promoter sequence design task presented in DDSM paper with A&E. As shown in Table 6 of the uploaded PDF, A&E with the Language Model as the AR component and DFM distilled as the DM component achieves the smallest MSE of 0.0262, outperforming DFM and DDSM. This confirms the effectiveness of A&E across various tasks.

Q2: Details about how exactly these sequences are constructed

We have detailed the construction process of the EPD dataset in the global response. Briefly, we downloaded all the promoter records (30 million) from EPD datasets. Each record is a triple of (sequence, species, cell types, expression levels). Sequences could be duplicated in these records, so we aggregated these records by sequence, producing 160K unique sequences (as mentioned in the line 218 of original paper). Each sequence is a promoter-downstream sequence centered around the transcription start site (TSS) of length 256 or 2048.

Q3: Why species-wise generation and not cell-type conditioning?

There are two reasons for setting species conditioning as a task. First, promoter sequences from different species follow different distributions, making it an excellent testbed for benchmarking various generative algorithms. In contrast, cell-type conditioning is less distinct since all cell types share the same genome, and regulatory elements can bind to specific cell types only under certain conditions. However, if enough samples are provided, each sequence will appear in all cell types. In theory, the optimal task format should be (species, cell type, expression) co-conditioning, and this paper presents the first step towards that goal.

Additionally, while we can train a single model for each species to generate promoters, limited data is available for certain species (e.g., P. falciparum in EPD), which is insufficient for training a generative model. By training on cross-species data, our model leverages genomic similarities across different species, enabling promoter sequence generation for species with limited data. One use case of unconditional promoter generation is expanding existing promoter libraries, crucial for synthetic biology. We expect future foundation models for genomics to make use of all available data, just like how LLMs are making use of the web-scale data, and being able to perform downstream tasks with contextual prompts.

Q4. Heterogeneity of the DNA sequences in the main results

Following your suggestion, we performed BLASTN on the two halves of the sequences sampled by A&E against promoter-like and protein-like segments from the training set. A larger score indicates higher similarity. We will add these additional results to our Appendix.

The results confirm the existence of heterogeneity in the generated and training sequences.

Q5. Is it better to use a masked language model instead of an AR model?

We recognize this as a valuable suggestion. A masked language model could potentially improve the quality of generated sequences by considering the bidirectional context $P(x_{i:j}|x_{0:i-1},x_{j+1:L})$. However, AR model with caches (storing intermediate results for previous generated tokens) could potentially be computationally efficient as it requires at most sequence length number of function evaluations. In contrast, a masked language model needs to consider the whole context each time.

Q6. How are probabilities for a sequence/sequence segment extracted from the DM?

$p^{DM}(\mathbf{x}_i)$ can be retrieved for most discrete diffusion algorithms (such as DDSM, DFM, DiscDiff). In DiscDiff, logits are available after the second stage decoder layer (Appendix A). In general, $p^{DM}(\mathbf{x})$ are usually stored as a variable logits $\in \mathbb{R}^{L\times4}$, where each row $p^{DM}(\mathbf{x}_i)$ represents probability distribution over {A,T,G,C} token. More details are available in our general response.

Q7. Discrepancies between the model names mentioned in lines 222-223 and those in Table 3.

We thank the reviewer for pointing it out. This is a typo, the BitDiffusion in Table 3 should be DNADiffusion. We retrain the DNADiffusion on our dataset and reported the results in Table 3.

Q8. How A&E generalises to other datasets/settings and evaluations

This is addressed with the additional experiment and more clarification about metrics and datasets in the general response. We believe motif distribution is essential to DNA generation, as it directly measures how DNA functions.

Q9. Tuning $T_{absorb}$ seems non-trivial.

We have found a default value of 0.85 effective in many settings. A sensitivity analysis of the algorithm on $T_{absorb}$ is provided in the general response.