Temperature-Annealed Boltzmann Generators

Thank you very much for your detailed and helpful feedback. We will address your points in the following:

• Comparison to FAB in accuracy and metrics: While Table 1 suggests TA-BG mainly improves in terms of the number of target evaluations, it’s important to note that the high-energy metastable regions constitute only a small part of the state space. Small NLL differences can thus have significant implications. For example, in the hexapeptide system, TA-BG and FAB show similar NLLs, but Ramachandran plots (Fig. 3) reveal that FAB fails to resolve the high-energy metastable region, while TA-BG succeeds. Therefore, quantitative metrics should be complemented with system-specific analyses, such as Ramachandran plots.
• Comparison to FAB in computational cost: We want to briefly discuss why our method is substantially more efficient than FAB. FAB uses annealed importance sampling (AIS) with intermediate Hamiltonian Monte Carlo (HMC) transitions to evaluate the α=2 divergence loss. While this yields accurate results (except for the hexapeptide), the AIS costs a significant amount of potential energy evaluations. Our method simply uses the reverse KLD at high temperature, which is significantly less costly to evaluate. While the annealing adds some additional cost, it is still very efficient since large buffers are used, where training samples are reused several times. Therefore, in total, our method is significantly cheaper than FAB.
• Scaling to larger systems: First, we would like to emphasize that scaling to the hexapeptide is already a big achievement in itself, as the previous SOTA method FAB was the only method that worked even on the smallest of our three systems studied. Despite this, we are confident that our method scales to larger systems. While it is true that the ESS drops down for the larger systems, this can be counteracted with a more expressive architecture, e.g., the one recently proposed in [3]. At the same time, even with the current flow architecture, one can counteract the drop in ESS with other measures:
  • Increase in the number of temperature steps or the number of drawn samples per annealing step, see our ablations in the answer to reviewer SzWG.
  • Additional fine-tuning steps, see our ablations in the answer to reviewer 1DKq.
  • One can use AIS instead of IS to estimate the loss in the annealing steps. While this comes at an additional cost, it can keep the variance of the loss low when increasing the intermediate AIS steps linearly with the dimensionality (see [4] for a theoretical analysis).
• Your question about data-driven vs data-free sampling of Boltzmann distributions is very interesting. First, we want to point out that the task of sampling the equilibrium Boltzmann distribution of proteins is significantly harder than just predicting folded structures, as is done in AlphaFold. While the references [1] and [2] that you cited tackle the task of equilibrium sampling, the results are only rough approximations of the true equilibrium distribution. Even though large MD datasets have become available recently, we believe that much more data is necessary to achieve good transferability that can replace direct sampling through MD or variational methods, such as the one proposed in our work.
  Another recent publication [5] trained transferable Boltzmann generators on a large dataset of dipeptide MD simulations, which is a very narrow chemical space. While the transferability was successfully shown, also here the transfer to new systems only resulted in a rough estimate of the true equilibrium distribution.
  We thus believe that a hybrid data / variational approach could be pursued in the future. In our method, MD data can be used to supplement the high-temperature pre-training, which will make finding new modes easier and may lead to quicker convergence. We agree that this discussion is very important and we will include a discussion of these aspects in our revised manuscript.
• The reported ESS in Table 1 is the ESS when sampling from the flow distribution (which has been annealed through multiple iterations to 300K) in the very end, reweighting to the 300K target.

We further note that we have now performed several new ablation studies, as detailed in our responses to reviewers SzWG and 1DKq. Lastly, we now additionally benchmarked against the 2D GMM system (see our rebuttal to reviewer CidM), where TA-BG outperforms FAB and all diffusion-based baselines.

With this, we hope you agree with us that our work is a significant achievement that the sampling community can build upon in the future. We hope to see more methods applied to the complex benchmark systems proposed in our work, while we are eager to scale our method to even larger systems next.

[3] Zhai et al. 2024, "Normalizing Flows are Capable Generative Models"
[4] Midgley et al. 2023, “Flow Annealed Importance Sampling Bootstrap”
[5] Klein et al. 2024, “Transferable Boltzmann Generators”

Thank you for your detailed and helpful feedback. We address your questions and concerns below:

• While Table 1 suggests TA-BG mainly improves in terms of the number of target evaluations, it’s important to note that the high-energy metastable regions constitute a small part of the state space. Small NLL differences can thus have significant implications. For example, in the hexapeptide system, TA-BG and FAB show similar NLLs, but Ramachandran plots (Fig. 3) reveal FAB fails to resolve the high-energy metastable region, while TA-BG succeeds. Therefore, quantitative metrics should be complemented with system-specific analyses like Ramachandran plots.
  In summary, TA-BG achieves better NLL across all systems, needs significantly fewer target evaluations, and is the only method to capture the complex interactions in the hexapeptide.
• FAB is, indeed, the only method successfully applied to learning the Boltzmann distribution of alanine dipeptide without prior knowledge (as also confirmed by Reviewer SzWG in his review). Of the works you cited:
  • [1] and [2] discuss transition path sampling between points A and B, which is a different task than sampling the full Boltzmann distribution.
  • [3] relies on data from the target distribution; we cover the data-free case.
  • [4] depends on a low-dimensional collective variable, typically non-trivial to obtain.
    We will cite these papers and discuss that they are related but ultimately solve a different task.
• While diffusion samplers have gained popularity, they have not yet been successfully applied to alanine dipeptide. Most diffusion sampling studies use synthetic benchmarks with less correlated transitions between minima than the molecular systems we examine. In our reply to reviewer 1DKq we further discuss why the necessity of using Cartesian coordinates might make applying diffusion models to sample molecular systems harder, but we can mostly speculate here. As reviewer 1DKq points out, diffusion models are "not quite there yet", but we hope to see a successful publication on this topic soon. FAB and our method remain the only approaches that scale to system complexities like alanine dipeptide, which we now significantly extended to the tetrapeptide and hexapeptide.
• We now include the 2D GMM system (as introduced by FAB). Since diffusion methods have been successfully applied here, this allows a direct comparison. We pre-trained the flow at T=30.0 with reverse KLD, annealed to T=1.0 in 7 steps, and performed one fine-tuning step. Results are shown in this figure (neural splines: https://ibb.co/7dkKCr76) and table (https://ibb.co/bgrDkVyD), using values reported by iDEM [5]. We note that we used two different flow architectures, REAL NVP (as originally used in FAB) and an improved neural splines architecture. We did not have enough time in the rebuttals to repeat the FAB experiments with neural splines, but we will prepare this for the revised manuscript (and also include error bars for our method). As shown, TA-BG outperforms FAB and all diffusion baselines on the 2D GMM system. We will further explore other sampling tasks, such as DW-4, and include them in the appendix.
  We emphasize our focus on molecular systems, which present more challenging and realistic benchmarks than commonly used synthetic systems. Our goal is to motivate the application of sampling methods, including diffusion, to these more difficult and application-relevant domains, supported by our release of the two new benchmark systems.
• We approximated training time (excluding evaluation) for FAB and TA-BG on the alanine dipeptide system. On our hardware, FAB takes ~18.2h, TA-BG requires ~17.1h. Thus, runtimes are comparable. Furthermore, inference time is identical, since the same model is used for all experiments.
  Force field evaluations in our current setup are relatively inexpensive. With more accurate and computationally costly forces, such as ML-based foundation models or DFT, the evaluation cost becomes dominant. In such cases, the higher sampling efficiency of our method will translate into significantly reduced computational cost.
• Thank you for pointing out the missing standard sampling and diffusion references; we will revise the manuscript accordingly.
• We have performed several new ablation studies (see responses to reviewers SzWG and 1DKq).

With this, we hope that we were able to convince you that our method and results are significant achievements and interesting to the sampling community. We hope to see more sampling approaches, including diffusion models, being applied and benchmarked on the complex and more application-relevant molecular systems proposed in our work.

[5] Akhound-Sadegh et al. 2024, “Iterated Denoising Energy Matching for Sampling from Boltzmann Densities”
[6] Durkan et al. 2019, “Neural Spline Flows”

Thank you for your positive, detailed, and constructive feedback.

• We address your points following the same numbering as in your review: 2. We will more clearly outline the advantages of our method. For example, explicitly stating the mass-covering nature of the forward KLD and why this is helpful in our case is a great suggestion. We will generally work on making our introduction of forward / reverse KLD more intuitive. 3. Thank you for these suggestions to improve the presentation of our method. We will work on these points to revise our manuscript. We will include the equation $T_i = T_\text{start} \left( \frac{T_\text{target}}{T_\text{start}} \right)^{\frac{i-1}{K-1}}$ to clarify the “geometric temperature progression”. Further, we will improve the explanation of the annealing procedure, including the fine-tuning and its motivation. 4. Choosing hyperparameters for a fair comparison of the different methods is not straightforward, as the final metrics and the total number of target evaluations form a tradeoff. Thus, you are right that saying our method “requires” a certain number of evaluations is imprecise; we will correct this.
  To better make the tradeoff visible, we included FAB hyperparameter variations in Appendix H of our original submission. Furthermore, we now performed ablation studies to make this tradeoff better visible for our method, see our answer to reviewer SzWG, and the fine-tuning ablation study discussed below. 5. One of the reasons why using Cartesian coordinates might make the problem harder is that with internal coordinates, we can constrain the angles and bond lengths to reasonable bounds. This is not possible with Cartesian coordinates, where even completely different chemical graphs can be formed. We suspect this is partly why diffusion-based sampling methods have not yet been successful in modeling molecular systems, since handling periodic torsions is less straightforward with diffusion than with flows, and thus Cartesian coordinates are often used. 6. Not clipping the importance weights results in numerical outliers visible in the reweighted Ramachandran plots. Therefore, we clipped the weights when determining the KLD of the Ramachandran plots, analogous to FAB. We further ran ablation experiments for alanine dipeptide to determine the impact of clipping the importance weights when resampling the buffer datasets during the annealing. The final NLL and ESS stayed within our error bounds, so clipping has little impact here. We will mention this in the revised manuscript.
• Further responses to your additional suggestions and questions: 7. We performed an ablation study to show the impact of fine-tuning. First, we look at the impact of the intermediate fine-tuning for the hexapeptide system. In the following figure, we show the ESS of the training buffer dataset in each iteration for the case with and the case without intermediate fine-tuning (both experiments have a final fine-tuning step): https://ibb.co/XxZJvr0r
  The buffer ESS of the fine-tuning steps that follow each annealing step is much higher, which is reasonable since we here reweight to the same temperature the flow was previously annealed to ($T_{i+1}=T_i$). One can further see that the buffer ESS generally drops down when not using intermediate fine-tuning. We summarize the impact of intermediate fine-tuning on the final metrics for all systems in this table (https://ibb.co/hjXk96b). As you can see, the impact is the largest for the hexapeptide, which is why we included intermediate fine-tuning only in this system for our main experiments.
  We further summarize the impact of the final fine-tuning step for all systems in the following table: https://ibb.co/XfcbxGb3 8. Testing other objectives (such as the α=2 divergence) for the annealing is a great idea. However, we believe this goes beyond the main claims we want to cover. We will add a short discussion of other possible objectives to our manuscript. 9. Thank you for pointing out the concurrent work [1]. We will mention it in our revised manuscript. 10. We will mention more advanced flow architectures as a path for the future [2,3]. 11. Thank you for your additional suggestions to improve our manuscript. We will include them in the revised version.

In addition to the molecular systems studied in our work, we now additionally benchmarked against the 2D GMM system (see our rebuttal to reviewer CidM), where TA-BG outperforms FAB and all diffusion-based baselines.

Overall, your comments were very helpful in improving our paper. Our approach outperforms the current SOTA (FAB) in this area and demonstrates scalability to more complex systems. We believe that this will be very helpful for the community to advance in this task, and we would appreciate it if you could revise your score based on our revision.

[3] Tan et al. 2025, “Scalable Equilibrium Sampling with Sequential Boltzmann Generators”

Thank you for your detailed feedback and the many helpful suggestions. We address them in the following:

We performed ablation studies to investigate the impact of hyperparameters, which we will include in our revised manuscript:

• We ran ablation experiments to determine suitable starting temperatures for each of the systems. At a given temperature, we performed 4 experiments to determine how many of these experiments show mode collapse in the Ramachandran plots when pre-training with the reverse KLD:
  https://ibb.co/5XgHtZYM
  There is a large range of temperatures where reverse KLD training is possible without mode collapse. This ablation was performed with a batch size of 1024 and 250k gradient descent steps. Increasing the temperature beyond the “critical temperature” where the experiments start not collapsing anymore allows us to pre-train cheaper with a smaller batch size and fewer gradient descent steps without mode collapse. A tradeoff exists: When increasing the temperature beyond the “critical temperature”, one can pre-train cheaper, but then the annealing will be more expensive.
  This ablation nicely extends our mode collapse discussion in the paper. The fact that alanine dipeptide can be sampled without mode collapse with the reverse KLD at only slightly increased temperature (375 K) is likely interesting and surprising to the sampling community in itself.
• We ran experiments to justify the geometric temperature schedule. We compare the geometric schedule with a linearly spaced temperature schedule. In the following figures, we show the ESS of each buffer dataset we train on during the annealing:
  https://ibb.co/TDfMwLbs (alanine dipeptide, one final fine-tuning step)
  https://ibb.co/gbBHb66W (hexapeptide, with intermediate fine-tuning steps)
  The buffer ESS “spikes” visible for alanine dipeptide in the end and for the hexapeptide in-between come from the fine-tuning steps, where $T_{i+1}=T_i$, yielding better overlap. One can see that the geometric temperature schedule yields an approximately constant buffer ESS, whereas the ESS drops down significantly toward the end for the linear schedule (since the temperature steps are too large!).
• Regarding details on the fine-tuning steps, these are performed with the same parameters as the annealing steps. The only difference is that we do not decrease the temperature, but rather reweight to the same temperature the flow was annealed to in the previous step, so $T_{i+1}=T_i$. We will make this clearer in the revised version of our manuscript.
  We further created an ablation to show the impact of fine-tuning. You can find details on this in our reply to reviewer 1DKq. To summarize, the final fine-tuning step is important for all systems, while the intermediate fine-tuning helps keep the ESS up for the hexapeptide.
• For alanine dipeptide, we ran ablation experiments to show the impact of the number of temperature steps on the final metrics. We summarize the results in the following table: https://ibb.co/NgPHGG9J
  As you can see, a tradeoff between increasing the number of intermediate temperatures (which yields more target evaluations) and improving the final metrics exists.
  In the following figure, we further show the buffer ESS over the course of training: https://ibb.co/Z6VGsLFj
  One can see that with smaller temperature steps, the buffer ESS during training increases.
• For alanine dipeptide, we ran ablations to show the impact of the number of samples drawn in each annealing iteration (https://ibb.co/jxwVY4M). Again, a tradeoff is visible between the number of target evaluations and the final NLL / ESS.

Furthermore, thank you for suggesting the additional baseline of reweighting directly from 1200 K to 300 K. You can see what this looks like (using $10^6$ samples) in the following figure: https://ibb.co/7JHtBb2v
We will include this additional baseline in our main text.

In general, the ablations show that a drop in ESS can be counteracted by using more intermediate temperatures, drawing more samples, or with additional fine-tuning steps. We are thus confident that we can scale our method to even larger and more complex systems, but we believe that already scaling to the hexapeptide is a significant milestone that we want to share with the community.

In addition to the molecular systems studied in our work, we now additionally benchmarked against the 2D GMM system (see our rebuttal to reviewer CidM), where TA-BG outperforms FAB and all diffusion-based baselines.

Again, thank you for your helpful comments, we believe to have addressed all suggestions and concerns of your review. The suggested ablations shed light on interesting dependencies present in our method. We are confident that our results, including the suggested ablations, are interesting for the sampling community and provide a foundation for future work. We hope that you can raise your score accordingly.