SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
This work addresses robustness to nuisance transformations in 3D data processing, which is important for applications like molecular modeling and object recognition, though it is incremental as it builds on existing attention and equivariance concepts.
The authors tackled the problem of ensuring stable performance in 3D point cloud and graph models under rotations and translations by introducing the SE(3)-Transformer, an equivariant self-attention network, which outperformed non-equivariant and equivariant baselines on datasets like ScanObjectNN and QM9.
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.