SE(3) Equivariant Augmented Coupling Flows

Summary

The manuscript proposes a normalizing flow architecture that is i.) equivariant wrt to the euclidian group SE(3) and permutation group S_n and ii.) is based on coupling blocks. Existing work on this specific equivariance has relied on continuous normalizing flows which rely on a Neural ODE instead of a coupling architecture to construct a bijective map. Coupling-based architectures have the advantage that they typically faster - both in training and sampling phase. The proposed architecture combines various previously explored ideas in an elegant manner. Specifically, it relies on augmented flows to realise the equivariance, similar to Klein et al and first proposed in by Huang et al., and achieves SE(3) symmetry by relying on zero centre of mass subspaces. The experimental part gives a detailed comparison to existing baselines and demonstrates the advantage of the proposed architecture over CNF-based flows as well as the proposal by Klein et al. Overall, I found this paper an insightful and solid contribution as well as a nice read. The proposed architecture is definitely worth exploring in my own research. I would encourage the authors to publish their code to maximize the impact of their paper.

Strengths

- Very clearly written manuscript - A thorough and extensive experimental analysis - Proposed architecture is a elegant combination of well-known mechanisms achieving full equivariance - Both likelihood and energy-based training are considered

Weaknesses

- The ms mainly combines previously proposed ideas (albeit in an innovative and elegant fashion). In particular, its proposal is relatively similar to the one of Klein et al.

Questions

- How is the KLD normalizer in the tables estimated? - Did you investigate if any mode-collapse is present for the reverse KL training? In particular, did you estimate the forward ESS, i.e. 1/E_p w where p is the target/data density and w the importance weight? - Cartesian projection seems to lead to numerical instabilities. Do you have an understanding of why they arise?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Yes, I think the limitations are sufficiently summarised although I would have appreciated a dedicated section.

Summary

The paper proposes a method of producing a special Euclidean group - SE(3) - equivariant method of training coupling flows by augmenting the input space then conducting standard normalizing flow transformations in a rotationally invariant space by means of an invertible transformation. This is applicable particularly to Boltzmann distribution problems. Update: I have read the rebuttal, and the additional details seem convincing. I have raised my score accordingly.

Strengths

The proposed SE(3) Equivariant Flows offer a specialized framework for dealing with atomical problems which utilize SE(3) structure. For the task of approximating the Boltzmann distribution, the model works quite well compared to some of the baselines established, and features higher speedup. In addition, it admits a FAB-version of training the model without maximum likelihood.

Weaknesses

Seeing the disparity between the NON E-ACF and the other models on the QM9 dataset, I would have liked to see some comparisons against the E-CNF DIFF baseline on the Alanine dipeptide baseline as well. In addition, the baseline speedup comparison is only against the diffusion model for which they claim a 500x speedup. Though this speedup might be due to the advantages proposed by the framework, a large portion of it might be due to the denoising process intrinsically taking up a lot of time. I would like to see the authors either talk about speedups compared to other standard CNFs, or along the lines of a performance/speed tradeoff that can be made with respect to the results shown in the tables.

Questions

In the paper, the primary baselines and comparisons are against CNFs, primarily due to the fact that existing literature on SE(3) equivariant flows have mainly through the CNF regime (though the authors do also mention equivariant residual flows as well). How does enforcing equivariance vs not enforcing affect subsequent performance on the benchmarks in the case of coupling flows?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

Yes.

Summary

This paper proposes a novel method for developing an SE(3) equivariant coupling flow that works on the Cartesian coordinates of atoms, crucially enabling quick sampling and density evaluation. The standard coupling architecture is unable to operate on the Cartesian coordinates with the Special Euclidean group SE(3) and permutation invariances of physical systems, hence the researchers introduce an advanced coupling flow that preserves these by using coordinate splits along added augmented dimensions. The method includes mapping atoms’ positions into learned SE(3) invariant bases at each layer, applying standard flow transformations such as monotonic rational-quadratic splines, and then reverting to the original basis. This proposed model demonstrated competitive performance when trained on the DW4, LJ13, and QM9-positional datasets, offering much faster sampling speeds. The researchers also became the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms. Furthermore, the study introduces a new approach to making coupling layers SE(3) equivariant by augmenting the input space with auxiliary variables that can be acted upon by SE(3). This allows the atom positions to be updated based on these auxiliary variables by first projecting the atoms into an SE(3)-invariant space, applying a standard normalizing flow transform and then projecting its output back onto the equivariant space. The researchers showcased that when trained by maximum likelihood, their flow matches the performance of both existing SE(3) CNFs and coupling flows operating on internal coordinates on molecular generation tasks. This flow is over 100 times faster for sampling than SE(3) CNFs. The study also demonstrates the flow's ability to be used in the energy-based training setting on the DW4 and LJ13 problems, showcasing its unique advantages over existing methods.

Strengths

This paper has a few strengths that I think are worth mentioning: 1. The overall idea of using an augmented normalizing flow and baking in symmetries in the product space is a novel direction and sensible idea. 2. In general the presentation of the paper is of high quality (except a few notational quirks). The writing is easy to follow and is supplanted by a great architecture diagram in Fig 1. 3. The experimental results show a lot of promise and validate the design of the flow.

Weaknesses

**Technical Details** I had a difficult time following the proof of equivariance it uses slightly non-standard notation for defining the group action and how it works. Specifically, if $\theta$ is invariant and the standard flow is invariant I'm not sure how this could lead to an equivariant flow. That is we often require an equivariant pushforward to map an invariant base density to an invariant target density. There are two ways you could convince me: 1.) Randomly sample rotations from $g \in SO(3)$ and numerically prove equivariance of the overall flow is satisfied. 2.) Or you could provide a lot more detail and use more standard notation. For example, I think you mean $\tau_{\theta}$ is an invariant density but is actually an equivariant pushforward map. Otherwise, I'm not sure how this goes through. The idea of using augmented variables is an interesting algorithmic choice, but using a standard coupling flow for $\tau_{\theta}$ is only 1 option. In fact, I believe you could use any flow architecture for $\tau_{\theta}$ as such it would be great to see if other mainstay flow architectures could be employed here (e.g. ResFlow, MintNet, BNAF). **Experiments** The biggest question I have from the experiments section is from Tab 1. It seems for DW4 training with samples is not competitive with E-CNF, but with training with energies, we find that the proposed approach clearly beats E-CNF. Can the authors provide more intuition and justification here? One of the main benefits of using symmetries as an inductive bias is that it leads to increased sample efficiency. This is not explored in the current manuscript. One could consider an additional experiment with samples where an equivariant model and a non-equivariant model are trained with various different dataset sizes. One would expect the equivariant model to be more performant in the low data regime. Can the authors try this experiment? **Minor** - In the text you use CoM-Shift but in equations, this is called ShiftCoM. Better consistency in naming conventions should improve readability.

Questions

1. Can you provide more details on the graph neural network $h$ in the main paper. It is not clear to me how this network outputs invariant parameters $\theta$. Is this the sum of all the node embeddings? 2. It is not clear to me how exactly the GNN $h$ outputs equivariant features. Are you using an EGNN and taking the velocity vectors? More detail on this in the main paper would be appreciated. 3. In Fig 1. The $\oplus$ operation is actually a concatenation rather than an addition right? This could be made more clear. 4. In line 210 it is stated "We first subtract the observations’ CoM from both the observed and augmented variable", but previously in line 134 it is stated, "Importantly, we do not restrict a to be zero-CoM." This appears to be a contradiction?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Summary

This paper tackles the problem of sampling molecule configurations using deep generative models. The learned distribution should be permutation invariant (in terms of permutations of the particles in $\mathbb{R}^3$) and SE(3) invariant (i.e., invariant to rotations and translations applied to the particles' cartestian coordinates). At the same time, the deep generative models should be fast to sample from and allow quick density evaluation. Continuous normalizing flows (and relatedly, diffusion models) can be built with these symmetry constraints, but are slow to sample from. Coupling flows are fast at sampling/density evaluation but existing approaches rely on internal coordinates rather than explicit cartesian coordinates. Generally, the goal of this paper is to endow coupling flows with the SE(3) and permutation symmetries of existing CNFs, while preserving the speed benefits of coupling flows. Coupling transforms pose a problem for SE(3)xS_n symmetry because splitting the coordinates will break either rotation or permutation symmetry, depending on which dimension one splits on. The paper gets around this problem by introducing augmented dimensions, which are instantiated as additional observations $a$. The original observations are transformed conditioned on the augmented variables and vice versa, with equivariance enforced by having a particular structure to the transformation of x given a (Eq 4). In experiments, the paper shows that the proposed method performs similarly to CNF based methods while preserving the computation benefits of coupling flows.

Strengths

This paper presents a solid and principled design for making coupling layers (in a normalizing flow) SE(3) and permutation equivariant, a property of practical interest in situations such as modeling molecular configurations. The experiments show that their approach preserves the computation benefits of coupling flows and the performance benefits of symmetry priors.

Weaknesses

Although the preliminaries are set up well, the presentation of the method starting from 3.1 is somewhat confusing and could benefit from more intuitive description (eg, of Prop 3.2) and concrete examples.

Questions

N/A

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The paper is not particularly clear on limitations of the proposed approach.

Summary

This paper constructs equivariant normalizing flows that are discrete to sidestep speed issues with equivariant continuous normalizing flows. The method does this by augmenting the input with some an additional variable to avoid representation issues.

Strengths

* The method is a reasonable extension of standard equivariant normalizing flows literature that seeks to address the fundamental problem of balancing expressivity and speed. * The method produces good results (ie likelihoods and other metrics are trained to a very good value while not using much compute).

Weaknesses

* I still don't understand the usage of $\gamma_r$. Line 160-161 says that $\gamma_r$ is a projection to an invariant space, but the definitions in Line 162 don't seem to accurately describe this (which should be $\gamma_r \cdot gx = \gamma_r \cdot x$ if I'm correct. Also, what is the definition of $\gamma_r^{-1}$, since I assume that necessarily $\gamma_r$ must got a lower-dimensional subspace. It would also help tremendously if the "Choice of projection $\gamma$" section was written to be more mathematically precise. * This "invariant feature projection" that $\gamma_r$ is known to be a rather weak form of constructing invariant functions. Some expressivity proof would be helpful to allay these types of concerns. * If $\gamma_r$ projects to something like S^2 angles, which I think might be happening for sphere-proj, then wouldn't one need a manifold-based flow model to accurately build a flow on this? See e.g. [1, 2]. * To evaluate densities, the paper computes using only 20 samples for a. This seems like a low amount. * Is there a reason why results aren't reported for DW4 or LJ13 for diffusion in table 1? I assume this is trained from samples, which should allow for diffusion. * In the appendix, diffusion seems to be trained for a very short amount of time (less than even the for the proposed fast ACF methods. From the literature, diffusion tends to take a while to train but this tends to improve log-likelihood. * For table 2, why can't the negative log-likelihood be compared? Also I think its KL divergence is lower than the proposed method (at least from the table). [1] https://arxiv.org/abs/2002.02428 [2] https://arxiv.org/abs/2006.10254

Questions

Nothing outside of the listed weaknesses.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

1 poor

Contribution

2 fair

Limitations

Yes

Further questions

Thank you for the clarifications. ## On Likelihoods Why does the choice of space make the likelihoods incomparable? Is it fundamentally mathematical, or more for a comparison purposes. It seems like they should be comparable (since the internal coordinate model is just the 3d coordinate system modulo symmetries).

Re: Rebuttal

I thank the authors for their responses. The notation is clearer in my mind now. I still contest that there are more clear ways of defining the group action which is in more line with the rest of the literature but that's not important. **Biggest Question** My biggest question, which I think the authors misunderstood is why in Table 1. do we see for DW4 E-CNF is way better (it's bolded) than your method? This is also true for the global response in the low data regime. I do not buy the biased MD simulations aspect here because it's orthogonal since you have a dataset---and are not training using energies. This falls perfectly into the current generative modeling setup and Maximum Likelihood framework. Thus I also disagree with the authors on the need for MintNet and ResidualFlows. It is true that coupling flows have fast inference but everyone in the literature knows that they have subpar performance---despite being universal density approximators. At present, I do not know how performant the presented flows are compared to E-CNF. Furthermore, to demonstrate scalability the experiments would need to be conducted on a much larger dataset which is not done in the current work, so while this claim might be true it is not demonstrated.