Research map / Image generation / Fast diffusion sampling
Amortizing intractable inference in diffusion models for vision, language, and control
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 relative trajectory balance (RTB), a data-free, asymptotically unbiased GFlowNet-style objective for fine-tuning diffusion priors to sample posteriors p(x)r(x), tested on classifier-guided image generation, text-to-image alignment, discrete-diffusion text infilling, and offline RL policy extraction.
Problem
Sampling from a posterior p(x)r(x) when the prior is a diffusion model and the constraint r(x) is a black-box function is intractable; existing guidance and RL fine-tuning methods are biased or mode-collapse.
Approach
Derives a relative trajectory balance constraint, Z·p_post(trajectory) = r(x1)·p_prior(trajectory), from the view of diffusion models as continuous GFlowNets, and trains a posterior diffusion model (initialized from the prior, often with LoRA) by minimizing its squared log-ratio over trajectories, optionally with off-policy exploration and a Langevin-style parametrization. It needs no samples from
Key results
- On MNIST/CIFAR-10 classifier-guided posteriors, RTB achieves high expected log r(x) (-0.17 on MNIST single-digit) with FID comparable to the best baselines and diversity comparable to classifier guidance, while DDPO mode-collapses.
- On story infilling with a SEDD discrete diffusion prior, RTB reaches BERTScore 0.156 versus 0.092 for prompting with x,y and 0.102 for autoregressive GFlowNet fine-tuning.
- On D4RL, RTB matches state-of-the-art with e.g. 100.40 on hopper-medium-replay and 93.57 on walker2d-medium-replay, but is below QGPO on several medium-expert tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune the diffusion prior with RTB rather than DPOK | Unconditional image diffusion posterior sampling · MNIST | FID: 1.2063±0.316 -> 1.1823±0.288 | Helped |
| Use RTB instead of DPOK | Unconditional image diffusion posterior sampling · MNIST even/odd | FID: 1.2536±0.206 -> 1.1794±0.171 | Helped |
| Use RTB instead of DPOK | Unconditional image diffusion posterior sampling · CIFAR-10 | FID: 0.5316±0.157 -> 0.4717±0.138 | Helped |
| Fine-tune the discrete diffusion prior with RTB rather than prompting the autoregressive model via GFlowNet | ROCStories infilling · ROCStories infilling | BLEU-4: 0.019±0.001 -> 0.025±0.002 | Helped |
| Train the posterior diffusion policy with RTB | Offline RL policy extraction · halfcheetah-medium-replay | average reward: 47.6±1.4 -> 48.11±0.56 | Helped |
| Train the posterior diffusion policy with RTB | Offline RL policy extraction · hopper-medium-replay | average reward: 100.7±0.6 -> 100.40±0.21 | Mixed |
| Train the posterior diffusion policy with RTB | Offline RL policy extraction · walker2d-medium | average reward: 86.0±0.7 -> 87.29±3.15 | Helped |
| Mix online and off-policy training instead of training online only (ablation) | RTB offline RL policy extraction · halfcheetah-medium-replay | average reward: 46.88±0.51 -> 48.11±0.56 | Helped |
| Mix online and off-policy training instead of training online only (ablation) | RTB offline RL policy extraction · hopper-medium-replay | average reward: 99.23±3.22 -> 100.40±0.21 | Helped |
| Mix online and off-policy training instead of training online only (ablation) | RTB offline RL policy extraction · walker2d-medium-replay | average reward: 94.01±0.28 -> 93.57±2.63 | Hurt |
| Remove KL regularization from RL fine-tuning | Image diffusion posterior sampling · CIFAR-10 | FID: 0.5316±0.157 -> 1.7686±0.589 | Hurt |
| Use conditional rather than unconditional posterior architecture (ablation) | MNIST even/odd posterior sampling · MNIST even/odd | FID: 1.1794±0.171 -> 0.9112 | Helped |
Surprising findings
- Pure RL fine-tuning without KL regularization achieves high reward but collapses to a few modes, and classifier-guidance baselines score low on expected log reward despite high diversity.
- RTB does not outperform all baselines on D4RL: on halfcheetah-medium-expert (74.93) and hopper-medium-expert (96.71) it trails methods such as QGPO and IDQL.
Limitations
- Simulation-based training is slow and memory-intensive.
- The trajectory-level loss gives no local credit assignment, so gradient variance is high.
- No error guarantees are provided for prior imperfection, amortization, or time discretization.
Experimental setup
Datasets: MNIST CIFAR-10 ROCStories D4RL halfcheetah/hopper/walker2d
Benchmarks: MNIST CIFAR-10 ROCStories infilling Stable Diffusion ImageReward prompts D4RL locomotion
Metrics: E[log r(x)] FID diversity (cosine distance) average reward BLEU-4 GLEU-4 BERTScore normalized return
Baselines: DPS LGD-MC DDPO DPOK Prompt (x) Prompt (x, y) GFN fine-tuning SFT BC CQL IQL Diffuser
Builds on: trajectory balance DPOK DDPO VarGrad IQL score entropy discrete diffusion
Closest papers by meaning
- Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models (ICML 2025 · similarity 0.90)
- Training Diffusion Models with Reinforcement Learning (ICLR 2024 · similarity 0.89)
- Prior-Guided Diffusion Planning for Offline Reinforcement Learning (NeurIPS 2025 · similarity 0.87)
- Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages (NeurIPS 2026 · similarity 0.87)
- Improved off-policy training of diffusion samplers (NeurIPS 2024 · similarity 0.87)
- What Matters When Repurposing Diffusion Models for General Dense Perception Tasks? (ICLR 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.