Research map / Probabilistic and scientific modeling / Variational and generative sampling
Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization
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 DGFS, a GFlowNet-inspired diffusion sampler that learns a flow function so training can use partial trajectory segments and intermediate learning signals, yielding more accurate log-normalizing-constant estimates than PIS, DDS, SMC, and VI baselines.
Problem
Sampling from unnormalized high-dimensional densities with diffusion-based controlled stochastic processes suffers from full-trajectory training requirements and terminal-only learning signals, causing credit assignment and high gradient variance issues.
Approach
Introduces an amortized learned flow (state flow) function trained with sub-trajectory balance losses over short partial segments, plus a forward-looking log-reward shaping that gives intermediate learning signals, connecting diffusion samplers to GFlowNets.
Key results
- DGFS achieves log-normalizer bias of 0.019 on MoG, 0.274 on Funnel, 0.904 on Manywell, 0.180 on VAE, and 8.974 on Cox, best among diffusion-based samplers (PIS, DDS).
- DGFS reduces gradient variance relative to PIS on Funnel, per the paper's gradient variance analysis.
- Without gradient information in the drift parameterization (DGFS-NN), DGFS remains competitive, reaching 0.023 on MoG.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use DGFS instead of PIS | diffusion sampler · MoG | absolute log normalizing factor estimation bias: 0.036 -> 0.019 | Helped |
| use DGFS instead of PIS | diffusion sampler · Funnel | absolute log normalizing factor estimation bias: 0.305 -> 0.274 | Helped |
| use DGFS instead of PIS | diffusion sampler · Manywell | absolute log normalizing factor estimation bias: 1.391 -> 0.904 | Helped |
| use DGFS instead of PIS | diffusion sampler · VAE | absolute log normalizing factor estimation bias: 2.049 -> 0.180 | Helped |
| use DGFS instead of PIS | diffusion sampler · Cox | absolute log normalizing factor estimation bias: 11.28 -> 8.974 | Helped |
| add intermediate forward-looking signals (ablation) | DGFS · VAE | absolute log normalizing factor estimation bias | Helped |
| use VP rather than VE diffusion modeling (ablation) | DGFS · VAE | absolute log normalizing factor estimation bias: 0.180 -> 0.86 | Hurt |
| remove target log-density gradient from drift parameterization (ablation) | DGFS · MoG, Funnel, Manywell, VAE, and Cox | absolute log normalizing factor estimation bias: MoG: 0.019 -> 0.023; Funnel: 0.274 -> 0.354; Manywell: 0.904 -> 2.659; VAE: 0.180 -> 0.496; Cox: 8.974 -> 109.0 | Mixed |
| anneal rollout variance coefficient from 1.1 to 1 during initial training (ablation) | DGFS (DGFS+) · MoG | absolute log normalizing factor estimation bias: 0.019 -> 0.012 | Helped |
| anneal rollout variance coefficient from 1.1 to 1 during initial training (ablation) | DGFS (DGFS+) · Funnel, Manywell, and Cox | absolute log normalizing factor estimation bias: Funnel: 0.274 -> 0.322; Manywell: 0.904 -> 0.932; Cox: 8.974 -> 10.23 | Hurt |
| anneal rollout variance coefficient from 1.1 to 1 during initial training (ablation) | DGFS (DGFS+) · VAE | absolute log normalizing factor estimation bias: 0.180 -> 0.162 | Helped |
Surprising findings
- The DDS implementation failed (NaN) on the Cox task in the PyTorch reproduction, so it is reported as N/A.
- The GFlowNet-style trajectory balance variant without the flow-function approach (TB) and the prior continuous GFlowNet method performed much worse on Cox (e.g., prior method ~2728 bias).
Limitations
- Results mostly on synthetic benchmark targets; the authors note the approach is not yet demonstrated on real scientific tasks.
- Intermediate signals use a straightforward forward-looking reward design that the authors say could be improved.
- Exploration with off-policy rollouts is only preliminary (DGFS+ does not uniformly improve).
Experimental setup
Datasets: MNIST
Benchmarks: Mixture of Gaussians Funnel Manywell VAE latent posterior Cox process
Metrics: absolute bias of estimated log normalizing constant standard deviation across 5 seeds
Baselines: SMC VI-NF CRAFT FAB w/ Buffer PIS DDS TB
Builds on: path integral sampler denoising diffusion sampler GFlowNet subtrajectory balance forward-looking GFlowNet
Closest papers by meaning
- Flow Factorization for Efficient Generative Flow Networks (AAAI 2025 · similarity 0.90)
- Analysis of Learning a Flow-based Generative Model from Limited Sample Complexity (ICLR 2024 · similarity 0.90)
- Stop the Sampler! Classifier-Based Adaptive Stopping for Sampling Kernels (NeurIPS 2026 · similarity 0.89)
- Annealing Flow Generative Models Towards Sampling High-Dimensional and Multi-Modal Distributions (ICML 2025 · similarity 0.89)
- Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling (NeurIPS 2024 · similarity 0.89)
- Beyond Squared Error: Exploring Loss Design for Enhanced Training of Generative Flow Networks (ICLR 2025 · similarity 0.89)
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.