# se3-transformer-public is the reference code for equivariant attention, with 2020 era dependencies

> se3-transformer-public is the official implementation of SE(3)-Transformers, built on DGL and PyTorch, with n-body, QM9 and a toy optimisation experiment included. It is a paper artefact rather than a maintained library, and the README says so by pointing you at an NVIDIA implementation instead.

**FabianFuchsML/se3-transformer-public** — code for the SE3 Transformers paper: https://arxiv.org/abs/2006.10503

- Repository: https://github.com/FabianFuchsML/se3-transformer-public
- Stars: 585 · Forks: 78
- Language: Python
- License: not declared
- Published: 2026-10-22 · Updated: 2026-10-22 · Language: en
- Canonical page: https://hysenlabs.com/projects/fabianfuchsml-se3-transformer-public

## The reference implementation, and the one the author points you away from

se3-transformer-public is the official implementation of SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks, the NeurIPS 2020 paper by Fuchs, Worrall, Fischer and Welling. The repository description says it holds the code for that paper, and the README asks you to cite the paper with a specific bibtex entry.

What makes the README worth reading is the first section, headed Update October 2021. It tells you to check out work by Alexandre Milesi and colleagues at Nvidia, who managed to speed up training of the SE(3)-Transformer by up to 21 times and reduce memory consumption by up to 43 times, with code in the NVIDIA DeepLearningExamples repository under DGLPyTorch and DrugDiscovery.

So, in the first section of the README, dated October 2021, the author of the official implementation is telling users to check out work elsewhere. That is unusual and worth taking at face value rather than reading past.

The reason is in the FAQ, under Speed. Constructing equivariant layers is computationally expensive, and one of the bottlenecks is the computation of the spherical harmonics. The author describes significant effort spent speeding that up, in part by parallelising it on the GPU, and names the core difficulty: the harmonics depend on relative distances and are therefore different for every point cloud.

That is the real constraint of this architecture, and it is why equivariant attention is expensive. The rest of this article is about what that means for anyone who wants to run the code.

## The dependency set dates the code precisely

The prerequisites section recommends dgl 0.5.3 and torch 1.7, and notes that the code was tested with the combinations dgl0.5.3 and torch 1.7, and dgl0.4.3.post2 and torch 1.4. The DGL install instructions are CUDA version specific, with the older 0.4 line needing a different wheel per CUDA version, for example dgl-cu90, dgl-cu101 and dgl-cu102 at version 0.4.3.post2. For the recommended line the README just says pip install dgl-cu110.

You can check which CUDA version your PyTorch is using with a one-liner, and the README links to the DGL install page for compatibility.

The requirements.txt file confirms the vintage. It pins absl-py 0.9.0, numpy 1.18.2, scipy 1.4.1, scikit-learn 0.22.2.post1, Cython 0.29.16, protobuf 3.11.3 and a set of other packages from the same period, and the heavy dependencies such as torch, wandb and packaging are commented out at the bottom, because the README asks for them to be installed separately.

The setup.py is barely more than a name and a version.

```python
from setuptools import setup, find_packages

setup(name='equivariant-attention', version='1.0', packages=find_packages())
```

The package is published as equivariant-attention at version 1.0, and the README asks you to install this repo in your virtual environment with pip install -e., adding the parenthetical that this is important for easier importing from parent folders. That is the install story: an editable install of a repository, not a package from an index.

The optional dependencies are worth noting too. Weights and Biases for experiment tracking, the lie_learn library from AMLab Amsterdam, which the README lists as an optional install, and line_profiler for testing the speed of different parts of the code.

## Fibers are how the model describes what it is allowed to do

The basic usage section is the most useful part of the README for understanding the architecture, and it starts with a toy model.

The central object is Fiber, described in the code comment as a representation of the structure of the activations. Its first argument is the number of degrees, zero based, and the second is the number of channels, called multiplicities, for each degree.

```python
# The maximum feature type is harmonic degree 3
num_degrees = 4
```

So a fiber with four degrees carries feature types 0, 1, 2 and 3, which in the comment's terms means the maximum feature type is harmonic degree 3. The code then builds three of them, one for the input, one for the middle of the network with 32 channels, and one for the output with 128 channels.

```python
fiber_in = Fiber(1, num_features)
fiber_mid = Fiber(num_degrees, 32)
fiber_out = Fiber(1, 128)
```

This is the mechanism that makes the network equivariant. Rather than passing a single tensor of activations around, the network carries a set of feature types, and the structure of that set says which operations are allowed to mix them. The comment notes it is possible to have a varying number of channels per feature type and to use arbitrary feature types, with the pointer to fibers.py for that functionality.

The model itself is then assembled as an ordered list of four blocks, and the comments enumerate them: a multihead attention block, a nonlinearity, a TFN layer without attention, and graph max pooling, with a fully connected layer to one output afterwards.

```python
model = nn.ModuleList([GSE3Res(fiber_in, fiber_mid),
                       GNormSE3(fiber_mid),
                       GConvSE3(fiber_mid, fiber_out, self_interaction=True),
                       GMaxPooling()])
fc_layer = nn.Linear(128, 1)
```

The naming is telling. GSE3Res is the residual attention block, GNormSE3 the normalisation, GConvSE3 the equivariant convolution.

## The forward pass computes a basis once, then passes it down

The second half of the usage example shows the data flow, and the design decision in it is the reason the model is expensive.

Before each forward pass you call get_basis_and_r, which the comment describes as computing the equivariant weight basis and relative positions of all the nodes in the graph. The variables are then passed as keyword arguments to the SE(3)-transformer layers.

```python
basis, r = get_basis_and_r(G, num_degrees-1)
```

The activations themselves are passed around as a dict, where the key is the feature type as an integer in string form, and the value is a Pytorch tensor in the DGL node feature representation.

```python
features = {'0': G.ndata['my_features']}
for layer in model:
    features = layer(features, G=G, r=r, basis=basis)
```

The final non-DGL layer works because GMaxPooling has already converted the features from the DGL node feature representation into the standard Pytorch tensor representation, so the fully connected layer can consume it.

The point to take away is that the basis and the relative positions are computed once per forward pass and threaded through every layer, rather than each layer recomputing them. That is an optimisation, and it comes with a constraint. The basis depends on the relative distances between nodes, which means it is specific to one graph, one point cloud. You cannot reuse it across samples without recomputing, which is exactly the bottleneck the author describes and the reason caching the spherical harmonics and basis vectors is listed as a promising direction.

It also means the batch structure of the experiments matters, because the per-sample basis has to be built and carried.

## What is included, and the dataset quirk that bites people

Experiment specific code is meant to go in the experiments/ directory, and the README says the project provides an implementation for the n-body experiment, the QM9 experiments, and a multi-step toy optimisation experiment loosely inspired by the protein structure prediction task. The repository layout matches, with an equivariant_attention/ directory for the model code, experiments/, utils/ and setup.py.

There is one documented gotcha, and it is in the FAQ under type issues with QM9 experiments. A user reported data type problems, and the fix was to add three lines just before line 184 of qm9.py.

```python
x=x.astype(np.float32)
one_hot=one_hot.astype(np.float32)
atomic_numbers=atomic_numbers.astype(np.float32)
```

Three separate arrays cast to float32, which is a good sign that the rest of the pipeline expects float32 tensors and the loader was not guaranteeing it. It is also a sign that the code assumes a fairly specific environment, since the fix is applied by hand in the experiment script rather than in the library.

The experiments are the point of the repository, since the model code alone is not much use without a task to run it on. QM9 is the molecular dataset, n-body is the synthetic setting, and the toy optimisation is a stand-in for the protein structure task that the paper was motivated by.

## Where the 21 times came from, and where you should start

The NVIDIA work that the README points at is the more useful reference for practical purposes, and the difference in approach is instructive.

The original code parallelised the spherical harmonics computation on the GPU, and the author is candid in the FAQ that constructing equivariant layers is expensive, that the harmonics are different for every point cloud, and that further speeding up this family of approaches and making them more memory efficient would be a big step forward.

The author's own explanation of why translation equivariance appears in so many architectures, whether or not the task has translational symmetry, is computational efficiency rather than any deep symmetry argument. Other symmetries, in that reading, cost more than they return. That is an argument you should weigh when you decide whether equivariant attention is the right model for your problem at all, and it is the author making it, not an outside reviewer.

The suggested direction for speeding things up is caching. Depending on the dataset and your system requirements, it might be possible to cache the spherical harmonics and basis vectors, which could speed things up tremendously because it addresses all the bottlenecks at once. The README notes that the line of thought continues past that point, and the text is cut off mid sentence, so the remaining ideas are not visible.

For a practical recommendation: the original repository is the thing to read if you want to understand the paper's formulation of the equivariant weight basis, and the NVIDIA implementation is the thing to run if you want results.

## Conclusion

Adopt se3-transformer-public if you are reproducing the SE(3)-Transformers paper, teaching equivariant attention, or you need the original formulation of the equivariant weight basis rather than a faster reimplementation. Do not adopt it as a library for new work, since the pinned dependency set is from 2020 and the README itself points at a faster NVIDIA implementation. Before you build, check that your CUDA version matches the DGL wheel you install, and decide whether the 21 times faster NVIDIA code should be the thing you start from instead.

## FAQ

### What is se3-transformer-public?

It is the official implementation of SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks, the NeurIPS 2020 paper, built on DGL and PyTorch. The repository includes the n-body, QM9 and a toy optimisation experiment alongside the model code.

### How do I install se3-transformer-public?

Install the repository in your virtual environment with pip install -e ., which the README notes is important for easier importing from parent folders. You also need PyTorch, DGL with a wheel matching your CUDA version, and packaging, with the DGL version recommended being 0.5.3 and torch 1.7.

### Is there a faster implementation of the SE(3)-Transformer?

Yes. The README's October 2021 update points to work by Alexandre Milesi and colleagues at Nvidia, who sped up training of the SE(3)-Transformer by up to 21 times and reduced memory consumption by up to 43 times, with the code in the NVIDIA DeepLearningExamples repository.

### What is a Fiber in the SE(3)-Transformer code?

The Fiber object represents the structure of the activations. Its first argument is the number of degrees, zero based, so num_degrees of 4 leads to feature types 0, 1, 2 and 3, and its second argument is the number of channels for each degree.

### Why is the SE(3)-Transformer slow to train?

The README says constructing equivariant layers is computationally expensive, with one bottleneck being the computation of the spherical harmonics, which depend on relative distances and are therefore different for each point cloud or graph. Caching the spherical harmonics and basis vectors is suggested as a way to address the bottlenecks.

## Sources

- [FabianFuchsML/se3-transformer-public on GitHub](https://github.com/FabianFuchsML/se3-transformer-public)
- [Issues](https://github.com/FabianFuchsML/se3-transformer-public/issues)
- [README](https://github.com/FabianFuchsML/se3-transformer-public/blob/master/README.md)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/fabianfuchsml-se3-transformer-public