[论文解读] Decoding Molecular Graph Embeddings with Reinforcement Learning

Steven Kearnes, Li Li+1 more|arXiv (Cornell University)|Apr 18, 2019

Machine Learning in Materials Science参考文献 23被引用 30

一句话总结

该论文提出RL-VAE，一种基于强化学习的图到图变分自编码器，能够从潜在嵌入中解码出化学上有效的分子图。通过采用基于马尔可夫决策过程（MDP）的解码器，并利用受潜在码条件化的价值函数，该模型在QM9测试集上实现了67%的重构准确率，表明基于强化学习的解码能够同时实现有效的分子生成与连续可优化的潜在空间。

ABSTRACT

We present RL-VAE, a graph-to-graph variational autoencoder that uses reinforcement learning to decode molecular graphs from latent embeddings. Methods have been described previously for graph-to-graph autoencoding, but these approaches require sophisticated decoders that increase the complexity of training and evaluation (such as requiring parallel encoders and decoders or non-trivial graph matching). Here, we repurpose a simple graph generator to enable efficient decoding and generation of molecular graphs.

研究动机与目标

• 解决在变分自编码器中从连续潜在嵌入生成化学上有效分子图的挑战。
• 弥合基于自编码器的潜在空间优化与能生成有效分子的生成模型之间的差距。
• 在确保解码过程中结构有效性的前提下，实现在分子设计中的连续可微分优化。
• 探索强化学习是否可作为图结构分子数据的有效且可扩展的解码器。

提出的方法

• 模型使用消息传递神经网络（MPNN）编码器将分子图映射为由均值μ和对数方差Σ参数化的潜在分布。
• 从学习到的分布中采样潜在向量，并通过使用双Q学习训练的强化学习智能体进行解码，以近似状态值函数。
• 解码器作为马尔可夫决策过程（MDP）运行，逐步构建分子，动作受到约束以确保化学有效性（例如，禁止移除键或在环中形成三重键）。
• 价值函数同时依赖于当前图状态和目标分子的潜在嵌入，从而通过策略学习实现重构。
• 经验回放缓冲区存储来自20步轨迹的经验转移，训练时使用128的批量大小进行经验回放，8的批量大小进行轨迹滚动。
• MDP通过禁止如键移除或错误的环形成等无效动作，确保所有生成的分子均有效。

实验结果

研究问题

• RQ1 基于强化学习的解码器是否能在保持化学有效性的前提下实现高重构准确率？
• RQ2 VAE学习到的潜在空间是否支持分子结构之间的平滑、连续插值？
• RQ3 与传统的自编码目标函数相比，使用时序差分学习预测Q函数在图生成中的表现如何？
• RQ4 将RL解码器基于目标分子的潜在嵌入进行条件化，在多大程度上提升了重构性能？

主要发现

• RL-VAE在QM9测试集上实现了67%的重构准确率，显著优于基线模型，如随机游走（0.00%）和贪婪解码（0.03%）。
• 当重构不完全时，输入与输出分子之间的Tanimoto相似度急剧下降，表明MDP编辑距离比SMILES等价性更能反映分子相似性。
• MDP解码器通过禁止如键移除或在环中形成三重键等无效动作，保证了化学上有效的分子生成。
• 潜在空间探索表明，潜在空间中余弦距离越大，分子相似性越低，提示嵌入空间具有局部平滑性。
• 采用折扣因子γ=0.99的模型性能显著优于γ=0，证实了在RL解码器中长期奖励设计的重要性。
• 该模型优于GVAE（51%），并接近JT-VAE（74%），表明其在更简单、更高效的解码机制下仍具备具有竞争力的重构性能。