Unified Generative Modeling of 3D Molecules with Bayesian Flow Networks
Authors: Yuxuan Song
Jingjing Gong
Hao Zhou
Mingyue Zheng
Jingjing Liu
Wei-Ying Ma
Venue: ICLR 2024 oral (2024)
Keywords: Drug Design
Molecule Generation
Deep Learning
Computational Biology
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: A new 3D molecule generative model based on Bayesian Flow Networks

Abstract: Advanced generative model (\textit{e.g.}, diffusion model) derived from simplified continuity assumptions of data distribution, though showing promising progress, has been difficult to apply directly to geometry generation applications due to the \textit{multi-modality} and \textit{noise-sensitive} nature of molecule geometry. 
This work introduces Geometric Bayesian Flow Networks (GeoBFN), which naturally fits molecule geometry by modeling diverse modalities in the differentiable parameter space of distributions. GeoBFN maintains the SE-(3) invariant density modeling property by incorporating equivariant inter-dependency modeling on parameters of distributions and unifying the probabilistic modeling of different modalities. 
Through optimized training and sampling techniques, we demonstrate that GeoBFN achieves state-of-the-art performance on multiple 3D molecule generation benchmarks in terms of generation quality (90.87\% molecule stability in QM9 and 85.6\% atom stability in GEOM-DRUG\footnote{The scores are reported at 1k sampling steps for fair comparison, and our scores could be further improved if sampling sufficiently longer steps.}). GeoBFN can also conduct sampling with any number of steps to reach an optimal trade-off between efficiency and quality (\textit{e.g.}, 20$\times$ speedup without sacrificing performance).

Published as a conference paper at ICLR 2024
UNIFIED GENERATIVE MODELING OF 3D MOLECULES VIA BAYESIAN FLOW NETWORKS
Yuxuan Song" " Jiigjing Gong', ; Hao Zhou', Mingyue Jingging Liu' & Wei- Wei-Ying Ma
1 Institute of AI Industry Research (AIR), Tsinghua University
2 Shanghai Institute of Materia Medica, Chinese Academy of Sciences songyuuun,, gong j ng ing,, zhouhao, www.wywwii... ying) @air eduuhhuua.edu
ABSTRACT
Advanced generative model (e.g., diffuuson model) derived from simplified conti- nuity assumptions of data distribution, though showing promising progress, has been difficult to apply directly to geometry generation applications due to the multi- modality and noise-sensitive nature of molecule geometry. This work introduces Geometric Bayesian Flow Networks (GeoBEN), which naturally fits molecule geometry by modeling diverse modalities the differentiable parameter space of distributions. GeoBFN maintains the SE-(3) invariant density modeling property by incorporating equivariant inter-dependency modeling on parameters of distributions and unifying the probabilistic modeling of different modalities. Through optimized training and sampling techniques, we demonstrate that GeoBEN achieves state-of- the-art performance on multiple 3D molecule generation benchmarks in terms of generation quality (90.87% molecule stability in QM9 and 85.6% atom stability GEOM-DRUG], GeoBFN can also conduct sampling with any number of steps to reach an optimal trade-off between efficiency and quality (e.g., 20 speedup without sacrificing performance).
INTRODUCTION
Molecular geometries can be represented three-dimensional point clouds, charac- terized by their Cartesian coordinates space and enriched with descriptive fea- tures. For example, proteins can be repre- sented proximity spatial graphs (Jing) et al. 2021) and molecules as atomic graphs 3D (Schütt al.] 2017). Thus, learning geometric generative models has the potential to benefit scientific discover- ies such as material and drug design. Re- cent progress in deep generative model- ing has paved the way for geometric gen- erative modeling. For example, Gebauer et al. 2019); Ji (2021) and Satorras et al. (2021a) use autoregressive models and flowebased models, respectively, for generating 3D molecules in-silico. Most recently, inspired by the huge success of diffusion model (DM) in image genera-
Recese Dittbution
Ruceiver Deatbbtion
Oy ohannel
Geas Notty channel
Pof-le3
Po
EOMN
EGNN
Pyt (9:)
Bayeckan
Bayesion Update
Upsanne
Sender
Sender
Sundir Distibution
Bundel Distibution
Figure The framework of GeoBFN
tion Meng al. (2022); [o et (2020)
and beyond Li et al. (2022). DM incorporating geometric symmetries has been widely explored in the field of geometry generation Hoogeboom al. (2022): Xu al. (2023).
Equual Contribution. Correspondence to Zhou(zhouhao @ dii..ssinguua.ddu
'The scores reported at 1k sampling steps for fair comparison, and our scores could be further improved if sampling sufficiently longer steps

---

Published as a conference paper at ICLR 2024
However, two major challenges remain in directly applying DM to molecule geometry: multi-modalle and noise sensitivity. The multi monliiimodalityy issue refers to the dependency on diverse data forms to effectively depict the atomic-level geometry of molecule. For instance, the continuous variable of atom coordinates is essential for describing the spatial arrangement, while either the discretised atom charge or categorical atom types are employed to completely determine the molecule's composition. Noise sensitivity refers to the fact that applying noise or perturbing the atom coordinates will not only change the value of the variable but also have significant impact on the relationship among different atoms as the Euclidean distances are also changed. Therefore, a small noise on atom coordinates could bring a sudden drop of the signal at the molecule level.
To alleviate these issues, Xu et al. (2023) introduces latent space for alleviating the inconsistency
of unified Gaussian diffusion on different modalities. Anand & Achim (2022) propose to use decomposed modeling of different modalities. [Peng et al. (2023) use different noise schedulers for different modalities accommodate noise sensitivity. However, these methods either depend on the sophisticated and artifact-filled design or lack of guarantee or constraint on the designed space.
In this work, we propose Geometric Bayesian Flow Networks (GeoBFN) to model 3D molecule geometry in a principally different way. Bayesian Flow Networks Graves et al. (2023) (BFN) is a novel generative model developed quite recently. Taking a unique approach by incorporating Bayesian inference modify the parameters of collection of independent distributions, brings fresh perspective to geometric generative modeling. Firstly, GeoBFN uses a unified probabilistic modeling formulation for different modalities in the molecule geometry. Secondly, regarding the variable of 3D atom coordinates, the input variance for BFNs is considerably lower than DMs, leading to better compatibility with the inherent noise sensitivity. Further, we bring the geometry symmetries into the Bayesian update procedure through an equivariant inter-dependency modeling module. We also demonstrate that the density function of implied generative distribution is SE-(3) invariant and the generative process of iterative updating is roto-transian iito-ranslational equivariant. Thirdly, with BFN's powerful probabilistic modeling capacity, 3D molecule geometry representation can be further optimized into representation with only two similar modalities: discretised charge and continuous atom coordinates. The mode-redundancy issue on discretised variable in the original BFNs is fixed by early mode-secking sampling strategy in GeoBFN.
With operating on the space with less variance, GeoBEN could sample with any number of steps which provides a superior trade-off between efficiency and quality, which leads to 20 speedup with competitive performance. Besides, GeoBEN is a general framework that can be easily extended to other molecular tasks. We conduct thorough evaluations of GeoBEN multiple benchmarks, in- cluding both unconditional and property-conditional molecule generation tasks. Results demonstrate that GeoBEN consistently achieves state-of-the-an generation performance on molecule stability and other metrics. Empirical studies also show significant improvement in controllable generation and demonstrate that GeoBEN enjoys a significantly higher modeling capacity and inference efficiency.
2 PRELIMINARIES
2.1 SE-(3) INVARIANT DENSITY MODELING
To distinguish geometry representation and the atomic property features, we use the tuple (x, h) to represent the 3D molecules. Note here x = (xt,... (x?,...,>> x*) E R x3 the atom coordinate matrix, and h = (h2 ...., h E R xd is the node feature matrix, e.g., atomic types and charges. Density estimation on the 3D molecules should satisfy specific symmetry conditions of the geometry. this work, we focus the transformations T, in the Special Euclidean group (SE-(3)), i.e., the group of rotation and translation in 3D space, where transformations T, can be represented by a translation t and an orthogonal matrix rotation R. Note for generative model molecule geometry with underlying density function pe( (x, h)), the likelihood should not be influenced by the rotation or translation of the entire molecule, which means the likelihood function should be SE--(3 invariant on the input coordinates, i.e., Po (x, h)) Pe ((Rx + t, h)).
2.2 BAYESIAN FLOW NETWORKS
The Bayesian Flow Networks (BFNs) are based on the following latent variable models: for learning the probability distribution po over g, series of noisy versions (",,"* yn,) of g introduced as
2

---

Published as a conference paper at ICLR 2024
latent variables. And then the variational lower bound of likelihood is optimized:
PΦ (g
y1..... Yn) Ps (Y1 (YI..... .n))
log (g) 2 Y1 Yum
log
v........... Nn g)
DD (yi...., ......))))))) + Ya....aa
log Ip (g Vi..... v.)l
(I)
And q is namely the variational distribution. The prior distribution of latent variables is usually organized autoregressively, L.c., Po Py .Y)) po(y1)po(y2 (ynn Yn-1- y1) which also implies the data generation procedure, i.e., Y1 + ... + Yn - g (Note: this procedure only demonstrates the generation order, yet does NOT imply Markov property for the following derivative).
One widely adopted intuition for the generation process is that the information of the data samples should progressively increase along with the above Markov chain, e.g., noisier images to cleaner images. The key motivation of BFNs is that the information along the latent variables should change as smoothly as possible for all modalities including discretized and discrete variables. To this end, BFNs operate on the distributions in the parameter space, in contrast to the sample space.
We introduce components of BFNs one by one (Fig 2a Firstly, the variational distribution q is defined by the following form:
Ps (y: g:0:)
ps (y; g; ags) is termed as the sender distribution, which could be seen as adding noise to the data according to a predefined accuracy Oig.
Secondly, for the definition of pa. BFNs will first transfer the noisy sample y to the parameter space, obtaining 0, then apply Bayesian update in the parameters space and transfer back to the noisy sample space at last. To clarify, 0 refers to the parameter of distributions in the sample space, e.g., the mean/variance for Gaussian distribution probabilities for categorical distribution. In the scope of BFNs, the distributions on the sample space are factorized by default, egg., p(g 0) I P (g(()) ((()).
Thirdly, a neural network & takes 0 as input and aims to model the dependency among different dimensions hence recover the distribution of the original sample g. The output of neural network (()) still lies in the parameter space, and we termed it as the parameter of output distribution po- where po (y () = - (y (0))
To map the noisy sample y to the input space, Bayesian update is applied to A:
0, + (..
(3)
h is called Bayesian update function. The distribution over (On...., On-1) is then defined by the Bayesian update distribution via marginalizing out y:
Ps (Op. 0n-1) = ((o) II pa (0: 0;-1;a4),
(4)
where p(On) is a simple prior for ease of generation, e.g., standard normal, and pa could be obtained from Eq.
Pe (0: 0;-1;0s) PR pn(y:10 8 being the Dirac delta distribution. PR(Y: 0,,1. the receiver distribution.
E
5 (0. h(O h(),),,
E po
ps(y, and is also called
0:-1:0)
At last we map (()) back to the noisy sample space by combining the known form, accuracy of Ps and marginalizing out y:
n
n II
Po (y1.....y.. Vn) po(r)) II Po Y(1x-1)
poly, 02-1)
E 1= po
(6)
150)
3

---

Published as a conference paper at ICLR 2024
h1
2/2
U1
yn
y1
y2
yi
Un
P
Graphical Model of BFN
b Graphical Model of Diffusion
Figure 2: Graphical View Comparison between BFN and Diffusion
where we use 0o:n-1 to abbreviate (fo.... 0n-1), and similar. Till now, we have defindd po(y....... .n), and pe (g :n....; simply po(g 0,)) on each sample, thus Eq can be estimated.
3 METHODOLOGY
3.1 SE-(3) INVARIANT GEOMETRY DENSITY MODELING
As discussed in Sec. 2.1 for generative model on the 3D molecule geometry, it is crucial to hold the SE-(3) invariant conditions. Recall the mathematics formula of the geometries = (x. h). we denote the latent variable, e.g. noisy samples, of g as y We are interested applying the SE-(3) invariant conditions to the probabilistic model Pş- To this end, we need to first reformulate the likelihood function:
Pa(g) Pp((x,)))
Polg yi. vA)Po dy)dy ... dy m
With H and y and the distributions defined the same way as above, we focus on the variables corresponds to x in the geometry ya. Then we have the following theorem:
Theorem 3.1. (SE-(3) Invariant Condition)
With the 0z, y", constrained in the zero Center of Mass(CoM) space (Köhler et al. 2020 Xu et al. 2022), the likelihood function Pá is translational invariant.
When the following properties are satisfied, the likelihood function Pe is roto-invariant:
Po (x' 1;)0) Po (R(x') R(0) Ps (y xx:a)= ps (R(y") R(x)):))
h(R(6 R(y)),
as): p(x' (0)) = (((()))005) y orthogonal R.
Proposition 3.2. With the condition Theorem. 3.1 satisfied, the evidence lower bound objective in Eq. i.e.
LvlB(x)= =
DKL (ps (: x; x;) [|pn( 0 1;)))))) loo (x 0
(8)
Po (05......
with the Bayesian update distribution Ps (Oi...., 0=-1) = IT= Pu 0 =
x; ax) similar to 03)], P x 0%) = -x Φ). Derivation
Eq. And An(. 00
(x OF
from Eq. to equation 8 is at Appendix then LVLB(X x) ) also SE-(3) invariant.
C.4 Pu (0: 0;-1, pu (RO: RO; Rx;
We leave the formal proof of Theorem.
3.1 and Proposition.
3.2 in Appendix
3.2 GEOMETRIC BAYESIAN FLOW NETWORKS
Then we introduce the detailed formulation of geometric Bayesian flow networks (GeoBFN) based on the analysis See. 3.1 For describing 3D molecule geometry g (x, h), various representations can be utilized for the node features. The atom types h, and atomic charges he are commonly employed, with the former being discrete (categorical) and the latter being discretized (integer). Together with the continuous variable, e.g., atom coordinates x, the network module in the modeling of the output distribution of GeoBEN could be parameterized with an equivariant graph neural network (EGNN) (Satorras et al. 2021b) 4:
(RO0
ht ghol) [R t, 0%0% 0%], VR,t
4

---

Published as a conference paper at ICLR 2024
where ((0", [0 1* 10" 0% 0%... And then introduce the necessary components to derive the objective in Eq. 8
Atom Coordinates and Charge he: For the continuous and discretized variables, the input
distribution is set as the factorized Gaussian distributions, where O del {1.p) the parameter of N (-1 pl P.P For simplicity, we take x as an example to illustrate the common parts of the two variables. And o is set as {0,1}, The sender distribution ps is also an isotropic Gaussian distribution:
ps(. x; eI) =N (x,a
(10)
Given the nice property of isotropic Gaussian (proof given by Graves et al. (202))) the simple form of Bayesian update function could be derived as:
11-10P-1 + ya
((Hi .y, {pe Here = Pi-1
+ aa, P
(11)
As shown in Eq. the randomness only exists in p. and the corresponding Bayesian update distribution in Eq. is as:
+ Mi-1Pi-1
Pu (0, (1-1, =N
(12)
Ps
pt
The above discrete-time Bayesian update could be easily extended continuous-time, with an accuracy scheduler defined as 3(t) = -e-- a (t') dt', E [0, Given the accuracy additive property of pu (proof given by Graves et al. (2023)
Pu (0.-1|0.-2,:;:..) Pu (0; 0,-1, 04)
E
pu (0: 0z-2, aa + an)), the Bayesian flow distribution could be obtained as:
pr x; pu (0" 0g. x: B()))
(13)
The key difference atom coordinates x and charges he lies in the design of the output distribution. For continuous variable the network module & directly outputs an estimated & (009, t). Hence for timestep t, the output distribution is
Po (x'
8 t; 0) = 6(x (009, t)
(14)
While for discretized variable he the network module will output two variables, Ph. and In with dimension equivalent he which implies distribution ((H,, II. I). With K-bins discretized variable, the support is split into buckets with each bucket & centered as kr 2 1 and left K
boundary as k; = ke K and right boundary as kū = ke + K- Then for each k, the probability is the mass from k to ky, i.e,, St. N(uh,, oi, I). And the first and last bins are curated by making sure the sum of the probability mass is 1. Then the output distribution
D II pod (b(((.) d=1
(b 0 0)))
0 t0))
(15)
where the function k(-) maps the variable to the corresponding bucket.
Atom Types h,: The atom types h, are discrete variables with K categories, where the corresponding parameter space lies in probability simplex thus the procedure slightly different from the others. The input distribution for h, is pr( 0) = I12_, (h.(d), where is number if variables. And the input prior 0h K: where 1 is the length KD vector whose entries are all K. The sender K
distribution, could be derived with the central limit theorem, lies the form of
ps(y h;:a) ar) = =N (y a (Ken, - -1,,
(16)
where a vector ones, I the identity matrix, and e, E RK is vector defined as the projection
from the class index to a length K one-hot vector (proof given by Graves al. (2023) ). In other
words, each element of ej is defined as (e,) = 5,k, where dsk is the Kronecker delta function. And
del' ( 0 E ERKD.
en,
5

---

Published as a conference paper at ICLR 2024
GesBEN
FDM
Figure The Bayesian Flow and Diffusion Process of GeoBFN and EDM.
The Bayesian update function could be derived as h (0 ΣK ( (proof given by
Graves al. (2023)
And similar to Eq. 13 the Bayesian flow distribution for h, is as:
Pr (0
h;;
8(0 softmax(y
(17)
With the network module 4 the oulput distribution could be obtained as
Po (A: (); = (softmax
Phe 0;1) = II (1!a) 00;
po (d)
(18)
Training Objective: By combining the different variables together, could obtain the unified continuous- caanuuus-timee/t loss for GeoBEN based on Eq. 25 to Eq. 41 in (Graves et al. 2023) as:
(t)
- gg) = L
L ((x, he: h₂)) = p0
g ((09,t)]²
a
a he
a (t) 2
E
x xx
lh $
lh T 11?
(19)
(-).)
2
2
Where &. is short for & (09, l). The joint Bayesian flow distribution is decomposed as:
py (8 (
g:t) (0 xx (f he: h;;
(20)
with a=, ah and h refer to the corresponding accuracy scheduler (details provided by Graves et al.) (2023)). And t defined the same as in Eq. [4] while 4 de- fined by the weighted average of different bucket centers with the output distribution in Eq. 15 as EE (,tkk, Ek (D) (l: 0, 1)ke: And for th is defined as the
Ek Po
(d) (A: 0,1)ex based on Eq. 18
Remark 3.3. The GeoBEN defined the above formulation satisfied the SE(3) condition in
Theorem. 3.1
Sampling GeoBEN will generate samples follow the graphical model the recursive procedure as illustrated in Fig. 2a e.g. g' ~ p0(-)0--) - ps(- a)) + 0, h(O, 1,
3.3 OVERCOME NOISE SENSITIVITY IN MOLECULE GEOMETRY
One key obstacle of applying diffusion models to 3D molecule generation is the noise sensitivity property of the molecule geometry. The property of noise sensitivity seeks to state the fact: When noise is incorporated into the coordinates and displaces them significantly from their original positions, the bond distance between certain connected atoms may exceed the bond length range[1]. Under these circumstances, the point cloud could potentially lose the critical chemical information inherently encoded in the bonded structures; Another perspective stems from the reality that when noise is added to the coordinates, the relationships (distance) between different atoms could alter more rapid pace, e.g. modifying the coordinates of one atom results in altering its distance to all other atoms. Thus, the intermediate steps' structure in the generation procedure of diffusion models the intermediate steps' structure might be uninformative. And the majority of the information being acquired in the final few steps of generation (as depicted in Fig.

---

Published as a conference paper at ICLR 2024
A fundamental belief underpinning GcoBFN is that smoother transformation during the generative process could result in a more favorable inductive bias according to Grrvess et al. 2023). This process occurs within the parameter space of GeoBFN, which is regulated through the Bayesian update procedure. Specifically, samples exhibiting higher degrees of noise are assigned lesser weight during this update (refer to Eq. TTP. This approach consequently leads to significant reduction in variance within the parameter space as (Graves et al.) (023). which in turn facilitates the smooth transformation of molecular geometries. As illustrated in Fig. S this is evidenced by the gradual convergence of the structure of the intermediary steps towards the final structure, thus underscoring the effectiveness of smoother transformation.
3.4
OPTIMIZED DISCRETISED VARIABLE SAMPLING
Previous research (Hoogeboom et al. 2022 Xu et al. 2023 Wu et al. 2022) utilizes both the atom types h, and charges h. to represent the atomic properties. The h. usually serves as an auxiliary loss for improving training which is not involved in determining the molecule graph during generation due to the insufficient modeling. However, there is redundant information between these two variables, since the h, and he variables have one-to-one mapping. e.g.the charge value 4 could be uniquely determined as the Carbon atom. We found that with advanced prohabilistic modeling on discretized data, GeoBFN could conduct training and sampling only with and h However, there exists counterexample for the objective in Eq. 19 and the output distribution during sampling as in Eq. [5] As shown in Fig 5 the boundary condition for clamping the cumulative probability function in the bucket could cause the mismatch, e.g., the true density should be centered in the center bucket while the ouput distribution instead put the most density in the first and last buckets which cause the mode-redumee reduuddedy as shown in upper-left in Fig. [] Though the weighted sum in Eq. 19 is optimized, the sampling procedure will rarely sample the center buckets. And such cases could be non-negligible in our scenarios, especially when the number of bins is small for low dimensional data. To alleviate this issue, we instead update the output distribution in the sampling procedure to:
K Σ
K
(0, NEAREST_CENTEE
(1) Po
1,tte)
Σ
Po
(k P
0 0.t)%o
(21)
Function VEAREST_CENTER compares inputs to the center bins K.
ke (1)
ke kDD)) and return
the nearest center for each input value. The updated distribution unbiased towards the training objective and also reduce the variance during generation which could be found in the trajectory of 5 5
EXPERIMENTS
4.1
EXPERIMENT SETUP
Task and Datasets We focus on the 3D molecule generation task following the setting of prior works (Gebauer et al.) 2019 Luo & Ji 2021 Satorras et 2021a Hoogeboom et al. 2022 Wu et al.) 2022). We consider both Unconditional Molecular Generation which assesses the capability to learn the underlying molecular data distribution and generate chemically valid and structurally diverse molecules and the Conditional Molecule Generation tasks which evaluate the capacity of generating molecules with desired properties. For Conditional Molecule Generation, we implement a conditional version GeoBEN with the details in the Appendix. The widely adapted QM9 (Ramakrishnan et al. 2014) and the GEOM-DRUG (Gebauer et al.] 2019 (021) with large molecules are used for the experiments. And the data configurations directly follow previous work (nndersonoon et al. 2019
Hoogeboom et al. 2022 Xu et al. 2023 2
Evaluation Metries The evaluation configuration follows the prior works (Hoogeboom al. 2022 Wu et al., 2022; Xu et al.] 2023). For the Unconditional Molecular Generation, the bond types are first predicted (single. double, triple. or none) based on pair-wise atomic distance and atom types in the 10000 generated molecular geometries (Hoogeboom et al., 2022). With the obtained molecular graph, we evaluate the quality by calculating both atom stability and molecule stability metrics. Besides, the validity (based on RDKIT) and uniqueness are also reported. Regarding the
"The official implementation http://gthhh..om/AAgggMMeefccnnn

---

Published as a conference paper at ICLR 2024
Table I: Results of atom stability, molecule stability, validity, validity x unniqueness (Vx (V,U), and novelty. A higher number indicates better generation quality, The results marked with an asterisk were obtained from our own tests. And GeoBFN, denote the results of sampling the molecules with specific number of steps
QM9
DRUG
Atom (%)
Mol Sta (9)
Valid (%)
(S6
Novelty (%)
Atom Sta (9)
Valid (%)
Data
99.0
95.2
97.7
97.7
86.5
99.9
ENF
85.0
4.9
40.2
39.4
G-Schnet
95.7
68.
85.5
80.3
77.7
EDM
98.7
82.0
91.9
90.7
58.0
81.3
92.6
EDM-Bridge
98.8
84,6
92.0
90.7
82.4
92.8
GEOLDM
89.4
92.7
57.0
99.3
GEOBEN
98.28 + 0.1
85.11
= 0,5
92.27 0.4
90.72 = 0.3
72.9
75.11
91.00
GROBEN 100
0.3
93.03
70.3
78.89
GEOBFN
98.78 + 0.8
88.42
+ 0.2
93.35 0.2
91.78
0.2
67.7
81.39
93,47
GROBEN 16
99.08
+ 0.06
90.87
+ 0.2
95.31 0.1.1
92.96
0.1
66.4
85.60
92.08
GEOBFN 2/
99.31
0.03
93.32
0.1
96.88
92.41
65.3
86.17
91.66
Table 2: Mean Absolute Error for molecular prop- erty prediction with 500 sampling steps. A lower number indicates a better controllable generation result.
Table 3: Ablation study, GeoBEN models
molecule charge settings, the sampling step set to 1,000.
Property
SHомо mcV
FLUMO
P
Ce
Charge Feature
Atom Stable (96)
Stuble (%)
Units
Bohr"
meV
meV
D
od
discretised basis continuous basis
99.08
90.87
QM
0,10
64
39
36
0.043
0.040
98.97
89. 94
Random
645
1.616
discrete
98.93
88.93
N Naoms
3,86
866
426
813
1,053
1,971
discrete + continuous
98.96
89.33
discrete + discretised
98.91
88.65
GEOLDM
2,37
587
340
522
1.108
1,025
GEOBEN
Conditional Molecule Generation, evaluate conditional version of GcoBFN QM9 with polarizability a. orbital energies θHоMO. ÉLUMO and their gap As, Dipole moment J4, and heat capacity C... Following previous work HHooeehoom al. 2022); Xu et al. (202)), the conditional GeoBFN is fed with range of property s to generate samples and the same pre-trained classifier is utilized measure the property of generated molecule as The Mean Absolute Error (MAE) between 8 and 8 is calculated to measure whether the generated molecules is related to the conditioned property.
6 properties:
Baselines GeoBFN is compared with several advanced baselines including G-Schnet (Gebauer et al. 2019), Equivariant Normalizing Flows (ENF) (Satorras al.) (021a) and Equivariant Graph Diffusion Models (EDM) with its non-equivariant variant (GDM) (Hoogeboom et al. 2022). Also with recent advancements, EDM-Bridge (Wu al. 2022) which improves upon the performance of EDM by incorporating well-designed informative prior bridges and also GeoLDM (Xu et al. 2023) where a latent space diffusion model is applied are both included. To yield a fair comparison, the method- -agostic configurations set the same. The implementation details could be found in Appendix.
Tompue Famebry
1.00
Original sampling algeritim
0.96
0.90
Malecule Slabilty
0.85
0.80
GeoBFN
Improved
sumping
0.75
EDM
alponitio
EDM-Bridge GEOLDM
bound
0.86
1000
2000
3000
4000
Figure 5: 2D Synthetic case of optimized synthetic example. the left columns, generated samples are in orange, and data points are blue.
Sampling Step
Figure 4: QM9 Molecule Stability wrt. Sampling
Steps
8

---

Published as a conference paper at ICLR 2024
4.2 MAIN RESULTS
The results of Unconditional Molecular Generation can be found in Tab. We could observe that in both the QM9 and GEOM-DRUG datasets, GeoBFN achieves new state-of-the-an performance regarding both the quality and diversity of the generated molecules which demonstrates the huge potential of GeoBFN on geometry generative modeling. The phenomenon demonstrates that the GeoBFN does not hold the tendency to collapse to the subset of training data which could imply probabilistic generalization ability and could be useful for several application scenarios; The Conditional Molecule Generation results can be found in Tab. [ GeoBEN consistently outperforms other baseline models by an obvious margin in all conditional generation tasks. This clearly highlights the effectiveness and generalization capability of the proposed methods.
4.3 ANY-STEP SAMPLING
One notable property of GeoBFN is that training with the continuous-time loss, e.g.. Eq. 19 the sampling could be conducted with any steps without incurring additional training overhead. As shown in Tab. GeoBEN could get superior performance compared to several advanced models with only 50 steps during sampling which brings 20 speed-up during sampling due to the benefit of low variance parameter space. As we could find in Fig[4] with the sampling steps increasing from 50 4600, the molecule stability could be further boosted to approach the upper bound, e.g., 94.25% with 4000 steps.
4.4 ABLATION STUDIES
We conduct ablation studies on the effect of input modalities in Tab. 3 We try different compositions and losses to represent the atom types, discretised basis refers to the case where the charge feature is
(x-1z)")
used with discretised and the Gaussian basis, i.e., (x)x) = cxp
is used as functional
20
embedding for charge; continous basis only differ in that the continous loss is utilized. The discrete refers to including the one-hot type representation; type and charge are included while continuous loss included; Similar the contiite- With only discretised variable utilized, the performance is superior to including the discrete variable which implies powerful probabilistic modeling capacity and the benefits of applying similar modality.
refers to both the one-hot
RELATED WORK
Previous molecule generation studies have primarily focused on generating molecules as 2D graphs (Jin et al. 2018 Liu et al. 2018 Shi et al.) 2020), but there has been increasing interest in 3D molecule generation. With the increasing interest in 3D molecule generation, G-Schnet and G- SphereNet (Gebauer et al. 2019 Luo & Ji 2021) respectively, employ autoregressive techniques to create molecules in step-by-step manner by progressively connecting atoms or molccular fragments. These frameworks have also been extended to structure-based drug design (Li et al. 2021, Peng et al., 2022 Powers al. 2022). There are approaches use atomic density grids that generate the entire molecule in single step by producing density over the voxelized 3D space (Masuda et al., 2020). Most recently, the attention has shifted towards using DMs for 3D molecule generation (Hoogeboom et al. 2022: Wu et al. 2022 Peng et al. 2023 Xu et al. (023). with successful applications in target drug generation (Lin et al.
2022), antibody design (Luo et al.
2022), and protein design (Anand & Achim 2022 Trippe al. (002)). However,
method is
based on the Bayesian Flow Network (Graves et al. (023) objective and hence lies a different
model family which fundamentally differs from this line of research in both training and generation.
6 CONCLUSION
We introduce GeoBFN, new generative framework for molecular geometry. GeoBEN operates in a differentiable parameter space for variables from different modalities. Also, the less variance in parameter space is naturally compatible with the noise sensitivity of molecule geometry. Given the appealing property, the GeoBFN achieves state-of-the -ortt--rtttttttt performance on several 3D molecule generation benchmarks. Besides, GeoBFN can also conduct sampling with arbitary number of steps to reach an optimal trade-off between efficiency and quality (e.g., 20x speedup without sacrificing performance).

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The authors thank Yanru Qu for the helpful discussions and proofreading of the paper, well as the anonymous reviewers for reviewing the draft. This work is supported by the National Science and Technology Major Project (2022ZD01 (75202)) Natural Science Foundation of China (62376133) and Guoqiang Research Institute General Project, Tsinghua University (No. 202IGQG1012).
REFERENCES
Namrata Anand and Tudor Achim. Protein structure and sequence generation with equivariant denoising diffusion probabilistic models. arXiv preprint arXiv:2205.15009, 2022.
Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural
networks. Advances in neural information processing systems, 32, 2019.
Niklas Gebauer, Michael Gastegger, and Kristof Schütt. Symmetrr-adapteer generation 3d point sets for the targeted discovery molecules. Advances in neural information processing systems, 32, 2019.
Niklas WA Gebauer, Michael Gastegger, Stefaan Hessmann, Klaus-Robed Müller, and Kristof Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. arXiv preprint (tX.i::200904822. 2021.
Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks. arXiv preprint arXiv:2308.07037, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffissio probabilistic models. arXiv preprint arXiv:2006. 1239, 2020.
Emiel Hoogeboom, Victor Gareia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning, pp. 8867-8887, PMLR, 2022.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. International conference on machine learning, pp. 2323-2332. PMLR, 2018.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference Learning Representations, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3nd International Conference on Learning Representations, 2014.
Jonas Köhler, Leon Klein, and Frank Noe. Equivariant flows: Exact likelihood generative learning for symmetric densities. In Proceedings of the 37th International Conference Machine Learning, 2020.
Xiang Lisa Li, John Thickstun, Ishaan Guirajani, Percy Liang, and Tatsunori Hashimoto. Diffusion- LM improves controllable text generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL /ettps://opereee
net /ttp:/ourmmmmi1111122.1 Ly
Yibo Li. Jianfeng Pei, and Luhua Lai. Structure- based de novo drug design using 3d deep generative models. Chemical science, (341):13664-1360 2021.
Haitao Lin, Yufei Huang, Meng Liu, Xuanjing Li, Shuiwang Ji, and Stan Zi. Diffbp: Generative diffusion of 3d molecules for target protein binding, arXiv preprint arXiv:2211.11214 2022.
Qi Liu, Miltiadis Allamanis, Mare Brockschmidt, and Alexander Gaunt. Constrained graph variational autoencoders for molecule design. In Advances in neural information processing systems, 2018.
10

---

Published as a conference paper at ICLR 2024
Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang. Jian Peng, and Jianzhu Ma. Antigen-specife specific antibody design and optimization with diffusion- based generative models for protein structures. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.). Advances in Neural Information Processing Systems, 2022. URL https:/ /oponeeeeeeeeeeeeeeeeeee htttps://openrevvieee
net. /t///fouuu forum?
[d=jSorrn2TT)
Youzhi Luo and Shuiwang Ji. An autoregressive flow model for 3d molecular geometry generation from scratch. In International Conference on Learning Representations, 2021.
Tomohide Masuda, Matthew Ragoza, and David Ryan Koes. Generating 3d molecular structures con- ditional on a receptor binding site with deep generative models. arXiv preprint arXiv:2010.14442
2020.
Chenlin Meng. Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun- Yan Zhu, and Stefano Ermon. SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2022. URL /ttp://////peeeeee.... www.ewe.eeeeeee..ee...e.eeeee ner/forum?
id=aBsCJcPu_
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Vito,o, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS- W, 2017.
Xingang Peng, Shitong Luo, Jiaqi Guan, Qi Xie, Jian Peng, and Jianzhu Ma. Pocket2mol: Efficient molecular sampling based on 3d protein pockets, In International Conference on Machine Learning. 2022.
Xingang Peng, Jiaqi Guan, Qiang Liu, and Jianzhu Ma. Moldiff: Addressing the atom-bond inconsistency problem in 3d molecule diffusion generation. arXiv preprint arXiv:2305.07508. 2023.
Alexander Powers, Helen H. Yu, Patricia Suriana, and Ron O. Dror. Fragment-band ligand generation guided by geometric deep learning on protein-ligan structure. bioRxiv, 2022. doi: 10. 1101/2022.03.1 17, 17.466633 URL httptttt/////tt 03/21/2022 03. 17. 484653
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Victor Gareia Satorras, Emiel Hoogeboom, Fabian B Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows for molecule generation in 3d. arXiv preprint arXiv:2/05.09016. 2021a.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International conference machine learning, pp. 9323-9332. PMLR, 2021b.
Kristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature communications, 8: 1310, 2017.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: flow-based autoregressive model for molecular graph generation. arXiv preprint a7X1v:2001.09382. 2020.
Brian L Trippe, Jason Yim, Doug Tischer, Tamara Broderick, David Baker, Regina Barzilay, and Tommi Jaakkola. Diffussio probabilistic modeling of protein backbones in 3d for the motif- scaffolding problem. arXiv preprint arXiv:2206.00 041 19, 2022.
Lemeng Wu. Chengyue Gong, Xingchao Liu. Mao Ye, and qiang liu. Diffusion-basee molecule generation with informative prior bridges. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https:: http:://///peeeeeeeeee net forum?
121 TKE
Minkai Xu, Lantao Yu, Yang Song. Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.0292. 2022.
11

---

Published as a conference paper at ICLR 2024
Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, and Jure Leskovec. Geometric latent diffusion models for 3d molecule generation. arXiv preprint arXiv:2305.01140, 2023.
A EXPLANATION OF THE DATA EXCHANGE PERSPECTIVE OF BAYESIAN FLOW NETWORKS
In this section, we provide a brief overview of the Bayesian Flow Networks Graves et al. (2023) from a data exchange perspective. Bayesian Flow Networks (BFNs) new class of generative model that operates on the parameters of a set of independent distributions with Bayesian inference. The meta elements of BFNs are the inpul distributions, sender distributions, and ouiput distributions. To start with, we denote the D-dimensional variable as = (m ,....
mm( "......"))) _MP, and (0()) (0(1) ..... O(D)) represent the parameters of a D-dimensional factorised distribution, i.e.. p(m 0) I P (m(4) ((()).
The basic logic of BFN could better explained by the following communication example between the sender, referred to as Alice, and the receiver Bob. Alice aims to transfer some data to Bod in progressive fashion, i.e., at each timestep, Alice corrupts the data according to some channel noise, and then the noisy sample is transferred. The sender distribution is then defined to describe the noise- adding procedure, which is also factorized distribution, ps(y m; a) = [T Ps (y(() m m a). a refers to the accuracy parameter, a = 0 refers to the information of the sample being totally destroyed by the noise, and with increase the noisy sample will contain more information of the original sample. Intuitively, the sender distribution could be approximately understood as adding noise to each dimension of the data independently.
After receiving the noisy sample Bob will first update an initial "guess" of what is the original sample behind the noisy sample, i.e. input distribution. Note that except for the noisy sample, Bob also knows the accuracy parameter a and noise formulation while not aware of the original sample m. Le., the noise level to create such a sample. The input distribution is initially a simple prior on the data space, e.g., a standard Gaussian, lies in the mean- -real family, (m 0) = -2 Pr (a(()) ((")). The parameter of input distribution will be updated through Bayesian inference, noted as O h (0.-1,y, This update usually lies in a simple form, e.g. additive or weighted average.
After updating the parameter of input distribution, Bob has an "assistant" which will help to provide a better guess on the original data which generates the observed noisy sample. The assistant aims exploit more context information between different dimensions, e.g., the relationship between different pixels in an image, in contrast to updating each dimension independently as in the inpul distribution. Empirically, the assistant could be implemented by a neural network which takes all parameters of input distribution for the prediction of parameters of each dimension, i.e, 4 (v(() (0, t)....,,' (DD) (0, t)).
The output distribution is then implied by the predicted parameter which lies in the formulation of (m | 1, I1 Po Po (m(d) (e)) (0, 1)). Then Bob could construct a distribution to approxi- mate the sender distribution at accuracy by combining the output distribution with the known noise form, accuracy, i.e., PR (- 0; t, ax) = (p (y x'; a). Such distribution is called receiver distribution. The "assistant" of BFNs v is minimize the KL divergence with a defined accuracy scheduler under different timesteps, i.e., DKL (ps (. m; as) ||pR (. 0; 1:-1;t-11, as)), which could also be interpreted transmission cost under the bits-back coding scheme.
B IMPLEMENTATION DETAILS
The bayesian flow network is implemented with EGNNs|Satorras et al. (2021b) by PyTorch (Paszke cet al. 2017) package. We set the dimension of latent invariant features k to 1 for QM9 and 2 for DRUG, which extremely reduces the atomic feature dimension. For the training of vector field network Vo: on QM9, we train EGNNs with layers and 256 hidden features with a batch size 64; and on DRUG, model uses SiLU activations. We train all the modules until convergence. For all the experiments, choose the Adam optimizer (Kingma & Ba 2014) with constant learning rate of 10-4 as
train EGNNs with 4 layers and 256 hidden features, with batch size 64. The
12

---

Published as a conference paper at ICLR 2024
default training configuration. The training on QM9 takes approximately 2000 epochs, and on DRUG takes 20 epochs.
C PROOF OF THEOREMS
In this Section, we provide the formal proof of the Theorem. 3.1 and Proposition. detailed derivations for Equations.
3.2 well as the
C.I DISCUSSION ON THE TRANSLATIONAL INVARIANCE
Remark C.I. It is important to distinguish it from the rotation invariant. The rotational invariant is defined p(x) p (Rz). while the translational is not p(z) p(x + as such distribution can not integrate into one and hence does not exist. Fortunately, the freedom of translation could be eliminated by only focusing
learning distribution on the linear subspace where the center of gravity is always zero. This is, for all configurations on Rax3 space, the density on the zero CoM space is utilized represent their density; It's important to note that the distribution is not defined for configurations outside the zero CoM space. However, it remains possible to leverage the distribution provide a eevautt-evallation (not probability density) on the configurations outside the zero CoM space. This is achieved by projecting them back into the subspace. The evaluation procedure for configurations out of zero CoM space could only get a quantity defined artificially instead of the true density of some real distribution, e.g. It is referred to "CoM-free density" in (Xu et al. 2022). Thus, there does not exist correctness issues.
C.I.1 ZERO CENTER OF MASS(CoM) IN THE GEOBEN
Here provide detailed diseussions optimizing the distribution in the zero CoM space. Recall the training objective in equation 8
LVLB(X)=
E
Σ DKLLDs (: x;);)) |pn[- (2-1ia;) - log Ps (x 0%)
(22)
Po 0:.
where P& (x (x)) = PO P 0 の). For learning distribution on the zero CoM space of Rax3 the ps x: a;), PR 10=1:aa as) and Pa (x (%)) are all defined and supported on the zero CoM space, here Σ X; = 0 and Σ 0= 0. other words, such distribution has no definition for variable E x2x3 if >x vi + O. We then express the likelihood function of an isotropic diagonal Gaussian distribution, which is originally defined on the zero CoM space ((n 1) 3-dimensional,, in the ambient space (n 3-dimensional) as (Hoogeboom et al., 2022]:
N₂. (a uu, 02) = (V2Ta) (n- 1)x3 exp
20 x 1ull?
(23)
Here o² is the variance which is equivalent for each dimension. Recall the Eq. 35 the (Graves et al.) 2023), which shows that (ri(ps (- x;a;) ||pn(- a=););)) takes the form of DRL (N (x, a; 'I) [ t), a²)) which is the KL divergence between diagonal Gaus- sian, then we derive the KL divergence for isotropic diagonal normal distributions of zero CoM with means represent on the ambient space. ps N ( and Pa A (μăg, o²1) on subspace, where fåj and piz (n 1) 3-dimension. Then the KI. between them could be represented
152 1 a2
DkI(alp)) =
(24)
There is an orthogonal transformation Q which transforms the ambient space μui E R where
f
Σ; #4 = 0 to the subspace in the way that 0
Qu. With || = || 0 || there is
|| اا2م 1| u2l Hence we have:
DKL (N (x, a; 1) (09 a) (1)
ai |xx 2x (09 .t)l
(25)
2
Hence demonstrate the correctness of our objective in Eq.
19
13

---

Published as a conference paper at ICLR 2024
C.1.2 PROOF THE TRANSLATIONAL INVARIANT DENSITY EVALUATION PROCEDURE.
Proof. For an n-atom molecule g - (x, h), the coordinate variable has the dimension of Note that with the zero Center of Mass mapping (Xu et al. 2022; Satorras et al.) 202 where we constrain the center of gravity as zero (E X 0). then variable x essentially lies in the (n 1) X 3- 3-dimensional linear subspace. The generative distributions Px mentioned in all of the related literature (Satorras et al. 2021a Hoogeboom al. 2022 Xu eu al. 2022 2023) is constrained in the zero Center of Mass space. This is, for samples in the ambient space, if -== F 0, Px is not defined. The translational invariant property of distribution px mentioned
referred to the fact that px (x) (x(x= + t) for all translation vector t is satisfied in the ambient space with dimension n x Actually, such conditions could not be satisfied in any space (Satorras et al. 2021a). The translational invariant condition actually refers to the invariant function which could evaluate the density of all the ambient space based on px. the evaluated density by f also referred to as "CoM-free
standard density" in (Xu et al. 2022). The function f is defined
f(x) = px (Qx)
(26)
where Q refers to the operation which maps the x to the zero Center of Mass space, egg. in our work Q is defined as
2is
X1
Q 1 IN (IN- N IN1k) S.t.
Qx
(27)
Xnn
that subtracting mean Σ from each Xs, where Is denotes the x identity matrix and 1 denotes the k-dimensional vector filled with 1s. Then the density evaluation function is translational invariant in the ambient space:
Ei (x.+ti) Xi + t;
f(x +t -xx(((++)) = Px(
... Xnn t,
n
(28)
E1s (xi-ti))
7
Note stands for translation vector, which implies that ti have:
-t;; tu Ce R Then
X1 + C Th (X++C)
1 X1
X1
f(x +t = px(
= px(
= xxxx)x [
(29)
[xn + C Er (x++C)
1i
Xn
Furthermore, the above proof has no constraint the distribution px. This is, for any distribution on the zero Center of Mass space, the corresponding evaluation function defined in Eq. 26 translational invariant.
C.2 PROOF THEOREM.
Given the above discussion the translational invariance, for simplicity, we could only focus on the rotation transformation.
Proof. Recall the graphical model in Fig. we could reformulate the density function in Eq. 7 as:
(x(x)= = / p (x 0
0 (P)P((( (f
OR (definition of marginal)
1 po(x (xx (()))()) I Pu (0; a;) do
(30)
Note that (x 8 Po Rx RO0;) po (R(x) R(O)); σ) due to the property of EGNN,
and p(0)) = p (RO) since 0 = 0. Then we prove that Pt (0; 10-1:a) satisfies the equiv-
ariant condition that pv (0, 0 0.110;;) pu (RO, RO o,). Recall that po (0, 0-11;0))
14

---

Published as a conference paper at ICLR 2024
E (0; h (1, a:); then we have:
poly: 1;a))
Pu (RØ; RR Po 110))
E
8 (RO; - h (RO (R0,
- |
= Po (y: R0;-1;a;) (RO: h (R0;- dy;
(31)
Then we apply interraton-by--susiiuure and replace the variable y, with new variable y i.e. y; = Ry, into the Eq. 31
po (y: 10: (RO: (RO0-11,y;; dy:
/ po (Ry' R0 (RO, (RO, Ry,, a4)) dRy
po (Ry' R0;-1:0;) 5 [RO, h (R0, Ry a;)) det Rd)
(32)
The rotation matrix R is SO(3) matrix, thus the det (R) And for the continuous coordinate variable, the update function h defined in Eq. [ is also equivariant:
RØ 1P:11 R0:-1P:-1
+ Ry,a;
h (RO:- Ry,,,,
R/ (0i-1));aa)))
(33)
Pi
Putting these conditions back the Eq. 32 we have that
Pu (RO: (R0--;;a)) / po (y: (0, (RO: - (RO, (RO,:I. ai)) dy:
/ po (Ry' 10;-1; (RO: Rhh (0,-1. |dyy
/ po (y/ Pc (0; 0;-1;a))
(34)
Hence the transitions on the space are the Markov and equivariant rotation as shown in Eq. 30 The initial state 0u is a zero vector which is rotation invariant. To derive the invariant property Pa. we will use the following Lemma. which is the direct application of Proposition in (Xu et al. 2022). We changed the notation to make it consistent with our literature.
Lemma C.2. (Xu et al. 2022) Let p (0,) be an {F(33-iivrriin invariant density function, i.e., p (0o) p (Tg (0n)). If Markov transitions p (0; 0;-1) are SE(3) i.e., p (0: 0;-1) = p (T, (0,) Tg (0;-1)), then we have that the density (0,) (0n) (0 0o) d0 is also SE(3)-invarians (I, stands for the SE(3) transformations))
For completeness, we also include the derivation of the lemma from Xu et al. 2022):
p (O.)) = (T, (T )..)) Tg (0o) d0wm
(T, (0o)) II IP (T, (0:) Ty (0:-1)) d0.n..
p (00) 1P (T, (0,) Tg (0;-)) d0
(invariant prior (0n))
(35)
(06) 11 (0, 0 1) d0o...
(equivariant kernels (0; (;-1))
p (0o) p (0 tmm. 0.) d0.n
(0.))
Given the invariant property of 0n and equivariant property of the transition py (0, 0;-1;a)) and the P 0%), could directly get the conclusion in Theorem. 3.1 Now finish the proof. []
15

---

Published as a conference paper at ICLR 2024
C.3 PROOF OF PROPOSITION. 3.2
Proof. Then we derive the invariant property of the variational lower bound in of the variational lower bounds in equation 8
(VLB(x)= = P%......0 1=1
DRe(ps (: x; x;); ||pa(- 0;-1:a4) - log po (x 0%).
To start with, we consider the first term:
E Po Po......
DRi(ps ( x; ||pp( 0f;a:))
O
(36)
n-1
Σ
E DKL (ps x; as) [|p[( (f;aa)) (o()
Note natural conclusion from the Theorem. 3.1 is that the SE(3) invariance property is not only satisfied in the marginal distribution of the last time step variable p(0,). but also for the distribution of any intermediate (0,). Such property could be justified based the condition in Lemma. C.2 Actually, the proof of Theorem. i. the above section does not specify the time steps, hence the marginal distribution of any time step could be proved in exactly the same way. Consider the i-th term in the KL part of Lvv.n (Rx):
E (0t)
(KI(Ps (- Rx; a,) [|R[( 0%;aa))
(37)
po (0,) (xL(ps ( Rx; os) [|pn( as))))
introduce the variable 8¹ similar Eq. 32 Le. 0, RO and then extend i-th term in the Eq. 36
Po [ DKL (ps
Rx; a:) [|p (- RO, (de))
(38)
As proved in the proof of Theorem. 3.1 Pe is invariant and hence p& [det(R.) = for SO(3) rotation matrix. And then we discuss the KL divergence term:
RD
(0 ); Also the
Ps (y Rx; as) dy
DRI(ps (: ax; [|p[(- ROO a.))
ps (y Rx; ax);
PR PR (y
RO ;0:)
=ʃ ps (Ry') Rx; a,) log
(Ry' Rx;a;)
det(R))ldy
Pn Pn Ryy RO ;a;)
= -f ps (y') x;; o;) log
ps (y'
(x))s)
dy's (x (- x; x;) [|pa( 0f ax))
(39)
pR(y' ;4;;)))
Note that ps (y' a;) ps RRy') RRx is due to that the sender distribution is isotropic; And for receiver distribution, the equivariant property that Pn (y' (x;as) = PR (Ry') Rx; a;) is guaranteed by both the parameterization of po with Equivariant Graph Neural Network and the isotropic ps- At last, we put the above conclusion back to Eq. (37] and we get that:
E
DKL(ps (- Rx; as) [|pn( 0%
Pa (0t)
/ Po (0.)) DrL(ps (: Rx; ||pn( as))
(40)
Po (0T) DKL(ps (- x; ||p[(: 0f: dd)
E (o():)
DEL (ps (: x; x;) [| 0% a4))
And here we prove the first term in L (Rx) is equivalent to LVLB(x). The second term could be derived in exactly the same way, and here finish the proof.
16

---

Published as a conference paper at ICLR 2024
C.4 DERIVATION OF EQUATION g
Note that the equation
(x)
Σ DKi(ps ( (x;aa)) |pn( 00 log Pa (x
0%)
(41)
Po (0
the extension formulation of Eq. To align the notation of Eq, [ and equation 8 we use x in
the following derivation. We first consider the term -Dxz (gllps .n)))..)))) in Eg Eq.
we put the
(VI....... Nn x)
ps (y: x:
(42)
And the Pe (y1.... vn)) in Eq. as
Ps (y1,...,,Nn)
E po (0
E i= Po (x( (x-00))1100))
[s(y)|x(:aa)
E Pa(Oon)
II
(43)
Putting them together into the KL divergence term, and then we get the
Ps (y: x;
DKz(alle .y........l
E log E
11 p (xaa..
I
PR(y,(00-1;a))
(0mm-1)
I Ps (y: x;) " - ( (1-1;a;)
E
log
(44)
po(0pm- IT ps (xcaa)
Σ
(ps(ps((
a:) [Pp[( 0 0;-1;:a).
(000
And we have derived the first term in equation
And the second term,
log p& (xly1, yn) log pa Oo.
..)) (Graphical Model in Fig.
= log pa x O..)
(Markov Property of 0)
(45)
And here we finish the derivation.
D DETAILS ON CONDITIONAL GENERATION EXPERIMENTS
D. PARAMETERIIZATION AND SAMPLING
For the conditional experiments, directly follow the conditional setting of previous litera- ture (Hoogeboom et al. (002).. We discuss the details of the parameterization and sampling in the following, For conditional experiments, add the property the extra input for the interde- pendency modeling module in Eq. 19 The conditional objective will be as:
L (g.
(t)
x
2x
Do (1)
2 lh,
Pn
a bi
|h; Pr IP)
(46)
Le (g)gtt))
2
2
Where Q. (09, t, c) is short for &. (05, c).
For the sampling procedure, the property and node number M will be firstly sampled from prior p(c, M) defined in (Hoogeboom et al. 2022 Here p(c, M) computed on the training partition as parametrized voo-dimeniional categorical distribution where the continuous variable c discretized into small uniformly distributed intervals. Then we could conduct generation as in Algorithm based on the conditional output distribution po (: 0, c, C,() base Φ00, c, t).
17

---

Published as a conference paper at ICLR 2024
D.2 EXPLANATIONS ON THE PROPERTIES TAB. 2
Polarizability: Tendency of a molecule to acquire an electric dipole moment when subjected to anexternal electric field.
НнооО: Highest occupied molecular orbital energy.
ELUMO: Lowest unoccupied molecular orbital energy.
Ae Gap: The energy difference between HOMO and LUMO.
μ. Dipole moment.
C. Heat capacity 298.15 K
E DETAILED ALGORITHMS FOR TRAINING AND SAMPLING
For better understanding of the whole procedure in training and sampling, we involve the detailed algorithms and implements of functions in Algorithm [ Algorithm and Algorithm
Algorithm 1 Functions for GeoBFN
function DISCRETISED CDF(
R.
E RR, R)
+ F(x)
+
g-p erf
ov
0
if < -1
G(x)
1
if 2
F(r)
otherwise
Return ((z) end function
function OUTPUT PREDICTION(12 Dx3 ph RR.t. [0, 1], Te. 7h Rt tmin R
# Emi set to 0.0001 by default if Iman then
x(0, t) <- 0
pi - 0
σ +-1 1
else
Input
(Ml th, t) to network, receive E(0, ih, In of, output
8 f- us
1-
E(0,
71
T=
An
1-7hh
I-
In 2
end if
for d E 1,
D. k E K do
Po (k
(d)
0:t)
DISCRETISED_CDF(/// (.0 k,) DISCRETISED_CDEE(// (d) .0(( kk)
end for
Return d (0,1t) po( 011) end function
18

---

Published as a conference paper at ICLR 2024
Algorithm 2 Training with continuous loss
Require: On.0n E R. number of bins K E N Input: coordinates a E Rxx3 normalized charges 1 1, t11
L~ 10,
" + 1- o2 7h + 1 oh
Me NNo1 Ya 1)
14 ~ ~Nih,, 7/ T))
a(0, to( 10:)) OUTPUTT (UTPUT_PREDICTION(IL: Mh, Na. th)
k(0,t) + EE (1)
(1)
po(k t.t)ker....
Ex. Ex (k 1
(D)
(x) f- In 21
x -0(0,)]1
(h) <- - 0n
21
h k(0,
Return ( 1 (1z
Algorithm 3 Sampling procedure
(1)
(D)
Function NEAREST CENTER compares inputs to the center bins K., and return the nearest center for each input value.
Require: Jz Tb E R+ number of steps N
Pa. PA t 0
Pz: Ph
for i = to N do
t t i-1
Yv < 12
2t
7hh t of
@00, t), Po
(.1
0,t) t OUTPUT SETPPU_PREDDCTION((,, Mhy Ta, 7h)
21/n
0x 2/n
a, + Jz
21/n
2/n
ah Th
Th
(0,t) 1) < NEAREST_CENTERE
Ex (1) (R: O.t)ke..... , Expo ( 1)))k(
yh ~ (K. t))
Ys (N 1), oz I)
Pz pn < P-1ts to. 0ana+an91
ta
Pston
PayPh < (px + Oz)) (pn + Oa)
end for
2(0, 1), 0;1) + OUTPUT PREDICTION( Ph: 1,1 - (D) (0, 1) + NEAREST NEAREST_CENTER Ex (k 0; (k 0: 1)1k
Return @00, 1), (0,
19


Reviews:
[{"review_id": "4zuEXKcAyD", "reviewer": "Reviewer_m4nX", "summary": "Bayesian Flow Networks (BFN) are a recently proposed generative model, which uses diffusion in the inference process (like diffusion models), but for the generative process maintains a latent distribution parameters over the data, and updates these with Bayesian updates. These models have as advantage that they can handle discrete and discretized variables, besides continuous variables.\n\nThe authors propose to use BFNs to sample molecules, consisting of a collection of atoms, each with a continuous position and discrete atom type (or discretised atom charge). As the authors use a prior equivariant neural network, the resulting sampler is invariant to rotations (and to translations via centering). The authors show state-of-the-art performance in several unconditional and conditional sampling tasks.", "questions": "- The boldness in the $V \\times U$ and Novelty columns of table 1 appears incorrect.\n- There appears to be an inconsistency in the definition of $p_U$ in Eq (5) of the manuscript and Eq (6) of [3]. Is the $y$ sampled from $p_O$ as the manuscript states, or from $p_S$ as [3] states?\n- The authors write \"For Conditional Molecule Generation, we implement a conditional version GeoBFN with the details in the Appendix\", but I can't find this in the appendix.\n- I'm quite surprised that sampling the charges via the discretized method outperforms sampling as discrete atom types. The atomic charge doesn't seems much like a continuum to me. Could the authors elaborate on this? Is it because the hydrogen / not hydrogen distinction is most important, which the discretized method is sensitive to?\n- The results of EDM in table 1 seem worse than those reported in the EDM paper. Why is this?\n\nTypo:\n- Thm 3.1: transitional -> translational\n\nRefs:\n- [1] Köhler, Jonas, Leon Klein, and Frank Noé. 2020. “Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities.”  http://proceedings.mlr.press/v119/kohler20a/kohler20a.pdf.\n- [2] Xu, Minkai, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. 2021. “GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation,”https://openreview.net/forum?id=PzcvxEMzvQC.\n- [3] Graves, Alex, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. 2023. “Bayesian Flow Networks.” http://arxiv.org/abs/2308.07037.\n\n\n---- \nScore raised following the discussion and the positive opinions of the other reviewers.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The method shows strong performance, exceeding prior methods.\n- It's great to see a molecular sampling method used that handles the continuous positions and discrete atom types so naturally.\n- The method improves consistently when more compute (=sampling steps) is used.", "weaknesses": "- I'm not so convinced about the translational equivariance of theorem 3.1. The concept \"Zero of Mass\" is not defined in the cited [1]. I suppose this is the space where $x$ has a zero center of mass. How does this affect $\\theta$ and $y$? [2] gives a detailed analysis about how to handle translation invariance in diffusion, but it's not so clear to me how this applies immediately to a BFN. The proof of theorem 3.1 in the present manuscript says nothing about translations.\n- The proposed method has limited novelty, as it combines a sampling method with an equivariant neural network to create an invariant sampler, as has been done many times previously, without other significant methodological innovation.\n\nIf the authors clear up the translational equivariance, I'll increase my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VMzjoHflrh", "reviewer": "Reviewer_xteH", "summary": "The authors apply a very recent generative modeling framework, [1], to the context of modeling biological molecules, and demonstrate improved performance on established molecular generation benchmarks.\n\n------ Post-rebuttal ------\n\nI believe my concerns about the theoretical presentation and clarity have been resolved. I believe this work is a solid contribution to the field, and I recommend its acceptance.\n\n[1] Bayesian Flow Networks. https://arxiv.org/pdf/2308.07037.pdf", "questions": "I could not find the definition of the subscripts $GeoBFN_{50}, GeoBFN_{100}, \\dots, GeoBFN_{2k}$ in Table 1, Section 4. Where are these defined?\n\nCan the authors clarify the meaning of this sentence in the last paragraph of Section 3.3: \"The underlying reason lies in the fact that the marginal of θi in GeoBFN is in an entropy-increase procedure, e.g., from δ distribution(θ0) to the data distribution(θn). While in diffusion-based models, the marginal is in an entropy-decrease fashion, e.g., from a high-variance Gaussian distribution N (0, I ) to the data distribution.\" This sentence is very unclear to me. Additionally, the reasoning does not entirely connect for me. Isn't $\\theta$ obtained from $\\mathbf{y}$, which is an inherently noisy variable (i.e. Eq. 2)? So why do the authors claim that it has low entropy?\n\nAre there any architectural / data preprocessing differences between the diffusion models (e.g. EDM) and the proposed GeoBFN? Are all improvements in performance attributable to the new training / sampling algorithm given by the Bayesian Flow Networks formulation [1]? Though the derivation is different, the training loss (Eq. 19) ultimately looks very similar to a diffusion model loss.\n\nCan the authors provide a clear formulation of the training algorithm (e.g., via a latex algorithmic block) so it is more clear what is being calculated?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The authors propose a novel model that conceptually simple to understand: the Bayes Flow Network introduced in [1] is applied to the context of molecular modeling by incorporating the equivariant structure of 3D molecular geometries.\n- The empirical capabilities of the model, and its improved performance over existing works, is vetted with established baselines for molecular generation tasks.", "weaknesses": "-  There appears to be some significant issues with the formulation of the model. First, key equations are not derived or justified in the text or appendix. For example, how is the variational bound of the probabilistic model (Eq. 8) derived? Furthermore, how does it lead to Eq. 19? What is $L_\\infty$? Is this the supremum norm? Why does minimizing this value lead to the correct parameters for the proposed model? It is difficult to verify the mathematical consistency of the model without these derivations. Second, the proof of Theorem 3.1 and Proposition 3.2 appear to be incomplete. For example, the proof of Theorem 3.1 ends mid-sentence. Moreover, I do not see anywhere a proof of translation invariance, only rotation invariance via the matrix $\\mathbf{R}$. Additionally, how does Lemma C.1 establish Eq. 23? There appears to be major steps that are skipped in this proof. \n\n- It is not entirely clear why Bayesian Flow Networks improve the performance of modelling 3D tasks (or perhaps specifically molecular generation tasks). The authors attempt to provide some intuition but I am not entirely convinced (see Questions).\n\nOverall, I believe the work is very interesting but the manuscript requires some significant polish before it can be accepted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0RDYz9nnjC", "reviewer": "Reviewer_mPXU", "summary": "The authors apply Bayesian Flow networks (BFN) to the problem of modelling the 3D coordinates of molecules. The authors show how to make the BFN SE(3) equivariant and incorporate an SE(3) equivariant GNN into the BFN architecture for this. Strong performance is obtained on QM9 and GEOM-drugs. Additionally the model allows for flexible specification of the number of steps for sample generation - allowing trading off accuracy for speed.", "questions": "In section 3.3 the text says\n> Hence, GeoBFN implies an objective with smoother information changes.\nWhat does \"objective\" refer to here? The text seems to be referring to sample generation (rather than training) but \"objective\" seems to imply a training objective? \nAdditionally the link between the marginals following an entropy increasing procedure implying an objective with smoother information changes is unclear to me - could the authors please elaborate?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- This paper is the first to apply BFN, a new class of generative models, to molecule generation.\n- Good results: On both QM9 and GEOM-drug BFN outperform diffusion models which are a strong baseline. \n- It is demonstrated that BFN achieve a better trade-off of sample quality vs sampling speed than diffusion models. \n- The lower variance of the parameter space that BFN operate (relative to diffusion models which operate in sample space) seems advantageous - this is nicely visualised in Figure 3.", "weaknesses": "- I found the text inside the section \"Overcome Noise Sensitivity In Molecule Geometry\" unclear (see Questions below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YZAWc15fSu", "reviewer": "Reviewer_iwqe", "summary": "This paper proposes the use of newly formulated Bayesian Flow Networks for use in generating small molecules.", "questions": "- Could you expand why larger numbers of generation steps seems to decrease novelty of the molecules?\n- Is it fair to say that atom stability is easier (requires less generation steps) than molecular stability?\n- Is it clear why the validity of the DRUG database samples seems to decrease with the number of generation steps, and is overall lower than some of the baselines? Is there something special about this dataset?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The authors provide a very readable and accessible introduction to Bayesian Flow Networks.\n- Well-formulated mathematical foundations for the proposed method.\n- The authors compare to a large variety of benchmarks, showing improvement compared to all of them.\n- The authors have run ablation studies that clarify the different design choices made.", "weaknesses": "- The noise-sensitivity section (3.3) is not very clear, the authors should describe in more detail the issue and why a variance-increasing versus variance-decreasing sampling procedure is an important design decision. The claim of \"smoother information changes\" especially seems intuitive yet subjective.\n- The authors should explain how the objective is changing in equation 21, as it is not clear to me how this is improving the issue with sparser sampling of the center buckets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)