CrystalDiT: A Diffusion Transformer for Crystal Generation
CrystalDiT:一种用于晶体生成的扩散 Transformer
Authors: Xiaohan Yi, Guikun Xu, Zhong Zhang, Liu Liu, Yatao Bian, Xi Xiao, Peilin Zhao
摘要
我们提出了 CrystalDiT,这是一种用于晶体结构生成的扩散 Transformer。它通过挑战架构复杂化的趋势,实现了最先进的性能。CrystalDiT 没有采用复杂的多流设计,而是使用统一的 Transformer,并施加了一种强大的归纳偏置:将晶格和原子性质视为一个单一的、相互依赖的系统。结合基于元素周期表的原子表示以及均衡训练策略,我们的方法在 MP-20 上达到了 8.78% 的 SUN(Stable, Unique, Novel;稳定、唯一、新颖)率,显著优于包括 FlowMM(4.21%)和 MatterGen(3.66%)在内的近期方法。值得注意的是,CrystalDiT 生成了 63.28% 的唯一且新颖结构,同时保持了可比的稳定性比例,这表明对于材料发现而言,架构简单性可能比复杂性更有效。我们的结果表明,在数据受限的科学领域中,经过精心设计的简单架构优于容易过拟合的复杂替代方案。
代码 — https://github.com/hanyi2021/CrystalDiT.git
1 引言
材料发现限制了能源与可持续发展领域的技术进步(Jain et al. 2013; Merchant et al. 2023)。传统筛选方法将探索限制在已知结构范围内(Curtarolo et al. 2012)。
生成模型通过直接提出新颖晶体结构,提供了一种具有变革性的替代方案。近期方法已经通过各种复杂途径取得了有前景的结果:ADiT 采用具有独立 VAE 和 DiT 组件的两阶段潜空间扩散(Joshi et al. 2025),FlowMM 使用黎曼流匹配来处理晶体对称性(Miller et al. 2024),MatterGen 则引入了带有等变评分网络和适配器模块的联合扩散过程(Zeni et al. 2025)。尽管这些方法展示了生成建模的潜力,但它们的架构复杂性也引发了关于其必要性与有效性的问题。
尽管 Transformer 架构已在多个不同领域取得成功(Feng et al. 2024; Yi et al. 2023; Gong et al. 2024; Chen et al. 2024; Wu et al. 2020; Fang et al. 2024; Chen et al. 2025; Shibo et al. 2025; Wang et al. 2025; Xu et al. 2021),但用于数据受限科学应用的设计原则仍然缺乏充分探索。
我们识别出当前方法中的两个关键问题。首先,评估指标强调与训练数据的结构相似性,从而无意中惩罚了发现所必需的新颖性(Xie et al. 2021)。针对传统有效性指标进行优化的模型擅长复现已知模式,但在探索方面表现不足。其次,在材料科学中,走向架构复杂化的趋势可能适得其反,因为该领域的数据集较小,并且偏向于已知稳定结构。不同于拥有大规模数据集的视觉或语言任务,复杂模型更可能产生过拟合,而不是实现泛化。
我们提出 CrystalDiT,用于探索简化架构是否能够在晶体生成中优于复杂替代方案。为了严格检验这一假设,我们同时开发了一种简单的统一架构和一种复杂的双流变体,以便进行直接比较。我们的主要贡献如下:
简化架构: 一个统一的扩散 Transformer 通过联合注意力机制处理所有晶体信息,这与我们采用独立处理路径并配备交叉注意力机制的双流方案形成对比。
化学表示: 一种使用**元素周期表位置(周期、族)**的二维原子编码,在不增加架构复杂性的情况下自然捕获化学关系。
均衡评估: 一种复合分数,明确优化了生成质量与发现潜力之间的权衡关系,解决了评估目标与实际需求之间的不匹配问题。
我们的发现表明,在数据受限的科学领域中,具有领域特定表示的、经过精心设计的简单架构优于复杂替代方案,这挑战了关于架构复杂性必要性的假设。
2 相关工作
晶体生成已经成为材料信息学领域的核心关注点。为了解决传统计算材料发现方法的局限性,例如其有效性不足和高计算成本(Pickard and Needs 2011; Curtarolo et al. 2012; Yamashita et al. 2018; Wang, Botti, and Marques 2021),近期方法越来越多地利用基于深度学习的生成模型。
CDVAE(Xie et al. 2021)是最早预测三个关键组成部分(原子的数量和类型以及晶格结构)的方法之一,作为材料结构的初始近似,随后使用基于扩散的方法进行细化(Song and Ermon 2019; Ho, Jain, and Abbeel 2020)。DiffCSP(Jiao et al. 2023)是第一个采用联合等变扩散范式(即联合扩散晶格和分数坐标)进行晶体结构预测(CSP)的方法。通过在原子类型上引入额外的离散扩散(Austin et al. 2021),该方法可以进一步扩展到处理从头晶体生成,而 MatterGen(Zeni et al. 2025)进一步改进了这种联合扩散范式。
扩散 Transformer(DiT)(Peebles and Xie 2023)通过其强大的注意力机制和自适应条件化,在学习稳定晶体结构方面表现出显著能力。然而,这种表达能力也在数据受限的科学领域中引入了显著的过拟合风险,在这些领域中,模型可能会记忆训练模式,而不是学习可泛化的材料发现原则。ADiT(Joshi et al. 2025)引入了一种两阶段方法:首先使用自编码器生成晶体结构的潜在表示,然后在潜空间中应用 DiT 架构进行结构生成。
流匹配技术(Lipman et al. 2022, 2024)近来作为扩散模型的一种高度有效的替代方案出现,这主要归因于其增强的推理效率以及在定义先验分布方面的灵活性。在这一背景下,FlowMM(Miller et al. 2024)在黎曼流形中利用联合黎曼流匹配(Chen and Lipman 2023),从而更好地处理晶体周期性。随后,FlowLLM(Sriram et al. 2024)将这一框架扩展到大型语言模型(LLMs),利用它们作为采样亚稳态材料化学式的先验分布,并通过黎曼流匹配生成其对应的结构构型。
在另一条研究线上,材料对称性在晶体结构生成中受到了显著关注。近期,DiffCSP++(Jiao et al. 2024)、SymmCD(Levy et al. 2025)和 WyFormer(Kazeev et al. 2025)通过将 Wyckoff 位置(Wyckoff 1922)纳入晶体生成任务取得了显著进展,使得能够生成具有确定对称性的晶体结构。尽管这些技术取得了进步,现有方法仍表现出不断增加的架构复杂性,并采用专门化的多组件设计。这促使我们研究统一的、简化的架构是否能够通过更有效地学习晶格—原子关系来实现更优性能。
3 方法
我们提出 CrystalDiT,这是一种简单而有效的用于晶体结构生成的扩散 Transformer 架构。我们的方法由四个关键组成部分构成:(1)一种新颖的二维原子表示,通过元素周期表定位捕获化学关系;(2)一种精简的扩散 Transformer 架构,通过统一注意力机制处理晶体结构;(3)一种替代传统验证的均衡模型选择策略;以及(4)一种用于推理的概率式原子解码过程。本节详细介绍我们方法中的每个组成部分。
3.1 晶体表示
有效的晶体结构表示对于训练能够生成稳定且新颖材料的扩散模型至关重要。传统方法要么使用基于图的表示,其复杂度随原子数量呈二次增长;要么依赖于忽略化学关系的原子序数编码。尽管近期工作如 UniMat(Yang et al. 2023)提出了基于元素周期表的表示,但其 4D 张量方法需要为元素周期表中每一种可能元素预先分配空间,导致稀疏表示,其中大多数位置仍未被占用。此外,它们的方法还需要定义每种元素类型的最大原子数量,并在化学维度上使用专门的张量操作。相比之下,我们的方法**提供了一种更紧凑的表示,只编码 20 个原子,并且能够通过标准 Transformer 操作进行统一处理**。
我们引入了一种简化但化学意义更强的晶体表示,该表示建立在元素周期表结构之上,同时纳入来自材料科学的领域特定洞见。
二维原子表示
我们不使用原子序数来表示原子,而是使用其在元素周期表中的位置来编码每个原子:周期(行)和族(列)。这种编码的动机来自一个基本原则:同一周期中的元素具有相似的电子壳层构型,而同一族中的元素表现出相似的化学性质。
具体而言,对于原子序数为 的原子,我们将其映射到一个元组 ,其中 表示周期, 表示族。这里, 且 对应一个特殊的“空原子”,用于表示少于 20 个原子的晶体中的空位置。对于有效元素, 且 ,镧系元素和锕系元素 对应元素周期表最下面2行 使用分数族号来保留它们的独特位置。然后,我们将这些值归一化到 :
这种表示具有若干优势:(1)它通过元素周期表中的空间邻近性自然捕获化学相似性;(2)它提供了一个连续嵌入空间,有利于扩散建模;(3)它在保留化学意义的同时降低了维度;(4)它通过位于原点 的空原子表示无缝处理可变大小的晶体结构。
归一化晶格参数化
对于晶格向量,我们采用一种归一化策略,以应对真实材料中晶格参数取值范围很宽的问题。给定晶格矩阵 ,我们使用 MP-20 数据集中观察到的最大长度尺度进行归一化:
其中,L_{\max} = 46.7425 \ \mathrm{\AA},这是我们对 MP-20 数据集(Xie et al. 2021)分析得到的结果。该数据集包含 45,231 个亚稳态晶体结构,最多含 20 个原子,覆盖 89 种元素类型。
完整结构表示
我们的完整晶体表示结合了归一化晶格向量和原子特征。对于具有 个原子的晶体(对于 MP-20,):
(归一化晶格)
(原子特征)
其中, 的每一行包含 ,分别表示归一化周期、归一化族以及分数坐标。对于少于 20 个原子的晶体,我们使用 的“空”原子进行填充。
3.2 架构设计
我们的 CrystalDiT 架构围绕这样一个原则进行设计:在晶体生成任务中,简单性能够带来更好的泛化。不同于复杂的多流架构,我们采用一种统一方法,通过单一、精简的 Transformer 路径处理所有晶体信息。
统一 DiT 架构
该模型由三个主要组件组成:(1)晶体结构嵌入Embedding,将晶格向量和原子特征映射到共享隐藏空间;(2)由18个DiT模块组成的序列,通过统一自注意力机制处理组合表征;以及(3)专用输出头可分别对原子组分和晶格组分生成噪声预测值。
核心启示在于,通过在单一注意力路径中共同处理原子和晶格特征,模型可以自然学习原子位置与晶格参数之间复杂的相互依赖关系,而不需要显式交叉注意力机制。这种统一方法施加了一种强归纳偏置:将晶格和原子性质视为一个单一的、相互依赖的系统,这与晶体性质源于原子组成和晶格几何之间相互作用的物理现实相一致。
核心启示在于:通过在单一注意力机制中同时处理原子特征与晶格特征,该模型能够自然地学习原子位置与晶格参数之间的复杂相互依赖关系,而无需依赖显式的交叉注意力机制。这种统一方法强化了强烈的归纳偏置——将晶格特性与原子特性视为一个相互依存的整体系统,这与晶体性质源于原子组成与晶格几何结构相互作用的物理现实完全吻合。
晶体结构被嵌入到维度为 的隐藏空间中。晶格向量和原子特征通过独立的线性嵌入层进行处理,并结合位置编码和类型编码。嵌入后的特征被拼接形成组合表示 (20 个原子 + 3 个晶格向量)。
该组合表示通过 个相同的 DiT 块进行处理,每个块均采用时间条件自适应层归一化(AdaLN)技术,根据扩散时间步长对特征进行调节。最后,专用输出头分别生成原子特征(5维)和晶格向量(每向量3维)的噪声预测结果。
架构比较
与 Figure 1 所示的统一方法相比,我们还实现了一种复杂的双流变体以进行直接比较。该架构采用级联处理方式:先由12个纯原子级 DiT 模块独立处理原子特征,随后是2个纯晶格模块处理晶格向量,最后是2个配备双向交叉注意力机制的联合模块用于融合不同数据流之间的信息。与我们从一开始就统一处理所有特征的简单方法不同,这种级联设计通过专用模块将原子级处理路径与晶格级处理路径分离。然而,这种架构复杂性在数据有限的晶体生成领域会导致过拟合现象——尽管个体有效性指标有所提升,但独特的新型晶体生成率却显著降低。
3.3 训练目标与模型选择
我们的训练方法结合了标准扩散目标与一种新颖的模型选择策略,该策略解决了晶体生成中的独特挑战。
扩散损失函数
我们采用具有 个时间步和线性噪声调度的高斯扩散过程。我们的模型学习预测在每个时间步 添加的噪声 ,使用如下加权损失函数:
其中:
我们设置 以平衡不同尺度,并使用 feature-specific 权重 来强调周期和族的预测。
均衡模型选择策略
传统检查点选择方法在晶体生成中面临根本性局限。验证损失虽然是机器学习中的标准做法,但无法捕获生成晶体的真实稳定性和发现潜力,因为它只度量重构保真度,而不是物理合理性。近期方法如 ADiT(Joshi et al. 2025)采用的替代方法是在每个检查点生成 1000 个结构,并选择具有最高有效性比例的模型。然而,这种策略会无意中促进对训练数据中已知稳定模式的过拟合,导致高有效性分数但低新颖性比例——这与材料发现的需求恰好相反。
我们的 Balance Score 提供了一种有原则的替代方案,在检查点选择过程中显式优化生成质量与发现潜力之间的权衡:
其中, 是控制生成质量与发现潜力之间权衡的超参数。更高的 值强调质量,以获得更可靠的结构;更低的 值则偏向于发现唯一且新颖的晶体。
UN Rate 度量生成结构中既唯一又新颖的比例:
Quality Composite 是四个归一化质量指标的几何平均:
每个组成评分根据经验观察到的范围归一化到 :
我们在训练过程中实现了一种多阶段检查点选择策略。我们从三个不同训练阶段中识别具有最高 Balance Score 的模型:早期(0–30%)、中期(31–60%)和后期(61–100%)阶段。这种方法认识到,质量与发现潜力之间的最优权衡可能在不同训练阶段出现,其中早期模型可能偏向探索,而后期模型更强调质量。详细训练协议见 Appendix C。
3.4 推理过程
在推理过程中,我们的模型通过标准 DDPM 采样过程生成晶体结构,随后使用概率式原子解码过程将连续预测转换为离散原子类型。
概率式原子解码
我们的模型预测原子周期和族的连续值,这些值必须映射到离散原子序数。我们将候选原子类型定义为所有有效元素加上位于 位置的空原子()。遵循 DDPM(Ho, Jain, and Abbeel 2020)中的离散解码器方法,我们采用基于高斯分布的概率映射。
对于预测的连续值 以及位于位置 的元素 ,映射概率为:
概率式原子解码
我们的模型预测原子周期和族的连续值,这些值必须映射到离散原子序数。我们为每个候选元素分配一个责任区域,并使用高斯积分计算概率。
对于预测的连续值 以及位于归一化位置 的候选元素 ,映射概率为:
其中, 控制映射的尖锐程度。积分边界定义每个元素的责任区域:
族也采用类似边界,其中 和 是离散化间隔。最终原子序数选择为:
对于位于 的空原子,其区域扩展到 ,从而能够自然处理空位置。
算法 1:带有概率式原子解码的 CrystalDiT 生成
输入:时间步 ,批大小 ,模型
输出:晶体结构
- 采样初始噪声:
- 对 :
- 结束循环
- 提取晶格:
- 提取原子:
- 对每个样本 :
- 初始化:,
- 对每个原子 :
- 使用式(16)–(20)计算
- 如果 :
- 结束条件判断;如果 ,则跳过(空原子)
- 结束原子循环
- 根据 和 构造
- 结束样本循环
4 实验与结果
我们通过全面实验,将我们的CrystalDiT方法与最先进的晶体生成方法进行对比评估。评估过程中采用成熟的跨方法比较指标;同时特别使用平衡得分(Balance Score)在模型训练阶段筛选验证点,以优化生成质量与发现潜力之间的平衡关系。
4.1 实验设置
统一评估协议
为确保公平比较,我们采用统一的测试方案对所有基线方法进行重新评估,该方案使用完全相同的评估指标、DFT计算参数及统计抽样流程。此方法可消除因原始文献中不同评估实现方式可能导致的差异。
数据集与预处理
我们使用 MP-20 数据集(Xie et al. 2021),该数据集包含来自 Materials Project 的 45,231 个亚稳态晶体结构,最多含 20 个原子,覆盖 89 种元素类型。按照标准实践,我们使用既定的训练/测试划分,并使用我们的二维原子表示和晶格归一化对结构进行预处理。
模型配置
我们评估三种模型变体:
CrystalDiT(Simple): 我们的主模型,采用统一注意力处理。架构:, 层,8 个注意力头。模型大小:330MB。
CrystalDiT(Complex): 用于比较的双流架构,具有独立的原子和晶格处理流,并带有交叉注意力机制。该架构使用级联处理:12 个仅处理原子的 DiT 块 → 2 个仅处理晶格的 DiT 块 → 2 个带有双向交叉注意力的联合 DiT 块。模型大小:370MB(参数数量控制为与简单版本相近)。
两个模型均训练 50,000 个 epoch,批大小为 256,学习率为 ,使用 8 块 V100 GPU 训练 4 天。
基线方法
我们与五种最先进方法进行比较:DiffCSP(Jiao et al. 2023)、FlowMM(Miller et al. 2024)、DiffCSP++(Jiao et al. 2024)和 MatterGen(Zeni et al. 2025),使用预训练检查点或官方实现;ADiT(Joshi et al. 2025)使用其官方仓库中预生成的 10,000 个结构。
评估协议
遵循 FlowMM(Miller et al. 2024)建立的评估框架,我们从每种方法生成 10,000 个结构,并计算全面指标:
有效性指标:
- 结构有效性:所有成对原子距离 \geq 0.5 \ \mathrm{\AA} 且晶体体积 \geq 0.1 \ \mathrm{\AA}^3 的晶体百分比。
- 组成有效性:通过 SMACT(Davies et al. 2019)满足电荷中性和电负性平衡的百分比,使用氧化态枚举和 Pauling 电负性规则。
分布指标:
- 密度距离():生成集合与测试集密度分布之间的 Wasserstein 距离。
- 元素距离():生成集合与测试集元素出现频率分布之间的 Wasserstein 距离。
发现指标:
- 唯一性:由 PyMatGen 的 StructureMatcher(Ong et al. 2013)判定为彼此不同的结构。
- 新颖性:通过 StructureMatcher 与 MP-20 训练集中任何晶体都不匹配的结构。
- UN Rate:同时唯一且新颖的结构所占比例。
稳定性评估: 我们采用 FlowMM(Miller et al. 2024)建立的协议进行 DFT 评估。由于计算资源限制,我们随机抽取 500 个 UN 结构用于稳定性评估(相比之下,FlowMM 评估所有结构)。为量化采样不确定性,我们对 CrystalDiT、FlowMM 和 MatterGen 独立重复该采样三次,并报告样本间的均值 ± 标准差。每个采样结构经历:
- 使用 CHGNet(Deng et al. 2023)机器学习势进行预弛豫;
- 使用 VASP 和 MPRelaxSet 参数进行 DFT 弛豫(Jain et al. 2013);
- 相对于 Matbench Discovery 凸包计算凸包上方能量(Riebesell et al. 2023);
- 分类:稳定(),亚稳态()。
最终发现指标:
- SUN Rate:UN Rate × UN 结构中的 Stable Rate。
- MSUN Rate:UN Rate × UN 结构中的 Metastable Rate。
详细评估参数和实现细节见 Appendix A。
4.2 主要结果
Table 1 展示了我们与最先进方法的全面比较。结果揭示了关于不同方法在晶体生成中有效性的若干关键洞见。关于基线结果的说明:我们使用统一协议重新评估得到的结果与文献中一些原始报告值不同。关于这些差异的详细分析见 Appendix A。
我们的简单 CrystalDiT 达到了最高的 SUN rate(8.78%)和 MSUN rate(25.94%),显著优于包括 FlowMM(4.21%)和 MatterGen(3.66%)在内的近期方法。值得注意的是,CrystalDiT 生成了 63.28% 的唯一且新颖结构,同时保持了可比的稳定性比例,这表明对于材料发现而言,架构简单性可能比复杂性更有效。不同方法表现出不同的权衡:ADiT 达到了较高的有效性分数,但由于过拟合而新颖性较差(37.08% UN rate);而 FlowMM 和 MatterGen 能生成新颖结构,但稳定性比例较低。
4.3 架构比较分析
我们的简单架构与复杂架构之间的比较揭示了关于科学应用中生成建模的基本洞见。尽管复杂双流架构具有复杂的交叉注意力机制,但其在发现指标上表现低于简单版本(6.36% vs 8.78% SUN rate)。双流架构的实现与分析见 Appendix B。
这一结果挑战了机器学习中关于架构复杂性的流行假设。复杂模型在单独质量指标上取得了更好的表现,但 UN rate 显著更低,这表明架构复杂性会促进晶体生成任务中的过拟合。简单的统一注意力机制似乎更有效地学习了可泛化模式,而不是记忆训练数据分布。
类似地,ADiT 的两阶段方法(自编码器 + 潜空间 DiT)达到了优异的有效性分数,但新颖性较差(37.08% UN rate),这表明如果不仔细考虑发现目标,架构复杂化在材料科学应用中可能会适得其反。
4.4 训练动态与模型选择
Table 2 展示了我们的简单 CrystalDiT 模型在不同训练 epoch 下的演化,说明了在训练过程中使用 Balance Score 进行检查点选择的重要性。注意,我们的 Balance Score 仅用于选择 CrystalDiT 模型训练中的最佳检查点,而不用于比较不同方法。所有跨方法比较均使用标准 SUN/MSUN 发现指标。
简单模型显示出清晰模式:随着训练进行,有效性指标提高,但 UN rate 从 80.72% 稳定下降到 57.37%。关注有效性的传统验证方法会选择最终检查点,但我们的 Balance Score 正确识别出具有更好发现潜力的较早检查点。这说明均衡评估对于材料发现应用具有关键重要性。
4.5 组成部分分析
我们通过在简单 CrystalDiT 架构上与传统的一维原子序数编码进行比较,评估了二维原子表示的贡献。一维原子表示的详细实现见 Appendix B。
如 Table 1 所示,我们基于二维元素周期表的表示(CrystalDiT Simple:8.78% SUN rate)显著优于一维编码(CrystalDiT 1D atomic:6.28% SUN rate)。虽然一维方法达到了更高的 UN rate(78.47% vs 63.28%),但生成结构表现出较低的稳定性比例(8.00% vs 13.87% stable rate),从而导致最终发现性能较差。这表明嵌入在元素周期表结构中的化学知识提升了生成晶体在材料发现中的质量。
我们还对架构深度进行了消融研究,发现 18 层在模型容量与泛化之间提供了最优平衡。详细结果见 Appendix C。
4.6 能量分布分析
Figure 2 展示了 CrystalDiT 与代表性基线方法 FlowMM 之间的能量分布比较。我们的分析显示,相比 FlowMM,CrystalDiT 在生成热力学有利结构方面表现出更强能力。CrystalDiT 在稳定区域()显示出明显峰值,表明我们的简化架构能够有效学习生成能量上有利的晶体结构。能量分布提供了关键验证:CrystalDiT 不仅生成了更多唯一且新颖结构,而且还确保这些结构更可能在实际材料应用中具有热力学可行性。
在与其他基线方法(ADiT、MatterGen、DiffCSP++)比较时,也观察到了类似模式,CrystalDiT 始终生成更多稳定和亚稳态结构。这些全面的能量分布分析见 Appendix C,进一步强化了我们的主要发现:在晶体生成任务中,具有适当领域知识的简单且设计良好的架构能够优于复杂替代方案。
4.7 扩展到更大结构
CrystalDiT 在 MPTS-52(最多 52 个原子)上达到了 6.73% 的 SUN rate,相比 MP-20 仅下降 2%,显示出有效的泛化能力。详细结果见 Appendix C。
5 结论
我们提出了 CrystalDiT,证明简化架构能够在晶体生成中显著优于复杂替代方案。我们的统一扩散 Transformer 在 MP-20 上实现了 8.78% 的 SUN rate,并通过三项关键贡献优于现有方法:(1)一种使用统一注意力机制的简化架构;(2)一种使用元素周期表位置的二维原子表示;以及(3)一种优化生成质量和发现潜力的均衡模型选择策略 Balance Score。
除生成孤立结构之外,未来研究的重要方向包括将我们的框架扩展到材料—分子相互作用系统(Bian, Wu, and Yan 2026; Bian et al. 2026),这可能促进用于催化、药物递送和能源应用的功能材料发现。此外,将目标性质约束纳入生成过程,将允许直接设计具有期望特性的材料,从而进一步推动 AI 驱动的材料发现。
6 致谢
本工作得到了广东省自然科学基金(资助号 2025A1515011946)和新加坡国立大学计算机学院(YB 的资助号 A-0010308-00-00)的支持。部分工作是在 Xiaohan Yi 和 Guikun Xu 任职于腾讯 AI Lab 期间完成的。我们感谢腾讯提供的计算资源,感谢 Tao Chen 富有洞见的讨论,并感谢匿名审稿人提出的建设性反馈。
A 评估指标
A.1 评估指标说明
我们的综合评估框架为每种方法生成 10,000 个晶体结构,并由两个主要阶段组成:(1)结构与成分评估,以及(2)热力学稳定性评估。本节提供详细的实现规范,以便复现。
阶段 1:结构与成分评估
初始结构解析与预过滤
所有 10,000 个生成的 CIF 文件都会先经过初始解析和系统性预过滤,并采用特定的有效性标准。预过滤过程按如下顺序进行检查:
- 使用 PyMatGen 的
CifParser解析每个 CIF 文件,以提取结构信息 - 提取晶格参数(长度和角度)以及原子信息
- 按如下顺序应用预过滤标准:
- 原子序数范围检查:所有原子序数必须位于
[1, 104] - 正晶格参数:所有晶格长度必须为正(
> 0) - NaN/Inf 值检查:晶格参数、角度或原子坐标中不能有 NaN 或无穷值
- PyMatGen
Structure构造:必须成功构造有效的Structure对象 - 最小体积检查:晶体体积必须
≥ 0.1 Å^3
- 原子序数范围检查:所有原子序数必须位于
- 任何一项预过滤标准不满足的结构都会被标记为无效,并给出具体错误原因
- 所有百分比计算都以 10,000 为分母,将预过滤掉的结构视为无效
预过滤作为一个全面的质量门控,确保只有具有物理意义的晶体结构进入后续的详细评估。具有无效原子组成、退化晶格参数、数值错误或不现实的小体积的结构都会被系统性排除。
结构有效性评估 stru_val
结构有效性使用基于距离的几何约束,遵循 FlowMM 协议(Miller et al. 2024):
\mathrm{Structural\ Validity} = \begin{cases} \mathrm{True} & \text{if } \min(\mathrm{distance\ matrix}) \ge 0.5 \ \mathrm{\AA} \text{ and } V_{\mathrm{cell}} \ge 0.1 \ \mathrm{\AA}^3 \\ \mathrm{False} & \text{otherwise} \end{cases}
其中,距离矩阵不包括自相互作用(对角线元素设为无穷大)。
成分有效性评估 comp_val
化学有效性使用 SMACT 框架,并带有超时保护:
- 从
Structure对象中提取元素组成 - 枚举所有元素的氧化态组合
- 应用电荷中性条件:
- 应用 Pauling 电负性检验
- 金属合金例外:所有金属结构都视为有效
- 每个结构设置 30 秒超时,以防止无限枚举
分布距离指标
性质分布保真度通过 Wasserstein 距离量化。对于同时通过结构和化学有效性测试的结构,我们随机采样 1000 个结构(若少于 1000 个则使用全部),并与 MP-20 测试集进行比较:
其中, 表示密度分布, 表示元素出现频率分布。
唯一性与新颖性评估
这两个指标都使用 PyMatGen 的 StructureMatcher 默认参数:
ltol: 0.2(晶格参数容差)stol: 0.3(位点位置容差)angle tol: 5°(晶格角度容差)primitive cell:True,scale:True,attempt supercell:False
化学系统分组优化: 结构首先按化学系统(按元素类型排序)分组,再进行两两比较。该优化将计算复杂度从 降低到 ,且经测试不影响结果。
单元素过滤: 在 UN 评估阶段,仅包含一种元素类型的结构会被排除,不参与唯一性和新颖性计算。该过滤仅用于 UN rate 计算,而不用于初始预处理阶段。
唯一性: 将结构与同一化学系统中所有其他生成结构进行比较。若某结构与任何其他结构都不匹配,则该结构是唯一的。
新颖性: 将生成结构与 MP-20 训练集在同一化学系统内进行比较,遵循 FlowMM 的化学系统优化方法。
阶段 2:热力学稳定性评估
候选结构选择与 CHGNet 预弛豫
从阶段 1 识别出的 UN 结构中,我们随机抽取 500 个结构进行稳定性评估。每个结构都使用增强参数进行 CHGNet 预弛豫:
- 最大弛豫步数:2,500(相比 FlowMM(Miller et al. 2024)中的 1,000)
- 结构优化:同时优化原子位置和晶格参数
DFT 候选过滤
根据 CHGNet 预测的凸包上方能量选择进行 DFT 计算的结构:
这一扩大阈值(相比 FlowMM 的 0.1 eV/atom)能够实现更全面的稳定性评估,约可从 500 个 UN 结构中选出 490 个进行 DFT 计算,而 FlowMM 仅评估 261 个结构。
DFT 弛豫协议
DFT 计算采用 VASP(Kresse and Furthmüller 1996),并使用 Materials Project Relaxation Set(MPRelaxSet)参数(Jain et al. 2013),包括 PBE-GGA 交换关联泛函、PAW 赝势,以及合适的 k 点采样和收敛判据。
凸包上方能量计算
最终稳定性评估使用最新的 Materials Project 凸包(2024 年 12 月版本),并采用 Materials Project 2020 兼容性修正:
最终发现指标
A.2 结果分析
MatterGen 评估差异分析
由于数据集特定的评估协议,我们对 MatterGen 的评估结果与原论文不同。Table 3 比较了使用不同参考数据集时的新颖性评估。
这种差异的产生是因为 MatterGen 原始评估使用了更大的 Alexandria-MP-20 数据集作为新颖性参考。Alexandria-MP-20 包含 607,683 个训练结构,而 MP-20 仅包含 27,136 个训练结构,因此新颖性评估更严格。由于我们评估的是在 MP-20 上训练的 MatterGen 模型,因此与 MP-20 训练集比较才是合适的新颖性评估。我们的结果(89.89% UN rate)代表了基于 MP-20 的模型性能的正确评估。
FlowMM 评估协议差异
由于两个方法学改进,我们对 FlowMM 的重新评估得到了比原论文更高的稳定性比例:
更新的凸包数据库: 我们使用最新的 Materials Project 凸包(2024 年 12 月版本),与原始 FlowMM 评估中使用的早期版本相比,它提供了更准确的热力学参考数据。
扩展的 DFT 候选选择: 我们的协议评估更多结构,以进行更全面的稳定性分析:
- FlowMM 协议:CHGNet eV/atom → 评估 261/500 个结构;
- 我们的协议:CHGNet eV/atom → 评估 490/500 个结构。(由于)
这一扩展评估更能代表 UN 结构总体中的稳定性比例,减少了因过于严格的预筛选而可能产生的采样偏差。我们观察到的更高稳定性比例,既反映了更好的热力学参考数据,也反映了更全面的总体采样。
B 消融研究
B.1 双流架构
为验证“在晶体生成中,架构简单性优于复杂性”这一假设,我们实现了一个带有交叉注意力机制、级联处理的复杂双流变体。该复杂架构作为与我们统一方法的直接比较对象。
架构设计
双流架构实现了一个复杂的三阶段级联处理范式,如 Figure 3 所示。经过晶体嵌入后,架构先将 20 个原子 token 和 3 个晶格向量 token 分别通过独立路径处理,然后再进行复杂的联合处理。
阶段 1:独立原子处理
原子路径通过 12 个专门的 DiT 块,使用纯自注意力机制处理 20 个原子 token。这使模型能够学习原子集合内部复杂的原子间关系和化学模式,而不受晶格几何约束的干扰。
阶段 2:独立晶格处理
并行地,3 个晶格向量 token 通过 2 个专门的 DiT 块进行自注意力处理。该独立路径使模型能够独立于原子组成信息,学习晶格特有的几何和晶体学模式。
阶段 3:复杂联合处理
最后阶段使用 2 个联合 DiT 块,在每个块内实现一个复杂的四步注意力机制:
- 初始自注意力:分别对原子和晶格特征进行自注意力
- 双向交叉注意力:同时进行晶格→原子和原子→晶格的信息交换
- 交叉后的自注意力:在跨模态信息融合后再进行额外自注意力
- 独立 MLP 处理:两个流分别通过独立前馈网络处理
每个联合块中的注意力流如下:
这种带有复杂交叉注意力机制的级联设计,相比我们的统一方法具有显著的架构复杂性,需要在多个处理流和注意力机制之间进行精细协调。
计算复杂度分析
我们在相同训练条件下,使用 8×V100 GPU、批大小 256,测量了两种架构的实际计算开销。
该复杂架构引入了 7.98% 的训练时间开销,主要来自联合块中的双向交叉注意力计算。虽然这一开销看起来可控,但性能下降却很显著:复杂模型的 SUN rate 为 6.36%,而简单架构为 9.62%,发现性能相对下降了 34%。
这一实证分析揭示了一个关键洞见:架构复杂性的计算开销相对较小,但对泛化能力的影响却很严重。交叉注意力机制虽然在计算上可处理,但会从根本上改变模型的学习动态,促使其过拟合训练数据模式。这表明,对于科学生成建模任务而言,瓶颈不在计算资源,而在于引导模型学习的架构归纳偏置。
交叉注意力机制分析
联合块中的双向交叉注意力机制旨在捕获原子位置与晶格参数之间错综复杂的关系。然而,我们的分析表明,这种复杂的交互建模会导致过拟合行为。复杂模型取得了更高的单项有效性分数(结构有效性 98.39%,成分有效性 89.44%),但新颖性比例显著更低,这表明交叉注意力机制可能促进对训练模式的记忆,而非对新晶体结构的泛化。
B.2 一维原子表示
为评估我们基于二维元素周期表编码的贡献,我们实现了一个使用传统一维原子序数表示的基线。
实现细节
一维表示直接使用原子序数对原子进行编码,并归一化到 [-1, 1]:
其中, 是原子序数, 表示 MP-20 数据集中的最大原子序数。无效原子(空位置)表示为 。
因此,原子特征向量变为 4 维:[z_norm, x, y, z],而不是我们 5 维的周期表示 [r_norm, c_norm, x, y, z]。这种编码保留了原子身份,但丢失了元素周期表结构中蕴含的内在化学关系。
化学知识整合
这两种表示的关键区别在于化学知识的整合方式:
一维原子序数: 由原子序数直接映射提供了精确的元素识别,但将每个元素视为彼此独立。元素之间的化学相似性(例如碱金属、过渡金属)没有在表示空间中被显式编码。
二维元素周期表: 我们的方法通过在归一化的(周期,族)空间中的空间邻近性来嵌入化学关系。具有相似性质的元素会自然聚集在嵌入空间中,为化学推理提供归纳偏置。
性能分析
一维表示的 UN rate 更高(78.47% vs 63.28%),但稳定性比例显著更低(UN 结构中的稳定率为 8.0% vs 15.2%)。这一权衡导致最终发现指标更差(SUN rate 为 6.2% vs 9.62%),表明嵌入在元素周期表结构中的化学知识增强了生成晶体的实际效用。
B.3 架构深度消融
为验证我们选择 18 个 DiT 块的合理性,我们针对不同模型深度进行了系统消融研究。Table 5 给出了综合结果。
分析
结果揭示了模型容量与泛化之间的清晰权衡:
浅层模型(6 层): 具有最高的 UN rate(82.4%),但 SUN rate 较低(5.78%),表明其探索过度而质量控制不足。有限容量无法捕获复杂的稳定性模式。
最优深度(18 层): 提供了最佳发现性能,SUN rate 为 8.78%。该深度在生成新颖结构与确保热力学稳定性之间实现了最优平衡。
深层模型(24 层): 与 18 层相比,表现出过拟合行为,UN rate 降低至 56.8%,SUN rate 也略降至 7.10%。过大的容量导致对训练模式的记忆,削弱了探索能力。
这一系统分析验证了我们的架构选择,并表明模型深度是晶体生成任务中的关键超参数。
B.4 归一化消融
为验证我们晶格归一化策略的重要性,我们在不使用 Sec 3.1.2 所述归一化的情况下训练 CrystalDiT。
结果
未归一化模型达到 49.4% 的 UN rate、4.70% 的 SUN rate 和 18.04% 的 MSUN rate,明显低于我们的归一化方法(63.3% UN、8.78% SUN、25.90% MSUN)。
分析
如果没有归一化,晶格参数值的宽范围(MP-20 中约从 3 Å 到 46 Å)会造成训练不稳定,并阻碍晶格—原子相关性的有效学习。这验证了在材料科学应用中的扩散模型里,适当的特征缩放至关重要。
B.5 元素分布分析
Figure 4 展示了我们三种模型变体的元素数量分布,揭示了关于化学多样性和复杂性处理的重要洞见。
分布特征
三种变体在整体上都呈现出相似模式,在 3 元组成附近达到峰值,这与 MP-20 训练数据分布一致。然而,仍存在细微但重要的差异:
CrystalDiT(Complex): 显示出最受约束的分布(均值:3.016 个元素,标准差:0.851),其中 50.38% 的结构恰好包含 3 种元素。较窄的分布表明复杂架构倾向于保守、类似训练数据的组成。
CrystalDiT(Simple): 表现出稍高的多样性(均值:3.199 个元素,标准差:0.925),包含 4 种及以上元素的结构更多(31.76%,而复杂模型为 24.62%)。更宽的分布表明在保持稳定性的同时,对化学空间有更好的探索。
CrystalDiT(1D atomic): 展现出最高的变异性(标准差:1.016),但这种多样性提升是以稳定性为代价的,这一点已在我们的主要结果中通过更低的稳定率体现出来。
化学复杂性与稳定性的权衡
元素分布分析揭示了晶体生成中的一个基本权衡:更高的化学多样性可以带来更多新颖发现,但必须与热力学稳定性保持平衡。我们的简单架构与二维元素周期表示实现了最优平衡,生成了适度多样、同时保持较高稳定率的化学组成。
复杂架构受限的分布与其过拟合行为一致——它生成的是更安全、更像训练数据的组成,但错失了新颖发现的机会。相反,一维表示在缺乏化学指导的情况下表现出过度多样性,导致不稳定且不切实际的结构。
这一分析进一步强化了我们的主要发现:架构简单性结合适当的化学知识(二维周期表示)是材料发现应用中最有效的方法。
C 额外结果
C.1 Balance Score 分析与训练动态
为了全面验证我们的 Balance Score 方法和训练策略,我们针对不同超参数设置下的训练动态和检查点选择进行了详细分析。
评估协议
我们的训练评估遵循系统化方案:每训练250个周期即保存一次检查点,并从每个检查点生成 1,000 个晶体结构以进行全面评估。所有训练动态图均采用10点移动平均值,以降低随机生成过程中固有的波动性(对应 Figure 5 & 6 中的 Smoothed line)。需注意,本分析中观察到的较高UN率(相较于主论文中的 10,000 种结构评估结果)是由于样本量较小所致,因为唯一性与新颖性指标会随着生成结构数量的减少而自然提升。
不同 alpha 值下的 Balance Score 验证
Figure 5 展示了 和 下 Balance Score 的演化,说明了质量—发现权衡参数的影响。两条曲线整体趋势相似,但在检查点选择偏好上有所不同。
关键差异出现在训练中期(30–60%)的选择: 选择 epoch 16,250,而 选择 epoch 24,750。这表明更高的 值会更强地强调质量分量,从而选择结构和成分有效性更好的较晚检查点。早期和后期阶段的选择在两种取值下保持一致,说明训练极端阶段的表现具有鲁棒性。
- 取值范围 ,若 更大,则 对 Balance Score 的结果影响更大。
训练动态与质量—发现权衡
Figure 6 展示了推动我们提出 Balance Score 方法的基本权衡。随着训练推进,我们观察到关键指标出现分化趋势,而传统验证方法无法捕捉这一点。
结构有效性从 95% 稳步提高到 98%,化学有效性从 80% 提升到 88%。然而,UN rate 呈现相反趋势,从约 99.5% 下降到 62%。这种分化凸显了传统基于有效性的模型选择方法的关键局限:擅长复现训练数据模式(高有效性)的模型,可能牺牲材料发现所需的探索能力(低 UN rate)。
所有指标中观察到的显著波动反映了扩散模型生成的内在随机性。这些波动使得我们的 10 点平滑窗口成为必要,并进一步强调了系统性检查点评估的重要性,而非依赖单点评估。
三阶段检查点选择的理由
我们的三阶段选择策略旨在应对质量—发现权衡的时间演化。早期阶段模型(epoch ≤ 15,000)具有较强的探索能力,UN rate 超过 80%,但有效性分数中等。后期阶段模型(epoch > 30,000)则获得了出色的有效性指标,但新颖生成能力下降。中期模型代表了最佳折中点,在质量和发现潜力之间保持平衡。
这一时间分析验证了我们的多阶段选择方法:没有任何单一训练阶段能持续产出材料发现的最优模型。相反,最佳检查点需要借助我们的 Balance Score 框架,在所有训练阶段进行系统评估后才能确定。
C.2 完整能量分布比较
热力学有利性分析
为了补充主文中与 FlowMM 的比较,Figure 7 提供了所有基线方法的完整能量分布分析。这些分布揭示了 CrystalDiT 在生成热力学有利晶体结构方面的优越能力。
在所有比较中,CrystalDiT 都展现出在热力学有利性方面的一致优势:
增强的稳定结构生成: 在热力学稳定区域(),CrystalDiT 的密度峰值显著高于所有基线方法。这表明其生成能量有利晶体结构的能力更强。
集中的能量分布: CrystalDiT 在 0 eV/atom 附近产生更尖锐的峰,说明与基线方法更宽、更分散的分布相比,其生成结构更接近热力学平衡。
减少高能结构: 在高能区域( eV/atom),CrystalDiT 生成的结构显著更少,说明它更能避免热力学不利的构型。
这些能量分布模式提供了关键验证:CrystalDiT 的简化架构不仅能生成更多唯一且新颖的结构,还能确保这些结构具有更优的热力学可行性,适用于实际材料应用。与所有基线方法比较时的稳定表现进一步强化了我们的结论:统一注意力机制在学习具有物理意义的晶体生成模式方面非常有效。
C.3 扩展到更大晶体结构
为评估 CrystalDiT 在超出 MP-20 20 原子限制之外的泛化能力,我们在 MPTS-52 上进行了实验。MPTS-52 是一个更具挑战性的数据集,包含最多 52 个原子的晶体结构。我们使用相同超参数训练 CrystalDiT(Simple),仅将最大原子数从 20 增加到 52。
结果
CrystalDiT 在 MPTS-52 上达到了 6.73% 的 SUN rate,尽管最大结构尺寸增加了 2.6 倍,相比 MP-20 性能仅下降 2.05 个绝对百分点。这一温和退化表明,我们的简化架构与元素周期表示有效捕获了可跨越不同结构复杂度迁移的基本化学原理,验证了该方法在发现更大、更复杂材料方面的实际价值。
D 可视化
D.1 生成的晶体结构
Figure 8 展示了由 CrystalDiT 生成的代表性晶体结构,这些结构同时满足高唯一性/新颖性和热力学稳定性。这些结构证明了模型能够在保持真实原子排列和空间群对称性的同时,发现多样的晶体化学。每个结构都代表了我们 SUN(Stable, Unique, Novel)评估中的一个成功示例,确认 CrystalDiT 能够生成具有真实发现潜力的材料。
这些可视化揭示了几个关键洞见:(1)CrystalDiT 可以跨越多样化的化学系统生成结构,从稀土化合物到过渡金属合金;(2)生成的结构具有正确的配位环境和真实的键长;(3)模型成功捕捉了多种空间群对称性,从高对称性的四方体系到低对称性的三斜结构;以及(4)化学组成遵循与材料科学原理一致的合理化学计量比。
这些示例所展示的结构多样性验证了我们方法在材料发现中的有效性。La2CeSn7 结构代表一种具有复杂层状排列的稀土锡化物,这种结构特征典型地存在于 RE2TSn7 系列化合物中;PmGa3Cu 呈现出稀土—镓—铜体系中常见的四方晶系结构;TbDy2In3Cu3 展示了 CrystalDiT 生成含有多种过渡金属的混合稀土化合物的能力;Tb3Dy3Pd8 则证明其能够成功生成复杂的铂族金属合金。
这些可视化结果确认,CrystalDiT 不仅在定量上实现了较高的 SUN rate,而且生成了化学上有意义且结构上真实的结构,使其成为实验合成和进一步材料性质探索的可行候选。