SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks

Abstract

We introduce the SE(3)-Transformer, a variant of the self-attention module for 3D point clouds and graphs, which is equivariant under continuous 3D roto-translations. Equivariance is important to ensure stable and predictable performance in the presence of nuisance transformations of the data input. A positive corollary of equivariance is increased weight-tying within the model. The SE(3)-Transformer leverages the benefits of self-attention to operate on large point clouds and graphs with varying number of points, while guaranteeing SE(3)-equivariance for robustness. We evaluate our model on a toy N-body particle simulation dataset, showcasing the robustness of the predictions under rotations of the input. We further achieve competitive performance on two real-world datasets, ScanObjectNN and QM9. In all cases, our model outperforms a strong, non-equivariant attention baseline and an equivariant model without attention.

References in corpus (10)

• Neural Message Passing for Quantum Chemistry
• A Structured Self-attentive Sentence Embedding
• Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
• 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data
• Spherical CNNs
• General $E(2)$-Equivariant Steerable CNNs
• 3D G-CNNs for Pulmonary Nodule Detection
• N-body Networks: a Covariant Hierarchical Neural Network Architecture for Learning Atomic Potentials
• Deep Scale-spaces: Equivariance Over Scale
• On the Limitations of Representing Functions on Sets

Cited by in corpus (43)

• Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges
• Everything is Connected: Graph Neural Networks
• E(n) Equivariant Graph Neural Networks
• Informing Geometric Deep Learning with Electronic Interactions to Accelerate Quantum Chemistry
• Self-Contrastive Learning with Hard Negative Sampling for Self-supervised Point Cloud Learning
• Machine-learning hidden symmetries
• Choose a Transformer: Fourier or Galerkin
• Rotation Invariant Graph Neural Networks using Spin Convolutions
• Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning
• DeltaConv: Anisotropic Operators for Geometric Deep Learning on Point Clouds
• LieTransformer: Equivariant self-attention for Lie Groups
• Neural Networks Enforcing Physical Symmetries in Nonlinear Dynamical Lattices: The Case Example of the Ablowitz-Ladik Model
• Towards out of distribution generalization for problems in mechanics
• Independent SE(3)-Equivariant Models for End-to-End Rigid Protein Docking
• On the Universality of Rotation Equivariant Point Cloud Networks
• Unsupervised deep learning for semantic segmentation of multispectral LiDAR forest point clouds
• Encoding protein dynamic information in graph representation for functional residue identification
• Robust Partial-to-Partial Point Cloud Registration in a Full Range
• CKConv: Continuous Kernel Convolution For Sequential Data
• Machine learning and invariant theory
• Geometric and Physical Quantities Improve E(3) Equivariant Message Passing
• Correspondence-Free Point Cloud Registration with SO(3)-Equivariant Implicit Shape Representations
• Leveraging SE(3) Equivariance for Self-Supervised Category-Level Object Pose Estimation
• Frame Averaging for Invariant and Equivariant Network Design
• Symmetry-Aware Actor-Critic for 3D Molecular Design
• Unraveling the Potential of Diffusion Models in Small Molecule Generation
• PEMP: Leveraging Physics Properties to Enhance Molecular Property Prediction
• Learning the Crystal Structure Genome for Property Classification
• Probabilistic Numeric Convolutional Neural Networks
• Scaling Graph Neural Networks to Large Proteins
• Trajectory Prediction using Equivariant Continuous Convolution
• Group Equivariant Stand-Alone Self-Attention For Vision
• Learning dissection trajectories from expert surgical videos via imitation learning with equivariant diffusion
• PRIN/SPRIN: On Extracting Point-wise Rotation Invariant Features
• GeoT: A Geometry-aware Transformer for Reliable Molecular Property Prediction and Chemically Interpretable Representation Learning
• Sparse Steerable Convolutions: An Efficient Learning of SE(3)-Equivariant Features for Estimation and Tracking of Object Poses in 3D Space
• Beyond permutation equivariance in graph networks
• Detect the Interactions that Matter in Matter: Geometric Attention for Many-Body Systems
• Autoequivariant Network Search via Group Decomposition
• Equivariant plug-and-play image reconstruction
• Knowledge Distillation of a Protein Language Model Yields a Foundational Implicit Solvent Model
• Physics-Augmented Learning: A New Paradigm Beyond Physics-Informed Learning
• Rotation Equivariant 3D Hand Mesh Generation from a Single RGB Image