[论文解读] Symmetry-Aware Actor-Critic for 3D Molecular Design

Gregor N. C. Simm, Robert Pinsler+2 more|arXiv (Cornell University)|Nov 25, 2020

Machine Learning in Materials Science参考文献 45被引用 9

一句话总结

本文提出了一种用于3D分子设计的对称性感知演员-评论家强化学习框架，该框架基于球面谐函数构建旋转协变表示，以建模原子位置和动作。通过将空间对称性直接编码到策略和价值网络中，该方法生成了质量更高、更稳定的3D分子结构，并在未见化学计量比和传统图神经网络方法难以处理的复杂配位几何结构上表现出更强的泛化能力。

ABSTRACT

Automating molecular design using deep reinforcement learning (RL) has the potential to greatly accelerate the search for novel materials. Despite recent progress on leveraging graph representations to design molecules, such methods are fundamentally limited by the lack of three-dimensional (3D) information. In light of this, we propose a novel actor-critic architecture for 3D molecular design that can generate molecular structures unattainable with previous approaches. This is achieved by exploiting the symmetries of the design process through a rotationally covariant state-action representation based on a spherical harmonics series expansion. We demonstrate the benefits of our approach on several 3D molecular design tasks, where we find that building in such symmetries significantly improves generalization and the quality of generated molecules.

研究动机与目标

• 解决基于图的分子生成方法缺乏3D空间信息且无法建模金属元素或复杂配位几何结构的局限性。
• 开发一种显式编码3D分子设计中旋转与平移对称性的强化学习框架，以提升样本效率与泛化能力。
• 通过在笛卡尔坐标系中建模原子位置并结合对称性感知表示，实现对复杂分子结构（如超价或过渡金属化合物）的生成。
• 证明通过旋转协变性引入归纳偏置，可使生成分子在有效性、多样性与稳定性方面优于基线强化学习智能体。

提出的方法

• 该方法采用马尔可夫决策过程，其中状态由已放置原子的3D画布和剩余原子的集合定义，动作包括选择元素并在3D坐标处放置。
• 利用球面谐函数构建旋转协变的状态表示，以编码原子位置，确保智能体策略在空间旋转下保持一致变换。
• 演员-评论家网络采用自回归策略，其原子位置分布由球面谐函数系数参数化，实现等变的动作选择。
• 通过策略梯度方法训练价值函数与策略，奖励函数旨在鼓励化学有效性、结构稳定性及目标分子组成。
• 通过确保状态表示与动作分布在旋转下可预测地变换，整个架构保持旋转协变性，从而维持对称性不变性。
• 该模型在3D分子设计任务上进行端到端训练，包括单袋与随机袋设置，且在化学计量比未知的情况下进行。

实验结果

研究问题

• RQ1 在化学计量比未知或分布外的情况下，状态-动作表示中的旋转协变性是否能提升3D分子生成的泛化能力？
• RQ2 与标准强化学习智能体相比，对称性感知表示学习如何影响生成3D分子结构的有效性、多样性和稳定性？
• RQ3 具备对称性先验的强化学习智能体在生成图神经网络模型无法触及的复杂配位几何结构方面，其能力达到何种程度？
• RQ4 使用球面谐函数建模动作分布是否能提升3D分子设计中的样本效率与收敛性？
• RQ5 当训练期间未提供确切化学计量比时，该智能体是否能无需微调即可泛化到更大、未见过的分子组成？

主要发现

• 所提出的对称性感知智能体Covariant在单袋设置下，于C7H10O2和C7H8N2O2任务中分别实现了0.65和0.71的平均回报，优于基线智能体。
• 在化学计量比未知的随机袋任务中，Covariant仍保持强劲表现（平均回报分别为0.65和0.88），证明其在无确切组成先验知识下具备出色的泛化能力。
• 与Internal基线相比，Covariant生成的分子结构更稳定、质量更高，回报方差更低，且在分布外袋上的表现更优。
• 该智能体成功生成了复杂3D结构，包括具有非有机成键模式的超价碘化合物等，这些结构对标准图基模型不可达。
• 最优基线智能体在大袋任务中难以生成有效结构，常产生多分子簇，而Covariant在各类配置下均保持高有效性与稳定性。
• Covariant在泛化到未见分子组成方面表现显著更优，其在分布外测试袋上的平均回报更高、方差更小，优于Internal智能体。