Equivariant Network Showdown: GATr vs SE(3)-Transformer vs SEGNN vs EGNN
In 3D geometric deep learning, a dozen+ equivariant network architectures emerged between 2021 and 2024. All claim to preserve rotation/translation invariance, but their implementations, use cases, and performance profiles differ dramatically. This post uses the same benchmark data to build a clear decision map.
The Four Contenders
| Architecture | Year | Core Idea | Group | Complexity | |---|---|---|---|---| | EGNN | 2021 | Encodes distances directly in message passing | E(n) | O(N²) | | SE(3)-Transformer | 2020 | Spherical harmonics + self-attention | SE(3) | O(N² × L³) | | SEGNN | 2022 | Steerable vectors + message passing | SE(3) | O(N² × L²) | | GATr | 2023 | Geometric algebra + Transformer | E(3) | O(N²) |
(N = number of nodes, L = feature order / spherical harmonic order.)
Round 1: N-Body Dynamics Prediction
Predict positions of 5 particles 1000+ timesteps ahead under rotation/translation.
Results (MSE, lower is better)
| Model | 1000 ts | 1500 ts | 2000 ts | Forward Time (ms) | |---|---|---|---|---| | Linear | 6.831 | 20.012 | 39.513 | 0.1 | | GNN | 1.077 | 5.059 | 10.591 | 2.1 | | Radial Field | 1.060 | 12.514 | 26.388 | 3.0 | | EGNN | 0.716 | 2.201 | 4.049 | 5.5 | | SE(3)-Transformer | 2.483 | 18.891 | 36.730 | 114.6 | | TFN | 1.544 | 11.116 | 23.823 | 27.2 | | SEGNN | 0.481 | 1.552 | 3.294 | 24.5 | | GATr | 0.320 | ~0.52 | ~1.0 | ~20 |
Sources: SEGNO (ICLR 2024) and GATr (NeurIPS 2023) papers.
Analysis
- EGNN — best value for money: near-SEGNN accuracy at O(N²) cost by injecting distances as scalar inputs in message passing, \(m_{ij} = \phi_e(h_i, h_j, ||x_i - x_j||^2)\). Downside: weaker expressivity ceiling, less elegant than harmonic-based approaches.
- SE(3)-Transformer — great theory, poor practice: cleanest math via spherical-harmonic attention, but O(L³) harmonic cost makes it 20x slower than EGNN (114.6 ms). Fine for prototyping new equivariant mechanisms; not for production.
- SEGNN — king of small datasets: steerable-vector message passing captures directional information; beats EGNN but is 4–5x slower. Best for scientific computing with < 10K samples (molecular simulation, materials science).
- GATr — all-round top performer: best accuracy while 5x faster than SE(3)-Transformer; geometric algebra makes equivariant linear layers much simpler (16-dim matrix products vs harmonic convolutions). Suited to scalable geometric tasks (large point clouds, robot planning).
- Scarce training data (hundreds of samples)
- Strict physical consistency required (molecular dynamics)
- Generalization to rotated/translated test data (robot planning)
- 3D point cloud classification/segmentation (e.g., ModelNet40): > 10K samples → PointNet++ / Point Transformer (equivariance not needed); otherwise → EGNN (simplest) or GATr (best accuracy).
- Molecular dynamics / quantum chemistry: strict conservation required → SEGNN (most stable) or GATr (most general); otherwise → PaiNN / DimeNet++ (faster specialized architectures).
- Robot motion planning: generative modeling (diffusion) needed → GATr (validated diffusion backbone); otherwise → EGNN (better real-time performance).
- High-energy physics / particle simulation: Lorentz equivariance needed → L-GATr (2024 extension); otherwise → SEGNN.
- Large-scale 3D scene understanding (e.g., autonomous driving): > 100K points → LaB-GATr (geometric tokenization) or standard Transformer; otherwise → GATr.
- Fastest to get started: EGNN — 200 lines of code and running.
- Most mathematically correct: SE(3)-Transformer — but budget 20x the time.
- Unbeatable on small datasets: SEGNN — results from just ~500 samples.
- Best long-term investment: GATr — geometric algebra is a universal language that extends everywhere.
Round 2: Molecular Property Prediction (QM9)
QM9: 13K small molecules, 12 quantum-chemical properties (MAE, lower is better).
| Model | α (bohr³) | Δε (meV) | μ (D) | C_v (cal/mol·K) | |---|---|---|---|---| | SchNet | 0.297 | 71 | 0.049 | 0.043 | | DimeNet++ | 0.047 | 45 | 0.028 | 0.027 | | PaiNN | 0.049 | 49 | 0.029 | 0.028 | | SphereNet | 0.056 | 54 | 0.031 | 0.027 | | EGNN | 0.071 | 48 | 0.028 | 0.032 | | SEGNN | 0.077 | 46 | 0.033 | 0.035 | | SE(3)-Trans. | 0.142 | 53 | 0.051 | 0.054 |
Sources: DuSEGO (2024), SEGNO (2024).
Equivariance Is Not a Silver Bullet
EGNN and SEGNN are not SOTA on QM9 — specialized non-equivariant architectures (DimeNet++, PaiNN, SphereNet) win. With 13K training samples, networks can *learn* rotational regularity; fixed organic small molecules don't require arbitrary 3D transforms; and purpose-built interaction layers fit molecular physics better.
Where equivariant networks genuinely excel:
Round 3: Data Efficiency
Equivariance's core selling point: same accuracy with less data. A standard Transformer must learn physics *plus* rotation/translation invariance from its samples; equivariant architectures have symmetries built into the structure, so every sample goes toward learning the physics itself — roughly like shrinking a 100-page syllabus to 10 pages.
Decision Tree: Which Should You Pick
Architecture Philosophy
| Dimension | EGNN | SE(3)-Transformer | SEGNN | GATr | |---|---|---|---|---| | Math foundation | Distance metrics | Spherical harmonics | Steerable representations | Geometric algebra | | Design philosophy | Pragmatism | Mathematical purity | Physics-driven | Unified framework | | Learning curve | ★★☆☆☆ | ★★★★★ | ★★★★☆ | ★★★☆☆ | | Code complexity | Low (~200 lines) | High (~2000 lines) | Medium (~800 lines) | Medium (~1000 lines) | | Community | Mature | Research-only | Small | Growing | | Scalability | Medium | Poor | Medium | Good |
Hidden Costs
Learning cost: EGNN ~2 hours (writable straight after the paper); SE(3)-Transformer ~2 weeks (harmonics, Clebsch-Gordan coefficients, representation theory); SEGNN ~1 week (steerable vectors, Wigner D matrices, irreps); GATr ~3 days (grades, multivectors, geometric product).
Debugging cost: EGNN debugs like a normal GNN; SE(3)-Transformer suffers from numerically unstable harmonics; SEGNN's steerable-vector shape mismatches are a nightmare; GATr's geometric product dimensions are error-prone.
Maintenance: EGNN largest community; SE(3)-Transformer largely unmaintained (e3nn covers some gaps); SEGNN author-maintained with infrequent updates; GATr officially maintained by Qualcomm AI Research with the best docs.
Future Trends
| Direction | EGNN | SE(3)-Trans. | SEGNN | GATr | |---|---|---|---|---| | E(3) extension | ✗ (E(n)) | ✗ (SE(3)) | ✗ (SE(3)) | ✓ (E(3)) | | Conformal extension | ✗ | ✗ | ✗ | ✓ (CGA) | | Lorentz extension | ✗ | ✗ | ✗ | ✓ (L-GATr) | | LLM extension | ✗ | ✗ | ✗ | partial (Versor, exploratory) | | O(N) complexity | ✗ | ✗ | ✗ | ✓ (Versor RRA) |