[论文解读] Data-Driven Approach to Encoding and Decoding 3-D Crystal Structures

Jordan Hoffmann, Louis Maestrati+4 more|arXiv (Cornell University)|Sep 3, 2019

Machine Learning in Materials Science参考文献 46被引用 24

一句话总结

该论文提出了一种数据驱动的深度学习框架，通过从原子位置导出的连续3D密度表示来编码和解码3D晶体结构。通过在120,000个旋转后的晶体原胞上使用3D卷积神经网络训练自编码器，该模型学习到一个平滑、连续的潜在空间，从而实现对分子的精确重建、插值以及具有受控几何和组成特性的条件生成。

ABSTRACT

Generative models have achieved impressive results in many domains including image and text generation. In the natural sciences, generative models have led to rapid progress in automated drug discovery. Many of the current methods focus on either 1-D or 2-D representations of typically small, drug-like molecules. However, many molecules require 3-D descriptors and exceed the chemical complexity of commonly used dataset. We present a method to encode and decode the position of atoms in 3-D molecules from a dataset of nearly 50,000 stable crystal unit cells that vary from containing 1 to over 100 atoms. We construct a smooth and continuous 3-D density representation of each crystal based on the positions of different atoms. Two different neural networks were trained on a dataset of over 120,000 three-dimensional samples of single and repeating crystal structures, made by rotating the single unit cells. The first, an Encoder-Decoder pair, constructs a compressed latent space representation of each molecule and then decodes this description into an accurate reconstruction of the input. The second network segments the resulting output into atoms and assigns each atom an atomic number. By generating compressed, continuous latent spaces representations of molecules we are able to decode random samples, interpolate between two molecules, and alter known molecules.

研究动机与目标

• 开发一种连续、可微分的3D晶体结构表示方法，以捕捉几何信息和原子组成信息。
• 实现对包含多达100个以上原子的复杂3D晶体原胞的生成建模，超越一维或二维分子表示的范围。
• 学习一个解耦的、连续的潜在空间，以支持插值、条件生成和性质控制。
• 证明所学习的潜在表示可用于生成具有受控原子类型和空间排列的物理上合理的3D结构。

提出的方法

• 使用原子位置在晶体原胞中构建3D密度场表示，采用平滑连续的电子密度代理函数。
• 在120,000个增强样本（旋转后的原胞）上训练3D卷积自编码器，以学习压缩的、连续的潜在空间。
• 使用独立的分割头将潜在表示解码为原子类型，通过将原子序数分配给预测的密度峰值实现。
• 通过随机旋转进行数据增强，以提高泛化能力并增强模型对方向的不变性。
• 通过调节瓶颈层中的最大密度值，将生成过程与最大原子种类关联起来。
• 采用VAE架构以支持潜在空间中的随机采样和插值，用于分子生成。

实验结果

研究问题

• RQ1 所学习的3D密度表示能否有效编码包含多达100个原子的晶体原胞的几何与组成复杂性？
• RQ2 所学习的潜在空间是否支持对新型晶体结构进行有意义的插值与受控生成？
• RQ3 该模型能否生成保持真实晶体中常见原子间距离分布的物理上合理的3D结构？
• RQ4 该模型在多大程度上可针对特定化学性质进行条件控制，例如避免重原子的存在？
• RQ5 在几何保真度和化学合理性方面，生成的结构与真实晶体结构相比表现如何？

主要发现

• 该模型能够以高保真度重建3D晶体结构，准确保留了不同原胞几何形状下原子的空间排列。
• 在所学习的潜在空间中进行插值可产生分子结构之间的平滑过渡，表明该表示是连续且有意义的。
• 从潜在空间中随机采样可生成合理的3D密度场，其原子间距离分布与真实晶体一致，经直方图分析验证。
• 通过调节瓶颈层实现的条件生成可控制最大原子序数的出现，展示了有针对性的设计能力。
• 该模型支持尺度不变的生成：改变潜在变量仅影响尺寸而不扭曲分子几何形状。
• 尽管在生成可进行DFT优化的结构方面仍存在局限，但该框架为未来基于等变架构和可微原子定位的3D分子生成研究奠定了坚实基础。