Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for 3D Molecule Generation

Thank you for your positive review! We answer specific questions below. Please also take a look at our common 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:

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.

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.

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.

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.

Typo: The first equation in Section 2 misses a translation $T$.

The features we look at here are invariant under translation. We have clarified this in the text below.

Thank you again!

Thank you for your thoughtful feedback! We answer specific questions below. Please also take a look at our common 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.

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

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.

Thank you for your great feedback and taking the time to review our paper! We answer specific questions below. Please also take a look at our common response.

it lacks an ablation study to show how the multiple channels influence the final performance. For example, if the channel is 3 or 4.

This has now been added in Section F.2 (synthetic task) and Section G.1 (QM9). There is a clear benefit when adding channels for smaller $l$.

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.

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.

There is no evidence that the proposed framework performs better on large molecules compared with EDM.

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

I think the submission should provide their own sampling time to support this claim.

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.

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.

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.

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).Since there are 64 spaced values for $r$, I am wondering how you solve them during the sampling process.

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.

Can you give some explanation as to why the EDM outperforms the proposed framework for many metrics?

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.

Thank you so much for the great review! We answer specific questions below. Please also take a look at our common 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.

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.

It appears to me that the position distribution is an instance of a Harmonic Exponential Family [1] on a manifold, which probably should be cited.

Thanks for sharing this relevant article. We have added this as a citation.

Can the authors clarify why Eq (4) is smooth wrt the radius, but has a Dirac delta in the direction component? This should mean that the KL divergence is infinite almost always? Or is the smoothening handled implicitly by a band-limited Fourier transform?

Yes, you are correct. This is an approximate Dirac Delta distribution that this constructed using the same spherical harmonic projection idea. We have a description of this in Appendix H.

Did the authors mean the group SE(3) excluding mirroring?

We actually meant E(3) itself, because our proofs handle the case of improper rotations. Our predicted distributions will invert as well under inversions about the origin (mirroring).

Can the authors comment on the runtime of autoregressive generation vs diffusion methods in terms of the number of atoms? Might it be that autoregressive methods are only faster for small molecules?

We are currently trying to measure this (atleast approximately). The scaling time of EDM seems to be quadratic in the number of atoms (due to fully connected graphs). We will have an update soon.

Thanks again!

As high frequencies may be necessary to precisely predict positions, it'd be great if the authors can show in an ablation study that l_\\max = 5 suffices.

Yes, we have added an experiment with Symphony variants trained for different values of l_\\max and number of position channels in Section G.1 (Figure 12)

We have edited our response above and the appendix with the changes you recommended. We hope that our additional experiments in Section G.2 help resolve your worries about the discretization. Please let us know if there are any more experiments you would like to see. Otherwise, we would greatly appreciate it if you could improve your score for our paper. Thank you again!

I forgot to add: please consider adding a reference to [1] in a future version. Your way of defining and learning a distribution appears closely related to theirs.

[1] Murphy et al, Implicit-pdf: Nonparametric representation of probability distributions on the rotation manifold, 2021

Thank you for your bringing our attention to this paper! We were unaware of their approach. We will update our paper soon to discuss the differences between Symphony and Simm, et al. (2021) in model design, training and sampling.

Ablation for $\\lmax$ and Position Channels

We trained a Symphony on QM9 varying $\\lmax$ for the focus and position embedder, in Section G.1 (Figure 12). Due to computational constraints, we trained these models for $1,000,000$ steps each, which is $8\\times$ lesser than the original model reported in Section 3.3.

For the focus embedder E3SchNet, we did not see a significant increase in validity when going from $\\lmax = 1$ to $\\lmax = 2$, so we kept $\\lmax = 2$.

For the position embedder NequIP, we find a large jump when going from $\\lmax = 1$ to $\\lmax = 2$. Further increasing $\\lmax$ seemed to help slightly. For computational reasons, we kept $\\lmax = 5$.

Note that using our novel multiple channels trick, we are able to represent signals of higher angular frequency with a lower $\\lmax$. This seems to help for $\\lmax = 1$ in particular. However, we believe that the average number of neighbors in QM9 is too low for additional position channels to help at higher $\\lmax$.

To better understand the effect of adding multiple position channels, we explore the synthetic task of learning a random distribution on the sphere in Section F.2 (Figures 10 and 11). Here, we see that adding multiple position channels significantly helps the expressivity of the model. In particular, a $\\lmax = 4$ signal can be represented with $4$ channels of $\\lmax = 2$. As the computational complexity of many E(3)-equivariant models scales super-linearly with increasing $\\lmax$, we believe the use of multiple channels to represent higher-order signals may be useful beyond autoregressive generation in Symphony.

Training Times with Different Radial Discretizations

Reviewer msmy asked us about the effect of radial discretization on the training speed of the model. To investigate this, we changed the number of radii in the radial discretization, keeping all other hyperparameters the same. 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:

Training Times with Different Angular Discretizations

Reviewer UhEf asked us about the effect of angular discretization on the training speed of the model. To investigate this, we changed the discretization of $\\phi$, keeping all other hyperparameters the same. r_\\theta was kept as $180$ as in the original model. 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:

We do see an increase in the training speed, but not a linear one.

Sampling Times

Reviewers XBL9 and cp6D requested further details on the sampling time for Symphony compared to other models. These were present in Section 4.4 of the original manuscript. We restate these numbers here, as measured on a single NVIDIA RTX A5000 GPU:

Comparison to Simm, et al. (2021)

We greatly appreciate the public comment by Niklas drawing our attention to Simm, et al. (2021) [1] which also uses spherical harmonic projections molecule generation with reinforcement learning. We were not aware of this work during the preparation of this manuscript. Here, we highlight some differences between Symphony and their model:

• Given coefficients $c_l$, Symphony constructs the probability distribution $p$:

\begin{aligned} p_r(r) &= \sum \pi_m \mathcal{N(\mu_m, \sigma_m^2)} \\ p(\theta, \phi \ | \ r) &= \frac{1}{Z}\exp \left(-\frac{\beta}{k} |f(r, \theta, \phi)|^2 \right ) \\ p(r, \theta, \phi) &= p_r(r)p_{\theta, \phi}(\theta, \phi \ | \ r) \end{aligned} $$ where $k = {\\sum_{l = 0}^{\\lmax} |c_l|^2}$. $Z$ in both cases refers to the partition constant which normalizes the distributions. The extra factor of $k$ in [1] "regularizes the distribution so that it does not approach a delta function". However, in our case, we want to exactly learn a delta distribution (or a mixture of them). We show that their regularization is not ideal for our setting in Section F.2 (Figure 10 and 11).

• Our use of multiple channels allows representing signals of higher angular frequency than possible with the architecture in [1] as demonstrated by our experiments in "Ablation for Position Channels". [1] also noted the relation between the peakedness of the distribution and $\\lmax$, which our multiple channels trick helps with.
• To sample from the position distribution, [1] uses rejection sampling with a uniform base distribution. Our sampling strategy is much more efficient; we normalize the logits across the grid points with a softmax, and then simply sample from the grid points as a categorical distribution.
• The respective tasks are quite different; [1] provides the model with a "bag" of atom types with respective counts, while Symphony learns the distribution of atom types and counts. Further, due to the reinforcement learning setup, they have to sample many times from their model distribution to learn to optimize the true reward. In our case, we can directly learn the model distribution with the KL-divergence relative to the target distribution, without any sampling. This is why Symphony is able to learn much more efficiently; the model in [1] requires over $40,000$ sampling steps to learn the geometry of a single molecule.
• Symphony is trained and evaluated against the QM9 dataset of $\\approx 130,000$ molecules. Due to computational costs, [1] only tests their model on $9$ molecules. This makes it particularly hard to compare the performance differences at scale.

References:

• [1]: Simm, G. N. C., Pinsler, R. Csányi, G. & Hernández-Lobato, J. M. "Symmetry-Aware Actor-Critic for 3D Molecular Design". In International Conference on Learning Representations, https://openreview.net/forum?id=jEYKjPE1xYN (2021).

Other Changes

We found a bug in our evaluation of G-SchNet for the experiments in Section 4.4. These numbers have been updated. We have also added results for Symphony at later checkpoints; we are investigating why the fragment completion numbers drop while the unconditional generation performance increases. We also updated Figure 3 for clarity.

I thank the authors for their response. I have some follow-up questions, mostly regarding your wording.

• An energy-based model is typically [1] defined as exactly as a function $f$ such that $p(x)\propto \exp f(x)$, without having an explicit model for $p$, so you're definitely learning an energy-based-model. Could you clarify? It might be that you're confusing the term "energy-based models" this with the contrastive divergence sampling-based training methods for energy based models, but that's only one specific training way of training energy based models.
• You write that you don't need the normalized distribution for your training objective and write "This also avoids the key issue of estimating the partition function that occurs with energy-based models.", but you do need to estimate Z, as it depends on the parameters of $f$. So you are computing the normalized $\log p$ for your objective, estimated via integration, correct? Of course, as your space is only 3D, using an energy-based-model and integrating is a feasible strategy, different from higher-dimensional problems. I think the wording in your reply and revised appendix is not accurate, though.
• You write that you have access to the true distribution $q$. Do you mean that you know the density? Or do you mean that you have access to samples? Could you clarify?
• If you're integrating over 4 million points for each sample during training, how does this affect training times?
• Finally, do you agree that the method introduced in sec 3.4 can be interpreted as a mixture of energy based models? The mixture coefficients are proportional to the partition function for each channel.

Thanks!

[1] Teh, Yee Whye, Max Welling, Simon Osindero, and Geoffrey E. Hinton. 2003. “Energy-Based Models for Sparse Overcomplete Representations.” JMLR