Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groups

We thank the reviewer for their thoughtful and detailed comments. We are glad that the reviewer found the contributions of our paper original and significant, and the experimental results comprehensive and informative. We would also like to thank the reviewer for the references provided, which have now been included in the paper. We address the weaknesses and minor comments below:

• W1: We completely agree with the reviewer - in fact better constructions of energy functionals for such a task will likely be the main direction of improvement. Effectively we would wish to align minimas of such an energy with the training set - doing this directly is not that simple and should be a fruitful direction for future work.
• W2: We agree that it could potentially be an interesting comparison - however ensuring fair conditions for such a comparison are unclear in relation to sizes of networks, dataset sizes and a multitude of other factors. Particularly, this falls slightly outside the scope of what canonicalization is meant to achieve, as the main benefit is the ability to turn already pre-trained models equivariant without retraining (or at most minor finetuning). Overall, definitely interesting, but may fall too far outside the scope of this paper.

Minor comments and suggestions:

• Typos have been fixed.
• “On the orbit distance constraint (Line 357), the authors may find [10] relevant, as the approach uses invariant polynomials for linearly reductive groups (Lines 2227-2228) to measure orbit distance.”
  • This paper [10] is very nice, and their method is in our language of weighted canonicalization. They are able to leverage the fact that they look at reductive linear groups acting algebraically on vector spaces, which allows them to use many classical results on invariant polynomials. Arbitrary lie group actions of the kind that we consider do not have such a nice theory, so it would be unfeasible to write it so cleanly. It is definitely a valuable comparison to make with our method though and we have added that.

As the discussion period draws to a close, we'd like to check in on whether you've had a chance to review our responses and have any follow up questions?

We hope that our reply clarifies and alleviates the reviewer’s concerns. If this is the case, we kindly ask the reviewer to consider raising their rating, given that they are acknowledging the novelty, the strengths and the contributions of our paper.

Thank you for the response and incorporating the suggestions on typos and reference. I have read other reviews and responses and would like to retain my supportive rating.

On W2, I believe it is reasonable to ask whether a new canonicalization method offers performance gains over existing equivariant networks (if available; like in Table 1) because, in the end of the day, the goal of canonicalization is equivariant learning, and equivariant networks are established methods that solve the same problem. But I agree that this would require substantial effort for the tasks in Section 4.3, especially if the argument in Line 48-51 applies.

• “Can you clarify what role does the Haar measure play in the construction of weighted canonicalizations for the non-compact case? Does unimodularity play a role here? Is the case where the modular function is unbounded pose any additional obstructions?”
  • The Haar measure plays no role in constructing weighted canonicalizations. It plays a role in proving weighted closed canonicalizations are the sequential closure of weighted orbit canonicalizations, but we only really need it to have an (nice) $\sigma$-finite measure such that we can decompose a weighted canonicalization into easier to deal with terms. Replacing the left Haar measure with the right Haar measure in this proof doesn’t affect the result, though the constructed approximation will be different.
• “Is it not possible to make use of a riemannian structure on the group and have the energy minimizer be defined in terms of geodesic distance? It seems the current proposal already looks to work within the tangent space of the groups involved.”
  • We agree that utilizing more information of the group can prove beneficial. It is unclear however how exactly the geodesic distance can be helpful in minimizing the energy. In order for the method to be an orbit canonicalization, the problem has to be of the form eqn (2) or similar (see Kaba et al. 2023). Could you clarify how exactly you think this may be used?
  • In addition, when moving away from the compact case, being able to put a reasonable Riemannian structure on a Lie group is not guaranteed. Namely, we would want a bi-invariant metric on the Lie group (i.e. preserved by left and right multiplication) because this means that the Riemannian exponential map defined by geodesics and the Lie algebra exponential map defined by the group structure will agree. This is not always possible for non-compact groups. Putting a non bi-invariant metric on them is always possible but this will not give us the information about the group structure that we might care about. In specific cases, when knowing the group one could imagine tailoring the Riemannian metric to give nice geodesic distances that play well with energy minimization but this would need to be done case-by-case if it is even possible.
• “Canonicalization methods IIUC deal with global invariance/equivariance, as opposed to e.g. lifting + regular/steerable convolutions which could deal with local equivariance (e.g. imagine 2 objects in an image rotated at different angles). I'm wondering if the authors would find this distinction worthy of highlighting.”
  • We agree that it is very interesting to consider local equivariance in addition to the global equivariance considered in our work, and we have added references to work in which this local equivariance of steerable convolutions is discussed and a discussion of the distinction between local and global equivariance. In the context of steerable convolutions, it is worth noting that the local equivariance is not unconditional, but depends on the spatial extent of the filters being sufficiently small compared to the distance between features of interest. We believe that a proper treatment of such local transformations would naturally lead to studying “global” equivariance with respect to infinite-dimensional subgroups of the diffeomorphism group, which may act differently at every point in the domain. We consider this a promising direction for future work on LieLAC and have noted this in the revised paper.

• Lines 236, 270, 288, 296:
  • All of these have now been fixed and clarified.
• It would be useful to make clearer for each proposed construction the limitation that exists with the current methodology:
  • In order to address exactly what limitations exist in current work we have added a table to summarise precisely what methods have what limitations. The consideration of the weak topology of canonicalizations acting on continuous functions comes from analogy with other weak topologies present in functional analysis. As far as we can tell no one else has used this sort of weak topology of the action on continuous maps, despite being the most natural choice.
• "Clarifying which non-compact Lie groups acting on which spaces have non-closed orbits" and "cases where the energy induces a 'reasonable' probability measure":
  • We appreciate the comments about the lack of clarity on these group actions. We have added a mathematical example to the theory section that highlights all of the problems that can occur when moving to non-compact Lie groups, both with regards to orbits and with our energy optimization framework.
• “discussion in Section 3 much too general and unstructured”:
  • We completely agree with the reviewers comment and have revised our manuscript accordingly to improve clarity and focus (as explained at the start of the response. We have now also moved one of the algorithms to the main text in order to make the approach explicit.
• “I don't quite understand what is being claimed here”:
  • We clarify that the goal of this paragraph was to emphasise that being able to globally parameterise the group is hopeless in the general non-compact lie group case. The noetherian property does not necessarily hold for these. Now we completely agree with the reviewer that knowing global structure (or even local via algebra solvability) can lead to global parameterisations - but this does not hold in general, and we may not have computational access to such global charts. We agree this was not clear, and have now rewritten this paragraph.
• “I'm left not understanding what the authors mean when they state that their framework requires less knowledge about the symmetry group.“
  • Yes, that is right - the global structure is assumed for [2], while [1] utilises being able to calculate $\exp( v )$ for some lie algebra vector $v$. Particularly note that theory of [1] only considers $v$ infinitesimal and the exact global exponential is not even needed, however it does require the exponential map for any $v$. For lie algebra canonicalization, when utilizing coordinate descent (algorithm 3) one only needs action of some basis $v_i$, which is the only thing normally derived when considering PDE symmetries.

As the discussion period draws to a close, we'd like to check in on whether you've had a chance to review our responses and have any follow up questions?

We hope that our reply clarifies and alleviates the reviewer’s concerns. If this is the case, we kindly ask the reviewer to consider raising their rating, given that they are acknowledging the novelty, the strengths and the contributions of our paper.

Thank you for the clarifications. I would be happy to update my score; I would encourage the authors to upload the revised version of the manuscript so as to be able to see that the requested changes have been addressed.

Now we completely agree with the reviewer that knowing global structure (or even local via algebra solvability) can lead to global parameterisations - but this does not hold in general, and we may not have computational access to such global charts.

I agree. Potentially a more precise minimal category of groups for which this (and your framework) holds are matrix Lie groups which are reductive/semi-simple (in the sense defined by Knapp).

Yes, that is right - the global structure is assumed for [2], while [1] utilises being able to calculate $exp(v)$ for some lie algebra vector $v$. Particularly note that theory of [1] only considers infinitesimal and the exact global exponential is not even needed, however it does require the exponential map for any.

I understand now, thank you for clarifying. My confusion likely lied with the question of whether [1] actually achieves global equivariance given that they work only with the injectivity radius of the exponential map - a consequence of losing global structure.

We agree that utilizing more information of the group can prove beneficial. It is unclear however how exactly the geodesic distance can be helpful in minimizing the energy. In order for the method to be an orbit canonicalization, the problem has to be of the form eqn (2) or similar (see Kaba et al. 2023). Could you clarify how exactly you think this may be used?

To clarify, I agree this is beyond the scope of the paper. IIUC the jumping off point where one heads towards weighed canonicalizations is given by the need to have a minimum on every orbit. The requirement is similar to one encountered under a framework I find similar to canonicalization - the deformable template model/random orbit model, and the connection seems unexplored. If the authors are curious they can review [5] (see e.g. chapters 4,9), in short the group is endowed with a Riemannian metric and the geodesic distance acts as an additional regularization loss, such that one not only finds the group element taking a sample back to the 'canonical sample' (template) but the transformation is one of 'lowest magnitude'. Your concerns relating to the existence of a metric which is only left/right invariant rather than by-invariant also appear in this framework.

[5] - Riemannian geometric statistics in medical image analysis, Pennec et. al 2020.

edit after new revision has been uploaded:

In order to address exactly what limitations exist in current work we have added a table to summarise precisely what methods have what limitations.

Is this referring to Figure 3 of the current revised version?

We thank the reviewer for their thoughtful feedback and appreciate their recognition of our work's strengths, particularly its theoretical grounding and potential for broad applicability. We would also like to thank the reviewer for the references provided, which have now been included in the paper. We address the specific concerns raised below:

• “ The paper states "Prior work on equivariant neural networks focuses on simple groups ... that are not rich enough to encode specific structure found in scientific applications". Does it apply to the recently introduced Clifford algebra and geometric algebra networks which aim to provide a more general and flexible framework for equivariant NNs? “
  • These statements do indeed apply to geometric algebra - Geometrical Algebras implicitly assume that the objects considered transform in a linear representation of isometries of the underlying space. In this case one can decompose all possible fields in an equivariant representation. A similar thing follows when we have an algebraic group acting algebraically on a vector space, because then there is a very rich theory of invariant polynomials - both theoretically and computationally. Unfortunately, the groups necessary for working with PDEs, while they seem to generally be algebraic, do not act algebraically (for instance, the heat equation has a square root and an exponential which are not algebraic). This is precisely why we needed to introduce this more general framework - this has now also been clarified further in the introduction.
• “Paper is uneasy to digest. For example, 4 definitions, 3 theorems, 3 propositions, and one lemma are stated on just two pages. Authors should think about how to make the flow of the paper more intuitive by either providing a more comprehensive background or by presenting a simplified formulation in the main text of the paper while providing the complete theoretical details in the appendix.”
  • We completely agree with the reviewers comment and have revised our manuscript accordingly to improve clarity and focus. We would appreciate any further feedback on this.
• “ The distance-based relaxation in 331-333 is not clear. Can authors elaborate more on what is meant here? “
  • Taking the heat equation as an example, we are interested in transformations that map $x$ into the interval $[0,2\pi]$. Thus, to enforce this, we can simply add an extra term $(\max(x_{min}-0,0))^2 + (\max(x_{max}-2\pi,0))^2$ to the energy.
• “ While the method initially appears generic, finding the minimizer of the energy function is challenging and it is not clear how well it can be done besides the cases presented in the paper. The optimization problem of energy minimization in Eq.2 requires a lot of design choices and heuristics such as adversarial regularization. “
  • We agree with the reviewer - a price has to be paid somewhere, when imposing equivariance. For simple groups many equivariant architectures have been proposed, making the our approach rather clunky. However, there exists a much larger class of problems (of which PDE symmetries is one) where no such architectures exist. This work enables achieving equivariance in such problems. While we do pay a price in the complexity of trying to find minimizers of an energy, this allows us to construct a method that works for many difficult cases of group transformations - something that is not present in the field at all.
  • It is also worth emphasising that the goal of this work was not to propose the best possible energy choice, which may turn out to be much simpler than the proposed constructions, but instead showcase that this is possible.

• “Experimental evaluation is rather limited and is focused on toy or simple tasks. Also, the reported PDE evolution errors are reported only on one instance of initial conditions.”
  • The reported metrics for PDE evolution are taken as an average over a number of different initial conditions. This has now been clarified in the paper. We agree that some of the experiments are toy example (particularly the 2D setting or Heat and Burgers evolution). However, the MNIST experiments consider a relatively complex group as an example, while the ACE Poseidon example considers a large scale network, with large images, illustrating the scaling of the method with respect to the problem size.
  • Could you elaborate on what kinds of tasks you would consider to be non-simple in this context? This would be very helpful in guiding our future research and ensuring the method's applicability to a wider range of problems.
• “Table 2 suggests the misalignment of canonicalized samples with the original data indeed takes place (LieLAC+Poseidon vs Poseidon+ft) which again outlines the challenge of minimizing the energy function. This requires finetuning the baseline neural solver to adapt to new OOD samples which can be a serious practical bottleneck.”
  • We would like to emphasise that kNN classification, MNIST, Heat and Burger experiments did not require any finetuning, as the tasks considered are not as complex and do not require the same level of precision. Particularly, even without finetuning in Table 2 we can see significant improvement on out of domain data. The base solver needs to be tuned in this case, as this is a regression task. What is more - the Poseidon model is large and pre-trained on a lot of data - making it already very good to begin with - finetuning was only necessary to show that the model is able to reach the same level of accuracy on both in and out of domain data. Table 2 does indeed suggest a slight misalignment - but it is important to note the emphasis on ‘slight’.

As the discussion period draws to a close, we'd like to check in on whether you've had a chance to review our responses and have any follow up questions?

We hope that our reply clarifies and alleviates the reviewer’s concerns. If this is the case, we kindly ask the reviewer to consider raising their rating, given that they are acknowledging the novelty, the strengths and the contributions of our paper.

Dear Reviewer, with the deadline so close, we'd like to check whether you've had a chance to review our responses and have any follow up questions?

We believe our clarifications address your initial concerns, and if you agree, we would appreciate it if you would consider raising your rating. We're grateful that you recognize the novelty and contributions of our work.

We thank the reviewer for their thoughtful feedback and appreciate their acknowledgment of the novelty and potential of our approach. We have carefully considered the reviewer's comments and have made significant revisions to the manuscript to address the concerns raised.

Regarding specific questions:

• “How the theoretical results and understandings provided in sections 2 and 3 are used in practice to train a canonicalizer network is not well-explained.”
  • We completely agree with the reviewers comment regarding clarity and we have revised the manuscript and more particularly sections 2 and 3 to improve clarity and focus. To help with presentation of section 2, we clarify precisely the contribution of this paper compared to other neural network based approaches. Section 2 separately summarises existing limitations of previous works, illustrated through an explicit example, showing how our framework overcomes them. Section 3 now combines both generic explanations, as well as specific constructions for the examples provided.
  • We would like to clarify that there is no canonicalization network. Constructing a canonicalization network necessarily requires one to use equivariant architectures, which may not exist (or may be too expensive) for the groups considered in this paper. Instead, the energy functional is parameterized as a network. We have added further detailed information about how the energy functionals were chosen for all problems considered.
  • We would also like to clarify that sections 2 and 3 concern themselves with formalizing equivalence of frames and canonicalizations for the case of non-compact lie groups in the most general case. Theory provided in those sections does not provide any constructive information towards how to select the energy. In fact there is no unique good way to achieve this, as emphasised in section 4, and the best anyone may be able to do is provide a way to do this, which works in practice.
  • The approach considered in this paper (for MNIST and ACE) has been to train a VAE on the training set of the operator and a (convex) adversarial regulariser based on infinitesimal generators of the group - we have shown that they are able to achieve equivariance and improved out of domain performance, emphasising that they do work. We add a note on all of the above to the paper.
• “The general experimental setup was not well-explained in the paper ...”
  • Thank you for the feedback - the training algorithm has now been explained in the paper in the section on energy selection and the canonicalization algorithm included in the main text. However, as mentioned above and in section 4, there is no unique algorithm. Different knowledge of the group requires one to utilise different algorithms 1,2,3. We have included information on this in the paper.
• “The reported results do not include standard deviations over seeds (for example, in tables 1 and 2), and the results for Heat and Burgers’ are missing from the main paper (although included in the appendix). The paper also doesn’t include any comparisons with other baselines, such as data or loss augmentation.”
  • Due to space restrictions, experimental results for heat and Burgers equations had to be moved to the appendix. In order to improve the presentation we have now also moved particular information about the heat and Burgers equation to the appendix also.
  • We agree that it would be beneficial to include comparison with data augmentation and we are currently running the relevant experiments to be included in the revised version. However, due to the expensiveness of data augmentation in training - we are likely to be unable to provide these before the end of the rebuttal period. This particularly emphasises the benefit of using canonicalization - as no extra training (at most - minimal finetuning) of the operators is needed.
  • In regards to standard deviations, we agree it is a good idea for these to be included. We are now re-running the same experiments with different seeds to be included in the revised version via standard deviations. In the same manner, we are likely to be unable to provide these before the end of the rebuttal period.
  • For loss augmentation - unfortunately main approaches do not have public code available, and while it is possible to reproduce them, they do not result in equivariant models, and no out of distribution generalization occurs. They instead result in more data-efficient models - see eg Akhound-Sadegh et al. 2023.
• “The paper could be strengthened by providing an understanding of how expensive these optimizations are and how generalizable the suggested approach is for more complex settings. “
  • We agree with this, and have now provided an extra explanation of the time requirements of canonicalization. However, it is worth emphasising that most of these have not been optimized for, and can most definitely be sped up further.

As the discussion period draws to a close, we'd like to check in on whether you've had a chance to review our responses and have any follow up questions?

We hope that our reply clarifies and alleviates the reviewer’s concerns. If this is the case, we kindly ask the reviewer to consider raising their rating, given that they are acknowledging the novelty, the strengths and the contributions of our paper.