Paper: Particle Guidance: non-I.I.D. Diverse Sampling with Diffusion Models (ICLR-2024)

Reviewer comment (Evaluation):
When the number of modes far exceeds the number of samples, it's evident that iid sampling and sampling with repulsion are similar, which seems to be the case for most data. For instance, in Figure 4 (a), diversity does not appear to be a severe issue.

Author response:
Thank you for raising this great point, we have added a subsection (B.1) to our appendix to discuss this point and refer to it in our main text. As suggested by the reviewer, when considering distributions with a number of modes, combinatorial analysis shows that IID sampling is particularly inefficient when the number of samples is in the same order as the number of modes. Intuitively, if the number of modes is much larger than the number of samples, then I.I.D. samples are still likely to capture separate modes; while if the number of samples is much larger than that of modes, the number of modes not captured by any sample is likely to be small.

This is indeed the case in the setting of molecular conformer generation, where for each molecule the energy function has a finite set of relevant minima (called conformers) that one wants to find. In the commonly used DRUGS dataset, each molecule has on average ~100 conformers. Molecular conformer generation methods aim to recover all these modes with as few samples as possible (commonly for benchmarks the number of samples extracted from each method is twice the number of conformers). Therefore, molecular conformer generation falls into this setting where I.I.D. sampling is inefficient.

We also discuss and point to other works showing the inefficiency of I.I.D. sampling when dealing with more general continuous distribution e.g. w.r.t. the variance of Monte Carlo estimation or improving diversity through determinantal point processes.

Our experiments on text-to-image generation with stable diffusion also support the thesis that while typically there is no clear mode-collapse, diversity can still be improved. We provide at this URL https://anonymous.4open.science/r/pg-images/ the (non-cherry picked) images generated from the first 50 text prompts of COCO using the same hyperparameters and the same random seed (as recommended by another reviewer), with a fixed guidance scale $w=8$. Obviously when comparing visually the differences are subjective, but, at our eyes, many of the generations with PG show clear improvements in sample diversity. For example, for the first prompt “A baby eating a cake with a tie around his neck with balloons in the background.”, conventional I.I.D. sampling tends to generate a brown-haired white child, while particle guidance generates diverse depictions of babies with varying hair and skin colors.

Reviewer comment (Experiments):
The authors should emphasize scenarios where the number of samples is close to the number of modes to highlight the setting's effectiveness. Molecular conformer might be an example, but since the reviewer is not an expert in that area, the authors could further explain.

Author response:
Thank you for raising this great point, we have added a subsection (B.1) to our appendix to discuss this point and refer to it in our main text. As suggested by the reviewer, when considering distributions with a number of modes, combinatorial analysis shows that IID sampling is particularly inefficient when the number of samples is in the same order as the number of modes. Intuitively, if the number of modes is much larger than the number of samples, then I.I.D. samples are still likely to capture separate modes; while if the number of samples is much larger than that of modes, the number of modes not captured by any sample is likely to be small.

This is indeed the case in the setting of molecular conformer generation, where for each molecule the energy function has a finite set of relevant minima (called conformers) that one wants to find. In the commonly used DRUGS dataset, each molecule has on average ~100 conformers. Molecular conformer generation methods aim to recover all these modes with as few samples as possible (commonly for benchmarks the number of samples extracted from each method is twice the number of conformers). Therefore, molecular conformer generation falls into this setting where I.I.D. sampling is inefficient.

We also discuss and point to other works showing the inefficiency of I.I.D. sampling when dealing with more general continuous distribution e.g. w.r.t. the variance of Monte Carlo estimation or improving diversity through determinantal point processes.

Our experiments on text-to-image generation with stable diffusion also support the thesis that while typically there is no clear mode-collapse, diversity can still be improved. We provide at this URL https://anonymous.4open.science/r/pg-images/ the (non-cherry picked) images generated from the first 50 text prompts of COCO using the same hyperparameters and the same random seed (as recommended by another reviewer), with a fixed guidance scale $w=8$. Obviously when comparing visually the differences are subjective, but, at our eyes, many of the generations with PG show clear improvements in sample diversity. For example, for the first prompt “A baby eating a cake with a tie around his neck with balloons in the background.”, conventional I.I.D. sampling tends to generate a brown-haired white child, while particle guidance generates diverse depictions of babies with varying hair and skin colors.

Reviewer comment (Theory):
The paper's theoretical foundation is weak and does not directly explain why particle guidance can increase diversity.

Author response:
Assuming that $\Phi$ is some measure of diversity then one can see particle guidance as locally improving the diversity of the samples at every step. However, for the general (fixed potential) setting it is hard to argue about global improvement or optimality as the effect on the final probability distribution depends on the data distribution itself (see Theorem 1). On the other hand, as we added to Section 7, using the maximum entropy theorem we can show that the learned potential particle guidance setup is globally optimal in the sense that it learns to sample the distribution minimizing the KL divergence with the independent distribution while satisfying a bound on the expected value of a diversity metric.

Reviewer comment (Experiments):
Has the paper investigated the impact of the number of particles on the final generation outcome?

Author response:
Yes in our experiments we found that in general when one is interested in generating a certain number of samples in the end it is best (in terms of diversity) to use that as batch size or divide it in batches as large as possible. In the stable diffusion experiments, we use a batch size of 4, as this is the number of samples commonly reported. For the torsional diffusion experiments, where some molecules have a much higher number of conformers to be generated, we use larger batch sizes up to 128. We see monotonic improvements in performance when scaling the maximum batch size from 32 to 64 to 128 to 256. However, having larger batch sizes has two sources of runtime overhead, (1) the computation of the potential that, if parameterized as a sum of pairwise kernels scales quadratically with the batch size (although this is typically very fast compared to score model computation), and (2) the requirement of sampling all the elements of the batch in parallel that, for large batch sizes, requires moving elements in and out of GPU. These were the reasons we chose to clip the batch size to 128 for the torsional diffusion experiments: if the number of conformers to generate n is smaller or equal to 128 we use that, otherwise, we cover n with batches of size 128. We have added these discussions in the appendix of the manuscript.

Reviewer comment (Theory):
Is there theoretical proof that particle guidance is better to some extent (even though this aligns with intuition)?

Author response:
Assuming that $\Phi$ is some measure of diversity then one can see particle guidance as locally improving the diversity of the samples at every step. However, for the general (fixed potential) setting it is hard to argue about global improvement or optimality as the effect on the final probability distribution depends on the data distribution itself (see Theorem 1). On the other hand, as we added to Section 7, using the maximum entropy theorem we can show that the learned potential particle guidance setup is globally optimal in the sense that it learns to sample the distribution minimizing the KL divergence with the independent distribution while satisfying a bound on the expected value of a diversity metric.

Reviewer comment (Experiments):
The paper uses an ODE solver. Would there be any changes if an SDE solver is used?

Author response:
No, in fact for the synthetic and torsional diffusion experiments, we use an SDE solver. We clarified this point in the manuscript.

Reviewer comment (Novelty):
While the angle of this paper (to promote efficiency and diversity) is novel, I found the methodology itself, namely equation 1, does not have too much _technical_ novelty, and seems to be following the line of "guidance" works, e.g. classifier guidance, reconstruction guidance [1], DPS[2], Pseudo-inverse diffusion sampler [3], universal guidance [4], etc.

Author response:
**Q: "I found the methodology itself, namely equation 1, does not have too much technical novelty, and seems to be following the line of "guidance" works"**

Although we chose to present our method from the perspective of guidance to make it as easy as possible to understand for researchers in the field, we believe that our methodology is clearly distinct from such previous works. In particular, we believe our work presents novel ideas w.r.t. this line of guidance works on several fundamental aspects: (1) the goal and motivation, (2) the potential being applied on a set of particles instead of individual ones (3) the form of the potential analyzed (either fixed or learned) (4) the theoretical analysis of effects and (5) the connections to methods in other disciplines.

Reviewer comment (Reproducibility):
One potential weakness is the memory constraint invoked by simulating multiple particles (and getting their gradients) at the same time. Can the authors comment on this point?

Author response:
**Q: "One potential weakness is the memory constraint invoked by simulating multiple particles (and getting their gradients) at the same time. Can the authors comment on this point? Another missing discussion is on the choice of the number of particles used in practice."**

Thank you for raising this point, we have now added this discussion in Appendix B.2 of the paper. In terms of memory, particle guidance does not create any significant overhead since the kernel computations can be aggregated per element. However, when running inference on GPU if n is larger than the batch size that fits the GPU memory when running score model inference, the data might have to be moved back and forth between RAM and GPU memory to enable synchronous steps for particle guidance, causing runtime overhead. 

In the case of our experiments on Stable Diffusion, the number of samples extracted (4) does not create significant overhead. However, in the setting of conformer generation on DRUGS different molecules can have very different numbers of conformers with some even having thousands of them. For efficiency, we therefore cap the size of n in particle guidance to 128 and perform batches of 128 samples until the total number of conformers is satisfied.

Reviewer comment (Theory):
Another missing discussion is on the choice of $n$, the number of particles used in practice. How does one determine $N$, and what is its implication on the joint distribution? For example, I hypothesize that for $n_1>n_2$, marginalizing out the last $n1-n2$ particles in the joint distribution of $n1$ does not recover the joint distribution of $n_2$. And when $n$ is much larger than the number of modes in the diffusion model that we are interested in, what would happen by running the proposed method (especially when using a large guidance weight). Would that result in generating particles that have very low density under the diffusion model?

Author response:
**Q: "...Would that result in generating particles that have very low density under the diffusion model?" … I think there is a missing discussion on whether such joint sampling of would recover marginally exact samples from a diffusion model"**

Thank you for raising this point. Indeed the marginal distribution of the sampled points under the particle guidance is different from the original data distribution. We fixed the comment on page 3 and now discuss this question of preservation of marginal distribution in the manuscript (see for example Section 7.2).

We note, that this (lack of preservation of the marginals) is the case also in all the other non-IID sampling methods discussed in related works (Section 4) and is, arguably, positive depending on the definition of optimality (see new components in Section 7). However, we do recognize that there are several domains where preserving marginal distribution is desirable (e.g. if using samples for some estimates). Therefore, in the setting of learned potential PG, we derive a training procedure (Section 7.2) for the joint potential that preserves marginal distributions (still in a maximum entropy optimal way).

Reviewer comment (Writing):
This point is regarding clarity. Similar to the above "marginalization" argument, I think there is a missing discussion on whether such joint sampling of $x_1, \cdots, x_n$ would recover marginally exact samples from a diffusion model, and I assume not (correct me if I'm wrong). But somehow I found the paper can be misleading in suggesting "yes", e.g. on page 3 "Intuitively, this will push our different samples to be dissimilar from one another while at the same time matching our distribution, improving sample efficiency."

Author response:
**Q: "...Would that result in generating particles that have very low density under the diffusion model?" … I think there is a missing discussion on whether such joint sampling of would recover marginally exact samples from a diffusion model"**

Thank you for raising this point. Indeed the marginal distribution of the sampled points under the particle guidance is different from the original data distribution. We fixed the comment on page 3 and now discuss this question of preservation of marginal distribution in the manuscript (see for example Section 7.2).

We note, that this (lack of preservation of the marginals) is the case also in all the other non-IID sampling methods discussed in related works (Section 4) and is, arguably, positive depending on the definition of optimality (see new components in Section 7). However, we do recognize that there are several domains where preserving marginal distribution is desirable (e.g. if using samples for some estimates). Therefore, in the setting of learned potential PG, we derive a training procedure (Section 7.2) for the joint potential that preserves marginal distributions (still in a maximum entropy optimal way).

Reviewer comment (Experiments):
While Theorem 1 is interesting, the expression derived for the joint density is not very interpretable and so one does not get a good grasp of what the modified backward process is targeting. It is quite unsatisfying that the authors did not add a toy experiment where they explicitly compare the law of the samples obtained from particle guidance with the initial law that is targeted. I believe that the authors should consider an experiment in which the modes do not have the same weights and then show what is actually the law that they are sampling from. I would expect this law to not have the correct statistical weights, which could be quite inconvenient.

Author response:
**Q: the expression derived for the joint density is not very interpretable and so one does not get a good grasp of what the modified backward process is targeting.**

Indeed the derived expression in Theorem 1 is not very interpretable, in Appendix A.1 we tried to add some intuition for the final expression, however, we recognize that the joint potential for arbitrary fixed potential particle guidance is hard to interpret. On the other hand, we added further intuition to the joint potential learned when applying the learned potential particle guidance (see the general response) as the distribution minimizing the KL divergence with the independent distribution while satisfying a bound on the expected value of a (diversity) metric.

**Q: I would expect this law to not have the correct statistical weights, which could be quite unconvenient.**

Thank you for raising this point. Indeed the marginal distribution of the sampled points under the particle guidance is different from the original data distribution. As now discussed in the paper (section 7.2 and appendix A.3), this is the case also in all the other non-IID sampling methods discussed in related works (Section 4) and is, arguably, positive depending on the definition of optimality (see new components in Section 7). However, we do recognize that there are several domains where preserving marginal distribution is desirable (e.g. if using samples for some estimates). Therefore, in the setting of learned potential PG, we derive a training procedure (Section 7.2) for the joint potential that preserves marginal distributions (still in a maximum entropy optimal way).

Reviewer comment (Experiments):
Failure cases of the proposed method are not discussed. Furthermore, as far as I can tell there is no real discussion on the choice of the potential and its parameters (besides in the toy example). How did the authors choose the potential in the other examples? Can a badly chosen potential perform worse than the original diffusion model?

Author response:
**Q: Failure cases of the proposed method are not discussed. Furthermore, as far as I can tell there is no real discussion on the choice of the potential and its parameters (besides in the toy example). How did the authors choose the potential in the other examples? Can a badly chosen potential perform worse than the original diffusion model?**

Thank you for the suggestion. We tried to clarify these points in the manuscript and added discussions in Appendix B on the failure cases of particle guidance and when this can provide benefits over I.I.D. sampling. We discuss and benchmark the different choices of potential both in Sections 5.1 (synthetic experiments) and 5.2 (text-to-image experiments). Indeed a badly chosen potential or in general one with a guidance weight too high can overly change the marginal likelihood and negatively impact the sampling quality (see Figure 5). However, in fixed potential particle guidance, its parameters can be easily fit at inference time, therefore, it is typically relatively inexpensive to test the optimal value of the guidance weight for the application of interest and the chosen potential. This leads to the prevention of too high guidance weights and the simple detection of bad potential when the optimal weight is close to 0.

Reviewer comment (Experiments):
Why did the authors choose such a small variance for the Gaussian mixture experiment?

Author response:
**Q: Why did the authors choose such a small variance for the Gaussian mixture experiment?**

We chose a small variance in the Gaussian mixture experiment to have separate modes in the distribution that could make the mode recovery analysis possible. While one can also analyze diversity and coverage in different ways for "more homogeneous" distributions, we believe this perspective makes it easier to understand. Moreover, it is closely related to the setup of conformer generation where molecules are characterized by a finite number of somewhat distinct possible conformations (conformers).