Source-linked AI summary
Generalizing Convolutional Neural Networks for Equivariance to Lie Groups on Arbitrary Continuous Data
Marc Finzi, Samuel Stanton, Pavel Izmailov, Andrew Gordon Wilson
TL;DR
Convolutional networks need equivariance beyond translations to handle non-image continuous data with diverse symmetries. The paper introduces LieConv, which uses group exponential and logarithm maps to construct equivariant convolutions across data types, and reports strong results including conserved physical quantities in Hamiltonian systems. Its scope is bounded by assumptions and computational considerations associated with the supported data and parameterizations.
Problem
Existing convolutional inductive biases are powerful for images, but broader spatial data may be non-grid and require equivariance to transformations beyond translations.
Method
LieConv parameterizes continuous convolutional kernels on Lie groups and requires only the group exponential and logarithm maps to specialize equivariance to a new group.
Results
The same architecture achieves state-of-the-art performance across demonstrated domains, including Hamiltonian modeling where equivariance corresponds to preserving physical quantities.
Takeaways & Limitations
LieConv provides one flexible architecture for images, molecular data, and dynamical systems, with equivariance especially impactful for conserving linear and angular momentum in Hamiltonian systems.
Takeaways & Limitations
The kernel parameterization is restricted to Lie groups whose exponential maps are surjective, and practical computation can be costly for continuous point-convolution evaluation.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper develops LieConv, a single convolutional framework for arbitrary continuous spatial data that can be specialized to specified Lie-group symmetries. The same architecture spans images, molecular data, and dynamical systems, achieving competitive or state-of-the-art results while supporting physical conservation properties.
- Method: The method represents continuous convolutional kernels as neural-network parameterized functions on the group, extending convolution beyond regular grids.This covers images as well as ball-and-stick molecules and dynamical-system coordinates.
- Method: LieConv can be made equivariant to a given Lie group by defining its exponential and logarithm maps.The paper considers Lie groups including translations, rotations, and scalings.
- Results: The same network architecture is used across transformation groups and data types, while achieving state-of-the-art performance in the demonstrated domains.The paper evaluates images, molecular data, and dynamical systems.
- Contribution: LieConv is a convolutional layer equivariant to transformations from Lie groups and applicable to arbitrary collections of coordinates and values.Models with LieConv can operate on nonhomogeneous spaces and arbitrary spatial data.
- Results: On rotMNIST and QM9, LieConv is competitive with existing methods and outperforms state-of-the-art methods on some QM9 tasks.The introduction reports competitive results in all evaluated QM9 cases.
- Results: LieConv outperforms state-of-the-art methods for modeling Hamiltonian dynamical systems, where equivariance corresponds to preserving physical quantities.The cited physical quantities include energy and angular momentum.
2. Related Work
Prior equivariant-network approaches either rely on image-grid transformations or require cumbersome representation theory, whereas LieConv reduces the work needed to support new groups.
- Existing approaches: Transforming standard kernels or feature maps provides exact equivariance for discrete groups but has primarily been used on image data.This approach uses the regular representation for discrete groups and has also been extended to vector fields and other representations.
- Existing approaches: Harmonic-analysis methods construct exactly equivariant kernels for continuous groups but require deriving the group’s representation theory and are limited to compact groups.The approach has been applied to spherical images, voxel data, and point clouds.
- LieConv: LieConv reduces the implementation effort for equivariance to a new group, enabling rapid prototyping.This is presented as an advantage over methods requiring group-specific representation-theoretic work.
- Related Lie-group methods: Related Lie-group methods use logarithm maps or B-spline kernels, but the cited prior method is not readily applicable to point clouds.The related approaches include skeletal-pose rotations and image problems.
3. Background
The background develops equivariance and Lie-group convolutions, then motivates continuous kernel parameterizations for data without regular grids and computational strategies for making them practical.
- Equivariance: Equivariance means transforming an input causes the output to transform correspondingly, preserving information rather than discarding it through strict invariance.Translation-equivariant convolution is the standard deep-learning example.
- Groups and Lie groups: Groups formalize composable transformations through closure, identity, inverses, and associativity; SO(2) illustrates these properties for rotations.Rotations can be parameterized using an exponential map from an infinitesimal generator.
- Groups and Lie groups: A Lie group is a smooth manifold with group structure, while its Lie algebra is the vector space of infinitesimal transformations at the identity.The exponential map maps Lie-algebra elements to group elements, and a logarithm map can provide an inverse when the exponential is surjective.
- Group convolution: Group convolution generalizes translation-equivariant convolution to functions defined on a transformation group using the group’s Haar measure.The background introduces this under a left-equivariance convention.
- Continuous data: For non-grid data, continuous convolution integrates learned filters over coordinate offsets and approximates the integral by discretization.A neural network can parameterize filters that map offsets to filter matrices.
- Computational efficiency: 20 billion values are required to evaluate one typical PointConv layer for CIFAR-10 with batchsize = 32, N = 32 × 32, cin = cout = 256, and n = 3×3.The PointConv trick reduces memory and computational requirements by ∼2 orders of magnitude.
4. Convolutional Layers on Lie Groups
LieConv extends convolution to arbitrary continuous spatial data by lifting inputs to Lie-group elements, learning kernels through Lie algebras, enforcing invariant local neighborhoods, and estimating integrals with equivariant Monte Carlo sampling.
- Overview: LieConv applies to arbitrary collections of coordinates and values, including images, molecules, point clouds, mechanical systems, and other spatial data.The framework is designed for domains such as R2 or R3 and does not require a regular grid.
- Kernel parameterization: The convolutional kernel is represented by an MLP on Lie-algebra coordinates, kθ(u) = ˜kθ(log u), for groups with surjective exponential maps.This replaces separate parameters for enumerated group elements with a continuous kernel function.
- Locality: Locality is enforced by evaluating the kernel only within an invariant neighborhood defined by a group distance d(u,v) based on the matrix logarithm and Frobenius norm.Left invariance ensures that transforming the inputs preserves the corresponding neighborhood, unlike fixed regions such as a 3 × 3 image patch.
- Discretization: A Monte Carlo estimator approximates the group convolution integral using samples from the local neighborhood, and is equivariant in distribution.Equivariance also holds deterministically when sampling locations are transformed along with the function.
- Discretization: The discretized convolution remains equivariant for coordinate-value inputs, and equivariant subsampling methods can further reduce computation.The computation can also be accelerated with the Efficient-PointConv trick.
- Lifting: Lifting maps input points into group elements, retaining orbit information when the group does not act transitively on the input space.For homogeneous spaces, lifting uses an origin and elements that map the origin to each point; stabilizer elements generate additional lifts and may be sampled with Haar measure.
3. X = Rd and G = T(k) : Embed(q(x)) = x[k+1:d]
For non-transitive translation groups, the method represents inputs with orbit information and incorporates both group coordinates and orbit coordinates into the learned convolutional filter.
- Scope: This construction addresses equivariance when the input space is not a homogeneous space of the group.The quotient Q = X/G contains distinct orbits rather than a single orbit shared by all inputs.
- Kernel inputs: The MLP can receive log(v−1u), qi, and qj, allowing the filter to use relative group coordinates together with orbit identifiers.An extended distance can additionally impose locality over orbits, although its orbit component need not be transformation-invariant.
- Lifting: The lifting algorithm chooses an origin and stabilizer for each orbit, then constructs group elements mapping sampled orbit points to each input coordinate.Each lifted tuple contains a group element, an orbit identifier, and the associated feature value.
- Convolution layer: The Lie-group convolution layer operates on matrix-orbit-value tuples and forms neighborhoods using the group distance.The algorithm computes each neighborhood before accumulating contributions from its members.
5. Applications to Image and Molecular Data
LieConv applies a shared bottleneck ResNet architecture to image and molecular data, representing spatial inputs as coordinates and values while flexibly incorporating group equivariance. It achieves competitive or state-of-the-art performance across these tasks, with SE(3) performing best in the HOMO ablation.
- Architecture: LieConv uses the same bottleneck ResNet architecture across image and molecular tasks, with a final global pooling layer.The architecture contains LieConv layers within bottleneck blocks.
- Data representation: Images are represented as 784 coordinate-value points in R2, while QM9 molecules use 3D atom coordinates and atomic charges.QM9 targets molecular properties invariant to translation, rotation, and reflection.
- Image data: LieConv incorporates equivariance to different groups without changing the method or network architecture, while remaining competitive with image-specialized methods.The image experiments use RotMNIST, whose digits are randomly rotated with rotations sampled from SO(2).
- Molecular data: On QM9, LieConv performs competitively with state-of-the-art methods and attains the lowest MAE on several molecular-property tasks.The HOMO ablation evaluates different equivariance groups, with the SE(3) network performing best.
6. Modeling Dynamical Systems
LieConv is applied to Hamiltonian dynamical systems to enforce conservation through continuous symmetries. Changing the imposed symmetries controls which momenta are conserved, while the method improves rollout and generalization performance relative to baselines.
- Modeling dynamical systems: LieConv models Hamiltonian dynamics by enforcing symmetries associated with conserved physical quantities such as linear and angular momentum.The approach evaluates particle systems connected by springs against FC, OGN, and HOGN models.
- Hamiltonian symmetries: Noether’s theorem links continuous symmetries in the Hamiltonian to corresponding conserved quantities, including momentum under translations and rotations.Hamiltonian systems use generalized coordinates and momenta with scalar Hamiltonian H.
- Exact conservation of momentum: Linear and angular momentum are conserved with relative error close to machine epsilon when LieConv imposes T(2) and/or SO(2) symmetry.The reported error is determined by the integrator tolerance.
- Exact conservation of momentum: Changing Hamiltonian-model invariances directly controls conservation of linear and angular momentum in predicted trajectories.The figure reports momentum behavior for LieConv models with different imposed symmetries.
- Results: LieConv outperforms HOGN on dynamics results and improves test MSE across training-set sizes as Hamiltonian, graph-network, and equivariance biases are added.Figure 8 reports 95% confidence intervals estimated across 3 trials.
7. Discussion
The discussion presents LieConv as a flexible convolutional layer for diverse data types and equivariance choices. It frames broader modalities and systems without exact conservation as directions for future work.
- Discussion: LieConv networks can handle diverse data types while allowing the model’s equivariance to be swapped flexibly.The demonstrated domains are images, molecular data, and dynamical systems.
- Future work: Future applications named for this approach include time-series, geostatistics, audio, meshes, and other heterogeneous or higher-dimensional data.The paper describes these as application domains and directions for future work.
- Future work: The authors suggest that HLieConv inductive biases may also benefit control and reinforcement-learning systems that do not exactly preserve energy or momentum.This is presented as a belief about possible future use rather than an established experiment.
- Motivation: More general data may be non-grid-structured, have different symmetries, or contain quantities that are difficult to combine using image-oriented representations.The discussion motivates decoupling dense-array computation from geometrically sensible operations.
A. Derivations and Additional Methodology
LieConv extends convolution to continuous group actions by operating in Lie-algebra coordinates, while generalized computation and equivariant subsampling make the layer practical across diverse groups and data.
- Generalized PointConv: The generalized PointConv computation reorders tensor operations to reduce memory and compute when intermediate channels are fewer than input-output channel pairs.The condition is |γ| < |α||β|; the resulting operations can use matrix multiplication after flattening channel indices.
- Coordinate Transforms: For Abelian homogeneous groups, LieConv becomes Euclidean convolution in logarithmic group coordinates, with coordinate-transform methods as special cases.This framework includes log-polar coordinates for rotation and scaling and hyperbolic coordinates for squeezes and scaling.
- Coordinate Transforms: The log-polar construction represents positive scalings and rotations through logarithmic radial and angular coordinates, with angular values projected modulo 2π.The resulting convolution is equivariant to the corresponding scaling and rotation transformations.
- General Framework: The same approach extends beyond single-orbit Abelian actions to groups that are larger or smaller than the input space, including equivariance to squeezes and scalings.The paper explicitly identifies coordinate-transform equivariance as a special case of the broader framework.
- Group Geometry: The distance d(u,v)=∥log(v^-1u)∥F is a geodesic distance under commuting-generator conditions, but is otherwise only a semi-metric because the triangle inequality may fail.The stated sufficient condition is that every generator commute with its transpose; SE(d) does not satisfy this property in general.
- Equivariant Subsampling: Quadratic convolution cost O(N^2) can be reduced to O(pN) by equivariantly subsampling a fixed number p of neighborhood quadrature points.The paper describes random selection and farthest point sampling as the two subsampling methods.
B.1. Equivariance Demo
The equivariance demo evaluates LieConv under transformed inputs and describes practical implementation choices, showing that equivariant models retain performance under the corresponding group transformations.
- Equivariance Evaluation: LieConv models trained on only 20k QM9 examples were evaluated on test sets randomly transformed by translations, rotations, and related groups without data augmentation.The experiment compares Trivial, T(3), SO(3), and SE(3) models on the HOMO task.
- Equivariance Evaluation: LieConv-G performance does not degrade when random G transformations are applied to the test set, with added equivariances especially important in the low-data regime.The accompanying table reports test MAE in meV for randomly transformed HOMO examples without data augmentation.
- Image Experiments: The image experiments compare LieConv group equivariances and baselines on RotMNIST classification error, using both standard dataset splits and an additional full-data setting with augmentation.The paper notes that some compared methods fold the validation set into training after tuning, whereas the main table uses 10k training examples.
- Implementation: For practical implementation, group elements and neighborhood coefficients can be precomputed, Monte Carlo neighborhoods can be sampled, and analytic exponential and logarithm maps can be used.SE(2), SO(3), and SE(3) require sampling multiple lifts for a point; SE(d) lifts use translations combined with uniformly sampled rotations.
- Implementation: The architecture uses ResNet-style bottleneck blocks with Swish activations and LieConv kernels parameterized by three-layer MLPs.Swish is twice differentiable, which the paper identifies as necessary for backpropagating through Hamiltonian dynamics.
C.4. Details for Hamiltonian Models Model Symmetries:
The Hamiltonian experiments use symmetry-aware models for spring dynamics, emphasizing conservation properties while accounting for model invariance, centering, and numerical constraints.
- Model Symmetries: Centering positions gives HLieConv-SO2* additional T(2) invariance and therefore SE(2) invariance, while the authors separately benchmark centered and uncentered variants.The authors state that invariance is sufficient for conservation laws, whereas layer equivariance can preserve information for modeling.
- Model Symmetries: Hamiltonian conservation requires both modeling the Hamiltonian and incorporating the relevant symmetry.The model table records invariance and equivariance properties for the relevant models and baselines.
- Dataset Generation: The spring datasets contain systems of six particles connected by springs, with masses and spring constants sampled for each system.Ground-truth trajectories are generated from the analytic Hamiltonian using RK4 integration and divided into transition segments.
- Training: The Hamiltonian models are trained with single-precision arithmetic, integrator tolerance 1e-4, cosine learning-rate decay, validation-MSE early stopping, and Adam.Training uses minibatches of 200 for 100 epochs; the HLieConv model takes about 20 minutes on a 1080Ti with 3k examples.
C.5. Details for Image and Molecular Experiments
The image and molecular experiments use fixed LieConv implementation settings while visualizing how different transformation groups define local neighborhoods over continuous input data.
- Image Experiments: RotMNIST models use Adam for 500 epochs, six bottleneck blocks, group-element downsampling, and Monte Carlo neighborhoods sampled with p=25 elements.The models use a first linear layer producing 128 channels and train in under 12 hours on a 1080Ti.
- Molecular Experiments: QM9 models use atom-type and charge features with raw atomic coordinates in angstroms, training a separate model for each molecular-property task.The experiments use the split from Anderson et al. and early stopping on validation MAE.
- Neighborhood Visualization: Figure 10 compares local neighborhoods under translations, rotations, and scalings for several groups, with SE(2) producing multiple group-element lifts per point.For several simpler groups, points correspond one-to-one with group-element and orbit tuples; SE(2) instead maps each point to multiple tuples.