Research map / Image generation / Fast diffusion sampling
Non-equilibrium Annealed Adjoint Sampler
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 NAAS, an SOC-based unbiased diffusion sampler whose reference dynamics is an annealed SDE between a Gaussian prior and the target, trained with lean adjoint matching and reciprocal adjoint matching with replay buffers.
Problem
Sampling from an unnormalized density with learned diffusion samplers; existing SOC-based samplers start from uninformative (static or OU) reference dynamics, which limits guidance toward the target and efficiency of learning the control.
Approach
Uses an annealed reference SDE driven by interpolated energy U_t between a Gaussian prior and the target as the base process in an SOC problem, and learns a prior over an extra backward interval [-1,0] with a second control, alternating adjoint matching and reciprocal adjoint matching so the terminal law is exactly the target without importance weighting.
Key results
- On MW54 (5d), NAAS reaches Sinkhorn 0.10 vs 0.32 for the next-best AS baseline.
- On MoS (50d), NAAS reaches Sinkhorn 394.55 and MMD 0.113, beating SCLD's 656.10 Sinkhorn.
- On GMM40 (50d), NAAS reaches Sinkhorn 496.48 vs 3787.73 for SCLD; on Funnel it is not best on Sinkhorn (132.30 vs 117.48 SMC-ESS).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use annealed reference dynamics with learned controls instead of AS | NAAS · MW54 (d=5) | Sinkhorn: 0.32 $\pm$ 0.06 -> 0.10 | Helped |
| use annealed reference dynamics with learned controls instead of CMCD-KL | NAAS · Funnel (d=10) | MMD: 0.095 $\pm$ 0.003 -> 0.076 $\pm$ 0.004 | Helped |
| use annealed reference dynamics with learned controls instead of SMC-ESS | NAAS · Funnel (d=10) | Sinkhorn: 117.48 -> 132.30 | Hurt |
| use annealed reference dynamics with learned controls instead of SCLD | NAAS · GMM40 (d=50) | Sinkhorn: 3787.73 $\pm$ 249.75 -> $\mathbf{496.48}$ $\pm$ 27.08 | Helped |
| use annealed reference dynamics with learned controls instead of DDS | NAAS · MoS (d=50) | MMD: 0.131 $\pm$ 0.001 -> $\mathbf{0.113}$ $\pm$ 0.070 | Helped |
| use annealed reference dynamics with learned controls instead of SCLD | NAAS · MoS (d=50) | Sinkhorn: 656.10 $\pm$ 88.97 -> $\mathbf{394.55}$ $\pm$ 29.35 | Helped |
| replace alternating optimization with joint optimization (ablation) | NAAS · MW54 | Sinkhorn: 0.10 -> 0.20 | Hurt |
| replace alternating optimization with joint optimization (ablation) | NAAS · MoS | Sinkhorn: 394.55 -> 597.99 | Hurt |
Surprising findings
- On Funnel Sinkhorn, NAAS (132.30) is slightly worse than SMC-ESS (117.48), so it is not best across all tasks despite the claim of state-of-the-art on all tasks.
Limitations
- Requires computing the Hessian of the target energy during sampling, adding overhead.
- Requires two networks trained in an alternating scheme.
- Not tested on larger-scale or real-world problems with more expressive annealing schedules.
Experimental setup
Datasets: MW54 Funnel GMM40 MoS alanine dipeptide
Benchmarks: MW54 Funnel GMM40 MoS alanine dipeptide
Metrics: Sinkhorn distance MMD KL divergence Wasserstein distance of energies
Baselines: SMC SMC-ESS CRAFT DDS PIS AS CMCD SCLD
Builds on: Adjoint Matching Adjoint Sampling Denoising Diffusion Sampler Path Integral Sampler SCLD
Closest papers by meaning
- Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching (ICML 2025 · similarity 0.93)
- Adjoint Schrödinger Bridge Sampler (NeurIPS 2025 · similarity 0.89)
- AdjointDEIS: Efficient Gradients for Diffusion Models (NeurIPS 2024 · similarity 0.88)
- On scalable and efficient training of diffusion samplers (NeurIPS 2025 · similarity 0.88)
- Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities (NeurIPS 2025 · similarity 0.88)
- NETS: A Non-equilibrium Transport Sampler (ICML 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.