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

Summary

GEqTrain is a configuration-driven framework that separates dataset semantics, model composition, and training objectives, enabling reuse of equivariant graph neural networks across diverse 3D scientific tasks. Raw data are mapped to typed fields, and models, losses, and workflows are assembled declaratively via Hydra configurations. The framework is demonstrated on three tasks within one software stack: coarse-grained-to-atomistic backmapping of biomolecular systems, NMR chemical shift prediction in molecular solids, and equivariant generative modeling. It also introduces GEqDiff, a generative extension based on equivariant flow matching, which treats user-defined equivariant fields as first-class targets and jointly transports Cartesian positions and non-scalar node fields up to l=3. Validated on a synthetic benchmark inspired by protein secondary-structure motifs, GEqDiff shows high-fidelity joint reconstruction of fields with heterogeneous transformation properties. GEqTrain aims to make equivariant modeling more reproducible, extensible, and reusable.