我们如何利用AI设计材料？

很高兴与@liun_online、Kostya Novoselov、@ylecun及团队合作，推出全新的生成技术Crys-JEPA。通过利用JEPA构建高质量的能量感知潜在空间，我们在MP20上实现了47.9%的VSUN。 https://t.co/mWrjazmiGo https://t.co/UdltCOr9ix

Crys-JEPA：通过嵌入筛选与生成精炼加速晶体发现

来源：https://arxiv.org/html/2605.14759 Nian Liu¹ Nikita Kazeev¹ Stephen Gregory Dale¹ Artem Maevskiy¹ Yuwei Zeng¹ Ryoji Kubo¹ Pengru Huang¹ Thomas Laurent² Yann LeCun³ ⁴ Kostya S. Novoselov¹ Xavier Bresson¹ {nianliu, yuweizeng, ryojikubo}@u.nus.edu, [email protected], [email protected] {kna, sdale, maevskiy, pengru, kostya, xaviercs}@nus.edu.sg ¹新加坡国立大学 ²洛约拉马利蒙特大学 ³纽约大学 ⁴Meta人工智能研究院

摘要

从头晶体生成的目标是发现不仅具有现实性，而且具备稳定性和新颖性的材料。然而，大多数现有生成模型通过最大化观测晶体的似然来训练，这鼓励生成的样本接近已知材料，但未必符合实际发现所需的标准。通过实证研究，我们发现当前的晶体生成模型深陷于稳定性-新颖性的权衡之中：向观测分布靠近会保留稳定性但限制新颖性，而向远离方向移动则会迅速破坏稳定性。这表明，能够同时发现稳定且新颖晶体的有用区域极为狭窄。为了摆脱这一权衡，我们引入了Crys-JEPA，一种面向晶体的联合嵌入预测架构，它学习了一个能量感知的潜在空间，能够保留生成能差异。在这个空间中，稳定性评估可以转化为基于嵌入的比较，与训练集中的晶体进行参照，从而减少对昂贵的能量评估和特定任务的外部参考的依赖。基于Crys-JEPA，我们进一步构建了一个筛选-精炼流水线：识别有前途的生成晶体，并将其重新引入以精炼生成模型。在MP-20和Alex-MP-20数据集上，我们的方法在V.S.U.N指标上分别取得了比基线高81.4%和82.6%的提升。

1 引言

发现新材料是推动太阳能电池、电池和催化等广泛应用领域进步的关键驱动力[7 (https://arxiv.org/html/2605.14759#bib.bib52)]。在材料类别中，晶体尤为重要，因为其周期性的原子排列产生了多样且可调的物理性质，使其成为计算材料设计的核心目标。这催生了从头生成（DNG）任务[48 (https://arxiv.org/html/2605.14759#bib.bib6)]，旨在不依赖预定义模板的情况下发现全新的晶体结构。近年来，DNG因深度生成模型（尤其是扩散[18 (https://arxiv.org/html/2605.14759#bib.bib3)]和流匹配[29 (https://arxiv.org/html/2605.14759#bib.bib51)]）而取得了显著进展。

大多数现有DNG模型[48 (https://arxiv.org/html/2605.14759#bib.bib6), 53 (https://arxiv.org/html/2605.14759#bib.bib12), 23 (https://arxiv.org/html/2605.14759#bib.bib40)]在训练中采用最大化观测数据对数似然的传统目标。然而，尚不清楚该目标在多大程度上能够改善实际中最关键的指标，即有效性（V）、稳定性（S）、唯一性（U）和新颖性（N），我们统称为V.S.U.N。在第三节（https://arxiv.org/html/2605.14759#S3）中，我们首先通过实证研究揭示了当前晶体生成模型通常表现出显著的稳定性-新颖性权衡。为了更好地理解这一现象，我们分析了这些指标如何随观测分布的密度景观变化。我们的结果表明，晶体生成对精度要求极高：向高密度区域靠近通常不足以提高新颖性，而向低密度区域偏移则可能破坏稳定性。换句话说，几乎没有有效的中间区域能让晶体同时兼具稳定性和新颖性。因此，尽管对数似然最大化鼓励生成样本停留在观测分布附近，但这些附近区域并不一定同时满足稳定性和新颖性。

值得注意的是，标准训练数据集主要由已知的稳定晶体组成。因此，在不改变底层生成骨干网络的情况下，缓解稳定性-新颖性权衡的一个自然方法是将满足V.S.U.N.的晶体（尤其是同时稳定且新颖的晶体）重新引入训练，使模型能够拟合更理想的分布。如果出于公平性考虑无法使用外部数据，另一种方法则是直接从模型生成中识别这些有前途的晶体。然而，这需要对生成样本进行可靠的V.S.U.N.评估。虽然有效性、唯一性和新颖性可以相对廉价地评估，但稳定性的评估则困难得多。稳定性是通过与参考集的生成能比较来定义的，这带来了两个基本挑战。首先，参考模糊性：参考的选择取决于具体任务，即使像Materials Project[20 (https://arxiv.org/html/2605.14759#bib.bib53)]这样的公共参考集也会随时间演变，导致在训练过程中无法预先知道最终的参考集。其次，计算成本：标准的稳定性评估通常依赖密度泛函理论（DFT），在大规模下成本高昂。因此，对所有生成的晶体进行DFT计算是不切实际的，尤其是一些生成样本质素较低。

在这项工作中，我们共同应对这两个挑战。为缓解参考模糊性，我们相对于训练集来评估稳定性。其关键直觉是，在常用参考集下，训练晶体本身已经是稳定的。因此，如果生成的晶体在相同化学体系下的生成能与训练晶体相当，那么它也大概率是稳定的。这大大降低了对显式外部参考的依赖。为避免昂贵的能量计算，我们进一步引入了一个能量感知的替代模型Crys-JEPA。具体而言，我们使用InfoNCE目标[36 (https://arxiv.org/html/2605.14759#bib.bib28)]预训练一个联合嵌入预测架构（JEPA）[27 (https://arxiv.org/html/2605.14759#bib.bib27)]，构建一个由生成能结构化的晶体潜在空间，使得生成能相近的晶体映射到邻近位置，而能量差异大的晶体则被良好分离。然后，我们使用基于Crys-JEPA嵌入的比较作为稳定性评估的代理，从而实现对生成晶体的高效筛选。基于这一替代模型，我们构建了一个简单的精炼循环：预训练一个基础生成模型，生成候选晶体，使用Crys-JEPA筛选有前途的样本，然后在选中的样本上微调基础模型。实验表明，这种筛选-精炼流水线显著提高了生成质量。

我们的贡献总结如下：

• • 我们识别了晶体从头生成中的稳定性-新颖性权衡，并表明其根源在于晶体同时保持稳定和新颖所需的极高精度。
• • 我们开发了Crys-JEPA，一种用于DFT稳定性评估的能量感知潜在替代模型，并展示了其能够驱动简单的精炼循环，从而提高生成质量。
• • 我们的方法在MP-20和Alex-MP-20数据集上持续优于强基线，通过DFT评估，V.S.U.N指标分别提升了81.4%和82.6%。

2 预备知识

晶体表示。

晶体C\mathbf{C}通过其基本重复单元（即晶胞）在三维空间中的周期性排列来定义。我们通常使用三个分量来描述晶胞：原子分数坐标X∈[0,1)N×3\mathbf{X} \in [0,1)^{N\times 3}、原子类型A∈R N\mathbf{A} \in \mathbb{R}^{N}，以及晶格矩阵L∈R3×3\mathbf{L} \in \mathbb{R}^{3\times 3}，其中NN为晶胞中的原子数。对于A\mathbf{A}，我们考虑前100种化学元素，并将每种原子类型编码为一个独热向量，得到one-hot(A)∈R N×100\text{one-hot}(\mathbf{A}) \in \mathbb{R}^{N\times 100}。为了表示晶格，我们采用先前工作中的重参数化方法[53 (https://arxiv.org/html/2605.14759#bib.bib12)]。具体地，我们通过奇异值分解对L\mathbf{L}进行分解，并将其重新表示为旋转矩阵和对称矩阵的乘积：

L=U L~,U=W V⊤,L~=V Σ V⊤,\mathbf{L} = \mathbf{U} \tilde{\mathbf{L}}, \quad \mathbf{U} = \mathbf{W} \mathbf{V}^{\top}, \quad \tilde{\mathbf{L}} = \mathbf{V} \mathbf{\Sigma} \mathbf{V}^{\top}, (1)

其中W\mathbf{W}和V\mathbf{V}分别为L\mathbf{L}的左右奇异矩阵，Σ\mathbf{\Sigma}是对角线上的奇异值。在此公式下，U\mathbf{U}对应一个旋转矩阵，而L~\tilde{\mathbf{L}}是对称正定矩阵。然后，我们取L~\tilde{\mathbf{L}}的上三角部分并将其展平为一个6维向量，记为L^=vec(triu(L~))∈R6\widehat{\mathbf{L}} = \mathrm{vec}(\mathrm{triu}(\tilde{\mathbf{L}})) \in \mathbb{R}^{6}。基于这些分量，晶体C\mathbf{C}中的第ii个原子由一个原子向量表示：

v i=[X i ‖ one-hot(A i) ‖ L^]∈R3+100+6,\mathbf{v}{i} = [\mathbf{X}{i} , | , \text{one-hot}(A_{i}) , | , \widehat{\mathbf{L}}] \in \mathbb{R}^{3+100+6}, (2)

该向量拼接了其坐标、元素类型和晶格表示。整个晶体的表示通过堆叠所有原子向量得到：

V=[v 1,…,v N]⊤∈R N×(3+100+6).\mathbf{V} = [\mathbf{v}{1}, \dots, \mathbf{v}{N}]^{\top} \in \mathbb{R}^{N\times (3+100+6)}. (3)

热力学稳定性。

对于晶胞内具有NN个原子的晶体C\mathbf{C}，总能量记为E tE_{t}，每个原子的总能量定义为E t/atom=E t/NE_{t/\mathrm{atom}} = E_{t}/N。假设C\mathbf{C}包含kk种不同的原子种类{T1,…,Tk}{T_{1}, \dots, T_{k}}。其化学体系记为T1T_{1}–T2T_{2}–⋯–TkT_{k}，成分由f=(f 1,…,f k)\mathbf{f} = (f_{1}, \dots, f_{k})表示，其中f i≥0f_{i}\ge 0且∑i f i=1\sum_{i} f_{i} = 1。我们将C\mathbf{C}与成分-能量空间中的一个条目P=(f,E t/atom)\mathcal{P} = (\mathbf{f}, E_{t/\mathrm{atom}})相关联。

DNG的主要需求是生成在热力学意义上稳定的晶体，这最初基于“生成能”来定义。设μ i ref\mu_{i}^{\mathrm{ref}}为物种T iT_{i}在相同计算设置下从其稳定元素相中取出的每原子元素参考能量。C\mathbf{C}的每个原子的生成能定义为

E f/atom(C)=E t/atom−∑ i=1 k f i μ i ref.E_{f/atom}(\mathbf{C}) = E_{t/\mathrm{atom}} - \sum_{i=1}^{k} f_{i} \mu_{i}^{\mathrm{ref}}. (4)

给定参考数据集R\mathcal{R}，我们使用R\mathcal{R}中与C\mathbf{C}属于相同化学体系或其任何子体系的所有条目构建一个相图。对于每个参考条目P j=(f j,E t/atom j)\mathcal{P}^{j} = (\mathbf{f}^{j}, E_{t/\mathrm{atom}}^{j})，其生成能为E f/atom j=E t/atom j−∑ i=1 k f i j μ i refE_{f/\mathrm{atom}}^{j} = E_{t/\mathrm{atom}}^{j} - \sum_{i=1}^{k} f_{i}^{j} \mu_{i}^{\mathrm{ref}}。凸包定义为成分-生成能空间的下凸包络。相应地，在成分f\mathbf{f}处的凸包生成能为

E f/atom hull(f)=min{λ j}∑ j λ j E f/atom j s.t. ∑ j λ j f j=f, ∑ j λ j=1, λ j≥0.E_{f/\mathrm{atom}}^{\mathrm{hull}}(\mathbf{f}) = \min_{{\lambda_j}} \sum_{j} \lambda_j E_{f/\mathrm{atom}}^{j} \quad \text{s.t.} \quad \sum_{j} \lambda_j \mathbf{f}^{j} = \mathbf{f}, ;; \sum_{j} \lambda_j = 1, ;; \lambda_j \ge 0. (5)

凸包上方的能量定义为

ΔE=E f/atom(C)−E f/atom hull(f).\Delta E = E_{f/\mathrm{atom}}(\mathbf{C}) - E_{f/\mathrm{atom}}^{\mathrm{hull}}(\mathbf{f}). (6)

在附录A (https://arxiv.org/html/2605.14759#A1)中，我们推导出ΔE\Delta E也可以通过每原子总能量来表示：

ΔE=E t/atom−E t/atom hull(f),\Delta E = E_{t/\mathrm{atom}} - E_{t/\mathrm{atom}}^{\mathrm{hull}}(\mathbf{f}), (7)

其中E t/atom hull(f)=∑ j λ j E t/atom jE_{t/\mathrm{atom}}^{\mathrm{hull}}(\mathbf{f}) = \sum_{j} \lambda_j E_{t/\mathrm{atom}}^{j}。在这项工作中，我们认为晶体是热力学稳定的，如果ΔE<ε\Delta E < \varepsilon，其中ε=0.1\varepsilon = 0.1 eV/atom，遵循先前的研究[53 (https://arxiv.org/html/2605.14759#bib.bib12)]。

3 稳定性与新颖性的权衡：一项实证研究

在本节中，我们研究对数似然最大化与V.S.U.N评估标准之间的关系，特别关注稳定性和新颖性。

发现晶体生成中的权衡。

我们在MP-20数据集[48 (https://arxiv.org/html/2605.14759#bib.bib6)]上复现了多种基线模型，包括CDVAE[48 (https://arxiv.org/html/2605.14759#bib.bib6)]、DiffCSP[21 (https://arxiv.org/html/2605.14759#bib.bib30)]、DiffCSP++[22 (https://arxiv.org/html/2605.14759#bib.bib46)]、FlowMM[33 (https://arxiv.org/html/2605.14759#bib.bib19)]、FlowLLM[42 (https://arxiv.org/html/2605.14759#bib.bib50)]、SymmCD[28 (https://arxiv.org/html/2605.14759#bib.bib47)]、ADiT[23 (https://arxiv.org/html/2605.14759#bib.bib40)]、CrysLLMGen (7B)[24 (https://arxiv.org/html/2605.14759#bib.bib48)]、SGEquiDiff[8 (https://arxiv.org/html/2605.14759#bib.bib45)]和MatterGen[53 (https://arxiv.org/html/2605.14759#bib.bib12)]。这些模型的详细描述可在附录E (https://arxiv.org/html/2605.14759#A5)中找到。

在评估中，我们重复采样10次，每次收集1,000个晶体。然后报告稳定性（S）和新颖性（N）的均值和标准差，以及两个组合指标，即S.U.N.和V.S.U.N.。然而，在本案例研究中，我们主要关注稳定性和新颖性。在评估期间，结构弛豫和能量预测使用MatterSim-v1-1M[52 (https://arxiv.org/html/2605.14759#bib.bib20)]进行。遵循先前工作[33 (https://arxiv.org/html/2605.14759#bib.bib19)]，我们使用MP-2023作为参考数据集，并将凸包能量低于0.1 eV/原子的晶体视为稳定[53 (https://arxiv.org/html/2605.14759#bib.bib12)]。

见图注

图1：(a) 当前晶体生成模型所呈现的稳定性与新颖性之间的权衡。(b) 当生成晶体从靠近观测分布的区域向远离方向移动时，稳定性、新颖性和S.U.N.的变化趋势，使用本节定义的代理距离进行测量。

结果总结在表1 (https://arxiv.org/html/2605.14759#S5.T1)中。我们进一步在图1 (https://arxiv.org/html/2605.14759#S3.F1)(a)中可视化了稳定性与新颖性之间的关系。例如，CDVAE[48 (https://arxiv.org/html/2605.14759#bib.bib6)]生成了高度新颖的晶体，但其稳定性相对有限。相反，