1. Bibliographic Information

• Title: Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groups
• Authors: Zakhar Shumaylov, Peter Zaika, James Rowbottom, Ferdia Sherry, Melanie Weber, and Carola-Bibiane Schönlieb.
• Affiliations: The authors are from the University of Cambridge and Harvard University. Their backgrounds span machine learning, geometric deep learning, and applied mathematics.
• Journal/Conference: This paper is a preprint available on arXiv. It was submitted in October 2024. As a preprint, it has not yet undergone formal peer review for a conference or journal.
• Publication Year: 2024
• Abstract: The paper introduces Lie aLgebrA Canonicalization (LieLAC), a novel method to make existing machine learning models equivariant to symmetries described by Lie groups. The primary challenge addressed is that many important symmetries (e.g., in Physics-Informed Neural Networks for PDEs) involve non-compact groups whose full structure is often unknown beyond their infinitesimal generators (the Lie algebra). LieLAC overcomes this by using only the Lie algebra to transform inputs into a "canonical" form before they are fed to a model. This process, a form of canonicalization, aligns inputs according to the allowed symmetries, effectively making a pre-trained, non-equivariant model behave equivariantly. The authors provide a theoretical foundation for this approach, connecting it to frame averaging, and demonstrate its effectiveness on tasks like invariant image classification and equivariant PDE solving with pre-trained models.
• Original Source Link:
  • Official Source: https://arxiv.org/abs/2410.02698v2
  • PDF: https://arxiv.org/pdf/2410.02698v2.pdf

2. Executive Summary

• Background & Motivation (Why):

  • Core Problem: Modern deep learning models often lack robustness when inputs are transformed (e.g., rotated, scaled, or sheared). Building models with built-in symmetries, known as equivariant neural networks, can solve this, but existing methods are often limited to simple, well-behaved (compact) groups like rotations. Many real-world problems, especially in physics and scientific computing, involve complex, non-compact Lie groups (like scaling or shearing groups) for which the full group structure is not even known. This makes it impossible to directly apply most equivariant architectures.
  • Why It Matters: In scientific domains like solving Partial Differential Equations (PDEs), symmetries are fundamental physical principles. Models that respect these symmetries are more generalizable, data-efficient, and physically plausible. For example, the solution to a heat diffusion problem shouldn't fundamentally change if the coordinate system is scaled or shifted. The inability to enforce these complex symmetries in neural networks is a major gap.
  • Fresh Angle: Instead of building a new equivariant model from scratch, the paper proposes a method, LieLAC, to make any pre-trained model equivariant. Crucially, LieLAC does not require knowledge of the full Lie group. It only needs the infinitesimal generators (the Lie algebra), which are much easier to derive and work with. It does this by finding a "canonical" version of any input by "undoing" the symmetry transformation.

• Main Contributions / Findings (What):

  1. A Novel Canonicalization Framework (LieLAC): The paper introduces LieLAC, an energy-based canonicalization method that works for arbitrary Lie groups, including non-compact ones. It operates by defining an "energy" function and using Lie group optimization techniques to find the group transformation that minimizes this energy, thereby mapping the input to a canonical representation.
  2. Theoretical Unification: It provides a unified theoretical framework that connects different approaches to equivariance, including frame averaging and canonicalization. It extends existing theory to handle continuous, non-compact groups by introducing concepts like weighted closed canonicalizations.
  3. Practical Application to Pre-trained Models: LieLAC is a "plug-and-play" module. It can be wrapped around any existing model (like a standard CNN or a large foundation model) to make it equivariant without retraining from scratch.
  4. Demonstrated Efficacy: The authors showcase LieLAC's effectiveness in two key areas:
     • Invariant Image Classification: Making a standard CNN robust to affine and homography transformations, outperforming specialized equivariant models.
     • Equivariant PDE Solvers: Applying LieLAC to pre-trained neural operators (DeepONet, POSEIDON) to solve the Heat, Burgers', and Allen-Cahn equations, showing dramatically improved accuracy on out-of-distribution data that has been transformed by the PDE's symmetry group.

3. Prerequisite Knowledge & Related Work

• Foundational Concepts:

  • Symmetry, Invariance, and Equivariance:
    • A function is invariant to a transformation if its output does not change when the input is transformed. For a function $f$ and transformation $g$, invariance means $f(g \cdot x) = f(x)$. Example: A digit classifier should be invariant to small rotations; "7" is still "7" if slightly tilted.
    • A function is equivariant if transforming the input results in an equivalent transformation of the output. For a function $\Phi$ and transformation $g$, equivariance means $\Phi(g \cdot x) = g \cdot \Phi(x)$. Example: In image segmentation, if you rotate the input image, the output segmentation mask should also be rotated by the same amount.
  • Lie Groups and Lie Algebras:
    • A Lie group is a set of continuous transformations (like rotations, translations, scaling) that is also a smooth manifold. This means we can use calculus on the group itself.
    • A Lie algebra is the "infinitesimal" version of a Lie group. It's the vector space of tangent vectors at the group's identity element. Think of it as describing the "velocity" or direction of transformations. For example, the Lie algebra of the 2D rotation group SO(2) is a 1D space representing infinitesimal rotation speed. The paper's key idea is that you only need the Lie algebra, not the full group.
  • Partial Differential Equations (PDEs): These are equations involving unknown functions of multiple variables and their partial derivatives. They are fundamental to describing physical phenomena like heat flow, fluid dynamics, and wave propagation.
  • Physics-Informed Neural Networks (PINNs): These are neural networks trained to solve PDEs. Instead of relying solely on data, their loss function includes a term that penalizes the network if its output violates the governing PDE. This embeds physical laws directly into the learning process.
  • Neural Operators: Unlike standard neural networks that map between finite-dimensional vectors, neural operators learn mappings between infinite-dimensional function spaces. For example, a neural operator can learn to map any initial condition function of a PDE to its solution function at a later time.

• Previous Works & Differentiation:

  • Equivariant Convolutions: A popular approach (e.g., Cohen & Welling, 2016) generalizes standard convolutions to work on groups. This requires defining signals on the group and performing group convolutions, which is difficult for continuous groups and grid-based data. LieLAC avoids this complexity.
  • Lie Algebra-based Networks: Some works (e.g., Finzi et al., 2020) use the Lie algebra to approximate group convolutions. These methods still require building custom architectures. LieLAC, in contrast, works with any existing architecture.
  • Canonicalization/Symmetrization: The idea of transforming an input to a canonical form is not new. However, prior work often had theoretical gaps, especially regarding non-compact groups, or practical issues in constructing the canonicalization map. This paper addresses these theoretical issues and provides a constructive, energy-based method.
  • Symmetry in PINNs: Previous attempts to use symmetries in PINNs often relied on data augmentation (generating new training samples by transforming existing ones) or adding a symmetry-violating term to the loss. LieLAC provides a way to enforce equivariance directly in the model's architecture (or rather, its wrapper) during inference.

4. Methodology (Core Technology & Implementation)

The core of the paper is the development of a practical and theoretically sound canonicalization framework, LieLAC, that works for arbitrary Lie groups.

Principles: From Frames to Weighted Closed Canonicalizations

The authors build a theoretical bridge from existing concepts to their new framework.

1. Frames and Canonicalizations: These are two ways to achieve equivariance.

   • Frames: For each input $x$, a frame $\mathcal{F}(x)$ is a set of group elements. Averaging a function's output over these transformed inputs makes it equivariant. This is computationally expensive and has issues with continuous groups.
   • Canonicalizations: For each input $x$, a canonicalization $\mathcal{C}(x)$ is a set of "canonical" points in the same orbit as $x$. The model only ever sees points from this canonical set. The key property is invariance: $\mathcal{C}(gx) = \mathcal{C}(x)$. This means no matter how you transform an input, it gets mapped to the same canonical representation.

2. The Problem with Non-Compact Groups: Standard canonicalization runs into two major problems with non-compact Lie groups (e.g., scaling):

   • Orbits may not be closed: An orbit $Gx = \{g \cdot x | g \in G\}$ might not contain its limit points. For example, scaling a point towards the origin creates a sequence whose limit (the origin) is not in the orbit. An energy function might not have a minimum on such an open set.
   • Orbits can be infinite: Averaging over an infinite set is not straightforward.

3. The Solution: Weighted Closed Canonicalizations: LieLAC addresses these issues by:

   • Introducing Weights: Instead of a finite set of points, they define a weighted canonicalization $\kappa_x$ as a probability measure on the input space $X$. This allows for continuous or infinite canonical sets.
   • Using Orbit Closures: To solve the non-closed orbit problem, they define the canonical set not on the orbit Gx, but on its closure $\overline{Gx}$. This guarantees that a minimum for a continuous energy function exists. This leads to the concept of a Weighted Closed Canonicalization (Definition 3.7), where the support of the measure $\kappa_x$ is contained within the orbit closure $\overline{Gx}$.

LieLAC: Energy-based Canonicalization

LieLAC is a constructive method to create such a canonicalization. The pipeline is:

1. Define an Energy Function: Choose an energy function $E: X \to [0, +\infty]$ on the input space. This function should assign low energy to "canonical" or "desirable" inputs. The choice of $E$ is domain-specific. Examples include:

   • The negative log-likelihood of a data distribution (if known).
   • The distance to the training data domain (e.g., ensuring inputs stay within the bounding box of the training data).
   • A measure of "simplicity" or "alignment" (e.g., aligning an image to be upright).

2. Find the Minimizing Set: For any given input $x$, find the set of points in its orbit closure that minimize the energy: $$\mathcal{M}_E(x) = \underset{y \in \overline{Gx}}{\arg\min} E(y)$$ This set $\mathcal{M}_E(x)$ is the canonical representation of $x$. By construction, it is invariant to group actions (i.e., $\mathcal{M}_E(gx) = \mathcal{M}_E(x)$).

3. Construct the Canonicalization: The final canonicalized output is an average over the points in $\mathcal{M}_E(x)$. If this set is finite, it's a simple average. If it's a continuous set, a measure (like the normalized Hausdorff measure) is used.

Practical Implementation

In practice, we need to find a group element $g$ that transforms the input $x$ to its canonical form, i.e., $g \cdot x \in \mathcal{M}_E(x)$. This is framed as an optimization problem: $$\text{find } g \in G \text{ that minimizes } E(g \cdot x)$$ Since $G$ is a Lie group, this can be solved using gradient-based optimization methods on the group or its Lie algebra.

• Algorithm 1: The paper proposes using standard Lie group descent schemes. The algorithm starts with an initial guess for a Lie algebra element $\xi_0$ and iteratively updates it using gradient descent on the energy function. A retraction map $\tau: \mathfrak{g} \to G$ is used to map the Lie algebra element back to a group element at each step.

  Algorithm 1: Canonicalization with a global retraction

  Data: Non-canonical input x parameters :N, steps ηi, ξ0 g, Energy E : X → R, Retraction

  τ : g → G Result: Canonicalized input x = g−1 · x; inverse g−1;

  canonicalizing group element g. # Do gradient descent on ξ g

  ξ ← ξ0;

  for i = 0 . . . N do | ξ ← ξ − ηi∇ξE(τ (ξ) · x)

  end

  return τ (ξ) · x; τ(ξ); [τ (ξ)]−1

• Optimization Challenges: This optimization is generally non-convex, even if the energy $E$ is convex, because the orbit Gx is often a non-convex manifold (e.g., a circle for rotations). This can lead to local minima or noisy gradients, requiring careful optimization strategies.

  该图像是包含函数曲线与热力图的图表，展示了不同变换作用下的函数形态及其时空演变。左图有函数曲线组，右图依次为真实解、变换前后函数的时间-空间分布，体现了对称变换在PDE求解中的应用。

• Figure 4 Analysis: This figure illustrates the non-convex optimization process. The left panel shows a target distribution. The plots on the right show the trajectory of an input sample as it is optimized (canonicalized). The "Energy over the group" plot (top middle) is highly non-convex, with many local minima, demonstrating the difficulty of the optimization. The "Energy over trajectory" plot (right) shows that the optimization successfully finds a low-energy state.

  The diagram below (Figure 3) summarizes the theoretical landscape, showing how the paper's contributions (WCan, WCCan) extend previous work on finite and compact groups to the general non-compact case.

  该图像是论文中Figure 8展示的TSNE二维降维图，分别显示了后仿射和后齐次变换正则化处理前后的MNIST数据分布，反映了LieLAC方法对输入数据表征的规范化效果。

5. Experimental Setup

The authors validate LieLAC on three types of tasks.

• Datasets:

  1. 2D Toy Example: A synthetic dataset of two concentric rings, designed to show how canonicalization can simplify a non-linearly separable problem. The symmetry group is 2D rotation, SO(2).
  2. Invariant Image Classification:
     • MNIST: The standard dataset of handwritten digits.
     • affNIST & homNIST: Versions of MNIST where the images are distorted by random affine transformations (rotation, scaling, shear, translation) and homography transformations (projective transformations).
  3. Equivariant PDE Solvers:
     • Heat Equation: A 1D PDE with a 6-dimensional non-compact symmetry group.
     • Burgers' Equation: A 1D non-linear PDE with a 5-dimensional non-compact symmetry group.
     • Allen-Cahn Equation: A 2D PDE with a 4-dimensional symmetry group (SE(2) - rotations and translations). The experiment uses the pre-trained POSEIDON foundation model.

• Evaluation Metrics:

  • Classification Accuracy: The standard percentage of correctly classified images for the MNIST tasks.
  • $L_2$ Relative Error: Used for the PDE tasks to measure the difference between the predicted solution $u_\theta$ and the true solution $u$.
    1. Conceptual Definition: It measures the normalized magnitude of the error vector. An error of 0.1 means the prediction is off by 10% on average, relative to the magnitude of the true solution.
    2. Mathematical Formula: $$\text{Error} = \frac{\| u - u_\theta \|_{L_2}}{\| u \|_{L_2}} = \frac{\sqrt{\int (u(x) - u_\theta(x))^2 dx}}{\sqrt{\int u(x)^2 dx}}$$
    3. Symbol Explanation:
       • u(x): The ground truth solution.
       • $u_\theta(x)$: The solution predicted by the neural network with parameters $\theta$.
       • $\|\cdot\|_{L_2}$: The $L_2$ norm, which represents the "length" or "magnitude" of a function, calculated by integrating the square of the function over its domain and taking the square root.

• Baselines:

  • Standard Models: CNN for image classification, DeepONet and POSEIDON for PDE solving. These models are not inherently equivariant.
  • Equivariant Models: affConv and homConv, which are specialized convolutional networks designed to be equivariant to affine and homography transformations, respectively.

6. Results & Analysis

LieLAC consistently improves the robustness and generalization of pre-trained models.

5.1 2D Toy Example

该图像是论文中展示的图1，属于示意图，展示了标准kNN分类与通过Lie代数范式化后kNN分类对内外环数据分布的决策边界影响，体现了范式化对分类边界的显著改善。

• Figure 1 Analysis: This figure clearly demonstrates the power of canonicalization. The original data (top left) consists of rings and scattered points. A standard k-NN classifier (bottom right) fails to separate the inner and outer rings, creating a nonsensical boundary. After applying LieLAC, all points are canonicalized (rotated to align with a specific axis). In this canonical space, the rings are perfectly separated, and the k-NN classifier easily learns the correct circular boundary (bottom left).

5.2 Invariant Classification

This experiment shows that LieLAC can make a standard CNN robust to complex transformations.

(Manual Transcription of Table 1) Table 1: MNIST test accuracy for Affine and Homography groups, compared with affConv and homConv of MacDonald et al. (2022).

• Core Results: A standard CNN performs well on MNIST but fails dramatically on the transformed affNIST and homNIST datasets (accuracy drops from ~98.5% to ~63-64%). By wrapping the exact same CNN with LieLAC, the accuracy on the transformed datasets soars to 97.2% and 96.0%. This is not only a massive improvement but also surpasses the performance of the specialized equivariant models (affConv and homConv).

  该图像是论文中展示的图表，展示了通过仿射变换和单应性变换作用于MNIST手写数字图像的原始图像、规范化图像、失真图像及失真后规范化图像对比。

• Figure 5 Analysis: This figure provides a visual confirmation. The first column shows original MNIST digits. The second shows their canonicalized versions (which look almost identical, as they are already near-canonical). The third column shows randomly distorted images. The fourth column shows the result of applying LieLAC to these distorted images: they are successfully transformed back into their canonical, upright forms. The model then easily classifies these canonicalized images.

5.3 Equivariant PDE evolution

This section is the key demonstration of LieLAC's utility for scientific machine learning.

该图像是图表，展示了数值偏微分方程演化的规范化流程，对比了多种算子在初始条件、真实演化和规范化处理后的表现，涉及的方程包括热方程、Burger方程和Allen-Cahn方程，图中还标注了误差值。

• Figure 2 Analysis: This figure illustrates the entire workflow for PDE solving.

  • Top Row: A standard neural operator ($O_θ$) is given an out-of-distribution initial condition. Its prediction is poor, with a high error.

  • Bottom Row: The same out-of-distribution initial condition is first passed through LieLAC's canonicalization step ($g_c^{-1}$). This transforms it into an in-distribution, "canonical" initial condition. The neural operator then makes an accurate prediction on this canonicalized input. Finally, the output is transformed back using the inverse group element ($g_c$) to get the correct solution in the original coordinate system. The error is drastically reduced.

    (Manual Transcription of Table 2) Table 2: $L_2$ relative error for Heat equation and Burgers' equation averaged over the time period. Heat and Burgers u achieved via symmetry transformations.

• Heat & Burgers' Equations Results: For both PDEs, the standard DeepONet performs well on in-distribution (ID) data but catastrophically fails on out-of-distribution (OOD) data created by symmetry transformations (e.g., error for Heat equation jumps from 0.05 to 0.66). LieLAC [DeepONet] reduces the OOD error by over an order of magnitude, making its OOD performance as good as its ID performance. This shows that LieLAC successfully generalizes the model's capabilities to the entire symmetry group orbit. Even when the baseline is trained with data augmentation, LieLAC still provides better and more consistent OOD performance.

  (Manual Transcription of Table 3) Table 3: Results for the Allen-Cahn Equation. Errors are $(\times 10^{-4})$.

• Allen-Cahn Equation & Foundation Models: This experiment tests LieLAC on POSEIDON, a large pre-trained model.
  • The base POSEIDON model fails on OOD data (error jumps from 6.93 to 75.76).
  • LieLAC [POSEIDON] significantly reduces the average error, but the ID error increases. This suggests a misalignment between the canonicalization energy minimum and the model's trained data distribution.
  • The best results are achieved via fine-tuning (ft). Fine-tuning the POSEIDON model on a small set of canonicalized data (LieLAC [POSEIDON + ft.]) yields strong performance on both ID and OOD data. This is more data-efficient than standard data augmentation. It shows that canonicalization provides a more structured way to fine-tune models for symmetry.

7. Conclusion & Reflections

• Conclusion Summary: The paper successfully introduces and validates LieLAC, a powerful and general framework for enforcing equivariance in neural networks under arbitrary, even non-compact, Lie groups. By relying only on the Lie algebra and an energy-minimization principle, it bypasses the need for full group knowledge or custom equivariant architectures. This allows it to be used as a wrapper around pre-trained models, dramatically improving their generalization and robustness to symmetry transformations, as demonstrated in both computer vision and scientific computing tasks.

• Limitations & Future Work: The authors are transparent about the limitations:

  • Inference Speed: The optimization step required for canonicalization at inference time is slow, causing a 5-30x slowdown compared to the base model. This makes it unsuitable for real-time applications without further optimization.
  • Non-Convex Optimization: The energy minimization is a non-convex problem that can be difficult to solve reliably.
  • Model Constraints: Applying LieLAC to existing foundation models can be difficult if they have hard-coded assumptions, such as fixed boundary conditions.

• Personal Insights & Critique:

  • High Practical Impact: The ability to confer equivariance upon any pre-trained model is a significant practical contribution. As foundation models become more prevalent, methods like LieLAC that adapt them for specific physical or geometric constraints will be increasingly valuable. Fine-tuning on a small set of canonicalized data is a particularly promising and efficient paradigm.
  • Generality is Key: The paper's main strength is its generality. By not being tied to a specific group structure (like SO(3) or SE(3)), it opens the door to using symmetries from a much wider range of physical problems that were previously inaccessible to equivariant deep learning.
  • The Trade-off: The primary trade-off is between universality and speed. While LieLAC is very general, its iterative optimization at inference time is a major bottleneck. Future work could explore learning the canonicalization map itself, perhaps with a separate neural network, to amortize the optimization cost and achieve faster inference.
  • Energy Function as an Inductive Bias: The choice of the energy function is a critical and non-trivial step. It's a powerful way to inject domain knowledge, but a poor choice could lead to suboptimal canonical forms. More systematic ways to design or learn these energy functions would be a valuable research direction.