Unified Generative Modeling of 3D Molecules with Bayesian Flow Networks

Q3. The reason why larger numbers of generation steps seem to decrease the novelty of the molecules

This is a great question! Note that by increasing the discretization steps, the sampled distribution will approach the distribution implied by the fully continuous objective in Eq. (19). The novelty is defined as the proportion of generated samples that are unique and unseen in the training dataset. This reveals the consistency of model distribution in Eq. (19) and data distribution.

Besides, the decrease in novelty could also raise a concern about an overfitting phenomenon. And we clarify the concerns in the following. As demonstrated in previous works [2,3], the concept of novelty must be examined alongside validity, uniqueness, and stability rates to get a comprehensive evaluation of the generative model. For a fair comparison with other models, e.g. EDM and GeoLDM, which use 1k steps to sample, we consider the performance of GeoBFN under the 1k steps. We could find that our method indeed demonstrates higher novelty compared to others. With the superior performance on the molecule quality metrics (stability, etc), our method also holds a higher novelty. This demonstrates that our proposed method shows better generalization ability and could better explore the molecule space instead of overfitting.

Q4. Relationship between atom stability and molecule stability

Yes, as the reviewer noticed, with fewer generation steps, the atom stability does not decrease obviously. Note that atom stability is defined as the proportion of atoms having the right valency, and molecular stability is defined as the proportion of generated molecules for which all atoms are stable [2]. Hence, molecule stability could be approximately seen as the $M$-th exponential to the atom stability where $M$ is the number of atom nodes.

Q5. 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?

We borrow the discussion towards validity and stability in [2]. "Note that the validity is evaluated using RDkit. The molecule is built to contain only heavy atoms and RDKit will add hydrogens to each heavy atoms in such a way that the valency of each atom matches its atom type. Hence invalid molecules mostly appear when an atom has a valency bigger than expected. Experimentally, it is observed that validity could artificially be increased by reducing the number of bonds. For example, predicting only single bonds was enough to obtain close to 100% of valid molecules on GEOM-DRUGS. On the contrary, the stability metrics directly model hydrogens and cannot be tricked as easily." As the reviewer noticed, the validity decreased with the number of generation steps and was even lower than some baselines on GEOM. This could reflect some intrinsic properties of GeoBFN. We hypothesize that the GeoBFN would tend to generate molecules with a more compact structure (as the generation variance is lower) and hence could get bigger valencies for atoms especially when the atom number is large. However, it is worth noticing that on the stability metrics which could not be tricked easily, GeoBFN still shows superior performance.

[1] CRC Handbook. 2007. "CRC Handbook of Chemistry and Physics", 88th Edition.

[2] Emiel Hoogeboom, Victor Garcia Satorras, Clement Vignac, Max Welling 2022 "Equivariant Diffusion for Molecule Generation in 3D"

[3] Minkai Xu, Alexander S. Powers, Ron O. Dror, Stefano Ermon, Jure Leskovec 2023 "Geometric Latent Diffusion Models for 3D Molecule Generation"

Thanks a lot for the constructive comments and. recognition of our work. We address all your concerns in the following.

Q1. 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.

Sorry for causing the confusion. Section 3.3 has been rewritten as suggested.

• Further clarification on noise sensitivity:

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 at a more rapid pace, e.g. modifying the coordinates of one atom results in altering its distance to all other atoms.

• Towards the "smoother information changes" and design decision:

We agree that such a claim could be subjective as the reviewer says and we have rephrased the presentation in the updated version to eliminate the misunderstanding. The opinion of "smoother information changes" is borrowed from the original BFN paper [2], which describes the generation process in the parameter space as regularized with the Bayesian update procedure. This is, noisier samples will be assigned with a smaller weight during the update (Eq. 11). Our objective here is to describe that during the diffusion process for generation, the intermediate steps' structure might be uninformative, with the majority of the information being acquired in the final few steps of generation (as depicted in Figure 3); On the other hand, as shown in Figure 3, the structure of the intermediate steps gradually converges to the final structure, suggesting a smoother increase in information. The key intuition is our belief that gradual changes could potentially provide a beneficial inductive bias for generation procedure design decisions as in [2]. The "entropy-increasing" and "entropy-decreasing" are removed for a more strict and objective presentation.

Q2. 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 the sparser sampling of the center buckets.

We apologize for any confusion caused. The original sampling process for the discretised variable of the BFN, as detailed in [2], is depicted in Algorithm 6. Here, the variable $k$ is sampled using the formula

We thank the reviewer for the insightful comment and the recognition of our work. We address your concerns in the following:

Q1. Toward the section "Overcome Noise Sensitivity In Molecule Geometry":

• What does "objective" refer to here?

And yes, as the reviewer mentioned, here we are aiming to discuss the generation process. The "objective" should change to the "generation process".

• The link between the "entropy increasing procedure" and "smoother information changes"

Sorry for causing the confusion. The opinion of "smoother information changes" is borrowed from the original BFN paper [1], which describes the generation process in the parameter space as regularized with the Bayesian update procedure. This is, noisier samples will be assigned with a smaller weight during the update (Eq. 11). And we realized that the "entropy increasing procedure" is not strict enough considering the differentiable entropy of continuous variables. Our objective here is to describe how during the diffusion process for generation, the intermediate steps' structure might be uninformative, with the majority of the information being acquired in the final few steps of generation (as depicted in Figure 3); On the other hand, as shown in Figure 3, the structure of the intermediate steps gradually converges to the final structure, suggesting a smoother increase in information. We have revised section 3.3 as suggested to prevent further misinterpretation.

[1] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, Faustino Gomez 2023 "Bayesian Flow Networks"

Q4. The architectural / data preprocessing differences to EDM and the attribution of the improvements of GeoBFN.

In our experiment, we ensured all components such as network architecture, data preprocessing, and evaluation pipeline were consistent with the settings in EDM to maintain fair comparisons. The performance enhancements were primarily due to the new training objective and sampling algorithm. As evident in Table 3's ablation studies, the performance was also improved with an optimized version of sampling. The training loss in Equation (19), as the reviewer mentioned, is similar to the diffusion objective. The distinguishing factors between BFNs and Diffusion models mainly lie in their differing modeling space and graphic models, as portrayed in Figure 2. The parameter space has the favorable characteristic of low variance, and the sampling process varies significantly from the diffusion model. In the latter, the final sample is forecasted and incorporated through a Bayesian update. The joint modality modeling and previous properties could primarily contribute to the improved performance.

Q5. Clear formulation of the training algorithm.

Thanks for the suggestion. We have included the full training and sampling algorithm as suggested in Appendix E of the updated draft.

[1] Diederik P Kingma, et. al. 2013. "Auto-Encoding Variational Bayes"

[2] Jonathan Ho, et. al. 2020"Denoising Diffusion Probabilistic Models"

[3] Yang Song, et. al. 2021. "Score-Based Generative Modeling through Stochastic Differential Equations"

[4] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, Faustino Gomez 2023 "Bayesian Flow Networks"

Q2. Definition of the subscripts: GeoBFN50, GeoBFN100,…, GeoBFN2k in Table 1, Section 4.

Apologies for the confusion caused. When training with a continuous time step objective, we have the flexibility to sample the molecule with arbitrary steps.

In our approach, we use the term $\text{GeoBFN}\_{k}$ to denote the process of sampling the molecules with a specific number of steps. For instance, $\text{GeoBFN}\_{50}$ refers to the sampling of a molecule with 50 steps. To address your suggestion, we have included an explanation of this terminology in the updated draft.

Q3. Clarification of the last paragraph of Section 3.3.

• Towards the "entropy" increase or "entropy" decrease:

Sorry for making the presentation unclear. After carefully checking the presentation in the mentioned sentence, the term "entropy" needs to be further clarified as the reviewer mentioned. Here we tend to distinguish the information changes between the BFNs and diffusion models. Entropy could be used as a good metric to analyze. For strict presentation, we need to limit the scoop of discussion on quantized space to guarantee the entropy is non-negative, e.g. images with a pixel value in {0,1,...,255}, instead of a continuous variable to avoid the discussion on differentiable entropy. In such cases, the entropy could be interpreted as the bits needed to describe/compress the distribution. Thus, in diffusion models, the sampling starts from a normal distribution, and with iterative refinement, the distribution approaches the data distribution which certainly has less entropy than the initial distribution("entropy" decrease); While in Bayesian Flow networks, the starting point of all sampling procedure is the same constant $\theta_0 = **0**$ where the entropy is zero. Here the sampled distribution will have entropy larger than the initial state("entropy" increase). The mentioned part is revised and the mentioned inaccurate parts are removed as suggested in the updated version to make it clear.

Essentially, we would like to provide insight by distinguishing the source of randomness of the two different models. In the diffusion model, with recent advancements such as probability flow ode [3], the randomness of the system could be seen as obtained from the initial prior distribution. While in BFN, the randomness of the system is added gradually at each step by conducting a Bayesian update with the noisy variable.

• To the question of why the variable $\theta$ has lower entropy (variance):

As we only claim the modeling on the $\theta$ space enjoys lower "variance", we suppose the reviewer is concerned about the low-variance property. We take the modeling procedure of BFN on the continuous variable as an example. And yes, as the reviewer mentioned, $\theta$ is indeed obtained from $y$ which is the noisy variable. Recall the Bayesian update function in Eq. (11), i.e. $\boldsymbol{\mu}\_i = \frac{\boldsymbol{\mu}\_{i-1} \rho_{i-1}+\mathbf{y} \alpha}{\rho_i}$. The variance of $y$ is $\alpha^{-1}$ as defined in Eq. (10). Therefore, for $y$ with large variance, the $\alpha$ is small and hence contributes less to update $\theta$. In other words, the update of $\theta$ is dominated by those $y$ with small variance. Such a property of Bayesian update helps get a lower variance on $\theta$ compared to $y$. A visualizational analysis can be found in Figure 5 of [4]. The above discussion has been included as suggested.

Thank you very much for the constructive suggestions and detailed comments. We address all your concerns in the following:

Q1. The presentation/formulation issues:

Thank you for bringing these issues to our attention, and we apologize for any confusion caused. We have addressed all of the mentioned issues in the following revised version.

• variational bounds in Eq. (8)

Actually, the variational bounds in Eq. (8) is exactly the detailed expansion of the variational bounds in Eq. (1), i.e.:

In Section 2.2, we derive the formulation of each component: $q(\mathbf{y}_1, \ldots, \mathbf{y}_n)$ in Eq. (2); $p\_\phi\left(\mathbf{y}_1, \ldots, \mathbf{y}_n\right)$ in Eq. (6); The reconstruction term $p\_\phi\left(\mathbf{g} \mid \mathbf{y}_1, \ldots, \mathbf{y}_n\right)$ is defined exactly as $p\_\phi\left(\mathbf{g} \mid \theta_n \right)$ in the framework. Push all this together into Eq. (1) we get the formulation of Eq. (8). And we include the key steps here.

A more formal and detailed derivation can be found in Appendix C.4 of the revised version

• derivation from Eq. (8) to Eq. (19) and explanation on $L_{\infty}$.

We apologize for any confusion caused. Our usage of the notation $L_{\infty}$ directly follows the notation in [1], which describes the objective function with an infinite number of continuous time steps. On the other hand, Equation (8) pertains to the objective function with a finite number of discrete time steps. The process of extending Equation (8) to account for continuous time steps can be found in [1] (specifically, from Equation (25) to Equation (41) for continuous variables). We added the reference as suggested to enhance the comprehensiveness of our paper.

• The correctness of the optimization objective $L_{\infty}$

As shown in the above discussion, the optimization objective $L_{\infty}$ can be viewed as an extension of the variational lower bound of the log-likelihood function to the case of infinite continuous time steps. Therefore, optimizing this objective essentially means maximizing the variational lower bounds of the log-likelihood function. Although there is a slight bias towards the maximum likelihood objective, such objectives are widely applied in training generative models, such as VAEs [1] and Diffusion Models [2]. We can observe the utilization of these objectives in practice.

• proof of Theorem 3.1 and Proposition 3.2 and the relationship between Lemma C.1 to Eq. (23)

Thanks a lot for pointing it out. The proof of Theorem 3.1 and Proposition 3.2 has been fully rewritten as suggested in the updated version in Appendix C.

• proof and discussion of translation invariance

The discussion and proof of translational invariance are added in Appendix C as the reviewer suggested. A more intuitive discussion can be found in the first part of our response to Reviewer m4nX's Q1.

Q4. The details of the Conditional Molecule Generation experiment.

Thanks for bringing this to our attention. The generation of conditional molecules is directly performed after conducting conditional experiments in our previous work [5,6]. During conditional training, we utilize the interdependency modeling network $\phi$ in Eq.(19), which takes the property $c$ as an additional input, i.e., $\Phi(\theta,t,c)$.

In the process of generating molecules conditionally, we first sample the property $c$ and the node number $M$ from a prior distribution $p(c,M)$ that is defined in [5]. The distribution $p(c,M)$ is computed on the training partition and is parameterized as a two-dimensional categorical distribution. The continuous variable $c$ is discretized into small, uniformly distributed intervals.

For a more comprehensive explanation, please refer to the updated information in the Appendix D.

Q5. Discussion on sampling the charges via the discretized method outperforms sampling as discrete atom types.

This is a great question! Firstly, our main motivation for using charges for sampling is that we only need to model and sample on two similar modalities: discretized data type and continuous data type. This approach reduces the redundancy of information between the charges and types compared to previous literature [5, 6], which utilized all three modalities for modeling.

Furthermore, the discretized charges allow for a more granular perspective of the distinctively varied electrostatic environment around different atoms. This introduces a more informed sampling framework that better integrates the chemical properties and reactivity.

Additionally, as the reviewer mentioned, the atomic charge may not present a clear continuum, and there are nuances that justify the use of discretization. We fully agree with the reviewer's proposition that distinguishing between hydrogen and non-hydrogen atoms could be another important way to better understand the benefits of our proposed sampling method. Hydrogen is unique in terms of its proton count and chemical behavior, and while a categorical approach can differentiate between hydrogen and other atoms, a discretized approach can leverage this distinction even further. This could inspire a fruitful feature direction for further exploration.

Q6. The results of EDM in Table 1 seem worse than those reported in the EDM paper.

The results of EDM are directly copied from the EDM paper. After carefully checking the number, we make sure it is consistent with the number reported in their paper [5] and follow-ups [6].

Q7. Typos and boldness in Table 1.

Thanks again for the detailed comment. The mentioned issues are fixed in the updated version.

[1] Köhler, Jonas, Leon Klein, and Frank Noé. 2020. “Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities.”

[2] Xu, Minkai, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. 2021. “GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation,”

[3] Victor Garcia Satorras, Emiel Hoogeboom, Fabian B. Fuchs, Ingmar Posner, Max Welling. 2021. "E(n) Equivariant Normalizing Flows"

[4] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, Faustino Gomez 2023 "Bayesian Flow Networks"

[5] Emiel Hoogeboom, Victor Garcia Satorras, Clement Vignac, Max Welling 2022 "Equivariant Diffusion for Molecule Generation in 3D"

[6] Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, Jure Leskovec 2023 "Geometric Latent Diffusion Models for 3D Molecule Generation"

We thank the reviewer for the detailed comments and insightful suggestions. We address your concerns in the following:

Q1. The translational equivariance of theorem 3.1.

Thank you for bringing this up, and we apologize for any confusion caused. The term "Zero of Mass" actually refers to a slightly abused concept of "Center-of-Mass Free" as discussed in Section 5 of [1], as well as "Zero Center of Mass" mentioned in [2]. We will now provide a thorough analysis of the translation invariance in Theorem 3.1, which is outlined below. The term SE(3)-invariant density is a concept that has been derived from previous research [2]. However, the notion of a "translational-invariant density" requires further clarification. In the normal Euclidean space, there is no distribution that can be truly translation-invariant, meaning that $p_X(x) = p_X(x+t)$ for all t, where t represents the translation vector. This condition implies that the probability density function is a constant function, which can not integrate into one [3].

In existing literature, the term "translational-invariant property" actually refers to two key ideas.

1. Firstly, it suggests that we can avoid modeling the freedom of translation of the geometries by learning a probability distribution only on the subspace where the center of mass is zero.
2. Secondly, it implies that it is possible to construct a function $f$ that can evaluate the density for samples in the original sample space, including samples with the center of mass is not zero.

Intuitively, this function $f$ maps the samples to the zero center of mass space and returns the density of probability defined on that space. It is important to note that this function f does not correspond to the probability density function of any real distribution and is often referred to as the "CoM-free standard density" as described in [2]. Achieving translational invariance can be easily accomplished by constraining the sample space of the generative model to lie in the zero center of mass space. This can be achieved by moving the center of mass to zero for the generated samples, denoted as $x$ in our context.

In Theorem 3.1, we have regularized all the parameters $\theta$, $y$, and $x$ to be centered around zero. Although this condition is sufficient to achieve translation invariance, it is not necessary as mentioned earlier. Specifically, regularizing the output of $\phi_x$ in the zero Center of Mass space should be sufficient to obtain this property. To ensure stable training, we also center all $\theta$ and $y$ around zero in the Center of Mass space.

For more detailed and formal proof, as well as further discussion, please refer to Appendix C.1 of the updated version.

Q2. Towards limited novelty

Our paper is motivated by the natural compatibility of BFN[4] with the two key challenges in molecule generation. Apart from the several innovations made for application on the specific task, we would like to highlight several contributions:

• We present an alternative approach to understanding BFN through a simplified, graphical model summary of the original BFN paper outlined in [4]. This paper was based on a "communication" perspective that could often prove challenging to grasp and apply in other fields. The summarization could act as a beneficial resource aiding the comprehension of the complex properties and further exploration of this sophisticated generative model.
• Another significant aspect is the demonstration of BFN compatibility with SE(3) invariant density modeling tasks, which is substantiated by the proof of Theorem 3.1. This realization is non-trivial given the fundamental differences between BNF and previous diffusion models.

Q3. The inconsistency of Eq. (5) in our paper and Eq. (6) in [4]:

Thank you for pointing it out, and we apologize for any confusion caused. In our paper, we use Eq. (5) to describe the components of the generative process or the generative model. On the other hand, in [4], Eq. (6) is employed to formalize the training objective.

In Eq. (5), the $p_O(y_i|\theta_{i-1};\alpha_i)$ should be changed to the so-called receiver distribution in [4], i.e.,
$p_R (y_i|\theta_{i-1};\alpha_i) = \underset{p_O\left(\mathbf{x}^{\prime} \mid \boldsymbol{\theta}_{i-1}; \phi \right)}{\mathbb{E}} p_S\left(\mathbf{y} \mid \mathbf{x}^{\prime} ; \alpha_i\right)$.

The typos have been fixed as suggested in the updated version.

I thank the authors for their follow-up response. My second and third question have been addressed. Regarding the first, I'm not sure I am fully convinced.

I understand that you want to define a distribution on the space $\mathbb R^{3 N - 3}$, of the atomic positions modulo the centre of mass. My question is: what is the density on that space (which you use in the loss)?

If I have an arbitrary density $p(x)$ on $\mathbb R^{3 N}$, and project it to a density $\hat p(\hat x)$ on $\hat x \in \mathbb R^{3 N - 3}$, then the value of the zero-CoM density is given by the integral:

This is not very practical and I understand you don't do this integration in practice. So how do you compute the densities in the zero CoM subspace?

As an example, I understand $p_S(y | g, \alpha)$ to apply a Gaussian to $g$. If we assume $g$ to be zero CoM subspace, then an isotropic Gaussian on $\mathbb R^{3 N}$ moves $y$ out of that subspace. So I suppose that your noise distribution is only supported in the zero CoM subspace? If so, does your likelihood include a factor $(2\pi)^{-(3N-3)/2}$? Similar for your distribution $p(g|\theta)$, which I suppose is a Gaussian with diagonal covariance matrix. Is that a diagonal $3N \times 3N$ matrix or a diagonal $(3N-3)\times (3N-3)$ matrix? If it's the former, I suppose that it doesn't project to a diagonal $(3N-3)\times (3N-3)$ matrix. Do you then explicitly compute the determinant of that matrix to find the likelihood?

Prior works have worked with zero CoM variables in a diffusion context, where I understand from [1] that things simplify, but I haven't seen a convincing argument why the same applies for the BFN.

[1] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. 2021. “GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation,”