[Paper Review] Generalizing Convolutional Neural Networks for Equivariance to Lie Groups on Arbitrary Continuous Data

Marc Finzi, Samuel Stanton+2 more|arXiv (Cornell University)|Feb 25, 2020

Machine Learning in Materials ScienceMaterials Science47 references73 citations

TL;DR

The paper introduces LieConv, a convolutional layer that achieves equivariance to transformations from any Lie group with surjective exponential map, enabling processing of arbitrary continuous spatial data (e.g., images, molecules, dynamical systems) with a single architecture. It demonstrates competitive or state-of-the-art results across image and molecular tasks and enables exact momentum conservation in Hamiltonian systems.

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.

Motivation & Objective

• Motivate and formalize a general framework for equivariant models on arbitrary continuous spatial data represented as coordinates and values.
• Develop LieConv, a convolutional layer equivariant to Lie group transformations via exp/log maps.
• Enable rapid prototyping by requiring only group exponential and logarithm maps to incorporate a new symmetry.
• Demonstrate the method across diverse domains (images, molecular data, dynamical systems) with a single architecture.

Proposed method

• Define LieConv as a convolution on Lie groups with a kernel k_theta mapped through the exponential/logarithm to the Lie algebra.
• Lift input data from space X to G orbits by selecting origins and sampling group elements (lifting procedure).
• Parameterize the kernel k_theta as a neural network over the Lie algebra to handle continuous group elements.
• Enforce locality using a group distance d(u,v)=||log(u^{-1}v)||_F to restrict the convolution support.
• Discretize the group convolution with a Monte Carlo estimator over a neighborhood to preserve equivariance in distribution.
• Extend the lifting to cases where the group acts with multiple orbits on X (X/G) and incorporate orbit information into the kernel.

Experimental results

Research questions

• RQ1 Can a single neural architecture, using Lie group theory, achieve equivariance to transformations from arbitrary Lie groups on non-grid data?
• RQ2 How can we efficiently lift arbitrary continuous spatial data to a group and define a local, differentiable, equivariant kernel over the group?
• RQ3 Does the proposed LieConv approach yield competitive or superior performance across image, molecular, and dynamical-system tasks while preserving physical symmetries?
• RQ4 Can exact conservation laws (e.g., linear and angular momentum) be achieved in Hamiltonian systems by enforcing appropriate Lie group symmetries?

Key findings

• LieConv achieves competitive results on RotMNIST and QM9, with state-of-the-art performance on some QM9 tasks.
• Using T(3), SE(3), and other Lie group equivariances improves performance on molecular property prediction compared to baselines.
• LieConv enables exact conservation of linear and angular momentum in Hamiltonian system modeling when appropriate symmetries (translations/rotations) are imposed.
• A single LieConv-based architecture can be applied across images, molecules, and dynamical systems with strong empirical performance.