Diffusion-based Molecule Generation with Informative Prior Bridges
Anonymous Author(s)
AfﬁliationAddress email
Abstract
AI-based molecule generation provides a promising approach to a large area ofbiomedical 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 togenerate high-quality and realistic molecules.
We propose a simple and novel approach to steer the training of diffusion-based generative models with physical andstatistics prior information.
This is achieved by constructing physically informeddiffusion bridges, stochastic processes that guarantee to yield a given observationat the ﬁxed terminal time.
We develop a Lyapunov function based method toconstruct and determine bridges, and propose a number of proposals of informativeprior bridges for both high-quality molecule generation and uniformity-promoted3D point cloud generation.
With comprehensive experiments, we show that ourmethod provides a powerful approach to the 3D generation task, yielding moleculestructures with better quality and stability scores and more uniformly distributedpoint clouds of high qualities.
Introduction
As exempliﬁed by the success of AlphafoldV2 [20] in solving protein folding, deep learning techniques have been creating new frontiers on molecular sciences [43].
In particular, the problem ofbuilding deep generative models for molecule design has attracted increasing interest with a magnitudeof applications in physics, chemistry, and drug discovery [ e.g. , 2, 3, 23].
Recently, diffusion-basedgenerative model have been applied to molecule generation problems [9, 17] and obtain superiorperformance.
The idea of these methods is to corrupt the data with diffusion noise and learn a neuraldiffusion 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ﬁcientlyincorporate 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, 22, 34] have shownpromising results by introducing inductive bias into the design of model architectures to reﬂectphysical constraints such as SE(3) equivariance.
In this work, we present a different paradigm ofprior incorporation tailored to diffusion-based generative models, and leverage it to yield substantialimprovement in both 1) high-quality and stable molecule generation and 2) uniformity-promotedpoint 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.
Weprovide a general Lyapunov approach for constructing and determining bridges and leverage it todevelop a way to systematically incorporate prior information into bridge processes.
Physics-informed Molecule Generation.
We apply our method to molecule generation.
We proposea number of energy functions for incorporating physical and statistical prior information.
Comparedwith existing physics-informed molecule generation methods [ e.g. , 9, 14, 26, 14], our method modiﬁesthe training process, rather than imposing constraints on the model architecture.
Experiments showthat our method achieves current state-of-the-art generation quality and stability on multiple testbenchmarks 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, 25, 48].
A largelyoverlooked 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 asmesh generation, challenging and inefﬁcient.
In this work, we leverage our framework to introduceuniformity-promoting forces into the prior bridge of diffusion generative models.
This yields a simpleand efﬁcient approach to generating regular and realistic point clouds in terms of both shape andpoint distribution.
Related works
Diffuse Bridge Process.
Diffusion-based generative models [16, 37, 38, 41] have achieved greatsuccesses in various AI generation tasks recently; these methods leverage a time reversion techniqueand can be viewed as learning variants auto-encoders with diffusion processes as encoders anddecoders.
Schrodinger bridges [7, 9, 42] have also been proposed for learning diffusion generativemodels that guarantee to output desirable outputs in a ﬁnite time interval, but these methods involveiterative proportional ﬁttings and are computationally costly.
Our framework of learning generativemodels with diffusion bridges is similar to that of [32], which learn diffusion models as a mixture offorward-time diffusion bridges to avoid the time-reversal technique of [40].
But our framework isdesigned to incorporate physical prior into bridges and develop a systematic approach for constructinga 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. 27, 29, 35, 36, 45, 46, 47] consider conditional conformal generation, whichtakes the 2D SMILE structure as conditional input and generate the 3D molecule conformationscondition on the input.
Another series of works [ e.g. , 13, 17, 24, 34, 44] focus on directly generatingthe atom position and type for the molecule unconditionally.
For these series of works, improvementsusually come from architecture design and loss design.
For example, G-Schnet [13] auto-regressivelygenerates the atom position and type one by one after another; EN-Flow [34] and EDM [17] adopt
E(n) equivariant graph neural network (EGNN) [34] to train ﬂow-based model and diffusion model.
These methods aim at generating valid and natural molecules in 3D space and outperform previousapproaches by a large margin.
Our work provides a very different approach to incorporating thephysical information for molecule generation by injecting the prior information into the diffusionprocess, rather than neural network architectures.
Point Cloud Generation.
A vast literature has been devoted to learning deep generative modelsfor real-world 3D objects in the form of point clouds.
[1] ﬁrst proposed to generate the point cloudby 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 [48]and diffusion-base models [5, 25, 26].
However, the existing works miss a key important priorinformation: the points in a point cloud tend to distribute regularly and uniformly.
Ignoring thisinformation causes poor generation quality.
By introducing uniformity-promoting forces in diffusionbridges, 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 modelswith prior bridges.
After introducing the training algorithm for deep diffusion generative models
3.1 Learning Diffusion Generative Models with Prior Bridges
Problem Deﬁnition.
We aim at learning a generative model given a dataset { x ( k ) } nk =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 covariancecoefﬁcient; s ✓ t : R d !
R d is parameterized as a neural network with parameter ✓ , and µ 0 is theinitialization.
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 thedistribution 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 yieldthe same terminal distribution but have different distributions of latent trajectories Z .
Hence, it isimportant 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 animputation process Q x for each x 2 R d , such that a draw Z ⇠ Q x yields trajectories that 1) areconsistent with x in that Z 1 = x deterministically, and 2) reﬂect important physical and statisticalprior 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 simplyan x -bridge.
Assume we ﬁrst generate a data point x
⇠ ⇧ ⇤ , and then draw a bridge Z ⇠ Q xpinned at x , then the distribution of Z is a mixture of Q x with x drawn from the data distribution:
Q ⇧ ⇤ :=
R
Q [Equation]
[Equation]
) ⇧ ⇤ (d x ) . 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 bridgecondition and incorporate important prior information (see Section 3.2).
Assuming this is done, using
Girsanov theorem [30], the loss function L ( ✓ ) can be reformed into a form ofdenoised score matchingloss of [ e.g. , 38, 40, 39]:
[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 ) withx =
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 atime-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 datapoint ˜
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 notincorporate problem-dependent prior information into the learning procedure.
This is also the caseof the other standard diffusion-based models [40], 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 [32], which provides a similar forward time bridge framework for learningdiffusion models, and it recovers the bridges in SMLD and DDPM as a conditioned stochasticprocess derived using the h -transform technique [10].
However, the h -transform method is limited toelementary stochastic processes that have an explicit formula of the transition probabilities, and cannot incorporate complex physical statistical prior information.
Our work strikes to construct and use abroader class of more complex bridge processes that both reﬂect problem-dependent prior knowledgeand satisfy the endpoint condition Q x ( Z 1 = x )
= 1 .
This necessitate systematic techniques forconstructing 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 anduser-friendly technique to design Q x in (1) to ensure the bridge condition Q x ( Z 1 = x
) =
1 whileleaving the ﬂexibility of incorporating rich prior information.
To achieve this, we ﬁrst develop ageneral criterion of bridges based on a Lyapunov function method which allows us to identify a verygeneral form of bridge processes; we then propose a particularly simple family of bridges that we usein 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) approximatelyfollows the gradient ﬂow of a Lyapunov function and the magnitude (or step size) or the gradient ﬂowshould 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 stepsize ↵ t should increase to inﬁnity as t !
1 sufﬁciently fast
to dominate the effect of the diffusionterm  t d W t and the perturbation ⌫ t d t term to ensure that U is minimized at time t = 1 .
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 [Equation]  s +  s )d s = + 1 .
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 154 ⌫ = 0 .
Hence simply introducing an extra drift term into bridge bridge yields that a broad family ofbridges 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 intuitionis 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 singularand 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 hasa form of s ✓ t ( z ) =
↵ f t ( z ) + ˜ s ✓ t ( z ) where ˜ s is a neural network (typically a GNN) and ↵ can beanother learnable parameter or a pre-deﬁned parameter.
Please refer to algorithm 3.2 and Figure 1 fordescriptions about our practical algorithm.
Input : Given a
, Q the bridge
, and a problem-dependent prior force
,v and a diffusion model P
Training : Estimate ✓ by minimizing L ( ✓ ) in (2) with stochastic gradient descent and time discretization.
Sampling : Simulate from P ✓ .
[Figure]
Molecule and 3D Generation with Informative Prior Bridges
We apply our method to the molecule generation as well as point cloud generation.
Informativephysical or statistical priors that reﬂects the underlying real physical structures can be particularlybeneﬁ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 markedpoints, in 3D Euclidean space.
In particular, we have x = [ x ri , x hi ] mi =1 , where x ri 2 R 3 is thecoordinate of the i -th atom, and x hi 2 { e 1 , . . .
, e k } where each e i =
[0 · · · 1 · · · 0] is the i -th basisvector of R k , which indicates the type of the i -th atom of k categories.
To apply the diffusiongenerative model, we treat x hi as a continuous vector in R r and round it to the closest basis vectorwhen we want to output a ﬁnal result or have computations that depend on atom types (e.g., calculatingan 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 maximumvalue.
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 thefollowing.
4.1 Prior Bridges for Molecule Generation
Previous prior guided molecule or protein 3D structure generation usually depends on pre-deﬁnedenergy or force [27, 47].
We introduce our two potential energies.
One is formulated inspired byprevious works in biology, and the other is an k nearest neighbour statistics directly obtained fromthe 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-moleculardynamics and free energy calculations.
It consists of a number of important forces, including thebond 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 ) =
P ij 2 bond ( x ) (Len( x rij )
` (ˆ x hi , ˆ x hj )) 2 , where Len( x rij )
=
x r i  x r j  , and bond ( x ) denotes the set of bonds from x , which is set to be the set of atom pairswith a distance smaller than 1 .
15 times the covalent radius; the ` 0 ( r, c ) denotes the expected bondlength between atom type r and c , which we calculate as side information from the training data.
• The angle energy is E angle ( x ) =
[Equation]
) 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 ri  x rj and x rk  x rj , and !
(ˆ x hijk ) is the expected angle between atoms of type ˆ x hi , ˆ x hj ,ˆ 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 i 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 ) =
 P ij q (ˆ x hi ) q (ˆ x hj ) /
x ri  x rj  , where  is Coulomb constant and q ( r ) denotes the point chargeof atom of type r , which depends on the number of protons.
Statistical Energy.
When accurate physic laws are unavailable, molecular geometric statistics, suchas bond lengths, bond angles, and torsional angles, etc, can be directly calculated from the data andshed important insights on the system [ e.g. , 8, 19, 28].
We propose to design a prior energy functionin 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 ofx ; for each pair of atom types r, c 2 [ k ] , ˆ µ rc and ˆ  2 rc denotes empirical mean and variance of lengthof rc -edges in the dataset; for each triplet r, c, r 0 2 [ k ] , ˆ µ rcr 0 and ˆ  2 rcr 0 is the empirical mean andvariance of angle betwen rc and cr 0 bonds.
Intuitively, depending on the atom type and order of the nearest neighbour, we force the atom distanceand angle to mimic the statistics calculated from the data.
We thus implicitly capture different kindsof interaction forces.
Compared with the AMBER energy, the statistical energy (7) is simpler andmore 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 exceptthat the points are un-typed so we only have the coordinates { x ri } .
One important aspect of pointcloud generation is to distribute points uniformly on the surface, which is important for producinghigh-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 separatethe points apart from each other.
We achieve this by minimizing the Riesz energy [15],
[Equation] [Equation]
KNN Distance Energy.
Similar to molecule design, we directly calculate the average distancebetween each point and its k nearest neighbour neighbour, and deﬁne the following energy:
[Equation]
(9) 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 wouldencourage the points to have similar average nearest neighbor distance and yield uniform distributionbetween points.
In common geometric setups, the valence of the point on the surface is 4, whichmeans 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, whichtargets producing generated samples in a higher quality.
We directly compare the performance andalso 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, weconduct experiments on molecule and macro-molecule generation experiments.
We follow [26] insettings and observe that our proposed prior bridge processes consistently improve the state-of-the-artperformance.
Diving deeper, we analyze the impact of different energy terms and hyperparameters.
Metrics.
Following [17, 34], we use the atom and molecular stability score to measure the modelperformance.
The atom stability is the proportion of atoms that have the right valency while themolecular 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 bondtypes, which is a common practice.
To demonstrate that our force does not only memorize the data inthe dataset, we further calculate and report the RDKit-based [21] novelty score.
we extracted 10,000samples to calculate the above metrics.
Dataset Settings QM9 [33] molecular properties and atom coordinates for 130k small moleculeswith 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 commonpractice in [17] 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 molecularconformations annotated by energy and statistical weight for over 450,000 molecules.
Each moleculecontains 44 atoms on average, with 5 different types of atoms.
Following [17, 34], we retain the 30lowest energy conformations for each molecule.
Training Conﬁgurations.
On QM9, we train the EGNNs with 256 hidden features and 9 layersfor 1100 epochs, a batch size 64, and a constant learning rate 10  4 , which is the default trainingconﬁguration.
We use the polynomial noise schedule used in [17] which linearly decay from 10  2 /Tto 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, aconstant learning rate 10  4 , and 10 epochs.
It takes approximately 10 days to train the model onthese two datasets on one Tesla V100-SXM2-32GB GPU.
We provide E(3) Equivariant Diffusion
Model (EDM) [17] and E(3) Equivariant Normalizing Flow (EN-Flow) [34] as our baselines.
Bothtwo are trained with the same conﬁgurations as ours.
[Figure] [Table]
[Figure] [Figure]
[Table] [Table]
Results: Higher Quality and Novelty.
We summarize our experimental results in Table 1.
Weobserve that (1) our method generates molecules with better qualities than the others.
On QM9, wenotice that we improve the molecule stability score by a large margin (from 82 .
to 84 . 6 )
and slightlyimprove the atom stability score (from 98 .
7 to 98 . 8 ).
It indicates that with the informed prior bridgehelps improves the quality of the generated molecules.
(2) Our method achieves a better noveltyscore.
Compared to E-GDM, we improve the novelty score from 65 .
7 to 68 . 8 .
This implies that ourintroduced energy does not hurt the novelty when the statistics are estimated over the training dataset.
Notice that although the GDM and EN-Flow achieve a better novelty score, the sample quality ismuch worse.
The reason is that, due to the metric deﬁnition, low-quality out-of-distribution sampleslead to high novelty scores.
(3) On the GEOM-DRUG dataset, the atom stability is improved from81 .
3 to 82 . 4 , which shows that our method can work for macro-molecules.
(4) We visualize andqualitatively 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 comparableresults on our tested benchmarks.
The computational load added by introducing prior bridges issmall.
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 timesteps in Table 2.
We observe that (1) with fewer time steps, the baseline EGM method gets worseresults than 1000 steps in Table 1.
(2) with 500 steps, our method still keeps a consistently goodperformance.
(3) with even fewer 50 or 100 steps, our method yields a worse result than 1000 stepsin 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, andcompare them on the QM9 dataset.
(1) We notice that our energy (7) gets better performance withlarger
k
when k  5 .
k = 7 achieves comparable performance as k = 5 .
Larger k also requiresmore 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 forcesoutperforms 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 onthe diffusion model for point cloud generation introduced by point cloud diffusion model [25] andcompare 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 pointcloud generation works [5, 25, 48, 49].
We construct the point clouds following the setup in [25],split the train, valid and test dataset in 80% , 15% and 5% and samples 2048 points uniformly on themesh surface.
Evaluation Metric.
We evaluate the generated shape quality in two aspects following the previousworks, including the minimum matching distance (MMD) and coverage score (COV).
These scoresare 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 exactlythe same experiment setup conﬁguration introduced in [25].
Thus, we use the same model architecture
[Figure]
and train the model in 100 diffuse steps with a learning rate 2 ⇥ 10  3 , batch size 128, and linearnoise schedule from 0 .
02 to 10  4 .
We initial ↵ with 0.1 and jointly learn it with the network.
For thesecond setup, to evaluate the better converge speed of our method, we decrease the diffuse step from100 to 10 with other settings the same.
For the diffusion model baseline, we reproduce the number bydirectly using the pre-trained model checkpoint and testing it on the test set provided by the ofﬁcialcodebase.
Result.
We show our experimental result in Table 4.
We see that (1)
In the 10 steps setup, all variantsof our approach are clearly stronger than the diffusion model.
With force added, our method withphysical 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 thediffusion model baseline.
(3) Since the test points are uniformly sampled on the surface, a better scoreindicates 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 tosome outlier position in sample generate samples, while statistic gap energy is more robust.
We alsoshow visualization samples in Figure 4.
[Table]
47.59 Bridge 13. 47.13 + Riesz 12. 47.42 + Statistic 12. Airplane
Diffusion [25] 3. 46.38 Bridge 3. 46.17 + Riesz 3. 46.23 + Statistic 3.
Conclusion and Limitations
We propose a framework to inject informative priors into learning neural parameterized diffusionmodels, with applications to both molecules and 3D point cloud generation.
Empirically, wedemonstrate that our method has the advantages such as better generation quality, less sampling timeand easy-to-calculate potential energies.
For future works, we plan to 1) study the relation betweendifferent types of forces for different domain of molecules, 2) study how to generate valid proteins inwhich the number of atoms is very large, and 3) apply our method to more realistic applications suchas antibody design or hydrolase engineering.
In both energy functions in (7) and (6), we do not add torsional angle related energy [18] mainlybecause it is hard to verify whether four atoms are bonded together during the stochastic process.
Weplan 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 previousdiffusion models [26], it takes a long time to train a model.
We attempted to speed the training up byusing a large batch size ( e.g. , 512, 1024) but found a performance drop.
An important future directionis to study methods to distribute and accelerate the training.