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 fixed 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 high-quality molecule 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.
Introduction
As exemplified by the success of AlphafoldV2 [16] in solving protein folding deep learning techniques have been creating new frontiers on molecular sciences [38].
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. , 1, 2, 19].
Recently, diffusion-based generative model have been applied to molecule generation problems [6, 13] 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 is to efficiently incorporate strong prior information to reflect the physical and problem-dependent statistical properties of the problems at hand.
In fact, a recent fruitful line of research [8, 18, 29] have shown promising results by introducing inductive bias into the design of model architectures to reflect physical constraints such as SE(3) equivariance.
In this work, we present a different paradigm of prior incorporation tailored to diffusionbased generative models, and leverage it to yield substantial improvement in high-quality and stable molecule generation.
Our contributions are summarized as follows.
Prior Guided Learning of Diffusion Models.
We introduce a simple and flexible 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 fixed 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. , 6, 11, 21, 11], our method modifies 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.
Related works
Diffuse Bridge Process.
Diffusion-based generative models [12, 32, 33, 36] 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 [4, 6, 37] have also been proposed for learning diffusion generative models that guarantee to output desirable outputs in a finite time interval, but these methods involve iterative proportional fittings and are computationally costly.
Our framework of learning generative models with diffusion bridges is similar to that of [27], which learn diffusion models as a mixture of forward-time diffusion bridges to avoid the time-reversal technique of [35].
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. 22, 24, 30, 31, 40, 41, 42] 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. , 10, 13, 20, 29, 39] 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 [10] auto-regressively generates the atom position and type one by one after another; EN-Flow [29] and EDM [13] adopt
E(n) equivariant graph neural network (EGNN) [29] to train flow-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.
Method
We first introduce the definition 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 example.
3.1 Learning Diffusion Generative Models with Prior Bridges
Problem Definition.
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 definition covariance coefficient; 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 ∈ [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 infinite 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 ∈ 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) reflect 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 first 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: Q
[Equation]
A key property of Q Π ∗ is that its terminal distribution equals the data distribution, i.e., Q Π ∗
= Π ∗ .
Therefore, we can learn the diffusion model P θ by fitting 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 [25], the loss function L ( θ ) can be reformed into a form ofdenoised score matching loss of [ e.g. , 33, 35, 34]:
[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 ∈ [0 , + ∞ ) and β t = (cid:82) 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 infinite 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 [35], 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 [27], 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 [7].
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 reflect 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 flexibility of incorporating rich prior information.
To achieve this, we first 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 modification to Brownian bridges.
Definition 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 ∀ z ∈ R d
and U 1 ( z ) =
if and only if z ∈ A .
Intuitively, a diffusion process Q is a bridge A , i.e., Q ( Z 1 ∈ A )
= 1 , if it (at least) approximately follows the gradient flow of a Lyapunov function and the magnitude (or step size) or the gradient flow should increase with a proper magnitude in order to ensure that Z t ∈
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 flow of U and ν is an extra perturbation term.
The step size α t should increase to infinity as t → 1 sufficiently 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 = 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 ) ∈ C 2 ( R d ) and U ( z, · ) ∈ C 1 ([0 , 1]) .
Then, Q A in (4) is an bridge to A , i.e., Q A ( Z 1 ∈ A )
= 1 , if the following holds:
1) U follows an (expected) Polyak-Lojasiewicz condition: E Q A [ U t ( Z t )] [Equation] z [Equation]
( Z t ) [Equation] [Equation] 0 , ∀ t. 2) Let β t =
E Q A [ ∇ z U t ( Z t )
⊤ ν t ( Z t )] , and γ t =
E Q A [ ∂ t U t ( Z t )
+ 12 tr( ∇ 2 z U t ( Z t ) σ 2 t ( Z t ))] , and ζ t =
exp( (cid:82) t
α s d s ) .
Then lim t ↑ 1 ζ t = + ∞ , and lim t ↑
[Equation]
Brownian bridge can be viewed as the case when U t ( z ) =
∥ x − z ∥ 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 [ ∥ f t ( Z t ) ∥ 2 ] < + ∞ and σ t > 0 , ∀ t , which is very mild condition and is satisfied for most practical functions.
The intuition is that the Brownian drift σ 2 t x − Z t β 1 − β t is singular and grows to infinite as t approaches 1 .
Hence, introducing an f into the drift would not change of the final 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-defined parameter.
Please refer to algorithm 3.2 and Figure 1 for descriptions about our practical algorithm.
Algorithm 1 Learning diffusion generative models.
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.
Informative physical or statistical priors that reflects the underlying real physical structures can be particularly beneficial for molecule generation as we show in experiments. [Figure]
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 ∈ R 3 is the coordinate of the i -th atom, and x hi ∈ { 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 final result or have computations that depend on atom types (e.g., calculating an energy function as we do in sequel).
Specifically, for a continuous x hi ∈ 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 ( · ) = −∇ 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-defined energy or force [22, 42].
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 [9] is a family of force fields for molecule simulation.
It is designed to provide a computationally efficient 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 ) = (cid:80) ij ∈ bond ( x ) (Len( x rij ) − ℓ (ˆ x hi , ˆ x hj )) 2 , where Len( x rij ) =
(cid:13)(cid:13) x ri − x rj (cid:13)(cid:13) , 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 ) = (cid:80) ijk ∈ 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 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 defined by E LJ ( x ) =
(cid:80) i ̸
=
j e ( (cid:13)(cid:13) x ri − x rj (cid:13)(cid:13) ) 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 ) = κ (cid:80) ij q (ˆ x hi ) q
(ˆ x hj ) /
(cid:13)(cid:13) x ri − x rj (cid:13)(cid:13) , 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. , 5, 15, 23].
We propose to design a prior energy function in bridges by directly calculate these statistics over the dataset.
Specifically, we assume that the lengths and angles of each type of bond follows a Gaussian distribution that we learn from the dataset, and define the energy function as the negative log-likelihood:
E stat ( x ) = (cid:88)
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [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 ∈ [ k ] , ˆ µ rc and ˆ σ 2 rc denotes empirical mean and variance of length of rc -edges in the dataset; for each triplet r, c, r ′ ∈
[ k ] , ˆ µ rcr ′
and ˆ σ 2 rcr ′ is the empirical mean and variance of angle betwen rc and cr ′ 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.
Experiment
[Table]
We verify the advantages of our proposed method (Bridge with Priors) in several different domains.
We first compare our method with advanced generators ( e.g. , diffusion model, normalizing flow, etc.) on molecule generation tasks.
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 efficiency and effectiveness of our bridge processes and physical energy, we conduct experiments on molecule and macro-molecule generation experiments.
We follow [21] 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 [13, 29], 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 [17] novelty score.
we extracted 10, samples to calculate the above metrics.
Dataset Settings QM9 [28] 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 [13] to split the train, validation, and test partitions, with 100K, 18K, and 13K samples.
GEOM-DRUG [3] 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 [13, 29], we retain thelowest energy conformations for each molecule.
Training Configurations.
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 configuration.
We use the polynomial noise schedule used in [13] which linearly decay from 10 − 2 /T [Table]
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) [13] and E(3) Equivariant Normalizing Flow (EN-Flow) [29] as our baselines.
Both two are trained with the same configurations as ours.
[Figure] [Figure]
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.
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 definition, 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.
[Table]
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 efficiency.
(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.
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 which the number of atoms is very large, and 3) apply our method to more realistic applications such as antibody design or hydrolase engineering.
In both energy functions in (7) and (6), we do not add torsional angle related energy [14] 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 [21], 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.
A Proofs
Proof of Proposition 3.3.
It is a direct result of Theorem A.1.
Theorem A.1.
Assume
[Equation]
We have Z 1 ∈ A with probability one if there exists a function U : R d × [0 , 1] → R such that 1) U ( · , t ) ∈ C 2 ( R d ) and U ( z, · ) ∈ C 1 ([0 , 1]); 2) U ( z, 1) ≥ 0 , z ∈ R d , and U ( z, 1) = 0 implies that z ∈ A , where A is a measurable set in R d ; 3) There exists a sequence { α t , β t , γ t : t ∈ [0 , 1] } , such that for t ∈ [0 , 1] ,
[Equation]
4) Define ζ t = exp( (cid:82) {+t 
+}
α s d s ) .
We assume
[Equation]
Proof.
Following d Z t = η ( Z t , t )d t + σ ( Z t , t )d W t , we have by Ito’s Lemma, d U
[Equation]
Z t , t ) = ∇ U
( Z t , t ) ⊤ ( η ( Z t , t )d t + σ ( Z t , t )d W t ) + ∂ t U
( Z t , t )d t + 1 2tr( ∇ 2 U ( Z t , t ) σ 2 ( Z t , t ))
d t, for t ∈ [0 , T ] .
Taking expectation on both sides,
[Equation]
Let u t =
E [ U ( Z t , t )] .
By the assumption above, we get
[Equation]
Following Grönwall’s inequality (see Lemma A.2 below), we have E
[ U ( Z 1 , 1)] =
u 1 = lim t ↑ 1 u t ≤ 0 if (8) holds.
Because U ( z, 1) ≥ 0 , this suggests that U ( Z 1 , 1) =
and hence Z 1 ∈ A almost surely.
Lemma A.2.
Let u t ∈ R and α t , β t ≥ 0 , and dd t u t ≤ − α t u t + β t , t ∈ [0 , T ] for T > 0 .
We have
[Equation]
, where ζ t = exp( (cid:90) t
α s d s ) .
Therefore, we have lim t ↑ T u t ≤ 0 if
[Equation]
Proof.
Let v t =
ζ t u t , where ζ t = exp( (cid:82) t
α s d s )
so ˙
ζ t = ζ t α t .
Then
[Equation]
So
[Equation]
and hence
[Equation]
To make lim t ↑ T u t ≤ 0 , we want
[Equation]
Corollary A.3.
Let d Z t = x − Z t 1 − t + ς t d W t with law Q .
This uses the drift term of Brownian bridge, but have a time-varying diffusion coefficient ς t ≥ 0 .
Assume sup t ∈ [0 ,T ] ς t < ∞ .
Then Q ( Z 1 = z ) = 1 .
Proof.
We verify the conditions in Theorem A.1.
Define U ( z, t ) =
∥ x − z ∥ 2 / 2 , and η ( z, t )
= x − Z t 1 − t .
We have η ( z, t ) ⊤ ∇ U ( z, t ) =
− U ( z, t ) /
( T − t ) .
So α t = 1 / ( T − t ) .
Also, ∂ t U ( z, t ) + 12 tr( ς 2 t ∇ 2 z U ( z, t ))
= 12 diag( ς 2 t
I d × d ) =
d 2 ς 2 t := β t ≤ C < ∞ .
Then ζ t = exp( (cid:82) t
α s d s ) = 11 − t → + ∞ as t ↑ T .
Also, (cid:82) t
ζ s β s d s ≤ C (cid:82) t
ζ s d s =
CT
(log( T ) − log( T − t )) .
So
[Equation]
Using Girsanov theorem, we show that introducing arbitrary non-singular changes (as defined below) on the drift and initialization of a process does not change its bridge conditions.
Proposition A.4.
Consider the following processes
[Equation]
Assume we have KL ( µ 0 || ˜ µ 0 ) <
+ ∞ and E Q [ (cid:82) T
∥ f t ( Z t ) ∥ 2 ] < ∞ .
Then for any event A , we have
Q ( Z ∈ A )
=
if and only if ˜ Q ( Z ∈ A ) = 1 .
Proof.
Using Girsnaov theorem [26], we have
[Equation]
Hence, we have KL ( Q || ˜ Q ) <
+ ∞ .
This implies that Q and ˜ Q has the same support.
Hence
Q ( Z ∈ A )
= 1 iff ˜ Q
( Z ∈ A )
= 1 for any measurable set A .
This gives an immediate proof of the following result that we use in the paper.
Corollary A.5.
Consider the following two processes:
[Equation]
Assume E Q x, bb ,f [ ∥ f t ( Z t ) ∥ 2 ] < + ∞ and σ t > 0 for t ∈ [0 , + ∞ ) .
Then Q x, bb ,f is a bridge to x .
B Model Details B.1 Model Architecture for Molecule Generation.
Following EGM [13], we apply an E(3) equivariant GNN network (EGNN) as our basic model architecture.
EGNNs are a type of graph neural networks that satisfies the equivariance constraint,
[Equation] R x [Equation] t
[Equation]
h [Equation] when x ′ , h ′ =
f ( x, h ) , (9) where x and h represent the 3D coordinates and additional features, orthogonal R stands for the random rotation and t ∈ R 3 is a random transformation.
One EGNN is usually made up of multiple stacked equivariant graph convolutional layers (EGCL), and every EGCL satisfies the equivariance constraint.
Denote N the number of nodes, x l and h l the coordinates and features for layer l ∈ { 0 , · · · , L } , we have
[Equation]
where h
= h, x 0 = x , d ij = ∥ x
li − x lj ∥ 2 , d ij + 1 is introduced to improve training stability, and ϕ e , ϕ h , ϕ x represents fully connected neural network with learnable parameters.
We refer the readers to the previous paper [29] for more details.
Scaling Features Following [13], we re-scale the data with additional scaling factors.
The atom type one-hot vector and atom charge value × .
25 and × 0 .
1 , respectively.
It significantly improves performance over non-scaled inputs, e.g. 47% relative improvements on molecule stability.