SymDiff: Equivariant Diffusion via Stochastic Symmetrisation

Thanks very much for your review, and we are glad to hear that you found the paper clear and high quality. Please see responses to your individual points below.

Second point

Unfair comparison: I didn't fully understand for SymDiff+++, ++, and +, why you cannot choose the same number of parameters of EDM. It seems your model SymDiff+ is already larger than EDM+++, but with worse performance.

Please see our overall response to all reviewers above. In summary, we believe that overall computational cost (computational time and vRAM) is what ultimately matters in practice, and since the neural net architectures here are so different (GNN vs. Transformer), parameter counts do not give a good measure of this. We have provided additional benchmarks above that demonstrate our method yields the same or improved performance at a lower computational cost.

First point

Performance gap: from Tab2 and 3, seems the model is not parameter efficient compared with baseline models. When reducing the model size, the performance will have a larger drop than the baseline model.

Thanks for pointing this out. We think that Table 2 and 3 were confusing, and have now streamlined into Table 3 in the updated manuscript. For convenience, we have copied Table 3 below.

The key point here is the following:

• The SymDiff model from Table 1 is the largest SymDiff model in Table 2 and 3.
• The EDM model from Table 1 is the smallest EDM model from Table 2 and 3.

In effect, these tables show the result of making the SymDiff model from Table 1 smaller, and the EDM model bigger. In essence, we have

$$\text{SymDiff}^{--} \prec \text{SymDiff}^{-} \prec \text{SymDiff} \prec \text{EDM} \prec \text{EDM}^{+} \prec \text{EDM}^{++}.$$

where $\prec$ denotes smaller computational cost and we have renamed the models compared to the first version as follows:

$

&\text{SymDiff}^{+++} \to \text{SymDiff}^{}, \ \text{SymDiff}^{++} \to \text{SymDiff}^{-}, \ \text{SymDiff}^{+} \to \text{SymDiff}^{--}, \\ 
&\text{EDM}^{+++} \to \text{EDM}^{++}, \ \text{EDM}^{++} \to \text{EDM}^{+}, \ \text{EDM}^{+} \to \text{EDM}^{}

$

So we are not comparing like-for-like: the smallest SymDiff model here (SymDiff--) is much smaller in terms of computational cost than the smallest EDM model (EDM).

From this perspective, SymDiff again is more computationally efficient. From the updated version of Table 3, SymDiff$^{-}$ still does better than EDM in terms of molecular stability, but has considerably lower computational cost. Likewise, SymDiff does better than even the largest EDM model, EDM$^{++}$, again even with considerably lower cost.

Thank you very much for the positive feedback and review score. We are glad that you found our paper to be well-written and useful to the research community. Please see below for our response to the points raised within your weaknesses and questions section.

Further comparisons with SymDiff

I suggest the authors include a comparative analysis between SYMDIFF and traditional methods ... Additionally, while the paper claims that SYMDIFF is computationally efficient, it doesn’t provide empirical evidence to support this ...

Please see our overall response to all the reviewers above for an extended comparison between SymDiff and the other methods that we benchmark against in the paper.

Miscellaneous questions

Do you have insights into how many training epochs are necessary for SYMDIFF to effectively learn equivariance in the molecular generation task?

The model used in SymDiff is (stochastically) equivariant by construction (Proposition 1 in the paper). The flexibility of our method is not that the model can learn to enforce equivariance to a greater or lesser extent, but that we have more architectural freedom within our model. This is because the equivariance in SymDiff does not require any intrinsically equivariant components. We show in the overall response above that this leads to a reduction in computational complexity while either maintaining or increasing performance compared to other methods which lack this architectural freedom.

How does the convergence speed of SYMDIFF compare to that of intrinsically equivariant models in learning equivariant properties?

From above, the model does not learn to enforce equivariance as this is already enforced by construction. However, we can consider the convergence rate of the models in terms of the metrics we use. For example, we observed that the convergence rate of $\text{SymDiff}$ (31M) in terms of molecular stability was comparable to the baseline EDM model and was faster when comparing against the larger variants of EDM from Table 3 in our paper.

In the context of molecular generation, do you consider the equivariant component of the task to be a minor or major challenge compared to learning the underlying chemical laws necessary for accurate atom placement?

Equivariance is an important inductive bias for handling geometric structure in molecular generation tasks as it provides a principled way to reduce the complexity of the generation task by processing the symmetries of molecules in a systematic fashion. However, equivariance can only handle global symmetries within the geometry and cannot directly inform the model about the chemical laws which are the main factor determining the structure and function of molecules. Hence, enforcing equivariance is an important yet less difficult challenge compared to learning the fundamental chemical laws which underlie molecular structure.

Related to that, it is not sufficiently well discussed, explained and analyzed how exactly this situation qualitatively changes when is learnt (see point above).

Please see our discussion of the distributions that $\gamma_\theta$ learns above.

Experimental results

Related to the last point above, overall, it seems that SymDiff can be seen as a sophisticated data augmentation strategy (as also pointed out by the authors).

As mentioned above, we would view this as (in a way) the other way around -- data augmentation is a special case of SymDiff. For example, Proposition 3 shows that data augmentation corresponds to SymDiff with the choice of $\gamma_\theta$ being the Haar measure, where we do not only transform the input but also transform the output with $\gamma_\theta$.

Further, the authors make the point that the comparison to EDM is fair. However, the very different network sizes still make a fair comparison somewhat difficult.

Please see our overall response to all reviewers, which responds to this point.

Miscellaneous question

In 4.2, the authors write that, given the size of GEOM-Drugs, they train a smaller model than on QM9. But the GEOM-Drugs dataset is larger and more complex than QM9. Wouldn't one want to train a larger, more expressive model in that case, and not a smaller one?

Yes, you are right that it would make sense to train a larger model for a more complex dataset. The reason for our choice was that we wanted to match [2] which trains a smaller model on GEOM-Drugs. Their motivation for training a smaller network was due to computational limitations, as they did not have enough compute to train a larger EDM model on GEOM-Drugs.

We have now ran our SymDiff model from the QM9 dataset on GEOM-Drugs. We include the updated results below and have also now included these results in the manuscript (Table 4).

[1] Villar, Soledad, et al. "Scalars are universal: Equivariant machine learning, structured like classical physics." Advances in Neural Information Processing Systems 34 (2021).

[2] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning (2022)

Firstly, thanks very much for your careful reading of our paper and insightful comments and questions. We are very glad you found the paper high quality, interesting, and easy to follow.

Please see individual responses to the points you have raised below.

More discussion of intuitions and more experimental analyses:

What sort of $\gamma_\theta$ and $f_\theta$ functions are we intuitively expected to learn in different situations?

What the model ends up learning in practice is to some extent opaque. However, we would argue that this is not necessarily unique to our method, but applies to most neural networks. Our main intuition is that:

• SymDiff is always guaranteed to produce an equivariant diffusion (Proposition 1), which is known to be a good inductive bias for many applications.
• SymDiff encompasses data augmentation as a special case (Proposition 3), which is known to be a sensible baseline.

Practically speaking, our approach is then simply to make the model as large and flexible as possible, and to let it choose how it wants to actually optimise the loss. Assuming we optimise it well, we can then expect it to do at least as well as data augmentation, although in practice it may choose to learn something else that performs better. This is borne out empirically by our experiments: SymDiff did work better than data augmentation (and moreover better than intrinsically equivariant architectures).

Moreover, is there anything preventing the model from simply learning $R = R_0 f_\theta (\ldots) = I$, this is, no ''augmentation'' at all?

This situation cannot occur. We have $R = I$ iff

$$f_\theta(R_0^T \cdot z_t, \eta) = R_0^T.$$

In other words, $f_\theta$ has to learn to map $R_0^T \cdot z_t \mapsto R_0^T$. However, this is not possible because there are many pairs of $z_t$ and $R_0$ that map to the same $R_0^T \cdot z_t$ in general.

One thing that does hold is that if the ``backbone'' of the reverse process $k_\theta(z_{t-1}|z_t)$ is already equivariant, then its symmetrised version is just $sym_{\gamma_{\theta}} (k_{\theta})(z_{t-1}|z_{t}) = k_{\theta}(z_{t-1}|z_{t}),$ so that symmetrisation has no effect (Proposition 5.1 of Cornish (2024)). This is a good thing in terms of expressiveness, since it implies that the symmetrised model can express at least as many equivariant kernels (by choosing $\theta$ appropriately) as the backbone $k_\theta(z_{t-1}|z_t)$ can.

I would like to thank the authors for the comprehensive reply to my review and I appreciate the additional explanations and results. Some of my concerns have been addressed.

The authors promised to run additional experiments and to update their results:

We do not have the final trained model yet, but we will update this before the deadline.

as well as

We also agree that the case where $\gamma_\theta$ is not applied during sampling is an interesting ablation. We will run this and report back.

I think these experiments would be helpful and insightful and I would like to kindly follow up whether the authors were able to run these experiments.

I also have one rather high-level follow-up question. I believe I now understand the framework reasonably well mathematically, and I agree that the framework can be seen as a generalization of regular data augmentation. Formally, the approach can be considered strictly more flexible. However, I still lack some intuitions: Why exactly would it ever be better to use such "learnt data augmentation" instead of augmenting with random rotations (or, equivalently, the Haar measure)? Intuitively, when being confronted with a problem exhibiting such rotation symmetry, one would assume that random rotation augmentation is exactly the right type of augmentation. Why should there be a preferential direction? However, the presented work does seem to improve upon this. Hence, what quality does the data distribution need to have such that such a more complex augmentation via a more complex $\gamma_\theta$ is beneficial? And, on the contrary, for what sort of data distribution would you expect that the ideal $\gamma_\theta$ is simply the random Haar measure, without refinement? Could you provide some examples and more intuitions?

Ablation with $\gamma_\theta$:

Here are the results of an ablation where we evaluated our SymDiff model three times where we now did not apply $\gamma_\theta$ at sampling time (SymDiff-GammaOff). Please find the results in the table below.

From the above, we see that not applying $\gamma_\theta$ during sampling causes only a small decrease in performance. This observation suggests that our SymDiff model has learnt an approximately equivariant $k_\theta$ in this case. This is in line with our earlier explanations.

Distribution of $\gamma_\theta$:

Please find an image of the distribution of samples from $\gamma_\theta$ from the model described earlier, which we call SymDiff-Scalars, here. We have now trained this model to completion, and the performance metrics for this model are reported in the table below.

From the above, we observe that SymDiff-Scalars overall performed worse than SymDiff, although it still outperformed EDM. Additionally, we see that the distribution of $\gamma_\theta$ has a much tighter spread than the one by SymDiff. These results further support our earlier discussion, showing that learning more complex $\gamma_\theta$ can also be beneficial.

We hope that these responses address your outstanding concerns (please let us know if not), and if so that you will consider increasing your score.

[1] Kaba, Sékou-Oumar, et al. "Equivariance with learned canonicalization functions." International Conference on Machine Learning. PMLR, 2023.

[2] Dym, Nadav, Hannah Lawrence, and Jonathan W. Siegel. "Equivariant frames and the impossibility of continuous canonicalization." arXiv preprint arXiv:2402.16077 (2024).

[3] Allen-Zhu, Zeyuan, Yuanzhi Li, and Yingyu Liang. "Learning and generalization in overparameterized neural networks, going beyond two layers." Advances in neural information processing systems 32 (2019).

[4] Kim, Jinwoo, et al. "Learning probabilistic symmetrization for architecture agnostic equivariance." Advances in Neural Information Processing Systems 36 (2023)

[5] Puny, Omri, et al. "Frame averaging for invariant and equivariant network design." arXiv preprint arXiv:2110.03336 (2021).

I would like to thank the authors for their additional explanations and experimental results. I think these results are indeed helpful to provide context and understanding. I would like to ask the authors to add these experiments as well as the explanation that considers the two extreme cases of $\gamma_\theta$ (I found this very helpful to get some intuitions) to the final paper, for instance in the appendix if space is a concern.

I am willing to raise my score and support publication of the work.