Iterative SE(3)-Transformers

TL;DR

This paper introduces an iterative version of the SE(3)-Transformer, an attention-based model that respects 3D symmetries, and explores its advantages over the single-pass approach in graph data tasks.

Contribution

It presents the first iterative SE(3)-Transformer, addressing challenges of iterative application and comparing its performance to the single-pass version.

Findings

Iterative SE(3)-Transformer can outperform single-pass models in certain tasks.

The paper provides implementation details and code for the iterative model.

Analysis suggests iterative models may better capture complex symmetries in 3D data.

Abstract

When manipulating three-dimensional data, it is possible to ensure that rotational and translational symmetries are respected by applying so-called SE(3)-equivariant models. Protein structure prediction is a prominent example of a task which displays these symmetries. Recent work in this area has successfully made use of an SE(3)-equivariant model, applying an iterative SE(3)-equivariant attention mechanism. Motivated by this application, we implement an iterative version of the SE(3)-Transformer, an SE(3)-equivariant attention-based model for graph data. We address the additional complications which arise when applying the SE(3)-Transformer in an iterative fashion, compare the iterative and single-pass versions on a toy problem, and consider why an iterative model may be beneficial in some problem settings. We make the code for our implementation available to the community.