晶体生成10倍加速！基于流模型的晶体结构生成模型CrystalFlow在昇思社区开源

# 01 背景

理论晶体结构预测是通过计算的手段寻找物质在给定的外界条件下最稳定结构的重要手段。传统结构预测方法依赖在势能面上广泛的随机采样来寻找最稳定结构，然而，这种方法需要对大量随机生成的结构进行局域优化，而局域优化通常需要消耗巨大的第一性原理计算成本，尤其在模拟多元素复杂体系时，这种计算开销会显著增加，从而带来巨大的挑战。

近年来，基于深度学习生成模型的晶体结构生成方法因其能够在势能面上更高效地采样合理结构而逐渐受到关注。这种方法通过从已有的稳定或局域稳定结构数据中学习，进而生成合理的晶体结构，与随机采样相比，不仅能够减少局域优化的计算成本，还能通过较少的采样找到体系的最稳定结构。在这项工作中，基于流模型发展了晶体结构生成模型CrystalFlow，该模型在MP20等基准数据集上达到优秀的水平，目前已在昇思社区开源。

项目地址：

https://gitee.com/mindspore/mindscience/tree/master/MindChemistry

# 02

模型介绍

CrystalFlow采用先进的流模型生成方案，有效地在未知的目标概率分布和已知的简单概率分布建立概率密度变换关系，以学习晶体结构数据库中合理的稳定结构特征，并生成新的合理构型。通过使用图等变消息传递网络，CrystalFlow有效处理了晶体结构内在的基本对称性，包括置换、旋转和周期平移不变性，实现对晶体结构的高效学习和建模，以及灵活的条件式生成策略。

图1 CrystalFlow模型架构图

CrystalFlow模型架构如图1所示。对于晶体结构而言，结构的组成包括周期晶格、原子坐标和元素种类等部分。对各个元素分别构造含时矢量场u_t，推动概率密度从简单分布向目标分布演化，模型则使用由参数θ参数化的神经网络v_(t,θ)对该矢量场进行建模。CrystalFlow通过选取与旋转无关的晶格表示，以及原子的分数坐标表示，将旋转自由度从晶体表示中解耦，简化了流模型对结构不变性的处理。进一步，CrystalFlow将以上各结构组成部分的变量分离，分别使用各自的概率流进行计算，再将各部分联合训练，完成对晶体结构整体的学习。

对于晶体结构预测问题，我们希望确定给定化学配比时的稳定结构，因此此时结构的元素组成保持不变，不参与概率流变换；而对于新材料设计问题，我们可能不希望预先固定结构的元素组成，此时需要将其参与概率流变换。CrystalFlow分别实现了两种情形下的模型训练与结构生成。另一方面，为了进行功能材料的逆向结构设计，CrystalFlow还实现了以外界压强或结构形成能等信息为目标的条件式结构生成，实现根据期望的特定物理化学性质对模型输出进行调控。

# 03 模型表现

目前，CrystalFlow已经在昇思社区开源。在晶胞内包含不超过20原子的MP20数据集和更复杂与挑战的MPTS52数据集上的基准测试如表1所示。通过在测试集上计算模型的匹配率（MR）和均方根误差（RMSE）来衡量模型能否找到目标稳定结构的能力，以及找到的结构距离真实目标结构的差异，以对模型性能的评估。在重复次数分别取1，20和100时，与其他方法相比，CrystalFlow模型的预测效果达到了最优水平，体现了其更强的预测能力。CrystalFlow在推理时间上也具有显著优势。CrystalFlow仅采用100积分步数（S）即可达到与基于Diffusion的模型相同的效果，推理时间降低了一个数量级，较现有晶体生成SOTA模型DiffCSP推理性能10倍加速。具体推理时间如表2所示。

表1 不同方法在MP20和MPTS52数据集上的结构预测表现

表2 不同方法在相同硬件上的推理时间对比。在MP20测试集上进行生成以进行评估。S表示不同方法使用的积分步数。

# 04 结语

CrystalFlow使用流模型构建了一种晶体结构生成方法，通过对晶格、原子位置和元素种类分别设计先验分布和概率流演化概率路径，实现在满足晶体结构置换、旋转和平移等不变性的要求下，高效生成晶体，并在多个生成任务的表现上达到最优水平。CrystalFlow作为一种极具前景且用途广泛的结构采样工具，有望深度整合到各类晶体结构预测方法与材料设计工作流程中。相关使用案例可在昇思社区查看，欢迎开发者们使用。

项目地址：

https://gitee.com/mindspore/mindscience/tree/master/MindChemistry

参考文献

[1] X. Luo et al., CrystalFlow: A Flow-Based Generative Model for Crystalline Materials, arXiv 2412.11693