Search papers, labs, and topics across Lattice.
GEqTrain is a configuration-driven framework designed to enhance the reuse of equivariant graph neural networks (GNNs) across diverse 3D scientific tasks by decoupling dataset semantics, model composition, and training objectives. By utilizing Hydra configurations, the framework allows for the retargeting of a shared equivariant backbone to new tasks with minimal adjustments, achieving competitive accuracy in applications such as biomolecular backmapping and NMR chemical shift prediction. Additionally, GEqDiff extends this framework to generative modeling, enabling the joint reconstruction of complex equivariant fields with high fidelity, thus reducing the software overhead associated with transitioning between different modeling tasks.
A single configuration change can retarget a powerful equivariant GNN framework to tackle diverse 3D scientific challenges with competitive accuracy.
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.
