Common Response for All Reviewers

We would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our "simple but powerful core idea" and "comprehensive experimentation" demonstrating "competitive performance" in recovering Boltzmann distributions for unseen proteins.

The main concerns about the paper shared by the reviewers can be broadly placed in the following categories:

• A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)
• A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)
• An ablation of the jump step (Reviewer E7vk)
• A lack of results on larger systems (Reviewers iqm5 and tyev)

We have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\approx$ 50 AA) and the shortest ATLAS proteins ($\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.

Comparison with Diffusion Models

To highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):

Essentially, we find that walk-jump sampling is $64\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.

We would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.

Physical Validity and Energy Analysis of JAMUN Samples

We have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.

We see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the overall quality of the JAMUN samples is high. For a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:

Next, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:

Ablation of the Jump Step

As requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far higher divergence metrics. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.

Results on Larger Systems

We recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.

We have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.

Importantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).

For the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware. For the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating transferability even in this setting.

We compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.

Response to Reviewer tyev

1. The idea is that walk-jump sampling shares many parallels with MD, but you are correct that we do not currently identify which specific aspects help the transferability of JAMUN. We would be happy to rephrase this in the paper. In a standard diffusion model, the noise level is continuously reducing during sampling which makes it hard to give a physical interpretation.
2. Our analysis of the force field energy of JAMUN samples should address these concerns; we do not do any post-processing of the JAMUN samples. It is true that there is some bias in the final conformational distributions, which we cannot fully correct for, but our analysis does show that we are able to get quite close.
3. The MSMs were made with the analysis code from MDGen (Jing, et al.) which we have modified and uploaded as part of our supplementary code. We utilize the same MSM hyperparameters (eg. TICA lag), but it is likely that those hyperparameters are not optimal for the Timewarp datasets.
4. We ran the exact commands (and checkpoints) used by the MDGen authors for sampling, according to their README. It is possible that the recommended sampling parameters/model is not ideal for this comparison.
5. We hope that you will find the results on the fast-folding proteins and ATLAS encouraging.

We would be very happy to improve the formatting and the clarity of the paper in the camera-ready version, if accepted.

Common Response for All Reviewers

We would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our "simple but powerful core idea" and "comprehensive experimentation" demonstrating "competitive performance" in recovering Boltzmann distributions for unseen proteins.

The main concerns about the paper shared by the reviewers can be broadly placed in the following categories:

• A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)
• A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)
• An ablation of the jump step (Reviewer E7vk)
• A lack of results on larger systems (Reviewers iqm5 and tyev)

We have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\approx$ 50 AA) and the shortest ATLAS proteins ($\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.

Comparison with Diffusion Models

To highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):

Essentially, we find that walk-jump sampling is $64\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.

We would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.

Physical Validity and Energy Analysis of JAMUN Samples

We have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.

We see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the overall quality of the JAMUN samples is high. For a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:

Next, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:

Ablation of the Jump Step

As requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far higher divergence metrics. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.

Results on Larger Systems

We recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.

We have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.

Importantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).

For the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware. For the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating transferability even in this setting.

We compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.

Response to Reviewer E7vk

1. Our detailed Posebusters analysis above shows that JAMUN captures bond lengths correctly; we hope you will find this sufficient, but we are happy to perform further analyses.
2. In the current framework, no, we indeed lose kinetic information and cannot map the dynamics in $Y$ to those in $X$.
3. and 4. We will fix the caption of Table 1 and also detail the JSD metric for clarity in the updated paper.
4. According to Figure 2 a), TBG generates 5000 structures, so we used 5000/20 = 250 structures to represent TBG (20x shorter). JAMUN generated 100000 structures.

Common Response for All Reviewers

We would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our "simple but powerful core idea" and "comprehensive experimentation" demonstrating "competitive performance" in recovering Boltzmann distributions for unseen proteins.

The main concerns about the paper shared by the reviewers can be broadly placed in the following categories:

• A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)
• A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)
• An ablation of the jump step (Reviewer E7vk)
• A lack of results on larger systems (Reviewers iqm5 and tyev)

We have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\approx$ 50 AA) and the shortest ATLAS proteins ($\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.

Comparison with Diffusion Models

To highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):

Essentially, we find that walk-jump sampling is $64\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.

We would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.

Physical Validity and Energy Analysis of JAMUN Samples

We have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.

We see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the overall quality of the JAMUN samples is high. For a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:

Next, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:

Ablation of the Jump Step

As requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far higher divergence metrics. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.

Results on Larger Systems

We recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.

We have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.

Importantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).

For the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware. For the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating transferability even in this setting.

We compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.

Response to Reviewer iqm5

1. We believe our preliminary experiments on the fast-folding proteins and ATLAS where we keep all the training and sampling hyperparameters fixed as before, should answer your question. It does seem like these hyperparameters are robust.
2. Yes, indeed the framework is general enough to support the sampling of collective variables (similar to how diffusion models can be learned for those data). Here, we wanted to utilize the connection between walk-jump sampling and MD which does operate in an all-atom space, to show that it might be possible to build transferable yet efficient samplers for molecular conformations.

Common Response for All Reviewers

We would like to thank all of our reviewers for their thoughtful feedback, especially the encouraging comments on our "simple but powerful core idea" and "comprehensive experimentation" demonstrating "competitive performance" in recovering Boltzmann distributions for unseen proteins.

The main concerns about the paper shared by the reviewers can be broadly placed in the following categories:

• A lack of comparison with diffusion models (Reviewers u7xj, E7vk and tyev)
• A lack of analysis on potential energy and physical validities of JAMUN samples (Reviewers E7vk and tyev)
• An ablation of the jump step (Reviewer E7vk)
• A lack of results on larger systems (Reviewers iqm5 and tyev)

We have now performed additional experiments and analyses which we believe considerably improve our work addressing many of the concerns above. In short, we have added a fair comparison to a diffusion model with the exact same architecture, performed a Posebusters analysis on JAMUN samples, identified the effect of the jump step in terms of Jenson-Shannon divergence, and finally, demonstrated that our model can be trained on the fast-folding proteins ($\approx$ 50 AA) and the shortest ATLAS proteins ($\approx$ 180 AA), which are an order-of-magnitude larger than the original systems studied in the paper. Some of these experiments are preliminary due to the limited amount of time available to us, but we believe that the reviewers will find our results very encouraging.

Comparison with Diffusion Models

To highlight the speedup of walk-jump sampling over diffusion, we train a diffusion model on the Timewarp 2AA dataset, as recommended by Reviewer tyev. Note that our model already contains noise conditional blocks so this requires no architectural changes. We then sample from this model using both diffusion and walk-jump and compare the Jensen-Shannon divergence (JSD) of the backbone torsions averaged over the Timewarp 2AA test set as a function of number of samples and number of function evaluations (NFE):

Essentially, we find that walk-jump sampling is $64\times$ faster with only minor loss in fidelity. We would like to emphasize that the diffusion results can likely be improved by additional tuning of both training and sampling. For simplicity we decided to start with a reasonable setup that is known to work well in other settings (EDM noise schedule of 64 steps from 0.01A to 10A using the Heun second order method from Karras, et al. 2022). However, even with additional optimization of the diffusion model, walk-jump is likely to remain more efficient because it works in a partially noised spaced instead of having to generate every sample from an uninformative Gaussian prior over many steps.

We would also like to highlight that the TBG model (Klein, et al. 2024) that we compare to in the paper is a flow-matching model with a Gaussian prior distribution, which is functionally equivalent to a diffusion model.

Physical Validity and Energy Analysis of JAMUN Samples

We have performed a physical validity analysis on the generated JAMUN samples on MDGen 4AA, using the popular Posebusters (Buttenschoen, et al. 2024) package. We randomly selected 20 unseen test peptides from the MDGen 4AA-Explicit dataset for this analysis.

We see that the bond lengths are correctly captured with high probability by JAMUN (as requested by Reviewer E7vk). Further, the overall quality of the JAMUN samples is high. For a finer-grained view into the performance across the 20 test peptides, we report the empirical CDFs of the pass rates:

Next, we compute the force field energies for JAMUN samples, and find that they overlap well with the reference MD:

Ablation of the Jump Step

As requested by Reviewer E7vk, we show the JSD of the first two TICA components using noisy Y-samples generated during the walk step before denoising. We compare this with the JSD of the first two TICA components in X-space. As we expect, at a non-trivial noise level of 0.4 Angstrom, the Y-samples are quite diffuse and hence have far higher divergence metrics. In fact, the modes present in X-space are not easy to recover. This analysis highlights the necessity of the jump step to accurately sample.

Results on Larger Systems

We recently identified the memory usage of the standard e3nn.FullyConnectedTensorProduct in the E3ConvNet as a big bottleneck when scaling to larger systems. We switched to the SeparableTensorProduct (as in Equiformer) to reduce the number of weights in the tensor product and hence, the memory utilization. We found that the MSE denoising loss of the model increased only slightly after this switch. However, we can now train and perform inference on much larger systems.

We have trained one model on the fast-folding proteins TrpCage (20 AA, 152 atoms) and Protein G (56 AA, 439 atoms) from Majewski, et al. 2023. We have also trained one model on the shortest 1000 proteins (upto 183 AA, 2000 atoms) from the ATLAS dataset (Y Vander Meersche, et al.). These models were trained on 4 RTX A100 GPUs for a day.

Importantly, we do not change any of the training or sampling hyperparameters (including the noise level, which we keep at 0.4 A).

For the 20AA protein Trp-Cage which has high helical content, we start a simulation at a fully extended conformation and recover the secondary structure in a sampling run of wall clock time around 5 minutes. This would take more than a day of simulation time for a classical MD simulation running on the same hardware. For the ATLAS model, we were able to successfully run walk-jump sampling on unseen validation set proteins of slightly larger lengths than the training set proteins, indicating transferability even in this setting.

We compute some simplistic metrics - fluctuations in native contacts and secondary structure content - and find that we recover the qualitative characteristics of the distributions. The PMFs on fraction of native contacts look qualitatively similar for all three proteins, but there is definitely a bias towards folded states in JAMUN. The correlation plot between MD and JAMUN values for total fraction of time that a native contact persists has an $R^2$ of 0.96, with JAMUN contacts consistently remaining for longer times. In summary, more training is definitely required, but we hope that these preliminary experiments show that JAMUN can be scaled up to biologically relevant molecules.

Response to Reviewer u7xj

1. and 3. The Walk-Jump process corresponds to performing Langevin dynamics in the smoothed $Y$ space, just like how classical MD performs Langevin dynamics in clean $X$ space. Langevin dynamics is an approximation for the friction and thermal fluctuations induced by solvent particles: https://en.wikipedia.org/wiki/Langevin_dynamics#Overview.
2. We usually train on the entire MD trajectory data. Our results on cyclic peptides in Appendix J are on post-processed equilibrium states.
3. Note that we perform Langevin dynamics to sample $Y$ at a fixed noise level. This is different from the standard diffusion SDE where the noise level is continuously reduced during sampling.

We would be happy to update Figure 1) for clarity to show the sampling scheme. We would also add our analysis that shows that diffusion models spend many sampling steps due to sampling from an uninformative Gaussian prior, whereas walk-jump sampling operates from a smoothed, partially noised distribution which avoids the need to sample many steps.