Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groups
Zakhar Shumaylov, Peter Zaika, James Rowbottom, Ferdia Sherry, Melanie Weber, Carola-Bibiane Schönlieb
TL;DR
Non-compact Lie group symmetries in PDEs hinder standard equivariant architectures. LieLAC introduces energy-based canonicalization that uses only the action of a Lie algebra basis $\mathfrak{g}$ to canonically align inputs, enabling equivariance for pre-trained models without full group parametrizations. The work unifies frames and canonicalizations under a weighted, continuous framework and provides group-descent strategies for both solvable and non-solvable, compact and non-compact cases, including practical guidelines on energy design and optimization. Empirically, LieLAC improves invariant image classification under affine/homography transformations and yields Lie point symmetry equivariant neural PDE solvers for Heat, Burgers', and Allen–Cahn equations, with notable gains in out-of-distribution generalization and compatibility with foundation-model pipelines.
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.