[论文解读] Symphony: Symmetry-Equivariant Point-Centered Spherical Harmonics for 3D Molecule Generation

Ameya Daigavane, Song Eun Kim+2 more|arXiv (Cornell University)|Nov 27, 2023

Machine Learning in Materials Science被引用 4

一句话总结

Symphony 提出了一种 $E(3)$-等变的自回归生成模型，用于3D分子几何结构建模，通过以点为中心的球谐投影同时捕捉原子的径向与角度分布。该模型在QM9数据集上优于以往的自回归模型，并在性能上接近扩散模型，同时保持更快的推理速度和对未见片段的强大泛化能力。

ABSTRACT

We present Symphony, an $E(3)$-equivariant autoregressive generative model for 3D molecular geometries that iteratively builds a molecule from molecular fragments. Existing autoregressive models such as G-SchNet and G-SphereNet for molecules utilize rotationally invariant features to respect the 3D symmetries of molecules. In contrast, Symphony uses message-passing with higher-degree $E(3)$-equivariant features. This allows a novel representation of probability distributions via spherical harmonic signals to efficiently model the 3D geometry of molecules. We show that Symphony is able to accurately generate small molecules from the QM9 dataset, outperforming existing autoregressive models and approaching the performance of diffusion models.

研究动机与目标

• 开发一种尊重分子 $E(3)$ 对称性的生成模型，包括原子体系中固有的全局与局部对称性。
• 通过球谐函数引入更高阶的 $E(3)$-等变特征，提升自回归生成中3D分子几何结构的表征能力。
• 在单一聚焦原子框架下，实现对新原子径向与角度分布的精确高效建模。
• 评估模型在未见分子片段上的泛化性能，超越对训练数据的记忆化。
• 在保持高有效性和几何精度的前提下，实现比扩散模型更快的推理速度。

提出的方法

• Symphony 采用以点为中心的自回归生成流程，每个新原子均相对于单一聚焦原子预测，该原子作为局部原点。
• 利用从球谐函数 $Y_{l,m}$ 衍生的高阶 $E(3)$-等变特征，表征候选原子的角分布与径向分布。
• 通过球谐信号参数化新原子位置的条件概率分布，实现对径向与角分量的联合预测。
• 采用 $E(3)$-等变特征的消息传递机制，确保在旋转与平移下的不变性，从而保持分子对称性。
• 通过张量积组合特征，生成 $SO(3)$ 的更高阶不可约表示，实现更丰富的几何建模。
• 提出一种基于双谱的新角度精度度量，用于评估生成分子中局部环境匹配程度。

实验结果

研究问题

• RQ1 球谐投影能否实现更精确且高效的自回归生成中3D分子几何结构建模？
• RQ2 相较于旋转不变模型，使用高阶特征的 $E(3)$-等变性在多大程度上提升了生成质量？
• RQ3 Symphony 在完成未见分子片段方面能实现多大程度的泛化，表明其超越记忆化的真正泛化能力？
• RQ4 尽管张量积带来更高的计算成本，Symphony 的推理速度与扩散模型及先前自回归模型相比如何？
• RQ5 与基于查表法相比，使用球谐函数能否减少键长分布中的离散化误差？

主要发现

• Symphony 在 QM9 所有指标上均优于 G-SchNet 与 G-SphereNet，原子类型计数的 Jensen-Shannon 散度（JSD）为 0.0003，局部环境计数的 JSD 为 0.0039。
• 对于最常见的10种局部环境，模型的双谱 MMD 达到 0.0039，显著低于 G-SchNet（0.0150）与 G-SphereNet（0.1016）。
• 在未见测试片段上，Symphony 的有效完成率（VCR）达到 98.66%，优于 G-SchNet（98.15%），展现出强大的泛化能力。
• 在单张 A5000 GPU 上，每分子推理时间为 0.293 秒，约为 EDM（0.930 秒/分子）的 3 倍快，尽管张量积带来了更高的计算开销。
• 模型仅需约 80 个周期即可训练完成，远少于 G-SchNet（330 个周期）与 EDM（1100 个周期），表明其数据效率更高。
• 尽管性能出色，Symphony 的键长分布仍因径向离散化较粗而略显不准确，影响了键级分配的精度。