Paper: Unified Generative Modeling of 3D Molecules with Bayesian Flow Networks (ICLR-2024)

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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

$$
\mathbf{k} \sim \text{DISCRETISED\\_OUTPUT\\_DISTRIBUTION}\left(\boldsymbol{\mu}, t, 1-\sigma_1^{2 t}\right) = 
\text{DISCRETISED\\_CDF}\left(\mu_x^{(d)}, \sigma_x^{(d)}, k_r\right) - \text{DISCRETISED\\_CDF}\left(\mu_x^{(d)}, \sigma_x^{(d)}, k_l\right)$$


In this equation, $\text{DISCRETISED\\_OUTPUT\\_DISTRIBUTION}$ denotes the probability of each bin and the sampling procedure essentially involves drawing a sample from a categorical distribution defined by $\text{DISCRETISED\\_OUTPUT\\_DISTRIBUTION}$. 

For instance, if the $\text{DISCRETISED\\_OUTPUT\\_DISTRIBUTION}$ is supported on {0,1,2} with corresponding probabilities {0.5,0,0.5}, the original sampling method would only draw from either {0} or {2}. However, in Eq. 21, we first compute the weighted average, which turns out to be $0.5\*0 + 0\*1 + 0.5\*2 = 1$, and subsequently select the nearest neighbor in the support set, which also yields the value $1$. It's important to note that this method aligns more closely with the training objectives articulated in Eq.19. We have revised the corresponding part as suggested.

Reviewer comment (Evaluation):
Could you expand why larger numbers of generation steps seems to decrease novelty of the molecules?

Author response:
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.

Reviewer comment (Theory):
Is it fair to say that atom stability is easier (requires less generation steps) than molecular stability?

Author response:
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.

Reviewer comment (Evaluation):
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?

Author response:
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.

Reviewer comment (Writing):
I found the text inside the section "Overcome Noise Sensitivity In Molecule Geometry" unclear (see Questions below).

Author response:
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"

Reviewer comment (Writing):
In section 3.3 the text says "Hence, GeoBFN implies an objective with smoother information changes." What 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?

Author response:
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"

Reviewer comment (Theory):
Additionally 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?

Author response:
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"

Reviewer comment (Theory):
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.

Author response:
**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.:

$$
\begin{align*}
\log p_{\boldsymbol{\theta}}(\mathbf{g}) & \geq  \underset{\mathbf{y}_1, \ldots, \mathbf{y}_n \sim q}{\mathbb{E}}\left[\log  \frac{p\_\phi\left(\mathbf{g} \mid  \mathbf{y}_1, \ldots, \mathbf{y}_n\right) p\_\phi\left(\mathbf{y}_1, \ldots, \mathbf{y}_n\right)}{q\left(\mathbf{y}_1, \ldots, \mathbf{y}_n \mid  \mathbf{g}\right)}\right] \\newline
& =-D\_{KL}\left(q \| p\_\phi\left(\mathbf{y}_1, \ldots, \mathbf{y}_n\right)\right)+\underset{\mathbf{y}_1, \ldots, \mathbf{y}_n \sim q}{\mathbb{E}} \log  \left[p\_\phi\left(\mathbf{g} \mid  \mathbf{y}_1, \ldots, \mathbf{y}_n\right)\right]
\end{align*}
$$

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. 

$$
\begin{align*}
D_{K L}\left(q\left(\boldsymbol{y}_1, \ldots, \boldsymbol{y}_n \mid\boldsymbol{x}\right) \| p\_\phi\left(\mathbf{y}_1, \ldots, \mathbf{y}_n\right)\right)&=\underset{\prod\_{i=1}^n p_S\left(\mathbf{y}_i \mid \mathbf{x} ; \alpha_i\right)}{\mathbb{E}} \log \frac{\prod\_{p\_\phi\left(\boldsymbol{\theta}\_{0: n-1}\right)}^n \prod\_{i=1}^n p_R\left(\mathbf{y}_i \mid \mathbf{x} ; \alpha_i\right)}{\mathbb{E}\left(\mathbf{y}_i \mid \boldsymbol{\theta}\_{i-1} ; \alpha_i\right)} \\newline
& =\underset{p\_\phi\left(\boldsymbol{\theta}\_{0: n-1}\right)}{\mathbb{E}} \underset{\prod\_{i=1}^n}{\mathbb{E}} \mathbb{p _ { S } ( \mathbf { y } _ { i } | \mathbf { x } ; \alpha _ { i } )} \log \frac{\prod\_{i=1}^n p_S\left(\mathbf{y}_i \mid \mathbf{x} ; \alpha_i\right)}{\prod\_{i=1}^n p_R\left(\mathbf{y}_i \mid \boldsymbol{\theta}\_{i-1} ; \alpha_i\right)} \\newline
& =\underset{p\_\phi\left(\boldsymbol{\theta}\_{0: n-1}\right)}{\mathbb{E}} \sum\_{i=1}^n D\_{KL}\left(p_S\left(\cdot \mid \mathbf{g} ; \alpha_i\right) \| p_R\left(\cdot \mid \boldsymbol{\theta}\_{i-1} ; \alpha_i\right)\right)
\end{align*}
$$

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.

Reviewer comment (Theory):
Can 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?

Author response:
**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 = \textbf{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.

Reviewer comment (Experiments):
Are 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.

Author response:
**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.

Reviewer comment (Reproducibility):
Can 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?

Author response:
**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.

Reviewer comment (Theory):
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.

Author response:
**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.

Reviewer comment (Presentation):
The boldness in the $V \times U$ and Novelty columns of table 1 appears incorrect.

Author response:
**Q7. Typos and boldness in Table 1.**

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

Reviewer comment (Reproducibility):
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.

Author response:
**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.