GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks

简介

GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks

Abstract 摘要

理解图的三维（3D）复杂结构对于准确建模各种属性至关重要，然而许多现有方法在完全捕捉这类系统中固有的复杂空间关系和对称性方面存在困难，尤其是在大规模、动态分子数据集中。这些方法通常必须在表达性和计算效率之间进行权衡，限制了它们的可扩展性。为了解决这一差距，我们提出了一种新的几何张量网络（GotenNet），该网络有效地建模了 3D 图的几何复杂性，同时在欧几里得群 E(3)下确保严格的等变性。我们的方法通过利用有效的几何张量表示，而不依赖于不可约表示或 Clebsch-Gordan 变换，直接解决表达性-效率权衡问题，从而降低计算开销。我们引入了一种统一的结构嵌入，结合了具有几何感知的张量注意力和分层张量细化，通过在高度可操纵特征上的内积运算迭代更新边表示，从而为各种任务提供灵活且高效的表示。我们在 QM9、rMD17、MD22 和 Molecule3D 数据集上评估了模型，所提出的模型在标量和高阶属性预测中均优于现有最佳方法，展示了跨不同数据集的卓越鲁棒性，并将 GotenNet 确立为 3D 等变图神经网络的多功能和可扩展框架。

Performance and Results 性能和结果

QM9 数据集：QM9 数据集是一个研究广泛的包含 130k 个小分子和 12 种不同的分子属性的数据集。

右侧的交互式雷达图展示了GotenNetGotenNet在 QM9 数据集的 12 种分子属性上的性能。您可以通过悬停在不同的点上查看每种属性的精确性能指标。该可视化比较了我们的模型（S、B 和 L）与最先进的基线

视觉化分析的关键观察结果：雷达图的外围代表更好的性能（更低的误差），我们的最大模型GotenNetLGotenNetL始终达到这些外围，证明了其优越的性能。值得注意的是，即使是我们的最小模型变体GotenNetSGotenNetS也优于许多基线，突显了我们的方法的效率。为了交互式探索，您可以通过图例中的不同模型来显示或隐藏其性能曲线，从而实现直接的视觉比较。

在可视化中显示的显著改进（点击比较）：我们所有的模型变体都实现了强大的性能提升。我们最大的模型GotenNetLGotenNetL实现了最先进的性能，α（极化率）降低了 30%的错误，μ（偶极矩）提高了 33%，HOMO-LUMO 间隙提高了 32%的准确性，以及 0K 时的内能降低了 24%的错误。即使是我们的基础模型 GotenNetBB也显示出显著的收益，极化率错误降低了 18%，偶极矩提高了 28%，HOMO-LUMO 准确性提高了 29%，以及U0U_0错误降低了 23%。我们最小的模型 GotenNetS_S保持了有竞争力的性能，极化率错误降低了 18%，偶极矩提高了 25%，HOMO-LUMO 准确性提高了 27%，以及U0_U_0错误降低了 16%。

这个交互式可视化有助于展示GotenNetGotenNet如何在所有属性上实现一致的改进，而不是在不同目标之间权衡性能。

Attention Visualization on QM9 QM9 上的注意力可视化

上方的分子可视化器可以用来发现给定分子中原子之间的计算注意力。可以可视化不同层和这些层中的不同注意力头部的注意力图。还可以发现不同分子属性的不同注意力模式。

MD22 数据集：MD22 包含四种主要生物分子和超分子的分子动力学（MD）轨迹，从含有 42 个原子的短肽到含有 370 个原子的双层纳米管。模拟轨迹以 400K 和 500K 的温度采样，分辨率为 1fs。使用 PBE+MBD 理论级别计算势能和力。

MD22 结果在不同分子尺度上显示出一致的改进。对于四肽，GotenNetBGotenNetB在 MACE 的能量预测上提高了 19.2%，在 QUINNet 的力预测上提高了 11.9%。在 DHA 上，我们在能量上比 ViSNet-LSRM 提高了 34.1%，在力上比 Equiformer 提高了 16.8%。对于像淀粉糖这样的大分子，GotenNetBGotenNetB在能量上比 ViSNet-LSRM 提高了 36.2%，在力上比 QUINNet 提高了 21.4%。对于复杂结构，改进更为明显：AT-AT 在能量预测上比 ViSNet-LSRM 好 27.8%，在力预测上比 QUINNet 好 23.6%。对于 AT-AT-CG-CG，我们在能量上比 ViSNet-LSRM 提高了 15.5%，在力上提高了 27.3%。在具有挑战性的富勒烯捕捉器上，GotenNetBGotenNetB在能量上比 Equiformer 好 22.4%，在力上比 ViSNet-LSRM 好 22.3%。最后，对于最大的系统，双层纳米管，我们在能量预测上比 ViSNet 提高了 35.8%，在力预测上比 Equiformer 提高了 31.3%。

值得注意的是，即使是我们的较小模型变体GotenNetSGotenNetS，与现有基线相比也取得了显著的改进。对于 Stachyose，GotenNetSGotenNetS在能量预测方面比 ViSNet-LSRM 提高了 28.8%，同时保持了比 QUINNet 更好的力预测，提高了 5.7%。在 DHA 上，它在能量方面比 ViSNet-LSRM 提高了 26.5%，在力方面与 Equiformer 的表现相当。对于像 AT-AT-CG-CG 这样的大型系统，GotenNetSGotenNetS在能量方面比 ViSNet-LSRM 提高了 15.1%，在力方面提高了 22.5%。最令人印象深刻的是，即使面对具有挑战性的双层纳米管，我们的较小模型在能量预测方面比 ViSNet 提高了 29.6%，同时保持了比 Equiformer 更好的力预测。

请考虑在项目有帮助的情况下引用以下作品：

代码语言：txt复制

@inproceedings{aykent2025gotennet,
  author = {Aykent, Sarp and Xia, Tian},
  booktitle = {The Thirteenth International Conference on Learning Representations},
  year = {2025},
  title={{GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks}},
  url = {},
  howpublished = {},
}

实战

sarpaykent/GotenNet：[ICLR 2025] GotenNet：重新思考高效的 3D 等变图神经网络 --- sarpaykent/GotenNet: [ICLR 2025] GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks

虚拟环境

代码语言：txt复制

git clone .git

cd GotenNet/

# Using conda
    conda create -n gotennet python=3.10
    conda activate gotennet

代码语言：txt复制

pip install -r requirements.txt -i .t
singhua.edu/simple

代码语言：txt复制

apt install tree

tree -L 3

代码语言：txt复制

.
|-- LICENSE
|-- README.md
|-- assets
|   `-- GotenNet_framework.png
|-- configs
|   |-- __init__.py
|   |-- callbacks
|   |   |-- default.yaml
|   |   `-- none.yaml
|   |-- datamodule
|   |   `-- qm9.yaml
|   |-- experiment
|   |   |-- qm9.yaml
|   |   `-- qm9_u0.yaml
|   |-- hydra
|   |   |-- default.yaml
|   |   `-- job_logging
|   |-- local
|   |-- logger
|   |   |-- comet.yaml
|   |   |-- csv.yaml
|   |   |-- default.yaml
|   |   |-- many_loggers.yaml
|   |   |-- mlflow.yaml
|   |   |-- neptune.yaml
|   |   |-- tensorboard.yaml
|   |   `-- wandb.yaml
|   |-- model
|   |   `-- gotennet.yaml
|   |-- paths
|   |   `-- default.yaml
|   |-- test.yaml
|   |-- train.yaml
|   `-- trainer
|       `-- default.yaml
|-- gotennet
|   |-- __init__.py
|   |-- datamodules
|   |   |-- __init__.py
|   |   |-- components
|   |   `-- datamodule.py
|   |-- models
|   |   |-- __init__.py
|   |   |-- components
|   |   |-- goten_model.py
|   |   |-- representation
|   |   `-- tasks
|   |-- scripts
|   |   |-- __init__.py
|   |   `-- train.py
|   |-- testing_pipeline.py
|   |-- training_pipeline.py
|   |-- utils
|   |   |-- __init__.py
|   |   |-- logging_utils.py
|   |   |-- pylogger.py
|   |   `-- utils.py
|   `-- vendor
|       `-- __init__.py
|-- pyproject.toml
`-- requirements.txt

22 directories, 39 files

以下是为该目录结构的每个目录和文件标注的注解：

代码语言：markdown复制

# 项目目录结构

├── LICENSE # 项目的许可证文件

├── README.md # 项目的说明文档

├── assets # 项目资源文件夹

│ └── GotenNet_framework.png # 项目框架的示意图

├── configs # 项目配置文件夹

│ ├── init.py # 配置模块的初始化文件

│ ├── callbacks # 回调配置文件夹

│ │ ├── default.yaml # 默认回调配置

│ │ └── none.yaml # 无回调配置

│ ├── datamodule # 数据模块配置文件夹

│ │ └── qm9.yaml # QM9数据集配置

│ ├── experiment # 实验配置文件夹

│ │ ├── qm9.yaml # QM9实验配置

│ │ └── qm9_u0.yaml # QM9 U0实验配置

│ ├── hydra # Hydra配置文件夹

│ │ ├── default.yaml # 默认Hydra配置

│ │ └── job_logging # 作业日志配置

│ ├── local # 本地配置文件夹

│ ├── logger # 日志记录器配置文件夹

│ │ ├── comet.yaml # Comet日志记录器配置

│ │ ├── csv.yaml # CSV日志记录器配置

│ │ ├── default.yaml # 默认日志记录器配置

│ │ ├── many_loggers.yaml # 多日志记录器配置

│ │ ├── mlflow.yaml # MLflow日志记录器配置

│ │ ├── neptune.yaml # Neptune日志记录器配置

│ │ ├── tensorboard.yaml # TensorBoard日志记录器配置

│ │ └── wandb.yaml # W&B日志记录器配置

│ ├── model # 模型配置文件夹

│ │ └── gotennet.yaml # GotenNet模型配置

│ ├── paths # 路径配置文件夹

│ │ └── default.yaml # 默认路径配置

│ ├── test.yaml # 测试配置

│ ├── train.yaml # 训练配置

│ └── trainer # 训练器配置文件夹

│ └── default.yaml # 默认训练器配置

├── gotennet # GotenNet项目核心模块

│ ├── init.py # GotenNet模块初始化文件

│ ├── datamodules # 数据模块文件夹

│ │ ├── init.py # 数据模块初始化文件

│ │ ├── components # 数据模块组件文件夹

│ │ └── datamodule.py # 数据模块定义文件

│ ├── models # 模型文件夹

│ │ ├── init.py # 模型模块初始化文件

│ │ ├── components # 模型组件文件夹

│ │ ├── goten_model.py # GotenNet模型定义文件

│ │ ├── representation # 表示学习相关文件夹

│ │ └── tasks # 任务相关文件夹

│ ├── scripts # 脚本文件夹

│ │ ├── init.py # 脚本模块初始化文件

│ │ └── train.py # 训练脚本

│ ├── testing_pipeline.py # 测试流程脚本

│ ├── training_pipeline.py # 训练流程脚本

│ ├── utils # 工具模块文件夹

│ │ ├── init.py # 工具模块初始化文件

│ │ ├── logging_utils.py # 日志工具文件

│ │ ├── pylogger.py # Python日志器文件

│ │ └── utils.py # 通用工具函数文件

│ └── vendor # 第三方组件文件夹

│ └── init.py # 第三方组件初始化文件

├── pyproject.toml # 项目配置文件 (Poetry)

└── requirements.txt # 项目依赖项文件

代码语言：markdown复制

**总计：22 个目录，39 个文件**

科技改变生活-雨落星辰 - 所有的伟大,都源于一个勇敢的开始

GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks