Source-linked AI summary
Interacting particle guidance for sampling reward-tilted generative priors
Adhithyan Kalaivanan, Zheng Zhao, Jens Sjölund, Fredrik Lindsten
TL;DR
Inference-time steering seeks samples from reward-tilted generative priors, but guidance methods lack sampling guarantees and practical SMC can suffer weight degeneracy. IPG uses an RKHS corrective drift to replace reweighting with unweighted transport, and the paper reports improvements over prior work across Gaussian-mixture, image-inpainting, and protein-inference tasks.
Problem
Inference-time guidance targets reward-tilted samples, but heuristic guidance lacks sampling guarantees and SMC weights can degenerate with the few particles available for generative priors.
Source: §1 · p. 1, §2.2 · p. 3Method
Results
Takeaways & Limitations
The method provides training-free reward-tilted sampling with negligible computational overhead compared with existing SMC methods.
Source: §1 · p. 2, §6 · p. 9In multimodal distributions, IPG may struggle to move particles across large energy barriers when intermediate reward gradients are much weaker than the prior score.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Inference-time steering adapts pretrained diffusion and flow-based models to new tasks, e.g., to generate samples from a conditional distribution or samples with desired properties, without retraining. This can be formalized as sampling from a reward-tilted generative prior. As exact sampling from this distribution is intractable, guidance-based methods rely on approximations producing biased samples, and sequential Monte Carlo (SMC) methods correct for this bias using importance weights. However, while exact in the large particle limit, SMC suffers from weight degeneracy and particle collapse in practice. We propose interacting particle guidance (IPG), which replaces reweighting with transport. The particles interact through an additional drift, derived from the Feynman--Kac PDE to cancel the reweighting term, and remain unweighted. Choosing the drift in a reproducing kernel Hilbert space yields a closed-form solution that is cheap to compute, with negligible overhead compared to SMC. We demonstrate the method on Gaussian mixtures with known posteriors, and on high-dimensional image inpainting and protein structure inference tasks.
1 INTRODUCTION
Inference-time steering targets reward-tilted generative priors, but practical SMC can suffer weight degeneracy and particle collapse. The paper proposes IPG, an unweighted interacting-particle method with a closed-form corrective drift, and reports improvements over prior methods across several tasks.
1 INTRODUCTION
- Inference-time steering reuses pretrained generative models to sample from reward-tilted priors for desired properties or observed data, without task-specific retraining.The target is an exponentially reward-tilted generative prior.Source: §1 · p. 1
- Guidance heuristics lack guarantees of sampling from the target, while SMC is exact in the large-particle limit but can collapse in practice with memory-limited particle counts.In practical implementations, reweighting and resampling may act like local best-of-N search.Source: §1 · p. 1
- An RKHS corrective drift can cancel the reweighting term, enabling unweighted interacting-particle sampling with a closed-form drift that is cheap to compute.The drift is derived from the Feynman–Kac PDE and preserves training-free inference-time steering.Source: §1 · pp. 1–2
Contributions
- IPG avoids particle collapse and outperforms prior methods on Gaussian mixtures with known posteriors, high-dimensional image inpainting, and protein structure inference.The method uses flow-based generative priors for the image inpainting and protein structure tasks.Source: §1 · p. 2
2 PRELIMINARIES
Diffusion and flow models transport samples from a simple base distribution toward the data distribution. Reward tilting defines intermediate target densities and yields a Feynman–Kac PDE that can be simulated with weighted particles, though few particles can make importance sampling perform poorly.
2.1 DIFFUSION AND FLOWS
- Diffusion probability-flow ODEs and flow-based models fit a transport view in which a learned vector field moves samples from a simple base distribution to the data distribution.The intermediate marginal densities satisfy a continuity equation.Source: §2.1 · p. 2
2.2 REWARD TILTING
- A reward-tilted target is p1(x) ∝ q1(x) exp(R(x)); an interpolating path pt(x) ∝ qt(x) exp(r(x, t)) connects it to the base distribution.The path uses r(x, 0) = 0 and r(x, 1) = R(x).Source: §2.2 · p. 2
- The intermediate target densities satisfy a Feynman–Kac PDE and can be simulated with weighted particle systems, but this importance-sampling approach performs poorly for small particle counts.With few particles affordable for generative priors, high-dimensional weights can degenerate and resampling can cause particle collapse.Source: §2.2 · pp. 2–3
3 METHOD
The method adds a corrective transport drift to offset reward-based reweighting in the Feynman–Kac dynamics. An RKHS parameterization yields an unweighted interacting particle system, with a control-variate variant.
3 METHOD
- The added drift compensates for reweighting, while Langevin terms can be included without changing the weights because they leave p_t invariant.These terms can be applied jointly with transport or as separate corrector steps.Source: §3 · p. 3
3.1 INTERACTING PARTICLE GUIDANCE
- When the corrective drift cancels the centered reward and Stein terms, log-weights remain constant and the particles can be unweighted.The drift is chosen to reduce weight variance; the exact cancellation condition is generally intractable, so the method minimizes its squared residual locally at each time.Source: §3 · p. 3, §3.1 · p. 3
- IPG uses a small regularization constant to ignore weight updates, yielding an unweighted interacting particle system.The corrective drift is chosen in an RKHS and obtained through a closed-form solution to the locally minimized objective.Source: §3.1 · p. 3
- IPG-CV jointly minimizes over the control-variate estimate of E_p_t[g_t], and its weight updates can likewise be ignored with small λ.The resulting method is named for its use of a Stein control variate.Source: §3.1 · p. 4
- Unlike Stein transport, IPG uses a generative-model path with generative drift and guidance; its corrective drift accounts only for the remaining reweighting term.The cited connection describes Stein transport as moving particles from a tractable prior to the posterior by the corrective drift alone.Source: §3.1 · p. 4
4 CONNECTIONS TO PRIOR WORK
The paper places inference-time guidance alongside Monte Carlo corrections and alternative corrective-drift parameterizations. It also relates its transport framework to invariant corrector steps used in predictor-corrector methods.
4.1 GUIDANCE AND MONTE CARLO APPROXIMATIONS
- Training-free guidance adds reward gradients but does not sample from the target distribution due to approximations, while SMC corrects guidance bias with weighted particles and resampling.Other approaches use flow maps for consistent Monte Carlo estimates or MCMC targeting the correct distribution.Source: §4.1 · pp. 4–5
4.2 PARAMETRIC APPROXIMATIONS OF THE CORRECTIVE DRIFT
- Corrective drifts have been parameterized by neural networks, predefined bases, or gradient fields; predefined-basis drifts may still require reweighting and resampling.In the reported comparison, neural-network drifts achieve marginally lower mean error and SWD but higher MMD than IPG and IPG-CV, at significantly greater computational cost.Source: §4.2 · p. 5
4.3 PREDICTOR-CORRECTOR METHODS
- Langevin terms in the transport can act as an infinitesimal invariant MCMC move, and may be replaced by any continuous-time process leaving p_t invariant.The paper also considers multiple corrector steps at fixed time in SMC or IPG, including interacting dynamics such as affine-invariant Langevin dynamics and suitably noise-perturbed SVGD.Source: §4.3 · p. 5
5 EXPERIMENTS
The experiments evaluate IPG and IPG-CV against SMC-based baselines on Gaussian mixtures, image inpainting, and protein structure inference. Across these tasks, the reported results show less particle collapse with IPG variants, with task-specific differences in sample quality and runtime.
5 EXPERIMENTS
- The evaluation covers 256-dimensional GMMs, 256 × 256 AFHQ Cat inpainting, and protein structure inference using Proteina as the generative prior.All experiments use an RBF kernel with bandwidth set by the median heuristic at each time step; the image and protein tasks use 16 particles.Source: §5 · pp. 5–6, §5.1 · p. 6, §5.2 · p. 7
- Across tasks, IPG variants outperform GMM baselines across every metric, yield diverse image samples with better best-of-N fit, and IPG-CV finds a designable structure in every protein test case.The corresponding results are summarized in Tables 1–3.Source: Table 1 · p. 6, §5.1 · p. 7, Table 2 · p. 7
5.1 GAUSSIAN MIXTURE MODELS
- Without resampling, baselines have only 1–3 effective particles out of 256; resampling gives ESS/N = 1 but does not ensure accurate samples.With a small regularization constant, IPG and IPG-CV have negligible weight updates and outperform baselines across every metric.Source: §5.1 · pp. 6–7
5.2 IMAGE INPAINTING
- IPG variants produce perceptually diverse inpainting samples and take 122.9s per test image, close to FKC at 122.1s and below DriftLite at 444.3s.Average PSNR, SSIM, and LPIPS can be imperfect measures of posterior sample quality; best-of-N metrics favor IPG variants on fit to the ground truth.Source: §5.2 · p. 8
5.3 PROTEIN STRUCTURES FROM PAIRWISE DISTANCE MEASUREMENTS
- IPG variants preserve diverse protein orientations near the expected 126.5° rotation diversity; IPG-CV produces at least one designable structure in every test case.All methods have high average RMSDgt, consistent with the Proteina prior’s low designability; IPG variants’ runtimes are nearly FKC’s and far below DriftLite’s.Source: §5.3 · p. 8, Figure 3 · p. 9
6 CONCLUSION
IPG uses a training-free corrective drift to sample reward-tilted generative priors, with a closed-form RKHS solution and negligible computational overhead. The authors report improvements over prior work and identify limits in crossing energy barriers.
6 CONCLUSION
- IPG’s RKHS corrective drift targets the Feynman–Kac reweighting term, enabling unweighted sampling with negligible overhead compared with SMC.The authors report significant improvements over prior work, including SMC-based formulations, across Gaussian mixtures, image inpainting, and protein structure inference.Source: §6 · p. 9
- IPG may struggle to move particles across large energy barriers when reward gradients are much weaker than the prior score.The authors suggest tuning the interpolating path or increasing λ; the latter reintroduces particle weights and resampling while retaining corrective transport.Source: §6 · p. 9
- The authors state that experimental implementation details are described in Section 5 and the appendix, and that code will be made public upon publication.They report using coding agents for experiments and generative AI tools to polish writing.Source: REPRODUCIBILITY STATEMENT · p. 10, AI USE STATEMENT · p. 10
A REWARD-TILTED FEYNMAN–KAC PDE
The derivation differentiates the reward-tilted interpolation over time and rewrites its evolution using the continuity equation for the prior. Normalization then centers the reweighting term under the intermediate target density.
A REWARD-TILTED FEYNMAN–KAC PDE
- Differentiating the reward-tilted interpolation and applying the continuity equation yields a transport term plus a reward-dependent reweighting term in the PDE for p_t.The displayed PDE includes the time derivative of the reward, the inner product of the prior velocity and reward gradient, and the normalizer derivative.Source: A REWARD-TILTED FEYNMAN–KAC PDE · p. 13
- Differentiating the normalization condition gives the time derivative of log Z_t as an expectation under p_t.The expectation is of the time derivative of the reward plus the prior velocity’s inner product with the reward gradient.Source: A REWARD-TILTED FEYNMAN–KAC PDE · p. 13
- Substituting the normalizer derivative defines g_t(x) as the reward’s time derivative plus the prior velocity’s inner product with its spatial gradient.This is the reward-dependent quantity used in the subsequent centered evolution equation.Source: A REWARD-TILTED FEYNMAN–KAC PDE · p. 13
B DERIVATION OF IPG AND IPG-CV DRIFTS
The appendix derives regularized RKHS drift estimates from weighted particles and describes a diffusion-reweighting variant. The latter can be represented as drift-reweighting and may therefore add no capabilities.
B DERIVATION OF IPG AND IPG-CV DRIFTS
- With uniform particle weights, the general weighted equations recover the paper’s IPG and IPG-CV equations; a small λ allows weight updates to be ignored.The appendix specifies weights of 1/N for this recovery.Source: B DERIVATION OF IPG AND IPG-CV DRIFTS · pp. 13–14
- For weighted particles, the RKHS drift is obtained from a regularized linear system; IPG-CV applies a weighted projection to the solution.The derivation approximates the expectation in the objective with weighted particles and uses a Tikhonov-regression minimizer.Source: B DERIVATION OF IPG AND IPG-CV DRIFTS · pp. 13–14
- The Gram-matrix entries for the drift are cheap to compute because they use the kernel, its gradients, and terms already used in particle updates.Source: B DERIVATION OF IPG AND IPG-CV DRIFTS · p. 14
C ADDING DIFFUSION-REWEIGHTING
- Adding a symmetric positive semi-definite diffusion tensor yields a modified PDE with a diffusion term and a corrected reweighting term.The formulation rewrites the diffusion contribution using the tensor and derivatives of the target density.Source: C ADDING DIFFUSION-REWEIGHTING · p. 14
- The diffusion contribution is equivalent to a drift-reweighting term, so diffusion-reweighting may offer no additional capabilities.The equivalent drift is defined using the diffusion tensor, the target log-density gradient, and the tensor divergence.Source: C ADDING DIFFUSION-REWEIGHTING · p. 14
D ADDITIONAL DETAILS ON THE GMM EXPERIMENT
The GMM experiment uses a 256-dimensional prior and a linear-Gaussian observation reward with a known analytical posterior. It evaluates posterior agreement and particle-weight behavior under specified sampling settings.
D ADDITIONAL DETAILS ON THE GMM EXPERIMENT
- The prior is a 256-dimensional Gaussian mixture with 10 equally weighted components, and the linear-Gaussian observation likelihood defines a reward with an analytically known posterior.The component means are drawn from a uniform distribution and centered to give the mixture mean zero.Source: D ADDITIONAL DETAILS ON THE GMM… · pp. 14–15
- All methods use 500 time steps and Langevin noise scale σ_t = 3; IPG and IPG-CV use λ = 10^-3 with an adaptive-bandwidth RBF kernel.The kernel bandwidth is selected at each step as the median of pairwise distances.Source: D ADDITIONAL DETAILS ON THE GMM… · p. 15
- The experiment reports final-time ESS from self-normalized weights and posterior-accuracy metrics on 256 particles after systematic resampling.Mean error, MMD², and sliced 2-Wasserstein distance are compared against exact-posterior samples.Source: D ADDITIONAL DETAILS ON THE GMM… · p. 15
D.1 REPLACING OVERDAMPED LANGEVIN WITH NOISY SVGD
The section tests noisy-SVGD replacements for the overdamped Langevin terms and explores neural-network corrective drifts. SVGD-like interactions alone remain behind IPG, while neural drifts add substantial computation without large performance gains.
D.1 REPLACING OVERDAMPED LANGEVIN WITH NOISY SVGD
- The proposed noisy-SVGD update replaces the overdamped Langevin terms with a suitably noise-perturbed SVGD interaction targeting p_t.The equations resemble the SVGD drift up to coefficients.Source: §D.1 · p. 15
- IPG and IPG-CV still outperform FKC and DriftLite with noisy-SVGD interactions, indicating those interactions do not replace the IPG corrective drift.The IPG variants perform as well as with overdamped Langevin.Source: §D.1 · p. 15
D.2 NEURAL NETWORK BASED CORRECTIVE DRIFTS
- The neural-drift experiments use MLP and particle-interacting transformer parameterizations, run as weighted particle systems without resampling.The objective is optimized at each time step; control-variate versions carry the “-CV” suffix.Source: §D.2 · p. 16
- Neural corrective drifts achieve marginally lower mean error and SWD than IPG and IPG-CV, but higher MMD.Computing the divergence and optimizing with inner iterations significantly increases computational cost without a corresponding large performance gain.Source: §D.2 · p. 16
E ADDITIONAL DETAILS ON IMAGE INPAINTING
The image-inpainting experiments use a pretrained rectified-flow prior with central-square or right-half masks. IPG and IPG-CV produce diverse posterior samples, while FKC and DriftLite suffer severe particle collapse.
E ADDITIONAL DETAILS ON IMAGE INPAINTING
- The prior is a 1-rectified flow trained on 256 × 256 AFHQ Cat images; observations use a centered square mask or a right-half mask.The centered square has side length 0.3× the image size.Source: E ADDITIONAL DETAILS ON IMAGE INPAINTING · p. 16
- Methods use 200 uniform time steps and a Langevin noise schedule σ_t = 1 − t; image quality and diversity are assessed with PSNR, SSIM, LPIPS, pixel-space covariance trace, and mean pairwise LPIPS.Values near zero on the two diversity measures indicate particle collapse.Source: E ADDITIONAL DETAILS ON IMAGE INPAINTING · pp. 16–17
- IPG and IPG-CV produce diverse inpainting samples, whereas FKC and DriftLite suffer severe particle collapse.The samples are visualized for both half-mask and box-mask tasks.Source: E ADDITIONAL DETAILS ON IMAGE INPAINTING · p. 17
F ADDITIONAL DETAILS ON PROTEIN BACKBONE STRUCTURE INFERENCE
The protein-inference experiments steer a pretrained flow model using noisy observations of a subset of residue-pair distances and assess both observation fit and structure designability. IPG and IPG-CV produce random orientations, while the imperfect prior can yield undesignable structures near reflected ground truths.
F ADDITIONAL DETAILS ON PROTEIN BACKBONE STRUCTURE INFERENCE
- Proteina supplies α-carbon coordinates and was trained on an AFDB subset containing synthetic structures up to 256 residues; observations use 3% of residue-pair distances with 1.0 Å Gaussian noise.The observed pairs are uniformly sampled distinct residue pairs.Source: F ADDITIONAL DETAILS ON PROTEIN… · p. 17
- Structure assessment measures RMSD to observed distances and rotationally aligned ground truth, and counts a structure as designable when its self-consistency RMSD is below 2 Å.Self-consistency is tested by predicting eight sequences per structure with ProteinMPNN and refolding them with ESMFold; “Any designable” is the fraction of test cases with at least one designable structure.Source: F ADDITIONAL DETAILS ON PROTEIN… · p. 17
- The experiments leave out the heuristic generative SDE with γ ≠ 1 because its marginals differ from q_t and its Langevin score cannot be estimated from v_t.The authors leave adapting the updates and reweighting term for that setting to future work.Source: F ADDITIONAL DETAILS ON PROTEIN… · p. 17
- IPG and IPG-CV produce random orientations as expected, while FKC and DriftLite show particle collapse; IPG diversity metrics for 30JI are bimodal because some samples lie near reflected ground truths.The reflected-near-ground-truth structures are undesignable, consistent with Proteina’s designability rate below 20%.Source: F ADDITIONAL DETAILS ON PROTEIN… · pp. 17–18, Figure 4 · p. 19