Research map / Image generation / Fast diffusion sampling

A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Diffusion theory and learned samplers · 0 citations (OpenAlex)

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.

Formulates inference-time steering of diffusion models under distribution-level rewards as targeting a tilted measure via a mean-field framework, and derives a weighted McKean-Vlasov particle sampler with Feynman-Kac-style reweighting.

Problem

Inference-time steering of diffusion models has principled guarantees only for pointwise rewards; distribution-level rewards (calibration to population data, diversity, MMD matching) are steered heuristically with no clear target distribution.

Approach

Treat distributional steering as sampling a measure-dependent exponential tilt mu* ∝ p_1 exp(Psi(x, mu*)), derive weighted McKean-Vlasov dynamics with log-weight updates including a time-derivative of the first variation, and approximate them with weighted interacting particles; pointwise FK steering is recovered as a special case.

Key results

• In the 1D bimodal Gaussian MMD tilt, mean-field steering attains the smallest objective gap at the correct lambda, while gradient steering samples from a measure depending on steering strength.
• For HIV-1 protease, the objective proxy is minimised at the correct steering coefficient lambda, with generations matching the 25%/75% DEER-inferred proportions proxy.
• Fixed-point iteration with K_FP=3 or 5 is reported as highly accurate, while K_FP=1 shows a slight bias.

Surprising findings

• Gradient steering alone is reported as often effective in practice, so exact distributional optimality is not always necessary.

Limitations

• Relies on resampling, which can cause duplicate images in image modalities and weight degeneracy.
• Requires a differentiable reward and extra cost for the time-derivative solver.
• Reweighting can only redistribute mass among regions populated by existing particles.

Experimental setup

Datasets: HIV-1 protease DEER-based state proportions adenylate kinase MD bond angles PDB 4OLE electron density

Benchmarks: 1D Gaussian mixture MMD tilting HIV-1 protease residue-55 distance distribution PDB 4OLE electron density cosine alignment

Metrics: objective gap J(mu)-J(mu*) cosine alignment proportion of samples in conformation B

Baselines: gradient steering

Builds on: Feynman-Kac steering gradient guidance Particle Guidance Boltz-2

Closest papers by meaning

• A General Framework for Inference-time Scaling and Steering of Diffusion Models (ICML 2025 · similarity 0.88)
• Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design (ICML 2025 · similarity 0.88)
• Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities (NeurIPS 2025 · similarity 0.88)
• Hierarchical Variational Policies for Reward-Guided Diffusion (NeurIPS 2026 · similarity 0.88)
• On scalable and efficient training of diffusion samplers (NeurIPS 2025 · similarity 0.88)
• Particle Guidance: non-I.I.D. Diverse Sampling with Diffusion Models (ICLR 2024 · similarity 0.87)