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

2026-07-21

AI Summary

Existing equivariant graph neural networks are tightly coupled to specific tasks, outputs, and training pipelines, limiting their reusability across diverse 3D scientific applications. This work proposes GEqTrain, a framework that decouples data semantics, model composition, and training objectives through declarative configuration, enabling rapid adaptation to new tasks by modifying only the configuration while leveraging a unified equivariant backbone and shared training infrastructure. GEqTrain is the first to flexibly support switching between prediction and generation tasks as well as between scalar and tensor outputs, and introduces GEqDiff for joint generation of heterogeneous equivariant fields up to spherical harmonic degree ℓ=3. The framework achieves competitive accuracy across multiple benchmarks—including biomolecular backmapping, solid-state NMR chemical shift prediction, and equivariant generative modeling—and demonstrates high-fidelity reconstruction of diverse equivariant fields on synthetic protein structures using GEqDiff.

Abstract

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.

Problem

• equivariant graph neural networks
• 3D scientific tasks
• model reuse
• configuration-driven framework
• retargeting

Innovation

• equivariant graph neural networks
• configuration-driven framework
• GEqTrain
• equivariant flow matching
• GEqDiff