LieTransformer: Equivariant self-attention for Lie Groups

Abstract

Group equivariant neural networks are used as building blocks of group invariant neural networks, which have been shown to improve generalisation performance and data efficiency through principled parameter sharing. Such works have mostly focused on group equivariant convolutions, building on the result that group equivariant linear maps are necessarily convolutions. In this work, we extend the scope of the literature to self-attention, that is emerging as a prominent building block of deep learning models. We propose the LieTransformer, an architecture composed of LieSelfAttention layers that are equivariant to arbitrary Lie groups and their discrete subgroups. We demonstrate the generality of our approach by showing experimental results that are competitive to baseline methods on a wide range of tasks: shape counting on point clouds, molecular property regression and modelling particle trajectories under Hamiltonian dynamics.

References in corpus (20)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• The NumPy array: a structure for efficient numerical computation
• Neural Message Passing for Quantum Chemistry
• Language Models are Few-Shot Learners
• Linformer: Self-Attention with Linear Complexity
• Group Equivariant Convolutional Networks
• Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
• Reformer: The Efficient Transformer
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
• Big Bird: Transformers for Longer Sequences
• On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups
• 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data
• Symplectic ODE-Net: Learning Hamiltonian Dynamics with Control
• General $E(2)$-Equivariant Steerable CNNs
• Hamiltonian Graph Networks with ODE Integrators
• Probabilistic symmetries and invariant neural networks
• Relevance of Rotationally Equivariant Convolutions for Predicting Molecular Properties
• Co-Attentive Equivariant Neural Networks: Focusing Equivariance On Transformations Co-Occurring In Data

Cited by in corpus (11)

• A Comprehensive Survey on Deep Graph Representation Learning
• An Efficient Lorentz Equivariant Graph Neural Network for Jet Tagging
• Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning
• CKConv: Continuous Kernel Convolution For Sequential Data
• Coordinate Independent Convolutional Networks -- Isometry and Gauge Equivariant Convolutions on Riemannian Manifolds
• EBM-Fold: Fully-Differentiable Protein Folding Powered by Energy-based Models
• E(n) Equivariant Normalizing Flows
• Group Equivariant Stand-Alone Self-Attention For Vision
• Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes
• Equivariant Networks for Pixelized Spheres
• Autoequivariant Network Search via Group Decomposition