Diffusion-based Molecule Generation with Informative Prior Bridges
Abstract
AI-based molecule generation provides a promising approach to a large area of biomedical sciences and engineering, such as antibody design, hydrolase engineering, or vaccine development.
Because the molecules are governed by physical laws, a key challenge is to incorporate prior information into the training procedure to generate high-quality and realistic molecules.
We propose a simple and novel approach to steer the training of diffusion-based generative models with physical and statistics prior information.
This is achieved by constructing physically informed diffusion bridges, stochastic processes that guarantee to yield a given observation at the ﬁxed terminal time.
We develop a Lyapunov function based method to construct and determine bridges, and propose a number of proposals of informative prior bridges for both high-quality molecule generation and uniformity-promoted 3D point cloud generation.
With comprehensive experiments, we show that our method provides a powerful approach to the 3D generation task, yielding molecule structures with better quality and stability scores and more uniformly distributed point clouds of high qualities.
Introduction
As exempliﬁed by the success of AlphafoldV2 [21] in solving protein folding deep learning techniques have been creating new frontiers on molecular sciences [45].
In particular, the problem of building deep generative models for molecule design has attracted increasing interest with a magnitude of applications in physics, chemistry, and drug discovery [ e.g. , 2, 3, 25].
Recently, diffusion-based generative model have been applied to molecule generation problems [9, 18] and obtain superior performance.
The idea of these methods is to corrupt the data with diffusion noise and learn a neural diffusion model to revert the corruption process to generate meaningful data from noise
A key challenge in deep generative models for molecule and 3D point generation is to efﬁciently incorporate strong prior information to reﬂect the physical and problem-dependent statistical properties of the problems at hand.
In fact, a recent fruitful line of research [11, 23, 36] have shown promising results by introducing inductive bias into the design of model architectures to reﬂect physical constraints such as SE(3) equivariance.
In this work, we present a different paradigm of prior incorporation tailored to diffusion-based generative models, and leverage it to yield substantial improvement in both 1) high-quality and stable molecule generation and 2) uniformity-promoted point cloud generation.
Our contributions are summarized as follows.
Prior Guided Learning of Diffusion Models.
We introduce a simple and ﬂexible framework for injecting informative problem-dependent prior and physical information when learning diffusionbased generative models.
The idea is to elicit and inject prior information regarding how the diffusion process should look like for generating each given data point, and train the neural diffusion model to imitate the prior processes.
The prior information is presented in the form of diffusion bridges which are diffusion processes that are guaranteed to generate each data point at the ﬁxed terminal time.
We provide a general Lyapunov approach for constructing and determining bridges and leverage it to develop a way to systematically incorporate prior information into bridge processes.
Physics-informed Molecule Generation.
We apply our method to molecule generation.
We propose a number of energy functions for incorporating physical and statistical prior information.
Compared with existing physics-informed molecule generation methods [ e.g. , 9, 14, 28, 14], our method modiﬁes the training process, rather than imposing constraints on the model architecture.
Experiments show that our method achieves current state-of-the-art generation quality and stability on multiple test benchmarks of molecule generation.
Uniformity-promoting Point Generation.
A challenging task in physical simulation, graphics, 3D vision is to generate point clouds for representing real objects [ e.g. , 1, 5, 27, 50].
A largely overlooked problem of existing approaches is that they tend to generate unevenly distributed points, which lead to unrealistic shapes and make the subsequent processing and applications, such as mesh generation, challenging and inefﬁcient.
In this work, we leverage our framework to introduce uniformity-promoting forces into the prior bridge of diffusion generative models.
This yields a simple and efﬁcient approach to generating regular and realistic point clouds in terms of both shape and point distribution.
Related works
Diffuse Bridge Process.
Diffusion-based generative models [17, 39, 40, 43, 24] have achieved great successes in various AI generation tasks recently; these methods leverage a time reversion technique and can be viewed as learning variants auto-encoders with diffusion processes as encoders and decoders.
Schrodinger bridges [7, 9, 44] have also been proposed for learning diffusion generative models that guarantee to output desirable outputs in a ﬁnite time interval, but these methods involve iterative proportional ﬁttings and are computationally costly.
Our framework of learning generative models with diffusion bridges is similar to that of [34], which learn diffusion models as a mixture of forward-time diffusion bridges to avoid the time-reversal technique of [42].
But our framework is designed to incorporate physical prior into bridges and develop a systematic approach for constructing a broad class of prior-informed bridges. 3D Molecule Generation.
Generating molecule in 3D space has been gaining increasing interest.
A line of works [ e.g. 29, 31, 37, 38, 47, 48, 49] consider conditional conformal generation, which takes the 2D SMILE structure as conditional input and generate the 3D molecule conformations condition on the input.
Another series of works [ e.g. , 13, 18, 26, 36, 46] focus on directly generating the atom position and type for the molecule unconditionally.
For these series of works, improvements usually come from architecture design and loss design.
For example, G-Schnet [13] auto-regressively generates the atom position and type one by one after another; EN-Flow [36] and EDM [18] adopt
E(n) equivariant graph neural network (EGNN) [36] to train ﬂow-based model and diffusion model.
These methods aim at generating valid and natural molecules in 3D space and outperform previous approaches by a large margin.
Our work provides a very different approach to incorporating the physical information for molecule generation by injecting the prior information into the diffusion process, rather than neural network architectures.
Point Cloud Generation.
A vast literature has been devoted to learning deep generative models for real-world 3D objects in the form of point clouds.
[1] ﬁrst proposed to generate the point cloud by generating a latent code and training a decoder to generate point clouds from the latent code.
Build upon this approach, methods have been developed using ﬂow-based generative models [50] and diffusion-base models [5, 27, 28].
However, the existing works miss a key important prior information: the points in a point cloud tend to distribute regularly and uniformly.
Ignoring this information causes poor generation quality.
By introducing uniformity-promoting forces in diffusion bridges, we obtain a simple and efﬁcient approach to generating regular and realistic point clouds.
Method
We ﬁrst introduce the deﬁnition of diffusion generative models and discuss how to learn these models with prior bridges.
After introducing the training algorithm for deep diffusion generative models, we discuss the energy functions that we apply to molecules and point cloud examples.
3.1 Learning Diffusion Generative Models with Prior Bridges
Problem Deﬁnition.
We aim at learning a generative model given a dataset { x ( k ) } n k =1 drawn from an unknown distribution ⇧ ⇤ on R d .
A diffusion model on time interval [0 , 1] is
[Equation]
where W t is a standard Brownian motion;  t : R d !
R d ⇥ d is a positive deﬁnition covariance coefﬁcient; s ✓ t : R d !
R d is parameterized as a neural network with parameter ✓ , and µ 0 is the initialization.
Here we use P ✓ to denote the distribution of the whole trajectory Z = { Z t : t 2 [0 , 1] } , and P ✓ t the marginal distribution of Z t at time t .
We want to learn the parameter ✓ such that the distribution P ✓ 1 of the terminal state Z 1 equals the data distribution ⇧ ⇤ .
Learning Diffusion Models.
There are an inﬁnite number of diffusion processes P ✓ that yield the same terminal distribution but have different distributions of latent trajectories Z .
Hence, it is important to inject problem-dependent prior information into the learning procedure to obtain a model
P ✓ that simulate the data for the problem at hand fast and accurately.
To achieve this, we elicit an imputation process Q x for each x 2 R d , such that a draw Z ⇠ Q x yields trajectories that 1) are consistent with x in that Z 1 = x deterministically, and 2) reﬂect important physical and statistical prior information on the problem at hand.
Formally, if Q x ( Z 1 = x ) = 1 , we call that Q x is a bridge process pinned at end point x , or simply an x -bridge.
Assume we ﬁrst generate a data point x
⇠ ⇧ ⇤ , and then draw a bridge Z ⇠ Q x pinned at x , then the distribution of Z is a mixture of Q x with x drawn from the data distribution:
[Equation]
A key property of Q ⇧ ⇤ is that its terminal distribution equals the data distribution, i.e., Q ⇧ ⇤ 1 = ⇧ ⇤ .
Therefore, we can learn the diffusion model P ✓ by ﬁtting the trajectories drawn from Q ⇧ ⇤ with the “backward” procedure above.
This can be formulated by maximum likelihood or equivalently minimizing the KL divergence:
[Equation]
Furthermore, assume that the bridge Q x is a diffusion model of form
[Equation]
where b t ( Z t | x ) is an x -dependent drift term need to carefully designed to both satisfy the bridge condition and incorporate important prior information (see Section 3.2).
Assuming this is done, using
Girsanov theorem [32], the loss function L ( ✓ ) can be reformed into a form ofdenoised score matching loss of [ e.g. , 40, 42, 41]: [Equation]  [Equation]
[Equation]
const ,
[Equation]
which is a score matching term between s ✓ and b .
The const term contains the log-likelihood for the initial distribution µ 0 , which is a const in our problem.
Here ✓ ⇤ is an global optimum of L ( ✓ ) if
[Equation]
This means that the drift term s ✓ t should be matched with the conditional expectation of b t ( z | x ) with x = Z 1 conditioned on Z t = z .
Remark 3.1.
The SMLD can be viewed as a special case of this framework when we take Q x to be a time-scaled Brownian bridge process:
[Equation]
where  t 2 [0 , + 1 ) and  t =
R t 0  2 s d s .
This can be seen by the fact that the time-reversed process ˜ Z t :=
Z 1  t follows the simple time-scaled Brownian motion d ˜ Z t =
1  t d ˜ W t starting from the data point ˜
Z
= x , where ˜ W t is another standard Brownian motion.
The Brownian bridge achieves
Z 1 = x because the magnitude of the drift force is increasing to inﬁnite when t is close to time 1 .
However, the bridge of SMLD above is a relative simple and uninformative process and does not incorporate problem-dependent prior information into the learning procedure.
This is also the case of the other standard diffusion-based models [42], such as denoising diffusion probabilistic models (DDPM) which can be shown to use a bridge constructed from an Ornstein–Uhlenbeck process.
We refer the readers to [34], which provides a similar forward time bridge framework for learning diffusion models, and it recovers the bridges in SMLD and DDPM as a conditioned stochastic process derived using the h -transform technique [10].
However, the h -transform method is limited to elementary stochastic processes that have an explicit formula of the transition probabilities, and can not incorporate complex physical statistical prior information.
Our work strikes to construct and use a broader class of more complex bridge processes that both reﬂect problem-dependent prior knowledge and satisfy the endpoint condition Q x ( Z 1 = x )
= 1 .
This necessitate systematic techniques for constructing a large family of bridges, as we pursuit in Section 3.2.
3.2 Designing Informative Prior Bridges
The key to realizing the general prior-informed learning framework above is to have a general and user-friendly technique to design Q x in (1) to ensure the bridge condition Q x ( Z 1 = x ) = 1 while leaving the ﬂexibility of incorporating rich prior information.
To achieve this, we ﬁrst develop a general criterion of bridges based on a Lyapunov function method which allows us to identify a very general form of bridge processes; we then propose a particularly simple family of bridges that we use in practice by introducing modiﬁcation to Brownian bridges.
Deﬁnition 3.2 ( Lyapunov Functions ) .
A function U t ( z ) is said to be a Lyapunov function for set
A ⇢ R d at time t = 1 if U 1 ( z )  0 for 8 z 2 R d and U 1 ( z )
=
if and only if z 2 A .
Intuitively, a diffusion process Q is a bridge A , i.e., Q ( Z 1 2 A )
= 1 , if it (at least) approximately follows the gradient ﬂow of a Lyapunov function and the magnitude (or step size) or the gradient ﬂow should increase with a proper magnitude in order to ensure that Z t 2 A at the terminal time t = 1 .
Therefore, we identify a general form of bridges to A as follows:
[Equation]
where ↵ t > 0 is the step size of the gradient ﬂow of U and ⌫ is an extra perturbation term.
The step size ↵ t should increase to inﬁnity as t !
1 sufﬁciently fast
to dominate the effect of the diffusion term  t d W t and the perturbation ⌫ t d t term to ensure that U is minimized at time t =
Proposition 3.3.
Assume U t
( z ) =
U ( z, t ) is a Lyapunov function of a measurable set
A at time 1 and
U ( · , t ) 2 C 2 ( R d ) and U ( z, · ) 2 C 1 ([0 , 1]) .
Then, Q A in (4) is an bridge to A , i.e., Q A ( Z 1 2 A )
= 1 , if the following holds:
1) U follows an (expected) Polyak-Lojasiewicz condition: E Q A [ U t ( Z t )]  kr z
U t ( Z t ) k 2 ]  0 , 8 t. 2) Let  t =
E Q A [ r z U t ( Z t ) > ⌫ t ( Z t )] , and  t =
E Q A [ @ t U t ( Z t )
+ 12 tr( r 2 z U t ( Z t )  2 t ( Z t ))] , and ⇣ t = exp ( R t
↵ s d s ) .
Then lim t " 1 ⇣ t = + 1 , and lim t "
[Equation]
⇣ s (  s +  s )d s = + 1 .
[Figure]
Energy
Brownian bridge can be viewed as the case when U t ( z ) =
k x  z k 2 / 2 and ↵ t =  2 t / (  1   t ) , and ⌫ = 0 .
Hence simply introducing an extra drift term into bridge bridge yields that a broad family of bridges to x :
[Equation]
In Appendix A.4 and A.5, we show that Q x, bb ,f is a bridge to x if E Q x, bb [ k f t ( Z t )
k 2 ] <
+ 1 and  t > 0 , 8 t , which is very mild condition and is satisﬁed for most practical functions.
The intuition is that the Brownian drift  2 t x  Z t  1   t is singular and grows to inﬁnite as t approaches 1 .
Hence, introducing an f into the drift would not change of the ﬁnal bridge condition, unless f is also singular and has a magnitude that dominates the Brownian bridge drift as t ! 1 .
To make the model P ✓ compatible with the physical force f , we assume the learnable drift has a form of s ✓ t ( z ) =
↵ f t ( z ) + ˜ s ✓ t ( z ) where ˜ s is a neural network (typically a GNN) and ↵ can be another learnable parameter or a pre-deﬁned parameter.
Please refer to algorithm 3.2 and Figure 1 for descriptions about our practical algorithm.
Input : Given a dataset { x ( k ) } , Q x the bridge in (5),
and a problem-dependent prior force f ,v and a diffusion model P ✓ .
Training : Estimate ✓ by minimizing L ( ✓ ) in (2) with stochastic gradient descent and time discretization.
Sampling : Simulate from P ✓ .
Molecule and 3D Generation with Informative Prior Bridges
We apply our method to the molecule generation as well as point cloud generation.
Informative physical or statistical priors that reﬂects the underlying real physical structures can be particularly beneﬁcial for molecule generation as we show in experiments.
In our problem, each data point x is a collection of atoms of different types, more generally marked points, in 3D Euclidean space.
In particular, we have x = [ x ri , x hi ] mi =1 , where x ri 2 R 3 is the coordinate of the i -th atom, and x hi 2 { e 1 , . . .
, e k } where each e i =
[0 · · · 1 · · · 0] is the i -th basis vector of R k , which indicates the type of the i -th atom of k categories.
To apply the diffusion generative model, we treat x hi as a continuous vector in R r and round it to the closest basis vector when we want to output a ﬁnal result or have computations that depend on atom types (e.g., calculating an energy function as we do in sequel).
Speciﬁcally, for a continuous x hi 2 R k , we denote by ˆ x hi =
I ( x hi = max( x hi ))
the discrete type rounded from it by taking the type with the maximum value.
To incorporate priors, we design an energy function E ( x ) and incorporate f t ( · ) =
r E ( · ) into the Brownian bridge (5) to guide the training process.
We discuss different choices of E in the following.
4.1 Prior Bridges for Molecule Generation
Previous prior guided molecule or protein 3D structure generation usually depends on pre-deﬁned energy or force [29, 49].
We introduce our two potential energies.
One is formulated inspired by previous works in biology, and the other is an k nearest neighbour statistics directly obtained from the data.
AMBER Inspired Physical Energy.
AMBER [12] is a family of force ﬁelds for molecule simulation.
It is designed to provide a computationally efﬁcient tool for modern chemistry-molecular dynamics and free energy calculations.
It consists of a number of important forces, including the bond energy, angular energy, torsional energy, the van der Waals energy and the Coulomb energy.
Inspired by AMBER, we propose to incorporate the following energy term into the bridge process:
[Equation]
• The bond energy is E bond ( x ) =
[Equation] 2 bond ( x ) (Len( x rij )
` (ˆ x hi , ˆ x hj )) 2 , where Len( x rij )
=
x ri  x rj  , and bond ( x ) denotes the set of bonds from x , which is set to be the set of atom pairs with a distance smaller than 1 .
15 times the covalent radius; the ` 0 ( r, c ) denotes the expected bond length between atom type r and c , which we calculate as side information from the training data.
• The angle energy is E angle ( x ) =
{+P ijk 2 angle ( x ) (Ang( x rijk )  ! 
+}
(ˆ x hijk )) 2 , where angle ( x ) denotes the set of angles between two neighbour bonds in bound ( x ) , and Ang( x rijk ) denotes the angle formed by vector x r i  x r j and x r k  x r j , and !
(ˆ x h ijk ) is the expected angle between atoms of type ˆ x h i , ˆ x h j , ˆ x hk , which we calculate as side information from the training data.
• The Lennard-Jones (LJ) energy is deﬁned by E LJ ( x ) =
P [Equation] 6 =
j e (  x ri  x rj  ) and e ( ` ) =
(  / ` ) 12  2(  / ` ) 6 .
The parameter  is an approximation for average nucleus distance.
• The nuclei-nuclei repulsion (Coulomb) electromagnetic potential energy is E Coulomb ( x ) =
[Equation] ij q [Equation] ˆ x hi
[Equation] x [Equation]  x ri  x rj  , where  is Coulomb constant and q ( r ) denotes the point charge of atom of type r , which depends on the number of protons.
Statistical Energy.
When accurate physic laws are unavailable, molecular geometric statistics, such as bond lengths, bond angles, and torsional angles, etc, can be directly calculated from the data and shed important insights on the system [ e.g. , 8, 20, 30].
We propose to design a prior energy function in bridges by directly calculate these statistics over the dataset.
Speciﬁcally, we assume that the lengths and angles of each type of bond follows a Gaussian distribution that we learn from the dataset, and deﬁne the energy function as the negative log-likelihood:
E
[Equation]
where knn ( x ) denotes the K-nearest neighborhood graph constructed based on the distance matrix of x ; for each pair of atom types r, c 2 [ k ] , ˆ µ rc and ˆ  2 rc denotes empirical mean and variance of length of rc -edges in the dataset; for each triplet r, c, r 0 2 [ k ] , ˆ µ rcr 0 and ˆ  2 rcr 0 is the empirical mean and variance of angle betwen rc and cr 0 bonds.
Intuitively, depending on the atom type and order of the nearest neighbour, we force the atom distance and angle to mimic the statistics calculated from the data.
We thus implicitly capture different kinds of interaction forces.
Compared with the AMBER energy, the statistical energy (7) is simpler and more adaptive to the dataset of interest.
4.2 Prior Bridges for Point Cloud Generation
We design prior forces for 3D point cloud generation, which is similar to molecule generation except that the points are un-typed so we only have the coordinates { x ri } .
One important aspect of point cloud generation is to distribute points uniformly on the surface, which is important for producing high-quality meshes and other post-hoc geometry applications and manipulations.
Riesz Energy.
One idea to make the point distribute uniformly is adding a repulsive force to separate the points apart from each other [15].
We achieve this by minimizing the Riesz energy [16],
[Equation] [Equation]
KNN Distance Energy.
Similar to molecule design, we directly calculate the average distance between each point and its k nearest neighbour neighbour, and deﬁne the following energy:
[Equation]
[Equation] where knn-dist i ( x ) = 1 K
P j 2 N K ( x i ; x ) ||
x ri  x rj || 2 denotes the average distance from x i to its
K nearest neighbors, and µ knn is the empirical mean of knn-dist i ( x ) in the dataset.
This would encourage the points to have similar average nearest neighbor distance and yield uniform distribution between points.
In common geometric setups, the valence of the point on the surface is 4, which means we set k = 4 .
Experiment
We verify the advantages of our proposed method (Bridge with Priors) in several different domains.
We ﬁrst compare our method with advanced generators ( e.g. , diffusion model, normalizing ﬂow, etc.) on molecule generation tasks.
We then implement our method on point cloud generations, which targets producing generated samples in a higher quality.
We directly compare the performance and also analyze the difference between our energy prior and other energies we discuss in Section 3.
5.1 Force Guided Molecule Generation
To demonstrate the efﬁciency and effectiveness of our bridge processes and physical energy, we conduct experiments on molecule and macro-molecule generation experiments.
We follow [28] in settings and observe that our proposed prior bridge processes consistently improve the state-of-the-art performance.
Diving deeper, we analyze the impact of different energy terms and hyperparameters.
Metrics.
Following [18, 36], we use the atom and molecular stability score to measure the model performance.
The atom stability is the proportion of atoms that have the right valency while the molecular stability stands for the proportion of generated molecules for which all atoms are stable.
For visualization, we use the distance between pairs of atoms and the atom types to predict bond types, which is a common practice.
To demonstrate that our force does not only memorize the data in the dataset, we further calculate and report the RDKit-based [22] novelty score.
we extracted 10, samples to calculate the above metrics.
Dataset Settings QM9 [35] molecular properties and atom coordinates for 130k small molecules with up to 9 heavy atoms with 5 different types of atoms.
This data set contains small amino acids, such as GLY, ALA, as well as nucleobases cytosine, uracil, and thymine.
We follow the common practice in [18] to split the train, validation, and test partitions, with 100K, 18K, and 13K samples.
GEOM-DRUG [4] is a dataset that contains drug-like molecules.
It features 37 million molecular conformations annotated by energy and statistical weight for over 450,000 molecules.
Each molecule contains 44 atoms on average, with 5 different types of atoms.
Following [18, 36], we retain thelowest energy conformations for each molecule.
Training Conﬁgurations.
On QM9, we train the EGNNs with 256 hidden features and 9 layers for 1100 epochs, a batch size 64, and a constant learning rate 10  4 , which is the default training conﬁguration.
We use the polynomial noise schedule used in [18] which linearly decay from 10  2 /T to 0.
We linearly decay ↵ from 10  3 /T to 0 w.r.t.
time step.
We set k = 5 (7) by default.
On
GEOM-DRUG, we train the EGNNs with 256 hidden features and 8 layers with batch size 64, a constant learning rate 10  4 , and 10 epochs.
It takes approximately 10 days to train the model on these two datasets on one Tesla V100-SXM2-32GB GPU.
We provide E(3) Equivariant Diffusion
Model (EDM) [18] and E(3) Equivariant Normalizing Flow (EN-Flow) [36] as our baselines.
Both two are trained with the same conﬁgurations as ours.
Results: Higher Quality and Novelty.
We summarize our experimental results in Table 1.
We observe that (1) our method generates molecules with better qualities than the others.
On QM9, we notice that we improve the molecule stability score by a large margin (from 82 .
to 84 . 6 ) and slightly improve the atom stability score (from 98 .
7 to 98 . 8 ).
It indicates that with the informed prior bridge helps improves the quality of the generated molecules.
(2) Our method achieves a better novelty score.
Compared to E-GDM, we improve the novelty score from 65 .
7 to 68 . 8 .
This implies that our introduced energy does not hurt the novelty when the statistics are estimated over the training dataset.
[Table]
Notice that although the GDM and EN-Flow achieve a better novelty score, the sample quality is much worse.
The reason is that, due to the metric deﬁnition, low-quality out-of-distribution samples lead to high novelty scores.
(3) On the GEOM-DRUG dataset, the atom stability is improved from 81 .
3 to 82 . 4 , which shows that our method can work for macro-molecules.
(4) We visualize and qualitatively evaluate our generate molecules.
Figure 3 displays the trajectory on GEOM-DRUG and
Figure 2 shows the samples on two datasets.
(5) Bridge processes and E-GDM obtain comparable results on our tested benchmarks.
The computational load added by introducing prior bridges is small.
Compared to EGM, we only introduce 8% additional cost in training and 3% for inference.
Result: Better With Fewer Time Steps.
We display the performance of our method with fewer time steps in Table 2.
We observe that (1) with fewer time steps, the baseline EGM method gets worse results than 1000 steps in Table 1.
(2) with 500 steps, our method still keeps a consistently good performance.
(3) with even fewer 50 or 100 steps, our method yields a worse result than 1000 steps in Table 1, but still outperforms the baseline method by a large margin.
Ablation: Impacts of Different Energies.
We apply several energies we discuss in Section 3, and compare them on the QM9 dataset.
(1) We notice that our energy (7) gets better performance with larger k when k  5 .
k = 7 achieves comparable performance as k = 5 .
Larger k also requires more computation time, which yields a trade-off between performance and efﬁciency.
(2) For (6), once removing a typical term, the performance drops.
In all the cases, applying additional forces outperforms the bridge processes baseline w/o.
force.
5.2 Force Guided Point Cloud Generation
We apply uniformity-promoting priors to point cloud generation.
We apply our method based on the diffusion model for point cloud generation introduced by point cloud diffusion model [27] and compare it with the original diffusion model as well as the case of bridge processes w/o.
force prior.
We observe that our method yields better results in various evaluation metrics under different setups.
Dataset.
We use the ShapeNet [6] dataset for point cloud generation.
ShapeNet contains 55 categories.
We select Airplane and Chair, which are the two most common categories to evaluate in recent point cloud generation works [5, 27, 50, 51].
We construct the point clouds following the setup in [27], split the train, valid and test dataset in 80% , 15% and 5% and samples 2048 points uniformly on the mesh surface.
Evaluation Metric.
We evaluate the generated shape quality in two aspects following the previous works, including the minimum matching distance (MMD) and coverage score (COV).
These scores are the two most common practices in the previous works.
We use Chamfer Distance (CD) and Earth
Mover’s Distance (EMD) as the distance metric to compute the MMD and COV.
Experiment Setup.
We train the model with two different conﬁgurations.
The ﬁrst one uses exactly the same experiment setup conﬁguration introduced in [27].
Thus, we use the same model architecture and train the model in 100 diffuse steps with a learning rate 2 ⇥ 10  3 , batch size 128, and linear
[Figure]
[Table] [Table] noise schedule from 0 .
02 to 10  4 .
We initial ↵ with 0.1 and jointly learn it with the network.
For the second setup, to evaluate the better converge speed of our method, we decrease the diffuse step from 100 to 10 with other settings the same.
For the diffusion model baseline, we reproduce the number by directly using the pre-trained model checkpoint and testing it on the test set provided by the ofﬁcial codebase. [Figure]
Result.
We show our experimental result in Table 4.
We see that (1)
In the 10 steps setup, all variants of our approach are clearly stronger than the diffusion model.
With force added, our method with physical prior achieves nearly the same performance as the 100 steps setup.
In the 100 steps setup, adding energy potential as prior improves the bridge process performance and further let it beat the diffusion model baseline.
(3) Since the test points are uniformly sampled on the surface, a better score indicates a closer point distribution to the reference set.
Further, when compare with Riesz energy (8), statistic gap energy (9) performs better.
One explanation is the Riesz energy pushes the points to some outlier position in sample generate samples, while statistic gap energy is more robust.
We also show visualization samples in Figure 4.
Conclusion and Limitations
We propose a framework to inject informative priors into learning neural parameterized diffusion models, with applications to both molecules and 3D point cloud generation.
Empirically, we demonstrate that our method has the advantages such as better generation quality, less sampling time and easy-to-calculate potential energies.
For future works, we plan to 1) study the relation between different types of forces for different domain of molecules, 2) study how to generate valid proteins in [Table] which the number of atoms is very large [Table]
and 3) apply our method to more realistic applications such as antibody design or hydrolase engineering
[Table] In both energy functions in (7) and (6) [Table] we do not add torsional angle related energy [19] mainly because it is hard to verify whether four atoms are bonded together during the stochastic process.
We plan to study how to include this for better performance in future works.
Another weakness of deep diffusion bridge processes are their computation time.
Similar to previous diffusion models [28], it takes a long time to train a model.
We attempted to speed the training up by using a large batch size ( e.g. , 512, 1024) but found a performance drop.
An important future direction is to study methods to distribute and accelerate the training.
Acknowledgements Authors are supported in part by CAREER-1846421, SenSE-2037267,
EAGER-2041327, and Ofﬁce of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML).
We would like to thank the anonymous reviewers and the area chair for their thoughtful comments and efforts towards improving our manuscript.