Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: The paper develops a new method for learning densities over SU(n), and evaluates the performance for learning some idealized densities over SU(2) and SU(3), showing improvement over prior work of Boyda et al.

main_review: Unfortunately the methods of this paper lie far outside my area of expertise so there is very little of substance I can offer in review.  However my guess is that, as currently written, it will be accessible and hence of interest only to a tiny fraction of the NeurIPS community.  This sentence for example seems to convey the main significance of the paper:
"Learning SU(n) gauge equivariant neural network flows is important for obtaining good flow-based samplers of densities on SU(n) useful for lattice quantum field theory [2]." 
I have no idea what this means, and my guess is that many others will be in the same boat.  If the authors wish to reach more than just the theoretical physics/quantum field theory community, I would suggest that they put more work into writing this in a way that others can appreciate and learn something from that they can possibly apply to their own work.

limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The paper investigates the construction of equivariant flows on Riemanian manifolds. It proposes to learn such flows by using continous flows, where the vector field is the gradient of an invariant function. Defining invariant functions is usually much easier than defining equivariant vector fields. This framework is then tested on S^2 with the action rotations along a fixed axis, and on SU(N) with the action of  SU(N) on itself by conjugation, for N=2,3.
main_review: Summary of the review: On the positive side, the approach of first building invariant functions in order to get equivariant vector fields is elegant. The experiments also show clear improvement over past work. On the negative side, I find the that authors over-claim their contribution, and Theorem 4 (which is crucial to prove that this approach is general enough) is badly written.


== Positive points ==
Moving the problems from "building equivariant maps to building invariant maps" is not just elegant, it is also of high practical value since building invariant functions is often much easier than building equivariant ones. This is particularly true when we also have the requirements that our functions should be invertible. The theory justifying this approach is laid out clearly, and the proof of Theorem 1 looks correct.

The experiments provided by the authors are convincing, and show clear improvements over the reference paper of Boyda et al.

== Negative points ==
A big over-claim I have found in the paper is the claim of "learning gauge invariant densities over SU(n) in the context of QFT". This claim appears in the abstract, and is repeated throughout the paper. This is totally incorrect. This papers addresses the issue of learning conjugation equivariant densities on SU(n), and tests it for n=2, 3. While learning such densities is a crucial step *towards* flows for QFT in Boyda et al, it is definitely not the whole story and this paper shows no QFT experiment. In short, the authors are right to use QFT as a motivation for their work, but they should *not* claim to be learning gauge invariant densities, or be doing QFT in this paper.

I tried really hard to understand Theorem 4 about the sufficency of flows generated by invariant functions, but failed to do so. The authors should rewrite this. For example (but this might not be the only problem), on line 219, how can then take a derivative with respect to time, when time is not used to define \rho or \pi?

In the experiment section, the authors build conjugation equivariant flows on SU(N) for N=2 or 3 only, when Boyda et al also showed results for SU(9). Did the authors find that there method was harder to implement for N > 3?


==========================
Update following rebuttal and discussions with authors/reviewers
==========================

The discussion with other reviewers and with the authors has been particularly fruitful. I am happy to raise my rating to "7. Good Paper - Accept", conditioned on the few points that were raised and discussed during the reviews and rebuttal period:
* make it clear (including in the abstract and introduction), that the paper does not contain any Lattice QCD or QFT experiment. Instead, lattice QFT is a motivation for learning SU(3) conjugation equivariant flows.
* as pointed out by other reviewers, the relationship of Theo 1-3 with Lemma 2 in [1] and Theo 2 in [2] should be made much clearer. In particular, Theo 1-3 are extensions to the Riemannian case of ideas that were previously explored in the Euclidean case.
* Theo 4 needs to be completely rewritten, using all the added explanations and improvements discussed during the rebuttal.

[1] Papamakarios, G., Nalisnick, E., Rezende, D.J., Mohamed, S. and Lakshminarayanan, B., 2019. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762.

[2] Köhler, J., Klein, L. and Noé, F., 2020, November. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning (pp. 5361-5370). PMLR.
limitations_and_societal_impact: Continuous ODE flows are powerful but can be quite expensive, and ODE solvers do not usually deal with manifolds. The paper would benefit from having a short discussion on these points. For example, how many integration steps were used? Does this number vary over training? What ODE solver was used (Euler, some Runge-Kutta, ...)? What technique did the authors use to ensure their ODE solver stayed on SU(N) or S^2 (doing x + dt * grad_x is not going to stay on the manifold)?
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: The paper proposes a general method to construct equivariant continuous normalizing flows on Riemannian manifolds. For this, the authors built on (more pedestrian) results of Koehler et al using the fact that an invariant potential can be used to obtain an equivariant vector field. This vector field can then in turn be used to define a continuous normalizing flow which is equivariant under (a certain  left or right action of) the symmetry group. Finally, the push-forward of a invariant density under the equivariant flow is a invariant density with respect to the symmetry group.  Specifically, the manuscript considers the case for which G is (a subgroup of) the isometry group of the manifold. 

The developed framework is then applied to SU(N) Lie groups for which the conjugation by SU(N) is (a subgroup of) the isometry group. This application is of relevance in the context of lattice field theory, i.e. numerical simulation of (mostly strongly interacting) quantum field theories, for which ML techniques have recently gained significant attention. 
main_review: The paper is well-written, insightful and rigorous. The underlying theory is clearly explained (although the presentation necessarily requires a solid background in Differential Geometry). The proofs of the main results are derived in a transparent and pedagogical manner in the Supplement while the main text focuses on summarizing these main results and their interplay for the proposed methods which makes the paper very readable. 

Personally, I find the experimental sections a bit weaker. The LQFT application is of high relevance but in this context one is mainly interested in gauge symmetries, i.e. each lattice site can be conjugated by different SU(N) group elements, while the corresponding density is invariant under such local symmetry actions. The authors mainly consider the case of a single lattice site. Of course, this is a very important building block for a fully gauge-invariant flow and thus still highly interesting. However, it would be preferable if the authors had demonstrated the superiority of their approach over existing results by Boyada et al by training a flow to approximate the density of e.g. 2d quenched SU(N) lattice gauge theory comparing for example the ESS or the error of the topological susceptibility for fixed number of samples. The reported results would suggest that their method also helps in this case but an experimental validation of this expectation would be desirable.


Apart from this weakness, I find the manuscript very convincing and insightful. I also want to encourage the authors to release their code such that the community can benefit from it.

A few minor suggestions:
L136-137: d and \partial an the equation should be used consistently in the numerator and denumerator
L189-190: larger spacing between the two equations would improve readability
L214, L216-218: distribution -> density
L506-507: Du -> D_u


========================== Update after rebuttal and discussions with authors/reviewers ==========================
I want to thank the authors and also the other reviewers for a very insightful discussion period. I want to strongly encourage the authors to implement the changes already suggested by the other reviewers. 
limitations_and_societal_impact: Both limitations and societal impact are adequately discussed.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper proposes an extension Continuous Manifold Normalizing Flows to Equivariant Manifolds where the the equivariance is with respect to an isometry group of the manifold. To this end, the paper lifts existing tools to a more manifold specific setting and empirically demonstrates the utility of the proposed approach on sampling from SU(N).  Overall, the paper is relatively well written and easy to follow. I found the main ideas, such as Theorems 1-3, interesting---albeit not extremely novel---and correct. I believe the empirical results, while a bit on the simpler end, are illuminating and see no reason to doubt the correctness of the numbers.
main_review: First I'd like to highlight that the paper is quite easy to read and I had little trouble following the main results including the proofs in the Appendix. The main highlight of this paper for me were the two invariant potentials introduced which I believe are more of a central contribution than the developed theory. 

This brings me the central complaint of the paper which is limited novelty. To be perfectly honest Theorems 1-3 are not sufficiently novel. In fact, I would not even elevate them to theorems but rather Lemmas as done in [1]. In fact [2] proves much of the results in Theorems 1-3 for group actions on Euclidean spaces. Lifting these to more general manifolds requires a little bit of technical nuance but I don't believe it is that difficult (e.g. instead of gradient, we use Riemannian gradient). Theorem 4 on the other seems to be more novel and interesting. I think it pokes at a larger but not fully explored issue on the universality of learning equivariant flows on manifolds. However, I believe this should be investigated further in the context of the proposed invariant potentials. A more powerful result would be to show that for a specific class of invariant potentials we can learn any equivariant density for some isometry group. Right if my understanding is correct the current result highlights that there exists an invariant potential where the KL is minimized and as a result this is a good learning target. 

Finally, the experiments are also a bit light. It would've been interesting to propose new manifolds where existing approaches would definitely not work. Also, I would argue that there are missing baselines for the S^2 example. Finally, a lot of the results in the appendix are quite interesting and would've made for a stronger experimental section.

Overall, I believe this paper in its current narrative leans too heavily on borrowed theory which is extended in a natural but straightforward manner. Consequently, the main results that I found interesting such as the two proposed invariant potential functions and Theorem 4 are not given the full attention that they deserve. Consequently, I'm currently leaning towards a weak reject but I am willing to change my score if a good argument can be made against the points above and below.


Detailed Feedback:
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
line 61 - Kohler et. al formulation for equivariant vector field is fairly general.
Section 3.2, this is technically correct but it would be more fitting and match the title if you decomposed the change of variable and explicitly show the role of the Riemannian metric

Theorem 1. Is a strict generalization from Kohler[2] & Papamakarios[1] review paper
Theorem 2-3 are direct consequences
R_g is a representation of the group. Is it a linear represenation? You must be more precise here.
- line 196 G-equivariant CNF, I can construct discrete NF's on manifold without this condition.
- you could also consider the prior density introduced in Kohler et. al which is on M-1 dimensional subspace.
- 5.2.2. you could also use uniform density here no?
- 359, Can you elaborate on what exactly you had to tune? This is an interesting point for the general community.


Point 2 line 45: We can always restrict the flow to a pre-defined support. This can be done by a projection operation. For example this is the underlying principle of a Restricted/Wrapped Gaussian. I don't believe this is an insurmountable limitation.

Point 3. Line 51. Can you expand more on this unnaturalness? Give a more precise characterization and explicate why exactly this is undesirable? It is possible to design flows over discrete structures in the conventional non-Manifold sense. Can we not leverage a similar design principle here? Please argue against this point.

Contribution 3 line 71. This is not really a new contribution but moreso the benefit of invoking symmetries when doing generative models. Other papers have also remarked the same observations. This paper also acknowledges this fact line 103-104


line 207. I disagree. It does not strictly reduce the problem as one can construct G-invariant densities but not going through a G-invariant potential function. This is a clear mistake. I believe what the authors intended to say that in this particular context it is sufficient to use G-invariant potential function. But indeed this is not necessary, for example I can construct an equivariant map that is not the gradient of an invariant potential function and still arrive at a G-invariant density. 

Theorem 4. I believe I don't fully agree with the interpretation here. I agree that there may exist a G-invariant potential for every target distribution on a closed manifold. But you must show that the construction of one is universal. This point is missing.


For 6.2 The obvious missing baseline is NF's on Spheres and Tori in addition to the Manifold flow.

Minor Presentation issues:
Line 141, you haven't defined what a symmetry is.
Line 146: This is obvious, and here $f$ is not distinguished as a flow. 
Overall 3.3 is a bit sloppy in notation and could be cleaned up a little bit.
X in line 174 is a vector field, earlier in the background it was a manifold.

Typo in line 506 in proof. I believe it should be D_u not Du R_g(x), ... in middle equality.

5.1 -> 5.2 This is confusing, first you assume that you're given an invariant potential function then immediately you design two new ones? From a narrative perspective this is not very clear.


Questions:
1. The prior for S^2 is only on the z-axis right ---i.e. its a 1-d uniform density?
2. 5.3 Isn't the conventional duality of Normalizing Flows. We either need to do density estimation or we need to generate high quality samples. I don't see how this is particularly unique to Manifold Flows. 

------------------ Post-Rebuttal Updates -----------------------------------------------------------------------------------------------------------------------------------------
After fruitful discussion with the other reviewers and the detailed author comments to my questions, I have raised my score from 5->6. This is conditioned on the authors making significant changes to the writing as mentioned in my original review. 

1.) The paper must do a better job of putting this work in context with [1] and [2] and clearly stating these are extensions of these ideas to the Riemannian case and not that they are proposed here for the first time. This is currently ambiguous in the original draft.

2.) Theorem 4 at submission time has many interesting nuances that were uncovered in the discussion. I suggest this theorem be re-written along with the proof which has a few technical issues that need to be addressed (e.g. smoothness in the limit) as pointed in the discussion.

I believe this paper has the potential to be very interesting to the community but the key detractor at the moment is the writing and more generally the unclear narrative in context with prior work. Once these issues are fixed---and if accepted---this paper would be interesting to the NeurIPS community.

[1] Papamakarios, G., Nalisnick, E., Rezende, D.J., Mohamed, S. and Lakshminarayanan, B., 2019. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762.

[2] Köhler, J., Klein, L. and Noé, F., 2020, November. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning (pp. 5361-5370). PMLR.

limitations_and_societal_impact: The authors have adequately addressed their limitations and discussed negative societal impact.
review_title: null

Meta-review: The paper proposes a theoretical framework for constructing normalizing flows on Riemannian manifolds that are equivariant with respect to the action of an isometry group. This is an interesting contribution and highly relevant for distribution modelling under known symmetries in non-Euclidean spaces, which comes up often in problems in physics.

Overall, the reviews are positive so I'm happy to recommend acceptance. There has been a fruitful discussion between the reviewers and the authors that I hope will improve the final version further.

**Comments to the authors**: I would like to see the improvements suggested by the reviewers in the final version of the paper. For clarity, I'm copying below the changes requested by the reviewers that I'd like to see implemented:

1. Make it clear that the paper doesn't contain an application to QFT or Lattice QCD, but uses these as motivation only.
2. Moderate the novelty of Theorems 1-3 and make it clear how they relate to similar results from previous work.
3. Revise Theorem 4 and its proof following the discussion with the reviewers.

I would also like to highlight the review by Reviewer e1R7 and specifically this statement: "as currently written, [the paper] will be accessible and hence of interest only to a tiny fraction of the NeurIPS community". Indeed, I'd encourage you to reflect carefully on how this work can be communicated effectively so that it can maximize its future impact.
Decision: Accept (Poster)