Paper: Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for 3D Molecule Generation (ICLR-2024)

Reviewer comment (Experiments):
The manuscript employs a coarse discretization of the radial component for predicting the relative positions of subsequent atoms, a limitation acknowledged in the conclusion section. It would be intriguing to explore the effects of employing a finer discretization. Specifically, would such a refinement significantly compromise the computational efficiency of the method? Furthermore, could this potentially improve the distribution of bond lengths?

Author response:
> would such a refinement significantly compromise the computational efficiency of the method?

We answer exactly this question below! The following is the approximate training speed of the model for a batch size of $16$ graphs, as measured on a single NVIDIA RTX A5000 GPU:

| Number of Radii| $32$ | $64$ | $128$ | $256$ |
| -------- | -------- | -------- | -------- | -------- |
|  Approximate Training Speed (steps/second)     |   $16.5$   | $13.8$ | $10.1$ | $6.6$ |

So while there definitely is a tradeoff, it does not seem to be a major bottleneck. We are currently exploring the effect of the radial discretization on the bond length metrics.

Reviewer comment (Theory):
In Equation (1), would it be more precise to describe this as a conditional distribution of $f$ given $S$? Does the embedder at $f_{n+1}$ use both $h^{\text{position}}$ and $h^{\text{focus}}$?

Author response:
> In Equation (1), would it be more precise to describe this as a conditional distribution of $f$ given $S$?

Hmm, it is just a matter of notation, whether to represent it as $p(f | S)$ or $p(f; S)$. In any case, the fragment $S$ is given at training time, but is sampled from the model's predictions at inference time. So, both options are okay.

> does the embedder at $f_{n + 1}$ use both $h^\text{position}$ and $h^\text{focus}$?

For predicting the focus and the target atom type, $h^\text{focus}$ is used.
For predicting the positions, $h^\text{position}_f$ is used. The two embedders never interact with each other directly.

Reviewer comment (Writing):
Typo: In the final equation on page 4, the term should correctly be denoted as $p^{\text{position}}$.

Author response:
> Typo: In the final equation on page 4, the term should correctly be denoted as $p^\text{position}$.

We believe this formulation is correct: we are parametrizing the logits $f^\text{position}$ of the distribution $p^\text{position}$. The logits are then exponentiated to get the distribution, as described in Section 3.3.

Reviewer comment (Reproducibility):
On the first line of page 6, the mechanism for predicting the STOP condition remains unspecified. Could you elaborate on this aspect?

Author response:
> On the first line of page 6, the mechanism for predicting the STOP condition remains unspecified. Could you elaborate on this aspect?

Great find! We have added this to Section 3.3 under Equation 5.

Reviewer comment (Evaluation):
Complexity: Spherical harmonic projections and inverse projections involve complex mathematical operations, potentially increasing the computational complexity of the model. There is no comparison regarding the time spent by the network among different autoregressive generation models, such as G-SchNet and G-SphereNet.

Author response:
> There is no comparison regarding the time spent by the network among different autoregressive generation models, such as G-SchNet and G-SphereNet.

This comparison is present in Section 4.4 of our manuscript. Symphony is slower than G-SchNet and G-SphereNet. Symphony is currently 3x faster than EDM, but we believe the sampling time for Symphony can be improved by avoiding some of JAX's limitations. We measured the inference times ourselves on the same NVIDIA RTX A5000 GPU.

Reviewer comment (Theory):
Approximation Error: Spherical harmonic projections are an approximation method and may introduce errors, especially when a limited number of spherical harmonics are used.

Author response:
> Approximation Error: Spherical harmonic projections are an approximation method and may introduce errors, especially when a limited number of spherical harmonics are used.

We show the effect of increasing $l_\max$ and the number of position channels for Symphony variants in Section G.1 (Figure 12). We did not feel that much utility could be obtained on going to $l_\max > 5$. We also measure the error due to the sampling grid by measuring the validity of molecules generated in Section G.2 (Figure 13).

Reviewer comment (Theory):
Challenges in Handling Symmetry: When dealing with the symmetry of molecules, ensuring that the predicted distributions are symmetrically equivariant (although proved) may increase the complexity of the model.

Author response:
> When dealing with the symmetry of molecules, ensuring that the predicted distributions are symmetrically equivariant(although proved) may increase the complexity of the model.

This is true, but we believe it is necessary that a model should respect the underlying symmetries of a problem. Ours is just one approach, however.

Reviewer comment (Theory):
I think the Property (3) in page 4 is very important. And many previous autoregressive models can not solve it well. Actually I think it is not difficult to define an E(3)-equivariant model currently as many previous works have solved it well. The diffusion model doesn't face the permutation-invariant problem. But it can be serious for the autoregressive model. The main paper claims the three properties have been proved in Theorem B.1. Unfortunately, I found the authors only proved the first two properties and ignored the property (3). I DON'T believe the current framework can guarantee the permutation-invariant even the embedder is E(3)-equivariant.

Author response:
Sorry for omitting this in our original manuscript. In fact, we had an error in our property (3): we meant to say that the focus distribution should be permutation-equivariant. We have added a proof of this in Theorem B.2. Note that we are not claiming that the entire generation process is invariant to the choice of ordering of nodes; this is almost impossible to guarantee for an autoregressive model. Note that G-SchNet and G-SphereNet both impose some sort of ordering on the way atoms are added. We are also planning to explore some alternate strategies compared to the nearest-neighbours logic in our paper.

Reviewer comment (Experiments):
I think many criticisms of the diffusion model such as EDM on page 6 are unfair. "To avoid recomputation of the neighbor lists during diffusion, current diffusion models use fully-connected graphs where all atoms interact with each other. This could potentially affect their scalability when building larger molecules." I believe it is even more challenging for the autoregressive model on the large molecules as the sequence could be too long. And the molecules in QM9 used in the experiment are very small. There is no evidence that the proposed framework performs better on large molecules compared with EDM.

Author response:
This is a fair criticism; we will plan to demonstrate Symphony on the larger GEOM-DRUGS dataset that has been used in other work.

Reviewer comment (Evaluation):
I think many criticisms of the diffusion model such as EDM on page 6 are unfair. "diffusion models are significantly slower to sample from, because the underlying neural network is invoked ≈ 1000 times when sampling a single molecule." It is true that for EDM it needs 1000 times denoising process. But actually, it runs very fast and supports batch sampling. I tested EDM on a single GPU A100, it takes 20 seconds to sample 50 molecules (a batch). I think the submission should provide their own sampling time to support this claim.

Author response:
We have provided the sampling times in our common response above (and in Section 4.4 of the paper). Symphony is currently 3x faster than EDM, but we believe the sampling time for Symphony can be improved by avoiding some of JAX's limitations. We measured the inference times ourselves on the same NVIDIA RTX A5000 GPU.

Reviewer comment (Reproducibility):
Though the definition of the position distribution is very clear and reasonable, the current submission doesn't provide any details on how to sample from it after training. People can know $f^{position}$ for any given position, but it is not easy to sample from an energy function. I doubt this sampling process can be very time-consuming.

Author response:
Sorry for omitting this, we have added this in our common response (and in Section D of the Appendix of the paper). Hopefully this clears up any confusion.

Reviewer comment (Experiments):
Why did you choose 64 uniformly spaced values from 0.9A to 2.1A? I think the model can learn the distribution automatically by equation (3).

Author response:
This should now be clear from Section D of the Appendix. Basically, we do not have a orthonormal basis over the space $(r, \\theta, \\phi)$,  only the space $(\\theta, \\phi)$. Discretization of $r$ was the simplest alternative; we plan to explore other basis functions (and normalizing flows) soon.

Reviewer comment (Reproducibility):
Since there are 64 spaced values for $r$, I am wondering how you solve them during the sampling process.

Author response:
This should now be clear from Section D of the Appendix. Basically, we do not have a orthonormal basis over the space $(r, \\theta, \\phi)$,  only the space $(\\theta, \\phi)$. Discretization of $r$ was the simplest alternative; we plan to explore other basis functions (and normalizing flows) soon.

Reviewer comment (Evaluation):
Can you give some explanation as to why the EDM outperforms the proposed framework for many metrics?

Author response:
Great question! We are also not super sure; the radial discretization that Symphony uses seems to hurt the bond length metrics. However, note that EDM is also trained for 10x longer than Symphony was. We are exploring if the gap can be bridged with further training.

Reviewer comment (Experiments):
In effect, the authors define an energy based model for the positions. It appears from the code that training and sampling of positions is done by discretizing $f^{position}$ on a grid. This seems like a shortcoming of the method, as I suspect that precise positioning is important in molecular generation. The paper should be transparent about this discretization. What type and resolution of grid is used? It'd be interesting to see how the resolution affects the sample quality. If the quality is highest with a very fine grid, it'd be great if the authors could consider and evaluate alternative learning and sampling strategies for energy based models (contrastive divergence / Hamiltonian Monte Carlo / Langevin dynamics etc).

Author response:
> in effect, the authors define an energy based model for the positions.

We have clarified this above; in the autoregressive setting with teacher-forcing, we know the target distribution exactly. Further the partition function can be easily estimated via numerical integration. Thus, we do not need any complicated methods for learning the position distribution, because minimizing KL divergence works well. 

> The paper should be transparent about this discretization. 

Thank you for mentioning this. As discussed in the common rebuttal, we have added this to Section D of the manuscript. We use a Gaussian Product grid for each value of $r$, which consists of a 1D Gauss-Legendre quadrature with $180$ points over $\\cos \\theta \\in [-1, 1]$, and a uniform grid of $359$ points over $[0, 2\\pi)$ for $\\phi$.

> It'd be interesting to see how the resolution affects the sample quality.

Yes, we have added this in Section G.2 (Figure 13) where we vary the resolution of the grid significantly. The model does not seem particularly unstable to these variations. In Figure 14, we show that the training is also not significantly affected by the resolution for a synthetic task of learning a single distribution on a sphere. We hope that these sets of experiments help alleviate your concern about the gridding procedure.

Reviewer comment (Theory):
While the approach predicting the sequence index for the focus point is in/equivariant to permutations of the resulting molecule, it still appears to me that an opportunity for additional permutation equivariance is missing. Why didn't the authors use a permutation equivariant network on $S^n$ to select one focus atom?

Author response:
>  Why didn't the authors use a permutation equivariant network on $\mathcal{S}^n$ to select one focus atom?

We are very sorry for the confusion. We actually do use a permutation equivariant network (GNN) on $\mathcal{S}^n$ to select the focus atom. Our Property (3) incorrectly stated that the focus distribution should be permutation-invariant; this has now been fixed to read permutation-equivariant. We also added Theorem B.2 in the Appendix to prove that this must be the case.