Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groups
Authors: Zakhar Shumaylov
Peter Zaika
James Rowbottom
Ferdia Sherry
Melanie Weber
Carola-Bibiane Schönlieb
Venue: ICLR 2025 Poster (2025)
Keywords: Canonicalization
Equivariance
Invariance
Lie algebra
Partial Differential Equations
Neural Operator
PINN
Neural PDE solver
Lie point symmetries
Frames
Frame Averaging
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: The quest for robust and generalizable machine learning models has driven recent interest in exploiting symmetries through equivariant neural networks. In the context of PDE solvers, recent works have shown that Lie point symmetries can be a useful inductive bias for Physics-Informed Neural Networks (PINNs) through data and loss augmentation. Despite this, directly enforcing equivariance within the model architecture for these problems remains elusive. This is because many PDEs admit non-compact symmetry groups, oftentimes not studied beyond their infinitesimal generators, making them incompatible with most existing equivariant architectures. In this work, we propose Lie aLgebrA Canonicalization (LieLAC), a novel approach that exploits only the action of infinitesimal generators of the symmetry group, circumventing the need for knowledge of the full group structure. To achieve this, we address existing theoretical issues in the canonicalization literature, establishing connections with frame averaging in the case of continuous non-compact groups. Operating within the framework of canonicalization, LieLAC can easily be integrated with unconstrained pre-trained models, transforming inputs to a canonical form before feeding them into the existing model, effectively aligning the input for model inference according to allowed symmetries. LieLAC utilizes standard Lie group descent schemes, achieving equivariance in pre-trained models. Finally, we showcase LieLAC's efficacy on tasks of invariant image classification and Lie point symmetry equivariant neural PDE solvers using pre-trained models.

Published as a conference paper at ICLR 2025
LIE ALGEBRA CANONICALIZATION: EQUIVARIANT NEURAL OPERATORS UNDER ARBITRARY LIE GROUPS
Zakhar Shumayloy".,, Peter Zaika* 1 James Rowbottom', Ferdia Sherry',
Melanie Weber", Carola-Bibiane Schönlieb'
University of Cambridge, Harvard University, "Equal contributions
ABSTRACT
The quest for robust and generalizable machine learning models has driven recent interest in exploiting symmetries through equivariant neural networks. In the context of PDE solvers, recent works have shown that Lie point symmetries can be a useful inductive bias for Pryysic-Inffomed Neural Networks (PINNs) through data and loss augmentation. Despite this, directly enforcing equivariance within the model architecture for these problems remains elusive. This is because many PDEs admit non-compact symmetry groups, oftentimes not studied beyond their infinitesimal generators, making them incompatible with most existing equivariant architectures. In this work, we propose Lie aLgebr Canonicalization (LieLAC), a novel approach that exploits only the action of infinitesimal generators of the symmetry group, circumventing the need for knowledge of the full group structure. To achieve this, we address existing theoretical issues in the canonicalization literature, establishing connections with frame averaging in the case of continuous non-compact groups. Operating within the framework of canonicalization, LieLAC can easily be integrated with unconstrained pre-traines models, transforming inputs to canonical form before feeding them into the existing model, effectively aligning the input for model inference according to allowed symmetries. LieLAC utilizes standard Lie group descent schemes, achieving equivariance in pre-trained models. Finally, we showcase LieLAC's s efficacy on tasks of invariant image classification and Lie point symmetry equivariant neural PDE solvers using pre-trained models.
INTRODUCTION
The deep learning revolution of the 2010s has demonstrated the remarkable effectiveness of deep learning (DL) models in various domains, from text generation to computer vision. This widespread success has led to the adoption of DL in numerous scientific fields, including biology (Abramson et al.) 2024). physics (Merchant et al. 2023) and engineering (Degrave et al. 2022). The training of such models often requires access to large amounts of high-quality data and computing resources, which can be challenging to obtain in practice. However, many tasks in these domains involve structured data, including data with symmetries induced by fundamental physical principles. Phyxics-informer neural networks (PINNs) seek leverage such structure to mitigate the aforementioned challenges.
DL models have been observed to significantly benefit from the inductive bias of symmetries. Notable examples include translation invariance in imaging data and time-series, translation and rotation symmetries of molecules, resulting in position and momentum conservation. Incorporating these symmetries into the design of neural networks can enhance their performance and generalization capabilities. This has been analyzed empirically (Esteves et al. 2023 Wang et al. (020) and through the lens of sharning-theored trade-offs (Mei et al. 2021 Bietti et al. 2022 Kiani al.. 2024b).
Prior work on equivariant neural networks focuses on "simple" groups, resulting in frameworks that are often not rich enough to encode the complex geometric structure found in scientific applications. For example, the heat equation in Appendix D.3 was recently shown by Koval & Popovych (2023) to admit a group of point symmetries SL(2, IR) X o (I,1)) which is not only non-compact, it also does not act in a representation equivariant architectures unusable. To overcome this, we introduce a novel framework that is able to encode wider range of symmetries, not requiring priori knowledge of the full group structure. This flexibility allows for integration with existing models, notably pre-trained models, with the potential of leveraging the benefits of geometric inductive biases beyond classical equivariant architectures.
the underlying jet space. This renders most of the currently available

---

Published as a conference paper at ICLR 2025
1.1 EQUIVARIANCE IN NEURAL NETWORKS
Incorporating group symmetries inside neural networks, in the form of equivariance, is one of the fundamental tasks geometric deep learning (Bronstein al. (021). The translational equivariance of convolutional neural networks (CNNs) (Fukushima & Miyake 1982 LeCun et al. 1989) has been suggested as an important reason for their successes in tasks such as image classification (Krizhevsky al. 2012 He et al., 2016). Because datasets may have symmetries more intricate than just translations, the past decade has seen a push towards generalizing this approach to allow for other (Lie group) symmetries, starting from works roto-translationa/e
equivariant CNNs (Dieleman et al. 2013 2016; Cohen & Welling 2016). to more recent methods, efficiently utilizing knowledge of the Lie algebra for more general groups (Kumar al.. 2024 MacDonall al.) [2022
Finzi et al. 2020a McNeela 2023) and computationally efficient methods (Mironenco & Forré 2024). Notably, convolution- based approaches equivariance satisfy type local equivariance 1o0 (Weiler& & Cesa 2019): with compactly supported filters, if features in the data are sufficiently separated, these networks are equivariant to transformations localized the features. Despite all
of the work in this direction, construction of efficient and accurate equivariant architectures is not simple, see Appendix [A]
More recently. there has been drive to- wards publicly available foundational mod- els (Bommasani et al.] (22T)) for vari-
ous tasks, pretrained on massive datasets.
These models, however, are often not equiv- ariant. For example, Alphafold (Abram- et al. 2024) longer utilises equivari- ant architectures, unlike the initial version (Jumper al. (02T). For this reason a
direction of research has been investigat- ing the possibility of utilizing the general- ization offered by equivariance, while be- ing able to use pre-trained models. The main two approaches to achieve this are canonicalization/ Kaba al. (2023); y; Mon- dal et al. (2023) and frame averaging Puny itt al. (2022): Basu et al. (2023) (also re- ferred to symmetrization). However, there still remain a number of open the- oretical questions regarding the possibility of continuous canonicalization (Dym et al.
Zar ULST muduutin
2024), connections between frame averag- ing and canonicalizations Ma et al. (2024) and constructibility of frames Dym et
Figure Effect of canonicalization on decision bound- aries in k-NN classification for separating the inner and the outer rings Section 4.1
(2024): Puny et al. (2022).
1.2 THYSICS-INFORMAE MACHINE LEARNING
In separate direction, deep neural networks being increasingly used in scientific computing, particularly in simulating physical systems modeled by partial differential equations (PDEs). This surge in interest is largely driven by limitations of traditional numerical solvers, they often struggle with demanding scenarios, such as those involving non-linear PDEs on complex domains necessitating fine discretization, et al. 2018). This is further exacerbated if there is need to solve PDEs across varying domain geometries, parameters, and initial/boundary conditions, requiring independent simulations for each unique configuration.
high-dimensional problems where the curse of dimensionality kicks in (Han
The main prominent approach that has emerged is based on integrating physical laws and constraints into the learning process to solve PDEs (Raissi et al. 2019), called physics informed -hysicc-informed neural networks (PINNs), bypasses the need for direct discretization. However, classical PINNs are not without
The idea of canonicalization in many ways is new, see Appendix D.1
2

---

Published as a conference paper at ICLR 2025
Tenal Coadica
Tour Bodd
Tre Podylle
The Bodeie
Eoe 244
Eow: 346
Eine 1280c22
(a) DeepONet operator for the Heat equation Appendix[D.. [D.3
DeepONet operator the Burg-
(c) POSEIDON operator for the
ers' equation Appendix [.44
Allen-Cahn equation Section 4.3.1
Figure 2: Canonicalization pipeline for numerical PDE evolution Section 4.3
limitations (Grossmann et al. 2023 Krishnapriyan et al. 2021). See Appendix A for a further discussion. To address the challenge of solving PDEs across diverse scenarios, Neural Operator models have emerged as a promising solution (Li et al. 2021a Raonic et al., 2024). Unlike standard DNNs, which approximate functions over finiteddimensionll vector spaces, neural operators learn mappings between fiiniteddimensiona function spaces. These learned operators map parameters or IBCs to corresponding PDE solutions. Combining such Neural Operator models with physics- informed loss has resulted in physics-informed neural operators (Li el al. 2021b; Wang et al. 20216), oftentimes improving generalization and physical validity, compared their physics un-informed variants. In what follows, we concern ourselves precisely with this class of approaches.
1.3 EQUIVARIANCE IN PINNS
Based on the discussions above, a natural question arises can generalization of physics-informed neural networks be improved by utilizing symmetries arising in the study of PDEs? A few strides have been made towards understanding ways in which to take advantage of the Lie point symmetries (LPS) of differential equations in neural solvers. The simplest ways add symmetry information into neural PDE solvers is through data augmentation (Brandstetter et al. 2022; Li et al. 2022). al.)
resulting in more expensive training. loss augmentation (Akhound-Sadegh et al. 2023
2023; Zhang et al.] 2023), resulting in more difficult training. Both of these, however, oftentimes result in improved generalization. Similar benefits have also been observed in Gaussian process based solvers Dalton et al. (2024). Alternatively, exact equivariance can be seeked. As an example, in Wann et al. (2021a), neural PDE solvers were combined with standard approaches designing equivariant architectures for simple symmetries such as translations, roto-translations and scalings. Still, differential equations may have significantly more complicated Lie point symmetry groups than the groups for which there are off-the-shed equivariant architectures. The concept of moving frames has been applied to the problem of using PINNs to solve ODEs with LPS in Arora et al. (2024). to invariantize the ODE in question, and apply PINN to this equation,
Here, the approach taken before mapping back to an (approximate) solution of the original ODE. Similarly, Lagrave & Tron (2022) uses equivariant neural networks first approximate differential invariants, before passing to a classical numerical scheme for evolving the solution. Overall, the literature so far has illustrated that Lie point symmetries can provide useful inductive biases for Neural PDE solvers, and from practical point of view can even be approximated from observed data (Gabel et al. 2024). However it still remains an open question, whether is possible to make existing physics-informen approaches equivariant. In this paper we directly tackle this question. Preliminaries are provided in Appendix A
1.4 OVERVIEW OF CONTRIBUTIONS
Motivated by limitations of existing equivariant architectures, we turn to energy-based canonical- ization means for inducing equivariance in existing models. This approach first introduced in Kaba et al. (2023), but remains largely unexplored. In this paper we revisit and extend this approach, and explore its utility in scientific machine learning applications. Noting that many existing
3

---

Published as a conference paper at ICLR 2025
definitions and results in the prior literature do not extend readily the general setting, we extend and refine existing definitions to handle non non-discrete and non-compace groups. Our key contribution is an energy-based canonicalization approach that provides unifying theoretical framework for constructing frames and canonicalizations (Section 2.2 To define energy-based canonicalization we introduce generalization of standard canonicalization allowing use with non-compact Lie groups, which we call weighted canonicalization. This new framework encapsulates prior work of Dy et al. (2024) on weighted frames, and our energy-based canonicalizations live naturally in this framework.
We emphasize that our proposed approach is generic, in that it has the potential to apply to wide range of models and learning tasks. The energy camonicalization framework developed here is readily applicable for an arbitrary group with an arbitrary action, including those not acting linearly, therefore being extendable to gauge and local symmetries. For Lie groups with smooth actions, one only requires knowledge of the action of some basis of the Lie algebra, without requiring global parametrization of the group.
This flexibility stems from the choice of energy functional, which directly influences empirical performance and the size of resulting canonicalizations. We offer guidelines for constructing energy functionals and for performing Lie group optimization in Section 3 We showcase the effectiveness of our method on image classification tasks under homography and affine group invariance in Section 4.2 We further discuss turning neural operators Lie point symmetry equivariant three examples, including the pre-trained POSEIDON foundation model (Herde et al. 2024) in Section 4.3
GENERAL FRAMEWORK FOR CANONICALIZATIONS AND FRAMES
For this section let G be Lie group acting smoothly manifold X. Let (x) be the bounded continuous R- valale functions on X. Let PMeas(X X) be the set of Radon probability measures on X, with respect to the Borel o-algebra on X. We will set Φ (X) some function to be acted on.
The two main ways of imposing equivariance on arbitrary models: frame averaging and canonical- ization, both arise from the study of the Reynolds operator. For finite groups G, defined as
Ra(o)(z)
Σ BEG
Eg x)..
G
Evaluating this equivariant function Ro( can become prohibitively expensive, requiring summation over the entire group for every input Therefore one seeks ways to restrict the range of this summand while preserving the equivariance of the resulting function. Frames are defined, for each element E X, as subset of the group G for which to average over, with equivariance condition detailed below to ensure equivariance of the resulting function. Canonicalizationn]]e are the opposite perspective, for each element & E X, directly providing subset of X over which to directly average the function over with an invariance condition to ensure the resulting function is equivariant. The presentation here was reduced to fit the page limit; for full description Appendix H
2.1 FRAMES
Definition 2.1. Frames are functions F - 2G {Ør such that for any x E X and g E G:
F(gx) gx))). We call the space of frames (xag
EBEFWW ogg As F(x) GzF(z) for any frame F, frame averaging is inherently limited to the case that the frames are finite at each point, i.e. that the F(x) is finite Vr E X. For frames this can only be the case when the group acts with finite stabilizers at all points. This leads to the following definition of Dym et al. (0044) A weighted frame H4 acts on 6 similarly by averaging (())(z) Jc Ó(g
frame F acts on I by frame averaging F())(z)= = şीक
-x) (g)
Definition 2.2.
Weighted Frames functions Mm X - PMMMs (G) s.t. Vr
X, G : 1'gz 9. (respectively: (fgz)=Hggr = (Fzzz )-.))>>>> where Tgr G - G/G gzz is the We call the space of weakly- cquivaraani weighted frames (Frag(X).
quotient map. See Lemma F.2
Original proposal of Kaba al. 2223) relied on these being singletons, requiring authors to introduce relaxed equivariance. Treating these as sets turns out to be natural for both energy canonicalizations (Section 2.3), and for providing (equivalence) isomorphism of (weighted) canonicalizations with (weighted) frames (Theorem G.3 Ma et al. (2024)
4

---

Published as a conference paper at ICLR 2025
2.2 CANONICALIZATIONS
We briefly recall the definitions of (finite) canonicalizations from Ma et al. 2024).
Definition 2.3. A canonicalization is function C X 2 (0) s.t. Ve E X,9 E G C(gx) C(τ). We call the space of canonicalizations (X))
2*
Definition 2.4. An orbit canonicalizations is canonicalization s.t. Vo X, C(x) C Gr. We call the space of orbit canonicalizations OCang(X)
- A canonicalization C acts & similarly to frames by (()) [ि(ā] Eyec(s) o(v). By the orbit- stabilizer theorem, orbit canonicalizations are equivalent frames. this isomorphism, frames are always larger by factor equal to the size of the stabilizer group. Thus from a computational perspective orbit canonicalizations are prefered, and we will consider the canonicalization perspective from now on. We shall also see that this is indeed the correct perspective to work with when G is non-compact Lie group.
2.3 ENERGY- BASED CANONICALIZATION
We are now in a position to introduce the definition of our energy-based canonicalization method. We start with an energy function E X + [0, +oo] which we then minimize over the orbits of the group action Gz. First, in the case where G' is finite, every orbit in X necessarily has minimum of E, and we note that energy minimization naturally results in a frame (Proposition G.I.:
FÆ(Z) = arg min E(g 9EG
(1)
Moving onto canonicalizations we want to note that intuitively the best energy would be some approximation of the negative log likelihood for the sampled distribution. Therefore finding the minimizers of the energy within an orbit has the effect of finding those examples in the orbit which are most likely to have been trained on. keeping with this intuition, we would like to define the energy minimizing canonicalization to be:
Cr((x) arg min E(y) yEGr
In the non-compact cases we wish to consider, this definition isn't suitable for two main reasons: the set of minima may not be finite or there may be no minima. Dealing with the first naturally leads us to defining weighted canonicalizations, in analogy with weighted frames:
Definition 2.5. A weighted canonicalization is a G-invariant function [il X - PMeas(X) We
call the space of weighted canonicalizations WCang(X).
Definition 2.6. A weighted orbit canonicalization is a weighted canonicalization x such that for every I E X, supp C Ga. We denote the set of weighted orbit canonicalizations as WXCang(X).
A weighted canonicalization K₂ acts é similarly as (x)) = [[x driz- As in the non-weighted case, we have a similar equivalence between weighted orbit canonicalizations and weakly-equivand weighted frames given by Theorem G.3 Weighted canonicalizations do not address the second issue however, as orbits of non-compact Lie group are not necessarily compact or even closed. The conpacepassssss of these orbits is easy to solve by making the standard assumption that E is topologically coercive: for any N E R* there is a compact K C X such that E > N. The non-closedness of orbits more subtle, as shown by the following example.
Example. Let X R and = Ro with its standard multiplicative Lie group structure. Note that G' is non-compact Lie group. Let E G act by t. tz. Then one sees that the orbits of this action are exactly the non-closed rays R5 emanating from origin as well as simply the origin. In practice, this might represent scale-invariant learning problem.
Now consider the energy minimization problem defined by taking the energy to be E(x) = | The global minima of this function is clearly the origin, which isn't 't in the domain. On the other hand, any of the non-closed rays have the origin as their limit points, therefore E does not have a minima on any of these orbits.
To get around the non-closedness issue introduce the following notion. This may seem like an artificial solution to the non-closedness problem, but in fact weighted closed canonicalizations is the most natural extension, as they are simply limits of weighted orbit canonicalizations by Theorem G.5
5

---

Published as a conference paper at ICLR 2025
Definition 2.7. A weighted closed canonicalization is weighted canonicalization X - PMeas(X) such that for every x E X, suupp Ke S Gz. Call the set of weighted closed canonicaliza- tions WXC.
By taking the closures of orbits, we solve the problem of E. not having minima, so are naturally led constructing weighted closed canonicalivations. First, we define the energy-minimizing set, M MΦ(r) = arg min ye Ely), set of minima of E on the closure of the orbit of x.
Proposition 2.1. (Proposition G.2 (a) G-invariant and non-empry.
If MM(x)) is finite for each X, then define the (non-weighted) closed canonicalization KE E WCCanG by taking (x( to be the normalized counting measure on Mı( as usual.
For general E, may be able to put any reasonable probability measure the Mn for example if Mx(r) has highly fractal geometry. We therefore must assume that E is such that for all x, the Hausdorff measure of MM (z) ( non-zero. As each My (z) is compact, the Hausdorff measure is finite, hence combined with this assumption we may normalize it to obtain a probability measure on MM(x)) which a weighted canonicalization (g(x)) E WCCang (X). Refer to Appendix for more discussion on the Hausdorff measure and on why this assumption is reasonable. Note that if (æ(r) is finite, then the normalized Hausdorff measure constructed agrees with the normalized counting measure. Hence this construction is generalization of the purely finite case.
2.4 CONTINUITY PRESERVING CANONICALIZATIONS
Using the setup above we may now write down our main diagram connecting all the different notions of canonicalization and framing together, shown in Figure 3
Faite comniiiiiiiee
Oliak
(
Cang
In practice, it is often desired for conti- nuity be preserved under canonicaliza- tion/frame averaging. In Dym et al. (2024) shown that there are many topolog- ical obstructions to constructing continu- ous canonicalization functions. In anal- ogy with Dym et al. (2024), we may now ask whether energy canonicalizations pre- sented here take continuous functions continuous functions. our framework this is simple condition to add.
Wonsoe
WOOmeeRf
WCneu
Gconged
Figure Connections between various notions introduced previously and in this work. Top row: finite frames and canon- icalizations, mapping vertically into weighted versions via normalized counting measures. Inside each, contained a se- quentially closed subspace of those that preserve continuity.
Definition 2.8. A weighted canonicalization is continuous if for every & CX( k(O) is continuous. Call the space of continuous weighted canonicalizations WCan (X). We similarly define continuous weighted (orbit) canonicalizations and weighted frames.
Proposition 2.2. A weighted canonicalization is continuous if and only if for every convergent sequence Fn + in X, the sequence of measures Kz,, converges weakly to Rg-
turns out, that for energy canonicalizations this completely provides a characterization of which energies give rise to continuous weighted canonicalizations
Lemma 2.3. (Lemma G.7 Let E be lower semi-continuoss energy function. Then Kr is continuous if und only for every convergent sequence Fn - x we have Li MF(x,) =M))..

---

Published as a conference paper at ICLR 2025
3 PRACTICAL CONSIDERATIONS
In Section 2.3 we proposed energy canonicalization constructive way of achieving equivariance, but a question remains how to choose the energy?? Unfortunately there is no unique way to choose the energy or the optimization routine to find the minima. In this section we instead outline key considerations to guide these choices, with Section 4.1 serving as a guiding example. In complex settings of Sections 4.2 and 4.3.1 we choose be a neural network, as explained in Appendix [ while for Section 4.1 and appendices D.3 and D.4 the energy is hand-crafted.
As we may not directly have access to parameterizations of the orbits for Equation 22. we instead implement group descent algorithms for Equation ( e.g. illustrated in Algorithm and then consider the resulting canonicalizations, as this are much smaller than frames and often.
In what follows, we clarify the interplay between the four primary problem aspects. We assume to be given a specific choice of group and action a domain, as well as infor- mation about the domain itself (either through samples or in the case of synthetic data, full knowledge is available). Based on these, we need to construct two things: (1) the en- ergy function, based on the domain information; and (2) the group-based optimization routine for the problem in Equa- tion P based on the domain information and the energy choice (potentially leading to problem relaxation as e.g. in Section 4.3.1 or improved initialization for iterative optimization as e.g. in Appendix D.3).
Algorithm Canonicalization
global retraction
Datu: fon-canonical input
parameters .N, steps y. & E g. Energy : X + IR,
Retraction r Result: Canonicalized inpur & - 9-1 La; inverse g" canonicalizing group element g. gradient descent on G.G.D
G
E Sai for - O....N I - ((T())))))
return n(g) (:) [+(]]-
Domain information Models trained on a specific domain oftentimes do not generalize beyond their training data. To ensure the canonicalized input remains within the domain of the trained model, Nomain--peciiee information should be directly incorporated into the energy function. To give concrete example, consider setting in which some neural network is trained on samples living on compact set M. In such setting we would only be interested in canonicalizations which are either close to the set, or live on the set itself. Thus, choosing the energy function to include the characteristic function of the set X.M is particularly suitable when the action is regular (or transitive). Otherwise, there may not be any elements in the set M that can be mapped to, in which case a relaxation using the distance function (z) minge M | ell can be used. Such choices turn out to be beneficial, especially in the setting of Section
In the case of data sampled from underlying distribution ~ Px, one natural choice would be to consider as energy function the negative log likelihood E(x) = loggox (2),, which may be available in closed form if trained on synthetic data. Otherwise, an approximation can be used. For this, there exist quite a few choices, including normalizing flows (Dinh et al., 2016) or score- based
diffusion models (Song et al. 2020) to approximate the true likelihood or its gradient. In this paper. for examples in Sections 4.2 and 4.3. we opt for a simpler variational autoencoder (Kingma 2013). using the ELBO, being a lower bound, instead of the likelihood for the energy.
Non-convexity and Flatness Even with a convex energy, minimization with respect to E G introduces non-convexity. The orbit Gz is often non-convex, e.g. see Figure 5 where this is circle. Invoking a parameterization of group element through its Lie algebra and considering its action on the space leads to further non-convexity. In the worst cases, such as those considered in PDE examples in Section 4.3 we may not have access to global group parameterization, significantly complicating the process. What is worse, the energy may exhibit many flat regions, with the gradients being noisy, as in Figure
Thus the problem requires the use of non-convex optimization methods. We find, that in all examples in Section 3] utilizing multi-initialization strategy with a fixed number of gradient steps is sufficient for convergence. Altenatively, relaxation to the problem in Equation (2 to be performed over x, subject to an orbit distance constraint can be used.
'we note that unlike Kaba et al. 2003)) there is no canonicalization network. Constructing canonicalization network requires one to use equivariant architectures, which may not exist (or may be too expensive) for the groups considered in this paper. Instead, when needed, the energy functional is parameterized as network.
7

---

Published as a conference paper at ICLR 2025
TSNE MNIST
TSAE homMMIST
TSNF Canonicalzee hoMMNIST
Figure 4: I-SNE plots: For MNIST, homNIST and canonicalized homNIST (Section 4.2
To address the problem of flat zero-likelihood regions, alternative energy functions can be considered, e.g. empirical distance E(x)= = Ex-px x)) However, in Sections 4.2 and 4.3. we make a simpler choice to utilize extra regularization. Due to the difficulty in handcrafting regularization for various specific cases, we utilize learned convex regularization, trained ad adversarially on real and partially transformed data, based on the recent proposals in adversarial regularization Mukherjee et al.. (2024); Shumaylov et al. (2024) 2023). Ideally, the regularization would also ensure uniqueness of minimizers to the problem in Equation 2 however is not clear how to achieve this. Generically, would want to construct an energy such that level sets are transverse to group orbits, similar to Panigrahi & Mondal (2024), which is indeed achieved by utilizing adversarial regularization.
Non- tou-parrameterizabiiity Moving from compact groups to non-compact presents further challenges. Even with access to the exponential mapping. it may not be surjective, making global parameterization of the group using the Lie algebra potentially impossible, if the group is not Noetherian. However, Lie algebras of many PDE symmetry groups solvable, which can be used for group parameterizations that simplify the optimization problem: if the Lie algebra is solvable, any group element can be decomposed using the derived subalgebras (Varadarajan 2013 Theorem 3.18.11) as train of exponentials. For example, for the ID heat equation (Appendix D.3) we can represent any element of the group as exp(a6%6) ... exp(a1 v1), reducing optimization over the group to optimization over (a1..... ac) E (ddim(s), In this case one can utilize Algorithm |
When the algebra is not solvable, one has to revert to Lie algebra based descent schemes. However, for Lie point symmetry groups, oftentimes only actions of specific basis vectors of the Lie algebra are available. making it necessary to use coordinate descent (CD) approach. CD has received considerable attention in classical optimization (see Wright (2015)) and has recently been studied in geometric settings such as Riemannian manifolds and Lie groups (Han et al., 2024). The algorithms for both of these cases can be found in Soarithms[aandl
a global parameterization is available, then classical schemes may be utilized in the same way (Lezcano Casado 2019 Absil et al. 2008).
EXPERIMENTS
4.1 2D TOY EXAMPLE
To illustrate the universality of the approach, we consider one simple two dimensional example without any machine-learned components. We consider the problem of classification under inherent rotation invariance illustrated on Figure 5 To mimic a realistic setting, we assume to have samples only. In this case, we can approximate the density using kernel density estimation with a Gaussian kernel, illustrated in the second column on Figure 5 As one would expect, without knowledge of the symmetries of the distribution, the approximation far from the ground truth. In what follows, the energy is chosen to be the negative log- likeliikdoood of the approximated distribution.
This example illustrates the main issue with energy based canonicalization, which also extensively arises in more complex examples: even when the energy is convex, it is not convex when viewed as function of g. This results in non-convex optimization problem with (potentially many) local
8

---

Published as a conference paper at ICLR 2025
minima, which can get stuck in when running standard descent schemes, and as discussed in Section one has to revert to more complex optimization routines. In this case, as shown in Figure [ it is possible to find global minimizers, despite the non-convexity, by rumning gradient descent for fixed number of iterations, but for number of different initializations. The power of canonicalization can then be seen in Figure T turning a non-equivariant kNN classifier rotation equivariant.
4.2 INVARIANT CLASSIFICATION
To illustrate the effectiveness of LieLAC in achieving invariance of pre-trained models, we first test it on in- variant image classification with respect to the groups of affine and homography transformations. We take MacDonale et al. (2022) as the baseline for compar- ison due the similarity in the required knowledge of the group required for implementation.
Table I: MNIST test accuracy for Affine and Homography groups, compared with affConv and homConv of Mac et al. (0022)
Name
MNIST
aFNIST
CNN
0.983
0.629
LicLAC [CNN]
0.979
0.972
afCCony
0.982
0.943
Here, take pre-trained convolutional classifer, trained on the original MNIST dataset. The combi- nation with the canonicalizer,
Name
MNIST
bomNIST
CNN
0.985
0.644
LicLAC [CNN]
0.982
0.960
described in Sec- tion 3 is then tested on the affine-perturbed MNIST (aftNIST) (Gu & Tresp 2020) and the werogrrphy-peettrrre MNIST (homNIST) (MacDonald
homCony
0,980
0.927
et al. 2022). The overall performance results for this example are reported in Table ! We can see that the canonicalization based approach achieves improved equivariance compared to the baseline, as illustrated by improved test accuracy on the datasets. We can also directly analyse achieved cquivariance, as it appears only through canonicalization. We provide examples, visualising achieved equivariance in Figure 6 We can also analyze the effect of canonicalization globally by considering t-SNE plot (Van der Maaten & Hinton 2008) in Figure [ showing that energy canonicalization turns the dataset into the usual MNIST form, clearly separating the classes, unlike homNIST.
4.3 EQUIVARIANT PDE EVOLUTION
In analogy with Akhound -Sanngh--aeghhh et al. (2023). ) we first consider two simple examples with non- trivial symmetry groups: the ID Heat and Burgers' equations, which also admit global parameterisations of the group elements (Koval & Popovych| 2023). Due to space restrictions, evaluation can be found in Appendix [D.3] for the Heat and Appendix D.4 for the Burgers' equation. Full description of
the experimental setting, further visualisations and theoretical background can also be found in the appendix. Examples of canonicalized solutions are shown in Figure 2
We explicitly emphasize that in this paper we only consider symmetries for initial value problems, where there is dependence of g x and g x' (I, t. u) = x' (z, t). In this case the action on I, t can be separated from the action on u, making the desired equivariance for the mhysics-informer neural operator to be:
Le. for (a', f', u') 9 (2, u), have
Og (uo To, (n) - g- Oglg uo. 9
-o,g -1 tol(g= -gs 1.) .t)
4.3.1 ALLEN-CAHN EQUATION
We test the POSEIDON (Herde et al. 2024) foundation model, fine-tuned on the Allen-Cahn equation (ACE) on data it was trained on, and transformed data (using Lie-point symmetry group of ACE).
Us Usz + Ugy Z2u (u² -
(4)
We consider solving Equation 4 on 2D domain with periodic boundary conditions for / E [0, 1].2 x E Ω [0, 1 This equation admits 4-dimensional Lie algebra, see Appendix D.5 with SE(() as the point symmetry group.
Choice of Energy In order to justify the choice of energy, we need to understand the POSEIDON (Herde et al. 2024) model itself. Firstly, 2o and to in Equation (3) are implicitly defined in the model to be fixed to the standard domain Zp = Z and to = O. As such, the only input is the initial condition wols- Therefore, are going to have to only consider transformations, that preserve xo and l. This

---

Published as a conference paper at ICLR 2025
can be imposed directly within the energy canonicalization framework described in Section by considering the characteristic functions as follows:
(x to, Mo, ,f, = X(aa (zo) + (40)((0) + X(o, (tr) + x2( + (uo)
Note, that this is able to significantly reduce the computational cost of optimization, as the group generated by in 28, resulting in rotations of the underlying domain, can be reduced to a discrete subgroup C2 I S2(2), as the only transformations such that g Q = Q. Further information on the
training and testing data is available in Appendix D.5
The overall performance is reported in Table 2
Finetuning As seen in prior works (Mon- dal et al. 2023). the canonicalized data
Tuble 2: ACE Tex: enor evaluated over 90 urjecturiiss for in dattribution (ID) and out of distrituiton (OOD), All erors are (x10 1).
may end up misaligned with the original
Name (x110*)
ID Emorrey
29
ODD Error sq. 50
Avg
Possmox
73.76
41.35
data. For this reason finetuning of the
LicLAC (PoSEIDON)
16.39 1. 5.42
29.15 5.64
22.91
model/canoniicalzation may be desired.
POSSIDOX + (t
5.45 1 2.87
20.05 1 1
14.27
LisLAC (PosuDDNN+ [t]
10.23 + 3.36
11.34 = 2.92
10.79
PossIDox + D. (data ang)
8:41 + 3.48
8.50 3:11
8:46
It is important to emphasize that finetuning
not necessary, and even without finetuning the out of distribution errors go down significantly, as seen in Table All of the examples above, including kNN (Section 4.1.. MNIST (Section (42), Heat and Burgers (Appendices D.3 and D.4p did not require finetuning. That is due to tasks not being as complex, and not requiring the same level of precision. POSEIDON, being pre-trained in a regressive manner, turns out be very good to begin with, and finetuning is performed to illustrate that baseline accuracy of data-augmentation (Brandstetter et al. 2022) can be matched.
While may seem that finetuning the resulting operator on canonicalized data is as hard as augmenting the original data, it turns out to not be, as fine-tuning can be done on very small data-pool (in this case only 100 trajectories). This stems from the energy minima in Equation (2) residing near, but not necessarily on, training data. If the orbits are well-sampled in training. finetuning is not necessary, e.g. seen in Section 4.2 However, based on the training data in eq. (29). the orbits under C4 are well-sampled, while the orbits under SE(2) are not, illustrated in Table wherein the error falls 2/3-fold after finetuning. An example is shown on Figure 2c with further illustrations in Figure 13
CONCLUSION
In this paper we provided unified theoretical framework for defining frames and canonicalizations over non-compact and infinite groups We showed that energy-based canonicalization can be used constructively to obtain both, and provide a Lie algebra based approach for turning any pre-trained network equivariant. We showcased that the proposed method is effective in turning image classifiers invariant, well as neural operators Lie point symmetry equivariant.
Future work fine-tuning foundation models. Due to the typically high cost of (re-) )training foundation models, further exploration of such applications is an important direction for future work.
Our example in Section 4.3.] illustrates the promise of the proposed framework for
Of particular interest are "realistic" scientific machine learning applications: In such applications collecting sufficiently large training data sets may be prohibitively expensive or technically impossible, i.e., available data not sufficient for retraining existing models/ using standard fine-tuning approaches. Encoding symmetries explicitly allows for more efficient fine-tuning, which may be more accessible in settings with little data.
Limitations Energy-basse canonicalization suffers from slow inference (resulting in 5 30 slowdown). Improvements to the optimization are possible via non non-convex optimization techniques (e.g., simulated annealing (Baudoin et al. 2008 Rutenbar 1989), particle swarm (Huang et al., 2023). consensus-bon (Torzeck] [02 Fornasier et al.] particularly applicable as Lie groups are oftentimes now-dinensional, while the ambient space not. The majority of these, however, not casily adaptable to the setting of Lie groups.
Implementing Lie point symmetries in existing neural operators is difficult due to hard-coded boundary conditions. Future work should focus on foundation models handling initial/boundary conditions and parameters, explore alternative PDE symmetries (Brandstetter al. 2022). This could include studying studying equivariance to local transformations.
symmetry groups, which would also naturally arise when
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors would like to thank Olga Obolenets, Lucy Richman, Hong Ye Tan and Yury Korolev for useful discussions and feedback. ZS acknowledges support from the Cantab Capital Institute for the Mathematics of Information, Christs College and the Trinity Henry Barlow Scholarship scheme. FS acknowledges support from the EPSRC advanced career fellowship EP/VO29428/L MW was supported by NSF Awards DMS-2406905 and CBET-21 SEEE-2112085 12085, and a Sloan Research Fellowship. CBS acknowledges support from the Philip Leverhulme Prize, the Royal Society Wolfson Fellowship. the EPSRC advanced career fellowship EP/VO29428/1, EPSRC Grants EP/S026045/1 and EP/T003553/1, EP/N014588/1 EP/T017961/L. the Wellcome Innovator Awards 215733/2/19/Z and 22 1333//2/22// the European Union Horizon 2020 research and innovation programme under the Marie Sktodowska- Cururre Grant agreement No. 777826 NoMADS, the Cantab Capital Institute for the Mathematics of Information and the Alan Turing Institute.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, pp. 1-3, 2024.
P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2008
Tara Akhound-Sadegh, Laurence
Johannes Brandstetter, Max Welling, and Siamak Ravanbakhsh. Lic point symmetry and physics-informed networks. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neu- ral Information Processing Systems, volume 36, pp. 42468 42481. Curran Associates, Inc., neurips..oee.ureeeee...mmm.mme. http://ppppppeeeeeepppp/////// 2023. URL http::/:/pppeeeee.......o...o/ea Eile
Shivam Arora, Alex Bihlo, and Francis Valiquette. Invariant physics-informed neural networks for ordinary differential equations, 2024. arXiv:2310.17053.
Sourya Basu, Prasanna Sattigeri, Karthikeyan Natesan Ramamurthy, Vijil Chenthamarakshan, Kush R Varshney, Lav R Varshney, and Payel Das. Equi-tuning: Group equivariant fine-tuning of pretrained models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 6788- 6796, 2023.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order cquivariant message passing neural networks for fast and accurate force fields. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave. K. Cho, and A. Oh (eds.), Ad- vances in Neural Information Processing Systems, volume 35, pp. 11423-1144333 1436. Curran Asso- ciates, Inc., 2022. URL http:: http::///oooeeeedeeoe neur htt.:/sss.rrrr..s....eeeeeeeeeee filess/aperr http:///c.......aaaaa.o......aa.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Korbbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1):2453, May 2022.
ISSN 2041-1723, doi: 10. 0011088/444477-02--2222
Fabrice Baudoin, Martin Hairer, and Josef Teichmann. Omstein-Uhlenbeck processes on Lie groups. Journal of Functional Analysis, 255(4):877-89) 2008
Gerd Baumann. MathLie a program of doing symmetry analysis. Matheraatics and Computers in Simulation, 48:205-223, 1998.
122231021
Erik Bekkers. B-Spline CNNs on Lie groups. In International Conference on Learning Representa- fions, 2020. URL wtt.:httt///peemmt www..eet/forummmiicmmm id= |
gBhk 1i

---

Published as a conference paper at ICLR 2025
Erik J. Bekkers, Maxime W. Lafarge, Mitko Veta, Koen A. J. Eppenhof, Josien P. W. Pluim, and Remco Duits. Roto-translation covariant convolutional networks for medical image analysis. In Alejandro F. Frangi. Julia A. Schnabel, Christos Davatzikos, Carlos Alberola-López, and Gabor Fichtinger (eds.), Medical Image Computing and Computer Assisted Intervention MICCAI 2018, pp. 440- 448, Cham, 2018. Springer International Publishing. ISBN 978-3-030-00928-11
Alberto Bietti, Luca Venturi, and Joan Bruna. On the sample complexity of learning under geometric stability. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances Neural Information Processing Systems, volume 34, pp. 18673-18684. Curran Associates, Inc., 2021.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint arXiv:2108.072255 2021.
Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmentation for neural PDE solvers. In International Conference on Machine Learning, pp. 2241-2256. PMLR, 2022.
Johann Brehmer, Pim De Haan, Sönke Behrends, and Taco S Cohen. Geometric algebra transformer. Advances in Neural Information Processing Systems, 36, 2024.
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. CoRR, abs/2104.13478, 2021. URL https
//arxiv. org/ abaaabaa/a2a/ 13447
Elena Celledoni, Matthias J Ehrhardt, Christian Etmann, Brynjulf Owren, Carola-Bibiane Schönlieb, and Ferdia Sherry. Equivariant neural networks for inverse problems. Inverse Problems, 37(8): 085006, 2021. ISSN 0266-5611, 1361-6420 doi: 10. 1008//136---0 01988//3116622/a11100
Dongdong Chen, Mike Davies, Matthias J. Ehrhardt, Carola- Bibbine Schonlieb, Ferdia Sherry, and Julian Tachella. Imaging With Equivariant Deep Learning: From unrolled network design to fully unsupervised learning. IEEE Signal Processing Magazine, 40(1):134-14) 2023. ISSN 1053-5888, 1558-0792. doi: 10. 10.11009....02.2205552222
Taco S. Cohen and Max Welling. Group Equivariant Convolutional Networks. In Proceedings of The 33rd International Conference on Machine Learning, pp. 2990-29999999 2016.
David Dalton, Dirk Husmeier, and Hao Gao. Physics and Lie symmetry informed Gaussian processes. In Forty-first International Conference Machine Learning. 2024.
Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):414 419, 2022.
Harm Derksen and Gregor Kemper. Computational Invariant Theory. Springer, 2015.
Sander Dieleman, Kyle W. Willett, and Joni Dambre. Rotation-invariant convolutional neural networks for galaxy morphology prediction. Monthly Notices of the Royal Astronomical Society. 450(2):1441-1459, 04 2015. ISSN 0035-8711. doi: 10.1093/mnras//tv632. URL https
//doi.org// 1ttt//0rm/a///////mmmmm.m stv62
Sander Dieleman, Jeffrey De Fauw, and Koray Kavukeuoglu. Exploiting cyclic symmetry in convolu- tional neural networks. In Maria Florina Balcan and Kilian Q. Weinberger (eds.). Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1889-18888 New York, New York, USA, 20-22 Jun 2016. PMLR. URL https: http::///pooeeediggs mlr prwss/eee/seeeeee v48 1/48/dies die wedileeaaan lemanneannnn 16 6. html
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. arXiv preprint a605v 10050000,,. 2016.
Nadav Dym, Hannah Lawrence, and Jonathan W. Siegel. Equivariant frames and the impossibility of continuous canonicalization, 2024, arXiv:2402.16077
12

---

Published as a conference paper at ICLR 2025
Carlos Esteves, Jean-Jacques Slotine, and Ameesh Makadia. Scaling spherical cnns. arXiv preprint arXiv:2306.05420, 2023.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. International Conference Machine Leurning, pp. 3165-3176. PMLR, 2020a.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to Lie groups on arbitrary continuous data. In Proceedings of the 37th International Conference Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 3165-3177 3176. PMLR, 2020b. URL mwr. press/ wttsss/vvviiie wwtps///ll//iinneeeeeee (ell19/innzi0a.... html
Massimo Fornasier, Timo Klock, and Konstantin Riedl. Consensus-basee optimization methods converge globally. SIAM Journal on Optimization, 34(3):2973-304- 2024.
Igor L Freire. Note on Lie point symmetries of burgers equations. Trends in Computational and Applied Mathematics, 112):151-157, 2010.
Kunihiko Fukushima and Sei Miyake. Neocognitron: A new algorithm for pattern recognition tolerant
of deformations and shifts in position. Pattern Recognition, 15(6):455 469, 1982. ISSN
doi: http:///ooooog/1/10000000011-333 com/ http:///noe//rraaaeee/////111333
URL http::ttt///ttttt hctpp://www..e wcw. wwtpw.ccieneeiree..o
Alex Gabel, Rick Quax, and Efstratios Gavves. Data-driven Lie point symmetry detection for
continuous dynamical systems. Machine Learning: Science and Technology, 2024.
John C Gower. Generalized Procrustes analysis. Psychometrika, 40:33-51, 1975.
Tamara G Grossmann, Urszula Julia Komorowska, Jonas Latz, and Carola-Bibiane Schönlicb.
Can physics-informec neural networks beat the finite element method? arXiv preprint 2.0i::2302.041007, 2023.
Nate Gruver, Mare Anton Finzi, Micah Goldblum, and Andrew Gordon Wilson. The lie derivative for measuring learned equivariance. In The Eleventh International Conference Learning Representations.
Jindong Gu and Volker Tresp. Improving the robustness of capsule networks to image affine transformations, 2020. (RRl/tttps::/axxii..rrg
07968
Andi Han, Pratik Jawanpuria, and Bamdev Mishra. Riemannian coordinate descent algorithms on matrix manifolds. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 17395-17415. PMLR, 21-27 Jul 2024. URL htttttss::/poceeeeeiiieeeo .aa.g.smm........... wwwpppesssss... han24c. html
Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34)88505-855) 8510, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, and Siddhartha Mishra. Poscidon: Efficient foundation models for PDEs, 2024. URL htt.ps://arxin o abs///a//// 2405. 19101
Hui Huang, Jinniao Qiu, and Konstantin Riedl. On the global convergence of particle swarm optimization methods. Applied Mathematics & Optimization, 88(2):30, 2023.
Nail H Ibragimov. CRC handbook of Lie group analysis of differential equations, volume CRC
press, 1995.
13

---

Published as a conference paper at ICLR 2025
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn /unyaauvunakkoo,, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873:5833-88)) 2021.
Sékou-Oumar Kaba, Amab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with Learned Canonicalization Functions. In Proceedings of the 40th International Conference on Machine Learning, pp. 15546-15566. PMLR, July 2023.
Frederic Kanter and Jan Lellmann. A flexible meta learning model for image registration. In International Conference on Medical Imaging with Deep Learning, pp. 638-652. PMLR, 2022.
Sharmila Karumuri, Rohit Tripathy, Ilias Bilionis, and Jitesh Panchal. Simulator- frree solution of high-diment nigh-dimensional stochastic elliptic partial differential equations using deep neural networks. Journal of Computational Physics, 404: 014:109120 2020.
Bobak Kiani, Thien Le, Hannah Lawrence, Stefanie Jegelka, and Melanie Weber. On the hardness of learning under symmetries. In The Twelfth International Conference on Learning Representations, 2024a. URL /ttetpps://operreeewww net. http://wwwwe//ootutmmmma rttzAAt
Bobak Kiani, Thien Le, Hannah Lawrence, Stefanie Jegelka, and Melanie Weber. On the hardness of learning under symmetries. In International Conference Learning Representations, 2024b.
Diederik P Kingma. Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114 2013..
Alexander Kirillov, An Introduction to Lie Groups and Lie Algebras. Cambridge Studies in Advanced Mathematies. Cambridge University Press, 2008.
Serhii D Koval and Roman O Popovych. Point and generalized symmetries of the heat equation revisited. Journal of Mathematical Analysis and Applications, 527(3): 2023.
Aditi Krishnapriyan, Amir Gholami, Shandian Zhe, Robert Kirby, and Michael W Mahoney. Char- acterizing possible failure modes physics-informed neural networks. Advances in Neural Information Processing Systems, 34:26548-2656/ 2021.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In F. Pereira, C.J. Burges, L. Bottou, and K.Q Wein- berger (eds.), Advances in Neural Information Processing Systems, volume 25. Curran Asso- ciates, Inc., 2012. wwR.httpp::/pprceeeeiiee 2012 20222221002298225999666666662282
neurips. cc/paper_ ffiles/paper. .dd
Harshat Kumar, Alejandro Parada-Mayorga, and Alejandro Ribeiro. Lie group algebra convolutional filters. IEEE Transactions on Signal Processing, 2024.
Pierre-Yve Lagrave and Eliot Tron. Equivariant neural networks and differential invariants theory for solving partial differential equations. In Physical Sciences Forum, volume 5, pp. 13. MDPI, 2022.
Y. LeCun, B. Boser, S. Denker, D. Henderson, R. E. Howard, W. Hubbard, and L. D. Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1(4):541-55,, 1989. doi: 10. 16.1ne2//cc..1/ 1989. 1.4.54..
Mario Lezcano Casado. Trivializations for gradient-based optimization manifolds. Advances in Neural Information Processing Systems, 32, 2019.
Jie- King Li, Hui Zhang, Ye Liu, Lei-Lei Guo, Li-Sheng Zhang, and Zhi- Zhi-Yong Zhang. Utilizing symmetry-en eehanncd physics-informed neural network to obtain the solution beyond sampling domain for partial differential equations. Nonlinear Dynamics, 111(23)21861-21187 2023.
Ye Li, Yiwen Pang, and Bin Shan. Physics-guided data augmentation for learning the solution operator of linear differential equations. In 2022 IEEE 8th International Conference on Cloud Computing and Intelligent Systems (CCIS), pp. 543-547, 2022. doi: 10.1 N0.1109/CCIS55298.2222. 20t.:/55777222222200000000 10016380.
14

---

Published as a conference paper at ICLR 2025
Zongyi Li. Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference Learning Representations, 2021a. URL
net/ forum? bettps:/eepenreeee hww.//rrttff i
d=c8P9NQVtmmD
Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Avizzadeneshell, and Anima Anandkumar. Physics-informee neural operator for learning partial differential equations. ACM/IMS Journal of Data Science, 2021b.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. DecpONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem operators. CoRR, abs/1910.03193, 2019.
03193
Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis. DeepXDE; A deep learning library for solving differential equations. SIAM Review, 63(1):208 228, 2021. doi: 10.1137/1MM127006.
Sebastian Lunz, Ozan Öktem, and Carola-Bibiane Schönlieb. Adversarial regularizers in inverse problems. Advances in neural information processing systems, 31, 2018.
George Ma Yifei Wang, Derek Lim, Stefanie Jegelka, and Yisen Wang. A canonization perspective on invariant and equivariant learning, 2024. URL
18378
Lachlan E. MacDonald, Sameera Ramasinghe, and Simon Lucey. Enabling equivariance for arbitrary Lie groups. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 8173-8182, IEEE, 2022. doi: 10. 1990CC 2022.00801. URL http:: org 10 wtt.:///....1111.222100500211010
2022. 2022.00801
Pertti Mattila. Geometry of Sets and Measures in Euclidean Spaces: Fractals and Rectifiability Cambridge Studies in Advanced Mathematics. Cambridge University Press, 1995.
Daniel McNeela. Almost equivariance via lie algebra convolutions. arXiv preprint arXiv:2310. 13.11434 2023.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. Mikhail Belkin and Samory Kpotufe (eds.), Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pp. 3351-3418. PMLR, 15-19 Aug 2021.
Amil Merchant, Simon Batzner, Samuel Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk. Scaling deep learning for materials discovery. Nature, 624(7990):80-85, 2023.
Mircea Mironenco and Patrick Forré. Lie group decompositions for equivariant neural networks. In The Twelfth International Conference on Learning Representations, 2024. URL https
7 /ooenrevvi net./ forum?idd
Arna Kumar Mondal, Siba Smarak Panigrahi, Oumar Kaba, Sai Mudumba, and Siamak Ravanbakhsh. Equivariant adaptation of large pretrained models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Artem Moskalev, Anna Sepliarskaia, Ivan Sosnovik, and Arnold Smeulders. Liegg: Studying learned lie group generators. Advances in Neural Information Processing Systems, 35:25212-25223, 2022.
S Mukherjee, S Dittmer, Z Shumaylov, S Lunz, O Öktem, and C-B Schönlieb. Data-driven convex regularizers for inverse problems. ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 13386-13390. IEEE, 2024.
Johannes Müller and Marius Zeinhofer. Position: Optimization in SciML should employ the function space geometry. In Forty-first International Conference Machine Learning, 2024.
15

---

Published as a conference paper at ICLR 2025
Dat Tien Nguyen, Jinwoo Kim, Hongseok Yang, and Seunghoon Hong. Learning symmetrization for equivariance with orbit distance minimization. In NeurlPS 2023 Workshop on Symmetry and Geometry in Neural Representations.
Peter Olver. Applications of Lie groups to differential equations, volume 107. Springer Science & Business Media, 1993.
Siba Smarak Panigrahi and Arnb Kumar Mondal. Improved canonicalization for model agnostic equivariance, 2024. URL https: /tttpp://arxi..... .rg. abs /abs/24455 14089
Omri Puny, Matan Atzmon, Edward Smith, Ishan Misra, Aditya Grover, Heli Ben-Hamu, and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022. URL
wwt.v/iwwww.weeeeeeeeeeeeeeeeeee net/forum?
id fd=z1em 105
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Bogdan Raonic, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucei, Rima Alaifari, Siddhartha Mishra, and Emmanuel de Bézenac. Convolutional neural operators for robust and accurate learning of pdes. Advances in Neural Information Processing Systems, 36, 2024.
David Ruhe, Johannes Brandstetter, and Patrick Forré. Clifford group equivariant neural networks. Advances in Neural Information Processing Systems, 36, 2024.
Rob A Rutenbar, Simulated anncaling algorithms: An overview. IEEE Circuits and Devices magazine, 5(1):19-26, 1989.
Gerald W. Schwarz. Smooth functions invariant under the action of a compact lie group. Topology, 14(1):63-68. 1975. ISSN 0040-9383. doi: https://oii.rg///
1016/0040-933775500006-- URL hntps:: http://wwwwcciccddeeeeeeeeeee com/ wwt.:/cccccceeceeeeeo..cce.eeeee att.ccle/p11/0049987775000
Zakhar Shumaylov, Jeremy Budd, Subhadip Mukherjee, and Carola-Bibiane Schönlicb. Provably convergent data-driven nonvex-nonconvex regularization. In NeurIPS 2023 Workshop on Deep Learning and Inverse Problems, 2023.
Zakhar Shumaylov, Jeremy Budd, Subhadip Mukherjee, and Carola-Bibiane Schönlieb. Weakly con- vex regularisers for inverse problems: Convergence of critical points and primal-dual optimisation. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scar- lett, and Felix Berkenkamp (eds.), Proceedings of the 4/st International Conference on Machine Learning, volume 235 of Procceddings of Machine Learning Research, pp. 45286-45314. PMLLR 21- 27 Jul 2024. URL httl:ittt://pooceeeeeee
ngs .mm.a....... mmlr http:.pprres//////////aaaala///
html
Justin Sirignano and Konstantinos Spiliopoulos. DGM: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1364, 2018.
Yang Song, Jascha Sohl-Dickstein Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.. 13466, 2020.
Jonas Spinner, Victor Bresó, Pim de Haan, Tilman Plehn, Jesse Thaler, and Johann Brehmer. Lorentz-equivariant geometric algebra transformers for high-energy physics. arXiv preprint arXiv:2405. 14806, 2024.
Luning Sun, Han Gao, Shaowu Pan, and TianXXu Wang. Surrogate modeling for fluid flows based on physicc-constraineed deep learning without simulation data. Computer Methods in Applied Mechanics and Engineering, 361:1 12732, 2020.
Claudia Totzeck. Trends in consensus-basee optimization. In Active Particles, Volume 3: Advances in Theory, Models, and Applications, pp. 201-226. Springer, 2021.
16

---

Published as a conference paper at ICLR 2025
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of machine learning research, 9(11), 2008.
Veeravalli S Varadarajan. Lie groups, Lie algebras, and their representations, volume 102. Springer Science & Business Media, 2013.
Soledad Villar, Weichi Yao, David Hogg, Ben Blum-Smith,, and Bianca Dumitraseu. Dimensionless machine learning: Imposing exact units equivariance. Journal of Machine Learning Research, 24 1109):1-32, 2023.
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arXiv preprint arXiv:2002.0306. 2020.
Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. In 9th International Conference on Learning Representatiinss, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a. URL /ttp:sss///peeeeeieeww
net beetfoorue Eorum? 1d wt
8H22KD
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of parametric partial differential equations with physics-informed DeepONets. Science advances, 7(40):eabi8605, 2021b.
Maurice Weiler and Gabriele Cesa. General equivvaanr steerable CNNs. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc; E. Fox, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httpss:/pprooeedinggs http://ss/pppppeeei..........oo. http:://ww....u..i....c.spaaaaaa File/ FLL2//66667771102222237777522220
files//paerr2211/
Marysia Winkels and Taco S. Cohen. Pulmonary nodule detection CT scans with equivariant CNNs. Medical Image Analysis, 55:15-26, 2019. ISSN 1361-8415. doi: http:://boi.org//0.000//
1.media.201.003.000. URL https: hwtp://wwww....oo www.::/ieiineeeeeeeemmce com www.o//cccenne/aatil..e...
pii/ 021/51351841118306033
Stephen J Wright. Coordinate descent algorithms. Mathematical programming, 151(1)3-334, 2015.
Lipei Zhang, Yanqi Cheng, Lihao Liu, Carola-Bibiane Schönlieb, and Angelica Aviles-Rivero. Biophysics informed pathological regularisation for brain tumour segmentation. arXiv preprint arXiv:2403.09136. 2024.
Zhi- Yong Zhang, Hui Zhang, Li-Sheng Zhang, and Lei-Lei Guo. Enforeing continuous symmetries in physics-informed neural network for solving forward and inverse problems of partial differential equations. Journal of Computational Physics, 492:1 12415, 2023.
Maksim Zhdanov, David Ruhe, Maurice Weiler, Ana Lucic, Johannes Brandstetter, and Patrick Forré. Clifford-steerable convolutional neural networks. arXiv preprint arXiv:2402.14730, 2024.
Yinhao Zhu, Nicholas Zabaras, Phaedon-Stelios Koutsourelakis, and Paris Perdikaris. Physics- constrained deep learning for high-dimensional surrogate modeling and uncertainty quantification without labeled data. Journal of Computational Physics, 394:56-85 2019.
17

---

Published as a conference paper at ICLR 2025
Appendices
Extended Background
19
A.1
Equivariant convolutions
19
A.2
Physics-Informa Machine Learning
19
A.3
Why Lie point symmetries?
20
A.4
Symmetries of Partial Differential Equations
21
A.5
Deep Learning For PDES
22
A.6
Lie Groups
22
Choosing the energy
23
Further MNIST Visualisation
24
Further PDE discussion
26
D.I
Canonicalization unifying view
26
D.2 On Deep Learning for PDES
26
D.3 Heat equation
27
D.4 Burgers' equation
30
D.5
Allen-Cahn equation
32
Measure Theory Background
35
More on Frames and Canonicalization
36
Weakiy-Equivvriant Frames
36
projections
38
Energies and Closed Canonicalizations
38
Proofs
38
H Overview of previous concepts
44
Weighted Canonicalization
45
H.2
Energy Minimizing Canonicalization
46
H.3 Almost everywhere continuity is not as impossible
46
Aggoritmms for solvable and non-solvable Lie algebras
48
18

---

Published as a conference paper at ICLR 2025
Energ IOROrOUTTION TENOOO
Autadin
Figure 5: LieLAC optimised sample from mix of Gaussian distributions.
A EXTENDED BACKGROUND
A.1 EQUIVARIANT CONVOLUTIONS
An important fact that comes up repeatedly in the design equivariant neural networks is that ordinary convolutions on Euclidean spaces be generalized to any group which an invariant measure (as is the case for any Lie group, for example), giving rise to equivariant linear maps in the form of group convolutions. Furthermore, converse results be established, showing that a linear equivariant map (under general conditions) must take the form of a group convolution. This is the approach taken in Cohen Welling (2016) for discrete groups of soo-translations acting on images. Since group convolutions are applied to signals defined on the group, it is necessary to apply an initial equivariant linear lifting operation convert signal on the original domain signal on the group. Incompatibility of continuous groups and grid-based discretizations of signals is a major issue methods that take this approach, and general groups, it is not necessarily clear how one should discretize the integral defining group convolution. For Lie groups, there have been various proposals for approximation methods for these integrals based on finite samples, including the works in [Bekkers) (2020): Finzi al. (2020b) and works following them, which focus more heavily on using the Lie algebra for approximating group convolutions (MacDonald et al. 2022 Kumar al. 2024). The corresponding methods have been shown result in more efficient training data and more robustness to transformations not seen during training, and have been applied successfully in areas such medical image analysis (Bekkers al. 2018 Winkels Cohen 2019).
integrate against
of
computational molecular sciences (Batzner et al. [2022] Batatia et al. (2022) and inverse problems in imaging (Celledoni al., 2002]: Chen et al.) (023). Beyond construction equivariant neural networks, information about the symmetries can be extracted directlty from data (Moskalev et al.
2022 Gruver al.).
It is worth noting that some specific groups allow for construction of more efficient equivariant architectures. For example, recent work on geometric and clifford algebra networks Zhdanov et al.) (2024); Ruhe al. (2024) provides efficient parameterization for groups of symmetries of vector spaces, like Euclidean (Brehmer et 2024) and Minkowski groups (Spinner al. 2024).
A.2 PHYSICS-INFORMED MACHINE LEARNING
In a separate direction, deep neural networks are being increasingly used in scientific computing. particularly in simulating physical systems modeled by partial differential equations (PDEs). This surge in interest is largely driven by limitations of traditional numerical solvers, as they often struggle with demanding scenarios, such those involving non-linear and non-smooth PDEs necessitating
19

---

Published as a conference paper at ICLR 2025
fine discretization, or high-dimensional problems where the curse of dimensionality kicks in (Han ett al. 2018). This is further exacerbated there is a need to solve PDEs across varying domain geometnes, parameters, and initial/boundary conditions, requiring independent simulations for each unique configuration.
The main prominent approach that has emerged is based integrating physical laws and constraints into the learning process to solve PDEs (Raissi et al. (019). These neural networks have been dubbed PINNs, and have been shown be useful for solving forward and inverse problems for PDEs (Raissi
et al. 2019 Sun et al. 2020 Zhu et al. 2019 Karumuri et al. 2020; Sirignnno & Spiliopoulos 2018
regularizing learning (Zhang et al. 2024..
PINNs are able to handle irregular domains and complex geometries, overcoming some of the challenges of classical methods. However, classical PINNs are not without limitations (Grossmann
et al. 2023 Krishnapriyan et al., ( The original method is unable to handle varying parameters or input and boundary conditions (IBCs), and often faces challenges with training due to unbalanced gradients and non-convex loss landscapes, requiring careful hyperparameter tuning preconditioning Müller & Zeinhofer (2024).
To address the challenge of solving PDEs across diverse scenarios, neural operator models have emerged as promising solution. Unlike standard DNNs, which approximate functions over finite- dimensional vector spaces, neural operators learn mappings between infinite-dimensione function spaces. These learned operators map parameters or IBCs to corresponding PDE solutions, overcoming the computational burden of performing independent simulations for each scenario. Combining such neural operator models with a physics informed loss has resulted in physics-infore informeed neural operators (Li al., bb2 Wang al. 20211), oftentimes improving generalization and physical validity, compared to their physics un-informed variants. In what follows, we concern ourselves precisely with this class of approaches.
A.3 WHY LIE POINT SYMMETRIES?
What is appealing about LPS, is that for given PDE, under certain regularity conditions, all such (simply connected) symmetries can be derived the form of one-parameter transformation groups (Olver 1993). For many classic PDEs these have been derived and can e.g. be found in (lbragimov 1995). There further exists symbolic software packages, c.g. Baumann (198), to derive symmetries for general PDEs.
20

---

Published as a conference paper at ICLR 2025
A.4 SYMMETRIES OF PARTIAL DIFFERENTIAL EQUATIONS
Partial differential equations (PDEs) are the most generic mathematical model used to describe dynamics of various systems in sciences. In what follows, will be working primarily with open subsets of euclidean space, but these ideas generically extend to more complicated spaces. The formal presentation here closely follows Olver Olver 1993), and an interested reader referred there for more in-depth discussions of the concepts.
This subsection introduces the concept of PDE point symmetries transformations that map solutions of given PDE to other solutions of the same PDE. To analyze these symmetries, traditionally focus is shifted to Jet spaces, where the problem of studying PDE symmetries is transformed into the simpler task of studying symmetries of algebraic equations. By employing one-parameter groups of transformations, it becomes possible to systematically derive all such symmetries by examining their infinitesimal actions.
Suppose we are considering a system A of n-th order differential equations involving, p independent variables z (z'...., x²) E X = IR,, and dependent variables (u....., u E U R PDEs define relations between various derivatives of the function, for denoting which we use the 8* f(a)
following multi-index notation for (fi... " jk), with 0, fz) 32/10222
Jet spaces and Prolongations: Consider f + U smooth and let Ug be the Euclidean space, endowed with coordinates v 0, (xx) so represent the above derivatives. Furthermore, (()) U x U1 x... Un be the Cartesian product space, whose coordinates represent all the derivatives of functions u = f(x) of all orders from 0 to n. The total space X x Ul7),, whose coordinates represent the independent variables, the dependent variables and the derivatives of the dependent variables up order n is called the nith order jet space of the underlying space X U. Given smooth function prolongation of f, defined by the equations u = f (x)). Thus pr(*) f is function from X to the space (?),,, and for each in X, pr(") f(r) is vector representing values of f and all its derivatives up to order n at the point x.
fz), there is an induced function ('") pp(")) f(x), called the r-th
PDE Symmetry groups: With the notation defined, we can now write down the PDE, given system of equations A (z, m()))= = 0, with A X x (m)) + R' smooth. The differential equations themselves tell where the given map A vanishes on X (m)), and thus determine a subvariety
SA = { {I, A w!) o} C X x U())
Solutions of the system will be of the form u f(x). A symmetry group of the system A will be local group of transformations, GA, acting some open subset M C X U in such that "G transforms solutions of A to other solutions of A
Prolongations of group actions and vector fields Now suppose G is local group of transforma- tions acting on an open subset M c x U of the space of independent and dependent variables. There an induced local action of G on the n-jet space (n),, called the n-th prolongation of G de- noted (*) G. This prolongation is defined so that transforms the derivatives of functions fa) into the corresponding derivatives of the transformed function f(ā). Consider now v a vector field on M, with corresponding (local) one- peeeaerreeter group exp(ev). The n-th prolongation of denoted (r) will be a vector field on the n-jett space (n), and is defined to be the infinitesimal generator of the corresponding prolonged one- mne-parameter group (r" (exp(<v)) In other words,
pr (n)
(r) [exp(Ev)] (z, u()))
ds s==
for any (x, (m) E (n). These can be derived using Olver 1993
Theorem 2.36).
Theorem A.I (2.31 in Olver 1993/). Suppose A (xr, (m)) = 0, is system of differential equations
of maximal rank defined over M X U. IFG is a local group of transformations acting on M, and
(M())
A (z, 22 an))
=
whenever
A A(x: (1) 0
for every infinitesimal generator of then G is symmetry group of the system.
21

---

Published as a conference paper at ICLR 2025
A.5 DEEP LEARNING FOR PDES
The abstract discussion of symmetries of PDEs above only concerned itself with transformations leaving the set of solutions invariant. In practice, we are often interested in problems, for which the solution is unique due to either initial or boundary conditions (IBCs). In this paper, we primarily concern ourselves with temporal PDEs the form of
A ue + [ [|] O, and w(0, 2) f(x),
JEΩ, E [0, IE 0Ω, E [0,
u(t, x) = g(t,z),
where u44, a:) E Rediu u is the solution to the PDE, denotes the time, and æ is vector of spatial coordinates. Q C Rddm is the domain and Dr a (potentially non-linear) differential operator. f() is known as the initial condition function and (f, describes the boundary conditions. As discussed in Section 1.2 we are interested in physics-informed approaches, wherein (t, x) is parameterised neural network uo (t, with parameters 6, and the network is trained by minimizing loss comprised of two parts: L(0) LPDE + Damu- The first term encourages the network to satisfy the PDE, while data-fit is supervised loss, enforcing IBC and ensuring good fit with the training data. Specific forms of the losses, and further discussion can be found in Appendix D.2
A.6 LIE GROUPS
In this section we briefly recall some necessary notions from Lie group theory. For a more compre- hensive introduction, see Kirillov (2008). A Lie group G' is group with smooth manifold structure such that both the inverse and multiplication maps are smooth. A Lie algebra is a finite dimensional vector space V with alternating bilinear map [. V x V + V called the Lie bracket, satisfying the Jacobi identity. Given a Lie group G. the tangent space at the identity = T.G naturally has the structure of Lie algebra. A 1-parameter subgroup of G is simply a group homomorphism IR -> G, where is group under addition. For every Lie algebra vector v there is a unique 1-parameter v. There is smooth map exp: g + G which sends v to &. (1). This has image contained inside Go. the connected component of the identity E G. If Go is compact then exp surjects onto Gg. but in the non-compact case it doesn't 't necessarily surject.
subgroup du such that Φ. (0) is the identity and d', (0),
For Lie group G acting smoothly on a manifold X we have a version of the orbit-stabilizer theorem.
Theorem A.2. The stabilizer G2 S G is a closed Lie subgroup, and the natural map G/G,, - X is an injective immersion with image Gr.
22

---

Published as a conference paper at ICLR 2025
B CHOOSING THE ENERGY
Based on discussion Section 3 we wish to construct an energy, which contains information about the domain (or dataset) and the energy levels are transverse orbits. The first we achieve by using a variational autoencoder, and choosing the ELBO as the energy itself; and second by training an appropriate (convex) adversarial regularizer. Thus the overall energy
C a.R for some constant >> 0 and terms below. We emphasize once again that this is simply construction, which have observed work well practice. Effectively all that is required is align minimas of such an energy with the training set doing this directly is not that simple and should be fruitful direction for future work.
Variational Autoencoders Variational Autoencoders (VAEs) (Kiggma 2013) are class of genera- tive models that learn the underlying probability distribution of a dataset. VAΣs consist of two main components:
This component maps an input data point to lower-dimensional latent representation z. Instead of directly mapping to a single point in latent space, the en- coder outputs parameters for probability distribution, typically a Gaussian distribution, g(>[æ=) = . (z; (x)). Here, (uz) and (x()) are the mean and variance of the latent distribution, respectively, and are outputs of the encoder network.
Encoder:
2. Decoder: This component takes a point in the latent space and maps it back to the original data space, reconstructing the input data f.. This is achieved through probability distribution p(z[2), also often modeled as Gaussian.
VAFs are trained by maximizing lower bound on the data log-likelihood, known as the Evidence Lower Bound (ELBO):
((x) (log - DRL (
(6)
The first term in the ELBO encourages the decoder to reconstruct the input data accurately. The second term is the Kullback-Leibler (KL) divergence between the learned latent distribution g(=(ar) and prior distribution p(z), typically standard normal distribution. This term acts as a regularizer, encouraging the latent space to be well-structured and preventing overfitting.
Adversarial Regularization Adversarial regularization Lunz et al. (2018) Mukherjee et al. (2024) a framework for solving inverse problems that incorporates a discriminator network to learn and impose data-driven regularization. The discriminator network is trained distinguish between unregularized reconstructions and real solutions obtained from ground truth data. The regularizer trained with the objective of favoring solutions that are similar to the ground-truth images in the training dataset and penalizing reconstructions with artifacts.
The adversarial loss function for the regularizer be formulated as:
Ladv (R) (x) Ex) (R())
where Rt is the regularizer, is ground truth image, 4 is noisy measurement, px and pa are the data distributions. This loss encourages the regularizer to produce small output when a true image is given input and large output when it is presented with unregularized reconstruction.
In our setting, the noisy measurement distribution is chosen to be = x, for x ~ px and g E G sampled randomly. Thus the regularizer learns to discern between true looking images and group transformed ones - leading to level sets lying transverse to the orbits.
Adversarial regularization has been shown to provides flexible data-driven approach to impose priors in inverse problems, which lends it readily to be used in our problem.
23

---

Published as a conference paper at ICLR 2025
t
4
4
4
4
9
9
9
S
9
9
9
9
3
3
3
3
3
3
3
3
7
7
7
7
7
7
2
2
2
2
2
2
2
2
(a) Affine MNIST examples
(b) Homography MNIST examples
Figure 6: MNIST canonicalization images for both Affine and Homography groups, as described in Section 4.2 From left to right: Original MNIST image; Canonicalization of the original image; Original image distorted by a random homography transformation; Canonicalization of the distorted image. This figure illustrates the equivariance properties of energy canonicalization.
C FURTHER MNIST VISUALISATION
Of particular interest are the failure modes of canonicalization, warranting further theoretical investi- gation. We present the failure modes in Figure 7c and can see from the examples provided in majority of the cases the reason why failure modes exists is due to existence of group transformations resulting in completely different digits! As such, the overall distribution is not exactly equivariant.
24

---

Published as a conference paper at ICLR 2025
Tue
Corerad
7
T=
E-091
1--444
t- MLM
T-39035
--10P%
3RAM
9
9
Class
Class
C-5SSE
D333
5-2523
C-123
3
3
3
3
3
3
3
3
3
Class
Class
7
7
7
7
7
7
7
F-TWE:
19225
P-TEE
3
3
2
2
2
2
2
2
2
2
Class 3
Class 2
--10/N
S
S
S
S
S
Class 2
Class
1-12.1)
2-
C-0113
C-1344
2
6
6
6
6
6
6
6
--
x* 16.2
--1tb2r
Class 2
Class
g
8
g
8
16
4
O
Class
Class
C-5013
C-96.
2-
C-135
4
4
4
4
4
4
E-3BB
t+ 95
E+ 18,25
Class 4
Class
9
9
9
9
9
g
9
3
3
3
3
3
3
3
3
Class 0
Class
4
7
7
7
7
7
7
7
7
4
L-1XX
L-10
C-19%
C- 19375
C- 19023
C- 1--8
Class 9
Class 4
7
Z
7
Z
7
7
1
Class
Class
(a) Affine MNIST examples
(b) Homography MNIST examples
(c) Misclassification examples
Figure Expanded Visualisations
25

---

Published as a conference paper at ICLR 2025
THe WART
THE Corear neweer
VIRTT
THE
(a) TSNE post-affine canonicalization
(b) TSNE post-homography canonicalization
Figure 8: TSNE plots
D FURTHER PDE DISCUSSION
D.I CANONICALIZATION AS UNIFYING VIEW
The idea of putting the evolutionary equation and initial into canonical form prior to evolving is not novel, and fact quite standard in numerical analysis textbooks. The first step in solving an ODE numerically is to normalize the equations and remove units and normalize the domain. For example, when solving a time evolution equation on the domain x E [0, 15]. one would first normalize the variable by rescaling it to be in [0, 1]. and solve on this domain instead. Floating point representations are much coarser for smaller values and as such, this can lead to lower numerical errors.
Another example is the recent work achieving exact unit equivariance which is precisely scaling group canonicalization (Villar et al. [2023). A similar problem also often arises in the problems of image registration (Kanter & Lellmann 2022) with an infinite-dimensional Lie algebra of diffeomorphisms. Another particular example being Procrustes analysis (Gower 1975). when considering only the homography group.
In this work, we present that canonicalization can be seen as the unifying approach, especially through the lens of energy canonicalization. Insights from the problems above can then be used improve the optimization, and used for PDE symmetry groups
D.2 ON DEEP LEARNING FOR PDES
In what follows we demonstrate the effectiveness of canonicalization across three example PDEs Heat, Burgers, Allen-Cahn with baseline models taken from three recent state of the art ohysice-informed machine learning papers Akhound-Sadeeh et al. (2023); Wang et al. (2021b); Herde et al. (2024) respectively. In each case take the pre-trained neural PDE solver and demonstrate failure mode when evaluating on out-of-distribution initial conditions. We then optimize for a group action derived from the Lie algebra based energy minimization and report results after the canonicalization technique is applied. Results and additional experimental details for each PDE are provided in the following sections. In PINNs, the PDE solution u(t, x)) parameterised parameters 0. The loss function is then compromised of two parts, with different weightings:
neural network us (t, x) with
(()) CPDE + Ydata Cdata-fit
The first term is the physics-informed objective of Raissi et al. (2019), ensuring the neural network satisfies the PDE. The PDE loss is the calculated as the residual, with derivatives are calculated using automatic differentiation, as calculated on finite set of points (f, r):Naam which are sampled from inside the domain [0, T] obtain the PDE loss:
Naom
CPDE
|10,uq(ts, 24) + Dx [ue(tz, x:)1l?
(8)
Naom
The second term, is a supervised loss enforcing initial and boundary conditions:
Nic
Nbc
Lasta-ii =
1
Σn
uo (0, - (z. 112
1
Σ Ilue (t/,
(t2) 12:
(9)
Nic
Noc
1=1
26

---

Published as a conference paper at ICLR 2025
with 1:Ne Ω sample at which the initial condition function f is evaluated. (t", 2²) 1:Nuc Nb points sampled on the boundary [0, d92. The two loss terms are then weighted by oPINN for PDE loss equation [ aand "data for data-fit equation 9 to be consistent with (aatkAAhh-d-aaeee Sadehh et al. (2023). Within the field neural operators, desirable that the trained network generalizes across various initial conditions, and the above mentioned loss is evaluated for number of initial conditions Itrain.
D.3 HEAT EQUATION
This and following sections primarily recap the relevant facts regarding the symmetry groups of the PDEs considered, which act to motivate the need for new equivariant architectures due to the non-compactness of the resulting symmetry groups, and the difficulty of constructing global parametrisations of either the action or the group considered.
D.3.1 POINT SYMMETRIES
Simplest PDE example is the heat equation, admitting non-trivial group of Lie point symmetries, beyond space isometries.
(10)
This equation admits the following 6-dimensional Lie algebra (note the typo in Akhound-Sadegh
al. 2023 see Olver
1993): Koval & Popovych 2023
1 _ud,, = 2t0: + +xxz= - zudw:
1
= 0z, U2 Oz,
(11)
2
U/ = t0 rud,,
U6 = t'd + txd
(x² + 2vt) udw.
(12)
2v
4u
The resulting one parameter flows are:
expfe1](c): = t
2 = I + C,
à = u,
exp(e>[(e): l ++ 6
T = I,
Ū u,
exp expluy] [e): i=
I,
e/v u.
exp(ea[(e) i= e
á e er
û e 1 'u,
expfes](c) i= t,
$ d,
à
222
exp(ve](e): i = T-rtt
5 T-a
v 11
The only non-zero Lie brackets are:
[U4, U2] -22,,
[Cs, Us! 20e,
[v2, vo] U4,
(13)
[U2, Ug] = U1,
[U4, 2/6] = 1).
[e4, 01 = -111
[va, u1] = -25
(14)
1
[UG, U1]
U3.
(15)
In particular, these would be enough for LieLAC, however in this particular example it is possible show more. To be precise, as shown in Koval & Popovych| (2023), the global structure, as explained below, of the symmetry group (the finite simply connected component) is GH SL((2 R) H(, where H(1, R) is the rank one polarized Heisenberg group.
The point symmetry pseudogroup G of the (1 + -iddieensional linear heat equation constituted by the point transformations of the form
Theorem D.I (Koval & Popovych| (2023) for v * 1).
at B
x Ao
5
7/+ 4
71. 18+8
(x A1t 20)
A1
27
ū = o 0
2/4
(u h(t,))))
4ul +))
2v
where a, B.Y, 5, >1, An and a are arbitrary constants with ad - By = and o # 0. and h is an arbitrary solution the heat equation.
27

---

Published as a conference paper at ICLR 2025
We now describe the product G. First note that there is subgroup R consisting of those elements for which Å1 = Ào = 0 and o = 1. This is evidently isomorphic to SL(2, Therefore
we shall group the a, 9, y, terms into matrix -(2 8 Further we note that there are two
connected components of this group, corresponding to sign of and the group can be written direct product corresponding to this sign
G G' Zn
(16)
Thus we may simply describe the product on G'. For this we may assume that the two elements we multiply have >> For clarity, we shall write out the product rule in terms of the logarithm of The product can now be written as
(A', X1, Yos In lnn') ) (A, AI, Ao, Ino) (A'A, adi + 7 Ao + +B + 6%::
ln Ind 1X6 - (aY) + xbX + 5X1))
Ao (a\) + 7X6))
We see that the subgroup F' consisting of those elements for which = 8 = 1, 0 and > 0 isomorphic to the rank one polarized Heisenberg group H R) via
01 0 0
2
A1
hn
Ao,
A
(17)
0
where we use the standard embedding of H(1, R) into the group x real valued matrices. Note, that while the parameterization here looks the same the one from Koval & Popovych| (2023), is
different (resulting in different o). as that paper, the Heisenberg group and the polarized Heisenberg group conflated, making the product rule unclear. Here fix specific form of the Heisenberg group as a subgroup of the group of 3 x real matrices, with standard matrix multiplication as the product rule. The product formula can then be described in terms of F and R well. Let Φ S22,, R) > Aut(H(1, R)) be the antihomomorphism defined by
A
ad1 + + BA1 + dNo
((A)
Ày In
+
Aood1 (
(18)
In
Haa (A0 + +
where write H(1,R) using the identification F H). Then we can write
G' SLL2, R) K H(1,R)
(19)
As the semidirect product with an antihomomorphism is rarely written out in literature, we write out explicitly what this entails. We have that the underlying set of G' is SL, x R) and the product is
(A', (A, ((A
(20)
In order to decanonicalize, we also require access to the inverse group element. Particularly, for our implementation we either find the action of the inverse of exp(e< exxp(((1t)) as but given access to a global parametrization of the group, we can
cxp(-e111) U2) write down an inverse given element the form of A, H. This (A-1, (A (H-1) Le. for 0,,3, 3,,,,, 7, 6, λo, X1, a, the inverse is
1
-Io (A1
8 -3 a (-x
A-
H
-o
0
1
Writing this out in coordinates this becomes
1
1
8. 1.3,-740 a, AA - ado, -8d1 + 7do: - In
dod1
1 + +7>0) -ado)
(-021
4
4
28

---

Published as a conference paper at ICLR 2025
(Original)
Heat (4 mag
Burgers (Original)
Model
ID
OOD 21
OOD
ID
D412
0OD
DeepONet
0,049884
0U72
8.6572-4
1255
0,0504 6,0014
0,0687 O 00:1
0 08332 O 0547
0.8369 + 0.0857
LieLAC [DeepONet)
04045---
05000 0.0003
0003
0.0637
Table 3: L2 relative eгror for Heat equation and Burgers' equation averaged over [0, T]. Original refers to model trained on in-distribution (ID) data. For heat equation this corresponds to with fixed A = in Equation (21). Data aug refers to model trained on out of distribution data achieved via symmetry transformations, resulting in A E {0.5, 5.0). For testing ID refers to A E {0.95, 1.05}. while OOD refers to ATTs E {0.5, 5.0}. For Burgers' equation. ID corresponds gaussian random fields as in Appendix D.4.2 while OOD refers to ID plus constant shift.
D.3.2 NUMERICAL RESULTS
For the numerical experiments, consider solving Equation TOp with diffusivity 0.1 on the spatial domain [0. 2T). with periodic boundary conditions for & E [0, 16). We train a physics
informed DeepONet (Lu al. 2019 Wang et al., (02Ib) with initial conditions of the form following Akhound-Sadegh et al. (2023 Brandstetter et al. (2022)
K
(x)
A sin (2 L Φk).
(21)
To focus on the effects of data augmentation and canonicalization, we train the model under two regimes. In the first regime (which we refer to as in-distribution), the amplitude parameter is fixed to
A₁ 1, with the other parameters as K 1, 4 2, Φb = O. analogous to PINN training. In the second regime, the amplitude parameter is sampled from A [1.55, 5.0]. representing a broader operator training distribution.
The model is then trained using the physics loss with sampling weights of oPINN = 150, Vdata 20, we use a learning rate of 0.001 and batch size of 8 and train for 100k epochs following AAkhound Sadegh et al. (023). Training data consists of N fttain initial condition functions. evaluated at Nic = 200 points at to evaluate ug Nb. = 100 points on the boundaries x = 0 or x = 2n to enforce the periodic boundary conditions and finally Ndom 500 points in the interior of the domain where the residual physics loss is enforced. The official implementation of DeepONet from the code and examples provided by the library DeepXDE (Lu et al. 2021)
used.
We report results for both training regimes, averaged over 10 random seeds, in Table 3 The first rows in both tables show the accuracy of the standard DeepONet. In both cases, the first column shows slight improvement with LieLAC canonicalizing to the average of the tight distribution. In Table the second column shows DeepONet's failure to generalize when tested on N fteet = 10 test functions sampled A ~ 2((0.5, 5.0), which is outside of the training range. Applying LieLAC, following the canonicalization pipeline shown Figure 2a] restores test accuracy to in-distribution levels. In Table 3 extending the training range to A ~ [[0.5, 5.0] improves generalization by exposing the model to a broader range of amplitudes. However, even in this setting, LieLAC outperforms data augmentation by leveraging the canonicalizing group action.
Choice of Energy Given the initial conditions in Equation 21 we see that the distribution of initial conditions can be described by the bounding box of the jet. The energy is therefore chosen to be the distance the trained domain. Letting jet, (ug, To, fo) and jetn = (1, S, 0), the energy is:
Eteat (No, to. Uo 2y. zz.ty)
dist (max(jeto), max(jotn)] + dist [min(joto). min(jet + dist [ry, Q] + dist [ty, 16]
(22)
D.3.3 FURTHER HEAT ILLUSTRATIONS
Figure shows canoniculicalisation of the initial conditions, out-of- and canonicalized DeepONet predictions.
29

---

Published as a conference paper at ICLR 2025
Truth
01 (ught)
LOS, bw
u(x;t)
R- 0.10
Figure 9: Canonicalization for the Heat equation given out of distribution ux( and improved
equivariant prediction
D.4 BURGERS' EQUATION
D.4.1 POINT SYMMETRIES
Burgers' Equation combines diffusion and non-linear advection, which can be related to the linear heat equation above via the Cole Hopf transformation. The nonlinearity of the equation results in more complex dynamics, e.g. via shock formation, while still admitting non-trivial Lie point symmetries.
Ut + uug 1u2z
0
(23)
This equation admits the following 5-dimensional Lie algebra (note the typo in Akhound-Sadegh
al. (2023)), see Freire (2010): Koval & Popovychi (2023):
U1 2z, C3 2t0; + xē2- ud.,, Us Po txx, + (x tu)
(24)
04 t0 + δ.. U2 = O, The resulting one parameter flows, for arbitrary constant, are:
exp v1 (c) i t,
I + e,
u.
exp 02 (c) i t+
I x;
u,
exp V3 (e)
[ [= e
I e
exp (e): [= l
I = x + el,
ú = u + 6,
exp (u]((e) i T-et
I = T-
= u(1 Et) + Ex,
The non-zero brackets are:
[Ug, U2] - 221
[V3, Ub]
20g
(02, Ug] U3:
(25)
[U2, v44 01,
[V3, v4] U4, [0g, vi] -1,,
(Us, vi] -
(26)
Once again, this would be enough for the LieLAC formulation, however in this particular example, we can say more:
Theorem D.2 (oova & Popovych 2023 for v + 1). The point symmetry group GB of the Burgers' equation consists of the point transformations of the form
i at 3 3
x
I
A
ū ū= ( 8)u 7a + \18 \07
7t 5
where a; 3, A1 and Au are arbitrary constants with aó 3 =
And similar to before, the global structure of the group can be inferred, with the group GB isomorphic to the group S2(,, R) K2 (R², +), where & SL(2, R) + Aut (IR?, +)) is the group antihomo- morphism defined by p(a, 8.7,6) 6) = (A1,An) + (Ai. Ao) (i B, 7, 6). Thus, the group G is connected.
30

---

Published as a conference paper at ICLR 2025
+ - (a
3/ 8 ) (21,20) + A0).
21 (a,3,7,0) --=
The standard conjugacy action of the subgroup 3 G the normal subgroup R GB is given by (A1, 2o) (o) e (,,8,7,88 7).). The inverse be found similar Appendix D.3
D.4.2 NUMERICAL RESULTS
For the numerical experiments, we consider solving Equation (23) on the domain 52 [0. 1], with periodic boundary conditions on & E [0, 1]. An example of a canonicalized solution is shown in Figure 2
with further illustrations in Figures I0 and 11
For Burgers' equation we take the code and pretrained Physics Informed DeepONet (PI- DeepONet) model weights provided by Wang et al.] (202Tb). This includes a model trained on Ny fitrain = 1000 initial conditions sampled from a Gaussian random field (GRF) parameterised as N (0 + 5" on the periodic domain ? [0, LJ x [0, TJ [0, 1] x [0, 1] discretized by a meshgrid with 101 x 101 points.
In order sample aut-of-distribution initial conditions from the GRF we apply uniform 0.2 shift to Up away from the mean of O. Figure shows surprisingly even this small perturbation is enough to break the pre-trained PI DeepONet. Interestingly in both in and out of distribution cases LieLAC learns a group action that is able to redistribute around zero mean, leading to more accurate prediction.
Choice of Energy Given the slightly different form of initial conditions to the heat equation, now sampled GRFs adjust the energy focus the bounding box of the domain and mean of uo- The energy is therefore chosen to be the distance to the training domain. Letting trunko = ( to) and trunko (2, 0), the energy is:
Earrgen (ro, lo. uo, k,, = dist [max(<a)) min(2n), 1))
(27)
+ dist [4o, Qt] + dist [mean[uo]] 0] + dist [xy. Q + dist [ty. S2]
D.4.3 FURTHER BURGERS' ILLUSTRATIONS
Figure shows canonicalicalisation of the initial conditions, duitr--disttibution and canonicalized DeepONet predictions.
40
u(x:t)
00
10
03
04x. ge
05
-4
6.4
04
62 3.00
600
6:3
853
033
Figure 10: Illustration of Burgers' equation for in distribution GRF up(x)
31

---

Published as a conference paper at ICLR 2025
Figure 11 shows canonicalicalisation of the initial conditions, out-of-distribution and canonicalized NeepONet predictions.
Truts
0.4
0.2
S- QUALITI SE" UEX 1-17)
0.0
0,20.
u(x,t
-02
3.
X
Figure 11: Canonicalization for Burgers' equation given out of distribution uz( and improved equivariant prediction
D.S ALLEN-CAHN EQUATION
The Lie point symmetries of Equation can be found either by hand using the computational software of Baumann 1998). Both yield 4-dimensional Lie algebra with basis vectors:
O, 1/2 = 0g, Vi = 0
yy + x.
(28)
Noting that there is coupling between the space, time and the dependent components, we identify this group as the connected component of the group of isometries of the underlying space. worth noting here that the full symmetry group can be observed be EE(). while the group
generated by vectors in 28 is SE(2), i.e. the identity connected component of E(2).
Training data in the form of:
The POSEIDON (Herde
2024 model is trained on random initial condition:
K
Mo(a, y) = Σ ag (²+))) sin((Tia) sin( Va, y E (0, 1),
(29)
+1=1
where K is an integer drawn uniformly at random from [16, 32]. and r ~ U and 0gy ~ 1.1] Based this, one obvious limitation of the training data is that all sampled initial conditions are zero on the boundary. However, due to the semi-group property based training Herde et al. (2024), as trajectories are sampled at different points, the values on the boundary are not guaranteed remain zero over time, As result, we expect the model to not be equivariant under the group of symmetries of ACE, as illustrated on Figure 12 but we do not have an obvious way to construct the energy. Based on the positive results in Section 4.2 we employ the same approach, by training variational autoencoder and an adversarial regulariser. For the experiments, we considered the TINY POSEIDON model from Herde et al. (2024).
Testing data To test the equivariance, we use a dataset modified initial conditions, achieved by transformations of Equation (29).
K
(0. = Σ (P+ sin(Tifa - mo}) sin sij - mnh), Va, y E (0, 1), (30)
where ).) denotes the sawtooth function, with parameters as in Equation 29 with To. yo ~ Up.
"Noe here that this is slightly different from the original paper. The form helow was chosen, the generated initial conditions provided by (Herde al. 2024) either due pre or post-processing
be periodic.
32

---

Published as a conference paper at ICLR 2025
CONREE
Fecttec
332
R2I
330
535
83
83
440
130
52
521
H0M
-30
Figure 12: Example of failure equivariance of the original POSEIDON model.
D.5.1 FURTHER ACE ILLUSTRATIONS
33

---

Published as a conference paper at ICLR 2025
Figure 13: Further canonicalization illustration for the ACE.
34

---

Published as a conference paper at ICLR 2025
E MEASURE THEORY BACKGROUND
Let X be a set. Then a o-algebra Σ is subset of 2* such that
XEE
If E then X X
For any countable collection (An)=- of sets in E, the union U A, EΣ
set A E is called measurable. A pair (X, 2 of set and a o-algebra called a measure space.
Suppose (X, Σ), (Y, E') measure spaces. Then a function f X - Y is measurable if for every measurable set E E' the set f (A) is measurable.
Given measure space, measure on it is set function μ E [0, oo] such that
m(0)= 0
For all countable collections (An) of pairwise disjoint sets in E
m U ==1
A,) = -0=-mAA ((An)
(31)
i=1
We may also define signed and complex measures, which simply replace the codomain of measure with either IR U (-oo, } or C respectively.
f (X,E) Σ) + (Y, E is measurable, and f is a measure on (X, E), we define the pushforward measure fa as follows
f.u(A) = P(F
(32)
There is always a measure on (X, Σ) which simply sends every measurable set to 0 we call this the zero measure. A finite measure is measure μ such that (xX) < co. A probability measure is
measure μ such that ((X) 1. Given a finite non-zero measure on (X, 2) we may always define
a probability measure p' on (X, 2) as follows
pl (A) (A(A)
(33)
xXX
We will often call this the normalized measure associated p.
If X is a topological space there is a canonical ralleerra associated to it, known as the Borel r-algebra B(X). It is the smallest o-algebra such that all open are measurable. Then any continuous function between and is measurable when X and Y are equipped with their respective Borel o-algebras. We often call the pair (X, B(X)) Borel space and when clear we will simply refer to this pair by X.
For topological space S and real number 2 there is associated Hausdorff measure of dimension which we shall denote Hg. To S we can also define what is known as the Hausdorff dimension, (Sm which always non-negative. There is then Hausdorff measure of that dimension, which we write as Hs. All Borel sets are measurable for this measure. If S is compact then (S) < oo, but may still be the zero measure. If it is not zero then we may as above construct the normalized Hausdorff measure. In this case this will be Radon measure. For S = [0, 1J" the Hausdorff dimension is n and the normalized Hausdorff measure is exactly the Lebesgue measure. Similarly for S finite and discrete, the Hausdorff dimension is 0 and the normalized Hausdorff measure is the normalized counting measure.
The Hausdorff measure is important to consider beacuse of Frostman's 's's Lemma Mattila (1995). Lemma E. (Frostman's Lemma). Lei be a Borel subset of R. and let s > O. Then the following equivalent.
(A) >
There is an unsigned Borel measure is on R.". satisfying ((A) and such that
p(B(I,r)) S
(34)
for all RR. O.
35

---

Published as a conference paper at ICLR 2025
Fintle canonicaliaatio
Veate mo01
(xa)
O'Cng" (X]
CCang"
Can (XX
nntt
OChna(X)
OCano(X))
(ano(X)
ooeaaan WPrag
WOCaa
WCCang (X)
WCang (X)
WF
WOCano(X)
(CCang(XX)
WCang(X)
G comppett & Finils framns
Grenecomped
Figure 14: This diagram illustrates connections between the various notions of frames and canonical- izations introduced previously and in this work. The top row the finite frames and canonicaliza- tions, which all map vertically into their weighted versions by taking normalized counting measures. Inside of each we have the sequentially closed subspaces of those that perserve continuity.
Let A be compact with Hausdorff dimension greater than zero. Then know that the Hausdorff measure at the Hausdorff dimension is finite. By Frostman's Lemma, the Hausdorff measure of A being zero means that we cannot find any probability measure on A that interacts nicely with the metric structure. This shows that our assumption of not allowing sets whose Hausdorff measure is the zero measure is reasonable.
F
MORE ON FRAMES AND CANONICALIZATION
Recall that a frame is defined as a function F: X + 2 such that for x E X and E G
F(gx) = gz(
(35)
Lemma F1. Let F. be frame for the action of a group G on a set X. Then Vo E X, F(x) is union left cosets of Gx C G.
Proof. Indeed, let e X and let g E Gz. Then
F(x) = F(gx) = gF())
(36)
Hence we see that
F(z) GxF(
(37)
Therefore we see that F(x) is union left cosets of G...
F.1 WEAKLY-EQUIVARIANT FRAMES
In Dym et al. 2024) weakly equivariant frames are introduced. These are maps μ X PMcas( G)
satisfying equivariance up to averaging over stabilizers.
In the review process it was noted that this bears similarity to Nguyen et al., and indeed, in our framework their method is exactly the weighted frame defined by miniziming the energy given by miniziming orbit distance. Because the symmetry groups they consider are compact and act algebraically, they are able to obtain more concrete descriptions of these frames than are able to in the general case. We do expect that in the compact case with algebraic actions, similar results should hold for energy minimizing canonicalizations.
First, for any x E X and Radon probability measure A C G,
Dym et al. (004) define for every measurable
(7)a(A) = JoEG. sAT(A)ds
(38)
36

---

Published as a conference paper at ICLR 2025
where ds is the normalized Haar measure on Gg. For frame j we let Pz (Ma)x. Then j is weakly cquivariant for all a E X and E G.
F g2 = 9wH4.
(39)
This definition does not extend to the case where the action has non-compact stabilizers.
To deal with this propose imposing what call equivariiac. Again consider simply map μ : X + PMeas(G). For point I E X, let Ty. G + G/G, be the quotient. We require that
(Tgz)=fgge = (To
(40)
Lemma F.2. If the stabilizer G2 is compact then (4) = (p")a if and only if (TT)aft = (T2)ap'.
Proof. For ease of notation let Ty. First suppose that (14), (p)2. Let A S G/G, be measurable set. Note that π (A) is a union of cosets of Ga. Hence it is invariant under multiplication by elements in Gz. Therefore for all G2 have s-1x-1(A) (A). Thus
(A) / 8.1(11 (A))ds f - , (AA)
(41)
JEGG
aEG. and similarly, (1() (A)) ((A) = pT (A)). Hence,
n( (A) =
(42)
Therefore, T.H4 Totí'. Now suppose that TT = Tapi". Then for any measurable set A C G, we have that
p( - (T(()) = nT, (A)) = n. (T())) = μ ( (T(())
(43)
We have
(T) U SA
(44)
sEG,
Using Fubini's theorem and the fact that both and are probability measure, we may write Equation (38 as follows
S s,p(()ds
(4)=
J.EG.
(45)
S I xa(sg) dp(g)ds
(46)
-c Jol Xa(sg) dsdp(g)
(47)
(48)
Consider the inner integral as function ob on G. We claim that & is Gi-invariant, Indeed, let g' = 6'9 for some s' E G.. Then
= s(Ag
(Ag = ((g)
(49)
as is the normalized Haar measure, hence is right-invariant. Therefore & factors through the quotient map æ. Let o' be the corresponding function on G/Gz, then & = o' Hence
[ ((g)dp(g)
(p)²
(50)
[
Jac.. ((()
(51)
Jaja.
((()))d(..)))))))) (())
(52)
[ ((g)du'(g) (1))
(53)
Corollary F22.1 If the G-action on X has compact stabilizers at each point, then π. equiivaaain equivalent to weak equivariance.
37

---

Published as a conference paper at ICLR 2025
F.2 ON PROJECTIONS
For a weighted orbit canonicalization ns, Pp B(X,R) + B(X,R)C turns out to be a projection, as
PaPa(o)(2) = f Pulo)(s)
{ GG Pu(o)(s) dpz(s) (x)
(54)
G-invariant
F.3 ENERGIES AND CLOSED CANONICALIZATIONS
By the construction of energy minimizing frames, we have function from energy functions to weighted closed canonicalizations. We now describe the image of this.
Theorem F.3. Let R E WCCang satisfying the following:
Kg is the normalized Hausdorff measure on its support.
For all E Gy, supp Ky supp Ks n Ga.
Then there is an energy function E such that KE = K.
Proof. Suppose is closed canonicalization satisfying the requirements of the theorem. Let (x) = 1 Xsupp Na (z).
Now consider the minimum sets
ME (x) min Ey: yEGr
(55)
We have that supp K₂ is non-empty closed subset of Ga, hence the minimum of E, over Gæ is zero. Therefore ME, (x) = (0) n Ga. In other words it is those y E Go such that y E supp Ky. Further, by assumption, for any Gz,
supp By supp Kr n Gy
(56)
Therefore supp Ky if and only if y E supp Hence we see that
ME. (x) = supp Kig
(57)
But the by the assumption that K2 simply the normalized Hausdorff measure on supp Kg, we get that the KE
G PROOFS
Proposition G.1. Fg defines an equivariant frame, i.e. FE(gr) gFFl#) any G, x E X.
Proof. Suppose that & E X' and g E G. Then
FE(gx)
arg min E(h -lgg) kEG
arg min E((gh) gz) kEG
arg min -lh-1z) kEG
-1x)
glarg min E(h E(h-1r)] = gF)).. hea
Here, we have simply used that group multiplication is bijective when either of the arguments is fixed.
Proposition G.2. ME(F)) is G-invariant and non-empty.
38

---

Published as a conference paper at ICLR 2025
Proof. Mz is G-invariant, the infimum is taken over the orbit Ga. and this orbit G-invariant. By the cocrcivity assumption on E, this is non-empty. Indeed, let x E X and let N be large. Let K be a compact set of X such that E(y) N for y E X. Taking N large enough, we can assume that K n Gz is non-empty. Then compact set K, hence is compact. Therefore E attains its minimum on this set, hence Ma is non-empty.
that Mx(x) S K n Gr. K n Gr is closed subset of the
Theorem G.3. Let G be Lie group acting smoothly X manifold. Let ja G - X be defined as g - g 1 x. Then we have a map a :WFrac(X) - WOCana defined by
= a(F)()) (Fa))((())
(58)
Then for any F,F E (Frag(X) such that aFF) = a(F)) we have that (T2). (TT) for every x E X where Tx :G G + G/Gz- If each G. has probability measure Pa, then ar has image exactly the set of weighted orbit canonicalizations with compact support.
Proof. First we check that a as defined above lands in WOC Let WFrac(X). Then (z) is G-invariant as for any measurable A C X.
(A)) = (F(z)(()))
= F( (gx) E A}
=(or)(((E((++- A))
mm((((c))(((bbeC))) (gx) A)))
=.2[(7T))((((MEE)))) (gx) E A]))
=g((TTa)((be))))) A))
E A))
Gl ((F)(()())
Thus we see that (F) E WCana(X), as (F) = F()) = 1, as F(x) is a probabil- ity measure. To check that it is a weighted orbit canonicalization, first note that it is clear that supp a (x) C Ga, as for any y & Ga we can find an open neighborhood U of y such that (U) = Ø which is always ((r)-measure zero. Suppose that y E Gr., F(x) is a weighted frame, hence it is supported on some compact set K C G. We claim that there is open set U containing y such that je (U) C GYK. Suppose not, then for every open neighborhood U of y there is an gu E K such that (g) U. Then the collection (gv) over the neighborhoods of y forms a net in K. As K is compact, this net must have a convergent subsequence. Call this subsequence (gn) and let g be the limit. As the action of G on X is continuous we have
g
lim In.
(59)
For any U open neighborhood of y may, by definition of net, find an N such that for any
n 2 N, In x E U. As X is Hausdorff, we then must have that the limit lim In x is simply y. But this contradicts y Ga. Hence there is some open neighborhood U of y such that (U) C G K. Then (U) is measure zero with respect to F(x). Hence (F)) = O. so & supp oFF) Therefore we that (FF E WXC
Suppose that F,F' are weighted frames such that ((F) = =(F'). Let x E X. Then for any measurable A X,
= = = F())
(60)
Both F and F are compactly supported, say on compact subsets K, K' respectively. Consider the natural map Jo G/G, + X. This is injective for group-theoretic reasons. Let ग + G/G. be the quotient map. Then = (K) and x x(K") are compact subsets of G/G.. as they are continuous images of compacts. Now as ((K) and π () are compacts which inject into a Hausdorff space X,
Jz restricts to homeomorphisms π (K) = J. (K)) and similar for K". Using Equation (60) we see that for any U C G/G,, open we have
(V)) =
(61)
39

---

Published as a conference paper at ICLR 2025
Therefore F and F² agree the open sets of G/Gz, but then that implies that they agree the entire Borel o-algebra of G G/Gg. Hence T.F= = F,F'' as desired.
Finally, let (OCann(X)) have compact support. Let z X. Then 2s is some probability measure on X with supp 1/ Ç K S Ga for some compact K _ Gz. By Theorem A.2 we have that homeomorphism on K, and using
Jz G/G + X is immersion with image Ga. Ju restricts this we can define a measure byg on G/G, which agrees with 2y under this homeomorphism. supported on the preimage K' of K which is also compact. π : G - G/G2, is principal bundle with fiber G.. Therefore, locally on G/G,, π is the trivial fibration. We may cover K' by finitely many open sets over which trivial call them UL....,ú, Un. Let V; (U,)) = U4 x Gz. Now define
is
V = VAUU) jCi
(62)
Then these are disjoint measurable sets which cover (K'). They are all trivially fibered over the base Now we define a measure v using these, v will be supported on π TK').. Let A be a measurable set in G. Firstly we may assume that A S TK').. Let A, An n VL. Then the A, disjoint measurable sets which cover A. So the value of y on A is determined by the value of v on each A,. We define this value be
m(A,) (v/2 x (A.)
(63)
where here we use the homeomorphism V/ = TV,) x G2 to put the product measure on this set. T = tyy and this equivariant as Ggn = g Hence v n. -quqqiivari
Note that for the above theorem to work, need to assume that each G₂ has probability measure. This is automatic in the case that each G is compact, as in that case we may take that normalized Haar measure it. It will also hold in the case that each G2 has non-zero Hausdorff measure.
Theorem G.4. The sequential closure of WOCang is WCCana(X)
Proof. First we show WCCanc(X) is sequentially closed. Suppose that (f*) is a sequence in WCCang(X) converging to f E WCang(X), Fix & E X and y E X' Gz. Then it is sufficient to show that y É supp fa. As X is normal Hausdorff, by Urysohn's lemma there exists continuous 6 X + [0, 1] such that of (Grx
and <((y)) = O.
As fr + f, for any O, we can find N such that for any 2 N we have the following
(x) f()) <
(64)
We may easily compute (z) by the fact that supp fn(x) C Gz.
[x( = /
d =1
(65)
Joz
Similarly,
f 2 df =0
(x())())
(66)
Gy
Hence, by taking the limit we compute that (z))(z) = 1 and f(0)(1) = 0. Let U = ((,, this is open neighborhood of y.
/ ddfz
1 1 =(0)
(67)
1/x0 d + Ju df
(68)
$ (x((U 2 fa(U)
(69)
=1 S(U) + (a)(
(70)
1 2 (())
(71)
40

---

Published as a conference paper at ICLR 2025
Hence, fu( 0. Therefore suppfz.
Clearly, WOCang (X) _ WXCann(X), hence the sequential closure of WOCang(X) is contained inside WCCang We now show the reverse inclusion. Fix an arbitrary metric on X, this is possible as X is assumed metrizable.
Suppose have a weighted closed canonicalization f- Let h be the left invariant Haar measure on G. Let μsx = ).,,, this is a measure supported on Gr. Further it is o-finite measure as G is locally compact. Therefore we can the Lebesgue decomposition theorem to decompose f() as
fa = J +
1 + J
(72)
where f is absolutely continuous with respect to Mas f is singular with respect to pa and continuous, and fxd is singular with respect to Me and discrete.
First we construct a sequence of weighted closed canonicalizations f converging to f such that upon decomposing each f" as above, there is no singular continuous factor.
fac is singular continuous with respect to Mas hence there is a measurable set S C Ga such that He (S) = O and (S%) 0. Fix point S. For 1 > 0, let K; = Byy, 1) n S. Then we have
U Ki =
(73)
By continuity of measure
lim f (Ki) (2) (2) =
(74)
Fix large enough such that f (Ki) f K2 is compact points 21,- Tp(n) such that the balls around these points cover K₁. Call these open balls U,. Let U U, Ujze Uz- Then each U/ contains only x; and none of the other zij. Define discrete measure fa which is non-zero only on the I, and taking the following values those points
may find finitely many
("xxi)) f (U()
(75)
Then one casily sees that
(*(S)= (2 + f2 (U)) +... (U)())
(76)
Now let = fac + fm fzd. This is then measure with singular discrete part with respect to We now check that f + f. Let E (X). We may assume that & is non-negative on X by splitting it up Φ o+ + + for ot non-negative and o' non-positive.
M4g-
f (d(fn()) f(x))
fa(o)(=) f())(x)
/ (d((fa())) - fx()))
(1 @dffa(()) - + +1 sq..
@dfac(x)
Therefore we see that, with max @
p(n)
C (77)
[=(())()- max Σ
Σ ((d)
+))+
0%
n
Let > Take N large enough such that for all N, and 1....p
mmx
min U²
Φl <
(78)
Ug
Recall that we can do this as there are finitely many UK, each of which is contained in the 1 ball
around r,. Plugging this into Equation (77p we get
p(x)
C
C + cfs((S)
(zx())()) - f(0)(=))
+
Σ
(U))= + cfse(Ka) S n n
(79)
41

---

Published as a conference paper at ICLR 2025
Taking n large and small gives the convergence result.
Now we show that we can approximate a weighted closed canonicalization with no singular continuous part by weighted orbit canonicalizations. As before, decompose f. far + fxd with respect Pa- For any n, choose (x) to be the # open ball around Gx Ga. f is singular and discrete, hence it supported on (possibly infinite) set of points I1, 22..... in Ga. For each z; choose a sequence of points (zin) in Gr converging to 24. We may choose these such that for a fixed n. cach of the Zin are different for every Further we assume that these are chosen such that for every
1
d(z;,
(80)
n
Then define the singular measure fr which is non-zero only on the for every i, taking value
("((m)) = - (1
(81)
Define the absolutely continuous measure fr by
-(U))
(82)
Then let
(xxx) f
(83)
This is then measure supported Gx and by construction it is G-invariant. Hence this is weighted orbit canonicalization. Now we show that f + f. Let o E (X) and æ E X. Then
[z) f())(=)
=
① df -
dffT
(84)
Gz
Gx
1f df -
-f @d
dffac
odf
+
(85)
Go
Ga
I
Φdf
+ Σ fo (f2i))(((((s)) 6(zm)
(86)
V.((o)
I (Va((V)(z)) + Σ ((xi))))) - p(
(87)
Now fix > 0. As fod is a finite measure, there is some I such that
((x))) <e
(88)
Thus get
I-1
(z) - [0
Fed (fx:)))(())) 6(Fm))
(89)
Σ fal ((z;)))((()) - 6(xm)) +20
(90)
Now choose n. large enough such that (vv(Vv(x)) < 6 and for each i = 1, 1,.... -- 1, 16(x)) ((xx)) Then we get a bound
(f'()(()) f(0)(=)) s 3Cc+ De
(91)
Hence, f" f.
Theorem G.5. (CCang()cco sequentially closed.
Proof. Suppose that (n"),, is a convergent sequence of continuous weighted canonicalizations,
converging to weighted canonicalization, μ. We will show that y is continuous. To do this, fix
42

---

Published as a conference paper at ICLR 2025
convergent sequence (Em)m n in X converging to x. Let ø E (X)).. Then we have the following bound
(92)
Fix c > O. As pn converges to μ we may find large enough n such that both the first and third terms are less than For this as un is continuous, may find large enough such that the middle term less than e as well. Hence the right hand side is bounded by 3e, so we see that j is continuous. O
Proposition G.6. Let Y be a topological space. Let f Y + [0, +oo be lower semi-continuous. Then the set
M (u inf f}
(93)
closed.
Proof. Let be the infimum of f. We have that M - (faa() As is the infimum over Y, see that in fact M = ([), Then as ʃ is lower semi-continuous, this set is closed in Y. D
Lemma G.7. Let E be a lower semi-continuous energy function. Then KE is continuous if and only for every convergent sequence In + we have Li ME( = (M)...
Proof. Let E X + j0, too] be an energy function such that the corresponding energy minimizing canonicalization Kg is continuous. Consider convergent sequence (Ix) X. converging say to T. Mg(a)) is closed inside Gr, being the set of minimizers of lower semi-continuous function. As Ga is closed in X, .Mg closed in X as well. As X is metrizable, we may choose countable neighborhood basis of _Mg(x), which we shall write as a decreasing sequence (Um) of open neighborhoods such that
nmm = MR(rr)
(94)
Let C. U this is closed. As X is normal Hausdorff, we may find om X - [0,1] continuous such that " ({0}) Cn and o ((1)) = MM(x).
Now as K is continuous we have that the measures (g(x,) weakly converge to (x)) Hence for every have that
om djz + -Sence
lim n
Om dlz =1
(95)
Mi((an)
Hence
f On dis, =1
lim n n-000
(96)
ME(Zw)
Let e > O. We may find N such that for all n > N, the left hand side is greater than 1 1-6. 6. As we must have that the minimum of on over MÆ(1x) is greater than 1- E. As Cn = on (10)), we then must have that ME(a,) C U,. Hence by the definition of Kuratowski convergence we have Li (M S Mg(x).
Ma, (M((n))))))))))))))mmm)(())))()
Now suppose that (M) Then we may above choose decreasing sequence (Um) ) of open neighborhoods of which form neighborhood basis. Then fix for each m, @m X [0. continuous such that o ((0)) = X\U and @ ((1)) {v). By the same estimates as before we must have that for large enough n, each Me(#n) intersects Um non-trivially. Henee y Li (....).
So Li (M (z)
Conversely suppose that we are given that for every convergent sequence H, - xr, Li ME(x)) Mz(2). Let Cf We may assume that ø is non-negative. Fix > 0. As Mx(w)) is compact, we may find an open set containing Mg (x) such that
max ((y) min (()) <E BEU
(97)
WEU
As Li Mg(w)) S ME(x), may find N large enough such that for all 2 N. Mg(an) C U. Then
min Φ(v) - / MeMe(x)
odpe S max Φ(v)
(98)
MM(x)
43

---

Published as a conference paper at ICLR 2025
and similarly
min (x,)
( S ME(an)
① opzz
max (EM
((v)
(99)
Hence we see both of these integrals are bounded between the min and max of Φ on Hence
(ustco) a dif - -Meetar)
@ ddtenn
S 2e
(100)
Therefore this converges and we have weak convergence.
H OVERVIEW OF PREVIOUS CONCEPTS
Frames and canonicalization both represent different ways to impose group equivariance or invariance on arbitrary functions. The most basic approach is simply average over all group elements, sometimes refered to as the Reynolds operator. For finite group G and a continuous bounded function é it is given by (x) = क Egec 6(g 11.11 xx). This can also be modified impose G-equivariance by changing the integrand to Ö0g x)) Frames allow one to not have to integrate over the full group while still giving invariant or equivariant function as output. Specifically, frame is function F X + 2 such that for any E X and G: F(gr) gF(x). In order to define an action on continuous groups, one must assume that the frame finite, i.e. that for every x E X the set F(x) is finite. This is immediate if the group is finite, but if the group is non non-finite this is a very strong restriction, as it implies that the stabilizer groups of the action are finite, as F(x) GzF((s). This is detailed further in Appendix F
In order to deal with this restriction, Dym et al. (024), introduced the notion of weighted weakly equivariant frames. These generalize frames to be G-equivariant maps from X probability measures on G. Then for any continuous bounded function & we may apply a weighted frame by p()) (z) Soz (0g-g11-00 11)x) (g) This is still not enough in the setting of non-compact groups as discussed the subsequent section Even worse, there an inherent limitation to using frames: the input to @ may end up outside the training distribution, resulting in worsened performance. The overall network would therefore need be retrained, which may be exponentially hard (Kiani et al.. 2024a).
To deal with this, Kaba et al. (023) proposed canonicalization as simply a G-invariant function f X X such that for each E X, fz) Gz, which upon pre-composing with function would make invariant. In order to preserve continuity of functions, one needs to impose that f itself is continuous. In Dym al. (2244) it is shown that there many topological obstructions to constructing such continuous canonicalization functions.
To get around these obstructions Ma al. (2024) proposes canonicalization now be set-valued G-invariant function C X - 2x for C(x) finite, acting on continuous functions o by (())(=) c(t) Evec(E) d(v). Contractive canonicalizations, which we refer instead or- bit canonicalizations. are a specific kind of canonicalization where for every x E X, C(x) C Gr. The orbit-stabiliag theorem applies show these equivalent to frames. Namely, there is map a which sends equivariant frames to orbit canonicalizations For an equivariant frame F the corresponding orbit canonicalization (FF) is defined as (x) -g-1 .zgg E F(x)). Ma et al. (2024 Theorem 3.1) show that for G finite, is an isomorphism between the set of equivariant frames and orbit canonicalizations. Energy- based canonicalization was first introduced in Kaba et al. (2023). In energy-based canonicalization, given energy function F, the canonicalizing element c defined as c(x) arg mingeg (g- As shown the subsequent section, considering the full set of minimizers, instead of a singleton, naturally gives rise to frame, and naturally results in set- valuvaleed canonicalizations. Extending these to non-compact groups naturally gives rise to weighted closed canonicalizations.
WWekly equivariant weighted frames also require compact stabiliser subgroups, while the equivalent cquivariance of can he defined for non-compact groups
44

---

via Published as a conference paper at ICLR 2025
H.1 WEIGHTED CANONICALIZATION
For this section let G be a Lie group acting smoothly on manifold X. Let C% be the bounded continuous R-- valuee functions on X. Let PMeas( X) the set Radon probability measures on X, with respect to the Borel o-algebra on X. Let B(X, IR) be the space of bounded functions on X and (X, R)G the subspace of G-invariant such functions.
Definition H.1. weighted canonicalization is a G-invariant function K[]] : X + PMeas(XX.. We call the space of weighted canonicalizations WCang(X).
Given a weighted canonicalization & we may define an operator P.
(Px(o)) - =/ ó driz
(101)
We may write the function (a( simply as k(Ø). As K2 PMcas(XX for each z E X, we have [o))( I 1l x(() is G-invariant as R is. We also define the function P₂(x) Px(
The entire space WCang(X) is much too general to be practical. We would like to consider much more natural subset from the perspective of canonicalization, namely orbit canonicalizations. We now define a corresponding weighted notion.
Definition H.2. A weighted orbit canonicalization is weighted canonicalization x X - PMeas(X) such that for every x E X. supp k(x) S Ga. We denote the set of weighted orbit canonicalizations as WOCana(X).
In particular, for weighted orbit canonicalizations both operators become projections, see Ap- pendix As in the finite case, there is a function a WWFrag - WOCan:(X), which surjective and injective up to weak-equivallnnee see Appendix F.
Given the action on continuous functions, we may equip WCang(X) with a certain weak topology. namely the coarsest topology such that all the Po for ó continuous and bounded become continuous functions. Specifically this means that a sequence of weighted canonicalizations K,, converges to к if and only if for every & (X) the sequence Ki, (+) (Ø) converges to k(() pointwise. For the canonicalization problem this is a very natural topology to consider, as we only care about canonicalization through their respective actions on continuous functions.
So far, we have presented a generalization of canonicalization, which is readily applicable to finite or compact Lie groups. However, non compact Lie groups may have non-closed orbits, and the space of weighted orbit canonicalizations is not closed under taking limits in the weak topology. To understand the closure under taking limits, we introduce the concept of weighted closed canonicalization.
Definition H.3. weighted closed canonicalization PMeas X) such that for every E X, ((ppp((() S Ga. Call the set of weighted closed canonical- izations (XC.ng(X).
weighted canonicalization X ->>
The following theorem shows that these allow to fully characterize the behaviour of weighted orbit canonicalizations under taking limits.
Theorem H.1. (Theorem G.4
The sequential closure of WOCann is WCCana
Using this theorem and the setup above we may now write down our main diagram connecting all the different notions of canonicalization and framing together, shown in Figure 3 In analogy with Dym et al. (2024), may now ask whether such canonicalizations take continuous functions to continuous functions. the presented framework this is simple condition to add.
Definition H.4. A weighted canonicalization K : X + PMeas(X) is continuous if for ev- ery Φ (X), k(δ) is continuous. Call the space of continuous weighted canonicalizations WCang (X). We similarly define continuous weighted orbit canonicalizations and continuous weighted canonicalizations.
Proposition H.2. A weighted canonicalization k is continuous if and only for every convergent sequence Fn + in X, the sequence of measures Ka., converges weakly to Ea-
Theorem H.3. (Theorem
G.5 WCan (X) ) is sequentially closed inside WXane(X).
45

---

Published as a conference paper at ICLR 2025
H.2 ENERGY MINIMIZING CANONICALIZATION
In Section so far, we have provided a generalized theoretical understanding of frames and canonical- ization, encompassing various scenarios with compact/non-compact and discrete/continuous groups. However, their relevance and how these can be used in practice remains unaddressed. We now move on to the concept of energy frames and canonicalizations, for general Lie groups. We define an energy function on X to be a non-constant function E : X + [0, too]. We further assume that E is topologically coercive: for any N E R there is a compact K S X such that F Flxik > N.
In the case where G' is finite, every orbit in X necessarily has minimum of E, and we note that energy minimization naturally results in frame
Fg (z)
arg min Elg E(g-'z)
(102)
gEG
Proposition H.4. FE defines an equivariant frame, i.e. FE(ge) = gFF(())))))))))))()(()) for any g E G, x E X.
However when G is not fimite this set not guaranteed to be finite. G is additionally non-compact,
then E' does not necessarily achieve its minimum on every orbit in X, making this optimization problem not well-defined. For this reason, we instead turn to weighted canonicalizations. By taking the closures of orbits, we solve the problem of E not having minima, so we are naturally led to constructing weighted closed canonicalizations. First, we define the energy-minimizing set, MÆ(r)= = {x' Gx E(x')= inf yEG= (v)), set of minima of E on the closure of the orbit of æ. Proposition H.5. (Proposition G.2 Mg(x) is G-invariant and non-empty.
If (MM(r) is finite for each x E X, then we define the (non-weighted) closed canonicalization Kg (CCang(X) by taking (z) to be the normalized counting measure on Mg(x)) as usual.
For a general E, we may not be able to put any reasonable probability measure on the set MM(z)), for example if Mg(x) has highly fractal geometry. We therefore must assume that E is such that for all I, the Hausdorff measure of MM is non-zero. As each MMg (z) is compact, the Hausdorff measure is finite, hence combined with this assumption we may normalize it obtain a probability measure on Mg(x) which a weighted canonicalization (g(r)) E WCCang (X). Refer to Appendix E for more discussion on the Hausdorff measure and on why this assumption is reasonable. Note that if Mg(r) is finite, then the normalized Hausdorff measure constructed agrees with the normalized counting measure. Hence this construction is a generalization of the purely finite case.
For lower cemi-continoous energy functions, we have the following characterization of which energies give rise continuous weighted canonicalizations.
Lemma H.6. (Lemma G.7 Let E be lower semi-continuous energy function. Then Kg is continuous if and only for every convergent sequence I,, + x we have Li MM(2n)) = =M)
We have the following theorem telling us which weighted closed canonicalizations arise from energy minimizations.
Theorem H.7. (Theorem F.3 Let K E WCCang be such that Ka is the normalized Hausdorff measure on its support and for all y E Gy, supp Ky = supp Ka. n Gr. Then there is an energy function E such that kg K.
H.3 ALMOST EVERYWHERR CONTINUITY IS NOT IMPOSSIBLE
Canonicalization Let X be space with a group G acting on it by homeomorphisms. Then a continuous canonicalization function is continuous G-equivariant map v : X + X which sends any element x to an element of the orbit Gr. In Dym et al. (2024) it is shown that in general these do not exist, as it much too strong of a condition. In fact, because a continuous canonicalization function will be right-inverse to the quotient map X + X/G, necessary condition to the existence of e is that the map क (X) + л (X/ (X/G) /G) induced by the quotient must be a surjection. If the quotient X X/G is covering map then this only holds if G is trivial. For concrete example of this, we can consider quotienting IR? by the translation action of Z Then Ti (R?) = 0 while T, (R? /2²) Zz, showing that continuous canonicalization function can't exist.
On the other hand, global existence of a continuous canonicalization function is unnecessary for machine learning applications. As long as we can construct canonicalization function which is
46

---

Published as a conference paper at ICLR 2025
continuous away from set of measure zero, this will suffice for applications. Specifically introduce the notion an almost everywhere (a.e.) continuous camonicalization function.
Definition H.5. An a.e. continuous canonicalization function is function v X - X such that there exists set S C X of measure-zero such that
(i G(S) C S
(ii) is continuous away from S
(iii) Vz 3, G st. v(x) gx
(iv) G-invariant
As with the regular canonicalization functions, and invariant functions. Given an arbitrary function f X + IR it is casy to see that
such a function to construct equivariant
(x) = f(v(x))
(103)
G-invariant. Further it is easy check that this sends a.e. continuous functions to a.c. continuous functions.
The added freedom of the S where can take arbitrary values allows avoid the topological restrictions that prevent the existence of globally continuous canonicalization functions. For simple example consider again the presentation of torus as the quotient of I2² by the translation action. Then an easy example of an a.e. continuous canonicalization function is the function
v(z, y) [[z) [B])
where [a] is the decimal part of a real number E R. This is discontinuous whenever either x or integers, but this set has measure zero
For another example, consider the action of the orthogonal group ((d) on the space of matrices pdxx, The action preserves the rank of a matrix A Rdxn, so the subset of full rank matrices is preserved by Od). This is an open dense set, and if we take to just be the Lebesgue measure, its complement has measure zero. So take S to be the of matrices with non-full rank. Let A be matrix of full rank. Then one of the d x d minors of A is invertible, we may assume that it the first one. Thus we can write = (Ao-A1) with Ao an invertible Gram-Schmidt procedure to Ap. we obtain An QR with Q orthogonal and R upper-triangular. Now we take v( v(A) Q" A. Q given by a polynomial in the elements of A, hence smooth. Further, if O (d), then we can decompose OAA OQR with OQ orthogonal. Hence,
d matrix. Applying the
(OAA) (OQ) of OA A v(A)
Therefore is continuous canonicalization function, with S the set of non-full rank matrices in Iddxn
Invariant Theory Finding and computing G-invariant functions a rich topic which comes up in many different situations. In algebraic geometry, wants find generators for G-invariant polynomials, i.e. a set of G-invariant polynomials oi...., on such that any other G-invariant polynomial f can be written as
f= T=a(o)....,
for some polynomial g. For linearly reductive groups acting on C", there is an algorithm for computing these generators, which is described in Derksen & Kemper (015).
In the smooth case, there is a classic result of Schwarz (Schwarz) 1975). telling us that if is a compact Lie group acting on smooth manifold X, with generators for the space of G-invariant polynomials being 0i, .I...., On, then any G-invariant smooth functions f can be written as
f ==g.o1......
where g an some smooth function on R.a
47

---

Published as a conference paper at ICLR 2025
Algorithm 2: Canonicalization function via Lie algebra descent
Data: non-canonical imput æ
parameters Nouter. Nimar, step-sizes 7j, Energy F. +
Result: Canonicalized input å = g 1 x; inverse of canonicalizing group element g g canonicalizing group element g.
9 - id:
steps (}:
for k: = O... -....N.ur Nouter do
# Do gradient descent on g
for i O... N...Ninnar do
& & V. expl -E)
end
steps.append(""
t g exp ()
#Recalibrate to be in the right tangent plane
end
id #Reconstruct the element g
for k
N Nouter... do
exp(steps(k))p
end
return g 1. x;; g g
I ALGORRITMMS FOR SOLVABLE AND NON-SOLVABLE LIE ALGEBRAS
As discussed in the main body, in certain cases it is possible to significantly simplify the algorithmic approach for minimizing Equation 2 We will distinguish specific cases:
If the group is compact or the algebra solvable either of these cases, we are able to parameterise the group globally using the Liealgebra. Denoting the basis E and the coefficients as Og IR, we have that we can represent any g E G
= exp Σ
is compact, or as
9 = II exp ayy;
if the Lie algebra is solvable with ordered to represent the nested subalgebras. In both of these cases we can reduce the optimization to be done over the vector coefficients a IRdi g, Both of these can be viewed specific types of global retraction mappings 7 - G. For this case Algorithm
the group is non-compact or the algebra is not solvable Such setting is particularly important for PDE-based groups, these naturally result in groups that non-compact, see section 4.3 In this case is no longer possible to find global retraction, and a natural question arises can the infinitesimal actions be enough? It turns out that the answer is yes, quite generically, resulting in Algorithm
Proposition I.1 (Proposition 1.24 from Olver 1993)). Let G be connected Lie group and U C G 9k : 9i U} be the set of k-foold
neighbourhood of the identity. Also, let UK {g1 92 products of elements U. Then
G
u*
In other words, every group element G can be written finite product of elements of U.
48

---

Published as a conference paper at ICLR 2025
Algorithm 3: Canonicalization function via coordinate Lie algebra descent
Data: Non-canonical imput æ:
parameters :N step-sizes V/, Energy E: X' + R, тı if using proximal
Result: Canonicalized input å = g g x;; inverse of canonicalizing group element g g canonicalizing group element g.
9 + id;
steps ():
for k: = O...N do
# Do coordinate descent on geg, by picking index j
j j+ pickkk( dimg
A 4 arg min, Elg e exp(-Xu;) . x):
# May be desirable to use proximal descent
# Xạ arg E(g exp((Aoy) x) + 2m |X
{ Ap
steps.append(():
t g exp(-E) #Recalibrate to be in the right tangent plane end
f id Reconstruct the element g for k 0 f exp(steps(k))g
end
return g .::
If the group only provides information about basis vector action most examples of PDE symmetries, integrating the action for generic E g impossible, and instead only the action with respect to certain basis {v:} known, In this case coordinate descent schemes must be used. For this use Algorithm 3
49


Reviews:
[{"review_id": "ESIWmFLGaZ", "reviewer": "Reviewer_n4j8", "summary": "The paper proposes a Lie algebra canonicalization mechanism for achieving equivariance with respect to a variety of Lie groups. The authors propose an extension of the energy based canonicalization mechanism, analyzing limitations of existing work within this framework and describing how the methodology can be extended to work with more exotic (i.e. non-compact, non-abelian) Lie groups, enabling application to a large class of learning tasks where equivariance might be desired. The paper connects frames, canonicalization and frame averaging and the authors show that non-compact groups can be approached via the energy minimization framework where an optimization process makes use of the Lie algebra generators. The methodology is evaluated both on standard invariant image classification tasks as well as physics-informed learning problems dealing with Lie point symmetries.", "questions": "Besides the questions in the weaknesses section:\n- 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?\n- 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.\n- 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The paper is mathematically well-grounded. While I think the presentation could be reorganized (see below) I appreciate the focus that the authors have on providing derivations and proofs for their claims. The objective of unifying (weighted) canonicalization mechanisms and frame theory could prove useful for practitioners in this area.\n- I am happy to see extensions to the canonicalization framework, especially ones which focus on making use of the underlying geometry and parametrization of the groups and spaces involved, as well as the treatment of non-compact groups which are under-explored.\n- I find the potential use-case of applying this methodology to pre-trained models still under-explored and worthy of further investigation.", "weaknesses": "I think the current presentation lacks focus and the paper could be vastly improved with the objectives of:\n- Highlighting clearly the limitations of past/current proposals and how these limitations are overcome, potentially alternating more concrete examples with more abstract limitations.\n- Presenting a formalized methodology for the entire (extended) framework that could be understood by practitioners with choices and pitfalls for the spaces, groups, energy functions, etc. involved. \n\nIn regards to improving clarity:\n- $WFra_{G}(X)$ should be defined before it is used on line 236.\n- On lines 270-271 is $N$ simply some value $\\in \\mathbb{R}$?\n- It should be made clearer that $C_{E}(x)$ (line 288-289) refers to the normalized counting measure on $M_{E}(x)$, since we are stating that $C_{E}$ is a probability measure, and then on line 296 we write $C_{E} = \\mu_{x}$.\n- It would be useful to make clearer for each proposed construction the limitation that exists with the current methodology, e.g. it seems from Definition 2.2 and the subsequent paragraphs one should understand that restricting the support to the orbits and equiping $WCAN_{G}$ with the coarsest topology was not proposed before? Similarly, it would be useful to make clearer what topological limitations appear for non-compact groups, and for which in particular (e.g. we are not just talking about translation). \n- In the same spirit clarifying which non-compact Lie groups acting on which spaces (transitively or not) have non-closed orbits would highlight more clearly the settings where the framework should be considered.\n- And similarly to the previous comment the cases where the energy $E$ induces a 'reasonable' probability measure on $M_{E}(x)$ could be contrasted with specific/concrete choices of groups and group actions.\n- Considering the main proposal of the paper is a general framework I find the discussion in Section 3 much too general and unstructured. It is again highlighted what limitations could exist when one chooses an energy function, however I think it would be much more useful to present a summary of the entire canonicalization framework (and the limitations that were addressed) with a clear outline of potential choices for input spaces and groups (and their actions), potentially in increasing generality (e.g. finite -> compact Lie group -> non-compact Lie group). Once the methodology is clear one highlights criteria for choosing the energy function. Some form of presentation similar to Algorithm 1 in the appendix could potentially appear in the main manuscript given that the energy function is a key component of the generalized framework. \n- \"What makes moving away from the compactness assumption even worse, is the move away from the matrix view, as any compact Lie group is a matrix group (Knapp, 1996, Chapter 4)\" - I don't quite understand what is being claimed here and in the next few sentences. A matrix Lie group is a Lie subgroup of $\\textnormal{GL}(n, \\mathbb{R})$ (or $\\mathbb{C}$). Non-compact matrix Lie groups can also be decomposed/expressed as a product of exponentials, and one can optimize using their Lie algebra elements, see e.g. [3] and [4].\n- I'm left not understanding what the authors mean when they state that their framework requires  less knowledge about the symmetry group. This is highlighted both in the presentation of the contributions as well as in the experiments, e.g. for invariant classification a comparison is done with [1], I assume as opposed to [1] and [2]? I don't have an issue with what results are cited, but it is not made clear what knowledge about the Lie group is needed for [2] that isn't needed for [1]? \n\nI think the authors should focus on improving/rewriting parts of Sections 2 and 3 with the focus of highlighting both the limitations their framework overcomes and providing a clearer presentation of how the methodology can be applied in different cases. \n\nSome of the expression/grammar in the appendix could also be improved, e.g. \"Liealgebra descent\", \"why failure modes exists\".\n\n[1] Enabling equivariance for arbitrary lie groups, MacDonald et. al 2022\n[2] Lie group decompositions for equivariant neural networks, Mironenco & Forre 2024\n[3] Trivializations for Gradient-Based Optimization on Manifolds, Mario Lezcano-Casado, 2019\n[4] Optimization algorithms on matrix manifolds, Absil et. al 2009", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aImgvVxg1X", "reviewer": "Reviewer_VdBH", "summary": "The contribution of this paper can be summarized as follows:\n- An extension of frames and canonicalization for neural network symmetrization [1-4] to non-compact Lie groups specified by their infinitesimal generators (Section 2.2 and 3),\n- A class of optimization-based algorithms for the above energy-based canonicalization using (coordinate) Lie algebra descent (Appendix H),\n- Applications of the proposed method to affine and homography group invariant MNIST classification (Section 4.2) and, importantly, neural operator modeling of three PDEs with known point symmetry groups by canonicalizing a pre-trained model (Section 4.3).\n\n[1] Puny et al. Frame averaging for invariant and equivariant network design (2021)\n\n[2] Kaba et al. Equivariance with learned canonicalization functions (2023)\n\n[3] Dym et al. Equivariant frames and the impossibility of continuous canonicalization (2024)\n\n[4] Ma et al. A canonicalization perspective on invariant and equivariant learning (2024)", "questions": "Please see the weaknesses section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- S1. The contributions of this paper can be understood from two perspectives. The first is extending frame/canonicalization approaches for neural network symmetrization to non-compact Lie groups specified by their infinitesimal generators. The second is improving the performance of pre-trained neural operators for PDEs by canonicalizing them in accordance to the point symmetry of a downstream PDE. Both are original and significant contributions as far as I am aware.\n- S2. The extension of frames and canonicalization considered in prior work [3, 4] to non-compact Lie groups, and the proposed Lie algebra descent algorithms that implement the idea, are original and technically sound as far as I am aware.\n- S3. Experimental results are shown for a comprehensive set of problems (synthetic, two computer vision problems, and three PDEs) with informative visualizations and support the validity of the approach.\n\n[3] Dym et al. Equivariant frames and the impossibility of continuous canonicalization (2024)\n\n[4] Ma et al. A canonicalization perspective on invariant and equivariant learning (2024)", "weaknesses": "- W1. In Section 2.2, the authors propose to treat non-weighted energy-minimizing closed canonicalization as weighted canonicalization by taking the normalized Hausdorff measure on energy minimizing set (Line 287-296). The resulting class of weighted closed canonicalization (Theorem 2.7) has a weakness that, with the energy function specified, it is not possible to adjust the weights of canonicalization from training data, unlike in related prior work [5-7]. This may have led to the reliance on carefully designed energy functions based on domain knowledge (Line 324-332 and Section 4.3) which leaves a room for improvement.\n- W2. For the ACE experiment (Section 4.3.3), the current comparison is made only between Poseidon and its canonicalization. A comparison to existing intrinsically symmetric methods [8, 9] would be informative and show the usefulness of canonicalization, since Poseidon can benefit from pre-training while intrinsically symmetric approaches cannot.\n\nMinor comments and suggestions\n\n- In Line 249, the notation for closure $\\bar{X}$ of a set $X$ is used without definition.\n- In Line 260, $\\mathrm{PMeas}$ -> $\\mathrm{PMeas}(X)$\n- In Line 976, euclidean -> Euclidean\n- In Line 1758, the union of -> a union of?\n- 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.\n\n[5] Mondal et al. Equivariant adaptation of large pretrained models (2023)\n\n[6] Kim et al. Learning probabilistic symmetrization for architecture agnostic equivariance (2023)\n\n[7] Zhang et al. SymDiff: Equivariant diffusion via stochastic symmetrisation (2024)\n\n[8] Arora et al. Invariant physics-informed neural networks for ordinary differential equations (2024)\n\n[9] Lagrave & Tron, Equivariant neural networks and differential invariants theory for solving partial differential equations (2022)\n\n[10] Nguyen et al. Learning symmetrization for equivariance with orbit distance minimization (2023)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZZEoWawp8S", "reviewer": "Reviewer_aU4A", "summary": "This paper extends the energy-based canonicalization approach, introduced in [1] to settings where the group in non-discrete and non-compact, exploiting only the infinitesimal generators of Lie algebras. They provide a general framework of constructing energy functionals which can be optimized using standard Lie group descent schemes.'\n \n[1] Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh.Equivariance with Learned Canonicalization Functions. In Proceedings of the 40th International Conference on Machine Learning, pp. 15546–15566. PMLR, July 2023.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The topic of constructing equivariant networks when one only has access to the infinitesimal generators is an extremely interesting and well-motivated direction of research. The approach this paper takes, which is adapting the canonicalization framework for these settings is novel and using the energy-based canonicalization seems like a promising direction for these cases.  The figures in the paper are also very useful and help provide an intuition for how canonicalization is helping.", "weaknesses": "- One of the main limitations of the paper is that 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. I found it difficult to follow how the energy functionals were constructed  in each of the cases and how much that approach is generalizable to more complex systems and PDEs. \n- The general experimental setup was not well-explained in the paper. For example, it would be better to include a main algorithm, provide a more detailed explanation of how the energy functional was constructed in each of the experiments settings, and clearly explain the experimental setup and training of each of the datasets. \n- 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. \n- 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.  \n\nOverall, I think the topic and direction explored is very interesting and the experiments show promising results but the paper was extremely  difficult to follow (especially understanding how the theoretical discussions are mapped to practical algorithms) which limits its impact.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NUiMLFiOzs", "reviewer": "Reviewer_CTGT", "summary": "The paper proposes LieLAC, a method to make pre-trained models equivariant to Lie group symmetries through input canonicalization via energy minimization. The key strength of the work is its well-grounded theoretical foundation and potential applicability to a wide range of symmetry groups and tasks. However, the practical challenges of finding an appropriate energy function, finding its minimizer, the requirement of finetuning canonicalized network and limited experimental evaluation raise concerns about broader applicability beyond simple demonstrated cases.", "questions": "I suggest authors address and elaborate on weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed canonicalization approach allows employing pre-trained non-equivariant models while making them equivariant with canonicalization.\n2. The paper provides a good overview of the related work on canonicalization and frame-averaging methods.\n2. Figure 1 provides an intuitive example of the effect of canonicalization on the decision boundary.\n3. The proposed approach is generic and is potentially applicable to a wide range of models and learning tasks.\n4. The paper provides strong theoretical support for the proposed method.\n5. Canonicalization and frame averaging methods have been demonstrated to be successful in various domains, including geometric and imaging modalities. The paper further extends the canonicalization framework to PDE/ODE modeling with more complex symmetry groups.\n6. The paper nicely elaborates on the challenges of choosing the appropriate energy function and it outlines key considerations in doing so.", "weaknesses": "1. The paper states \"Prior work on equivariant neural networks focuses on simple groups ... that are not reach 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?\n2. 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.\n3. The distance-based relaxation in 331-333 is not clear. Can authors elaborate more on what is meant here?\n4. 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.\n5. 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.\n6. 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.\n\nSome important missed related work on equivariance:\n- M. Zhdanov et al. Clifford-steerable convolutional neural networks NeurIPS 2023.\n- D. Ruhe et al. Clifford group equivariant neural networks NeurIPS 2023.\n\nSome important missed related work on Lie groups in NNs:\n- A. Moskalev et al. Liegg: Studying learned lie group generators. NeurIPS 2022.\n- N. Gruver et al. The Lie Derivative for Measuring Learned Equivariance. ICLR 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)