GEqTrain: 一种配置驱动的框架, 用于在三维科学任务中重定向等变图神经网络

摘要

GEqTrain是一个配置驱动框架，通过分离数据集语义、模型组合和训练目标，使等变图神经网络能够跨不同3D科学任务重用。它将原始数据映射为类型化字段，通过Hydra配置声明式组装模型和训练流程。在单一软件栈中演示了三个不同任务：粗粒化到原子级生物分子反映射、分子固体NMR化学位移预测和等变生成建模。还引入了生成扩展GEqDiff，基于等变流匹配，可联合传输笛卡尔位置和最高l=3的非标量节点场，在合成基准上验证了高保真重建。该框架旨在提高等变建模的可复现性和可扩展性。