Research map / Probabilistic and scientific modeling / Variational and generative sampling
NETS: A Non-equilibrium Transport 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 NETS, a continuous-time annealed Langevin sampler augmented with a learned drift (trained via PINN or action-matching losses without backpropagating through the SDE) whose weights are reduced in variance, and shows improved ESS on GMM, funnel, MoS, and lattice φ⁴ benchmarks.
Problem
Sampling unnormalized distributions that are non-log-concave, where annealed importance sampling and Langevin methods suffer high-variance importance weights or slow mixing.
Approach
Augment the annealed Langevin SDE with a learned additional drift that transports samples toward the target, so Jarzynski importance weights have much lower variance; the drift is the minimizer of PINN or action-matching objectives that can be estimated without backpropagating through the SDE, and the PINN loss upper-bounds KL divergence.
Key results
- On the 40-mode GMM (d=2), NETS-PINN-resample reaches 0.993 ± 0.004 ESS versus 0.734 for the best prior baseline (iDEM).
- On the 50-d Mixture of Student-T, NETS-PINN-resample achieves MMD 0.030 vs 0.093 for FAB, and W2 400 vs 1204.
- On lattice phi^4 at L=20, NETS is almost 2 orders of magnitude more statistically efficient than AIS in ESS.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add one resampling step to NETS-PINN (ablation) | NETS-PINN · 40-mode GMM (d=2) | ESS: 0.979 ± 0.002 -> 0.993 ± 0.004 | Helped |
| Increase diffusion from ε=0 to ε=4 (ablation) | NETS-PINN · 40-mode GMM (d=2) | ESS: 0.954 ± 0.003 -> 0.979 ± 0.002 | Helped |
| Increase diffusion from ε=0 to ε=4 (ablation) | NETS-PINN · 40-mode GMM (d=2) | W₂: 3.55 ± 0.57 -> 3.14 ± 0.46 | Helped |
| Add resampling during generation (ablation) | NETS-PINN · Neal's Funnel (d=10) | MMD: 0.033 ± 0.002 -> 0.027 ± 0.003 | Helped |
| Add resampling during generation (ablation) | NETS-PINN · Neal's Funnel (d=10) | W₂: 388.91 ± 141.5 -> 343.78 ± 65.25 | Helped |
| Add resampling during generation (ablation) | NETS-PINN · Mixture of Student-T (d=50) | MMD: 0.032 ± 0.001 -> 0.030 ± 0.000 | Helped |
| Add resampling during generation (ablation) | NETS-PINN · Mixture of Student-T (d=50) | W₂: 482.393 ± 174.6 -> 400.076 ± 59.31 | Helped |
| Use annealed Langevin dynamics without learned transport (ablation) | 40-mode GMM sampler · 40-mode GMM (d=2) | ESS: 0% | Hurt |
| Increase diffusivity when using learned transport (ablation) | NETS-PINN · 8-mode high-dimensional GMMs (d=36, 64, 128, 200) | ESS: null | Helped |
| Add learned transport to AIS | AIS · 8-mode GMM (d=200) | ESS: ≈ 0 -> 60% | Helped |
| Add resampling to NETS (ablation) | NETS · Neal's Funnel and Mixture of Student-T | MMD: null | Helped |
| Use NETS instead of AIS without transport | AIS · Lattice φ⁴ theory at and past the phase transition | ESS: null | Helped |
Limitations
- Lattice phi^4 near phase transitions requires 1500-2000 integrator steps, so slowdowns persist.
- On Neal's funnel NETS is slightly worse in W2 than GMMVI and FAB.
- Larger diffusivity demands smaller step sizes and more integrator steps (up to K=2000 for eps=80).
Experimental setup
Benchmarks: 40-mode GMM Neal's funnel Mixture of Student-T (50-d) lattice phi^4 theory
Metrics: ESS W2 MMD
Baselines: FAB PIS DDS iDEM pDEM CMCD AIS GMMVI
Builds on: annealed importance sampling controlled monte carlo diffusions physics-informed neural networks jarzynski equality
Closest papers by meaning
- Scalable Discrete Diffusion Samplers: Combinatorial Optimization and Statistical Physics (ICLR 2025 · similarity 0.90)
- Sequential Controlled Langevin Diffusions (ICLR 2025 · similarity 0.89)
- Transport meets Variational Inference: Controlled Monte Carlo Diffusions (ICLR 2024 · similarity 0.88)
- Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling (NeurIPS 2026 · similarity 0.88)
- Sampling in Unit Time with Kernel Fisher-Rao Flow (ICML 2024 · 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.