Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groups

## Abstract

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.

## 1 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., ), physics (Merchant et al., ) and engineering (Degrave et al., ). 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. *Physics-informed neural networks* (PINNs) seek to leverage such structure to mitigate the aforementioned challenges.

In the recent deep learning literature it has been recognized that encoding symmetries as inductive bias into model architectures can give rise to computational benefits.
Notable examples include translation invariance properties in imaging data or sequences in time-series, as well as inherent 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., ; Wang et al., ) and through the lens of learning-theoretic trade-offs (Mei et al., ; Bietti et al., ; Kiani et al., ).
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. We introduce a novel framework that encodes a wider range of symmetries into model architectures. Moreover, we do not require a 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.

### 1.1 Equivariance in Neural Networks

Incorporating group symmetries inside neural networks, in the form of *equivariance*, is one of the fundamental tasks in *geometric deep learning* (Bronstein et al., ). The translational equivariance of convolutional neural networks (CNNs) (Fukushima & Miyake, ; LeCun et al., ) has been suggested as an important reason for their successes in tasks such as image classification (Krizhevsky et al., ; He et al., ). 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 on roto-translationally equivariant CNNs (Dieleman et al., , ; Cohen & Welling, ), to more recent methods, efficiently utilising knowledge of the Lie algebra for more general (Kumar et al., ; MacDonald et al., ; Finzi et al., ; McNeela, ) and computationally efficient methods (Mironenco & Forré, ).

Figure: Figure 1: Effect of canonicalization on decision boundaries in $k$-NN classification for separating the inner and the outer rings Section 5.1.
Refer to caption: x1.png

##### Equivariant convolutions

An important fact that comes up repeatedly in the design of equivariant neural networks is that ordinary convolutions on Euclidean spaces can be generalized to any group on which we can integrate against 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 can 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 () for discrete groups of roto-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 to convert a signal on the original domain to a signal on the group. Incompatibility of continuous groups and grid-based discretizations of signals is a major issue in methods that take this approach, and for general groups, it is not necessarily clear how one should discretize the integral defining a 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 (); Finzi et al. () and works following them, which focus more heavily on using the Lie algebra for approximating group convolutions (MacDonald et al., ; Kumar et al., ).
The corresponding methods have been shown to result in more efficient use of training data and more robustness to transformations not seen during training, and have been applied successfully in areas such as medical image analysis (Bekkers et al., ; Winkels & Cohen, ), computational molecular sciences (Batzner et al., ; Batatia et al., ) and inverse problems in imaging (Celledoni et al., ; Chen et al., ).

More recently, there has been a drive towards publicly available foundational models (Bommasani et al., ) for various tasks, pretrained on massive datasets. These models, however, are often not equivariant. For example, Alphafold 3 (Abramson et al., ) no longer utilises equivariant architectures, unlike the initial version (Jumper et al., ). For this reason a direction of research has been investigating the possibility of utilizing the generalization offered by equivariance, while being able to use pre-trained models. The main two approaches to achieve this are canonicalization(^1^11The idea of canonicalization in many ways is not new, see Section C.1.) Kaba et al. (); Mondal et al. () and frame averaging Puny et al. (); Basu et al. () (also referred to as symmetrization). However, there still remain a number of open theoretical questions regarding the possibility of continuous canonicalization (Dym et al., ), connections between frame averaging and canonicalizations (Ma et al., ) and constructability of frames (Dym et al., ; Puny et al., ).

### 1.2 Physics-Informed Machine Learning

Figure: (a) DeepONet operator for the Heat equation Section 5.3.1.
Refer to caption: x23.png

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 as those involving non-linear and non-smooth PDEs necessitating fine discretization, or high-dimensional problems where the curse of dimensionality kicks in (Han et al., ). This is further exacerbated if there is a need to solve PDEs across varying domain geometries, parameters, and initial/boundary conditions, requiring independent simulations for each unique configuration.

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., ). These neural networks have been dubbed PINNs, and have been shown to be useful for solving forward and inverse problems for PDEs (Raissi et al., ; Sun et al., ; Zhu et al., ; Karumuri et al., ; Sirignano & Spiliopoulos, ), and regularizing learning (Zhang et al., ).

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., ; 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 or preconditioning (Müller & Zeinhofer, ).

To address the challenge of solving PDEs across diverse scenarios, neural operator models have emerged as a promising solution. Unlike standard DNNs, which approximate functions over finite-dimensional vector spaces, neural operators learn mappings between infinite-dimensional 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-informed neural operators (Li et al., ; Wang et al., ), oftentimes improving generalization and physical validity, compared to their physics-uninformed variants. In what follows, we concern ourselves precisely with this class of approaches.

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. What is appealing about LPS, is that for a given PDE, under certain regularity conditions, all such (simply connected) symmetries can be derived in the form of one-parameter transformation groups (Olver, ). For many classic PDEs these have been derived and can e.g. be found in (Ibragimov, ). There further exists symbolic software packages, e.g. Baumann (), to derive symmetries for general PDEs. The simplest ways to add symmetry information into neural PDE solvers is through data augmentation (Brandstetter et al., ; Li et al., ), resulting in more expensive training, or loss augmentation (Akhound-Sadegh et al., ; Li et al., ; Zhang et al., ), 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., ). Alternatively, exact equivariance can be seeked. As an example, in Wang et al. (), neural PDE solvers were combined with standard approaches to 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-shelf 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. (). Here, the approach taken was to invariantize the ODE in question, and apply a PINN to this equation, before mapping back to an (approximate) solution of the original ODE. Similarly, Lagrave & Tron () uses equivariant neural networks to 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 a practical point of view can even be approximated from observed data (Gabel et al., ). However it still remains an open question, whether it is possible to make existing physics-informed approaches equivariant. In this paper we directly tackle this question.

### 1.3 Overview of Contributions

Motivated by limitations of existing equivariant architectures, we turn to *energy-based canonicalization* as a means for inducing equivariance in existing models. This approach was first introduced in Kaba et al. (), 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 definitions and results in the prior literature do not extend readily to the general setting, we extend and refine existing definitions to handle non-discrete and non-compact groups. Our key contribution is an energy-based canonicalization approach that provides a unifying theoretical framework for constructing frames (Section 3).

We emphasize that our proposed approach is generic, in that it has the potential to apply to a wide range of models and learning tasks.
We offer guidelines for constructing energy functionals and for performing Lie group optimization in Section 4. We showcase the effectiveness of our method on image classification tasks under homography and affine group invariance in Section 5.2. We further discuss turning neural operators Lie point symmetry equivariant on three examples, including the pre-trained Poseidon foundation model (Herde et al., ) in Section 5.3.

## 2 Preliminaries

##### Lie groups

In this section we briefly recall some necessary notions from Lie group theory. For a more comprehensive introduction, see Kirillov (). A Lie group $G$ is a group with a smooth manifold structure such that both the inverse and multiplication maps are smooth. A Lie algebra is a finite dimensional vector space $V$ with an alternating bilinear map $\left[\cdot,\cdot\right]:V\times V\rightarrow V$ called the Lie bracket, satisfying the Jacobi identity. Given a Lie group $G$, the tangent space at the identity $\mathfrak{g}=T_{e}G$ naturally has the structure of a Lie algebra. A 1-parameter subgroup of $G$ is simply a group homomorphism $\mathbb{R}\rightarrow G$, where $\mathbb{R}$ is a group under addition. For every Lie algebra vector $v$ there is a unique 1-parameter subgroup $\phi_{v}$ such that $\phi_{v}(0)$ is the identity and $\phi_{v}^{\prime}(0)(\partial_{t})=v$. There is a smooth map $\operatorname{exp}:\mathfrak{g}\rightarrow G$ which sends $v$ to $\phi_{v}(1)$. This has image contained inside $G_{0}$, the connected component of the identity $e\in G$. If $G_{0}$ is compact then $\operatorname{exp}$ surjects onto $G_{0}$, but in the non-compact case it doesn’t necessarily surject.

For a Lie group $G$ acting smoothly on a manifold $X$ we have a version of the orbit-stabilizer theorem.

###### Theorem 2.1 .

The stabilizer $G_{x}\subseteq G$ is a closed Lie subgroup, and the natural map $G/G_{x}\rightarrow X$ is an injective immersion with image $Gx$.

###### Theorem 2.1 .

##### PDE Symmetries

When considering PDE symmetries, we primarily work with open subsets of Euclidean space, but can be naturally extended. The formal presentation here closely follows Olver (), and an interested reader is referred there for more in-depth discussions of the concepts or Section A.1 for extra details. This subsection introduces the concept of PDE point symmetries – transformations that map solutions of a given PDE to other solutions of the same PDE.

Suppose we are considering a system $\Delta(x,u^{(n)})=0,$ of $n$-th order differential equations with independent variables $x\in X=\mathbb{R}^{p}$, and dependent variables $u\in U=\mathbb{R}^{q}$, with $x,u^{(n)}$ representing a vector in the jet space defined as $X\times U^{(n)}=X\times U\times U_{1}\times\cdots\times U_{n}$, where $U_{n}$ is endowed with coordinates $\partial_{J}u^{\alpha}(x)$ representing all the derivatives up to order $n$.
A symmetry group of the system $\Delta$ will be a local group of transformations, $G^{\Delta}$, acting on some open subset $M\subset X\times U$ in such that $G^{\Delta}$ transforms solutions of $\Delta$ to other solutions of $\Delta$. There is an induced local action of $G^{\Delta}$ on the $n$-jet space $M^{(n)}$, called the $n$-th prolongation of $G$ denoted $\mathrm{pr}^{(n)}G^{\Delta}$. The prolongation is defined to transform the derivatives of functions $u=f(x)$ into the corresponding derivatives of the transformed function $\tilde{u}=\tilde{f}(\tilde{x})$. Considering (local) one-parameter groups generated by vector fields on $M$, and their prolongations allows one to classify symmetries of a given PDE by considering vector field infinitesimal actions, see Olver (, Theorem 2.31).

##### 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 in the form of

$$ $\displaystyle\Delta={u}_{t}+\mathcal{D}_{{x}}[{u}]=0,\text{ and }$ $\displaystyle{u}(0,{x})=f({x}),\;$ $\displaystyle{x}\in\Omega,$ $\displaystyle t\in[0,T]$ $\displaystyle{u}(t,{x})=g(t,{x}),$ $\displaystyle{x}\in\partial\Omega,$ $\displaystyle t\in[0,T]$ $$

where $u(t,x)\in\mathbb{R}^{\operatorname{dim}u}$ is the solution to the PDE, $t$ denotes the time, and $x$ is a vector of spatial coordinates. $\Omega\subset\mathbb{R}^{\operatorname{dim}x}$ is the domain and $\mathcal{D}_{x}$ is a (potentially non-linear) differential operator. $f(x)$ is known as the initial condition function and $g(t,x)$ describes the boundary conditions. As discussed in Section 1.2, we are interested in physics-informed approaches, wherein $u(t,x)$ is parameterised as a neural network $u_{\theta}(t,x)$ with parameters $\theta$, and the network is trained by minimizing a loss comprised of two parts: $\mathcal{L}(\theta)=\mathcal{L}_{\text{PDE }}+\mathcal{L}_{\text{Data}}$.
The first term encourages the network to satisfy the PDE, while data-fit is a supervised loss, enforcing IBC and ensuring a good fit with the training data.
Specific forms of the losses, and further discussion can be found in Section C.2.

## 3 Relating canonicalization and frames

Frames and canonicalization both represent different ways to impose group equivariance or invariance on arbitrary functions. The most basic approach is simply to average over all group elements, sometimes refered to as the Reynolds operator. For a finite group $G$ and a continuous bounded function $\phi$ it is given by

$$ $\mathcal{R}_{G}(\phi)(x)=\frac{1}{|G|}\sum_{g\in G}\phi(g^{-1}\cdot x).$ $$

This can also be modified to impose $G$-equivariance by changing the integrand to $g\cdot\phi(g^{-1}\cdot x)$. Frames allow one to not have to integrate over the full group $G$ while still giving a invariant or equivariant function as output. Specifically, a frame is a function $\mathcal{F}:X\rightarrow 2^{G}\backslash\varnothing$ such that for any $x\in X$ and $g\in G$: $\mathcal{F}(gx)=g\mathcal{F}(x)$. In order to define an action on continuous groups, one must assume that the frame is finite, i.e. that for every $x\in X$ the set $\mathcal{F}(x)$ is finite. This is immediate if the group is finite, but if the group is non-finite this is a very strong restriction, as it implies that the stabilizer groups of the action are finite, as $\mathcal{F}(x)=G_{x}\mathcal{F}(x)$. This is detailed further in Appendix E.

In order to deal with this restriction, Dym et al. (), introduced the notion of *weighted weakly equivariant frames*. These generalize frames to be $G$-equivariant maps from $X$ to probability measures on $G$. Then for any continuous bounded function $\phi$ we may apply a weighted frame $\mu_{x}$ by
$\mu(\phi)(x)=\int_{G}\phi(g^{-1}\cdot x)\,d\mu_{x}(g)$. This is still not enough in the setting of non-compact groups as discussed in the subsequent section(^2^22*Weakly* equivariant weighted frames also require compact stabiliser subgroups, while the equivalent $\pi_{*}$ equivariance of E.1 can be defined for non-compact groups.). Even worse, there is an inherent limitation to using frames: the input to $\phi$ may end up outside the training distribution, resulting in worsened performance. The overall network would therefore need to be retrained, which may be exponentially hard (Kiani et al., ).

To deal with this, Kaba et al. () proposed canonicalization as simply a $G$-invariant function $f:X\rightarrow X$ such that for each $x\in X$, $f(x)\in Gx$, which upon pre-composing with a function would make it invariant. In order to preserve continuity of functions, one needs to impose that $f$ itself is continuous. In Dym et al. () it is shown that there are many topological obstructions to constructing such continuous canonicalization functions.

To get around these obstructions Ma et al. () proposes canonicalization to now be set-valued $G$-invariant function $v:X\rightarrow 2^{X}\backslash\varnothing$ for $v(x)$ finite, acting on continuous functions $\phi$ by

$$ $v(\phi)(x)=\frac{1}{|v(x)|}\sum_{y\in v(x)}\phi(y).$ $$

Contractive canonicalizations, which we refer to instead as *orbit canonicalizations*, are a specific kind of canonicalization where for every $x\in X$, $v(x)\subseteq Gx$. The orbit-stabilizer theorem applies to show these are equivalent to frames. Namely, there is a map $\alpha$ which sends equivariant frames to orbit canonicalizations. For an equivariant frame $\mathcal{F}$ the corresponding orbit canonicalization $\alpha(\mathcal{F})$ is defined as
$\alpha(\mathcal{F})(x)=\{g^{-1}\cdot x|g\in\mathcal{F}(x)\}$.
Ma et al. (, Theorem 3.1) show that for $G$ finite, $\alpha$ is an isomorphism between the set of equivariant frames and orbit canonicalizations.
Energy-based canonicalization was first introduced in Kaba et al. ().
In energy-based canonicalization, given an energy function $E$, the canonicalizing element $c:X\to G$ is defined as

$$ $c(x)\in\underset{{g\in G}}{\arg\min}\;E\left(g^{-1}x\right).$ $$

As shown in the subsequent section, considering the full set of minimizers, instead of a singleton, naturally gives rise to a frame, and naturally results in set-valued canonicalizations. Extending these to non-compact groups naturally gives rise to weighted closed canonicalizations.

### 3.1 Weighted Canonicalization

Figure: Figure 3: This diagram illustrates connections between the various notions of frames and canonicalizations introduced previously and in this work. The top row are the finite frames and canonicalizations, 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.
Refer to caption: x5.png

For this section let $G$ be a Lie group acting smoothly on a manifold $X$. Let $C^{0}_{b}(X)$ be the bounded continuous $\mathbb{R}$-valued functions on $X$. Let $\operatorname{PMeas}(X)$ be the set of Radon probability measures on $X$, with respect to the Borel $\sigma$-algebra on $X$. Let $B(X,\mathbb{R})$ be the space of bounded functions on $X$ and $B(X,\mathbb{R})^{G}$ is the subspace of $G$-invariant such functions.

###### Definition 3.1 .

A *weighted canonicalization* is a $G$-invariant function $\mu_{[\cdot]}:X\rightarrow\text{PMeas}(X)$. We call the space of weighted canonicalizations $\text{WCan}_{G}(X)$.

Given a weighted canonicalization $\mu$ we may define an operator $\mathcal{P}_{\mu}:B(X,\mathbb{R})\rightarrow B(X,\mathbb{R})^{G}$ via

$$ $(\mathcal{P}_{\mu}(\phi))(x)=\int_{X}\phi\,d\mu_{x}$ (1) $$

We shall often write the function $\mathcal{P}_{\mu}(\phi)$ simply as $\mu(\phi)$. As $\mu_{x}$ is a probability measure for each $x\in X$, we see that $|(\mathcal{P}_{\mu}(\phi))(x)|\leq\max_{X}\phi$. $\mu(\phi)$ is $G$-invariant as $\mu$ is. We also define the function $P_{\phi}(\mu)=\mathcal{P}_{\mu}(\phi)$.

The entire space $\operatorname{WCan}_{G}(X)$ is much too general to be practical. We would like to consider a much more natural subset from the perspective of canonicalization, namely orbit canonicalizations. We now define a corresponding weighted notion.

###### Definition 3.2 .

A *weighted orbit canonicalization* is a weighted canonicalization $\mu:X\rightarrow\text{PMeas}(X)$ such that for every $x\in X$, $\operatorname{supp}\mu(x)\subseteq Gx$. We denote the set of weighted orbit canonicalizations as $\text{WOCan}_{G}(X)$.

In particular, for weighted orbit canonicalizations both operators become projections, see Section E.2. As in the finite case, there is a function $\alpha:\operatorname{WFra}_{G}(X)\rightarrow\operatorname{WOCan}_{G}(X)$, which is surjective and injective up to weak-equivalence, see Section E.1.

Given the action on continuous functions, we may equip $\text{WCan}_{G}(X)$ with a certain weak topology, namely the coarsest topology such that all the $P_{\phi}$ for $\phi$ continuous and bounded become continuous functions. Specifically this means that a sequence of weighted canonicalizations $\mu_{n}$ converges to $\mu$ if and only if for every $\phi\in C^{0}_{b}(X)$ the sequence $\mu_{n}(\phi)$ converges to $\mu(\phi)$ 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 a weighted closed canonicalization.

###### Definition 3.3 .

A *weighted closed canonicalization* is a weighted canonicalization $\mu:X\rightarrow\text{PMeas}(X)$ such that for every $x\in X$, $\operatorname{supp}\mu(x)\subseteq\overline{Gx}$. Call the set of weighted closed canonicalizations $\text{WCCan}_{G}(X)$.

The following theorem shows that these allow us to fully characterize the behaviour of weighted orbit canonicalizations under taking limits.

###### Theorem 3.1 .

(Theorem F.4)
The sequential closure of $\operatorname{WOCan}_{G}(X)$ is $\operatorname{WCCan}_{G}(X)$

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. (), we may now ask whether such canonicalizations take continuous functions to continuous functions. In the presented framework this is a simple condition to add.

###### Definition 3.4 .

A weighted canonicalization $\mu:X\rightarrow\operatorname{PMeas}$ is continuous if for every $\phi\in C^{0}_{b}(X)$, $\mu(\phi)$ is continuous. Call the space of continuous weighted canonicalizations $\operatorname{WCan}^{\text{cts}}_{G}(X)$. We similarly define continuous weighted orbit canonicalizations and continuous weighted canonicalizations.

###### Proposition 3.2 .

A weighted canonicalization $\mu$ is continuous if and only if for every convergent sequence $x_{n}\rightarrow x$ in $X$, the sequence of measures $\mu_{x_{n}}$ converges weakly to $\mu_{x}$.

###### Theorem 3.3 .

(Theorem F.5)
$\operatorname{WCan}^{\mathrm{cts}}_{G}(X)$ is sequentially closed inside $\operatorname{WCan}_{G}(X)$.

###### Definition 3.1 .

###### Definition 3.2 .

###### Definition 3.3 .

###### Theorem 3.1 .

###### Definition 3.4 .

###### Proposition 3.2 .

###### Theorem 3.3 .

### 3.2 Energy Minimizing Canonicalization

With the relevant notions now defined, we can now define energy minimizing canonicalization for non-compact Lie groups. We define an energy function on $X$ to be a non-constant function $E:X\rightarrow[0,+\infty]$. We further assume that $E$ is topologically coercive, thus for every $N$ there is a compact set $K\subseteq X$ such that $E(x)>N$ for every $x\in X\backslash K$.

In the case where $G$ is finite, every orbit in $X$ necessarily has a minimum of $E$, and we note that energy minimization naturally results in a frame

$$ $\mathcal{F}_{E}(x)=\operatorname*{arg\,min}_{g\in G}E(g^{-1}x)$ (2) $$

###### Proposition 3.4 .

$\mathcal{F}_{E}$ defines an *equivariant frame*, i.e. $\mathcal{F}_{E}(gx)=g\mathcal{F}_{E}(x)$ for any $g\in G,\;x\in X$.

However when $G$ is not finite this set is not guaranteed to be finite. If $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, $\mathcal{M}_{E}(x)=\{x^{\prime}\in\overline{Gx}\,|\,E(x^{\prime})=\inf_{y\in
\overline{Gx}}E(y)\}$. This is the set of minimums of $E$ on the closure of the orbit of $x$.

###### Proposition 3.5 .

(F.2)
$\mathcal{M}_{E}(x)$ is $G$-invariant and non-empty.

If $\mathcal{M}_{E}(x)$ is finite for each $x\in X$, then we define the (non-weighted) closed canonicalization $\mathcal{C}_{E}\in\textnormal{WCCan}_{G}(X)$ by $\mathcal{C}_{E}(x)=\mathcal{M}_{E}(x)$. This we can then treat as a weighted canonicalization by taking the normalized counting measure on $\mathcal{M}_{E}(x)$ as usual.

For general $E$, we may not be able to put any reasonable probability measure on the set $\mathcal{M}_{E}(x)$, for example if $\mathcal{M}_{E}(x)$ has highly fractal geometry. We therefore must assume that $E$ is such that for all $x$, the Hausdorff measure of $\mathcal{M}_{E}(x)$ is non-zero. As each $\mathcal{M}_{E}(x)$ is compact, the Hausdorff measure is finite, hence combined with this assumption we may normalize it to obtain a probability measure on $\mathcal{M}_{E}(x)$ which we call $\mu_{x}$. Refer to Appendix D for more discussion on the Hausdorff measure and on why this assumption is reasonable.
Now we construct a weighted canonicalization by $\mathcal{C}_{E}(x)=\mu_{x}$.

Note that if $\mathcal{M}_{E}(x)$ 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 semi-continuous energy functions, we have the following characterization of which energies give rise to continuous weighted canonicalizations.

###### Lemma 3.6 .

(Lemma F.7)
Let $E$ be a lower semi-continuous energy function. Then $\mathcal{C}_{E}$ is continuous if and only if for every convergent sequence $x_{n}\rightarrow x$ we have $\textnormal{Li }\mathcal{M}_{E}(x_{n})=\mathcal{M}_{E}(x)$.

We have the following theorem telling us which weighted closed canonicalizations arise from energy minimizations.

###### Theorem 3.7 .

(Theorem E.3)
Let $\mu\in\operatorname{WCCan}_{G}(X)$ be such that $\mu_{x}$ is the normalized Hausdorff measure on its support and for all $y\in\overline{Gy}$, $\operatorname{supp}\mu_{y}=\operatorname{supp}\mu_{x}\cap\overline{Gx}$. Then there is an energy function $E$ such that $\mathcal{C}_{E}=\mu$.

###### Proposition 3.4 .

###### Proposition 3.5 .

###### Lemma 3.6 .

###### Theorem 3.7 .

## 4 Practical considerations

So far, in Section 3, we have provided a generalized theoretical understanding of frames and canonicalization, encompassing various scenarios with compact/non-compact and discrete/continuous groups. However, how these can be used in practice remains unaddressed. Unfortunately there is no unique way to choose the energy or the optimization routine for Equation 2, and we instead outline key considerations to guide these choices, with Section 5.1 serving as a guiding example.

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 on a domain, as well as information 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 energy function, based on the domain information; and (2) the group-based optimization routine for the problem in Equation 2, based on the domain information and the energy choice (potentially leading to problem relaxation as e.g. in Section 5.3.3 or improved initialization for iterative optimization as e.g. in Section 5.3.1).

##### 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, domain-specific information should be directly incorporated into the energy function. To give a concrete example, consider a setting in which some neural network is trained on samples living on a compact set $\mathcal{M}$. In such a setting we would only be interested in canonicalisations, 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 $\chi_{\mathcal{M}}$
is particularly suitable when the action is regular (or transitive). Otherwise, there may not be any elements in the set $\mathcal{M}$ that can be mapped to, in which case a relaxation using the distance function $d_{\mathcal{M}}(x)=\min_{y\in\mathcal{M}}\|x-y\|$ can be used. Such choices turn out to be beneficial, especially in the setting of Section 5.3.

In the case of data sampled from an underlying distribution $x\sim p_{\mathcal{X}}$, one natural choice would be to consider as energy function the negative log likelihood $E(x)=-\log p_{\mathcal{X}}(x)$, 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., ) or score-based diffusion models (Song et al., ) to approximate the true likelihood or its gradient. In this paper, for examples in Sections 5.3.3 and 5.2 we opt for a simpler variational autoencoder (Kingma, ), using the ELBO instead of the likelihood for the energy.

##### Non-convexity and Flatness

Even with a convex energy, minimization with respect to $g\in G$ introduces non-convexity. The orbit $Gx$ is often non-convex, e.g. see Figure 4, where this is a circle. Invoking a parameterization of a 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 5.3, we may not have access to a global group parameterization, significantly complicating the process. What is worse, the energy may exhibit many flat regions, with the gradients being noisy, as e.g. in Figure 4.

Thus the problem requires the use of non-convex optimization methods. There exist many proposals based on multi-particle systems like simulated annealing (Baudoin et al., ; Rutenbar, ), particle swarm (Huang et al., ) and consensus-based (Totzeck, ; Fornasier et al., ) optimization. Oftentimes the groups are low-dimensional, while the ambient space is not, making such particle based approaches applicable. The majority of these, however, are not easily adaptable to the setting of Lie groups. In all examples in Section 4, we instead utilize a multi-initialization strategy for fixed step gradient descent, which turns out to be sufficient for convergence. It is also worth mentioning that in many examples it may also be possible to relax the optimization to be performed over $x$, subject to an orbit distance constraint, potentially convexifying the problem.

To address the problem of flat zero likelihood regions, alternative energy functions can be considered, e.g. empirical distance $E(x)=\mathbb{E}_{X\sim p_{\mathcal{X}}}d(X,x)$. However, in Sections 5.3.3 and 5.2 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 adversarially on real and partially transformed data, based on the recent proposals in adversarial regularization Mukherjee et al. (); Shumaylov et al. (, ). Ideally, the regularization would also ensure uniqueness of minimizers to the problem in Equation 2, however it is not clear how to achieve this. Generically, we would want to construct an energy such that level sets are transverse to group orbits, similar to Panigrahi & Mondal (), which is indeed achieved by utilizing adversarial regularization.

##### Group non-compactness

Moving from compact groups to non-compact presents further challenges. Even with access to the exponential mapping, the mapping may not be surjective, making parameterization of the group using the Lie algebra rather complex. 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, , Chapter 4).
Lie algebras of many PDE based symmetry groups turn out to be solvable, which can be used for simple group parameterisations, leading to simplifications of the optimization problem. To be precise, if the Lie algebra is solvable, then we can always decompose any group element in a train of exponentials using the derived subalgebras (Varadarajan, , Theorem 3.18.11). For example, in the example of the heat equation Section 5.3.1, we can represent any element of the group as $\exp(a_{6}v_{6})\dots\exp(a_{1}v_{1})$. As a result, optimization over the group can be reduced to optimization over $a_{6},\dots,a_{1}$, which is a $\dim\mathfrak{g}$-dimensional Euclidean space.

When the algebra is not solvable, one has to revert back 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. As a result, one has to revert to coordinate descent (CD) based schemes. CD arises quite naturally within the context of classical optimization (e.g. see Wright () for an overview), and in the recent years started appearing in the context of Riemannian manifolds, and more recently for certain Lie groups (Han et al., ). The algorithms for both of these cases can be found in Appendix G.

## 5 Experiments

### 5.1 2D Toy Example

Figure: Figure 4: LieLAC optimised sample from a mix of Gaussian distributions.
Refer to caption: x26.png

To illustrate the universality of the approach, we consider one simple two dimensional example without any machine-learned components. We consider the problem of classificaiton under inherent rotation invariance illustrated on Figure 4.
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 4. As one would expect, without knowledge of the symmetries of the distribution, the approximation is far from the ground truth. In what follows, the energy is chosen to be the negative log-likelihood 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 no longer convex when viewed as a function of $g$. This results in a non-convex optimization problem with (potentially many) local minima, which one can get stuck in when running standard descent schemes, and as discussed in Section 4, one has to revert to more complex optimization routines. In this example, as illustrated in Figure 4, it turns out to be possible to find global minimizers, despite the non-convexity, by simply running gradient descent for a fixed number of iterations, but for a number of different initializations. The power of canonicalization can then be seen in Figure 1, turning a non-equivariant kNN classifier rotation equivariant.

**Table 1: Affine and Homography groups on MNIST classification**
| Name | MNIST Test Acc. | affNIST Test Acc. |
| --- | --- | --- |
| CNN | 0.985 | 0.629 |
| LieLAC [CNN] | 0.979 | 0.972 |
| affConv (MacDonald et al., ) | 0.982 | 0.943 |

### 5.2 Invariant Classification

Figure: (a) Affine MNIST examples
Refer to caption: x27.png

Figure: Figure 6: t-SNE plots: For MNIST, homNIST and canonicalized homNIST
Refer to caption: x29.png

To illustrate the effectiveness of LieLAC in achieving invariance of pre-trained models, we first test it on invariant image classification with respect to the groups of affine and homography transformations. We take MacDonald et al. () as the baseline for comparison due to the similarity in the required knowledge of the group required for implementation.

Here, we take a pre-trained convolutional classifer, trained on the original MNIST dataset. The combination with the canonicalizer, as described in Section 4, is then tested on the affine-perturbed MNIST (affNIST) (Gu & Tresp, ) and the homography-perturbed MNIST (homNIST) (MacDonald et al., ). The overall performance results for this example are reported in Table 1. We can see that the canonicalization based approach achieves improved equivariance compared to MacDonald et al. (), as illustrated by improved test accuracy on the datasets. We can also directly analyse achieved equivariance, as it appears only through canonicalization. We provide examples, visualising achieved equivariance in Figure 5. We can also analyze the effect of canonicalization globally by considering a t-SNE plot (Van der Maaten & Hinton, ) in Figure 6, showing that energy canonicalization turns the dataset into the usual MNIST form, clearly separating the classes, unlike homNIST.

### 5.3 Equivariant PDE evolution

In analogy with Akhound-Sadegh et al. (), we first consider two simple examples with non-trivial symmetry groups: the 1D Heat equation and the 1D Burgers’ equation, which also admit global parameterisations of the group elements (Koval & Popovych, ). This section primarily recaps 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.

We explicitly emphasise that in this paper we only consider symmetries for initial value problems, where there is no dependence of $g\cdot x$ and $g\cdot t$ on $u$. I.e. for $(x^{\prime},t^{\prime},u^{\prime})=g\cdot(x,t,u)$, we have $x^{\prime}(x,t,u)=x^{\prime}(x,t)$. In this case the action on $x,t$ can be separated from the action on $u$, making the desired equivariance for the physics-informed neural operator to be:

$$ $\mathcal{O}_{\theta}[u_{0},x_{0},t_{0}](x_{f},t_{f})=g\cdot\mathcal{O}_{\theta }[g^{-1}\cdot u_{0},g^{-1}\cdot x_{0},g^{-1}\cdot t_{0}](g^{-1}\cdot x_{f},g^{ -1}\cdot t_{f})$ (3) $$

For heat and Burgers’ equations, examples of canonicalized solutions are shown in Figure 2. Specific settings are provided in Sections C.3 and C.4, with examples in Sections C.3.1 and C.4.1.

#### 5.3.1 Heat equation

The heat equation is the simplest example of a PDE admitting a non-trivial group of Lie point symmetries:

$$ $u_{t}-\nu u_{xx}=0$ (4) $$

This equation admits a 6-dimensional solvable Lie algebra, see Section C.3, with a global action:

###### Theorem 5.1 ( Koval & Popovych ( 2023 ) for $\nu\neq 1$ ) .

The finite point symmetry group of the $(1+1)$-dimensional linear heat equation Equation 4 is constituted by the point transformations of the form

$$ $\tilde{t}=\frac{\alpha t+\beta}{\gamma t+\delta},\quad\tilde{x}=\frac{x+ \lambda_{1}t+\lambda_{0}}{\gamma t+\delta},\quad\tilde{u}=\sigma\sqrt{|\gamma t +\delta|}\exp\left(\frac{\gamma\left(x+\lambda_{1}t+\lambda_{0}\right)^{2}}{4{ \nu}(\gamma t+\delta)}-\frac{\lambda_{1}}{2{{\nu}}}x-\frac{\lambda_{1}^{2}}{4{ {\nu}}}t\right)u,$ $$

where $\alpha,\beta,\gamma,\delta,\lambda_{1},\lambda_{0}$ and $\sigma$ are arbitrary constants with $\alpha\delta-\beta\gamma=1$ and $\sigma\neq 0$.

This illustrates the global structure of the symmetry group (the finite simply connected component) as
$G^{H}=SL(2,\mathbb{R})\ltimes_{\phi}H(1,\mathbb{R}),$
where $H(1,\mathbb{R})$ is the rank one polarized Heisenberg group. Definition of the antihomomorphism $\phi$, and the full derivation can be found in Section C.3.

###### Theorem 5.1 ( Koval & Popovych ( 2023 ) for ν ≠ 1 𝜈 1 \nu\neq 1 italic_ν ≠ 1 ) .

#### 5.3.2 Burgers’ Equation

The viscous 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.

$$ $u_{t}+uu_{x}-\nu u_{xx}=0$ (5) $$

This equation admits a 5-dimensional solvable Lie algebra, see Section C.4, with a global action:

###### Theorem 5.2 ( Koval & Popovych ( 2023 ) for $\nu\neq 1$ ) .

The point symmetry group $G^{\mathrm{B}}$ of the Burgers Equation 5 consists of the point transformations of the form

$$ $\tilde{t}=\frac{\alpha t+\beta}{\gamma t+\delta},\quad\tilde{x}=\frac{x+ \lambda_{1}t+\lambda_{0}}{\gamma t+\delta},\quad\tilde{u}=(\gamma t+\delta)u- \gamma x+\lambda_{1}\delta-\lambda_{0}\gamma,$ $$

where $\alpha,\beta,\gamma,\delta,\lambda_{1}$ and $\lambda_{0}$ are arbitrary constants with $\alpha\delta-\beta\gamma=1$.

And similar to before the global structure of the group can be inferred as
$G^{B}=\mathrm{SL}(2,\mathbb{R})\ltimes_{\breve{\varphi}}\left(\mathbb{R}^{2},+\right)$
with the antihomomorphism $\breve{\varphi}$ defined in Section C.4.

###### Theorem 5.2 ( Koval & Popovych ( 2023 ) for ν ≠ 1 𝜈 1 \nu\neq 1 italic_ν ≠ 1 ) .

#### 5.3.3 Allen-Cahn Equation

We test the Poseidon (Herde et al., ) foundation model, fine-tuned on the Allen-Cahn equation (ACE) on similar data to which it was trained on, and on transformed data (using Lie-point symmetry group of ACE).

$$ $u_{t}=u_{xx}+u_{yy}-\epsilon^{2}u\left(u^{2}-1\right)$ (6) $$

We consider solving Equation 6 on a 2D domain with periodic boundary conditions for $t\in[0,1],x\in\Omega=[0,1]^{2}$. This equation admits a 4-dimensional Lie algebra, see Section C.5, with $\mathrm{SE}(2)$ 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., ) model itself. Firstly, $x_{0}$ and $t_{0}$ in Equation 3 are implicitly defined in the model to be fixed to the standard domain $x_{0}=\Omega$, and $t_{0}=0$. As such, the only input is the initial condition $u_{0}|_{\Omega}$. Therefore, we are going to have to only consider transformations, that preserve $x_{0}$ and $t_{0}$. This can be imposed directly within the energy canonicalization framework as described in Section 4 by considering the characteristic functions as follows:

$$ $E_{\textsc{AC}}(x_{0},t_{0},u_{0},x_{f},t_{f})=\chi_{\{\Omega\}}(x_{0})+\chi_{ \{0\}}(t_{0})+\chi_{[0,1]}(t_{f})+\chi_{\Omega}(x_{f})+E(u_{0})$ (7) $$

Note, that this is able to significantly reduce the computational cost of optimization, as the group generated by $v_{4}$ in 26, resulting in rotations of the underlying domain, can be reduced to a discrete subgroup $\mathrm{C}_{4}\leq\mathrm{SO}(2)$, as the only transformations such that $g^{-1}\Omega=\Omega$. Further information on the training and testing data is available in Section C.5. The overall performance is reported in Table 2.

**Table 2: ACE Test error**
| Name | ID Test Error eq. 27 | OOD Test Error eq. 28 | Average |
| --- | --- | --- | --- |
| Poseidon | $6.448\text{\times}{10}^{-04}$ | $7.619\text{\times}{10}^{-03}$ | $4.132\text{\times}{10}^{-03}$ |
| LieLAC [Poseidon] | $1.592\text{\times}{10}^{-03}$ | $2.916\text{\times}{10}^{-03}$ | $2.254\text{\times}{10}^{-03}$ |
| LieLAC [Poseidon+ ft.] | $9.667\text{\times}{10}^{-04}$ | $1.143\text{\times}{10}^{-03}$ | $1.055\text{\times}{10}^{-03}$ |
| Poseidon + ft (can). | $7.952\text{\times}{10}^{-04}$ | $2.021\text{\times}{10}^{-03}$ | $1.408\text{\times}{10}^{-03}$ |

##### Finetuning

As seen in prior works (Mondal et al., ), the canonicalized data may end up misaligned with the original data. For this reason finetuning of the model/canonicalization is required. While it may seem that finetuning the resulting operator on canonicalized data is as hard as augmenting the original data (Brandstetter et al., ), it turns out to not be, as fine-tuning can be done on a very small data-pool (and in this case was done only on 100 trajectories). Intuitively that is because the minima of the optimization problem Equation 2, may not exactly fall into the training data, but a well-chosen energy will have minima close to it. If the orbits in the training data are well-sampled, then there is no need for finetuning, as was e.g. seen in Section 5.2. However, based on the training data in Equation 27, the orbits under $\mathrm{C}_{4}$ are well-sampled, while the orbits under $\mathrm{SE}(2)$ are not. This is illustrated in Table 2, wherein the error falls 2/3-fold after finetuning. An example is shown on Figure 2(c), with further illustrations in Figure 14.

## 6 Conclusion

In this paper we provided a 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, as well as Neural Operators Lie point symmetry equivariant.

##### Future work

Our example in Section 5.3.3 illustrates the promise of the proposed framework for 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.

Of particularly 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

There are a number of limitations however. Energy-based canonicalization tends to be rather slow, and improvements to the optimization procedure could significantly speed up the evaluation, as discussed in Section 4. In addition to that, Lie point symmetries tend to be rather problematic to impose in existing neural operators, as oftentimes boundary conditions are hard-coded into a given model, thus a real foundation model, operating over initial conditions, parameters and boundary conditions should be developed. Alternatively, other types of PDE symmetries should be considered (Brandstetter et al., ).

### Acknowledgments

The authors would like to thank Lucy Richman, Hong Ye Tan and Yury Korolev for useful discussions and feedback. FS acknowledges support from the EPSRC advanced career fellowship EP/V029428/1. MW was supported by NSF Awards DMS-2406905 and CBET-2112085, and a Sloan Research Fellowship. CBS acknowledges support from the Philip Leverhulme Prize, the Royal Society Wolfson Fellowship, the EPSRC advanced career fellowship EP/V029428/1, EPSRC Grants EP/S026045/1 and EP/T003553/1, EP/N014588/1, EP/T017961/1, the Wellcome Innovator Awards 215733/Z/19/Z and 221633/Z/20/Z, the European Union Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie Grant agreement No. 777826 NoMADS, the Cantab Capital Institute for the Mathematics of Information and the Alan Turing Institute.

## Appendix A Appendix

### A.1 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, we 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 (), and an interested reader is referred there for more in-depth discussions of the concepts.

This subsection introduces the concept of PDE point symmetries – transformations that map solutions of a 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 $\Delta$ of $n$-th order differential equations involving $p$ independent variables $x=\left(x^{1},\dots,x^{p}\right)\in X=\mathbb{R}^{p}$, and $q$ dependent variables $u=\left(u^{1},\dots,u^{q}\right)\in U=\mathbb{R}^{q}$.
PDEs define relations between various derivatives of the function, for denoting which we use the following multi-index notation for $J=\left(j_{1},\ldots,j_{k}\right)$, with
$\partial_{J}f(x)=\frac{\partial^{k}f(x)}{\partial x^{j_{1}}\partial x^{j_{2}}
\cdots\partial x^{j_{k}}}.$

##### Jet spaces and Prolongations:

Consider $f:X\rightarrow U$ smooth and let $U_{k}$ be the Euclidean space, endowed with coordinates $u_{J}^{\alpha}=\partial_{J}f^{\alpha}(x)$ so as to represent the above derivatives. Furthermore, set $U^{(n)}=U\times U_{1}\times\cdots\times U_{n}$ to 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\times U^{(n)}$, whose coordinates represent the independent variables, the dependent variables and the derivatives of the dependent variables up to order $n$ is called the $n$-th order jet space of the underlying space $X\times U$.
Given a smooth function $u=f(x)$, there is an induced function $u^{(n)}=\operatorname{pr}^{(n)}f(x)$, called the $n$-th prolongation of $f$, defined by the equations $u_{J}^{\alpha}=\partial_{J}f^{\alpha}(x).$
Thus $\mathrm{pr}^{(n)}f$ is a function from $X$ to the space $U^{(n)}$, and for each $x$ in $X$, $\operatorname{pr}^{(n)}f(x)$ is a vector representing values of $f$ and all its derivatives up to order $n$ at the point $x$.

##### PDE Symmetry groups:

With the notation defined, we can now write down the PDE, given as a system of equations $\Delta\left(x,u^{(n)}\right)=0,$ with $\Delta:X\times U^{(n)}\rightarrow\mathbb{R}^{l}$ smooth. The differential equations themselves tell where the given map $\Delta$ vanishes on $X\times U^{(n)}$, and thus determine a subvariety

$$ $S_{\Delta}=\left\{\left(x,u^{(n)}\right):\Delta\left(x,u^{(n)}\right)=0\right \}\subset X\times U^{(n)}$ $$

Solutions of the system will be of the form $u=f(x)$. A symmetry group of the system $\Delta$ will be a local group of transformations, $G^{\Delta}$, acting on some open subset $M\subset X\times U$ in such that “$G$ transforms solutions of $\Delta$ to other solutions of $\Delta$”.

##### Prolongations of group actions and vector fields

Now suppose $G$ is a local group of transformations acting on an open subset $M\subset X\times U$ of the space of independent and dependent variables. There is an induced local action of $G$ on the $n$-jet space $M^{(n)}$, called the $n$-th prolongation of $G$ denoted $\mathrm{pr}^{(n)}G$. This prolongation is defined so that it transforms the derivatives of functions $u=f(x)$ into the corresponding derivatives of the transformed function $\tilde{u}=\tilde{f}(\tilde{x})$.
Consider now $\mathrm{v}$ a vector field on $M$, with corresponding (local) one-parameter group $\exp(\varepsilon\mathrm{v})$. The $n$-th prolongation of $\mathrm{v}$, denoted $\mathrm{pr}^{(n)}\mathrm{v}$, will be a vector field on the $n$-jet space $M^{(n)}$, and is defined to be the infinitesimal generator of the corresponding prolonged one-parameter group $\mathrm{pr}^{(n)}[\exp(\varepsilon\mathrm{v})]$. In other words,

$$ $\left.\operatorname{pr}^{(n)}\right|_{\left(x,u^{(n)}\right)}=\left.\frac{ \mathrm{d}}{\mathrm{d}\varepsilon}\right|_{\varepsilon=0}\operatorname{pr}^{(n )}[\exp(\varepsilon\mathbf{v})]\left(x,u^{(n)}\right)$ $$

for any $\left(x,u^{(n)}\right)\in M^{(n)}$. These can be derived using Olver (, Theorem 2.36).

###### Theorem A.1 (2.31 in Olver ( 1993 ) ) .

Suppose $\Delta\left(x,u^{(n)}\right)=0,$ is a system of differential equations of maximal rank defined over $M\subset X\times U$. If $G$ is a local group of transformations acting on $M$, and

$$ $\operatorname{pr}^{(n)}\mathrm{v}\left[\Delta\left(x,u^{(n)}\right)\right]=0, \quad\text{ whenever }\quad\Delta\left(x,u^{(n)}\right)=0$ $$

for every infinitesimal generator $\mathrm{v}$ of $G$, then $G$ is a symmetry group of the system.

###### Theorem A.1 (2.31 in Olver ( 1993 ) ) .

## Appendix B Further Mnist Visualisation

Figure: (a) Affine MNIST examples
Refer to caption: x30.png

Figure: (a) t-SNE post-affine canonicalization
Refer to caption: x32.png

#### B.0.1 Failure modes

Of particular interest are the failure modes of canonicalization, warranting further theoretical investigation. We present the failure modes in Figure 9, and as we 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.

Figure: Figure 9: Misclassification examples
Refer to caption: x34.png

## Appendix C Further PDE discussion

### C.1 Canonicalization as a unifying view

The idea of putting the evolutionary equation and initial into a canonical form prior to evolving is not novel, and in fact quite standard in all numerical analysis textbooks. The first step in solving an ODE numerically is to normalise the equations and remove units and normalise the domain. For example, when solving a time evolution equation on the domain $x\in[0,15]$, one would first normalise the $x$ 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 on achieving exact unit equivariance - which is precisely scaling group canonicalisation (Villar et al., ). Similar problem also often arises in the problems of image registration (Kanter & Lellmann, ) with an infinite dimensional Lie algebra of diffeomorphisms. Another particular example being Procrustes analysis (Gower, ), 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 to improve the optimization, and used for PDE symmetry groups.

### C.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 physics-informed machine learning papers (Akhound-Sadegh et al., ; Wang et al., ; Herde et al., ) respectively. In each case we take the pre-trained neural PDE solver and demonstrate a failure mode when evaluating on out-of-distribution initial conditions. We then optimize for a group action derived from the Lie algebra based energy minimisation 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)$ is parameterised as a neural network $u_{\theta}(t,x)$ with parameters $\theta$. The loss function is then compromised of two parts, with different weightings:

$$ $\mathcal{L}(\theta)=\alpha_{\mathrm{PINN}}\mathcal{L}_{\text{PDE }}+\gamma_{ \mathrm{data}}\mathcal{L}_{\text{data-fit }}$ $$

The first term is the physics-informed objective of Raissi et al. (), 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 a finite set of points $(t,{x})_{1:N_{\mathrm{dom}}}$ which are sampled from inside the domain $[0,T]\times\Omega$ to obtain the PDE loss:

$$ $\mathcal{L}_{\mathrm{PDE}}=\frac{1}{N_{\mathrm{dom}}}\sum_{i=1}^{N_{\mathrm{ dom}}}\left\|\partial_{t}u_{\theta}(t_{i},{x_{i}})+\mathcal{D}_{{x}}\left[u_{ \theta}(t_{i},{x_{i}})\right]\right\|_{2}^{2}$ (8) $$

The second term, is a supervised loss enforcing initial and boundary conditions:

$$ $\mathcal{L}_{\text{data-fit }}=\frac{1}{N_{\mathrm{ic}}}\sum_{i=1}^{N_{\mathrm {ic}}}\left\|u_{\theta}\left(0,{x}_{i}^{0}\right)-f\left(x_{i}^{0}\right) \right\|_{2}^{2}+\frac{1}{N_{\mathrm{bc}}}\sum_{i=1}^{N_{\mathrm{bc}}}\left\|u _{\theta}\left(t_{i}^{b},{x}_{i}^{b}\right)-g\left(t_{i}^{b},x_{i}^{b}\right) \right\|_{2}^{2},$ (9) $$

with $\left({x}^{0}\right)_{1:N_{\mathrm{ic}}}\in\Omega$ sample at which the initial condition function $f$ is evaluated. $\left(t^{b},{x}^{b}\right)_{1:N_{\mathrm{bc}}}$ are $N_{\mathrm{bc}}$ points sampled on the boundary $[0,T]\times\partial\Omega$. The two loss terms are then weighted by $\alpha_{\mathrm{PINN}}$ for PDE loss 8 and $\gamma_{\mathrm{data}}$ for data-fit 9 to be consistent with Akhound-Sadegh et al. ().

### C.3 Heat equation

Simplest PDE example is the heat equation, admitting a non-trivial group of Lie point symmetries, beyond space isometries.

$$ $u_{t}-\nu u_{xx}=0$ $$

This equation admits the following 6-dimensional Lie algebra (note the typo in Akhound-Sadegh et al. ()), see Olver (); Koval & Popovych ():

$$ $\displaystyle v_{1}=\partial_{x},\quad v_{2}=\partial_{t},\quad v_{3}=\frac{1} {\nu}u\partial_{u},\quad v_{4}=2t\partial_{t}+x\partial_{x}-\frac{1}{2}u \partial_{u},$ (10) $\displaystyle v_{5}=t\partial_{x}-\frac{1}{2\nu}xu\partial_{u},\quad v_{6}=t^{ 2}\partial_{t}+tx\partial_{x}-\frac{1}{4\nu}\left(x^{2}+2\nu t\right)u\partial _{u}.$ (11) $$

The resulting one parameter flows are:

$$ $\begin{array}[]{llll}\exp[v_{1}](\epsilon):&\tilde{t}=t,&\tilde{x}=x+\epsilon, &\tilde{u}=u,\\ \exp[v_{2}](\epsilon):&\tilde{t}=t+\epsilon,&\tilde{x}=x,&\tilde{u}=u,\\ \exp[v_{3}](\epsilon):&\tilde{t}=t,&\tilde{x}=x,&\tilde{u}=\mathrm{e}^{ \epsilon/\nu}u,\\ \exp[v_{4}](\epsilon):&\tilde{t}=\mathrm{e}^{2\epsilon}t,&\tilde{x}=\mathrm{e} ^{\epsilon}x,&\tilde{u}=\mathrm{e}^{-\frac{1}{2}\epsilon}u,\\ \exp[v_{5}](\epsilon):&\tilde{t}=t,&\tilde{x}=x+\epsilon t,&\tilde{u}=\mathrm{ e}^{-\frac{1}{4\nu}\left(\epsilon^{2}t+2\epsilon x\right)}u,\\ \exp[v_{6}](\epsilon):&\tilde{t}=\frac{t}{1-\epsilon t},&\tilde{x}=\frac{x}{1- \epsilon t},&\tilde{u}=\sqrt{|1-\epsilon t|}\mathrm{e}^{\frac{\epsilon x^{2}}{ 4\nu(\epsilon t-1)}}u,\\ \end{array}$ $$

The only non-zero Lie brackets are:

$$ $\displaystyle{\left[v_{4},v_{2}\right]=-2v_{2},\quad[v_{4},v_{6}]=2v_{6},\quad \left[v_{2},v_{6}\right]=v_{4},}$ (12) $\displaystyle{\left[v_{2},v_{5}\right]=v_{1},\quad\left[v_{4},v_{5}\right]=v_{ 5},\quad\left[v_{4},v_{1}\right]=-v_{1},\quad\left[v_{6},v_{1}\right]=-v_{5},}$ (13) $\displaystyle{\left[v_{5},v_{1}\right]=\frac{1}{2}v_{3}.}$ (14) $$

In particular, these would be enough for LieLAC, however in this particular example it is possible to show more.
In particular, as shown in Koval & Popovych ().

###### Theorem C.1 ( Koval & Popovych ( 2023 ) for $\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\nu}\neq 1$ ) .

The point symmetry pseudogroup $G$ of the $(1+1)$-dimensional linear heat equation (1) is constituted by the point transformations of the form

$$ $\displaystyle\tilde{t}=\frac{\alpha t+\beta}{\gamma t+\delta},\quad\tilde{x}= \frac{x+\lambda_{1}t+\lambda_{0}}{\gamma t+\delta},$ $\displaystyle\tilde{u}=\sigma\sqrt{|\gamma t+\delta|}\exp\left(\frac{\gamma \left(x+\lambda_{1}t+\lambda_{0}\right)^{2}}{4{\color[rgb]{1,0,0}\definecolor[ named]{pgfstrokecolor}{rgb}{1,0,0}{\nu}}(\gamma t+\delta)}-\frac{\lambda_{1}}{ 2{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\nu}}}x- \frac{\lambda_{1}^{2}}{4{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor} {rgb}{1,0,0}{\nu}}}t\right)(u+h(t,x)),$ $$

where $\alpha,\beta,\gamma,\delta,\lambda_{1},\lambda_{0}$ and $\sigma$ are arbitrary constants with $\alpha\delta-\beta\gamma=1$ and $\sigma\neq 0$, and $h$ is an arbitrary solution of the heat equation.

We now describe the product on $G$. First we note that there is a subgroup $R$ consisting of those elements for which $\lambda_{1}=\lambda_{0}=0$ and $\sigma=1$. This is evidently isomorphic to $\mathrm{SL}(2,\mathbb{R})$. Therefore we shall group the $\alpha,\beta,\gamma,\delta$ terms into a matrix $A=\begin{pmatrix}\alpha&\beta\\
\gamma&\delta\end{pmatrix}$. Further we note that there are two connected components of this group, corresponding to sign of $\sigma$, and the group can be written as a direct product corresponding to this sign

$$ $G=G^{\prime}\times\mathbb{Z}_{2}$ (15) $$

Thus we may simply describe the product on $G^{\prime}$. For this we may assume that the two elements we multiply have $\sigma>0$. For clarity, we shall write out the product rule in terms of the logarithm of $\sigma$. The product can now be written as

$$ $\displaystyle(A^{\prime},\lambda_{1}^{\prime},\lambda_{0}^{\prime},\ln\sigma^{ \prime})\cdot(A,\lambda_{1},\lambda_{0},\ln\sigma)=(A^{\prime}A,\lambda_{1}+ \alpha\lambda_{1}^{\prime}+\gamma\lambda_{0}^{\prime},\lambda_{0}+\beta\lambda _{1}^{\prime}+\delta\lambda_{0}^{\prime},$ $\displaystyle\ln\sigma+\ln\sigma^{\prime}+\frac{1}{4}(\lambda_{0}^{\prime} \lambda_{1}^{\prime}-(\alpha\lambda_{1}^{\prime}+\gamma\lambda_{0}^{\prime})( \beta\lambda_{1}^{\prime}+\delta\lambda_{0}^{\prime}))-\frac{\lambda_{0}}{2}( \alpha\lambda_{1}^{\prime}+\gamma\lambda_{0}^{\prime}))$ $$

We see that the subgroup $F$ consisting of those elements for which $\alpha=\delta=1,\beta=\gamma=0$ and $\sigma>0$ is isomorphic to the rank one polarized Heisenberg group $\mathrm{H}(1,\mathbb{R})$ via

$$ $(\lambda_{1},\lambda_{0},\ln\sigma)\mapsto\begin{pmatrix}1&-\frac{1}{2}\lambda _{1}&\ln\sigma\\ 0&1&\lambda_{0}\\ 0&0&1\end{pmatrix}$ (16) $$

where we use the standard embedding of $\mathrm{H}(1,\mathbb{R})$ into the group of $3\times 3$ real valued matrices. Note, that while the parameterization here looks the same as the one from Koval & Popovych (), it is different (resulting in different $\phi)$, as in that paper, the Heisenberg group and the polarized Heisenberg group are conflated, making the product rule unclear. Here we fix a specific form of the Heisenberg group as a subgroup of the group of $3\times 3$ real matrices, with standard matrix multiplication as the product rule.
The product formula can then be described in terms of $F$ and $R$ as well. Let $\phi:\mathrm{SL}(2,\mathbb{R})\rightarrow\text{Aut}(\mathrm{H}(1,\mathbb{R}))$ be the antihomomorphism defined by

$$ $\phi(A)=\begin{pmatrix}\lambda_{1}\\ \lambda_{0}\\ \ln\sigma\end{pmatrix}\mapsto\begin{pmatrix}\alpha\lambda_{1}+\gamma\lambda_{0 }\\ \beta\lambda_{1}+\delta\lambda_{0}\\ \ln\sigma+\frac{1}{4}\lambda_{0}\lambda_{1}-\frac{1}{4}(\alpha\lambda_{1}+ \gamma\lambda_{0})(\beta\lambda_{1}+\delta\lambda_{0})\end{pmatrix}$ (17) $$

where write $\mathrm{H}(1,\mathbb{R})$ using the identification $F\cong\mathrm{H}(1,\mathbb{R})$.
Then we can write

$$ $G^{\prime}=\mathrm{SL}(2,\mathbb{R})\ltimes_{\phi}\mathrm{H}(1,\mathbb{R})$ (18) $$

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^{\prime}$ is $\mathrm{SL}(2,\mathbb{R})\times\mathrm{H}(1,\mathbb{R})$ and the product is

$$ $(A^{\prime},H^{\prime})\cdot(A,H)=(A^{\prime}A,\phi(A)(H^{\prime})H)$ (19) $$

Particularly for our implementation we can either find the action of the inverse of $\exp(e_{6}v_{6})\dots\exp(e_{1}v_{1})$ as $\exp(-e_{1}v_{1})\dots\exp(-e_{6}v_{6})$, but here we shave the exact parameterisation of the group, so might as well keep consistent. So, we need to write down an inverse of a given element in the form of $A,H$. This is $\left(A^{-1},\phi(A^{-1})\left(H^{-1}\right)\right)$.
I.e. for $\alpha,\beta,\gamma,\delta,\lambda_{0},\lambda_{1},\sigma$, the inverse is

$$ $A^{-1}=\begin{pmatrix}\delta&-\beta\\ -\gamma&\alpha\end{pmatrix}\quad H^{-1}=\begin{pmatrix}1&\frac{1}{2}\lambda_{1 }&-\ln\sigma-\frac{1}{2}(\lambda_{1}\lambda_{0})\\ 0&1&-\lambda_{0}\\ 0&0&1\end{pmatrix}$ $$

Thus, $(\lambda_{1},\lambda_{0},\ln\sigma)\to(-\lambda_{1},-\lambda_{0},\ln\sigma^{-1
}+\frac{1}{2}(\lambda_{1}\lambda_{0}))\to(-\delta\lambda_{1}+\gamma\lambda_{0}
,\beta\lambda_{1}-\alpha\lambda_{0},\ln\sigma^{-1}-\frac{1}{2}\lambda_{1}
\lambda_{0}+\frac{1}{4}\lambda_{0}\lambda_{1}-\frac{1}{4}(-\delta\lambda_{1}+
\gamma\lambda_{0})(\beta\lambda_{1}-\alpha\lambda_{0}))$
The inverse is
$\delta,-\beta,-\gamma,\alpha,\beta\lambda_{1}-\alpha\lambda_{0},-\delta\lambda
_{1}+\gamma\lambda_{0},-\ln\sigma-\frac{1}{4}\lambda_{0}\lambda_{1}-\frac{1}{4
}(-\delta\lambda_{1}+\gamma\lambda_{0})(\beta\lambda_{1}-\alpha\lambda_{0})$

###### Theorem C.1 ( Koval & Popovych ( 2023 ) for ν ≠ 1 𝜈 1 \color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{\nu}\neq 1 italic_ν ≠ 1 ) .

##### Numerical results

For the numerical experiments, we consider solving the Heat equation Equation 4 with diffusivity $\nu=0.1$ on the spatial domain $[0,2\pi]$, with periodic boundary conditions for $t\in[0,16]$.

We train a physics-informed DeepONet (Lu et al., ; Wang et al., ) with initial conditions of the form following Akhound-Sadegh et al. (); Brandstetter et al. ()

$$ $u_{0}(x)=\sum_{i=1}^{K}A_{k}\sin\left(2\pi l_{k}x/L+\phi_{k}\right).$ (20) $$

However we focus on a reduced task sampling parameters tightly distributed parameters as $K=1$, $l_{k}=2$, $\phi_{k}=0$ and $A\sim\mathcal{U}[0.95,1.05]$. The model is then trained using the PINN loss with sampling weights of $\alpha_{\mathrm{PINN}}=150$, $\gamma_{\mathrm{data}}=20$, we use a learning rate of 0.001 and batch size of 8 and train for 100k epochs following Akhound-Sadegh et al. (). We also sample $Nf_{\mathrm{train}}$ initial condition functions, sampling $N_{\mathrm{ic}}=200$ points at $t=0$ to evaluate $u_{0}$, $N_{\mathrm{bc}}=100$ points on the boundaries $x=0$ or $x=2\pi$ to enforce the periodic boundary conditions and finally $N_{\mathrm{dom}}=500$ points in the interior of the domain where the residual physics loss is enforced. We used the official implementation of DeepONet from the code and examples provided by the library DeepXDE (Lu et al., ).

We report results in Table 4. First we show the accuracy of the standard DeepONet trained in the above setting. Next we demonstrate the failure mode of a DeepONet by sampling $Nf_{\mathrm{test}}=100$ test functions sampled from $K=1$, $l_{k}=2$, $\phi_{k}=0$ and $A\sim\mathcal{U}[2.5,5.0]$. Specifically this induces a scaling to the amplitude of the initial conditions which corresponds to the $\tilde{u}$ symmetry in Theorem C.1. Finally we apply LieLAC following the steps of the canonicalization pipeline as shown in Figure 2(a). As demonstrated in the results and Figure 2(a) it is clear that the pretrained DeepONet is unable to generalise to samples of out-of-distribution initial conditions to those seen in training. However LieLAC is able to learn a canonicalizing group action that allows a pretrained model to perform as if in distribution.

##### Choice of Energy

Given the initial conditions in Equation 20, 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 to the trained domain. Letting $\operatorname{jet}_{0}=(u_{0},x_{0},t_{0})$ and $\operatorname{jet}_{\Omega}=(1,\Omega,0)$, the energy is:

$$ $\displaystyle E_{\text{Heat}}(x_{0},t_{0},u_{0},x_{f},t_{f})=\operatorname{ dist}\left[\operatorname{max}(\operatorname{jet}_{0}),\operatorname{max}( \operatorname{jet}_{\Omega})\right]$ $\displaystyle+\operatorname{dist}\left[\operatorname{min}(\operatorname{jet}_{ 0}),\operatorname{min}(\operatorname{jet}_{\Omega})\right]$ (21) $\displaystyle\quad+\operatorname{dist}\left[x_{f},\Omega\right]+\operatorname{ dist}\left[t_{f},16\right]$ $$

**Table 3: $L^{2}$ relative error for Heat equation**
| Name | ID Test Acc. | OOD Test Acc. |
| --- | --- | --- |
| DeepONet | 0.0425 | 0.6013 |
| LieLAC [DeepONet] | 0.0461 | 0.0425 |

#### C.3.1 Further heat illustrations

Figure: Figure 10: Canonicalization for the Heat equation given out of distribution $u_{0}(x)$ and improved equivariant prediction
Refer to caption: x35.png

### C.4 Burgers’ equation

The viscous 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.

$$ $u_{t}+uu_{x}-\nu u_{xx}=0$ $$

This equation admits the following 5-dimensional Lie algebra (note the typo in Akhound-Sadegh et al. ()), see Freire (); Koval & Popovych ():

$$ $\displaystyle v_{1}=\partial_{x},\quad v_{3}=2t\partial_{t}+x\partial_{x}-u \partial_{u},\quad v_{5}=t^{2}\partial_{t}+tx\partial_{x}+(x-tu)\partial_{u},$ (22) $\displaystyle v_{2}=\partial_{t},\quad v_{4}=t\partial_{x}+\partial_{u}.$ $$

The resulting one parameter flows, for $\epsilon$ arbitrary constant, are:

$$ $\begin{array}[]{llll}\exp[v_{1}](\epsilon):&\tilde{t}=t,&\tilde{x}=x+\epsilon, &\tilde{u}=u,\\ \exp[v_{2}](\epsilon):&\tilde{t}=t+\epsilon,&\tilde{x}=x,&\tilde{u}=u,\\ \exp[v_{3}](\epsilon):&\tilde{t}=\mathrm{e}^{2\epsilon}t,&\tilde{x}=\mathrm{e} ^{\epsilon}x,&\tilde{u}=\mathrm{e}^{-\epsilon}u,\\ \exp[v_{4}](\epsilon):&\tilde{t}=t,&\tilde{x}=x+\epsilon t,&\tilde{u}=u+ \epsilon,\\ \exp[v_{5}](\epsilon):&\tilde{t}=\frac{t}{1-\epsilon t},&\tilde{x}=\frac{x}{1- \epsilon t},&\tilde{u}=u(1-\epsilon t)+\epsilon x,\\ \end{array}$ $$

The non-zero brackets are:

$$ $\displaystyle{\left[v_{3},v_{2}\right]=-2v_{2},\quad[v_{3},v_{5}]=2v_{5},\quad \left[v_{2},v_{5}\right]=v_{3},}$ (23) $\displaystyle{\left[v_{2},v_{4}\right]=v_{1},\quad\left[v_{3},v_{4}\right]=v_{ 4},\quad\left[v_{3},v_{1}\right]=-v_{1},\quad\left[v_{5},v_{1}\right]=-v_{4}}$ (24) $$

Once again, this would be enough for the LieLAC formulation, however in this particular example, we can say more:

###### Theorem C.2 ( Koval & Popovych ( 2023 ) for ${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\nu}\neq 1$ ) .

The point symmetry group $G^{\mathrm{B}}$ of the Burgers’ equation (9) consists of the point transformations of the form

$$ $\tilde{t}=\frac{\alpha t+\beta}{\gamma t+\delta},\quad\tilde{x}=\frac{x+ \lambda_{1}t+\lambda_{0}}{\gamma t+\delta},\quad\tilde{u}=(\gamma t+\delta)u- \gamma x+\lambda_{1}\delta-\lambda_{0}\gamma,$ $$

where $\alpha,\beta,\gamma,\delta,\lambda_{1}$ and $\lambda_{0}$ are arbitrary constants with $\alpha\delta-\beta\gamma=1$.

And similar to before the global structure of the group can be inferred, with the group $G^{\mathrm{B}}$ isomorphic to the group $\mathrm{SL}(2,\mathbb{R})\ltimes_{\breve{\varphi}}\left(\mathbb{R}^{2},+\right)$, where $\breve{\varphi}:\mathrm{SL}(2,\mathbb{R})\rightarrow\operatorname{Aut}\left(
\left(\mathbb{R}^{2},+\right)\right)$ is the group antihomomorphism defined by $\breve{\varphi}(\alpha,\beta,\gamma,\delta)=\left(\lambda_{1},\lambda_{0}
\right)\mapsto\left(\lambda_{1},\lambda_{0}\right)\varrho_{1}(\alpha,\beta,
\gamma,\delta)$. Thus, the group $G^{\mathrm{B}}$ is connected.

$$ $\varrho_{1}=(\alpha,\beta,\gamma,\delta)_{\alpha\delta-\beta\gamma=1}\mapsto \left(\begin{array}[]{cc}\alpha&\beta\\ \gamma&\delta\end{array}\right),\quad\left(\lambda_{1},\lambda_{0}\right) \mapsto\left(\lambda_{1},\lambda_{0}\right).$ $$

The standard conjugacy action of the subgroup $\breve{F}<G^{\mathrm{B}}$ on the normal subgroup $\breve{R}\triangleleft G^{\mathrm{B}}$ is given by $(\tilde{\lambda}_{1},\tilde{\lambda}_{0})=\left(\lambda_{1},\lambda_{0}\right)
\varrho_{1}(\alpha,\beta,\gamma,\delta)$.

###### Theorem C.2 ( Koval & Popovych ( 2023 ) for ν ≠ 1 𝜈 1 {\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\nu}\neq 1 italic_ν ≠ 1 ) .

##### Numerical results

For the numerical experiments, we consider solving Equation 5 on the domain $x\in\Omega=[0,1]$, with periodic boundary conditions on $t\in[0,1]$. An example of a canonicalized solution is shown in Figure 2, with further illustrations in Figures Figure 11 and Figure 12. The energy choice and further details are shown in Section C.4.

For Burgers’ equation we take the code and pretrained Physics-Informed DeepONet (PI-DeepONet) model weights provided by Wang et al. (). This includes a model trained on $Nf_{\mathrm{train}}=1000$ initial conditions sampled from a Gaussian random field (GRF) parameterised as $\mathcal{N}\left(0,25^{2}(-\Delta+5^{2}I)^{-4}\right)$ on the periodic domain $\Omega=[0,L]\times[0,T]=[0,1]\times[0,1]$ discretized by a meshgrid with $101\times 101$ points.

In order to sample out-of-distribution initial conditions from the GRF we apply a uniform $0.2$ shift to $u_{0}$ away from the mean of 0 0. Figure 12 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, even though this involves a reverse and reflection in the $x$ axis leading to a more accurate prediction.

##### Choice of Energy

Given the slightly different form of initial conditions to the heat equation, now sampled as GRFs we adjust the energy to focus on the bounding box of the domain and mean of $u_{0}$. The energy is therefore chosen to be the distance to the training domain. Letting $\operatorname{trunk}_{0}=(x_{0},t_{0})$ and $\operatorname{trunk}_{\Omega}=(\Omega,0)$, the energy is:

$$ $\displaystyle E_{\text{Burgers}}(x_{0},t_{0},u_{0},x_{f},t_{f})$ $\displaystyle=\operatorname{dist}\left[\operatorname{max}(\operatorname{trunk} _{0}),\operatorname{max}(\operatorname{trunk}_{\Omega})\right]$ (25) $\displaystyle\quad+\operatorname{dist}\left[\operatorname{min}(\operatorname{ trunk}_{0}),\operatorname{min}(\operatorname{trunk}_{\Omega})\right]$ $\displaystyle\quad+\operatorname{dist}\left[x_{f},\Omega\right]+\operatorname{ trunk}\left[t_{f},16\right]+\operatorname{dist}\left[\operatorname{mean}[u_{0} ],0\right]$ $$

#### C.4.1 Further burgers illustrations

Figure: Figure 11: Illustration of Burgers’ equation for in distribution GRF $u_{0}(x)$
Refer to caption: x37.png

Figure: Figure 12: Canonicalization for Burgers’ equation given out of distribution $u_{0}(x)$ and improved equivariant prediction
Refer to caption: x39.png

### C.5 Allen-Cahn equation

The Lie point symmetries of Equation 6 can be found either by hand or using the computational software of Baumann (). Both yield a 4-dimensional Lie algebra with basis vectors:

$$ $\displaystyle v_{1}=\partial_{t},\quad v_{2}=\partial_{x},\quad v_{3}=\partial _{y},\quad v_{4}=-y\partial_{x}+x\partial_{y}$ (26) $$

Noting that there is no coupling between the space, time and the dependent components, we can identify this group as the connected component of the group of isometries of the underlying space. It is worth noting here that the full symmetry group can be observed to be $\mathrm{E}(2)$, while the group generated by vectors in 26 is $\mathrm{SE}(2)$, i.e. the identity connected component of $\mathrm{E}(2)$.

##### Training data

The Poseidon (Herde et al., ) model is trained on random initial conditions(^3^33Note here that this is slightly different from the original paper. The form below was chosen, as the generated initial conditions provided by (Herde et al., ) either due to pre or post-processing turn out to not be periodic.) in the form of:

$$ $u_{0}(x,y)=\sum_{i,j=1}^{K}a_{ij}\cdot\left(i^{2}+j^{2}\right)^{-r}\sin(\pi ix )\sin(\pi jy),\quad\forall x,y\in(0,1),$ (27) $$

where $K$ is an integer drawn uniformly at random from $[16,32]$, and $r\sim U_{[0.7,1.0]}$ and $a_{ij}\sim U_{[-1,1]}$. Based on 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 of Herde et al. (), as trajectories are sampled at different points, the values on the boundary are not guaranteed to remain zero over time. As a result, we expect the model to not be equivariant under the group of symmetries of ACE, as illustrated on Figure 13, but we do not have an obvious way to construct the energy.
Based on the positive results in Section 5.2, we employ the same approach, by training a variational autoencoder and an adversarial regulariser. For the experiments, we considered the Tiny Poseidon model from Herde et al. ().

##### Testing data

To test the equivariance, we use a dataset of modified initial conditions, achieved by transformations of Equation 27.

$$ $u_{0}(x,y)=\sum_{i,j=1}^{K}a_{ij}\cdot\left(i^{2}+j^{2}\right)^{-r}\sin(\pi i \{x-x_{0}\})\sin(\pi j\{y-y_{0}\}),\quad\forall x,y\in(0,1),$ (28) $$

where $\{\cdot\}$ denotes the sawtooth function, with parameters as in Equation 27, with $x_{0},y_{0}\sim U_{[0,1]}$.

#### C.5.1 Further ACE illustrations

Figure: Figure 13: Example of failure of equivariance of the original Poseidon model.
Refer to caption: x41.png

Figure: Figure 14: Further canonicalization illustration for the ACE.
Refer to caption: x42.png

## Appendix D Measure Theory Background

Let $X$ be a set. Then a $\sigma$-algebra $\Sigma$ is a subset of $2^{X}$ such that

- •
$X\in\Sigma$
- •
If $A\in\Sigma$ then $X\backslash A\in\Sigma$
- •
For any countable collection $(A_{n})_{n=1}^{\infty}$ of sets in $\Sigma$, the union $\bigcup_{n=1}^{\infty}A_{n}\in\Sigma$

A set $A\in\Sigma$ is called measurable. A pair $(X,\Sigma)$ of a set and a $\sigma$-algebra is called a measure space. Suppose $(X,\Sigma),(Y,\Sigma^{\prime})$ are measure spaces. Then a function $f:X\rightarrow Y$ is measurable if for every measurable set $A\in\Sigma^{\prime}$ the set $f^{-1}(A)$ is measurable.

Given a measure space, a measure on it is a set function $\mu:\Sigma\rightarrow[0,\infty]$ such that

- •
$\mu(\varnothing)=0$
- •
For all countable collections $(A_{n})_{n=1}^{\infty}$ of pairwise disjoint sets in $\Sigma$
$\mu(\bigcup_{n=1}^{\infty}A_{n})=\sum_{i=1}^{\infty}\mu(A_{n})$
(29)

We may also define signed and complex measures, which simply replace the codomain of a measure with either $\mathbb{R}\cup\{-\infty,+\infty\}$ or $\mathbb{C}$ respectively.

If $f:(X,\Sigma)\rightarrow(Y,\Sigma^{\prime})$ is measurable, and $\mu$ is a measure on $(X,\Sigma)$, we define the pushforward measure $f_{*}\mu$ as follows

$$ $f_{*}\mu(A)=\mu(f^{-1}(A))$ (30) $$

There is always a measure on $(X,\Sigma)$ which simply sends every measurable set to 0 0, we call this the zero measure. A finite measure is a measure $\mu$ such that $\mu(X)<\infty$. A probability measure is a measure $\mu$ such that $\mu(X)=1$. Given a finite non-zero measure $\mu$ on $(X,\Sigma)$ we may always define a probability measure $\mu^{\prime}$ on $(X,\Sigma)$ as follows

$$ $\mu^{\prime}(A)=\frac{\mu(A)}{\mu(X)}$ (31) $$

We will often call this the normalized measure associated to $\mu$.

If $X$ is a topological space there is a canonical $\sigma$-algebra associated to it, known as the Borel $\sigma$-algebra $\mathcal{B}(X)$. It is the smallest $\sigma$-algebra such that all open sets are measurable. Then any continuous function between $X$ and $Y$ is measurable when $X$ and $Y$ are equipped with their respective Borel $\sigma$-algebras. We often call the pair $(X,\mathcal{B}(X))$ a Borel space and when clear we will simply refer to this pair by $X$.

For a topological space $S$ and a real number $s\geq 0$ there is an associated Hausdorff measure of dimension $s$ which we shall denote $\mathcal{H}^{s}_{S}$. To $S$ we can also define what is known as the Hausdorff dimension, $\dim_{\mathcal{H}}(S)$, which is always non-negative. There is then a Hausdorff measure of that dimension, which we write as $\mathcal{H}_{S}$. All Borel sets are measurable for this measure. If $S$ is compact then $\mathcal{H}_{S}(S)<\infty$, but it 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 a Radon measure. For $S\cong[0,1]^{n}$ 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 0 and the normalized Hausdorff measure is the normalized counting measure.

The Hausdorff measure is important to consider because of Frostman’s Lemma (Mattila, ).

###### Lemma D.1 (Frostman’s Lemma) .

Let $A$ be a Borel subset of $\mathbb{R}^{n}$ and let $s>0$. Then the following are equivalent.

- •
$\mathcal{H}^{s}(A)>0$
- •
There is an unsigned Borel measure $\mu$ on $\mathbb{R}^{n}$ satisfying $\mu(A)>0$ and such that
$\mu(B(x,r))\leq r^{s}$
(32)
for all $x\in\mathbb{R}^{n},r>0$.

Let $A$ be compact with Hausdorff dimension greater than zero. Then we 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.

###### Lemma D.1 (Frostman’s Lemma) .

## Appendix E More on Frames and Canonicalization

Recall that a frame is defined as a function $\mathcal{F}:X\rightarrow 2^{G}\backslash\varnothing$ such that for $x\in X$ and $g\in G$

$$ $\mathcal{F}(gx)=g\mathcal{F}(x)$ (33) $$

###### Lemma E.1 .

Let $\mathcal{F}$ be frame for the action of a group $G$ on a set $X$. Then $\forall x\in X$, $\mathcal{F}(x)$ is the union of left cosets of $G_{x}\subseteq G$.

###### Proof.

Indeed, let $x\in X$ and let $g\in G_{x}$. Then

$$ $\mathcal{F}(x)=\mathcal{F}(gx)=g\mathcal{F}(x)$ (34) $$

Hence we see that

$$ $\mathcal{F}(x)=G_{x}\mathcal{F}(x)$ (35) $$

Therefore we see that $\mathcal{F}(x)$ is the union of left cosets of $G_{x}$.
∎

###### Lemma E.1 .

###### Proof.

### E.1 Weakly-Equivariant Frames

In Dym et al. () weakly equivariant frames are introduced. These are maps $\mu:X\rightarrow\operatorname{PMeas}(G)$ satisfying equivariance up to averaging over stabilizers. First, for any $x\in X$ and Radon probability measure $\tau$ they define for every measurable $A\subseteq G$,

$$ $\langle\tau\rangle_{x}(A)=\int_{s\in G_{x}}s_{*}\tau(A)ds$ (36) $$

where $ds$ is the normalized Haar measure on $G_{x}$. For a frame $\mu$ we let $\overline{\mu}_{x}=\langle\mu_{x}\rangle_{x}$. Then $\mu$ is weakly equivariant if for all $x\in X$ and $g\in G$,

$$ $\overline{\mu}_{gx}=g_{*}\overline{\mu}_{x}$ (37) $$

This definition does not extend to the case where the action has non-compact stabilizers.

To deal with this we propose imposing what we call $\pi_{*}$-equivariance. Again we consider simply a map $\mu:X\rightarrow\operatorname{PMeas}(G)$. For a point $x\in X$, let $\pi_{x}:G\rightarrow G/G_{x}$ be the quotient. We require that

$$ $(\pi_{gx})_{*}\mu_{gx}=(\pi_{gx})_{*}g_{*}\mu_{x}$ (38) $$

###### Lemma E.2 .

If the stabilizer $G_{x}$ is compact then $\langle\mu\rangle_{x}=\langle\mu^{\prime}\rangle_{x}$ if and only if $(\pi_{x})_{*}\mu=(\pi_{x})_{*}\mu^{\prime}$.

###### Proof.

For ease of notation let $\pi=\pi_{x}$. First suppose that $\langle\mu\rangle_{x}=\langle\mu\rangle_{x}$. Let $A\subseteq G/G_{x}$ be a measurable set. Note that $\pi^{-1}(A)$ is a union of cosets of $G_{x}$. Hence it is invariant under multiplication by elements in $G_{x}$. Therefore for all $s\in G_{x}$ we have $s^{-1}\pi^{-1}(A)=\pi^{-1}(A)$. Thus

$$ $\langle\mu\rangle_{x}(A)=\int_{s\in G_{x}}s_{*}\mu(\pi^{-1}(A))ds=\int_{s\in G _{x}}\mu(\pi^{-1}(A))ds=\mu(\pi^{-1}(A))$ (39) $$

and similarly, $\langle\mu^{\prime}\rangle_{x}(\pi^{-1}(A))=\mu^{\prime}(\pi^{-1}(A))$. Hence,

$$ $\mu(\pi^{-1}(A))=\mu^{\prime}(\pi^{-1}(A))$ (40) $$

Therefore, $\pi_{*}\mu=\pi_{*}\mu^{\prime}$.

Now suppose that $\pi_{*}\mu=\pi_{*}\mu^{\prime}$. Then for any measurable set $A\subseteq G$, we have that

$$ $\mu(\pi^{-1}(\pi(A)))=\pi_{*}\mu(\pi(A))=\pi_{*}\mu^{\prime}(\pi(A))=\mu^{ \prime}(\pi^{-1}(\pi(A))$ (41) $$

We have

$$ $\pi^{-1}(\pi(A))=\bigcup_{s\in G_{x}}sA$ (42) $$

Using Fubini’s theorem and the fact that both $s$ and $\mu$ are probability measure, we may write Equation 36 as follows

$$ $\displaystyle\langle\mu\rangle_{x}$ $\displaystyle=\int_{s\in G_{x}}\,s_{*}\mu(A)ds$ (43) $\displaystyle=\int_{s\in G_{x}}\int_{G}\chi_{A}(sg)\,d\mu(g)ds$ (44) $\displaystyle=\int_{G}\int_{s\in G_{x}}\chi_{A}(sg)\,dsd\mu(g)$ (45) $$

Consider the inner integral as a function $\phi$ on $G$. We claim that $\phi$ is $G_{x}$-invariant. Indeed, let $g^{\prime}=s^{\prime}g$ for some $s^{\prime}\in G_{x}$. Then

$$ $\phi(g^{\prime})=s(Ag^{-1}(s^{\prime})^{-1})=s(Ag^{-1})=\phi(g)$ (47) $$

as $s$ is the normalized Haar measure, hence is right-invariant. Therefore $\phi$ factors through the quotient map $\pi$. Let $\phi^{\prime}$ be the corresponding function on $G/G_{x}$, then $\phi=\phi^{\prime}\circ\pi$. Hence

$$ $\displaystyle\langle\mu\rangle_{x}$ $\displaystyle=\int_{G}\phi(g)d\mu(g)$ (48) $\displaystyle=\int_{G/G_{x}}\phi^{\prime}([g])d(\pi_{*}\mu)([g])$ (49) $\displaystyle=\int_{G/G_{x}}\phi^{\prime}([g])d(\pi_{*}\mu^{\prime})([g])$ (50) $\displaystyle=\int_{G}\phi(g)d\mu^{\prime}(g)=\langle\mu^{\prime}\rangle_{x}$ (51) $$

∎

###### Corollary E.2.1 .

If the $G$-action on $X$ has compact stabilizers at each point, then $\pi_{*}$-equivariance is equivalent to weak equivariance.

###### Lemma E.2 .

###### Proof.

###### Corollary E.2.1 .

### E.2 On projections

For a weighted orbit canonicalization $\mu$, $\mathcal{P}_{\mu}:B(X,\mathbb{R})\rightarrow B(X,\mathbb{R})^{G}$ turns out to be a projection, as

$$ $\displaystyle\mathcal{P}_{\mu}\mathcal{P}_{\mu}(\phi)(x)$ $\displaystyle=\int_{X}\mathcal{P}_{\mu}(\phi)(s)\,d\mu_{x}(s)$ $\displaystyle=\int_{Gx}\underbrace{\mathcal{P}_{\mu}(\phi)(s)}_{\text{G- invariant}}\,d\mu_{x}(s)=\mathcal{P}_{\mu}(\phi)(x)$ (52) $$

### E.3 Energies and Closed Canonicalizations

By the construction of energy minimizing frames, we have a function from energy functions to weighted closed canonicalizations. We now describe the image of this.

###### Theorem E.3 .

Let $\mu\in\operatorname{WCCan}_{G}(X)$ satisfying the following:

- •
$\mu_{x}$ is the normalized Hausdorff measure on its support.
- •
For all $y\in\overline{Gy}$, $\textnormal{supp }\mu_{y}=\textnormal{supp }\mu_{x}\cap\overline{Gx}$.

Then there is an energy function $E$ such that $\mathcal{C}_{E}=\mu$.

###### Proof.

Suppose $\mu$ is a closed canonicalization satisfying the requirements of the theorem. Let $E_{\mu}(x)=1-\chi_{\operatorname{supp}\mu_{x}}(x)$.

Now we consider the minimum sets

$$ $\displaystyle\mathcal{M}_{E_{\mu}}(x)=\min_{y\in\overline{Gx}}E_{\mu}(y)$ (53) $$

We have that $\textnormal{supp }\mu_{x}$ is a non-empty closed subset of $\overline{Gx}$, hence the minimum of $E_{\mu}$ over $\overline{Gx}$ is zero. Therefore $\mathcal{M}_{E_{\mu}}(x)=E_{\mu}^{-1}(0)\cap\overline{Gx}$. In other words it is those $y\in\overline{Gx}$ such that $y\in\textnormal{supp }\mu_{y}$. Further, by assumption, for any $y\in\overline{Gx}$,

$$ $\textnormal{supp }\mu_{y}=\textnormal{supp }\mu_{x}\cap\overline{Gy}$ (54) $$

Therefore $y\in\textnormal{supp }\mu_{y}$ if and only if $y\in\textnormal{supp }\mu_{x}$. Hence we see that

$$ $\mathcal{M}_{E_{\mu}}(x)=\textnormal{supp }\mu_{x}$ (55) $$

But the by the assumption that $\mu_{x}$ is simply the normalized Hausdorff measure on $\textnormal{supp }\mu_{x}$, we get that the $\mathcal{C}_{E_{\mu}}=\mu$.
∎

###### Theorem E.3 .

###### Proof.

## Appendix F Proofs

###### Proposition F.1 .

$\mathcal{F}_{E}$ defines an *equivariant frame*, i.e. $\mathcal{F}_{E}(gx)=g\mathcal{F}_{E}(x)$ for any $g\in G,\;x\in X$.

###### Proof.

Suppose that $x\in\mathcal{X}$ and $g\in G$. Then

$$ $\displaystyle\mathcal{F}_{E}(gx)$ $\displaystyle=\operatorname*{arg\,min}_{h\in G}E(h^{-1}gx)$ $\displaystyle=\operatorname*{arg\,min}_{gh\,\text{s.t.}\,h\in G}E((gh)^{-1}gx)$ $\displaystyle=\operatorname*{arg\,min}_{gh\,\text{s.t.}\,h\in G}E(h^{-1}x)$ $\displaystyle=g[\operatorname*{arg\,min}_{h\in G}E(h^{-1}x)]=g\mathcal{F}(x).$ $$

Here, we have simply used that group multiplication is bijective when either of the arguments is fixed.
∎

###### Proposition F.2 .

$\mathcal{M}_{E}(x)$ is $G$-invariant and non-empty.

###### Proof.

$\mathcal{M}_{E}(x)$ is $G$-invariant, as the infimum is taken over the orbit $\overline{Gx}$, and this orbit is $G$-invariant. By the coercivity assumption on $E$, this is non-empty. Indeed, let $x\in X$ and let $N$ be large. Let $K$ be a compact set of $X$ such that $E(y)>N$ for $y\in X\backslash K$. Taking $N$ large enough, we can assume that $K\cap\overline{Gx}$ is non-empty. Then we see that $\mathcal{M}_{E}(x)\subseteq K\cap\overline{Gx}$. $K\cap\overline{Gx}$ is a closed subset of the compact set $K$, hence is compact. Therefore $E$ attains its minimum on this set, hence $\mathcal{M}_{E}(x)$ is non-empty.
∎

###### Theorem F.3 .

Let $G$ be a Lie group acting smoothly on $X$ a manifold. Let $j_{x}:G\rightarrow X$ be defined as $g\mapsto g^{-1}\cdot x$. Then we have a map $\alpha:\operatorname{WFra}_{G}(X)\rightarrow\operatorname{WOCan}_{G}(X)$ defined by

$$ $\alpha(\mathcal{F})(x)=(j_{x})_{*}(\mathcal{F}(x))$ (56) $$

In addition, for any $\mathcal{F},\mathcal{F^{\prime}}\in\operatorname{WFra}_{G}(X)$ such that $\alpha(\mathcal{F})=\alpha(\mathcal{F^{\prime}})$ we have that $(\pi_{x})_{*}\mathcal{F}=(\pi_{x})_{*}\mathcal{F^{\prime}}$ for every $x\in X$ where $\pi_{x}:G\rightarrow G/G_{x}$.
If each $G_{x}$ has a probability measure $P_{x}$, then $\alpha$ has image exactly the set of weighted orbit canonicalizations with compact support.

###### Proof.

First we check that $\alpha$ as defined above lands in $\operatorname{WOCan}_{G}(X)$. Let $\mathcal{F}\in\operatorname{WFra}_{G}(X)$. Then $\alpha(\mathcal{F})(x)$ is $G$-invariant as for any measurable $A\subseteq X$.

$$ $\displaystyle\alpha(\mathcal{F})(gx)(A)$ $\displaystyle=(j_{gx})_{*}(\mathcal{F}(gx))(A)$ $\displaystyle=\mathcal{F}(gx)(\{h\in G\;|\;h^{-1}\cdot(gx)\in A\}$ $\displaystyle=\mathcal{F}(gx)(\{h\in G\;|\;h^{-1}\cdot(gx)\in A\})$ $\displaystyle=\pi_{*}(\mathcal{F}(gx))(\pi(\{h\in G\;|\;h^{-1}\cdot(gx)\in A\}))$ $\displaystyle=\pi_{*}g_{*}(\mathcal{F}(x))(\pi(\{h\in G\;|\;h^{-1}\cdot(gx)\in A \}))$ $\displaystyle=g_{*}(\mathcal{F}(x))(\{h\in G\;|\;h^{-1}\cdot(gx)\in A\})$ $\displaystyle=\mathcal{F}(x)(\{h\in G\;|\;(gh)^{-1}\cdot(gx)\in A\})$ $\displaystyle=\mathcal{F}(x)(\{h\in G\;|\;h^{-1}\cdot x\in A\})=\alpha( \mathcal{F})(x)(A)$ $$

Thus we see that $\alpha(\mathcal{F})\in\operatorname{WCan}_{G}(X)$, as $\alpha(\mathcal{F})(x)(X)=\mathcal{F}(x)(G)=1$, as $\mathcal{F}(x)$ is a probability measure. To check that it is a weighted orbit canonicalization, first note that it is clear that $\textnormal{supp }\alpha(\mathcal{F})(x)\subseteq\overline{Gx}$, as for any $y\notin\overline{Gx}$ we can find an open neighborhood $U$ of $y$ such that $j_{x}^{-1}(U)=\varnothing$ which is always $\mathcal{F}(x)$-measure zero. Suppose that $y\in\overline{Gx}\backslash Gx$. $\mathcal{F}(x)$ is a weighted frame, hence it is supported on some compact set $K\subseteq G$. We claim that there is an open set $U$ containing $y$ such that $j_{x}^{-1}(U)\subseteq G\backslash K$. Suppose not, then for every open neighborhood $U$ of $y$ there is an $g_{U}\in K$ such that $j_{x}(g)\in U$. Then the collection $(g_{U})$ over the neighborhoods of $y$ forms a net in $K$. As $K$ is compact, this net must have a convergent subsequence. Call this subsequence $(g_{n})$ and let $g$ be the limit. As the action of $G$ on $X$ is continuous we have

$$ $g\cdot x=\lim_{n\rightarrow\infty}g_{n}\cdot x$ (57) $$

For any $U$ an open neighborhood of $y$ we may, by definition of a net, find an $N$ such that for any $n\geq N$, $g_{n}\cdot x\in U$. As $X$ is Hausdorff, we then must have that the limit $\lim g_{n}\cdot x$ is simply $y$. But this contradicts $y\notin Gx$. Hence there is some open neighborhood $U$ of $y$ such that $j_{x}^{-1}(U)\subseteq G\backslash K$. Then $j_{x}^{-1}(U)$ is measure zero with respect to $\mathcal{F}(x)$. Hence $\alpha(\mathcal{F})(x)(U)=0$, so $y\notin\textnormal{supp }\alpha(\mathcal{F})$. Therefore we see that $\alpha(\mathcal{F})\in\operatorname{WOCan}_{G}(X)$.

Suppose that $\mathcal{F},\mathcal{F}^{\prime}$ are weighted frames such that $\alpha(\mathcal{F})=\alpha(\mathcal{F}^{\prime})$. Let $x\in X$. Then for any measurable $A\subseteq X$,

$$ $\mathcal{F}(x)(j_{x}^{-1}(A))=\alpha(\mathcal{F})(x)(A)=\alpha(\mathcal{F}^{ \prime})(x)(A)=\mathcal{F^{\prime}}(x)(j_{x}^{-1}(A))$ (58) $$

Both $\mathcal{F}$ and $\mathcal{F}^{\prime}$ are compactly supported, say on compact subsets $K,K^{\prime}$ respectively.
Consider the natural map $\overline{j_{x}}:G/G_{x}\rightarrow X$. This is injective for group-theoretic reasons. Let $\pi:G\rightarrow G/G_{x}$ be the quotient map. Then $\pi(K)$ and $\pi(K^{\prime})$ are compact subsets of $G/G_{x}$, as they are continuous images of compacts. Now as $\pi(K)$ and $\pi(K)^{\prime}$ are compacts which inject into a Hausdorff space $X$, $\overline{j_{x}}$ restricts to homeomorphisms $\pi(K)\cong\overline{j_{x}}(\pi(K))$ and similar for $K^{\prime}$. Using Equation 58 we see that for any $U\subseteq G/G_{x}$ open we have

$$ $\mathcal{F}(x)(\pi^{-1}(U))=\mathcal{F}(x)(\pi^{-1}(U))$ (59) $$

Therefore $\pi_{*}\mathcal{F}$ and $\pi_{*}\mathcal{F^{\prime}}$ agree on the open sets of $G/G_{x}$, but then that implies that they agree on the entire Borel $\sigma$-algebra of $G/G_{x}$. Hence $\pi_{*}\mathcal{F}=\pi_{*}\mathcal{F^{\prime}}$ as desired.

Finally, let $v\in\operatorname{WOCan}_{G}(X)$ have compact support. Let $x\in X$. Then $v_{x}$ is some probability measure on $X$ with $\textnormal{supp }v_{x}\subseteq K\subseteq Gx$ for some compact $K\subseteq Gx$. By Theorem 2.1 we have that $\overline{j_{x}}:G/G_{x}\rightarrow X$ is a immersion with image $Gx$. $\overline{j_{x}}$ restricts to a homeomorphism on $K$, and using this we can define a measure $v_{x}^{\prime}$ on $G/G_{x}$ which agrees with $v_{x}$ under this homeomorphism. It is supported on the preimage $K^{\prime}$ of $K$ which is also compact. $\pi:G\rightarrow G/G_{x}$ is a principal bundle with fiber $G_{x}$. Therefore, locally on $G/G_{x}$, $\pi$ is the trivial fibration. We may cover $K^{\prime}$ by finitely many open sets over which $\pi$ is trivial, call them $U_{1},\dots,U_{n}$. Let $V_{i}=\pi^{-1}(U_{i})\cong U_{i}\times G_{x}$. Now define

$$ $V_{i}^{\prime}=V_{i}\backslash(\bigcup_{j<i}V_{j})$ (60) $$

Then these are disjoint measurable sets which cover $\pi^{-1}(K^{\prime})$. They are all trivially fibered over the base. Now we define a measure $\nu$ using these, $\nu$ will be supported on $\pi^{-1}(K^{\prime})$. Let $A$ be a measurable set in $G$. Firstly we may assume that $A\subseteq\pi^{-1}(K^{\prime})$. Let $A_{i}=A\cap V_{i}^{\prime}$. Then the $A_{i}$ are disjoint measurable sets which cover $A$. So the value of $\nu$ on $A$ is determined by the value of $\nu$ on each $A_{i}$. We define this value to be

$$ $\nu(A_{i})=(v_{x}^{\prime}\times P_{x})(A_{i})$ (61) $$

where here we use the homeomorphism $V_{i}^{\prime}\cong\pi(V_{i}^{\prime})\times G_{x}$ to put the product measure on this set. $\pi_{*}\nu=v_{x}^{\prime}$, and this is equivariant as $G_{gx}=g^{-1}G_{x}$. Hence $\nu$ is $\pi_{*}$-equivariant.
∎

Note that for the above theorem to work, we need to assume that each $G_{x}$ has a probability measure. This is automatic in the case that each $G_{x}$ is compact, as in that case we may take that normalized Haar measure on it. It will also hold in the case that each $G_{x}$ has non-zero Hausdorff measure.

###### Theorem F.4 .

The sequential closure of $\operatorname{WOCan}_{G}(X)$ is $\operatorname{WCCan}_{G}(X)$

###### Proof.

First we show $\operatorname{WCCan}_{G}(X)$ is sequentially closed. Suppose that $(f^{n})$ is a sequence in $\operatorname{WCCan}_{G}(X)$ converging to $f\in\operatorname{WCan}_{G}(X)$. Fix $x\in X$ and $y\in X\backslash\overline{Gx}$. Then it is sufficient to show that $y\notin\operatorname{supp}f_{x}$. As $X$ is normal Hausdorff, by Urysohn’s lemma there exists continuous $\phi:X\rightarrow[0,1]$ such that $\phi(\overline{Gx})=1$ and $\phi(\{y\})=0$.

As $f^{n}\rightarrow f$, for any $\varepsilon>0$, we can find $N$ such that for any $n\geq N$ we have the following

$$ $\|f_{n}(\phi)-f(\phi)\|<\varepsilon$ (62) $$

We may easily compute $f_{n}(\phi)(x)$ by the fact that $\operatorname{supp}f_{n}(x)\subseteq\overline{Gx}$.

$$ $f_{n}(\phi)(x)=\int_{\overline{Gx}}\phi\,df^{n}_{x}=1$ (63) $$

Similarly,

$$ $f_{n}(\phi)(y)=\int_{\overline{Gy}}\phi\,df^{n}_{y}=0$ (64) $$

Hence, by taking the limit we compute that $f(\phi)(x)=1$ and $f(\phi)(y)=0$. Let $U=\phi^{-1}([0,1/2])$, this is an open neighborhood of $y$.

$$ $\displaystyle 1=f(\phi)(x)$ $\displaystyle=\int_{X}\phi\,df_{x}$ (65) $\displaystyle=\int_{X\backslash U}\phi\,df_{x}+\int_{U}\phi\,df_{x}$ (66) $\displaystyle\leq f_{x}(X\backslash U)+\frac{1}{2}f_{x}(U)$ (67) $\displaystyle=1-f_{x}(U)+\frac{1}{2}f_{x}(U)$ (68) $\displaystyle=1-\frac{1}{2}f_{x}(U)$ (69) $$

Hence, $f_{x}(U)=0$. Therefore $y\notin\text{supp}f_{x}$.

Clearly, $\operatorname{WOCan}_{G}(X)\subseteq\operatorname{WCCan}_{G}(X)$, hence the sequential closure of $\operatorname{WOCan}_{G}(X)$ is contained inside $\operatorname{WCCan}_{G}(X)$. We now show the reverse inclusion. Fix an arbitrary metric on $X$, this is possible as $X$ is assumed metrizable.

Suppose we have a weighted closed canonicalization $f$. Let $h$ be the left invariant Haar measure on $G$. Let $\mu_{x}=(i_{x})_{*}h$, this is a measure supported on $\overline{Gx}$. Further it is a $\sigma$-finite measure as $G$ is locally compact. Therefore we can use the Lebesgue decomposition theorem to decompose $f(x)$ as

$$ $f_{x}=f^{\mathrm{ac}}_{x}+f^{\mathrm{sc}}_{x}+f^{\mathrm{sd}}_{x}$ (70) $$

where $f^{\mathrm{ac}}_{x}$ is absolutely continuous with respect to $\mu_{x}$, $f^{\mathrm{sc}}_{x}$ is singular with respect to $\mu_{x}$ and continuous, and $f^{\mathrm{sd}}_{X}$ is singular with respect to $\mu_{x}$ and discrete.

First we construct a sequence of weighted closed canonicalizations $f^{n}$ converging to $f$ such that upon decomposing each $f^{n}$ as above, there is no singular continuous factor.

$f^{\mathrm{sc}}$ is singular continuous with respect to $\mu_{x}$, hence there is a measurable set $S\subseteq\overline{Gx}$ such that $\mu_{x}(S)=0$ and $f^{\mathrm{sc}}_{x}(S^{c})=0$. Fix a point $y\in S$. For $l>0$, let $K_{l}=\overline{B(y,l)}\cap\overline{S}$. Then we have

$$ $\bigcup_{l}K_{l}=\overline{S}$ (71) $$

By continuity of measure

$$ $\lim_{l\rightarrow\infty}f^{\mathrm{sc}}_{x}(K_{l})=f^{\mathrm{sc}}_{x}( \overline{S})=f^{\mathrm{sc}}_{x}(S)+f^{\mathrm{sc}}_{x}(\overline{S}\cap S^{c })=f^{\mathrm{sc}}_{x}(S)$ (72) $$

Fix $l$ large enough such that $f^{\mathrm{sc}}_{x}(K_{l})>f_{x}^{\mathrm{sc}}-\frac{1}{n}$. $K_{l}$ is compact so we may find finitely many points $x_{1},\dots,x_{p(n)}$ such that the $\frac{1}{n}$ balls around these points cover $K_{l}$. Call these open balls $U_{i}$. Let $U_{i}^{\prime}=U_{i}\backslash\bigcup_{j<i}U_{j}$. Then each $U_{i}^{\prime}$ contains only $x_{i}$ and none of the other $x_{j}$. Define a discrete measure $f_{n}^{\prime}$ which is non-zero only on the $x_{i}$ and taking the following values at those points

$$ $\hat{f}^{n}(\{x_{i}\})=f^{\mathrm{sc}}_{x}(U_{i}^{\prime})$ (73) $$

Then one easily sees that

$$ $\hat{f}^{n}(S)=f^{\mathrm{sc}}_{x}(U_{1}^{\prime})+f^{\mathrm{sc}}_{x}(U_{2}^{ \prime})+\dots+f^{\mathrm{sc}}_{x}(U_{p(n)}^{\prime})=f^{\mathrm{sc}}_{x}(K_{l})$ (74) $$

Now let $f^{n}=f^{\mathrm{ac}}_{x}+\hat{f}^{n}+f^{\mathrm{sd}}_{x}$. This is then a measure with no singular discrete part with respect to $\mu_{x}$. We now check that $f^{n}\rightarrow f$. Let $\phi\in C^{0}_{b}(X)$. We may assume that $\phi$ is non-negative on $X$ by splitting it up as $\phi=\phi^{+}+\phi^{-}$ for $\phi^{+}$ non-negative and $\phi^{-}$ non-positive.

$$ $\displaystyle f_{n}(\phi)(x)-f(\phi)(x)$ $\displaystyle=\int\phi\,d(f_{n}(x)-f(x))$ $\displaystyle=\int\phi\,d(f^{\mathrm{sc}}(x)-f_{n}^{\prime}(x))$ $\displaystyle=\sum_{i=1}^{p}(\int_{U^{\prime}_{i}}\phi\,df^{\mathrm{sc}}(x)- \phi(x_{i})f_{n}^{\prime}(U^{\prime}_{i}))+\int_{S\backslash K_{n}}\phi\,df^{ \mathrm{sc}}(x)$ $$

Therefore we see that, with $C=\max\phi$,

$$ $|f_{n}(\phi)(x)-f(\phi)(x)|\leq\max(\sum_{i=1}^{p(n)}f^{\mathrm{sc}}(U^{\prime }_{i})(\max_{U^{\prime}_{i}}\phi-\phi(x_{i})),\sum_{i=1}^{p}f^{\mathrm{sc}}(U^ {\prime}_{i})(\phi(x_{i})-\min_{U^{\prime}_{i}}\phi))+\frac{C}{n}$ (75) $$

Let $\epsilon>0$. Take $N$ large enough such that for all $n>N$, and $i=1,\dots p$

$$ $|\max_{U_{i}^{\prime}}\phi-\min_{U_{i}^{\prime}}\phi|<\epsilon$ (76) $$

Recall that we can do this as there are finitely many $U_{i}^{\prime}$, each of which is contained in the $\frac{1}{n}$ ball around $x_{i}$. Plugging this into Equation 75 we get

$$ $|f_{n}(\phi)(x)-f(\phi)(x)|\leq\frac{C}{n}+\sum_{i=1}^{p(n)}f^{\mathrm{sc}}(U_ {i}^{\prime})\epsilon=\frac{C}{n}+\epsilon f^{\mathrm{sc}}(K_{n})\leq\frac{C}{ n}+\epsilon f^{\mathrm{sc}}(S)$ (77) $$

Taking $n$ large and $\epsilon$ 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_{x}=f^{\mathrm{ac}}_{x}+f^{\mathrm{sd}}_{x}$ with respect to $\mu_{x}$. For any $n$, choose $V_{n}(x)$ to be the $\frac{1}{n}$ open ball around $\overline{Gx}\backslash Gx$. $f^{\mathrm{sd}}$ is singular and discrete, hence it is supported on a (possibly infinite) set of points $x_{1},x_{2},\dots$ in $\overline{Gx}$. For each $x_{i}$ choose a sequence of points $(x_{in})$ in $Gx$ converging to $x_{i}$. We may choose these such that for a fixed $n$, each of the $x_{in}$ are different for every $i$. Further we assume that these are chosen such that for every $i$

$$ $d(x_{i},x_{in})<\frac{1}{n}$ (78) $$

Then define the singular measure $\hat{f}^{n}$ which is non-zero only on the $x_{in}$ for every $i$, taking value

$$ $\hat{f}^{n}(\{x_{in}\})=f^{\mathrm{sd}}_{x}(\{x_{i}\})$ (79) $$

Define the absolutely continuous measure $\tilde{f}^{n}$ by

$$ $\tilde{f}^{n}(U)=f^{\mathrm{ac}}(U\backslash V_{n}(x)).$ (80) $$

Then let

$$ $f^{n}(x)=\hat{f}^{n}+\tilde{f}_{n}$ (81) $$

This is then a measure supported on $Gx$ and by construction it is $G$-invariant. Hence this is a weighted orbit canonicalization. Now we show that $f^{n}\rightarrow f$. Let $\phi\in C^{0}_{b}(X)$ and $x\in X$. Then

$$ $\displaystyle|f^{n}(\phi)(x)-f(\phi)(x)|$ $\displaystyle=|\int_{\overline{Gx}}\phi\,df_{x}-\int_{Gx}\phi\,df^{n}_{x}|$ (82) $\displaystyle=|\int_{\overline{Gx}}\phi\,df^{\mathrm{ac}}_{x}-\int_{Gx}\phi\,d \tilde{f}^{n}_{x}|+|\int\phi\,df^{\mathrm{sd}}_{x}-\int\phi\,d\hat{f}_{x}^{n}|$ (83) $\displaystyle=|\int_{V_{n}(x)}\phi\,df^{\mathrm{ac}}_{x}|+|\sum_{i}f^{\mathrm{ sd}}_{x}(\{x_{i}\})(\phi(x_{i})-\phi(x_{in})|$ (84) $\displaystyle\leq Cf^{\mathrm{ac}}_{x}(V_{n}(x))+\sum_{i}f^{\mathrm{sd}}_{x}( \{x_{i}\})|\phi(x_{i})-\phi(x_{in})|$ (85) $$

Now fix $\epsilon>0$. As $f^{\mathrm{sd}}_{x}$ is a finite measure, there is some $I$ such that

$$ $\sum_{i=I}^{\infty}f^{\mathrm{sd}}_{x}(\{x_{i}\})<\epsilon$ (86) $$

Thus we get

$$ $\displaystyle|f^{n}(\phi)(x)-f(\phi)(x)|$ $\displaystyle\leq Cf^{\mathrm{ac}}_{x}(V_{n}(x))+(\sum_{i=1}^{I-1}+\sum_{i=I}^ {\infty})f^{\mathrm{sd}}_{x}(\{x_{i}\})|\phi(x_{i})-\phi(x_{in})|$ (87) $\displaystyle\leq Cf^{\mathrm{ac}}_{x}(V_{n}(x))+\sum_{i=1}^{I-1}f^{\mathrm{sd }}_{x}(\{x_{i}\})|\phi(x_{i})-\phi(x_{in})|+2C\epsilon$ (88) $$

Now choose $n$ large enough such that $f^{\mathrm{ac}}(V_{n}(x))<\epsilon$ and for each $i=1,\dots,I-1$, $|\phi(x_{i})-\phi(x_{in})|<\epsilon$. Then we get a bound

$$ $|f^{n}(\phi)(x)-f(\phi)(x)|\leq 3C\epsilon+D\epsilon$ (89) $$

Hence, $f^{n}\rightarrow f$.
∎

###### Theorem F.5 .

$\operatorname{WCan}_{G}(X)^{\mathrm{cont}}$ is sequentially closed.

###### Proof.

Suppose that $(\mu^{n})_{n}$ is a convergent sequence of continuous weighted canonicalizations, converging to a weighted canonicalization $\mu$. We will show that $\mu$ is continuous. To do this, fix a convergent sequence $(x_{m})_{m}$ in $X$ converging to $x$. Let $\phi\in C^{0}_{b}(X)$. Then we have the following bound

$$ $|\mu(\phi)(x)-\mu(\phi)(x_{m})|\leq|\mu(\phi)(x)-\mu_{n}(\phi)(x)|+|\mu_{n}( \phi)(x)-\mu_{n}(\phi)(x_{m})|+|\mu_{n}(\phi)(x_{m})-\mu(\phi)(x_{m})|$ (90) $$

Fix $\epsilon>0$. As $\mu_{n}$ converges to $\mu$ we may find a large enough $n$ such that both the first and third terms are less than $\epsilon$. For this $n$, as $\mu_{n}$ is continuous, we may find $m$ large enough such that the middle term is less than $\epsilon$ as well. Hence the right hand side is bounded by $3\epsilon$, so we see that $\mu$ is continuous.
∎

###### Proposition F.6 .

Let $Y$ be a topological space. Let $f:Y\rightarrow[0,+\infty]$ be lower semi-continuous. Then the set

$$ $M=\{y\in Y|f(y)=\inf_{Y}f\}$ (91) $$

is closed.

###### Proof.

Let $\alpha$ be the infimum of $f$. We have that $M=f^{-1}(\{\alpha\})$. As $\alpha$ is the infimum over $Y$, we see that in fact $M=f^{-1}([0,\alpha])$. Then as $f$ is lower semi-continuous, this set is closed in $Y$.
∎

###### Lemma F.7 .

Let $E$ be a lower semi-continuous energy function. Then $\mathcal{C}_{E}$ is continuous if and only if for every convergent sequence $x_{n}\rightarrow x$ we have $\textnormal{Li }\mathcal{M}_{E}(x_{n})=\mathcal{M}_{E}(x)$.

###### Proof.

Let $E:X\rightarrow[0,+\infty]$ be an energy function such that the corresponding energy minimizing canonicalization $\mathcal{C}_{E}$ is continuous. Consider a convergent sequence $(x_{n})$ in $X$, converging say to $x$. $\mathcal{M}_{E}(x)$ is closed inside $\overline{Gx}$, being the set of minimizers of a lower semi-continuous function. As $\overline{Gx}$ is closed in $X$, $\mathcal{M}_{E}$ is closed in $X$ as well. As $X$ is metrizable, we may choose a countable neighborhood basis of $\mathcal{M}_{E}(x)$, which we shall write as a decreasing sequence $(U_{m})$ of open neighborhoods such that

$$ $\bigcap U_{m}=\mathcal{M}_{E}(x)$ (92) $$

Let $C_{m}=X\backslash U_{m}$, this is closed. As $X$ is normal Hausdorff, we may find $\phi_{m}:X\rightarrow[0,1]$ continuous such that $\phi_{m}^{-1}(\{0\})=C_{n}$ and $\phi_{m}^{-1}(\{1\})=\mathcal{M}_{E}(x)$.

Now as $\mathcal{C}_{E}$ is continuous we have that the measures $\mathcal{C}_{E}(x_{n})$ weakly converge to $\mathcal{C}_{E}(x)$. Hence for every $m$ we have that

$$ $\lim_{n\rightarrow\infty}\int_{\mathcal{M}_{E}(x_{n})}\phi_{m}\,d\mu_{x_{n}} \rightarrow\int_{\mathcal{M}_{E}(x)}\phi_{m}\,d\mu_{x}=1$ (93) $$

Hence

$$ $\lim_{n\rightarrow\infty}\int_{\mathcal{M}_{E}(x_{n})}\phi_{n}\,d\mu_{x_{n}}=1$ (94) $$

Let $\epsilon>0$. We may find $N$ such that for all $n>N$, the left hand side is greater than $1-\epsilon$. As $\mu_{x_{n}}(\mathcal{M}_{E}(x_{n}))=1$, we must have that the minimum of $\phi_{n}$ over $\mathcal{M}_{E}(x_{n})$ is greater than $1-\epsilon$. As $C_{n}=\phi_{n}^{-1}(\{0\})$, we then must have that $\mathcal{M}_{E}(x_{n})\subseteq U_{n}$. Hence by the definition of Kuratowski convergence we have $\textnormal{Li }\mathcal{M}_{E}(x_{n})\subseteq\mathcal{M}_{E}(x)$.

Now suppose that $y\in\mathcal{M}_{E}(x)$. Then we may as above choose a decreasing sequence $(U_{m})$ of open neighborhoods of $y$ which form a neighborhood basis. Then fix for each $m$, $\phi_{m}:X\rightarrow[0,1]$ continuous such that $\phi_{m}^{-1}(\{0\})=X\backslash U_{m}$ and $\phi_{m}^{-1}(\{1\})=\{y\}$. By the same estimates as before we must have that for large enough $n$, each $\mathcal{M}_{E}(x_{n})$ intersects $U_{m}$ non-trivially. Hence $y\in\textnormal{Li }\mathcal{M}_{E}(x_{n})$. So $\textnormal{Li }\mathcal{M}_{E}(x_{n})=\mathcal{M}_{E}(x)$.

Conversely suppose that we are given that for every convergent sequence $x_{n}\rightarrow x$, $\textnormal{Li }\mathcal{M}_{E}(x_{n})=\mathcal{M}_{E}(x)$. Let $\phi\in C^{0}_{b}(X)$. We may assume that $\phi$ is non-negative. Fix $\epsilon>0$. As $\mathcal{M}_{E}(x)$ is compact, we may find an open set $U$ containing $\mathcal{M}_{E}(x)$ such that

$$ $\max_{y\in U}\phi(y)-\min_{y\in U}\phi(y)<\epsilon$ (95) $$

As $\textnormal{Li }\mathcal{M}_{E}(x_{n})\subseteq\mathcal{M}_{E}(x)$, we may find $N$ large enough such that for all $n\geq N$, $\mathcal{M}_{E}(x_{n})\subseteq U$. Then

$$ $\min_{y\in\mathcal{M}_{E}(x)}\phi(y)\leq\int_{\mathcal{M}_{E}(x)}\phi\,d\mu_{x }\leq\max_{y\in\mathcal{M}_{E}(x)}\phi(y)$ (96) $$

and similarly

$$ $\min_{y\in\mathcal{M}_{E}(x_{n})}\phi(y)\leq\int_{\mathcal{M}_{E}(x_{n})}\phi \,d\mu_{x_{n}}\leq\max_{y\in\mathcal{M}_{E}(x_{n})}\phi(y)$ (97) $$

Hence we see both of these integrals are bounded between the min and max of $\phi$ on $U$. Hence

$$ $\left|\int_{\mathcal{M}_{E}(x)}\phi\,d\mu_{x}-\int_{\mathcal{M}_{E}(x_{n})} \phi\,d\mu_{x_{n}}\right|\leq 2\epsilon$ (98) $$

Therefore this converges and we have weak convergence.
∎

###### Proposition F.1 .

###### Proof.

###### Proposition F.2 .

###### Proof.

###### Theorem F.3 .

###### Proof.

###### Theorem F.4 .

###### Proof.

###### Theorem F.5 .

###### Proof.

###### Proposition F.6 .

###### Proof.

###### Lemma F.7 .

###### Proof.

### F.1 Almost everywhere continuity is not as impossible

##### Canonicalization

Let $X$ be a space with a group $G$ acting on it by homeomorphisms. Then a continuous canonicalization function is a continuous $G$-equivariant map $v:X\rightarrow X$ which sends any element $x$ to an element of the orbit $Gx$. In Dym et al. () it is shown that in general these do not exist, as it is much too strong of a condition. In fact, because a continuous canonicalization function will be a right-inverse to the quotient map $X\rightarrow X/G$, a necessary condition to the existence of $v$ is that the map $\pi_{1}(X)\rightarrow\pi_{1}(X/G)$ induced by the quotient must be a surjection. If the quotient $X\rightarrow X/G$ is a covering map then this only holds if $G$ is trivial. For a concrete example of this, we can consider quotienting $\mathbb{R}^{2}$ by the translation action of $\mathbb{Z}^{2}$. Then $\pi_{1}(\mathbb{R}^{2})=0$ while $\pi_{1}(\mathbb{R}^{2}/\mathbb{Z}^{2})=\mathbb{Z}^{2}$, showing that a 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 a canonicalization function which is continuous away from a set of measure zero, this will suffice for our applications. Specifically we introduce the notion of an almost everywhere (a.e.) continuous canonicalization function.

###### Definition F.1 .

An a.e. continuous canonicalization function is a function $v:X\rightarrow X$ such that there exists a set $S\subset X$ of measure-zero such that

- (i)
$G(S)\subseteq S$
- (ii)
$v$ is continuous away from $S$
- (iii)
$\forall x\in S,\exists g\in G\;\text{st. }v(x)=gx$
- (iv)
$v$ is $G$-invariant

As with the regular canonicalization functions, we can use such a function $v$ to construct equivariant and invariant functions. Given an arbitrary function $f:X\rightarrow\mathbb{R}$ it is easy to see that

$$ $\mathcal{I}_{v}(f)(x)=f(v(x))$ (99) $$

is $G$-invariant. Further it is easy to check that this sends a.e. continuous functions to a.e. continuous functions.

The added freedom of the set $S$ where $v$ can take arbitrary values allows us to avoid the topological restrictions that prevent the existence of globally continuous canonicalization functions. For a simple example consider again the presentation of a torus as the quotient of $\mathbb{R}^{2}$ by the translation action. Then an easy example of an a.e. continuous canonicalization function is the function

$$ $v(x,y)=([x],[y])$ $$

where $[a]$ is the decimal part of a real number $a\in\mathbb{R}$. This is discontinuous whenever either $x$ or $y$ are integers, but this set has measure zero.

For another example, consider the action of the orthogonal group $O(d)$ on the space of matrices $\mathbb{R}^{d\times n}$. The action preserves the rank of a matrix $A\in\mathbb{R}^{d\times n}$, so the subset of full rank matrices is preserved by $O(d)$. This is an open dense set, and if we take $\mu$ to just be the Lebesgue measure, its complement has measure zero. So take $S$ to be the set of matrices with non-full rank. Let $A$ be a matrix of full rank. Then one of the $d\times d$ minors of $A$ is invertible, we may assume that it is the first one. Thus we can write $A=(A_{0}A_{1})$ with $A_{0}$ an invertible $d\times d$ matrix. Applying the Gram-Schmidt procedure to $A_{0}$, we obtain $A_{0}=QR$ with $Q$ orthogonal and $R$ upper-triangular. Now we take $v(A)=Q^{T}A$. $Q$ is given by a polynomial in the elements of $A$, hence $v$ is smooth. Further, if $O\in O(d)$, then we can decompose $OA$ as $OQR$ with $OQ$ orthogonal. Hence,

$$ $v(OA)=(OQ)^{T}OA=Q^{T}O^{T}OA=Q^{T}A=v(A)$ $$

Therefore $v$ is an a.e. continuous canonicalization function, with $S$ the set of non-full rank matrices in $\mathbb{R}^{d\times n}$.

###### Definition F.1 .

##### Invariant Theory

Finding and computing $G$-invariant functions is a rich topic which comes up in many different situations. In algebraic geometry, one wants to find generators for $G$-invariant polynomials, i.e. a set of $G$-invariant polynomials $\sigma_{1},\dots,\sigma_{n}$ such that any other $G$-invariant polynomial $f$ can be written as

$$ $f=g(\sigma_{1},\dots,\sigma_{n})$ $$

for some polynomial $g$. For linearly reductive groups acting on $\mathbb{C}^{n}$, there is an algorithm for computing these generators, which is described in Derksen & Kemper ().

In the smooth case, there is a classic result of Schwarz (Schwarz, ), telling us that if $G$ is a compact Lie group acting on a smooth manifold $X$, with generators for the space of $G$-invariant polynomials being $\sigma_{1},\dots,\sigma_{n}$, then any $G$-invariant smooth functions $f$ can be written as

$$ $f=g(\sigma_{1},\dots,\sigma_{n})$ $$

where $g$ is an some smooth function on $\mathbb{R}^{n}$.

##### 

## Appendix G Algorithms 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 to specific cases:

##### If the group is compact or the algebra is solvable

In either of these cases, we are able to parameterise the group globally using the Lie algebra. Denoting the basis as $v_{i}\in\mathfrak{g}$, and the coefficients as $\alpha_{i}\in\mathbb{R}$, we have that we can represent any $g\in G$ as

$$ $g=\operatorname{exp}{\sum_{i}\alpha_{i}v_{i}},$ $$

if $G$ is compact, or as

$$ $g=\prod_{i}\operatorname{exp}{\alpha_{i}v_{i}},$ $$

if the Lie algebra is solvable with $v_{i}$ ordered as to represent the nested subalgebras. In both of these cases we can reduce the optimization to be done over the vector of coefficients $\alpha\in\mathbb{R}^{\operatorname{dim}\mathfrak{g}}$. Both of these can be viewed as specific types of global retraction mappings $\tau:\mathfrak{g}\rightarrow G$. For this case we use Algorithm 1

Figure: Algorithm 1 Canonicalization with a global retraction

##### If the group is non-compact or the algebra is not solvable

Such a setting is particularly important for PDE-based groups, as these naturally result in groups that are non-compact, see section 5.3. In this case it is no longer possible to find a global retraction, and a natural question arises - can the infinitesimal actions ever be enough? It turns out that the answer is yes, quite generically, resulting in Algorithm 2:

###### Proposition G.1 (Proposition 1.24 from Olver ( 1993 ) ) .

Let $G$ be a connected Lie group and $U\subset G$ a neighbourhood of the identity. Also, let $U^{k}\equiv\left\{g_{1}\cdot g_{2}\cdot\ldots\cdot g_{k}:g_{i}\in U\right\}$ be the set of $k$-fold products of elements of $U$. Then

$$ $G=\bigcup_{k=1}^{\infty}U^{k}.$ $$

In other words, every group element $g\in G$ can be written as a finite product of elements of $U$.

Figure: Algorithm 2 Canonicalization function via Lie algebra descent

###### Proposition G.1 (Proposition 1.24 from Olver ( 1993 ) ) .

##### If the group only provides information about basis vector action

In most examples of PDE symmetries, integrating the action for a generic $v\in\mathfrak{g}$ is impossible, and instead only the action with respect to certain basis $\{v_{i}\}$ is known. In this case coordinate descent schemes must be used. For this we use Algorithm 3.

Figure: Algorithm 3 Canonicalization function via coordinate Lie algebra descent

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: Besides the questions in the weaknesses section:
- 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?
- 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.
- 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.
soundness: 3
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.
- 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.
- I find the potential use-case of applying this methodology to pre-trained models still under-explored and worthy of further investigation.
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.
weaknesses: I think the current presentation lacks focus and the paper could be vastly improved with the objectives of:
- Highlighting clearly the limitations of past/current proposals and how these limitations are overcome, potentially alternating more concrete examples with more abstract limitations.
- 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. 

In regards to improving clarity:
- $WFra_{G}(X)$ should be defined before it is used on line 236.
- On lines 270-271 is $N$ simply some value $\in \mathbb{R}$?
- 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}$.
- 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). 
- 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.
- 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.
- 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. 
- "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].
- 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]? 

I 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. 

Some of the expression/grammar in the appendix could also be improved, e.g. "Liealgebra descent", "why failure modes exists".

[1] Enabling equivariance for arbitrary lie groups, MacDonald et. al 2022
[2] Lie group decompositions for equivariant neural networks, Mironenco & Forre 2024
[3] Trivializations for Gradient-Based Optimization on Manifolds, Mario Lezcano-Casado, 2019
[4] Optimization algorithms on matrix manifolds, Absil et. al 2009

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Please see the weaknesses section.
soundness: 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.
- 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.
- 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.

[3] Dym et al. Equivariant frames and the impossibility of continuous canonicalization (2024)

[4] Ma et al. A canonicalization perspective on invariant and equivariant learning (2024)
summary: The contribution of this paper can be summarized as follows:
- 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),
- A class of optimization-based algorithms for the above energy-based canonicalization using (coordinate) Lie algebra descent (Appendix H),
- 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).

[1] Puny et al. Frame averaging for invariant and equivariant network design (2021)

[2] Kaba et al. Equivariance with learned canonicalization functions (2023)

[3] Dym et al. Equivariant frames and the impossibility of continuous canonicalization (2024)

[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.
- 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.

Minor comments and suggestions

- In Line 249, the notation for closure $\bar{X}$ of a set $X$ is used without definition.
- In Line 260, $\mathrm{PMeas}$ -> $\mathrm{PMeas}(X)$
- In Line 976, euclidean -> Euclidean
- In Line 1758, the union of -> a union of?
- 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.

[5] Mondal et al. Equivariant adaptation of large pretrained models (2023)

[6] Kim et al. Learning probabilistic symmetrization for architecture agnostic equivariance (2023)

[7] Zhang et al. SymDiff: Equivariant diffusion via stochastic symmetrisation (2024)

[8] Arora et al. Invariant physics-informed neural networks for ordinary differential equations (2024)

[9] Lagrave & Tron, Equivariant neural networks and differential invariants theory for solving partial differential equations (2022)

[10] Nguyen et al. Learning symmetrization for equivariance with orbit distance minimization (2023)

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 1
questions: See above.
soundness: 3
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.
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.'
 
[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.
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. 
- 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. 
- 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. 
- 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.  

Overall, 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.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: I suggest authors address and elaborate on weaknesses.
soundness: 3
strengths: 1. The proposed canonicalization approach allows employing pre-trained non-equivariant models while making them equivariant with canonicalization.
2. The paper provides a good overview of the related work on canonicalization and frame-averaging methods.
2. Figure 1 provides an intuitive example of the effect of canonicalization on the decision boundary.
3. The proposed approach is generic and is potentially applicable to a wide range of models and learning tasks.
4. The paper provides strong theoretical support for the proposed method.
5. 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.
6. The paper nicely elaborates on the challenges of choosing the appropriate energy function and it outlines key considerations in doing so.
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.
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?
2. 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.
3. The distance-based relaxation in 331-333 is not clear. Can authors elaborate more on what is meant here?
4. 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.
5. 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.
6. 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.

Some important missed related work on equivariance:
- M. Zhdanov et al. Clifford-steerable convolutional neural networks NeurIPS 2023.
- D. Ruhe et al. Clifford group equivariant neural networks NeurIPS 2023.

Some important missed related work on Lie groups in NNs:
- A. Moskalev et al. Liegg: Studying learned lie group generators. NeurIPS 2022.
- N. Gruver et al. The Lie Derivative for Measuring Learned Equivariance. ICLR 2022.

Meta-review: This paper proposes a novel extension of canonicalization-like approaches to non-compact Lie groups, a significant advancement in building symmetric neural networks. Frame averaging and canonicalization are foundational techniques, and extending these mechanisms to a broader class of groups addresses a critical gap in the literature. The approach is well-grounded theoretically, with rigorous justification, and the proposed method is both novel and impactful. Experimental results convincingly demonstrate its effectiveness. While initial concerns were raised regarding the clarity of the presentation, the authors resolved most issues during the rebuttal phase with thoughtful revisions, improving the manuscript's accessibility.

Decision: Accept (Poster)