Particle Guidance: non-I.I.D. Diverse Sampling with Diffusion Models
Authors: Gabriele Corso
Yilun Xu
Valentin De Bortoli
Regina Barzilay
Tommi S. Jaakkola
Venue: ICLR 2024 poster (2024)
Keywords: diffusion models
score-based models
diversity
guidance
conformer generation
image generation
Primary area: generative models
TL;DR: Enforcing diversity with non-IID sampling of diffusion models via guidance from a joint-particle time-evolving potential

Abstract: In light of the widespread success of generative models, a significant amount of research has gone into speeding up their sampling time. However, generative models are often sampled multiple times to obtain a diverse set incurring a cost that is orthogonal to sampling time. We tackle the question of how to improve diversity and sample efficiency by moving beyond the common assumption of independent samples. We propose particle guidance, an extension of diffusion-based generative sampling where a joint-particle time-evolving potential enforces diversity. We analyze theoretically the joint distribution that particle guidance generates, how to learn a potential that achieves optimal diversity, and the connections with methods in other disciplines. Empirically, we test the framework both in the setting of conditional image generation, where we are able to increase diversity without affecting quality, and molecular conformer generation, where we reduce the state-of-the-art median error by 13% on average.

Published as a conference paper at ICLR 2024
PARTICLE GUIDANCE: NON-I.I.D. DIVERSE SAMPLING WITH DIFFUSION MODELS
Gabriele Corso"l, Yilun Xu', Valentin de Bortoli", Regina Barzilay', Tommi Jaakkola 'CSAIL, Massachusetts Institute of Technology. PENS, PSL University
ABSTRACT
In light of the widespread success of generative models, a significant amount of research has gone into speeding up their sampling time. However, generative mod- els are often sampled multiple times to obtain a diverse set incurring cost that is orthogonal to sampling time. We tackle the question of how to improve diversity and sample efficiency by moving beyond the common assumption of independent samples. We propose particle guidance, an extension of diffusion-basee tive sampling where point-particle time-evolving potential enforces diversity. We analyze theoretically the joint distribution that particle guidance generates, how to learn potential that achieves optimal diversity, and the connections with meth- ods in other disciplines. Empirically, we test the framework both in the setting of conditional image generation, where we are able to increase diversity with- out affecting quality, and molecular conformer generation, where we reduce the state-of fofte-of-theatte ---thhe-rtttt art median error by 13% on average.
genera-
INTRODUCTION
Deep generative modeling has become pervasive in many computational tasks across computer vi- sion, natural language processing, physical sciences, and beyond. In many applications, these mod- els are used to take a number of representative samples of some distribution interest like Van Gogh's 's style paintings or the 3D conformers of small molecule. Although independent samples drawn from a distribution will perfectly represent in the limit of infinite samples, for finite num- ber, this may not be the optimal strategy. Therefore, while deep learning methods have so far largely focused on the task of taking independent identically distributed (I.ID.) samples from some dis- tribution, this paper examines how one can use deep generative models to take a finite number of samples that can better represent the distribution of interest.
In other fields where finite-samples approximations are critical, researchers have developed various techniques to tackle this challenge. In molecular simulations, several enhanced sampling methods, like metadynamies and replica exchange, have been proposed sample diverse sets of low-energy structures and estimate free energies. In statistics, Stein Variational Gradient Descent (SVGD) an iterative technique match distribution with finite set of particles. However, these methods not able to efficiently sample complex distributions like images.
Towards the goal of better simite-samples generative models, that combine the power of recent ad- vances with sample efficiency, we propose general framework for sampling sets of particles using diffusion models. This framework, which we call particle guidance (PG). is based on the use of a time-evolving potential guide the inference process. We present different strategies instan- tiate this new framework: the first,fxed potential particle guidance, provides ready-to-use potentials that require further training and have little inference overhead; the second, learned potential par- ticle guidance, requires a training process but offers better control and theoretical guarantees.
The theoretical analysis of the framework leads to two key results. On one hand, we obtain an expression for the joint marginal distribution of the sampled process when using any arbitrary guidance potential. On the other, we derive simple objective one can use to train a model to learn time-evolving potential that exactly samples from joint distribution of interest, We show this provides optimal joint distribution given some diversity constraint and it can be adapted to the addition of further constraints such as the preservation of marginal distributions. Further, also
"Correspondence gcorso tt. e.edu

---

Published as a conference paper at ICLR 2024
I.I.D. sampling
Particle Guidance sampling
g² Vx, log
|-f+ g² V log (x)) dt dw
dx;
dt g dw
+ gg² Vxy log (xr ... Xn)
38
33
38
Score Particle Guidance
Lower
Higher
diversity
diversity
Figure 1: Comparison of I.I.. and particle guidance sampling. The center figure represents each step, with the distribution in pink and the samples as yellow crosses, where particle guidance uses not only the score (in blue) but also the guidance from joint-potential (red), leading it to discover different modes (right-hand samples those the left). At the bottom, Van Gogh cafe images samples generated with Stable Diffusion with and without particle guidance. A more detailed dis- cussion the suboptimality of L.LD. sampling is presented in Appendix B.1.
demonstrate the relations of particle guidance to techniques for non-I.I.D sampling developed in other fields and natural processes and discuss its advantages.
Empirically, we demonstrate the effectiveness of the method in both synthetic experiments and two of the most successful applications of diffusion models: text-to-image generation and molecular conformer generation. In the former, we show that particle guidance can improve the diversity of the samples generated with Stable Diffusion [Rombach et al., 2021] while maintaining quality comparable that of I.I.D. sampling. For molecular conformer generation, applied to the state-of- the-art method Torsional Diffusion [Jing et al., 2022], particle guidance is able to simultancously improve precision and coverage, reducing their median error by respectively 19% and 8%. In all settings, we also study the critical effect that different potentials can have on the diversity and sample quality.
2 BACKGROUND
Diffusion models Let p(z) be the data distribution we are interested in learning. Diffusion models [Song al., 2021] define forward diffusion process that has p as the initial distribution and is described by dx = f(x, t)dt + ((w),w, where w is the Wiener process. This forward diffusion process is then reversed using the corresponding reverse diffusion SDE dx [-f(x, T + g(T - V) log - dtt + gT - t t)dw (using the forward time convention) where the evolving score V. (x) is approximated with a learned function (x,
One key advantage of diffusion models over the broad class of energy-based models [Teh et al., 2003] is their finite-time time sampling property for taking a single sample. Intuitively. by using of smoothed-out probability distributions diffusion models are able overcome energy barri- ers and sample every mode in finite time as guaranteed by the existence of the reverse diffusion SDE [Anderson, 1982]. In general, for the same order of discretization error, reverse diffusion SDE can efficiently sample from data distribution in much fewer steps than Langevin dynamics in energy-based models. For instance, Theorem I of Chen et al. [2022] shows that, assuming accurate learning of score, the convergence of diffusion SDE is independent of the isoperimetry constant of the target distribution. Langevin dynamics mixing speed can be exponentially slow if the spectral gap/isoperimetry constant small. This critical property is orthogonal the efficiency in the num- ber of samples one needs to generate to cover distribution; in this work, we aim to achieve sample efficiency while preserving the finite-lime sampling of diffusion models.
2

---

Published as a conference paper at ICLR 2024
Diffusion models were extended to Riemannian manifolds by De Bortoli et al. [2022], this formula- tion has found particular success [Jing et al., 2022; Corso et al., 2022; Yim al., 2023] scientific domains where data distributions often lie close to predefined submanifolds [Corso, 2023].
Classifier guidance (CG) [Dhariwal & Nichol, 2021] has been another technical development that has enabled the success of diffusion models on conditional image generation. Here a classifier po po(y) xx,; f), trained to predict the probability of x, being obtained from a sample of class y, is used to guide a conditional generation of class y following:
dx [- (') + g(t)) (sx(x, + V log (y) t' d + g(t)ddw where T
where in theory should be 1, but, due overspreading of the distribution, researchers often set it to larger values. This, however, often causes the collapse of the generation to a single or few modes, hurting the samples' diversity.
3 PARTICLE GUIDANCE
Our goal is to define a sampling process that promotes the diversity of finite number of samples while retaining the advantages and flexibility that characterize diffusion models. Let p(x) be some probability distribution of interest and Vx log Pa (x) be the score that we have learned reverse the diffusion process dx = f(x, t)dt + g(t) Similarly to how classifier guidance is applied, we modify the reverse diffusion process by adding the gradient of potential. However, we are now sampling together a whole set of particles XI, Xn, and the potential log 4, is not only a function of the current point but permutation invariant function of the whole set:
dx; 1 -(x,, (t())) Vx, Po (x,) + Vx, log By
(x1 (X1,..., x.)) dt g(t")dw.
where the points are initially sampled I.ID. from a prior distribution pr. We call this idea particle guidance (PG). This framework allows one impose different properties, such diversity, the set of particles being sampled without the need to retrain a new score model operating directly on the space of sets.
We will present and study two different instantiations of this framework:
1. Fixed Potential PG where the time-evolving joint potential is handcrafted, leading to very efficient sampling of diverse sets without the need for any additional training. We present this instantiation in Section 5 and show its effectiveness on critical real-world applications of diffusion models in Section 6.
2. Learned Potential PG where we learn the time-evolving joint potential to provably opti- mal joint distributions. Further, this enables direct control of important properties such as the preservation of marginal distributions. We present this instantiation in Section 7.
CONNECTIONS WITH EXISTING METHODS
As discussed in the introduction, other fields have developed methods to improve the tradeoff be- tween sampling cost and coverage the distribution of interest. this section, we will briefly introduce four methods (coupled replicas, metadynamics, SVGD and electrostatics) and draw con- nections with particle guidance.
4.1 COUPLED REPLICAS AND METADYNAMICS
In many domains linked to biochemistry and material science, researchers study the properties of the physical systems by collecting several samples from their Boltzmann distributions using molecular dynamics or other enhanced sampling methods. Motivated by the significant cost that sampling each individual structure requires, researchers have developed a range of techniques improve sample efficiency and speed by, for example, reducing the correlation of subsequent samples from slow- mixing Markov chains. The most popular of these techniques are parallel sampling with coupled replicas and sequential sampling with metadynamics.
As the name suggests, replica methods involve directly taking n samples of a system with the dif- ferent sampling processes, replicas, occurring in parallel. In particular, coupled replica methods
3

---

Published as a conference paper at ICLR 2024
[Hummer & Köfinger, 2015; adding, like particle guidance, an extra potential & to the energy function enforce diversity or match experimental observables. This results in energy-based sampling procedures that target:
Pasarkar et al., 2023] create dependency between the replicas by
p(x X.) (x1,
II p(x.)
Metadynamics [Laio & Parrinello, 2002: Barducci et al., 2008] was also developed to more ficiently sample the Boltzmann distribution of a given system. Unlike replica methods and our approach, metadynamics is a sequential sampling technique where new samples are taken based on previously taken ones to ensure diversity, typically across certain collective variables of interest s(x). In its original formulation, the Hamiltonian the htth sample augmented with potential
||s(x) s1
1 = H Σ exp jck
where H is the original Hamiltonian, xi are the previously sampled elements and w and a param- eters set priori. Once we take the gradient and perform Langevin dynamics to sample, we obtain dynamics that, with the exception of the fixed Hamiltonian, resemble those of particle guidance in Eq. where
(
||0(x;) - ,x23)11
Vx, log (x(, ;X,) f V
N.W
Σ exp J&
20
Although they differ in their parallel sequential approach, both coupled replicas and metadynam- ies can broadly classified as energy-based generative models. As seen here, energy- based models offer simple way of controlling the joint distribution one converges to by simply adding poten- tial to the energy function. On the other hand, however, the methods typically employ an MCMC sampling procedure, which lacks the critical finite-time sampling property of diffusion models and significantly struggles to cover complex probability distributions such as those of larger molecules and biomolecular complexes. Additionally, the MCMC typically necessitates a substantial number of steps, generally proportional
polynomial of the data dimension [Chewi et al., 2020]. With particle guidance, we instead aim to achieve both properties (controllable diversity and finite time sampling) at the same time. We can simulate the associated SDE/ODE with total number of steps that is independent of the data dimension.
4.2 SVGD
Stein Variational Gradient Descent (SVGD) [Liu & Wang, 2016] is well-established method in the variational inference community iteratively transport set of particles to match a target distribu- tion. Given set of initial particles (xf...x%)) J. it updates them at every iteration as:
-- x (x)((x))
where
((x)
Σ [x(x) (x(xxx)) x) x) log K(x,(
where i is some (similarity) kernel and the step size. Although SVGD was developed with the intent of sampling set of particles that approximate some distribution p without the direct goal of obtaining diverse samples, SVGD and our method have close relation.
This relation between our method and SVGD can be best illustrated under specific choices for drift and potential under which the probability flow ODE discretization of particle guidance can be ap- proximated as (derivation in Appendix A.5):
(+tt Re x (x(x)))xxxx) where (xx)
Σ (x, x) Vx log + (xf..(
(3)
Comparing this with Eq. 2. we can a clear relation in the form of the two methods, with some key distinctions. Apart from the different constants, the two methods use different terms the total score component. Interestingly both methods use smoothed-out scores, however, on the one hand, particle guidance uses the diffused score the specific particle x,, Vx, log (x()), while on the other, SVGD smoothes it out by taking a weighted average of the score of nearby particles weighted by the similarity kernel (E, k(x,, x,) Vx, log logp(x,))/(2 k A(x,, x,))
4

---

Published as a conference paper at ICLR 2024
The reliance of SVGD other particles for the "smoothing of the score". however, causes two related problems, firstly, it does not have the finite-time sampling guarantee that the time evolution of diffusion models provides and, secondly, it suffers from the collapse to few local modes near the initialization and cannot discover isolated modes in data distribution [Wenliang & Kanagawa, 2020]. This challenge has been theoretically [Zhuo et al., 2018] and empirically [Zhang et al., 2020] studied with several works proposing practical solutions. In particular, relevant works use an annealing schedule to enhance exploration [D' Angelo & Fortuin, 2021] or use score matching kernel for SVGD [Chang et al., 2020]. Additionally, we empirically observe that the score smoothing in SVGD results in blurry samples in image generation.
to obtain
4.3 ELECTROSTATICS
Recent works [Xu et al., 2022; 2023b] have shown promise in devising novel generative models inspired by the evolution of point charges in high-dimensional electrii fields defined by the data dis- tribution. It becomes natural therefore to ask whether particle guidance could be seen as describing the evolution of point charges when these only attracted by the data distribution but also repel one another. One can show that this evolution can indeed be seen as the combination of Poisson Flow Generative Models with particle guidance, where the similarity kernel is the extension of Green's 's function in N+1 1-dimensional space, i.e., We defer more details to Appendix A.6.
put in the same electric field such that they not
k(I, y) x 1/|l= y||l-1
5 FIXED POTENTIAL PARTICLE GUIDANCE
In this section, will present and study simple, yet effective, instantiation of particle guidance based on the definition of the time-evolving potential as a combination of predefined kernels. As will see in the experiments in Section 6, this leads to significant sample efficiency improvements with no additional training required and little inference overhead.
To promote diversity and sample efficiency. in our experiments, we choose the potential log t to be the negative of the sum of a pairwise similarity kernel between each pair of particles log 2r (x(X),,..,,, = = 24 Eis j Kio Xx), obtaining:
dx; - f(x + g (vx, log (x.) or Vx. Σ ke (x))) dt g(t")d 3=1
(4)
Intuitively, the kernel term will push our diffferent samples to be dissimilar from one another while at the same time the score term will try to match our distribution. Critically, this does not come significant additional runtime as, in most domains, the cost of running the pairwise similarity kernels is very small compared to the execution it allows the use of domain-specific similarity kernels and does not require training any additional classifier or score model. We can also view the particle guidance Eq. (4) as a sum of reverse-time SDE and guidance term. Thus, to attain more expedited generation speed, the reverse-ting SDE can also be substituted with the probability flow ODE [Song et al., 2021].
the large score network architecture. Moreover,
5.1 THEORETICAL ANALYSIS
To understand the effect that particle guidance has beyond simple intuition, we study the joint dis- tribution of sets of particles generated by the proposed reverse diffusion. However, unlike methods related energy-based models (see coupled replicas, metadynamics, SVGD in Sec. 4) analyzing the effect of the addition of a time-evolving potential log & in the reverse diffusion is non-trivial.
While the score component in particle guidance is the score of the sequence of probability distri- butions (x,,,......,, because, for arbitrary time-evolving potential L,, this sequence of marginals does not correspond to diffusion process. One strategy used by other works in similar situations [Du et al., 2023] relies on taking, after every step at the end, a number of Langevin steps reequilibrate and move the distribution back towards pr. This, however, increases significantly the runtime cost (every Langevin step requires score evaluation) and is technically correct only in the limit of infinite steps leaving uncertainty in the real likelihood of our samples. Instead, in Theorem we use the Feynman- Kac theorem to derive a formula for the exact reweighting that particle guidance has on a distribution (derivation in Appendix A)
x,, I[m=1 Pr. (x:), we are not necessarily sampling exactly po
5

---

Published as a conference paper at ICLR 2024
Theorem 1. Under integrability assumptions, sampling xi ...,, 2. from pr and following the par- ticle guidance reverse diffusion prucess, we obtain samples from the following joint probability distribution at time 1. 0:
po(x1..... - expl- for ((t) log 2 (X.)), log + A log (d.]
with Z (explicit in the appendix) such that
I N Po(x;) = E[2],
(X)))e(0,r] is a stochastic process driven by the equation
dx. = (J(X,, g(1)2V lo (X.))d g())w,
Xo - (xi
Hence the density po can be understood as a reweighting of the random variable that represents LLD. sampling.
Riemannian Manifolds. Note that our theoretical insights can also extended to the manifold framework. This is direct consequence of the fact that the Feynman-Kac Kac theorem can be extended to the manifold setting, see for instance Benton et al. [2022].
5.2 PRESERVING INVARIANCES
The objects that we learn to sample from with generative models often present invariances such as the permutation of the atoms in a molecule or the roto- ooto-translation of conformer. To simplify the learning process and ensure these are respected, it is common practice to build such invariances in the model architecture. In the case of diffusion models, obtain distribution that is invariant the action of some group such as that of rotations or permutations, it suffices to have an invariant prior and build score model that is G-equivariant [Köhler et al., 2020; Xu et al., 202 l]. Similarly, our case, we are interested distributions that are invariant to the action of on any of the set elements (see Section 6.2), we show that a sufficient condition for this invariance to be maintained that the time-evolving potential 2 itself invariant to G-transformations of any of its inputs (see Proposition in Appendix A.4).
6 EXPERIMENTS
Fixed potential particle guidance can be implemented on top of any existing trained diffusion model with the only requirement of specifying the potential/kernel to be used in the domain. We present three sets of empirical results in three very diverse domains. First, in Appendix C, we work with a synthetic experiment formed by a two-dimensional Gaussian mixture model, where we can vi- sually highlight some properties of the method. In this section instead, consider text-to-image and molecular conformer generation, two important tasks where diffusion models have established new ofte-of-thhear performances, and show how, in each of these tasks, particle guidance can pro- vide improvements in sample efficiency pushing the -uuaritty-quality Pareto frontier. The code guidance.
available at https://gii thu....tthuu...g. cot
6.1 TEXT-TO-IMAGE GENERATION
In practice, the most prevalent text-to-image diffusion models, such as Stable Diffusion [Rombach et al., 2021] or Midjourney, generally constrain the output budget to four images per given prompt. Ideally, this set of four images should yield a diverse batch of samples for user selection. However, the currently predominant method of classifier-free guidance [Ho, 2022] tends to push the mini-batch samples towards typical mode enhance fidelity, the expense of diversity.
To mitigate this, we apply the proposed particle guidance to text-to-image generation. We use Stable Diffusion v1.5, pre-trained on LAION-5B [Schuhmann et al., 2022] with resolution of 512 x 512, as our testbed. We apply an Buler solver with 30 steps to solve for the ODE version of particle guidance. Following [Xu et al., 2023a], we use the validation set in COCO 2014 [Lin et al., 2014] for evaluation, and the CLIP [Hessel al., 202 1/Aesthtiiiiicc score [Team, 2022] (higher is better) to assess the text-image alignment/visual quality, respectively. To evaluate the diversity within each batch of generated images corresponding to a given prompt, introduce the in-batch similarity

---

Published as a conference paper at ICLR 2024
score. This metric represents the average pairwise cosine similarity of features within an image batch, utilizing the pre-trained DINO [Caron et al., 2021] as the feature extractor. Contrasting the FID score, the in-batch similarity score specifically measures the diversity of batch of images generated for a given prompt. We use a classifier-free guidance scale from the trade-off curve between the diversity and CLIP/Aesthetic score. in line with prior works [Xu implement the RBF kernel the down-sampled pixel space (the latent space of the VAE encoder-) in Stable Diffusion, as well as the feature space of DINO. Please refer to Appendix E.1 for more experimental details.
10 visualize
al., 2023a; Saharia et al., 2022]. For particle guidance,
In-batch similarit Score
0.64
L...
In-batch similariiy score
0.64
L.LD
0.63
(pixel)
0.63
(pixel)
PG (fealure)
PG (feature)
0.62
0.62
0.61
0.61
0.60
0.60
0.59
0.59
0.318
0.319
0.320
0.321
0.322
5.34
5.35
5.36
5.37
5.38
5.39
5.40
5.41
CLIP score (VIT-g/14)
Aesthetic soore
(a) In-batch similarity versus CLIP score
(b) In-batch similarity versus Aesthetic score
Figure 2: In-batch similiarity score versus (a) CLIP ViT-g/14 score and (b) Aesthetic score for text- to-image generation 512 x 512 resolution, using Stable Diffusion v1.5 with varying guidance scale from 6 to 10.
(a) L.D.
(b) PG
(c) Training data
(d) LLD.
(e) PG
Figure 3: Text prompt: (a.b) "A baby eating cake with tie around his neck with balloons in the
background" (COCO); (c.d.e) "VAN GOGH CAFE TERASSE copy.jpg", with original training data in (c).
As shown in Fig. 2(a) and Fig. 2(b), particle guidance (PG) consistently obtains a better (lower) in-batch similarity score in most cases, given the same CLIP/Aesthetic score, with classifier- freee guidance scale ranging from
10. Conversely, we observe that while in-batch similarity score of I.D.. sampling improves with the reduced classifier-fi free guidance scale, particle guidance continues surpass I.I.D. sampling in terms of CLIP/Aesthetic score given the same in-batch similarity. When the potential is the similarity kernel applied in the feature space, particle guidance notably attains compared to I.LD. sampling or to the approach in the original
lower in-batch similarity downsampled pixel space. This suggests that utilizing a semantically meaningful feature space is more appropriate for determining distances between images.
In Fig. 3, we further visualize generated batches four images per prompt by L.I.. sampling and particle guidance (feature) with the same random seeds, when fixing the classifier-free guid- ance scale to 9. We can see that particle guidance improves the visual diversity the generated batch. Interestingly, particle guidance can also help to alleviate the memorization issue of Stable Diffusion [Somepalli et al., 2023]. For example, given the text prompt of a painting from LAION dataset, particle guidance (Fig. 3(c)) avoids the multiple replications of the training data in the LLD.. setting (the top-left and the bottom-rige images in Fig. 3(d)). We provide extended samples in Ap- pendix F. and additionally show that SVGD (Eq. 2) fails to promote diversity, instead yielding set of blurry images.
6.2 MOLECULAR CONFORMER GENERATION
Molecular conformer generation is key task in computational chemistry that consists of finding the set of different conformations that a molecule most likely takes in 3D space. Critically it is often important to find all or most of the low-energy conformers each can determine different behavior
7

---

Published as a conference paper at ICLR 2024
Table 1: Quality of generated conformer ensembles for the GEOM-DRUGS test set in terms of Coverage (%) and Average Minimum RMSD (A). We follow the experimental setup from [Ganea
et al., 2021], for experimental details and introduction of the baselines please refer to Appendix D.
Recall
Precision
Coverage
AMR +
Coverage f
AMR L
Method
Mean
Med
Mean
Med
Mean
Med
Mean
Med
RDKit ETKDG
38.4
28.6
1.058
1.002
40.9
30.8
0.995
0.895
OMEGA
53,4
54.6
0.841
0.762
40.5
33.3
0.946
0.854
GeoMol
44.6
41.4
0.875
0.834
43.0
36.4
0.928
0.841
GeoDiff
42.1
37.8
0.835
0.809
24.9
14.5
1.136
1.090
Torsional Diffusion
72.7
80.0
0.582
0.565
55.2
56.9
0.778
0.729
TD w/ particle guidance
77.0
82.6
0.543
0.520
68.9
78.1
0.656
0.594
(e.g. by binding to a protein). This necessity is reflected in the metrics used by the community that look both coverage (also called recall) and precision over the predictions.
Over the past few years, molecular conformer generation has been extensively studied by the ma- chine learning community, with well-established benchmarks [Axeirod & and several generative models designed specifically for this task [Ganea et al., 2021; Xu al., 2021; Jing et al., 2022]. However, all these methods are based on training generative model to generate single samples and then running this model several times (more than 200 average in the standard GEOM-DRUGS dataset) to generate a large number of I.I.D. samples.
2022]
As discussed before, however, this strategy is suboptimal to generate representative sets of samples and cover the distribution. Therefore, we take state-of-the-art conformer generation model, torsional diffusion, and, without retraining the model itself, we show that we can obtain significant improvements in both coverage and precision via particle guidance.
Torsional diffusiin [Jing et al., 2022] defines the diffusion process over the manifold defined by changes in torsion angles from some initial conformer because of the relative rigidity of the remain- ing degrees of freedom. Given this observation, we also define the guidance kernel on this manifold as RBF kernel over the dihedral angle differences.
Another important consideration when dealing with molecular conformers is given by the permuta- tion symmetries that characterize several molecules: conformers that appear different might be very similar under permutations of the order of the atoms that do not change the bond structure. To max- imize the sample efficiency and avoid generating similar conformers, we make the kernel invariant to these transformations. For this, we employ the simple strategy to take the minimum value of the original kernel under the different perturbations (formalized in Appendix D).
Table shows that by applying particle guidance to SDE-based reverse process of torsional dif- fusion (see Appendix D for details) we are able balance coverage and precision being able to obtain, without retraining the model, significantly improved results on both metrics with 8% and 19% simultancous reductions respectively in recall and precision median AMR.
7 LEARNED POTENTIAL PARTICLE GUIDANCE
While the fixed potential particle guidance seen so far is very effective in improving the diversity of samples with little overhead, it hard argue about the optimality of the resulting joint distribution. This is because of the complexity of the expression obtained in Theorem and its dependence on the data distribution itself. Furthermore, in some domains, particularly in scientific applications, researchers need to control the distribution that they are sampling. This is necessary, for example, to apply correct importance weights compute free energy differences. While Theorem allows us to theoretically analyze properties of the distribution, the joint and marginal distributions remain largely intractable.
In this section, we analyze how we can sample from desired joint probability distribution by leam- ing a tailored time-evolving potential for particle guidance. Using the maximum entropy theorem [Csiszár, 1975]. show that the distribution satisfying bound the expected value of
8

---

Published as a conference paper at ICLR 2024
(diversity) metric du while minimizing the KL divergence with the independent distribution is:
po (Xi, xc Lo x., ((a)
p(x,)
(5)
where 3 is a function of a, the value of the bound on Ep Po].
7.1 TRAINING PROCEDURE
We now have to learn time-evolving potential Φ that when used part of the particle guidance framework generates Po (we assume Po chosen such that 3(a) 1). To achieve this, we mandate that the generation process of particle guidance in Eq. adheres to the sequence of marginals (x ..., x (x .... x IT; Pr(x{) and learn b satisfy this evolution. Under mild
assumptions, using Doob h--theory (derivation in Appendix A.2), we show that we can learn the Φ by the following objective:
0* arg min B P ( Po .., xn)
where Pejo is the Gaussian perturbation kernel in diffusion models. Importantly, here the initial x are sampled independently from the data distribution so this training scheme can be easily executed in parallel to learning the score of Pt.
7.2 PRESERVING MARGINAL DISTRIBUTIONS
While the technique discussed in the previous section is optimal in the maximum entropy perspec- tive, it does not (for arbitrary 4)) preserve the marginal distributions of individual particles, i.e. marginalizing X; over does not recover p. Although not critical in many settings and not re- spected, for a finite number of particles, neither by the related methods in Section nor by the fixed potential PG, this is an important property in some applications.
Using again the maximum entropy theorem, we can show that the distribution satisfying a bound on the expected value of (diversity) metric 4' and preserving the marginal distribution while minimizing the KL divergence with the independent distribution can be written as:
Po(x1.... x x (x(,, x, ((()) (xx))xx)
1=1
for some scalar function over individual particles 79. In Appendix A.3, we derive a new training scheme to learn the parameters of 79. This relies on setting the normalization constant an arbitrary positive value and learning values of # that respect the marginals. Once ?9 is learned, its parameters can be frozen and the training procedure of Eq. 6 can be started.
CONCLUSION
In this paper, we have analyzed how one can improve the sample efficiency of generative models by moving beyond I.L.. sampling and enforcing diversity, critical challenge in many real appli- cations that has been largely unexplored. Our proposed framework, particle guidance, steers the sampling process of diffusion models toward more diverse sets of samples via the definition of time-evolving joint potential. We have studied the theoretical properties of the framework such as the joint distribution it converges to for an arbitrary potential and how to learn potential functions that sample some given joint distribution achieving optimality and, if needed, preserving marginal distributions. Finally, we evaluated its performance two important applications of diffusion mod- els text-to-image generation and molecular conformer generation, and showed how in both cases it able to push the Pareto frontier of sample diversity vs quality.
We hope that particle guidance can become valuable tool for practitioners to ensure diversity and fair representation in existing tools even beyond the general definition of diversity directly tackling known biases of generative models. Further, we hope that our methodological and theoretical con- tributions can spark interest in the research community for better soint-particl sampling methods.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Xiang Cheng, Bowen Jing. Timur Garipov, Shangyuan Tong, Renato Berlinghieri, Saro Passaro, Simon Olsson, and Hannes Stärk for their invaluable help with discussions and review of the manuscript. We also thank the anonymous reviewers for their useful feedback and suggestions. Concurrently with our work, Alex Pondaven also developed a similar idea'.
This work was supported by the NSF Expeditions grant (award 1918839: Collaborative Research: Understanding the World Through Code), the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the Abdul Latif Jameel Clinic for Machine Learning in Health, the D'TRA Discovery of Medical Countermeasures Against New and Emerging (DOMAÑE) threats pro- gram, the DARPA Accelerated Molecular Discovery program, the NSF Al Institute CCF-2112663 the NSF Award 21 34795, the GIST-MIT Research Collaboration grant. MIT-DSTA Singapore col- laboration, and MIT-IBM Watson AI Lab.
REFERENCES
Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Ap-
plications, 1982.
Simon Axelrod and Rafael Gomez-Bombarelli. Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 2022.
Alessandro Barducci, Giovanni Bussi, and Michele Parrinello. Well-tempered metadynamics: smoothly converging and tunable free-energy method. Physical review letters, 2008
Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising markov models. arXiv preprint 1 1.00599, 2022.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEFE/CVF international conference on computer vision, 2021.
Wei-Cheng Chang, Chun-Liang Li, Youssef Mroueh, and Yiming Yang. Kernel stein generative modeling, arXiv preprint arXiv:2007.03077 2020.
Sitan Chen, Sinho Chewi, Jungshian Li, Yuanzhi Li, Adil Salim, and Anru R. Zhang. Sam- pling is as easy as learning the score: theory for diffusion models with minimal data as- sumptions. ArXiv, abs/2209. 11215, 2022. URL https: /ttp://///a...maat.....// semant www.hhcaccccc.cccooooooo.o......
CorpuSID:252438990
Sinho Chewi, Chen Lu, Kwangjun Ahn, Xiang Cheng, Thibaut Le Gouic, and Philippe Rigol- let. Optimal dimension dependence of the metropolis-adjusted langevin algorithm. ArXiv, abs/2012.122100, 2020.
Gabriele Corso. Modeling molecular structures with intrinsic diffusion models. arXiv preprint arXiv:2302.122555 2023.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Dif- fusion steps, twists, and turns for molecular docking. arXiv preprint arXlv:2210.0177,, 2022.
Imre Csiszár. I-divergence geometry of probability distributions and minimization problems. The annals of probability, pp. 146-158,11 1975.
Francesco D' Angelo and Vincent Fortuin. Annealed stein variational gradient descent, arXiv 2021.
preprint
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet. Riemannian score-based generative modeling. urXiv preprint arXiv:2202.02763 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 2021.
10

---

Published as a conference paper at ICLR 2024
Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and meme. In International Con- ference on Machine Learning. PMLR, 2023.
Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer en- sembles. Advances in Neural Information Processing Systems, 2021.
John Geweke.
Antithetic acceleration of monte carlo integration in bayesian inference. Journal of Econometrics, 38(1 (11-2))73--8, 1988.
Paul CD Hawkins and Anthony Nicholls. Conformer generation with omega: learning from the data set and the analysis of failures. Journal of chemical information and modeling, 2012.
Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Joseph Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning. In Conference on Empirical Methods in Natural Language Processing. 2021.
Jonathan Ho. Classifier-fre free diffusion guidance. ArXiv, abs/2207.1259. 2022.
Gerhard Hummer and Jürgen Köfinger. Bayesian ensemble refinement by replica simulations and reweighting. The Journal of chemical physics, 2015.
John Ingraham, Max Baranov, Zak Costello, Vincent Frappier, Ahmed Ismail, Shan Tie, Wujie Wang, Vincent Xue, Fritz Obermeyer, Andrew Beam, et al. Illuminating protein space with programmable generative model. BioRxiv, 2022.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola, Torsional diffu- sion for molecular conformer generation. arXiv preprint arXiv:2206.0177 2022.
Ioannis Karatzas and Steven Shreve. Brownian motion and stochastic calculus. Springer Science & Business Media, 1991.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances Neural Information Processing Systems, 2022.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. International conference on machine learning, pp. 5361-5370. PMLR, 2020.
Alex Kulesza, Ben Taskar, et al. Determinantal point processes for machine learning. Foundations and Trends Machine Learning. 5(2-3)::23--86, 2012.
Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the national academy of sciences, 2002.
Greg Landrum et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling, 2013.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference on Computer Vision, 2014.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
Amey Pasarkar, Gianluca Bencomo, Simon Olsson, Adji Bousso Dieng. Vendi sampling for molecular simulations: Diversity force for faster convergence and better exploration. The Journal of Chemical Physics, 159( 14). 2023.
Philipp Pracht, Fabian Bohle, and Stefan Grimme. Automated exploration of the low-energy chem- ical space with fast quantum chemical methods. Physical Chemistry Chemical Physics, 2020.
Robin Rombach, Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. 2022 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
11

---

Published as a conference paper at ICLR 2024
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Bureu Karagol Ayan, Seyedeh Sara Mahdavi, Raphael Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. ArXiv, abs/2205. 14887, 2022.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarezyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. ArXiv, abs/2210.084022, 2022.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular confor- mation generation. In International Conference on Machine Learning, 2021.
Richard Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matrices. Annals of Mathematical Statistics, 35:876-879, 1964. URL https://api.
semant cscho 1 ar. .ar..rrr.ao.....e.............e. wtrg/CorpusTDy 1208467144
Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Diffusion art or digital forgery? investigating data replication in diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- fional Conference on Learning Representations, 2021.
LAION-AI Team. christ http:pppcchhhhaaaomoooooreeeeeee
Laion-aesthetics predictor v2.
https: //g/ thub. 2022.
Yee Whye Teh, Max Welling, Simon Osindero, and Geoffrey E Hinton. Energy-basee models for sparse overcomplete representations. Journal of Machine Learning Research, 2003.
Oleg Trott and Arthur J Olson. Autodock vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, 2010.
Li Kevin Wenliang and Heishiro Kanagawa. Blindness of score-based methods to isolated compo- nents and mixing proportions. 2020.
Minkai Xu, Lantao Yu, Yang Song. Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: geo- metric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2021.
Yilun Xu, Ziming Liu, Max Tegmark, and Tommi Jaakkola. Poisson flow generative models. Ad- vances in Neural Information Processing Systems, 2022.
Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, and Tommi Jaakkola. Restart sampling for improving generative processes. arXiv preprint arXiv:2306. 14878, 2023a.
Yilun Xu, Ziming Liu, Yonglong Tian, Shangyuan Tong, Max Tegmark, and T. Jaakkola. Pfgm++: Unlocking the potential of physics- generative models. In International Conference on Machine Learning, 2023b.
Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Amaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se(3) diffusion model with application to protein backbone generation. arXiv preprint, 2023.
Jianyi Zhang, Ruiyi Zhang, Lawrence Carin, and Changyou Chen. Stochastic particle-opimizzaion sampling and the non-asymptotic convergence theory. In International Conference on Artificial Intelligence and Statistics. PMLR, 2020.
Jingwei Zhuo, Chang Liu, Jiaxin Shi. Jun Zhu, Ning Chen, and Bo Zhang. Message passing stein variational gradient descent. In International Conference Machine Learning. PMLR, 2018.
12

---

Published as a conference paper at ICLR 2024
A DERIVATIONS
A.1 JOINT DISTRIBUTION UNDER PARTICLE GUIDANCE
In this section, we provide the proof of Theorem 1. First, we restate the Feynman-Kat Kac theorem. Let u: [O,T] x Rdd such that for any E [0, 7) and x E RRd we have
0,u(t, x) + (b(t, x). Vu((,=)) + (1/)) x),, V - V(t, + f(t, = 0, (8) with u(T, x) = 2). Then, under integrability and regularity assumptions, see Karatzas & Shreve [1991] for instance, we have
u(0 2) = ELL expl- So VT,, (rr X, )rr + expl- So (r,, (1)
X z].
(9)
with (T.. क(T, 2)) and dx, (,)d X,)d,, In the rest of this section, we derive the specific case of Theorem 1.
We recall that the generative model with particle guidance is given by (Pe)xE(o.T) and is associated with the generative model
dÝ, = (-f
t) 9T (so -- t) + V log -r- (Y))( + 9 - t)dw. (10)
We also recall that the generative model without particle guidance is given by (2))7)) and is associated with the generative model
dY, {-f(Y,, - t) + 9(T t)) (Y,,T t) g(T t)dw.
(11)
Using the Mokkerr-lanck equation associated with equation 11 we have for any x E (124)N
0.q. (z) + div(1 -, + 9(T - so
(2) - (o(()))))((((--00000)) 0. (12)
This can also be rewritten as
(x))(x))
Vu(()) + -
+ -,
(13)
(14)
Denoting u, Qr we have
(x)) + (f)
1), (z. (0)))))) div(l- g(t)ssa( (z) = 0.
(15)
(16)
Note that since u, Qr-t. we have that u, Pt with the conventions from 3. Now combining this result with equation 8 and equation with V zt) = div([-((-::)) g(t))so(tt (z) and f = we have that
ua(x)= = =E[]]
(17)
with and
7 =expp= 00~ V(T. X(rrXx)drrloolf
(18)
dx, = - -(0) ()dd + g(O)dw
(19)
with Xo x. We now consider a similar analysis in the case of the generative with particle guid- ance. Using the Fokker-Planck equation associated with equation 10 we have for any E (Rd)
0. (x)+ ((,,vv(((((,T()))/))( (T,,-)+ (sa(( T-))),,(++VVVVVV log -
0.
(00)
This can also be rewritten as
(x))(()) +-) (( (006
(21)
(22)
+ g - (0ongg/--.ttggggggglla///ag--/g0
O.
(23)
Denoting 0; -r we have
(x) +
t). ( -div((-f(;,)) t) + (t)))))(tt)))))))ddd
(24)
(25)
-((t) (x)) V (z) (og
(x))))(())) 0.
(26)
13

---

Published as a conference paper at ICLR 2024
Following the convetion of 3, we have that än Po. Now combining this result with equation 8 and equation 9 with V(t, div(f- + s(1) so(; (x) +
g(t)²((V log (z)), Vlog (logpr-((2)) + A log 2x(z)) and f = 0 we have that
ún(x) [[]),
(27)
with
2 = = - 00¹ V(T, X, (drlpx)))
(28)
and
dx, (S(X,,) 900ggg((())))))))))))))) 80 (X,, ()dl + g(()dw.
(29)
again with X, We conclude the proof upon noting that
2 = Z expl- for s()) log $ (X.), V log in (X.)) + A log t; (X.))d)
(30)
Interpretation We introduce the Stein operator in Liu & Wang [2016] given by for any v : (RT) + (R N
An interpretation of this reweighting term can be given through the lens of SVGD.
by
Ap, (V.) = V W. + V.,
(31)
Using V.- = V loo Lr, we get that
(r = (V log P., logg. + A log D..
(32)
The squared expectation of this quantity w.r.t. distribution on (I is the Kernel Stein Discrep- ancy (KSD) between q and Pr given the kernel log Dr.
A.2 SAMPLING PREDEFINED JOINT DISTRIBUTION
For ease of derivation via the Doob h-transform, we temporarily reverse the time from f to T t. Here, pr treated as the data distribution, and o is regarded as the potential, as specified by users. We now consider another model. Namely, we are looking for a generative model P. with E [0. T] such that for any means that we are considering a factorized model for all times / with the additional requirement that at the final time T, the model is given by pr Py with Φ known. This is to be compared with Theorem 1. Indeed in Theorem while the update on the generative dynamics is explicit (particle guidance term). the update on the density is not. In what follows, we are going to see, using tools from Doob hetrrassfrmm theory, that we can obtain an expression for the update of the drift in the generative process when considering models of the form P:: = P.
7' we have P. p. 4 with Φ given by the user. In layman's 's terms, this
More precisely, we consider the following model. Let PT Pr and for any 8, E [0,T) with < and x!!= (x1)²_ (Rdy"" and x!! {xx] (Rdy" we define
Prjs [x (x! Paja (x]""x be (x. 1:m) (xx!.")))
(33)
with 4, which satisfies for any x! (IRd)
0,
t)² (x ",
0,
(34)
with ọ given. Note that equation 34 expresses that o satisfies the backward Kolmogorov equation. Under mild assumptions, using Doob h-theory, we get that there exists (X:)tejorry such that for any t e [0, I] we have Y, Pr and for any E [0,7]
dÝ, (-fT-1(Y, (Y) + g(T - [) log Pe Y.) + Vlog (V.))ld + gT( - f)dw.
(35)
The main difficulty is to compute 4, for any E [0, T]. Under mild assumptions, solutions to the backward Kolmogorov equation 34 are for any t E [0, [ by
(x{x (m) = ((2r((Y))))) Yr, x = Φr (Yr = x "xr"))rr(((()) (x} 1: (xx "x
(36)
where we have
dY, = (- T( 1) log (Y:) g(T - -)dd.
(37)
14

---

Published as a conference paper at ICLR 2024
This means that (Y)) is given by the original generative model, with time-dependent marginals Pt- The expression equation 36, suggests to parameterize D, by 26 and consider the loss function
((() = Ev, ExyyErerere (T) 4. (Y))
(38)
Then, we can define global loss function (( = 0fd (0)) where A is some weight. One problem with this original loss function is that it requires sampling and integrating with respect to Y, which requires sampling from the generative model.
Recall that we reverse the time from to T' - at the beginning. Reverse back to the original convention in the main text, Eq. (39) can be expressed as
& (0) = Ex xxx (o)o) (2( (1.)
(39)
A.3 PRESERVING MARGINAL DISTRIBUTIONS
For arbitrary time evolving potentials (x1, Xn) sampling using particle guidance does not pre- serve the marginals p(x:) * JxI,...x: 1. p(x1, .... xn,
In many domains this is not required and none of the methods discussed in Section 4 have this property for finite number of particles, however, in domains where, for example, one wants to obtain unbiased estimates of some function this property may be useful.
While the technique discussed in Section 7 allows us to use any potential .x....... choosing to in such a way that preserves marginals is hard for non-trivial potentials and distributions. Therefore, we propose learn a non-trivial marginal preserving &6 from the data in the following way. Let (x x,) (x1, Xn) II. (xx)), where 4% is some predefined joint potential that, for example, encourages diversity the joint distribution and 7o a learned scalar function operating on individual points that counterbalances the effect that 46 has on marginals while maintaining its effect on sample diversity.
How do we learn such scaling function "9 to preserve marginals? By definition, the individual marginal distribution is (consider = w.l.o.g.):
Sa..... xn (6(1,, Xx) II, p (7x)
dxn.
(x1)
Ix (6('', xx,
II. pl (0) (x)) dxx, ...dx,
p(x1) ) Exg M1an mx.) 1 (/(())x11 (x1, x,,) II (x(x))
Exq 11, p(x()
1 (X1, X₂,) 11. 70
Ex Bec...xn.~ll. p(x.)
P(x.)
[x) (x.,..... xn)
= -(x1)
Ex Ex;,,.,,-III, mx))
x(,]]
Therefore pi = =p if and only if the fraction is always equal to 1 (intuitively for the marginal to be maintained on average the potential should have no effect). Assuming that the joint po- tential &6 is invariant to the permutation to its inputs. Then one can also prove that, P1 = p if and only if xx. p(x.) xx,,) is equal to a positive constant C for any æ1- We can then minimize the following regression loss learn the scalar function yo, by matching E. Ex Mes p(x.) (x)......,, xn) and C::
min Ex,
m(x.)) (%6 xn,) II [0(x,)] iff1
7 (x1)
However, achieving this objective necessitates costly Monte Carlo estimations for the expectation over n - independent samples. Moreover, obtaining an unbiased estimator for the full-batch gra- dient in mini-batch setup is challenging. To bypass this issue, we implement a greedy update rule that optimizes the value (x(x1) on the single sample x1. The greedy update can be viewed as continuous extension of the Iterative Proportional Fitting (IPF) for symmetry matrices. The essence of the IPF algorithm lies determining scaling factors for each row and column of given ma- trix, ensuring that the sum of each row and column (the marginals) aligns with a specified target value. IPF is known for its uniqueness and convergency guarantees [Sinkhomn, 1964]. Let's denote
15

---

Published as a conference paper at ICLR 2024
6 (x,,,,..,,,, st stop.grad((((11) .... xn) f (x(x))), the objective for the greedy update rule
is as follows:
min (x()))
Exi (Ex xn Meer p(xe) 01 (x1,..., x.)),.),),,,)_.......))))))).)..
(x())))
The gradient w.r.t in the objective above is:
2C
( (01(()) (x,,,...,,,)) 70(x1) )
= 8-88xx X2 II p(x)))
Voro(z1)
(40)
(x1)
where B is the learning rate. The corresponding update by stochastic gradient is:
0' - 3-
2C
(x((x)..., x,)
( Vora(1;;)
C
)
(41)
(x:)
(x(x;)
i=1
20
=0 3
(o) (x,(x),..,,x))) Xn) -n)-c)) Voru(z;)
(42)
(x,)
for X1, X1...., Xn ~ P(x)) The stochastic gradient is an unbiased estimator of the gradient in Eq. (40).
A.3.1 EMPIRICAL SYNTHETIC EXPERIMENTS
We demonstrate this training paradigm in a synthetic experiment using mixture of Gaussian dis- tributions in 2D. In particular, p to be mixture of 7 Gaussians with the same variance with placed as shown Figure 4.A. The middle Gaussian has a weight that is four times that of the others.
A.
Original distribution
8. Marginal of (orig nal potential)
1.0
0.5
0.0
-0.5
-0.-
-1.-
-1.0
-0.5
1.0
-1.0
-0.5
0.0
1.0
Marginal ((9) (learned potential)
D.
Learned values of y
140
0.5
100
0.0
60
40
-1.0
-1.0
-0.5
0.0
0.5
1.0
Figure 4: Synthetic experiment on learning a potential that preserves marginal distributions. The description of each plot can be found in the text.
Figure 4.B shows the marginal distribution when sampling particles with joint distribution:
(x(x)...., xn) x (xi, X., I p(x))
where b6 is a measure of diversity that in this case we take to be the exponential of the negative of the mean of pairwise Fuclidean RBF similarity kernels. Table 2 highlights how sampling from
16

---

Published as a conference paper at ICLR 2024
this modified joint distribution significantly increases the diversity of the samples, especially when comes to the expected value of log & itself. However, as clear from Figure 4.B, the marginal distribution is significantly changed with the middle mode being sampled only 13% of the times instead of 40% and even the shape of the outer modes being altered.
To fix this we learn a function "Yo, here parameterized simply as the set of values on a grid repre- senting the domain. We follow the training scheme presented in Eq. (41) obtaining the function presented in Figure 4.D. As evident from the plot, this has the effect overweighing samples in the central mode, while downsampling samples in the outside modes, especially when coming from the outer parts. In Figure 4.C we plot the marginal of the resulting distribution:
( xn) x P/(x x,.) II ((x;)7((x))
Although the marginals closely match, po still has more diversity in its sets of samples compared to I.I.D. sampling of p. however, as seen in Table 2 the level of diversity in po is lower than that po. highlighting the "cost" of imposing the preservation of marginals.
Table 2: Values of different observables under different joint probability distributions. For every method we take 5000 samples (of 10 particles), samples from po and po were obtained reweighting 50000 samples of the independent I.I.D. distribution.
Observable
L.D.
Po
Number of modes recovered at every sample Expected value of log 16
4.9
5.9
5.3
-37.3
-31.8
-36.3
A.4 INVARIANCE OF PARTICLE GUIDANCE
Proposition 1. Let G be the group of rotations permutations of a of vectors. Assuming that Pr (x) is a G-invariant distribution, the learned score s(x,t) t) and f(x, t) are G-equivariant and the potential log (x)) is G-invariant to a transformation of any of its inputs, then the resulting distribution we sample from will also be G-invariant to a transformation of any of the elements of the set.
Note that in this section we will derive this specific formulation for the group of rotations or permu- tations and the Brownian motion Euclidean space. For more general statement on Lie groups G and Brownian motions associated with given metric, one could generalize the result from Yim al. [2023] Proposition F.2.
Proof. For simplicity, we will consider Fuler discretization steps going with time from T to 0 (as used in experiments), however, the proposition applies in the continuous setting too:
po (x 1) (1)
x = Px(x (t-1)
(1) xx + f(x!" t) g" (sx (1) t) + log 2 (1)
x
(x(.))))
where ((0, I). Without loss of generality since the whole method is invariant to permutations of the particles, consider x,, to be the particle which we apply Tg the transformation of an arbitrary group element g.
Since by assumption log (x log t, [t-11 ( (t)
we have pa (1x/011/))))))))) ((()))
On the other hand, since log o (x is invariant to G transformations of x its gradient w.r.t the same variable will be G-equivariant. Therefore:
pu (x (x) (x,
(T - + V x log 2, (x ((()))
=Px(T,(x{}11) Tx(xxx))))))))))) + ) (fxx()))))
- -TT (v V ( log (xx(1)
- (xe(Txxx (x--1)
x! + f(x!9,t) t) (x00x(xx) V
log 2 (x (t)
= po(x (t- (x
17

---

Published as a conference paper at ICLR 2024
where between lines 2 and 3 we have used the equivariance assumptions and in the latter two the properties elements of G.
Putting these together, we follow similar derivation of Proposition from Xu et al. [2021]]
pe 1:n-11 T,
(0)
T
=/ plxxmm T: (x, ())
(3)
II
po (t-1) mm (x
1| TT(
= /
p(xx
X | Txx (x
iin
(1
11
Tx
pa
(II p(x))) (IP(T))
(x 11x))))))
((x
II po (t (t-1))))))))))))) x(xx))) = po (x()1)
KKn t=iin
A.5 PARTICLE GUIDANCE AS SVGD
In this section, we derive the approximation of Eq. 3 starting from the probability flow ODE equiv- alent of Eq. 4 under the assumptions of no drift f(x L) = and using the following form for (x1 In) = (2e,,
where k, is similarity kernel based on the Euclidean
distance (e.g. RBF kernel).
1 dr; ((an)) -
-
- (t)(t) vv, logpi(z;)) Ve, log Σ ke(Ez, _(((t)a)))>>>) dt
1
V a; Eis ke (2((, Eij ke W.))
9 (t)dt V (x;)
Now we can simplify the numerator using the fact that KG is symmetric and approximate the denom- inator assuming that different particles will have similar average distances other particles:
VN, Ezke(znrz)) E, ke)
- (1)de Vz; log))(a.)
(1)dd
Σ (e,, logrt() Vz, Va, ke 2 S(x))
where S(x;) = E, Re (N; Iz)). Now we can use the fact that V ke = - -Va, Ke
because the kernel only depends on the Euclidean distance between the two points:
(t)dt
Σ ke (e 2,) (eg Vay (e
2 S(x:)
n2
Letting cx(x;)
(t)(t we obtain Eq. 3: ((a))))
At N o + (((x:)))(()) where (xx) n1
(x), logp((z) Vzzhe(a).2)]
18

---

Published as a conference paper at ICLR 2024
A.6 PARTICLE GUIDANCE POISSON FLOW GENERATIVE MODELS
In this section, we consider the more general Poisson Flow Generative Models++ [Xu et al., 2023b] framework in which the N-dimensional data distribution is embedded into N D-dimensional space, where D a positive integer (D 1/D - O recover PFGM [Xu et al., 2022) models). The data distribution is interpreted as positive charge distribution. Each particle indepen- dently follows the electric field generated by the N-dimensional data distribution p(z)) embedded in NN + D--dimensional space. One can similarly do particle guidance in the PFGM++ scenarios, treating the group of particles negative charges, not only attracted by the data distribution but also exerting the mutually repulsive force. Formally, for the augmented data the ODE in PFGM+ (Eq.4 in Xu et al. [2023b]) is
dx (), dr Elr,rr
where Eg and Er are the electric fields for different coordinates:
y r(r,r)= SN+D-1(1) / (| ul12 p²) 2 x p(y)
r F(x,r), Sx+D- / (1|-- y/1 r2) "
Ni (yyddy
Note that when
DVD,D oo, the ODE is de F(z.r) dr E(as).
VD Va Po and the
framework degenerates diffusion models.
Now we consider the repulsive forces among a set of (uniformly weighted) particles with the same anchor variables {2,,, r> only the electric field in the coordinate changes (the component in the coordinate is zero). Denote the new electric field component as Eg:
B(x,,T)== =
E(zi,r), r)
I L
SN+ (1)
jfi (1æ 3y 11 X1
attractive force by data
repulsive force between particles
The corresponding new ODE for the i-th particle is
dr; M(x, dr E(r,, r)r
SN+D-1(1) n E Cif
Zy
E(ZiTT) E(x,,r),
(1lary as.l3 (12)
E(s;,
n-1 Ejtt (1la vi p(() dy
r) E( r)r
(Ia- 2
prodicted by pro-trained models
particle guidance
E(z;,T) Fl r)r
Σ jti F.
2 - 2
SV+D-1 1 PD PD-1Sp (r.))
where p, is the intermidate distribution, and S,, the surface area of n--pheee. Clearly, the direction of the guidance term can be regarded as the sum of the gradient of N + D-dimensional Green's function G(x, v) x 1/1læ 2 RDD-2 up some scaling factors:
V G(),) =
33 Ez
|r.
x; N+D
19

---

Published as a conference paper at ICLR 2024
A.7 COMBINATORIAL ANALYSIS OF SYNTHETIC EXPERIMENTS
Proposition 2. Let us have a random variable taking value equiprobably between N distinct bins. The expectation of the proportion of bins discovered (i.e. sampled at least once) after N samples is which tends to 1 - -1/e as N tends to infinity.
1
N N
Proof. Let be the number of samples in the 1s bin. The proportion of discovered bins is equal to:
P(n: > 0) P(n; = 0) =1 N N
1 N
Σ m2
E Tm,
=
N
the limit of - 00:
N -N-1 N
lim
-
= =1 - y
N
N
since (using L' L'Hôpital's rule):
N -
log(
N-1
1/N2
logy = Jim N log
lim N
lim
N
N
1/N
N-000
-1/N2
Therefore for N 10 we would expect 10 * (1 - 0.9 = 6.51, which corresponds to what is observed in the empirical results of Section
Proposition 3. (Coupon collector's problem) Let us have random variable taking a value equiprobably between N distinct bins. The expectation of the number of samples required to dis- cover all the bins is N HN, where HN is the N harmonic number, which is O log N) as N tends to infinity.
Proof. Len Luj be the number of samples it takes to go from j to bins discovered. We are therefore interested in [[[].
B[Lbbb-1] N (j #1 j
N
[[L
- N
Therefore:
N
E[ E
Σ E [[1 N Σ
N
1
N
HN
(j 1
j=1 j
Since H is (log N), then E[IN o) is O(N log N).
For N = 10, 11ool ~ 29.29.
B DISCUSSIONS
B. SUBOPTIMALITY OF LLD.. SAMPLING
Combinatorial analysis Let us consider again the setting of a random variable taking a value equiprobably between N distinct bins. In Fig. 5, we plot the expected number of modes (or bins) captured as a function of the number steps as derived in Appendix A.7. This suggests that the region where I.I... sampling is considerably suboptimal in these regards (capturing the modes of the distribution) when the number of samples is comparable with the number of modes: the number modes is much larger than I.I.D. samples are still likely to capture separate modes, and the number of samples is much larger the number of uncaptured modes is likely to be small.
20

---

Published as a conference paper at ICLR 2024
Number of modes captured (N=1000)
1000
I.I.D.
Optimal
800
number of modes
600
400
200
1000 number of samples
3500
4000
Figure 5: Plot of the functions N(1 ( N "") and min(æ, N) for N = 1000 represent- ing, respectively, the expected number of modes captured by L.D., sampling a distribution with N equiprobable modes and the optimal coverage.
Literature There is vast literature that has studied the suboptimality of L.I.. sampling from a statistical perspective and proposed different solutions. For example, in the field of Bayesian inference, antithetic sampling [Geweke, 1988] has been proposed way to reduce the variance of Monte Carlo estimates. Determinantal Point Processes [Kulesza et al., 2012] have also been widely studied as technique to improve the diversity of samples.
B.2 RUNTIME AND MEMORY OVERHEAD
The runtime overhead due to the addition particle guidance to the inference procedure largely de- pends the potential that is used and on the size of the In particular, while computing kernels tends to be significantly cheaper than running the score model, the number of kernel computations scales quadratically with while the number of score model executions at every step is n.
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 further overhead.
In the case of our experiments on Stable Diffusion, the number of samples extracted 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 in particle guidance to 128 and perform batches of 128 samples until the total number of conformers is satisfied.
Other limitations A badly chosen potential in general one with a guidance weight too high can overly change the marginal likelihood and negatively impact the sampling qual- ity. As an example Fig. 6 the use of a particle guidance parameter four times larger than the best one caused various aliasing artifacts on the image.
However, in fixed potential particle guidance, its parameters can be easily fit at inference time, therefore, it typically rela- tively inexpensive to test the optimal value the guidance for the application of interest and the chosen potential. This leads the prevention too high guidance weights and the simple detection bad potential when the optimal weight close to O.
Figure 6: Example of too large PG weight causing aliasing arti- facts.
21

---

Published as a conference paper at ICLR 2024
C SYNTHETIC EXPERIMENTS
To show visually the properties of particle guidance and its effect on sample efficiency, we use two-dimensional Gaussian mixture model. In particular, we consider mixture of N = 10 identical Gaussian distributions whose centers are equally spaced over the unit circle and whose variance is 0.005. These Gaussians form a set of approximately disjoint equal bins. As we are interested in inference, model is trained and the true score of the distribution is given as an oracle.
As expected if one runs normal I.L.. diffusion, the sample falls in one bin at random. Taking ten samples, shown in Fig. some of them will fall in the same bin and some bins will be left unfound. The empirical experiments confirm combinatorial analysis (see Appendix A.7) which shows that the expected number of bins discovered with N = 10 samples is only 6.5 and it takes on average more than 29 samples discover all the bins.
Euclidean kemne
Radial kernel
1e+03
10.0
M.D.
9.5
sampling
1e+02
9.0
Repulsion weight
8.5
7.5
Radial kernel
0.1
7.0
0.01
1.0
0.1
0.01
1.0
0.01
Noise weight
Noise weight
Figure 7: Left: plot of random samples (in blue) of the (wo-dimensional Gaussian mixture distribu- tion (density depicted in red). I.I.D. samples often recover the same modes, while particle guidance with radial kernel captures all modes. Right: average number of modes recovered with 10 samples as function of the weight given by the diffusion noising terms and the potential weight when using an RBF kernel with Euclidean and radial distances respectively. As expected with little weight to the potential terms we obtain approximately modes recovered in line with the I.I.. diffusion performance. Further increasing the potential weight on the Euclidean creates instability.
In many settings this behavior is suboptimal, and would want model discover all the modes of the distribution with as few samples as possible. Using the straightforward application of particle guidance with simple RBF kernel based on the squared Euclidean distance, we are able encourage diversity obtaining, on average, the discovery of nearly 9 bins on average from 10 samples (see Fig. 7).
Intrinsic diffusion models [Corso, 2023] have shown significant improvements when diffusion mod- els operate on the submanifold where the data lies. Similarly, here building into the kernel the degrees of freedom over which the diversity lies helps the particle guidance to effectively distribute samples over the distribution. We know that the different modes are distributed ion, and thus we build an RBF kernel based on the angle difference w.r.t. the origin. Using this power-dimensiona. kernel enables us to consistently discover all modes of the distribution. This submanifold observation aligns well with the practice of methods such as metadynamics where the kernels defined
radial fash-
dower-dimensional collective variables of interest.
D MOLECULAR CONFORMER GENERATION EXPERIMENTS
D.1 DATASET, METRICS AND BASELINES
Dataset We evaluate the method for the task of molecular conformer generation using the data from GEOM [Axelrod & Homez-Bombbaelli. 2022], collection datasets that has become the standard benchmark for this task in the machine learning community. In particular, focus
22

---

Published as a conference paper at ICLR 2024
GEOM-DRUGS, the largest, most pharmaceutically relevant and widely used dataset in GEOM which consists of 304k drug-like molecules. For each of these molecules, an ensemble of conformers was generated with metadynamics in CREST [Pracht et al., 2020], procedure that gives accurate structures but is prohibitive in thg-thhoughhut applications, costing an average of 90 core-hours molecule. To be able to use existing pretrained models we rely on the experimental setup and splits introduced by Ganea et al. [2021] and used by several papers afterward. As do not retrain the score model, we do not use the training set, instead, we finetune the inference parameters for particle guidance and the other ablation experiments on random subset of 200 molecules out of 30433 from the validation set.
Evaluation metries To evaluate conformer generation methods we want to test the ability of method to generate set of conformers that are both individually good poses (precision) and as set cover the distribution of true conformers (recall). For this we employ the same evaluation setup and metrics used by several papers in the field starting from Ganea et al. [2021]. In this setup, methods asked to generate twice as many conformers in the original ensemble and then the so-called Average Minimum RMSD (AMR) and Coverage (COV) are measured for precision (P) and recall (R). For K = 2L let {CT L] and {CL)LE(L be respectively the sets of ground truth and generated conformers:
COV-R = L
[1E [L.L] 3K [L.,K), C7) <
(43)
AMR-R
1
Σ
min RMSD(CR. Ci) kEIL..]
L
IE[1..]]
where is the coverage threshold (set to 0.75A for the GEOM-DRUGS experiments). Swapping ground truth and generated conformers in the equations above we obtain the precision metrics.
Baselines As baselines we report the performances of previous methods as measured by Ganea al. [2021] and Jing et al. [2022]. Cheminformatics conformer prediction methods rely rules and heuristics derived from chemical structures to fix the local degrees of freedom and then use combination of search and template techniques to set the more flexible degrees of freedom like torsion angles. The most accurate and widely used such methods include the open-source software RDKit ETKDG [Landrum et al., 2013] and the commercial tool OMEGA [Trott & Olson, 2010; Hawkins & Nicholls, 2012].
Before the already introduced Torsional Diffusion [Jing et al., 2022], a number of other machine learning approaches were proposed for this task, among these: GeoMol [Ganea et al., 2021] uses a GNN to sample directly from a random seed local structures around each atom and then torsion angles, GeoDiff [Xu et al., 2021] defines equivariant diffusion model over atom coordinates, and CGCF [Shi et al., 2021] learns an energy-based model over the space pairwise distance matrices.
D.2 PARTICLE GUIDANCE SETUP
Reverse diffusion As discussed in Section 6.2, we applied particle guidance to torsional diffusion, as this is currently considered be sfate-of-the- and uses, like most ML-based methods before, I.D.. sampling during inference. We define the particle guidance kernel to operate directly on the implicit hypertorus manifold where torsional diffusion defines the diffusion process, this, at the same time, makes the kernel lower dimensional and it involves a minor modification to the existing inference procedure. The reverse diffusion process that we apply is:
dr, 1 (T 1) 8(Ti, T.T 1) di BT (1) (() L,T 1) dl 9T dw
dillusion ODE
Langevin diffusion SDE
-T- g (T -) V. log ( )dt
2
particle guidance
23

---

Published as a conference paper at ICLR 2024
where follow the idea from Karras et al. [2022] of dividing the different components of the reverse diffusion and tuning their individual parameters. The potential was chosen be:
att Σ (T) 12
where (T(T,,Ty) = exp( ||T, - -1511
log t (TI.
(44)
2n
h4
where the difference of each torsion angle is computed to be (- , T.. Ott, 8r "t and he are inference hyperparameters that are 'logarithmically interpolated" between two end values chosen with hyperparameter tuning (using T = 1), e.g. at = exp(t log(a1) + (1 log()n)).
Permutation invariant kernel Since the kernel operates on the torsion angle differences it is naturally invariant to SE(3) transformations, i.e. translations or rotations, the conformers space. Moreover, as illustrated in Fig. 2 of Jing et al. [2022]. while exact torsion angle values depend on arbitrary choices of neighbors or orientation (to compute the dihedral angle) differences in torsion angles are invariant to these choices. However, one transformation that the kernel in Equation 44 is not invariant to are permutations of the atoms in the molecule. Many of these permutations lead to isomorphic molecular graphs where however each of the torsion angles may now refer to a different make the kernel invariant to these by taking the minimum over the values of the kernel under all such permutations:
dihedral. To maximize the sample efficiency
(T(Ts; TT)
min (T, PTT,) =Ell
where is the set of all permutations that keep the graph isomorphic (but do change the torsion angles assignment) and Ps is the permutation matrix corresponding some permutation и. In practice, these isomorphisms can be precomputed efficiently, however, to limit the overhead from applying the kernel multiple times, whenever there are more than 32 isomorphic graphs leading to change dihedral assignments we subsample these to only keep 32.
Batch size The number of conformers one has to generate is given, for every molecule, by the benchmark (2L) and can vary significantly. To avoid significant computational overheads, we use batches of up to = 128 until all the conformers for that particular molecule are generated.
D.3 FULL RESULTS
We provide in Table 3 again the results reported in Table 1 with the additions of other baselines and ablation experiments. In particular, as ablations, on top of running non-invariam particle guidance ie. without the minimization low-temperature sampling, another variation the inference-time procedure that has been proposed for diffusion models that we applied as described below.
the permutations described in the previous section, we also test
tewmtempertture sampling of some distribution p(x) with tem- perature - < consists of sampling the distribution (x(x) x p(x)^. This helps mitigate the overdispersion problem by concentrating more on high-likelihood modes and trading off sample diversity for quality. Exact low-temperature sampling is intractable for diffusion models, however, various approximation schemes exist. We use an adaptation of Hybrid Langevin- Reverre Time SDE proposed by Ingraham et al. [2022]:
Low-temperature sampling.
dr = --(>;-
A 2
sa.o(C.t) (1)) dt + VI+ g(t) dw with X aa + o²/\
od dd 0t
where À (the inverse temperature), 1/ and az are parameters that can be tuned.
E EXPERIMENTAL DETAILS ON STABLE DIFFUSION
SETUP
In this section, we detail the experimental setup Stable Diffusion. We replace the score func- tion (Vx; log Per (x;)) in the original particle guidance formula (Eq. (4)) with the classifier-free guidance formula [Ho, 2022]:
8 3(x;, e, (') w: v.. log (xx,, c) (1 w) Vx, log (x.:)((x))
24

---

Published as a conference paper at ICLR 2024
Table 3: Quality of generated conformer ensembles for the GEOM-DRUGS test set in terms of Coverage (%) and Average Minimum RMSD (A). Minimizing recall and precision refers to the hyperparameter choices that minimize the respective median AMR on the validation set.
Recall
Precision
Coverage
AMR +
Coverage 1
AMR 1
Method
Mean
Med
Mean
Med
Mean
Med
Mean
Med
RDKit ETKDG
38.4
28.6
1.058
1.002
40.9
30.8
0.995
0.895
OMEGA
53,4
54.6
0.841
0.762
40.5
33.3
0.946
0.854
CGCF
7.6
0.0
1.247
1.225
3.4
0.0
1.837
1,830
GeoMol
44.6
41.4
0.875
0.834
43.0
36.4
0.928
0.841
GeoDiff
42.1
37.8
0.835
0.809
24.9
14.5
1.136
1.090
Torsional Diffusion
72.7
80.0
0.582
0.565
55.2
56.9
0.778
0.729
TD w/ low temperature
minimizing recall
73.3 68.0
77.7
0.570
0.551
66.4
73.8
0.671
0.613
minimizing precision
69.6
0.617
0.604
72.4
81.3
0.607
0.548
TD w/ non-invariant PG minimizing recall minimizing precision
75.8 58.9
81.5
0.542
0.520
66.2
72.4
0.668
0.607
56.8
0.730
0.746
76.8
88.8
0.555
0.488
TD w/ invariant PG minimizing recall
77.0
82.6
0.543
0.520
68.9
78. 1
0.656
0.594
minimizing precision
72.5
75.1
0.575
0.578
72.3
83.9
0.617
0.523
where с symbolizes the text condition, w E R+ is the guidance scale, and Vx, log pe (x, c)/ log (x. is the conditional/uccondttinaa scores, respectively. As prob- ability ODE with classifier-free guidance is prevailing method employed in text-to-image models (Saharia et al., 2022]. we substitute the reverse-time SDE in Eq. (4) with the marginally equivalent ODE, Assuming that f(x,, O, the new backward ODE with particle guidance
dx; [ ( (0(())5(xxx,,, 5(xx,c,t) ar Vx.
ky 2x) dt
Following SVGD [Liu & Wang, 2016], we employ RBF kernel (T(T,, Tj) = exp( IT- 12 with hú m log n, where m, is the median of particle distances. We implement the kernel both in the original down-sampled pixel space (the latent of VAE) or the feature space of DINO- DINO-VIT-b/8 VIT-b/8 [Caron al., 2021]. Defining the DINO feature extractor JDINO, the formulation in the feature space becomes:
dx, = 1 10) (((x),)))) ar VxO ke (xDINo(x)), soIN((())) )
where the input to the DINO feature extractor 9DINO to the xo-prediction: x! x, + ( t'), ) as Xo-prciiction lies in the data manifold rather than noisy images. (t) is the standard deviation of Gaussian perturbatio kernel given time f in diffusion models. The gradient w.r.t x can be calculated by forward-mode auto-diff. We hypothesize that defining Euclidean distance in the feature space is markedly more natural and effective compared the pixel space, allowing the repulsion in a more semantically meaningful way. Our experimental results in Fig, 6.1 corroborate the hypothesis.
To construct the data for evaluation, we randomly sample 500 prompts from the COCO validation set [Lin et al., 2014]. For each prompt, we generate batch of four images. To get the aver- score versus in-batch similarity score curve, for I.I.D. sampling, we use w {6. 7.5, 8.5, 9}. We empirically observed that particle guidance achieved much lower in- batch similarity score (better diversity) than IID sampling. As diversity typically improves with smaller guidance weights [Ho, 2022; Saharia et al., 2022], we chose a set of smaller guidance weights for L.D.D. sampling further improve its diversity, keeping it in the relatively similar
age CLIP score/
25

---

Published as a conference paper at ICLR 2024
range particle guidance. Hence for particle guidance, use set of larger guidance scales: w E (7.5, 8, 9, 9.5, 10}. Indeed, the experimental results suggest that even though I.I.D. sampling used smaller guidance weight to promote diversity, its in-batch similarity score was still worse than that of particle guidance. We set the hyper-parameter ay to (t) in particle guidance (feature) and (t) (t)² in particle guidance (pixel), We use an Euler solver with 30 NFE in all the experiments.
E.2 IN-BATCH SIMILARITY SCORE
We propose in-batch similarity score capture the diversity of a small set of generated samples (xi,..., .x] given a prompt c
( gDINO (x,)
In-batch similarity score(x)...... Xn)) n
INDNO Xx (2kk]l22 |g (DDNO (x1)
To save memory, we use the DINO-VIT-s/8 [Caron et al., 2021] the feature extractor JDINO-
F EXTENDED IMAGE SAMPLES
In Fig. 8-Fig. 12, visualize samples generated by the LLD. sampling process, particle guidance in the pixel space, and particle guidance the DINO feature space, on four different prompts. For Fig. 8-Fig. 10, we select the prompts in Somepalli et al. [2023]. with which Stable Diffusion is shown to replicate content directly from the LAION dataset. We also include the generated samples of SVGD -ggddagndecee in which we replace the particle guidance term with SVGD formula (Eq. (2))- In Fig. 13, we observe that SVGD generally leads blurry images when increasing the guidance scale a₃, This is predictable as the guidance term in SVGD involves a weighted sum of scores of nearby samples, which will steer the samples toward the mean of samples.
Further, we provide at this URL https: / annnymous . 4 oppen. arrienne/r/ppp images/ the (non-cherry picked) images generated from the first 50 text prompts of COCO using the same hyperparameters, the same random seed. We set the guidance scale to commonly used w When comparing visually the differences are subjective, but, to our eyes, many of the generations with PG show clear improvements in sample diversity. For example, for the first prompt "A baby eating cake with a tie around his neck with balloons the background.", conventional iid sampling tends to generate brown-haired white child, while particle guidance generates diverse depictions of babies with varying hair and skin colors.
26

---

Published as a conference paper at ICLR 2024
MARIELL'
MARTALL
(a) LLD.
(b) particle guidance (pixel)
(c) particle guidance (feature)
Figure Text prompt: Captain Marvel Exclusive Cexp Poster Released Online By Marvel
(a) LLD.
(b) particle guidance (pixel)
(c) particle guidance (feature)
Figure 9: Text prompt: Portrait of Tiger in black and white by Lukas Holas
(a) LLD.
(b) particle guidance (pixel)
(c) particle guidance (feature)
Figure 10: Text prompt: VAN GOGH CAFE TERASSE copy.jpg
27

---

Published as a conference paper at ICLR 2024
(a) LLD.
(b) partiele guidance (pixel)
(c) particle guidance (feature)
Figure 11: Text prompt: A transparent sculpture of duck made out of glass
(a) LLD.
(b) particle guidance (pixel)
(c) particle guidance (feature)
Figure 12: Text prompt: unicorn in snowy forest
(a) 0r 0.1
(b) 0+ 0.3
(c) 0 =1
02
Figure 13: SVGD guidance, with varying
28


Reviews:
[{"review_id": "gMJs7QfX0a", "reviewer": "Reviewer_iSmF", "summary": "The aim of the methodology developed in this paper is to promote sample diversity when sampling I.I.D. from a diffusion model. Indeed, when drawing a $M$ i.i.d. samples from a multimodal density with $M$ modes, it is unlikely to get a sample on each one of the modes, even when all the modes have the same weight. To promote sample diversity when sampling from a diffusion model, the authors propose to modify the backward process by adding a repulsion term ensuring that samples drawn at each step of the diffusion process are as dissimilar as possible.", "questions": "Why did the authors choose such a small variance for the Gaussian mixture experiment?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1- The problem that this paper tries to solve is quite original.  \n2- The paper is well written and the idea in itself is interesting. It is also quite nice that authors managed to give an explicit formula for the the joint density targeted by their modified backward process. The connection with other works that make use of a repulsion term is also a nice addition.  \n3- The experiments are well explained and sound.", "weaknesses": "1- 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 unsatisying that the authors did not add a toy experiment where they explicitely 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 unconvenient. \n\n2- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kRs1pFMH6s", "reviewer": "Reviewer_rfgv", "summary": "This paper presents a method to sample from a joint distribution of a diffusion model combined with some potential function, with a specific focus on potentials that encourage diversity. This is done by attaching an additional \"guidance\" term to the diffusion SDE that moves the generation towards high-potential region, which shares the similar idea of prior works on guidance, e.g. classifier guidance. The paper analyzes the theoretical properties of the implied joint distribution, and connections to related works. Empirically, the paper demonstrates the superiority of the proposed method in Gaussian synthetic example, text-to-image generation and molecular conformation generation. In these examples, the proposed method has scored better in diversity and other downstream metrics compared to the standard IID samping.", "questions": "1. In abstract, can you explain which part of the paper discusses \"its implications on the choice of potential\"? \n\n2. The \"finite-sampling\" property of diffusion models does not seem to be accurate. Apparently, sampling from the diffusion SDE (in a discrete manner, and without running infinite Langevin steps) will also accumulate error. It seems like the authors were referring this property to getting to every \"mode\" of the distribution in finite steps. If this is the case, I wonder the authors can be more clear about this point, and provide evidence or reference to support this claim?\n\n3. Can you further explain why \"Hence the density p0ˆ can be understood as a reweighting of the random variable Z\" (the sentence under equation 5 on page 4)? Furthermore, can you provide intuition on what the random variable $Z$ encode?\n\n4. Maybe I missed this part, but for all experiments, what are the $n$, number of particles you used, and how do you determine them?\n\n5. In the preamble part of section 5, the reference to the text-to-image experiment is missing.\n\n6. In Figure 3, how is \"varying guidance scale from 6 to 10\" reflected/used in the figure and this experiment?\n\n7. Table 1 seems to miss the reference and descriptions of competing methods. Furthermore, I wonder why the authors didn't experiment the proposed method on geodiff model. Is that possible? \n\n8. Section 6 seems really interesting! I wonder if there is any practical challenge in instantiating that paradigm?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well-written and easy to follow.\n\n- The motivation of this paper is very clear and seems to be important. \n\n- The proposed method presents simple and effective way to overcome the limitation of IID sampling, which is the inefficiency in exploring all possible modes. \n\n- The empirical comparison is comprehensive and convincing.", "weaknesses": "- 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. \n\n- 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? \n\n- 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?\n\n- This point is regarding clarify. 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.\" \n\n[1] Video Diffusion Models. Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, David J. Fleet \nhttps://arxiv.org/abs/2204.03458\n[2] Diffusion Posterior Sampling for General Noisy Inverse Problems. Hyungjin Chung, Jeongsol Kim, Michael T. Mccann, Marc L. Klasky, Jong Chul Ye\nhttps://arxiv.org/abs/2209.14687\n[3] Pseudoinverse-Guided Diffusion Models for Inverse Problems. Jiaming Song, Arash Vahdat, Morteza Mardani, Jan Kautz\nhttps://openreview.net/forum?id=9_gsMA8MRKQ\n[4] Universal Guidance for Diffusion Models. Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, Tom Goldstein\nhttps://arxiv.org/abs/2302.07121", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8hUUqnXC7C", "reviewer": "Reviewer_ufFd", "summary": "The work investigates how to increase the diversity of one batch of samples for diffusion models. Specifically, the authors introduce a method named particle guidance, which is based on the gradient of a crafted time-varying potential field. The authors conduct various experiments to demonstrate the effectiveness of the method and provide an interesting analysis of the proposed method. However, some statements need to be clarified and more details about the experiments are needed.", "questions": "See above", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The topic, increasing diversity of generated samples, is important and crucial.\n- The theoretical analysis conducted in the main paper and appendix are non-trivial and interesting.\n- Experiments presented in the paper look good and show visual improvements.", "weaknesses": "Will the marginal distribution of Eq-1 be the same as that of the original diffusion model? It appears that the proposed method alters the marginal distribution. I am concerned that this shift in distribution may not be desirable in many applications.\n\nIt seems that most experiments conducted in this work focus on small batch sizes. I am interested in the authors' discussion regarding the scalability and effectiveness of the proposed method for large batch size.\n\nThere doesn't seem to be a principled approach to designing the proposed potential field besides the method present in Sec 6, which demands non-trivial training.\n\nThe illustrated plots in Figure 1 are a bit misleading to me. Why are the initial points concentrated in one mode? In a high-dimensional setting, the chance of sampling close-by Gaussian points is low.\n\nWhy were different guidance weights chosen for the IID and particle guidance experiments?\n\nFor fair evaluation purposes, could the authors post un-cherry-picked images on an anonymous website, for example, the first 50 text prompts of MS COCO or PartiPrompts? This should be done with the same hyper-parameters and random seed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8LXK5VocW5", "reviewer": "Reviewer_AGeB", "summary": "This paper proposes the use of particle guidance for sampling in diffusion-based models, emphasizing that particle guidance increases diversity without reducing quality.", "questions": "The authors are encouraged to further elucidate in what practical situations iid might fail (e.g., when the number of modes is close to the number of samples).\n\nHas the paper investigated the impact of the number of particles on the final generation outcome?\n\nIs there theoretical proof that particle guidance is better to some extent (even though this aligns with intuition)?\n\nThe paper uses an ODE solver. Would there be any changes if an SDE solver is used?", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Diversity is a crucial property in generative modeling and sampling. Whether for multi-modal sampling or the regeneration of real data, we aim to cover each mode. The problem studied in this paper is significant.\n\nThe paper's approach is succinct and clear, making it easy to follow.", "weaknesses": "Despite the author's belief that diffusion models may suffer from mode collapse, previous experience suggests that mode collapse in diffusion models is not particularly severe [1]. If mode collapse is not a leading factor in most cases, it might affect the significance of this study. 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. The main problem in Figure 4 (d) seems to be overfitting to the training data (an artificial setting). I think 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 I am not an expert in that area, the authors could further explain.\n\nThe paper's theoretical foundation is weak and does not directly explain why particle guidance can increase diversity.\n\n\n[1] A Study on Sample Diversity in Generative Models: GANs vs. Diffusion Models", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)