Research map / Probabilistic and scientific modeling / Variational and generative sampling
Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions
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 Annealing Flow, a continuous normalizing flow trained block-wise over annealed intermediate densities with a dynamic optimal-transport (Wasserstein-2 regularized) objective, for sampling high-dimensional multi-modal distributions with few steps and little MCMC.
Problem
Sampling from high-dimensional, multi-modal distributions where MCMC mixes slowly and existing normalizing-flow samplers suffer mode collapse, need many time steps, or rely on score estimation and MC corrections.
Approach
Train a CNF block-by-block along an annealing path of intermediate densities, using a dynamic OT objective (KL to next annealed density plus discretized W2 path-length regularization) instead of score matching; the optimal velocity equals the score difference between consecutive densities.
Key results
- On ExpGauss with 1024 modes in 50D, AF explores all 1024 modes (True: 1024) while CRAFT explores 473.2, LFIS 298.2, PGPS 813.6, PT under 10 and AIS 125.6 (20 time steps).
- AF uses 12 time steps for GMMs and 6 for Bayesian logistic regression, versus 128-256 steps for CRAFT, LFIS and PGPS in comparisons.
- AF generates 10,000 samples in 2.1 seconds on average in 50D experiments; storage of 2.1MB for 20 networks in 50D ExpGauss.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AF instead of SVGD | ExpGauss-1024, d=50 · ExpGauss-1024, d=50 | number of modes explored: 916.4 -> 1024 | Helped |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Cancer | accuracy (%): 98.9 ± 1.2 -> 97.9 ± 1.1 | Hurt |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Diabetes | accuracy (%): 76.3 ± 2.2 -> 76.3 ± 2.1 | No clear effect |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Heart | accuracy (%): 88.4 ± 3.0 -> 88.5 ± 2.7 | No clear effect |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Australian | accuracy (%): 85.0 ± 3.0 -> 86.6 ± 1.2 | Helped |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Ijcnn1 | accuracy (%): 91.8 ± 0.2 -> 92.0 ± 0.1 | No clear effect |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Svmguide3 | accuracy (%): 80.6 ± 1.0 -> 80.0 ± 1.0 | Hurt |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · German | accuracy (%): 77.5 ± 1.7 -> 78.0 ± 1.7 | No clear effect |
| Use AF instead of the strongest reported baseline | Bayesian logistic regression · Splice | accuracy (%): 83.1 ± 1.5 -> 86.9 ± 1.7 | Helped |
| Replace the endpoint negative-log-density term with its first-order Taylor approximation (ablation) | Annealing Flow · Gaussian Mixture Model, Funnel, and ExpGauss experiments | sampling performance | Helped |
| Remove annealing steps or use very few annealing steps (ablation) | Annealing Flow · High-dimensional distributions with widely separated modes | sampling performance | Hurt |
| Remove dynamic OT objective with W₂ regularization (ablation) | Annealing Flow · High-dimensional and multi-modal sampling | training stability | Hurt |
Surprising findings
- Standard NF/annealing baselines such as CRAFT, LFIS and PGPS fail on the truncated normal distribution, producing no meaningful samples.
- PGPS performs poorly without Langevin adjustments after each step.
Limitations
- Relies on annealing schedule choice and number of intermediate densities; ablations show performance degrades with very few annealing steps in some settings.
- Theoretical results (score-difference optimality) require vanishing step size h_k to 0, which is only approximate in practice.
Experimental setup
Datasets: Bayesian logistic regression benchmark datasets (Diabetes, Cancer, Heart, Australian, Ijcnn1, Svmguide3, German, Splice)
Benchmarks: GMM-6-8 GMM-8-10 GMM-10-12 wGMM-10-12 ExpGauss-1024 ExpGaussUV Bayesian logistic regression datasets
Metrics: Mode-Weight MSE number of modes explored accuracy log-posterior MMD Wasserstein distance sample-variance MSE
Baselines: Parallel Tempering AI-Sampler SVGD MIED CRAFT LFIS PGPS
Builds on: Continuous Normalizing Flow Benamou-Brenier dynamic optimal transport Annealed Importance Sampling Jordan-Kinderlehrer-Otto scheme
Closest papers by meaning
- Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization (ICLR 2024 · similarity 0.89)
- Latent Bayesian Optimization via Autoregressive Normalizing Flows (ICLR 2025 · similarity 0.88)
- Markovian Flow Matching: Accelerating MCMC with Continuous Normalizing Flows (NeurIPS 2024 · similarity 0.88)
- Faster Inference of Flow-Based Generative Models via Improved Data-Noise Coupling (ICLR 2025 · similarity 0.87)
- D-Flow: Differentiating through Flows for Controlled Generation (ICML 2024 · similarity 0.87)
- Discrete Neural Flow Samplers with Locally Equivariant Transformer (NeurIPS 2025 · similarity 0.87)
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.