PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling

arXiv:2006.10503v3cs.LGstat.ML

TL;DR

SE(3)-Transformer addresses the challenge of learning from variable-sized 3D point clouds and graphs while respecting roto-translation structure. It builds equivariant self-attention with invariant attention weights and equivariant value embeddings, and experiments report consistent gains over both non-equivariant attention and equivariant models without attention.

  • Problem

    Learning from irregular, variable-sized 3D point clouds requires handling global rotations and translations without changing the output inappropriately.

  • Method

    The SE(3)-Transformer combines self-attention with invariant attention weights and equivariant value embeddings for 3D point clouds and graphs.

  • Results

    The model consistently outperformed conventional attention and equivariant models without attention across the reported experiments, with equivariance providing mathematical robustness to input rotations.

  • Takeaways & Limitations

    Equivariant attention supports robust point-cloud and graph processing while retaining data-adaptive filters and scalability through neighborhoods.

  • Takeaways & Limitations

    Tensor field network comparison kernels impose strong angular constraints, leaving only radial learnable degrees of freedom.

Abstract

from arXiv · show

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.

1 Introduction

The paper introduces SE(3)-Transformer, a self-attention mechanism for 3D point clouds and graphs that enforces equivariance to roto-translations. This design aims to improve robustness while retaining self-attention’s flexibility and applicability to variable-sized data.

  • Motivation: SE(3)-Transformer adapts self-attention specifically to 3D point clouds and graphs while enforcing equivariance constraints.The method targets nuisance transformations such as global rotations and translations.
  • Motivation: Equivariance makes transformations of the input manifest as equivalent transformations of the output and reduces the space of learnable functions.The paper connects this restriction to increased weight-tying and richer invariance.
  • Contributions: The mechanism guarantees invariance to global rotations and translations while remaining equivariant to permutations of input point labels.
  • Contributions: The model addresses angularly constrained filters in concurrent SE(3)-equivariant neural networks.
  • Implementation: The implementation includes a spherical-harmonics speedup of up to 1000× on GPU and approximately 22× for one ScanObjectNN forward-pass comparison.

2 Background And Related Work

The background frames self-attention, permutation equivariance, graph neighborhoods, and SE(3)-equivariant representations as foundations for the proposed architecture. It also identifies angular constraints in tensor field network kernels as a relevant limitation.

  • Self-attention: Self-attention uses query, key, and value embeddings to compute outputs from input-dependent attention weights and value embeddings.
  • Self-attention: Permutation equivariance ensures that permuting point labels permutes the output without making attention depend arbitrarily on input ordering.
  • Graph attention: Neighborhoods replace all-to-all attention with nearest-neighbor attention, reducing quadratic scaling and representing the point cloud as a graph.
  • Equivariance: A function is equivariant when transforming its input and then applying the function equals applying the function and transforming its output.
  • Representations: For 3D rotations, type-0 features are invariant and type-1 features transform according to ordinary 3D rotation matrices.
  • Tensor field networks: Tensor field networks map point clouds equivariantly using kernels whose angular form is fixed by equivariant basis functions, leaving learnable freedom primarily in the radial direction.
  • Tensor field networks: This angular constraint motivates attention-based mechanisms that can add data-dependent angular modulation to equivariant kernels.

3 Method

The SE(3)-Transformer combines invariant attention weights, equivariant value messages, graph neighborhoods, and self-interaction layers to process point clouds and graphs equivariantly. Its attention weights also modulate the angular profile of tensor-field kernels while neighborhoods reduce computational complexity.

  • SE(3)-Transformer pipeline: The method constructs graph neighborhoods, equivariant edge functions, SE(3)-equivariant messages, and aggregation steps.Neighborhoods may come from nearest-neighbor methods or predefined structures such as molecular bonds.
  • SE(3)-Transformer pipeline: Neighborhoods reduce attention complexity from quadratic to linear in the number of points.
  • Equivariant attention: Attention weights αij are SE(3)-invariant, while value messages are SE(3)-equivariant and propagate information between neighboring nodes.Invariant weights can be formed from normalized inner products of equivariant query and key vectors.
  • Angular modulation: Attention weights act as data-dependent multipliers of tensor-field kernels, adding degrees of freedom in the angular direction.This provides angular modulation beyond the underlying tensor-field convolution.
  • Self-interaction: The model includes linear and attentive self-interaction layers that exchange information across channels or representation degrees.Attentive self-interaction replaces learned scalar weights with MLP-produced attention weights that are SE(3)-invariant.
  • Node and edge features: Node features enter through point features, while edge features can replace node features or enter scalar radial-network inputs.

4 Experiments

The experiments evaluate SE(3)-Transformer on equivariant particle prediction, real-world point-cloud classification, and molecular regression, comparing it with non-equivariant attention and equivariant non-attention baselines. Results report rotation-equivariant predictions and competitive performance across the evaluated tasks.

  • Experimental setup: The model is evaluated on N-body prediction, ScanObjectNN object classification, and QM9 molecular property regression.The three datasets test equivariance, noisy large point clouds, and rich graph structures, respectively.
  • N-Body Simulations: The N-body task predicts particles’ relative locations and velocities 500 time steps ahead, with rotated inputs requiring correspondingly rotated outputs.The experiment uses five charged particles and formulates the task as direct regression rather than iterative multi-step prediction.
  • N-Body Simulations: The SE(3)-Transformer outperforms both Set Transformer and Tensor Field Network baselines while remaining fully rotation equivariant up to computational precision.Set Transformer supplies attention without rotation equivariance, whereas Tensor Field Network supplies equivariance without attention.
  • N-Body Simulations: Figure 3 compares conventional and rotation-equivariant attention using t = 0 inputs, t = 500 ground truth, predictions, and predictions across 10-degree input rotations.Dashed curves denote predictions from a perfectly equivariant model.
  • ScanObjectNN: On ScanObjectNN, the x-axis represents the maximum test-set rotation around a random axis in the x-y-plane.Without train/test rotations, adding the z-coordinate as a scalar improves performance significantly; under rotation-invariant training, the model learns to ignore it.
  • ScanObjectNN: Table 2 reports ScanObjectNN classification accuracy for the object-only category, averaging SE(3)-Transformer performance over five runs with 0.7% standard deviation.The table concerns the object-only category of the ScanObjectNN benchmark.
  • QM9: Table 3 reports QM9 mean absolute error, separating non-equivariant models from equivariant models and averaging SE(3)-Transformer results over five runs.The table’s two sections distinguish the model families used for comparison.

5 Conclusion

The paper presents an attention-based architecture for point clouds that is robust to rotations and translations, while combining equivariance with data-adaptive attention and scalable neighborhoods. Across experiments, attention improves accuracy and training stability, and equivariance adds performance and a mathematical robustness guarantee.

  • 5 Conclusion: The architecture is guaranteed to be robust to input rotations and translations, reducing the need for training-time data augmentation and sensitivity to coordinate-frame choices.The conclusion describes these as consequences of the architecture’s design.
  • 5 Conclusion: Self-attention provides anisotropic, data-adaptive filters, while neighborhoods support scalability to large point clouds.The architecture also treats attention as a data-dependent nonlinearity within equivariant networks.
  • 5 Conclusion: Adding attention to a roto-translation-equivariant model consistently improves accuracy and training stability, with attention essential for convergence in large neighborhoods.The comparison is reported against an equivariant model without attention.
  • 5 Conclusion: Adding equivariance constraints to conventional attention improves performance in all experiments while providing a mathematical guarantee of robustness to input rotations.This conclusion compares equivariant and conventional attention mechanisms.

Broader Impact

The paper presents a mathematically motivated attention mechanism for point-cloud learning and identifies molecular classification as a potential application area. It highlights possible use in drug research while stressing the need for wet-lab testing.

  • The SE(3)-Transformer is a mathematically motivated attention mechanism for deep learning on point-cloud problems.
  • Its molecular classification and regression suitability creates a potential application in drug research.
  • One investigated application is early-stage classification of molecules for suitability to inhibit the reproductive cycle of coronavirus.
  • Computational screening can filter promising compounds from databases, but this research still requires intensive wet-lab testing.

C Accelerated Computation of Spherical Harmonics

The paper accelerates spherical-harmonics computation to address a scaling bottleneck in irreducible point-cloud networks. Its implementation uses GPU parallelization and memoized recursion for associated Legendre polynomials.

  • Spherical harmonics are a computational bottleneck for point-cloud methods based on irreducible representations, limiting prior approaches or sacrificing detail and exact equivariance.
  • The implementation computes spherical harmonics on the GPU, enabling the SE(3)-Transformer and Tensor Field Network baselines to run in reasonable time.
  • The associated Legendre polynomial component is accelerated with recursive computation and memoization of intermediate results.
  • Stored intermediate computations reduce average associated-Legendre-polynomial computation time by approximately 10x on CPU.

D Experimental Details

The experiments compare SE(3)-Transformers with Tensor Field Network baselines under specified sampling, neighborhood, representation, and training settings. The setup emphasizes stable training and invariant final representations.

  • SE(3)-Transformer and Tensor Field Network: ScanObjectNN experiments used up to 200 points per sample and neighborhoods of up to 40 nearest neighbors.
  • SE(3)-Transformer and Tensor Field Network: The ablation networks used four hidden equivariant layers, five channels each, and representation degrees up to 2.
  • SE(3)-Transformer and Tensor Field Network: For the Tensor Field Network baseline, stable training required three channels, 100 input points, and neighborhood size 10, despite degrees up to 3.
  • SE(3)-Transformer and Tensor Field Network: Quantitative comparisons used 128 sampled points and neighborhood size 10, with degrees up to 3 and five channels per degree for the SE(3)-Transformer.
  • Training setup: The final encoder maps each point to 64 degree-0 channels, followed by point pooling and a 64-dimensional hidden-layer MLP.
  • Input representations: Relative x-y-z neighbor positions are supplied as degree-1 fields, while the '+z' variants deliberately break SE(3) equivariance with additional scalar z-position fields.

D.1.2 Number of Input Points

On ScanObjectNN, increasing the input beyond 128–200 points did not improve performance, and the authors suspect global pooling contributes to this behavior.

  • The authors suspect that global pooling explains why additional input points did not improve performance.
  • With 128 points, the SE(3)-Transformer reached 85.0 ± 0.7%, exceeding PointCNN, PointGLR, and DGCNN under the same point limit.

D.1.3 Sample Complexity

Experiments on ScanObjectNN indicate that equivariance improves performance across training-set sizes, with SE(3)-Transformer outperforming a non-equivariant attention baseline.

  • SE(3)-Transformer outperformed the non-equivariant Set Transformer for all tested numbers of ScanObjectNN training samples.The results support a sample-complexity advantage for the equivariant model.
  • Equivariance is associated with smaller sample complexity, meaning that less training data may be needed.
  • The DeepSet baseline was reduced to 128 hidden units per layer and used no dropout for ScanObjectNN.
  • The Set Transformer baseline uses an induced set attention block encoder, pooling by multihead attention, and an MLP.

D.2 Relational Inference

The relational-inference experiment uses a 3D charged-particle simulation to compare equivariant attention with alternative architectures, including a linear non-interacting baseline.

  • The dataset simulates trajectories of 5 charged, interacting particles in three dimensions, with charges sampled as Bernoulli trials.The simulations used randomly sampled initial positions inside a [−5, 5]3 box and removed bounding boxes during simulation.
  • Each model was trained for 100,000 steps with batch size 128 using Adam and a separately tuned fixed learning rate.
  • The SE(3)-Transformer used 4 equivariant layers with degrees {0, 1, 2, 3}, 3 channels per degree, and attentive self-interaction.
  • The model used one attention head, producing one attention weight per point pair across all channels and degrees.
  • The baselines all preserve permutation invariance, but only the Tensor Field network and linear baseline are SE(3) equivariant.
  • The linear baseline propagates particles independently, so improvements beyond it indicate relational reasoning.

D.3 QM9

The QM9 setup applies equivariant attention to molecular graphs for chemical-property regression, using a multi-layer architecture with TFN components and structured atomic and bond features.

  • QM9 contains 134k organic molecules with up to 29 atoms, five atomic species, four bond types, and atom positions in ångströms.
  • The network uses 7 multihead attention layers, norm nonlinearities, a TFN layer, max pooling, and two linear layers.
  • The TFN baseline uses the same architecture but replaces each multihead attention block with a TFN layer of identical output shape.
  • Node features combine a five-species one-hot encoding with proton count, while edge features combine bond type and interatomic distance.
  • QM9 evaluation reports results on six regression tasks that do not require thermochemical energy subtraction.
  • The QM9 architecture table defines dout as the number of output feature types and C as channels per feature type.
  • Feature tuples such as (4, 32) encode degrees 0 through 3 with 32 channels per feature type.
  • Representation degree affects performance significantly but saturates: degrees {0, 1} to {0, 1, 2} improves performance substantially, while degree 3 adds only small improvements.Higher degrees increase memory usage and computation time; the authors recommend up to degree 2 when resources are constrained.
Loading 2006.10503v3…