Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities

We thank the reviewer for the detailed and constructive feedback. We greatly appreciate the positive evaluation of our work. In particular, the reviewer finds the proposed method "creative", "interesting", "addressing many current questions" and addressing "important challenges". In what follows, we address the concerns raised and provide additional evaluation requested by the reviewer.

Dataset benchmarks

Dataset benchmarks. It would be useful to include at least one experiment on a synthetic dataset such as a mixture of 50 Gaussians in 2 dimensions.

We thank the reviewer for their suggestion. We now include in this rebuttal the comparison on the synthetic Mixture of Gaussians dataset, training at $T_L=4$ and annealing to $T_S=1$. The comparison is presented in the table below, where we take the results for baselines from [1] and add the performance of PITA. Notably, PITA significantly outperforms all the considered baselines.

Table R3: Performance of methods on the mixture of 40 Gaussians in 2 dimensions.

Comparison with SMC in the experiments.

We thank the reviewer for suggesting the inclusion of a vanilla SMC baseline with a Langevin kernel and a classical annealing schedule. We have incorporated this baseline in our experiments, and the results are reported in Table R1 (see our response to Reviewer JqJ5).

Comparison with SMC in the text. Again, the most natural counterpart to the author's PITA algorithm is a basic SMC that follows a temperature-annealing path. It is also an algorithm most people in sampling are familiar path. Discussing how PITA and vanilla SMC differ would be very helpful to understanding PITA, as in the summary of this review.

Thank you for this constructive suggestion! We will add the corresponding discussion in the text.

Questions about Ablations

For inference-time annealing, is it really helpful to do run an SMC through the diffusion path from the previous temperature, raised to a certain power? Can we not just directly reweigh the samples from $\pi(x)^\beta$ to $\pi(x)^{\beta+ \Delta \beta}$ and then learn a diffusion model?

Following the reviewer's suggestion, we report the performance of inference-time annealing with importance sampling from $\pi(x)^{\beta+ \Delta \beta}$ using $\pi(x)^\beta$ as a proposal distribution, in Table R2 of our response to Reviewer Ldmh.

What is the difference in performance if PITA was run only with inference-time annealing and no model training?

The inference-time annealing procedure heavily relies on the parametric models and is designed to sample from the unnormalized density of the energy-based model, i.e. $q_t(x)\propto \exp(-\gamma U_t(x;\eta))$. Thus, sampling entirely without training would not produce a meaningful distribution. Although one could perform sampling without training on the samples from the intermediate annealed densities, this approach consistently underperforms (see ablation study in Appendix E, line 858).

Questions about Clarity of presentation

Some things that could be more clearly explained: at a given temperature $\beta$, there are two probability paths that are being used: first, the diffusion path from the previous temperature, raised to a new power ("inference-time annealing"), and then the diffusion path at the new temperature ("model training").

We thank the reviewer for their insightful suggestion. We agree with the reviewer that the proposed phrasing can help clarify the two probability paths used in the PITA algorithm. We will update the corresponding subsections of the paper to better emphasize this distinction.

why model training is useful in addition to inference-time annealing?

The inference-time annealing procedure is tightly coupled with the parametric models and is designed to sample from the unnormalized density of the energy-based model, $q_t(x)\propto \exp(-\gamma U_t(x;\eta))$. Training ensures that the parametric models capture these distributions accurately, which in turn makes the inference-time annealing effective. Jumping too far in temperature with inference-time annealing reduces performance of the method due to increased variance of the importance sampling weights. Model training is useful to "reset" these weights at a new lower temperatue.

Figure 1 is unclear to me. The two taxes of diffusion and temperature are very helpful. But the arrows are very confusing. I did not understand from the arrows that two probability paths are being used at a given temperature: the old diffusion path raised to a power ("inference-time annealing") and the new diffusion path ("model training").

We appreciate the reviewer's feedback. In the current version, the diffusion path is represented by the horizontal line, as the standard diffusion does not involve any change in temperature. In contrast, the inference-time annealing path changes the temperature of marginals, which is illustrated by the curved line leading to the marginal $\pi(x)^\beta$ at a higher value of $\beta$. We hope this clarifies the interpretation of Figure 1.

Lines 175-177: "the DSM objective is not sufficient for training a good score model close to the target distribution [...] it has no information about the target distribution". This is not very clear: near the target distribution, the score model in Eq 13 is trained using samples that are almost from the target distribution. How is this compatible with the statement: "it has no information about the target distribution"?

As $t \to 1$, equation (13) becomes

$\mathbb{E}_{x_t,x} \lambda (t) ||x-D_t(x_t;\theta)||^2$

$\to \mathbb{E}_{x_t,x}\lambda(t)|| \sigma_t^2 s_t(x;\theta)||^2$

because $\sigma_t \to 0$ and $x_t \to x$. This leads to unstable training of the score function $s_{t\approx 0}(x;\theta)$ with a finite dataset, as it converges to modeling the score of a uniform mixture of Gaussians with tiny variance around the dataset. This limitation of the DSM objective is well known, and we follow Target Score Matching [2] to address it. We will clarify this point in the text.

Lines 61-62: "temperature annealing mixes modes by lowering high energy barriers, while diffusion paths avoid mass teleportation." This is a strong statement which would benefit from citing a few references. While it may be clear to experts, it is not obvious to any reader. For example, if one has not done the derivations in closed-form of the diffusion path from a Gaussian to a mixture of Gaussians, it is not obvious that mass is not teleported.

This is a valuable suggestion. We will add the discussion of the original work on diffusion models with SDEs [3], which proposed the reverse-time SDE and probability-flow ODE, demonstrating that diffusion paths can be simulated without probability mass teleportation. As well as [4] which nicely contrasts temperature and diffusion interpolations visually.

Concluding remarks

We thank the reviewer again for the thorough review and valuable feedback which strengthens the clarity and results of our paper through new ablations. We hope that our rebuttal definitely clarifies the all the great points raised in the review, and we politely invite the reviewer to ask any further questions or potentially endorse this paper more strongly if they are satisfied.

References

[1] OuYang, RuiKang, Bo Qiang, and José Miguel Hernández-Lobato. "Bnem: A boltzmann sampler based on bootstrapped noised energy matching." arXiv preprint arXiv:2409.09787 (2024).

[2] De Bortoli, Valentin, Michael Hutchinson, Peter Wirnsberger, and Arnaud Doucet. "Target score matching." arXiv preprint arXiv:2402.08667 (2024).

[3] Song, Yang, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).

[4] Máté, Bálint, and François Fleuret. "Learning interpolations between boltzmann densities." arXiv preprint TMLR (2023).

Thank you authors.

I have one remaining concern about the experiment "Performance of methods on the mixture of 40 Gaussians in 2 dimensions".

Sometimes, the particles are initialized using a Gaussian distribution with a large enough covariance to "cover" all the modes. But this assumes prior knowledge on where the modes are located and makes the "search" problem easier.

Could the authors run the 40 Gaussians experiments, and initialize the particles with the density of one of the modes only (so as to not assume that the user knows where the other modes are). This should make the mode "search" problem harder and fairer and would better illustrate whether PITA better finds modes compared to DDS, PIS, FAB, and iDEM.

We thank the reviewer for their time and effort in assessing our paper. We are pleased to hear that the reviewer finds the key ideas in PITA of learning a sequence of diffusion models through annealing is "novel and conceptually solid". Moreover, we appreciate that the reviewer finds our formulation "rigorous" in conjunction with prior literature, and the overall paper "well written" with detailed supplemental experiments. We now turn to the main points of clarification raised by the reviewer below.

Cost of energy evaluations in comparison to standard diffusion sampler baselines

We acknowledge the reviewer's point that the high cost of energy function evaluations for standard diffusion sampler baselines, such as iDEM and DDS, is not provided in the current manuscript. We first note that the expensive energy evaluation cost of standard diffusion samplers like iDEM has been previously reported in prior works like "Adjoint Sampling" [1], Table 1, and in concurrent work "Adjoint Schrodinger Bridge Sampler" [2], Fig. 4.

Moreover, the scalability of past diffusion samplers to molecular systems, in the Cartesian coordinates setting, remains challenging, if not downright infeasible. Since iDEM is simulation-free and more scalable than DDS, we bolster such folklore knowledge with private correspondence with the iDEM authors, who told us that iDEM failed to scale to ALDP even with substantial effort. We further attempted to reproduce these statements and trained iDEM on ALDP ourselves in Vacuum instead of implicit solvent, and observed that iDEM is unable to successfully sample in this easier setting with most modes missing and a poor energy distribution. We will update the manuscript with a discussion on this point and will provide new Rama plots for iDEM in the appendix.

Evaluation on small peptides

The reviewer makes an astute observation that the molecular systems considered in this paper are small peptides, and real-world application domains may demand conformation sampling at a much larger scale than considered here. We agree with the reviewer on this point. However, we would like to definitively clarify that prior to our work, PITA, no other generative sampler has been demonstrated in Cartesian coordinates for even ALDP, much less AL3. In addition, there is a burgeoning literature of combining generative models with classical MCMC techniques in an effort to "learn to sample" complex Boltzmann distributions. While preliminary, these methods promise the full benefit of deep learning, such as amortization of cost, faster than MD sampling of new points, and transferability to new, unseen molecular systems within a reasonable family. Consequently, in this literature, we position PITA as the current most scalable diffusion sampler, with expectations that future work---given the rapid progress of generative models for sampling---will enable scaling to more meaningful systems of interest. Thus, we believe PITA represents an inflection point wherein future work on generative samplers will be tested beyond even simpler synthetic energy functions such as 2D GMMs and Lennard-Jones potentials.

We thank the reviewer for raising this point. We will include a discussion about the current limitations of diffusion samplers due to current scaling challenges in the conclusion.

Further Questions

Also as in persistent in this field, comparisons are made only on the basis of energy evaluations rather than wall-clock time.

We acknowledge the reviewer's comment regarding computational cost being tied to energy evaluation rather than simply wall clock time. In our Boltzmann sampling, the wall-clock time heavily depends on the quality of the energy approximation. For instance, using an energy function that takes into account quantum effects is orders of magnitude more expensive. Consequently, we report the number of energy evaluations in line with standard practice in the field.

Why is it necessary to distill the score model into an energy-based model? It should be possible to importance sample the target temperature with path RNDs computed only using the SDE drift terms. I would have guessed the weights at intermediate times, enabled by an explicit U_t, allow for frequent resampling using SMC - but the authors don't seem to do this.

We appreciate the reviewer's nuanced comments. Parameterizing and learning the energy-based model is necessary for the importance sampling with the target density $\pi^\beta(x)$ after the inference-time annealing. Indeed, note that the inference-time annealing samples from the density $\propto\exp(-\gamma U_t(x;\eta))$, which does not perfectly match the target density. We discuss this in lines 156-161 and Appendix B.

Furthermore, we perform an additional experiment using this method, which we label as "FKC" in Table R2 of our response to Reviewer Ldmh. We note that this method is only correct if the learned scores perfectly align with $U_t$, which is difficult to achieve in practice.

Authors should at least mention the vast body of prior literature on temperature replica exchange simulations.

Thank you for your comment. We will include further discussion on temperature replica exchange methods and their relation to our work on PITA.

Authors notably omit ESS metrics in final comparisons, despite the key role of importance sampling in the training method.

Thank you for the note. ESS is not applicable here because our method uses SMC and resamples at every step. While ESS could be considered in settings without resampling, the performance differs drastically with and without resampling in our task. Therefore, we believe that reporting ESS for the AIS method would not accurately reflect the performance of our method as a whole, and we instead rely on other metrics. We will add a note to clarify this in the text.

The development of a robust diffusion-based importance sampler of molecular systems in Cartesian coordinates is of independent interest and has not been reported in prior comparable work (iDEM, adjoint sampling). Authors should elaborate on the details and performance of the importance sampling weights in inference time annealing.

We thank the reviewer for this thoughtful suggestion and for recognizing the broader interest in this direction. We agree and will definitely elaborate on these points in the updated draft.

What is the meaning of the number of energy evaluations for the MD-diff baseline?

For the MD-diff baseline, we run a Molecular Dynamics (MD) simulation with the predefined budget of energy evaluations to collect the training data at the target temperature and then train a diffusion model on the collected samples. The number of energy evaluations corresponds to the budget for the MD simulation.

To best contextualize the performance of the method, baselines of normalizing flows trained on MD trajectories (i.e., SBG) should be shown.

We provide results of the normalizing flow trained on the MD trajectory data used for our MD-Diff baseline, which we call "MD-NF" [Table R2 Reviewer Ldmh]. We note that while this model can capture the energy distribution well, it fails to capture all the modes, as is evident from the TICA and Ramachandran (Rama) plots that will be included in the updated manuscript.

Concluding remarks

We thank the reviewer again for valuable comments and detailed feedback. We hope that our rebuttal addresses all of their concerns and encourages the reviewer to potentially consider updating their score if they so deem it. We are also available to answer any further questions that arise.

References

[1] Havens, Aaron, et al. "Adjoint sampling: Highly scalable diffusion samplers via adjoint matching." ICML (2025).

[2] Liu, Guan-Horng, et al. "Adjoint Schrodinger Bridge Sampler." arXiv preprint arXiv:2506.22565 (2025).

I appreciate the authors' response. However, I find aspects of the response not fully satisfactory and would like to continue the discussion.

Q: Why is it necessary to distill the score model into an energy based model? It should be possible to importance sample the target temperature with path RNDs computed only using the SDE drift terms. I would have guessed the weights at intermediate times, enabled by an explicit U_t, allow for frequent resampling using SMC - but the authors don't seem to do this.

A: We appreciate the reviewer's nuanced comments. Parameterizing and learning the energy-based model is necessary for the importance sampling with the target density after the inference-time annealing.

First of all, this does not address my question because importance sampling against a target density is also possible with path RNDs that make use only of the SDE drift (do authors agree?). More curiously however it seems like the authors do in fact resample as stated later in the rebuttal and (upon re-inspection) in the main text, though this is not stated prominently. This raises concerns that the method is not presented as clearly as it could be. To remediate this point I would ask that the authors produce pseudocode for the inference algorithm and agree to include it in the revision.

Q: Authors notably omit ESS metrics in final comparisons, despite the key role of importance sampling in the training method.

A: ESS is not applicable here because our method uses SMC and resamples at every step. While ESS could be considered in settings without resampling, the performance differs drastically with and without resampling in our task.

I interpret this to mean that the ESS is poor without resampling. I would not be surprised, and would not hold it against the authors --- there are no diffusion samplers that obtain good ESS on these systems. However it is important to transparently show these metrics so we can have common knowledge of where to improve.

We thank the reviewer for the detailed and constructive feedback. We appreciate the positive assessment of the proposed method. Namely, the fact that the reviewer finds it to be "intuitive", "well-motivated", and "addressing existing challenges in diffusion-based methods". In what follows, we address the concerns raised and answer their questions grouped by theme.

Empirical comparison with classical baselines

We acknowledge the reviewer concern that our paper could benefit from the addition of classical baselines(SMC and Replica Exchange/Parallel Tempering). We agree with this feedback and have included the suggested baselines in this rebuttal, as shown in Table R1 in our response to Reviewer JqJ5.

Computational overhead and efficiency of training multiple diffusion models

We appreciate the reviewer’s comment that the discussion of the computational overhead and efficiency of training separate models amortized across temperatures may not have been sufficiently clear in the manuscript. We now attempt to add increasing clarity to these points. We begin by first noting that in practice we find that sequential fine-tuning of the model demonstrates the best performance---a point thats outlined also in lines 199-201 of the main paper. More precisely, for the subsequent training of a diffusion model at a lower temperature $1/\beta_{i+1}$ we fine-tune the model trained at the previous temperature $1/\beta_{i}$. Notably, this strategy allows for reducing the number of required parameters (compared to amortization over temperatures) and the number of training iterations (compared to training from scratch for every temperature).

Furthermore, we wish to highlight that by cost we include both the training time as well as the number of target energy function evaluations needed, which is an important cost associated to real world Boltzmann distributions for molecular systems. We will clarify these computational efficiency gains and their overhead in the updated draft of the main paper.

Questions

I wonder why the authors chose to use two separate models for energy function and score (or equivalently, denoiser)? In principle the score function can be obtained by taking the gradient of the energy function. Is this separation primarily to avoid the computational and memory burden of backpropagation?

The reviewer makes an interesting observation. Indeed, the score function can be obtained as a gradient of the energy function. There are three reasons why we use the a separate score model.

1. As noted by the reviewer the computational burden is reduced, and in particular this scheme allows us to avoid the computation of the laplacian of the energy-based model for computing the weights.
2. Training separate score and energy models is the current SOTA in energy model training following [1]. We get the score function for free given our training scheme.
3. The score function in general produces better samples than the gradient of the energy function, therefore is a better proposal distribution in practice.

We think this is a cool and important component of PITA and will add additional discussion in the text.

Have the authors conducted an ablation study on the four types of losses used in Algorithm 1? Are all of them necessary? Do they receive equal weight in the final objective? Are they trained under the same optimizer settings?

The four losses in Algorithm 1 consist of two groups: 2 losses for training the score model and 2 losses for training the EBM. For training the score model when both training samples and the target unnormalized density are available, we follow the best practice in the field so far [2]. For training the energy-based model, we follow [1] and distill it from the score model. The novel loss introduced here is the "energy pinning" loss, which we ablate in Appendix E (see Table 6 for comparison). We hope this now sufficiently clarifies and the reviewers question.

Have the authors considered a baseline where a diffusion model is trained at a high temperature, and then sampling is performed using the Feynman-Kac PDE framework at the target (lower) temperature?

We have updated our results below to include

1. Training a diffusion model and only rescaling the score to anneal from 1200K to 300K (ScoreScaling),
2. Using the Feynman-Kac Correctors [5] framework to anneal from 1200K to 300K (FKC).

As can be seen in the table below, neither of these frameworks can successfully capture the lower temperature distribution on this task.

Table R2: Performance of methods for the ALDP sampling task. The starting temperature is $T_L$ = 1200 K, annealed to target $T_S$ = 300 K. Metrics are calculated over 10k samples and standard deviations over 3 seeds.

How do you tune the hyper-parameters for all methods?

For all methods we start with the parameters from SOTA methods in the literature then perform coordinate-wise grid searches over the most relevant parameters to our task using the best performing setting on a validation set at temperature 1200k. For loss weights we simply approximately balanced the magnitude of the losses on the training set, e.g. we used a weight of 0.01 for the target score matching loss to ensure it is on the same scale as the other losses.

Do you have any guidance in setting the annealing schedules for new tasks?

For simplicity, we follow the standard practice in the literature and consider the geometric annealing schedule [3]. This, however, can be potentially improved by incorporating modern advanced schedule optimization techniques [4]. We leave exploration of annealing schedules, including dynamic or learned ones, as interesting directions for future work. In theory, the standard practice of annealing to match a target ESS should also work in our setting.

Do diffusion models trained at different temperatures require different optimization settings or hyperparameter tuning? If yes, do you have practical guidance on that?

We find that our approach performs consistently well across different temperatures without additional hyperparameter tuning. Specifically, we keep the training settings (e.g., learning rate, number of iterations) constant across all temperature levels.

This paper also explores diffusion-based sampling with access to the target-mode information and may serve as a useful point of comparison: LEARNED REFERENCE-BASED DIFFUSION SAMPLING FOR MULTI-MODAL DISTRIBUTIONS Maxence Noble∗, Louis Grenioux∗, Marylou Gabrie & Alain Oliviero Durmus

We thank the reviewer for pointing out the relevant piece of literature. We have added a discussion of this work into the related work section.

Minor and writing-related comments

We thank the reviewer for the detailed minor writing comments.

In this expression "U_t = -log p_t(x) + const", is the constant independent of t?

Yes, that is the log normalization constant.

Eq 10: Can you explain the right equation? What does the "\approx" mean and why it holds?

It is approximate as $U_t$ is a learned model and $\pi$ is the true distribution.

Table 2 and 3: does the reported #energy evals for the proposed method include the evaluations used in the initial MCMC sampling phase? Could the authors provide the MCMC configuration details (e.g., number of steps, acceptance rates), and are the final results sensitive to these settings?

Yes, the reported number of energy evaluations includes those in the initial sampling phase. Note that we collected samples using MCMC for LJ-13 and Molecular Dynamics for the peptide experiments. These configurations are provided in Appendix G.3. Since we use a fairly converged chain, the final results are not very sensitive to these configuration details.

Concluding remarks

We thank the reviewer again for their detailed feedback that helped us improve our work. We very politely invite the reviewer to increase their score if they feel our rebuttal was successful in answering all their raised concerns. We also remain available for further discussion and for addressing any lingering doubts. Please do let us know.

References

[1] Thornton, James, Louis Béthune, Ruixiang Zhang, Arwen Bradley, Preetum Nakkiran, and Shuangfei Zhai. "Composition and control with distilled energy diffusion models and sequential monte carlo." arXiv preprint arXiv:2502.12786 (2025).

[2] De Bortoli, Valentin, Michael Hutchinson, Peter Wirnsberger, and Arnaud Doucet. "Target score matching." arXiv preprint arXiv:2402.08667 (2024).

[3] Del Moral, Pierre, Arnaud Doucet, and Ajay Jasra. "Sequential monte carlo samplers." Journal of the Royal Statistical Society Series B: Statistical Methodology 68, no. 3 (2006): 411-436.

[4] Syed, Saifuddin, Alexandre Bouchard-Côté, Kevin Chern, and Arnaud Doucet. "Optimised annealed sequential Monte Carlo samplers." arXiv preprint arXiv:2408.12057 (2024).

[5] Skreta, Marta, et al. "Feynman-kac correctors in diffusion: Annealing, guidance, and product of experts." ICML (2025).

[7] Karras, T., et al. "Elucidating the design space of diffusion-based generative models." NeurIPS (2022).

We thank the reviewer for their valuable feedback and the time spent. We appreciate that the reviewer finds the paper to be "well-written" and "the integration of Feynman–Kac theory into diffusion inference to be interesting and novel". In what follows, we address the potential misunderstandings and concerns raised.

Distinguishing the novelty of the two interpolants

We acknowledge the reviewer's comment that the core novelty of our two introduced interpolants may not have been sufficiently distinguished from existing literature. We now clarify this important point further.

The core novelty of the two interpolants in this work [...] is not sufficiently distinguished from existing literature.

From my understanding, the method reads like a combination of known components, and the authors may wish to clarify why this framework yields fundamentally new insights or performance advantages beyond prior annealed sampling or diffusion-EBM hybrids.

This framework yields performance advantages beyond prior annealed sampling or diffusion-EBM hybrids by being the first model to perform consistent simultaneous annealing in both the diffusion and temperature axes.

One of the core novelties of this work is the Sequential Monte Carlo (SMC) scheme for the inference-time annealing of diffusion models, which enjoys both the efficiency of expressive diffusion models and the consistency of SMC estimators. This combination of diffusion and temperature annealing in a consistent manner is what distinguishes PITA from the existing literature. In particular, compared to the previous literature, this procedure has the following merits:

1. Unlike Langevin Monte Carlo [1], which converges only in the limit of an infinite number of steps, it provides a consistent estimator w.r.t. the number of simulated particles. This fact allows us to curate new lower temperature (consistent) datasets to train new diffusion models---a fact that is not possible with simple Langevin Monte Carlo without incurring bias.
2. The critical advantage of PITA is that by combining both interpolation strategies, we achieve more stable and performative training of the final diffusion sampler. This combination, to the best of our knowledge, is novel and in particular benefits from our inference time SMC scheme that is tailor-made for diffusion models. As we stress, in the main paper, no existing diffusion sampler is able to tackle the molecular energy systems we consider in Cartesian coordinates. This is because training such samplers without data is extremely challenging and numerically unstable due to the poor estimation of the Stein score, as well as finding the modes in a complex energy landscape to estimate the score in the first place.

Annealing has long been used in MCMC [1] and has been applied to diffusion models [2–4].

The reviewer makes an astute observation that Annealing has been applied to diffusion models. In new experiments, we demonstrate that PITA performs better than Annealing in MCMC, and existing annealed diffusion models.

1. We first add two baselines of annealed MCMC with sequential Monte Carlo and parallel tempering (PT) strategies. As it can be seen in Table R1 below, PITA performs better overall in capturing the distribution given the same budget of energy evaluations

Table R1: Performance of the SMC and PT baselines on ALDP. For SMC, we take 10 annealing temperatures between 1200K and 300K using a geometric schedule, taking 166 steps per temperature with 30k particles. For PT, we use the same 4 annealing temperatures as PITA, running 25 MD steps between particle exchanges for a total of 50k iterations.

2. Second, we add a baseline of "score scaling", which is the typical strategy in classifier-free guidance approaches to approximately anneal diffusion models [Table R2 Reviewer Ldmh]. PITA substantially outperforms this method. We note that PITA is consistent where score scaling is not, which explains PITA's superior performance. Score scaling does not produce correct marginals both theoretically and empirically, as also demonstrated in the recent Feynmann-Kac Correctors paper [11], Section 5.2.

The idea of diffusion smoothing looks similar to annealed importance sampling [5] and was later extended to deep energy-based models [6–8].

Compared to Annealed Importance Sampling (AIS) [5], it follows the diffusion path raised to the power $\beta$ (i.e. $p_t(x)^\beta$) instead of the geometric averaging (i.e. $\pi(x)^{t\beta}p_0(x)^{(1-t)}$). The use of diffusion paths is crucial for our setting as it reduces the variance of importance weights.

In addition, we note that previous works on deep energy-based models [6-8] combine the standard inference with Langevin dynamics. This means that the inference procedure converges only in the limit of an infinite number of iterations. This is in stark contrast with the proposed consistent estimator used in PITA, which, as described above, allows for the consistent curation of lower temperature datasets.

Scope of Evaluation

We value the reviewer's concern that the scope of evaluation may initially appear limited. However, we would like to politely push back against this characterization. We begin by recalling that our paper is concerned with sampling from an unnormalized density---a problem domain that is well known to be significantly harder than conventional generative modelling. This is because, unlike generative modelling, we are afforded no samples from the target distribution, and yet we must draw i.i.d. samples from it. Due to this challenge, a majority of work has focused on smaller molecular systems in bespoke internal coordinates, which are of smaller dimensionality and are provided through exact domain knowledge for a given system. In contrast, energy functions in Cartesian coordinates are more challenging as they represent the full dimensionality of the molecular system and, as a result, are much larger. As a result, we wish to again reiterate this subtle but important point, existing diffusion samplers have failed to even scale to Alanine Tripeptide (AL3) with 99 dimensions, which our method, PITA, effectively tackles. Consequently, we argue that such an empirical demonstration goes beyond the current benchmarks in the literature.

We thank the reviewer again for their comment, and we will further update the manuscript to highlight the challenge of sampling in Cartesian coordinates.

Questions

Since the method relies on both a score model and an energy model at each temperature level, is the computational and training cost substantially higher than standard diffusion models or Boltzmann generators?

Training a score model and an energy model incurs higher cost than a standard diffusion model, which only has to train a score model. Training the energy model is around 2x the number of flops as the score model for 3x the cost. However, since training the energy model has substantial benefits performance-wise over only training a score model, this additional cost is a necessary price to pay.

The proposed method appears to be reminiscent of annealed Langevin dynamics and predictor-corrector samplers discussed in [9][...]. Could the authors discuss the relationship and trade-offs? What are the advantages of your approach?

First, we would like to address the potential confusion about the usage of the term "annealing". Although it is sometimes used for different noise scales of the target distribution (as in [9]), in this paper, we use annealing exclusively for temperature annealing (i.e., changing the power of the marginal density $p_t(x)^\beta$). Note that changing the temperature $p_t(x)^\beta$ and adding the noise $p_t(x) = \int dy\; \mathcal{N}(x|y,\sigma_t^2) * p_0(y)$ are completely different operations resulting in different marginals and requiring different models for their simulation. We will add the corresponding discussion to the text to clarify this terminology.

Second, the Predictor-Corrector scheme proposed in [9] samples from the correct marginals only in the limit of an infinite number of corrector steps (Langevin dynamics), while our proposed inference-time annealing uses a consistent SMC estimator in finite steps by incorporating the importance weights.

In [10], the authors proposed an energy model with adaptive temperature control by reparameterizing the energy function [...]. Could the authors comment on this approach and discuss the advantages of PITA compared to such adaptive-temperature energy models?

We thank the reviewer for their question. The work in question [10] does not consider sampling from the temperature annealed marginals. Instead, at inference, it runs samples from the marginals with different noise levels, which is similar to the Predictor-Corrector scheme, and cannot be used directly for consistent sampling.

Concluding remarks

We thank the reviewer again for their feedback. We hope that our response and additional experiments clarify all the great points raised in the review and help the reviewer understand the core novelty of our contributions and caliber of our empirical evaluation better. If the reviewer is satisfied with all of our responses, we very politely encourage the reviewer to potentially upgrade their score if they deem it appropriate. We are also more than happy to answer any follow-up questions the reviewer may have.

Additional References

[11] Skreta, Marta, et al. "Feynman-kac correctors in diffusion: Annealing, guidance, and product of experts." ICML (2025).