Title: Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities
Venue: NeurIPS 2025
Accepted: yes
Keywords: Boltzmann Sampling

Diffusion Models

Molecular Conformations

Sequential Monte Carlo
Abstract: Sampling efficiently from a target unnormalized probability density remains a core challenge, with relevance across countless high-impact scientific applications. A promising approach towards this challenge is the design of amortized samplers that borrow key ideas, such as probability path design, from state-of-the-art generative diffusion models. However, all existing diffusion-based samplers remain unable to draw samples from distributions at the scale of even simple molecular systems. In this paper, we propose Progressive Inference-Time Annealing (PITA) a novel framework to learn diffusion-based samplers that combines two complementary interpolation techniques: I.) Annealing of the Boltzmann distribution and II.) Diffusion smoothing. PITA trains a sequence of diffusion models from high to low temperatures by sequentially training each model at progressively higher temperatures, leveraging engineered easy access to samples of the temperature-annealed target density. In the subsequent step, PITA enables simulating the trained diffusion model to *procure training samples at a lower temperature* for the next diffusion model through inference-time annealing using a novel Feynman-Kac PDE combined with Sequential Monte Carlo. Empirically, PITA enables, for the first time, equilibrium sampling of $N$-body particle systems, Alanine Dipeptide, and tripeptides in Cartesian coordinates with dramatically lower energy function evaluations.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe goal is to draw samples from a distribution knowing its unnormalized density $\\pi(x)$ — the authors design a way to additionally get approximate samples to make the sampling easier. The idea is to sequentially draw samples from tempered distributions $\\pi(x)^{\\beta}$, increasing the inverse temperature $\\beta$ from $0$ (excluded) to $1$, as do classical Sequential Monte Carlo (SMC) algorithms. \n\n**Repeat** until $\\beta = 1$:\n\n1. Increase $\\beta$.\n\n2. Use the unnormalized density $\\pi(x)^\\beta$ to get approximate samples. \n\n3. Use the unnormalized density $\\pi(x)^\\beta$ and approximate samples to get better samples. To do so, train a diffusion model to learn and simulate energies of $\\pi_t(x)^{\\beta}$ where $t$ indexes the diffusion. \n\nSMC algorithms typically do Steps 1 and 2. The authors' algorithm distinguishes itself by adding Step 3. \n\nMoreover, the authors' implementation of Step 2 is also more sophisticated than SMC. The first time, step 2 is achieved by running an MCMC: it should converge relatively fast given that the multi-modality is flattened away. The rest of the time, step 2 is achieved by running an SMC (or FK-PDE) on the diffusion path at the previous temperature, raised to a power.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\nThe authors' method is creative and interesting and seems to be improve over previous models on biological problems. It addresses many current questions:\n\n- how to combine probability paths\n\nThe diffusion path avoids mass teleportation: it is usually simulated from samples. \nThe temperature path preserves the mode locations: it is usually simulated from an unnormalized density. \n\n- how to obtain \"good enough\" samples from an unnormalized density in order to train a diffusion model to get \"better\" samples\n\nThis is an important challenge in the literature of diffusion samplers, as explained for example in *He et al., No Trick, No Treat: Pursuits and Challenges Towards Simulation-free Training of Neural Samplers*.\n\n- how to use both combine samples *and* an unnormalized density to draw samples\n\n## Weaknesses\n\n**Dataset benchmarks**. It is encouraging to see sampling methods being used to make progress on real-world datasets. However, given the relevance and centrality of the authors' method to the sampling literature, it would be useful to include at least one experiment on a synthetic dataset such as a mixture of 50 Gaussians in 2 dimensions. This is because the sampling community as a whole does not have technical expertise in biological problems to judge the experiments section, because other competing methods are often benchmarked on a toy dataset, and also because it provides intuition on the performance of the method. \n\n**Comparison with SMC in the experiments**. The most natural competitor to the author's PITA algorithm is a basic SMC that follows a temperature-annealing path. I appreciate the fact that some of the baselines (e.g. Temperature Annealed Boltzmann Generators) are very similar to SMC, but a vanilla SMC with a Langevin Kernel and a classical schedule for choosing the intermediate temperatures would be a convincing baseline in the experiments. \n\n**Ablation study.** \n\nQ1. 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^\\beta$ to $\\pi^{\\beta + \\Delta \\beta}$ and then learn a diffusion model?\n\nQ2. What is the difference in performance if PITA was run only with inference-time annealing and no model training?\n\n**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. \n\n**Lack of clarity**. I am quite familair with the topic and it took me many reads to understand what the algorithm what doing. Some things that could be more clearly explained:\n- 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\"). \n- why model training is useful in addition to inference-time annealing? \n- 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\"). \n\n**Unclear statements**. \n\nLines 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 $x_t$ that are almost from the target distribution. How is this compatible with the statement: \"it has no information about the target distribution\"?\n\nLines 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.\n\nQUESTIONS:\nCould the authors address the different points in the weaknesses section above?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a scheme for training a diffusion sampler of a molecular Boltzmann distribution given access to the energy function and a set of high-temperature samples. In the scheme, diffusion models trained at the kth temperature are used to importance sample the target distribution at the (k+1)th temperature. These samples are used to train the next diffusion model. The process repeats until the desired temperature is reached.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe core idea proposed in the paper - that of bootstrapping a sampler by learning a sequence of distributions (using diffusion models) that anneal towards the target - is novel and conceptually solid. The formulation is rigorous and the connections with existing and prior literature are laid out well. The paper is overall well written and the supplementary experiments are detailed.\n\n**Concerns or questions**\n\n* The paper makes repeated claims about using less energy evaluations than previous methods, but these comparisons are not shown (on alanine peptides). Indeed, rather few prior methods are mentioned - standard diffusion sample baselines (DDS, iDEM, etc), while conceivably worse, are not shown in the results.\n\n* As in persistent in this field, evaluations are made in settings of questionable practical utility - the small system size conceals problems that may make applications to larger systems impossible. In particular (1) poor coverage of dominant modes at high temperatures, e.g. phase transitions; (2) increasing number of intermediate temperatures required for effective importance sampling; (3) divergence computations to calculate the particle weights (4) distillation of the energy based models.\n\n* Also as in persistent in this field, comparisons are made only on the basis of energy evaluations rather than wall-clock time.\n\n* 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.\n\n* Authors should at least mention the vast body of prior literature on temperature replica exchange simulations.\n\n* Authors notably omit ESS metrics in final comparisons, despite the key role of importance sampling in the training method.\n\n* 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 \n\n* What is the meaning of the number of energy evaluations for the MD-diff baseline?\n\n* To best contextualize the performance of method, baselines of normalizing flows trained on MD trajectories (i.e., SBG) should be shown.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel method that combines annealing and diffusion-based samplers to sample from unnormalized target distributions. The key idea is to begin at a high-temperature distribution, where mixing is easier for standard MCMC methods, and progressively anneal toward the target distribution. At each annealing step, the approach alternates between two components: (i) learning a diffusion-based sampler via time-dependent score and energy models at that temperature, and (ii) generating samples using a Feynman-Kac PDE framework for the next, lower temperature. \n\nIn empirical study, the method is evaluated on a series of synthetic particle systems. The proposed method generally produces qualitatively and quantitatively better results than existing baselines. On the most challenging task, Alanine Tripeptide, the method achieves good mode coverage but underperforms on certain evaluation metrics.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- The proposed method is intuitive and well-motivated, leveraging a progressive annealing path and coarse simulation data (by MCMC) to address challenges in existing diffusion-based methods.\n\n- The proposed method empirically scales to systems where prior diffusion-based approaches struggle.\n\nWeakness:\n\n- Some parts of the writing could be improved—there are inconsistencies and instances of imprecise or vague explanations (see specific comments). \n\n- The paper lacks empirical comparisons to classical baselines such as replica exchange or sequential Monte Carlo (with temperature annealing), making it difficult to assess the overall significance of the proposed approach. \n\n- The proposed approach can be less flexible as it requires training diffusion models at different temperature, and relies on specifying the initial temperature and annealing schedule.\n\nQUESTIONS:\n- If I understand correctly, the method involves training separate models at multiple temperature levels. Are these models trained independently, or is it amortized across temperatures? It would also be helpful to clarify how this overhead compares to other methods in terms of computational efficiency and scalability.\n\n- 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? \n\n- 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?\n\n- 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?\n\n- How do you tune the hyper-parameters for all methods? \n\n- Do you have any guidance in setting the annealing schedules for new tasks?\n\n- Do diffusion models trained at different temperatures require different optimization settings or hyperparameter tuning? If yes, do you have practical guidance on that?\n\n-This paper [1] also explores diffusion-based sampling with access to the target-mode information and may serve as a useful point of comparison:\n\n[1] LEARNED REFERENCE-BASED DIFFUSION SAMPLING\nFOR MULTI-MODAL DISTRIBUTIONS\nMaxence Noble∗, Louis Grenioux∗, Marylou Gabrie & Alain Oliviero Durmus\n\nMinor and writing-related: \n\n-  In this expression \"U_t = -log p_t(x) + const\", is the constant independent of t?\n\n- Eq 10: Can you explain the right equation? What does the \"\\approx\" mean and why it holds?\n\n- Algorithm 1: the loss function names in the optimizer lines are inconsistent with those in previous lines\n\n- Algorithm 1: It would be helpful to make the algorithm self-contained by specifying inputs such D_t and U_t. \n\n- 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? \n\n- Table 3: what is \"PITA\" vs \"PITA (w/o relaxation)\"?\n\n- On page 5 line 164: \"one subsection\" instead of \"two\"?\n\nLIMITATIONS:\nOne limitation acknowledged by the authors is the need to train both score and energy networks, which increases the training complexity \n\nAdditional imitations include\n\n- The need to train these models at several temperature levels, further raising the training cost. \n\n- Lack of comparison to some of the classical sampling approaches.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Progressive Inference-Time Annealing (PITA), a generative framework that progressively trains a cascade of diffusion models to sample from low-temperature Boltzmann distributions. The core idea is to combine two interpolation strategies: (1) temperature annealing, which flattens energy barriers and enables better mode mixing; and (2) diffusion smoothing, which enables efficient resampling at inference time via Feynman–Kac-weighted reverse SDEs. By alternating between training a diffusion model at a given temperature and then resampling to a colder temperature using learned score and energy models, PITA can produce samples that align closely with the Boltzmann target at low temperatures.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe paper is well-written and presents a clean formulation of its algorithmic ideas. The integration of Feynman–Kac theory into diffusion inference is interesting and, to my knowledge, novel in the generative modeling literature. Moreover, the approach achieves significant gains in sample quality and coverage compared to prior Boltzmann generator and diffusion methods, especially in physically meaningful molecular domains.\n\n**Weaknesses**\n\n1. The core novelty of the two interpolants in this work (temperature annealing and diffusion smoothing) is not sufficiently distinguished from existing literature. It seems that both techniques have been widely studied in the MCMC, diffusion-based modeling or energy-based modeling literature.\n2. Annealing has long been used in MCMC [1] and has been applied to diffusion models [2–4].\n3. The idea of diffusion smoothing looks similar to annealed importance sampling [5] and was later extended to deep energy-based models [6–8].\n4. 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.\n5. The scope of the evaluation is limited. While the authors demonstrate promising results on small physical systems like LJ-13 and alanine peptides, it is unclear whether this method generalizes to broader settings (e.g., image or text generation, or large protein folding datasets). The method is currently tailored for low-dimensional molecular systems where accurate energy evaluations are feasible.\n\n\n**References:**\n\n[1] Bayesian learning via stochastic gradient Langevin dynamics, ICML 2011\n\n[2] Generative modeling by estimating gradients of the data distribution, NeurIPS 2019\n\n[3] Denoising diffusion probabilistic models, NeurIPS 2020\n\n[4] Improved techniques for training score-based generative models, 2020\n\n[5] Annealed importance sampling, Statistics and computing, 2001\n\n[6] Learning energy-based models by diffusion recovery likelihood, ICLR 2021\n\n[7] A Langevin-like sampler for discrete distributions, ICML 2022\n\n[8] Score-based diffusion meets annealed importance sampling, NeurIPS 2022\n\nQUESTIONS:\n1. 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?\n2. The proposed method appears to be reminiscent of annealed Langevin dynamics and predictor-corrector samplers discussed in [9], which also gradually anneal the target distribution by adjusting the noise scale. Could the authors discuss the relationship and trade-offs? What are the advantages of your approach?\n3. In [10], the authors proposed an energy model with adaptive temperature control by reparameterizing the energy function, enabling built-in temperature adjustment in the generative model. Could the authors comment on this approach and discuss the advantages of PITA compared to such adaptive-temperature energy models?\n\n**References:**\n\n[9] Score-based generative modeling through stochastic differential equations, ICLR 2021\n\n[10] Learning energy-based models by cooperative diffusion recovery likelihood, ICLR 2024\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-06"}