Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts

We thank the reviewer for the detailed feedback and constructive suggestions. We are glad that they find the proposed idea 'novel' and its derivation 'supported by valid theory'. We are also happy to hear that our method 'contributes to the understanding of Diffusion Models.' We provide the suggested comparisons and answer the reviewer's questions below.

Testing the metrics on large-scale datasets (at least 1k prompts)

We thank the reviewer for the constructive feedback which has led to some interesting findings and a deeper understanding of this work. We evaluate FKC on two large-scale tasks: on ImageNet-1K using EDM2 [4] and on the Geneval benchmark [1] using SDXL, which evaluates the algorithms on a predefined set of prompts and measures adherence to the prompts. These additional large-scale experiments have led to two new insights:

1. FKC improves image quality and prompt adherence in ambient space but not latent space. Point 1 is supported by the following two tables. We can see substantial improvement in the performance of FKC on an ambient diffusion model but not on a latent diffusion model.

• EDM2 We first conduct a new set of image experiments using the EDM2 model [4], which directly generates outputs in pixel-space. We compare generations from the baseline method, which uses CFG, with generations guided by FKC on ImageNet-1K and evaluate 10k samples using two metrics: CLIP Score and ImageReward [5], which reflects how closely they align with human preferences.

We find that incorporating FKC improves both scores and qualitatively generates better images; we include examples in this PDF. We use default settings of $\beta=1.4$ for both models.

• SDXL Geneval The table below consists of three parts: Geneval benchmarks, performance of FKC on the prompts from Geneval (553 prompts), and on 1k prompts (553 + newly generated). We also rerun SDXL with the same hyperparameters as our algorithm for a fair comparison.

2. FKC improves performance in the ambient space across tasks whereas improvements in latent space diffusion models are limited. This is supported by Table 2 and Figure 2 of the PDF. Here, we present new and stronger results for the molecule generation task in the coordinate space instead of the latent space.

Is it possible to add particle interactive force to the proposed FKC method to mitigate the sampling covariance?

Indeed, one could combine the proposed method with the Stein Variational Gradient Descent, targeting the intermediate marginals (annealed density or the product of densities) similar to [2]. This would require adding a term to the weights, which, theoretically, should reduce the variance of the weights.

Сonnection to the Fisher-Rao gradient flow

The reweighting equation $\partial_t p_t^w(x) = p_t^w(x) \bar{g}_t(x)$ does correspond to the Fisher-Rao gradient flow of a linear functional $G[p_t] = \int {g}_t(x) p_t(x) dx$ (after constraining $p_t^w$ to remain normalized). The Wasserstein Fisher-Rao gradient flow has a form similar to the Feynman-Kac PDE with $\sigma_t=0$ (deterministic evolution). We have included a note about this in the Appendix of the updated revision (see e.g. Thm. 3.1 in [3]).

Closing remarks

We thank the reviewer for suggesting this comparison, which significantly improves our empirical study. We hope that our answers address all the important questions raised by the reviewer, and we are more than happy to consider any additional questions or further suggestions.

[1] Ghosh, Dhruba et al. "Geneval: An object-focused framework for evaluating text-to-image alignment."
[2] Corso, Gabriele et. al. "Particle guidance: non-iid diverse sampling with diffusion models."
[3] Lu, Yulong, et al. "Accelerating Langevin sampling with birth-death."
[4] Karras, Tero et al. "Analyzing and Improving the Training Dynamics of Diffusion Models."
[5] Xu, Jiazheng et al. "ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation"

We thank the reviewer for the constructive feedback. We are glad that the reviewer finds our paper to be 'well-structured and complete' and the experimental design to be 'reasonable'. We also would like to thank the reviewer for bringing up reference [1] (Du et al.), as this is an excellent reference for clarifying our contributions.

Before addressing the concerns raised, we would like to clarify a potential misunderstanding about the proposed method.

1. The main objective of the proposed method is to modify the target density at the inference time (e.g. by sampling from the annealed distribution or product of densities) rather than 'sampling from the training data and improving sampling efficiency'.
2. In complete agreement with [1], our work shows that simulating the reverse SDEs with modified scores does not sample from the target densities. This is exactly the goal of introducing the Feynman-Kac corrector scheme, which, as we prove, allows for consistent sampling from the target densities by re-weighting and re-sampling. We have added a self-contained proof of Eq. 9 in the updated Appendix to further emphasize this claim.

Next, we would like to address the salient concerns raised by the reviewer individually:

the choice of t_\max, t_\min

The choice of the interval [t_\min, t_\max] depends on the hyperparameters and, especially, the noise schedule, but does not require significant tuning. In practice, we choose the interval based on the fraction of unique samples resampled at every iteration. We report the corresponding plot in Fig. 1 of the PDF for the annealing experiments. Note that the scale of the weights is proportional to the noise $\sigma_t$ (see Prop 3.2), which results in a low number of unique samples close to $t=1.0$ due to the Variance Exploding schedule.

For molecule generation experiments in Table 3, we selected t_\max based on a validation task by going over the grid $[0.6,0.7,0.8,0.9,1.0]$; t_\min was always set to $0$.

We have also added a new set of molecule experiments for a harder set of tasks using a model that directly predicts the 3D coordinates of the atoms in a molecule from [2] (Zhou et al.) and docks molecules to a set of target protein pairs. Again, we use a validation task to sweep over t_\max values, which we report in Table 2 of the PDF.

Ablation study for the resampling methods

We perform the ablation study of the two considered resampling methods in Table 6 of the Appendix. We find that systematic resampling always performs better or comparably and opt to use it throughout the rest of the empirical study.

Computational cost of the resampling step

The additional computational cost of the resampling step is negligible compared to the reverse SDE simulation. Indeed, all the weights depend only on scores, which are already evaluated at the forward pass. We purposefully avoid the computation of the divergence operators when deriving the integration scheme (see Lines 176-189 L). The cost of the systematic resampling step is equivalent to generating one uniform random variable, which is negligible.

Our proposed method is notably more convenient than [1], which requires changing the diffusion model parameterization and performing additional MCMC steps or Metropolis-Hastings tests.

$\beta = 7.5$ for FKC on SDXL

As requested by the reviewer, we provide results for both values of $\beta$ on the Geneval benchmark [3] (Dhruba Ghosh et al.). See the detailed image-generation study in response to Reviewer 65RJ.

Derivation of equation (6)

Equation (6) can be understood by the separation of variables. Indeed, $\frac{\partial_t p_t^w(x)}{p_t^w(x)}=\partial_t\log p_t^w(x)=\bar{g}_t(x)\implies p_t^w(x)=p_0^w(x)\exp(\int_0^t ds \bar{g}_s (x)),$ where the exponent part corresponds to the update of the weights $dw_t=\bar{g}_t(x)dt$. Note that in eq. (9), $w$ appears as log-weights for Self-Normalized Importance Sampling.

Note that Eq. (6) is meant to introduce the reweighting evolution. In later sections, we derive the weights which preserve the particular target distribution evolution under simulation or transport by an SDE with given drift/diffusion.

Closing remarks

Again, we thank the reviewer for their questions, which gave us the opportunity to clarify and improve our work. We hope that our answers fully address all the important questions raised by the reviewer, and we are happy to consider any additional questions or further suggestions.

We kindly ask the reviewer to consider increasing their score if our responses address their concerns satisfactorily.

We thank the reviewer for their detailed feedback! We are thrilled to hear that they find our paper to be a 'useful contribution to the literature', 'well written and organized' and a new tool for the diffusion sampling community. Below, we address the reviewer's questions.

What is the interval [t_\min, t_\max] in the experiments used? Does it require significant tuning?

The choice of the interval [t_\min, t_\max] depends on the hyperparameters and, especially, the noise schedule, but does not require significant tuning. In practice, we chose the interval based on the fraction of unique samples resampled at every iteration during the resampling step. We report the corresponding plot in Fig. 2 of the PDF for the annealing experiments. Note that the scale of the weights is proportional to the noise $\sigma_t$ (see Prop 3.2), which results in a low number of unique samples close to $t=1.0$ where the variance of the noise is the largest due to the Variance Exploding schedule. We have added corresponding plots and discussion to the manuscript.

Why is dividing the weights by some constant factor $T$ is necessary? I suppose it makes sense that the particle resampling may be especially high variance in high dimensions?

Indeed, as the reviewer points out, the variance of the weights grows with the number of dimensions due to being proportional to the norm of the score. For the latent molecule generation (8000 dimensions), we had to divide the weights by $T=100$ and clip lowest and largest $20\\%$ of the weights in order to achieve reasonable variance of the weights. To verify that the need for such heuristics is caused by the dimensionality, we did a new empirical study for molecule generation in the coordinate space, where the dim. is 3 * number of atoms in the molecule (<100). Here, we neither divided nor clipped the weights and only chose the interval t_\max based on a validation set. We present new results below and have added them to the manuscript.

Is early stopping the right phrase to use here?

Thank you for suggesting a better explanation! The closest analog to selecting [t_\min, t_\max] is the reducing the time interval for the integration of diffusion models, i.e. instead of integrating $t\in[0,1]$, one usually integrates $t\in[ε,1]$, where ε is small constant. We have changed this explanation in the manuscript.

Practical details on image generation

For image generation, we used Stable Diffusion XL with a variance-preserving SDE and an Euler–Maruyama solver, running for 100 steps. All the computations are done with float16 precision, which is crucial for stable generation via SDXL. All images are generated in 1024x1024 resolution. In the FKC portion of our experiments, we did not apply weight rescaling; moreover, we resampled all 64 particles after each time step throughout the entire time interval.

New molecule experiments

We did a new set of molecule experiments using a model that directly predicts the 3D coordinates of the atoms in a molecule from [1] on a harder set of tasks of generating molecules that dock to a pair of proteins simultaneously, expanding on the promising results from Table 4 of the original submission. For 14 proteins pairs, we generated 32 molecules at 5 different molecule sizes (160 molecules per pair) using FKC product and found that it improved nicely over two SOTA methods:

---

Table: Docking scores of generated ligands for 14 protein target pairs (P₁, P₂). Lower docking scores are better.

---

[1] Zhou et al. "Reprogramming Pretrained Target-Specific Diffusion Models for Dual-Target Drug Design." NeurIPS (2024). [2] Guan et al. ICLR (2023).

Implementation details

For molecules, we chose the time interval based on a validation set of 1 protein pair at 5 molecule lengths (32x5 generated molecules). We kept $t_{\text{min}}$ at 0 and sweep over values of $t_{\text{max}}$ for $\beta=2$, SDE type="Target score", and FKC="on". We show this and additional ablations over $\beta$ and SDE type in this PDF. Setting $t_{\text{max}}$ to 0.6 gives a good tradeoff in terms of generating molecules that perform well vs. maintaining diversity, and so we proceed with $t_{\text{max}} = 0.6$.

We thank the reviewer for their time, feedback, and positive appraisal of our work. We are heartened to hear that the reveiwer feels that the paper is "theoretically sound" and found that the paper was "generally well-structured". We now address the questions and suggestions raised by the reviewer.

When comparing Eq.(5) and Eq.(7), an additional term is added. I am curious whether this arises from the evolving normalizing constant $Z_t$ in the context of sampling scenario? In contrast, does the vanilla FP equation in Eq. (5) describe the evolution of distributions while keeping $Z_t$ unchanged?

Both PDEs in eqs. (5) and (7) preserve the normalization of the density in the sense that $\int dx ~ p_t(x) = 1$ for all $t$, although the evolution of density changes with the addition of weighting terms.

Note that the normalization constants $Z_t$ defined in eq. (13) change in time to guarantee that the density on the left hand side is normalized, e.g. $\int dx ~ p_t^{\text{prod}}(x) = 1$. This is due to defining the density up to a constant, i.e. $p_t^{\text{prod}}(x) \propto q_t^1(x)q_t^2(x)$.
Note that, while the integral term $\int g_t(x) p_t(x) dx$ in eq. (6) ensures normalization for the density evolution PDE, our SMC resampling can proceed with access to only the unnormalized weights $g_t(x)$ (see eq. (9))

To obtain FK PDE in Eq.(7), is it simply the addition of the FP equation and the reweighting equation? I guess there may be something missing before the definition of the FK PDE.

FK PDE in eq. (7) is a more general type of PDEs than the FP equation in eq. (5). Indeed, the main difference between these PDEs is the reweighting term (the last term of FK PDE, i.e. eq. 6) that allows for simulation of processes not possible with FP equation. For instance, FK PDE can change the weights of disconnected modes by reweighting the samples without transporting them, while FP PDE would have to transport samples from one mode to another (via the vector field or noise) to adjust the relative weights. We will emphasize this point in the final version of the paper.

We would like to thank the reviewer for their time and feedback. We hope our answers here allow the reviewer to continue to positively endorse our paper, and we would happily clarify any addition questions which may arise.