[Paper Review] SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks

Fabian B. Fuchs, Daniel E. Worrall+2 more|arXiv (Cornell University)|Jun 18, 2020

3D Shape Modeling and Analysis43 references223 citations

TL;DR

Introduces SE(3)-Transformer, an SE(3)-equivariant self-attention model for 3D point clouds and graphs, enabling robust rotation/translation handling and competitive performance on N-body, ScanObjectNN, and QM9.

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.

Motivation & Objective

• Motivate robust learning on 3D point clouds/graphs under nuisance 3D roto-translations.
• Develop an SE(3)-equivariant self-attention mechanism that preserves pose information.
• Achieve scalability to large, irregular point sets via neighbourhoods and attention.
• Demonstrate improved robustness and accuracy over non-equivariant and non-attentive baselines.
• Provide an efficient implementation and open-source code for broader use.

Proposed method

• Define neighbourhoods on point clouds to form a graph.
• Design SE(3)-invariant edge attention weights and SE(3)-equivariant edge messages.
• Incorporate an attention mechanism that preserves SE(3) equivariance across layers.
• Introduce a self-interaction (linear or attentive) layer to combine intra-node features.
• Provide an attentive self-interaction variant to enhance expressiveness while maintaining equivariance.

Experimental results

Research questions

• RQ1 Can SE(3)-equivariant attention improve robustness to rotations/translations in 3D data?
• RQ2 Does attention add value beyond tensor-field network convolutions in SE(3) settings?
• RQ3 How does the model scale with neighborhood size and input point count?
• RQ4 What is the trade-off between linear and attentive self-interaction in SE(3) frameworks?
• RQ5 Is the approach competitive on real-world 3D tasks like object classification and molecular property prediction?

Key findings

• SE(3)-Transformer outperforms a non-equivariant attention baseline (Set Transformer) and an equivariant model without attention across tasks.
• The model achieves rotation-equivariant predictions with high fidelity (Delta_EQ near zero in several experiments).
• On N-body simulations, SE(3)-Transformer attains the best MSE for position and velocity and demonstrates strong equivariance.
• On ScanObjectNN, adding a z-coordinate as an input boosts accuracy when rotations vary, showing flexibility to symmetry assumptions.
• On QM9, the SE(3)-Transformer yields competitive mean absolute error against other equivariant models and non-equivariant baselines.
• The implementation includes a fast, GPU-accelerated spherical harmonics module improving training and inference efficiency.