Technical Summary: URGE – Simple Approximation and Derivative Free Inference-Time Scaling for Diffusion Models via Sequential Monte Carlo on Path Measures

1. Problem Statement

Modern diffusion models and flow-based generative methods are typically trained to capture the data distribution $p_{\text{data}}$. However, deployment often requires steering generation toward specific downstream objectives (e.g., physical validity, text-image alignment) without retraining. This is formalized as sampling from a reward-tilted posterior $q(x) \propto p_{\text{data}}(x)r(x)$.

Current inference-time scaling methods, particularly guidance, attempt to achieve this by modifying the drift of the sampling SDE with a gradient term derived from a guidance potential $G$. While empirically effective, guidance introduces a fundamental mathematical limitation: it does not guarantee sampling from the exact tilted distribution $q(x)$. Exact sampling requires the Doob's $h$-transform, which involves solving a backward Kolmogorov equation and is often intractable in high dimensions.

Existing alternatives, such as FK-Corrector and AFDPS, attempt to correct the bias of guided generation using Sequential Monte Carlo (SMC) by reweighting particles. However, these methods rely on infinitesimal generators that require access to higher-order derivatives of the reward function (specifically the score $\nabla \log p_t$, the Laplacian $\Delta r$, and the Hessian of the guidance). These requirements incur significant computational overhead and render these methods inapplicable to black-box reward models (e.g., neural reward models used in text-to-image generation) where derivatives are unavailable or expensive.

The core challenge addressed by this work is: Can we design an approximation-free inference-time scaling procedure that eliminates the bias of suboptimal guidance without requiring derivative information of the reward?

2. Methodology: URGE

The authors propose URGE (Unbiased Resampling via Girsanov Estimation), a derivative-free inference-time scaling algorithm that performs resampling over generation trajectories (path space) rather than individual generated samples (particle space).

Core Mechanism

URGE leverages the Girsanov theorem to perform a change of measure on the path space of the diffusion process.

1. Guided Generation: Multiple trajectories $\{X^{(i),G}_t\}$ are simulated using a guided SDE with an arbitrary guidance potential $G$ (which may be imperfect).
   $$dX^G_t = (v(X^G_t, t) + V(t)^2 \nabla_x G(X^G_t, t))dt + V(t)dW_t$$
2. Pathwise Reweighting: Instead of computing weights based on infinitesimal generators (which require derivatives), URGE assigns a multiplicative weight to each entire trajectory based on the Radon-Nikodym derivative between the target reward-tilted path measure $Q$ and the guided path measure $P_G$.
   The importance weight for a trajectory over $[s, t]$ is:
   $$w^{\text{URGE}}_{s,t} = \exp\left( \int_s^t -V(\tau)\nabla_x G(X^G_\tau, \tau)^\top dW_\tau - \frac{1}{2}\int_s^t V(\tau)^2 \|\nabla_x G(X^G_\tau, \tau)\|^2 d\tau + \int_s^t dr(X^G_\tau, \tau) \right)$$
   Crucially, this formulation does not require $\nabla r$, $\Delta r$, or $\nabla \log p_t$. It only requires the reward function $r$ and the guidance $G$.
3. Resampling: At discrete time steps, particles are resampled proportionally to their path weights. This implements a "birth-death" mechanism where high-reward trajectories are replicated and low-reward ones are eliminated, driving the empirical distribution toward the true posterior $q$.

Theoretical Equivalence

The paper establishes a theoretical equivalence between path-space resampling (URGE) and particle-space resampling (FK-Corrector/AFDPS).

• The authors prove that the marginalized generator of the URGE path-space process, when conditioned on the terminal state, recovers the exact infinitesimal reweighting term used by FK-Corrector and AFDPS.
• This confirms that URGE is approximation-free: it converges to the same target distribution as the theoretically optimal (but derivative-heavy) particle-space methods, but achieves this by lifting the weight construction to the path space, thereby bypassing the need for explicit derivative calculations.

3. Key Contributions

1. Novel Mixture-of-Experts Framework: The authors introduce a method that integrates SMC with generative processes by resampling over trajectories rather than samples. This allows for flexible weight construction via Girsanov change of measure without requiring reward derivatives.
2. Derivative-Free Implementation: URGE eliminates the need for score functions, Hessians, or Laplacians of the reward. This makes it applicable to black-box reward models (e.g., CLIP, ImageReward, DPO-tuned models) where gradients are inaccessible.
3. Theoretical Equivalence: The paper proves that URGE's pathwise reweighting, when marginalized, exactly recovers the weighting structure of state-of-the-art derivative-based methods (FK-Corrector/AFDPS), guaranteeing that both schemes produce the same approximation-free terminal law.
4. Empirical Superiority: URGE consistently outperforms existing baselines (including FK-Steering and AFDPS variants) across synthetic and real-world tasks while being significantly simpler to implement and fully gradient-free.

4. Experimental Results

The authors evaluate URGE on three benchmarks:

• Gaussian Mixture Model (Synthetic):

  • URGE achieves the lowest error across all metrics (MMD, SWD, Mean L2, Covariance Frob) compared to Pure Guidance, AFDPS, and FK-Steering.
  • It demonstrates superior convergence efficiency, achieving better results with fewer particles than baselines.
  • URGE remains robust across varying dimensions and discretization steps, showing less sensitivity to step size than competing methods.

• Inverse Problems (Image Restoration):

  • Tested on Gaussian Deblurring, Motion Deblurring, Super Resolution, and Box Inpainting using ImageNet-256 and FFHQ-256.
  • URGE achieves performance comparable to AFDPS (which requires derivatives) and significantly outperforms other baselines in PSNR and LPIPS.
  • URGE demonstrates better computational scaling, achieving shorter inference times than AFDPS-ODE while maintaining high reconstruction quality.

• Text-to-Image Generation:

  • Tested on Stable Diffusion v1.5 and SDXL with DPO-tuned variants.
  • URGE outperforms FK-Steering and the base model on aesthetic quality (ImageReward, HPS) and prompt fidelity (GenEval, CLIP-Score).
  • Notably, URGE allows a smaller model (SD v1.5) to match or approach the performance of a larger model (SDXL) when using particle scaling, whereas FK-Steering shows diminishing returns.
  • URGE exhibits monotonic reward improvement with increasing particle count, unlike FK-Steering.

5. Significance and Claims

The paper claims that URGE provides a training-free, approximation-free ensemble-based posterior sampling framework. Its primary significance lies in unifying path-based and state-based weighting strategies. By lifting the resampling mechanism from particle space to path space, URGE:

• Removes the derivative bottleneck: It enables the use of complex, black-box reward models in diffusion guidance without the computational cost of calculating high-order derivatives.
• Ensures theoretical correctness: It guarantees convergence to the true reward-tilted posterior, correcting the inherent bias of standard guidance methods.
• Offers practical efficiency: It achieves state-of-the-art generation quality with simpler implementation and competitive computational overhead, making it suitable for large-scale, real-world generative tasks.

The authors emphasize that while previous works (FK-Corrector, AFDPS) provided theoretical insights into correcting guidance, they were limited by their reliance on derivative information. URGE bridges this gap, offering a practical, derivative-free solution that retains the theoretical guarantees of approximation-free sampling.