Response to Reviewer qKai

We would like to thank the reviewer for the time and effort they spent on reviewing our work. We are appreciative of the fact that the reviewer found ENSEMBLE demonstrates "strong performance in equilibrium sampling tasks on unseen peptide systems", and suggests "good generalization capabilities". We now address the concerns raised by the reviewer.

Baselines

However, many recent deep learning methods developed for molecular simulation and conformation sampling, such as Timewarp [1], ConfDiff [2], MDGEN [3], AlphaFlow [4], UniSim [5], etc., are not considered. A comparison with these methods is necessary to substantiate the claim.

We thank the reviewer for suggesting these additional baselines on conformation sampling. We include additional baselines in the table below, using the same 1e4 energy evaluation budget as in Table 1 of the main paper.

In each case we use the pretrained model weights and sampling configurations provided by the original authors. In the case of Timewarp the model only supports peptides of lengths 2AA and 4AA hence no results are presented for 8AA. Following the authors procedure, for BioEmu the sidechains are added using hpacker followed by energy minimization. Additionally, for UniSim we employ the energy minimization step during the trajectory generation. In both cases a sweep was performed to optimize the energy minimzation budget allocated per-sample with the number of samples generated.

We would like to emphasize that only Ensemble (with SNIS), ECNF (with SNIS), and Timewarp (with Metropolis-Hastings acceptance) are able to draw asymptotically unbiased samples from the target distribution. Empirically we observe an a significant improvement of ENSEMBLE over baselines with respect to all metrics, except the macrostructure (Torus-W2 and TICA-W2) metrics on the largest systems (8AA). We however, emphasize that the E-W2 achieved by Ensemble to be significantly lower on the 8AA systems than the next best method.

Practical applicability

Moreover, the lack of validation on larger and well-studied molecular systems undermines the credibility of the practical applicability of the method.

We would like clarify that these are by far the largest systems that have been studied in a transferable setting for Boltzmann generator type models - a 4-fold increase given the only prior transferable Boltzmann generator operated on dipeptides only. We believe our work therefore significantly enhances the practical value and impact of this model class. While Boltzmann Generators as a model class are inherently less scalable than vanilla generative models, they offer the key advantage of exact likelihoods for rigorous statistical re-weighting. We argue that solving the fundamental challenge of length-transferability for peptides up to 8 residues is a critical and necessary step forward, establishing the foundation needed before this rigorous approach can be scaled to larger, canonical protein systems. Our work pushes the practical frontier for this entire class of models, making the future application to larger systems a more tangible research goal.

Evaluation

please report the ESS per second (ESS/s) for each method (including MD trajectories)

In the following table we report effective samples per second for ENSEMBLE and baseline methods. We will report ESS/s for MD trajectories and the additional baseline methods during the discussion period.

This raises the question of whether the observed performance gains are due to increased model size rather than architectural improvements.

We thank the reviewer for their concern. We would like to clarify that in our view the model size increase is actually a substantial benefit of ENSEMBLE over previous Boltzmann generator type approaches. Specifically, because our likelihood computation is efficient, we can scale the model significantly more without computational difficulties that might severely limit sampling efficiency. We further draw attention to the sampling wall-time results presented in Figure 22 of the supplementary material, which indicate ENSEMBLE to be significantly faster to sample from than the ECNF, with diminishing increase in wall-time over the standard TarFlow as the system size is increased.

Chignolin

Following the reviewer's suggestion, in the table below, we report the performance of the method on the 9-residue peptide YQNPDGSQA introduced by [1] and the well-studied Chignolin small protein. The reported metrics demonstrate that the proposed method outperforms all of the baselines. We further explore an additional step in the self-refinement procedure where the first round of proposal samples are not reweighed using SNIS, but instead undergo constrained energy minimization before they are used for self-refinement training. The self-refined model is then used to generate proposal samples which are reweighed using SNIS. Note that the training data only contains systems of length up to 8 amino acids, making clear the ability of Ensemble to extrapolate beyond it's training data distribution, despite the baselines being trained on larger systems.

Novelty

We thank the reviewer for this crucial feedback. We argue our technical novelty lies not in creating a flow from scratch, but in solving a long-standing challenge: enabling a single, all-atom model with efficient likelihood calculation to transfer zero-shot across variable-length peptides

This required novel architectural solutions for dimensional transferability and adaptive conditioning, which are non-trivial extensions that allow a Transformer-based flow to generalize across system sizes for the first time. The significance of this breakthrough is validated by our empirical results, where ENSEMBLE unlocks a previously unattainable capability and achieves a >4000x speedup over transferable baselines. We will revise the paper to frame our contribution more clearly around solving this specific, challenging problem.

Clarity

Could the authors clarify which forcefield was used in this computation?

The amber-14 forcefield was used through the OpenMM package. For details on the simulation please see Table 5 in the appendix.

Yes, energy computation has a significant impact on the model's inference speed, this is why we set a fixed number of energy evaluations budget across all methods. While we use a relatively efficient energy function here, we note that our method can easily be adapted to any energy function, some of which may be significantly more expensive than even our fairly large ENSEMBLE model to run. Given the favourable performance of ENSEMBLE with respect to energy evaluations demonstrated in Figure 1, any increase in energy function complexity can be anticipated to translate into further-still wall-time performance advantages.

Concluding remarks

We thank the reviewer again for their time. We believe we have answered, to the best of our ability, all the great questions raised by the reviewer. We hope our answers allow the reviewer to consider potentially upgrading their score if they see fit. We are also more than happy to answer any further questions.

[1] Xiongwu Wu and Bernard R Brooks. "Beta-hairpin folding mechanism of a nine-residue peptide revealed from molecular dynamics simulations in explicit water." Biophys J. (2004)

Thank you for the prompt response. Most of my concerns have been addressed, and I have a few minor suggestions remaining:

• I agree with the author’s explanation that the transformer architecture offers better scalability compared to GNN-based models. I hope the author will later support this point with actual references and evidence.

• The author mentions in the rebuttal that the performance of the ECNF model based on the GNN architecture will degrade with larger systems. However, from Table 2, it can be observed that the performance of ENSEMBLE also deteriorates as the system size increases. Additionally, the experimental results for ECNF on 4AA and 8AA are missing in Table 2, which makes it difficult to compare the performance degradation of ECNF and ENSEMBLE on larger molecular systems. This undermines the necessary argumentation for the author’s viewpoint. I believe this point requires further clarification.

Response to Reviewer 7hoH

We thank the reviewer for their time and effort in reviewing our work. We are pleased that the reviewer describes this paper as "the first demonstration of a transferable Boltzmann generator" and as "an important milestone in the field." In what follows, we address the concerns raised and provide suggested evaluations.

Figure 1. If test peptides have been simulated for up to 1us, then the MD results should be available for wall-clock times up to O(10) GPU-hrs, but authors have concealed this part of the curve.

Figure 1 does not contain MD trajectories up to 1$\mu$s; instead, the baseline MD trajectories there are simulated up to a fixed energy evaluation budget ($10^6$ energy evaluations). Note that simulating 1$\mu$s would require $10^9$ energy evaluations. The long MD simulations (1$\mu$s long) are used only as ground truth samples for metrics evaluation, and use a less frequent frame storing interval to prevent excessive data accumulation (making them unsuitable for use as a baseline across many orders of magnitude of energy evaluation budget). To avoid potential confusion, in the text, we will refer to different types of MD trajectories as train MD, eval MD, and ground true MD, which are used, respectively, as training data, as a baseline, and as the ground true distribution of samples.

However, we agree with the reviewer that the "stopping" of the MD baseline in Figure 1 is a missed opportunity for comparison between our method and MD at larger wall-time budgets. Hence we have rerun the MD baseline using a decaying frame storing interval up to 1e8 energy evaluations. Given we cannot communicate an updated figure, in the below table we present the metrics at various wall time budgets, in comparison to ENSEMBLE, where a spline was fit to enable arbitrary wall times to be queried.

It is very hard to conceptualize TICA and T-W2s. Authors should use JSDs or build markov state models and report recovery of markov states.

Following the reviewer's suggestion, we're currently running the evaluation using JSD and will share the updated results during the discussion period.

With very small ESS, it is unclear if SNIS is effective at estimating observables better than the unweighted proposal with small sample sizes. Authors should also report results without reweighting.

The suggested comparison is provided in Table 10 (Appendix C.1). This comparison demonstrates that the proposal samples multiple outliers (unphysical structures having very high values of energy), which are efficiently filtered out via SNIS.

It seems that the model still generates many high-energy samples. Authors should report statistics - what are the mean, median, 95th percentile energies of the unweighted samples? How does this change from tetrapeptides to octopeptides? This is important to understand the further scalability of the method.

We thank the reviewer for the excellent point. We provide the requested statistics for the unweighted proposal samples for sequence lengths of 4 and 8.

The results highlight a notable increase in energy with longer sequences, reflecting the greater complexity of the larger system. We believe this further motivates the need for reweighting strategies at higher sequence lengths.

It would be helpful to understand the impact of scaling the training data from 1k to 10k training peptides. There is a missed opportunity to explore scaling laws wrt both longer simulations and more training peptides.

We thank the reviewer for their suggestion. We agree that this could be quite useful for future work. Given the computational requirement of proper scaling plots in this domain, we will add this experiment in the camera-ready version of the paper.

Authors should show results on chignolin - the only peptide of nontrivial length considered to have been simulated to convergence - since it is a common reference point with many other works.

Following the reviewer's suggestion, in the table below, we report the performance of the method on the 9-residue peptide YQNPDGSQA introduced by [1] and Chignolin. The reported metrics demonstrate that the proposed method outperforms all of the baselines. Note that the training data only contains systems of length up to 8 amino acids, which suggests that the method generalizes across systems of previously unseen length.

We thank the reviewer again for the detailed feedback and constructive suggestions on the improvement of the manuscript. We hope we address their main concerns, and we remain available throughout the discussion period for further clarifications and suggestions for improvement.

We sincerely thank the reviewer for their thoughtful feedback and for their willingness to raise their score based on the additional results. We appreciate the constructive suggestions, which have helped strengthen the quality of our work.

Upon receiving those results, and if authors agree to include all rebuttal results in the revision

We confirm that we will include all rebuttal results in the revised version of the manuscript. We thank the reviewer for emphasizing this point and for encouraging a more complete presentation of our findings.

I look forward to seeing the JSD and MSM results in the discussion period

As requested, we report below the Jensen-Shannon Divergence (JSD) results for the 2AA, 4AA, and 8AA systems. These results assess the recovery of discrete states by our generative model compared to reference MD simulations. We follow the procedure outlined in [1]. Specifically:

• We perform k-means clustering on the reference MD trajectories using the top 2 time-independent components (TICs), which are computed from a featurization of pairwise distances and dihedral angles.
• The MD trajectories are clustered into 20 discrete states.
• Using the fitted k-means model, we assign cluster labels to both the reference MD and the generated samples.
• We then construct histograms over these cluster assignments and compute the JSD between the reference and generated distributions.
• We report this JSD metric for all methods, including new baselines requested by reviewer qKai in the following table, where all methods are given a budget of 10^4 energy evaluations.

This metric aligns with the results of Table 2 in the main paper, in which ECNF (trained only on 2AA) is the strongest method on the simplest (2AA) test dataset, but Ensemble is the strongest method on the more challenging 4AA and 8AA test datasets. The outperforming of ECNF on the 2AA dataset can be attributed to the simpler task (without sequence length transferability), and we note the ECNF++ (trained on 2AA -> 4AA) to be notably inferior on both 2AA and 4AA compared to Ensemble. These results will also be included in the updated manuscript.

Concerning the reviewers request of MSM results, while exploring this metric we found that it is not possible to compute an MSM metric for Ensemble, Tarflow, ECNF, and Boltzmann Generator methods in general. The work of [1] reports two additional metrics based on the MSM transition matrix fit on these discrete states: “Transition Negative Log Likelihood” and “Fraction of Valid Paths”. We clarify that Ensemble (as well as TarFlow and ECNF), generate uncorrelated proposal samples which are then reweighted using e.g. SNIS. This implies there is no notion of “transitions” between states, and therefore it is not possible to compare to a reference transition matrix as in [1]. This is in contrast to methods such as MD, TimeWarp, UniSim, which generate samples in a time-sequential manner and hence could be evaluated in such a way.

We once again thank the reviewer for their time reviewing our work, and their insightful suggestions for improvement.

[1] Raja, Sanjeev, et al. "Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional." arXiv preprint arXiv:2504.18506 (2025).

Response to Reviewer VbSw

We thank the reviewer for their detailed feedback and their time spent. We appreciate their positive evaluation, in particular, the reviewer describes the proposed approach as "scalable and transferable", the empirical study as "extensive", and the presentation as "clear and well-organized". In what follows, we provide requested clarifications.

(Quality, minor) Performance on 2-residue (2AA) systems. As acknowledged by the authors, the method underperforms on 2AA systems compared to other lengths. While this is a minor issue and does not detract significantly from the overall contributions, it may indicate areas where the model's generalization could be improved.

The fact that Ensemble slightly underperforms on 2-residue systems is the consequence of ECNF being trained only on 2 residues. Indeed, ECNF++ -- our attempt to scale this approach to 4 residues (following the improved training techniques from [1]), demonstrates worse performance across both 2 and 4 residues.

(Quality, minor) Efficiency claim could be better quantified. While the method outperforms baselines under fixed energy evaluation budgets, adding a simple experiment showing performance as a function of energy evaluations (e.g., performance-vs-budget curves on one system) would strengthen the claim of efficiency and allow for better interpretability of trade-offs

In Fig. 1 (page 2), we report the performance (across 30 unseen 4-residue systems) as a function of energy evaluations and GPU walltime. Notably, the proposed method outperforms Molecular Dynamics (MD) simulations for the same evaluation budget, as the latter fails to capture all metastable states (see the 3rd plot in the top row).

Minor issues

We thank the reviewer for a thorough reading of the manuscript. We will fix those typos.

We thank the reviewer again for the detailed feedback and constructive suggestions on the improvement of the manuscript. We hope we address their main concerns, and we remain available throughout the discussion period for further clarifications and suggestions for improvement.

[1] Tan, Charlie B., Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. "Scalable equilibrium sampling with sequential boltzmann generators." arXiv preprint arXiv:2502.18462 (2025).**

Response to Reviewer EUPF

We thank the reviewer for their time and effort in reviewing our work. We are glad that the reviewer describes the scale and generalization to other molecular systems as "notable", the empirical evaluation as "strong" and "significant", and the presentation as "clear". In what follows, we address the concerns raised and provide suggested evaluations.

More detailed experimental analysis on all of the modifications in Section 3.1 and each of their roles would help further readers to understand the roles of each individual component of the proposed architecture. What is the impact of each of the proposed modifications for the resulting architecture? Which components had the most impact for the empirical results?

In Section 4.3, we describe the ablation study of all the architectural changes introduced in our work. In particular, Table 3 compares the results of the original TarFlow architecture with the proposed architecture that includes the "adapt + transition" blocks as well as the proposed chemistry-aware permutations. We observe that both changes provide a performance boost; notable in particular given the chemistry-aware permutations imply no increase in training or inference wall-time. The walltime increase of "adapt + transition" is only pronounced on the smaller systems and is not significant on the 8AA dataset (Figure 22, supplementary material).

What practical obstacles arise from either extending to or evaluating on systems beyond octopeptides?

Following the reviewer's suggestion, in the table below, we report the performance of our method on the 9 residue peptide introduced by [1] and the well-studied small protein Chignolin. The reported metrics demonstrate that the proposed method outperforms all of the baselines. Note that the training data only contains systems of length up to 8 amino acids, which suggests that the method generalizes across systems of previously unseen length.

We further explore an additional step in the self-refinement procedure for Chignolin, where an initial set of 1000 samples from Ensemble undergo constrained energy minimization with restraints on the macrostructure. The resulting configurations are then used to finetune our model before running the self-refinement training.

We thank the reviewer again for the detailed feedback and constructive suggestions on the improvement of the manuscript. We hope we address their main concerns, and we remain available throughout the discussion period for further clarifications and suggestions for improvement.

[1] Xiongwu Wu and Bernard R Brooks. "Beta-hairpin folding mechanism of a nine-residue peptide revealed from molecular dynamics simulations in explicit water." Biophys J. (2004)