We implemented a custom SE(3)-Transformer in PyTorch using the Deep Graph Library (DGL) [1]} and made it available on Github [2]}. For details about the original SE(3)-Transformer implementation, refer to Fuchs et al. [3]}. The differences in the this implementations include an attention block without skip connections and usage of DGL's message passing functions. As shown in Figure 1, we constructed a basic network with three SE(3)-Transformer multi-head attention layers. Each multi-head attention layer passes its input vector to each of its heads. The layer then concatenates the outputs from all the heads. The concatenated outputs are then shrunk using an equivariant linear layer and the output is returned. We trained six different SE(3)-Transformers with architectures and test performances summarized in Table 1.
