Generalizing Convolutional Neural Networks for Equivariance to Lie Groups on Arbitrary Continuous Data

TL;DR

The paper tackles the challenge of generalizing convolutional networks to be equivariant to continuous Lie group transformations on arbitrary continuous data. It introduces LieConv, a convolutional layer whose kernel is defined on a Lie group via exp/log maps, with lifting of inputs, locality constraints, and Monte Carlo discretization to perform group convolutions without relying on regular grids, and it supports multiple orbits. The approach is demonstrated on image data (RotMNIST), molecular data (QM9), and dynamical systems, achieving competitive or state-of-the-art results and, in Hamiltonian settings, exact conservation of linear and angular momentum when appropriate symmetries are enforced. The work offers a practical, extensible framework for rapid prototyping of equivariant models across diverse modalities, with open-source code to enable broad adoption.

Abstract

The translation equivariance of convolutional layers enables convolutional neural networks to generalize well on image problems. While translation equivariance provides a powerful inductive bias for images, we often additionally desire equivariance to other transformations, such as rotations, especially for non-image data. We propose a general method to construct a convolutional layer that is equivariant to transformations from any specified Lie group with a surjective exponential map. Incorporating equivariance to a new group requires implementing only the group exponential and logarithm maps, enabling rapid prototyping. Showcasing the simplicity and generality of our method, we apply the same model architecture to images, ball-and-stick molecular data, and Hamiltonian dynamical systems. For Hamiltonian systems, the equivariance of our models is especially impactful, leading to exact conservation of linear and angular momentum.

Figures (10)

• Figure 1: Many modalities of spatial data do not lie on a grid, but still possess important symmetries. We propose a single model to learn from continuous spatial data that can be specialized to respect a given continuous symmetry group.
• Figure 2: Visualization of the lifting procedure. Panel (a) shows a point $x$ in the original input space $\mathcal{X}$. In panels (b)--(f) we illustrate the lifted embeddings for different groups in the form $[u, q]$, where $u \in G$ is an element of the group and $q \in \mathcal{X} / G$ identifies the orbit (see Section \ref{['sec:nonhomogeneous']}). For SE($2$) the lifting is multi-valued.
• Figure 3: A visualization of the local neighborhood for $\mathbb{R}^* \times$ SO($2$), in terms of the points in the input space. For the computation of $h$ at the point in orange, elements are sampled from colored region. Notice that the same points enter the calculation when the image is transformed by a rotation and scaling. We visualize the neighborhoods for other groups in Appendix \ref{['sec:app_nbhd']}.
• Figure 4: Orbits of SO($2$) and T($1$)$_y$ containing input points in $\mathbb{R}^2$. Unlike T($2$) and SE($2$), not all points are not contained in a single orbit of these small groups.
• Figure 5: A visual overview of the LieConv model architecture, which is composed of $L$ LieConv bottleneck blocks that couple the values at different group elements together. The BottleBlock is a residual block with a LieConv layer between two linear layers.

Theorems & Definitions (1)

• Definition 1