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

Code

• github.com/FabianFuchsML/se3-transformer-public
  OfficialIn paperpytorch★ 574
• github.com/rosettacommons/rosettafold
  pytorch★ 2,236
• github.com/lucidrains/se3-transformer-pytorch
  pytorch★ 325
• github.com/trachote/predict_elastic_tensor
  pytorch★ 1
• github.com/NVIDIA/DeepLearningExamples/tree/master/DGLPyTorch/DrugDiscovery/SE3Transformer
  pytorch★ 0

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.

Tasks

Reproductions