Research map / Image generation / Fast diffusion sampling
Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PITA, which progressively trains diffusion samplers from high to low temperature, using inference-time annealing via a Feynman-Kac PDE with SMC to generate lower-temperature training samples for Boltzmann distributions.
Problem
Efficient equilibrium sampling from unnormalized Boltzmann densities of molecular systems in Cartesian coordinates, where existing diffusion-based samplers need too many energy evaluations and fail at peptide scale.
Approach
Combine temperature annealing with diffusion smoothing: train diffusion models sequentially from high to low temperature, and at each step use a learned energy model with a Feynman-Kac PDE plus SMC resampling to anneal the trained diffusion model's marginals and produce training samples for the next, colder model.
Key results
- On LJ-13, PITA reaches Distance-W2 of 0.04 versus 1.21 for TA-BG and 0.127 for iDEM.
- On ALDP at 300K, PITA achieves Energy-W2 of 1.615 versus 3.704 for MD-Diff and 86.176 for TA-BG, with 5e7 energy evaluations.
- On AL3, PITA gets Rama-KL 1.209 and Energy-W2 2.592, outperforming baselines on energy metrics but worse on TICA-W1/W2.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use progressive inference-time annealing instead of TA-BG | LJ-13 sampling · LJ-13 | Distance-W2: 1.21 ± 0.02 -> 0.04 ± 0.00 | Helped |
| use progressive inference-time annealing instead of TA-BG | LJ-13 sampling · LJ-13 | Energy-W2: 61.47 ± 0.12 -> 2.26 ± 0.21 | Helped |
| use progressive inference-time annealing instead of iDEM | LJ-13 sampling · LJ-13 | Geometric-W2: 1.61 ± 0.01 -> 1.65 ± 0.00 | Hurt |
| use progressive inference-time annealing instead of MD-Diff | ALDP sampling · ALDP | Tica-W2: 0.579 $\pm$ 0.004 -> 0.379 $\pm$ 0.028 | Helped |
| use progressive inference-time annealing instead of MD-Diff | ALDP sampling · ALDP | Rama-KL: 1.308 $\pm$ 0.072 -> 4.773 $\pm$ 0.460 | Hurt |
| use progressive inference-time annealing instead of MD-NF | ALDP sampling · ALDP | Energy-W2: 1.198 $\pm$ 0.069 -> 1.615 $\pm$ 0.053 | Hurt |
| use progressive inference-time annealing instead of TA-BG | ALDP sampling · ALDP | Energy-W2: 86.176 $\pm$ 0.104 -> 1.615 $\pm$ 0.053 | Helped |
| use progressive inference-time annealing instead of MD-Diff | AL3 sampling · AL3 | Tica-W2: 0.426 $\pm$ 0.010 -> 0.952 $\pm$ 0.055 | Hurt |
| use progressive inference-time annealing instead of MD-Diff | AL3 sampling · AL3 | Rama-KL: 9.662 $\pm$ 0.085 -> 1.209 $\pm$ 0.144 | Helped |
| add relaxation (ablation) | PITA · AL3 | Rama-KL: 8.535 $\pm$ 0.254 -> 1.209 $\pm$ 0.144 | Helped |
| add relaxation (ablation) | PITA · AL3 | Energy-W2: 87.695 $\pm$ 0.294 -> 2.592 $\pm$ 0.107 | Helped |
| add relaxation (ablation) | PITA · AL3 | Tica-W2: 0.999 $\pm$ 0.043 -> 0.952 $\pm$ 0.055 | Helped |
Surprising findings
- PITA performs worse than baselines on TICA-W1 and TICA-W2 for AL3, suggesting it does not recover the correct relative mode weights despite better mode coverage.
- MD-NF achieves better Energy-W1 and Energy-W2 than PITA on ALDP (0.551 and 1.198 versus 1.530 and 1.615).
Limitations
- Requires training an additional energy-based model, which is notoriously difficult.
- Simultaneous training of score and energy models adds computational and memory cost.
- Importance weights have high variance without resampling; log ESS is low with pure importance sampling.
Experimental setup
Datasets: LJ-13 MCMC samples Alanine Dipeptide MD trajectories Alanine Tripeptide MD trajectories
Benchmarks: LJ-13 sampling Alanine Dipeptide (ALDP) at 300K Alanine Tripeptide (AL3) at 300K
Metrics: Distance-W2 Energy-W2 Geometric-W2 Rama-KL TICA-W1 TICA-W2 Energy-W1 T-W2 number of energy evaluations
Baselines: iDEM Adjoint Sampling TA-BG MD-Diff MD-NF FKC Score Scaling SMC PT
Builds on: Boltzmann Generators Temperature Annealed Boltzmann Generators Feynman-Kac correctors Target score matching EDM
Closest papers by meaning
- On scalable and efficient training of diffusion samplers (NeurIPS 2025 · similarity 0.89)
- Composing Diffusion Priors with Explicit Physical Context via Generative Gibbs Sampling (NeurIPS 2026 · similarity 0.89)
- Transport meets Variational Inference: Controlled Monte Carlo Diffusions (ICLR 2024 · similarity 0.88)
- A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models (NeurIPS 2026 · similarity 0.88)
- Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts (ICML 2025 · similarity 0.88)
- Non-equilibrium Annealed Adjoint Sampler (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.