# GEqTrain: A Configuration-Driven Framework for Retargeting Equivariant Graph Neural Networks Across 3D Scientific Tasks

arXiv: https://arxiv.org/abs/2607.19083 · HTML: https://mlsanity.com/paper/2607.19083/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 59 references that are in the map.
- First posted on arXiv 21 Jul 2026.

## Builds on (most cited first, 20 of 20)

- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Structure-based Drug Design with Equivariant Diffusion Models](https://mlsanity.com/paper/2210.13695/) (2022)
- [Machine Learning Force Fields](https://mlsanity.com/paper/2010.07067/) (2020)
- [Graph neural networks for materials science and chemistry](https://mlsanity.com/paper/2208.09481/) (2022)
- [TorchMD: A deep learning framework for molecular simulations](https://mlsanity.com/paper/2012.12106/) (2020)
- [State-specific protein-ligand complex structure prediction with a multi-scale deep generative model](https://mlsanity.com/paper/2209.15171/) (2022)
- [SchNetPack 2.0: A neural network toolbox for atomistic machine learning](https://mlsanity.com/paper/2212.05517/) (2022)
- [Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for 3D Molecule Generation](https://mlsanity.com/paper/2311.16199/) (2023)
- [Chemically Transferable Generative Backmapping of Coarse-Grained Proteins](https://mlsanity.com/paper/2303.01569/) (2023)
- [DiAMoNDBack: Diffusion-denoising Autoregressive Model for Non-Deterministic Backmapping of Cα Protein Traces](https://mlsanity.com/paper/2307.12451/) (2023)
- [Clifford Group Equivariant Diffusion Models for 3D Molecular Generation](https://mlsanity.com/paper/2504.15773/) (2025)
- [Constraint Decoupled Latent Diffusion for Protein Backmapping](https://mlsanity.com/paper/2410.13264/) (2024)

## Abstract (arXiv)

Equivariant graph neural networks provide a powerful modeling language for three-dimensional scientific data, but their reuse is often limited by implementations tied to specific tasks, outputs, and training regimes. We present GEqTrain, a configuration-driven framework that separates dataset semantics, model composition, and training objectives. Raw data are mapped to typed node-, edge-, and graph-level fields, while model stacks, losses, and training workflows are assembled declaratively through Hydra configurations. A shared equivariant backbone and training infrastructure can therefore be retargeted to a new task primarily through configuration. We demonstrate this flexibility on three different problems handled within one software stack: coarse-grained-to-atomistic backmapping of biomolecular systems, prediction of NMR chemical shifts in molecular solids, and equivariant generative modeling. Our aim is not to surpass individually optimized task-specific systems, but to show that a shared representation and training infrastructure can achieve competitive accuracy across qualitatively different tasks at the cost of a configuration change. We further introduce GEqDiff, a generative extension based on equivariant flow matching. GEqDiff treats user-defined equivariant fields as first-class generation targets, jointly transporting Cartesian positions and non-scalar node fields spanning representations up to l=3 within a single equivariant flow. We validate this capability on a controlled synthetic benchmark inspired by protein secondary-structure motifs, showing that fields with heterogeneous transformation properties can be reconstructed jointly and with high fidelity. By reducing the software overhead of moving between predictive and generative, scalar and tensorial settings, GEqTrain aims to make equivariant modeling more reproducible, extensible, and reusable.