Discovery and recovery of crystalline materials with property-conditioned transformers

本文介绍了 CrystaLLM-{\pi}，这是一种条件自回归框架，它将连续属性表示直接集成到 Transformer 注意力机制中，以实现晶体材料的稳健逆向设计，成功展示了从 X 射线衍射图谱恢复结构以及生成具有目标带隙的新型稳定光伏候选材料的能力。

要点

想象一下，你正在尝试发明一种新型太阳能电池板，或者仅通过观察阴影来推测某种神秘晶体的外观。长期以来，科学家们不得不依靠“试错法”，这种方法既缓慢又昂贵。最近，计算机开始利用“生成式人工智能”来协助设计这些材料，这就像一位能够发明新食谱的厨师。

然而，当前的“人工智能厨师”存在一个问题。如果你要求它们：“给我做一个含糖量正好为 20% 的蛋糕”，它们往往会感到棘手。它们可能会试图将"20%"拼写出来（例如拼成"t-w-e-n-t-y"），从而破坏了食谱的连贯性；或者，因为它们过于专注于糖分数字，而忘记了如何正确地烘焙蛋糕。

本文介绍了一种名为CrystaLLM-𝜋（读作"CrystaLLM-pi"）的新型人工智能系统，它解决了这一问题。以下是其工作原理，采用简单的类比进行说明：

1. 问题所在：“离散”与“连续”的冲突

将人工智能想象成一位正在弹奏钢琴的音乐家。钢琴琴键（音符）是离散的——你只能按下 C 或 C#，无法按下两者之间的音符。

• 旧方法：为了指示人工智能制造具有特定属性（如特定的“带隙”或密度）的材料，旧方法迫使人工智能将该数字当作单词处理。这就像要求音乐家通过逐个拼写音符名称的字母来演奏特定的音符。这种方法笨拙、令人困惑，且往往导致音乐（即材料）听起来错误或不稳定。
• 新方法（CrystaLLM-𝜋）：该系统不再拼写数字，而是为音乐家提供了一个连续旋钮。你只需将旋钮调至所需的精确设置，人工智能在演奏时便能直接感知该设置。它无需停下来思考数字，而是直接“感知”你想要的氛围。

2. 解决方案：两个新的“旋钮”（前缀与残差）

研究人员构建了两种具体方法，将这些旋钮连接到人工智能的大脑（基于一种称为 Transformer 的人工智能类型）：

• “前缀”方法（幽灵音符）：想象人工智能正在写一个故事。前缀方法在故事的开头添加几个“幽灵音符”，向人工智能低语目标属性。这些音符不会改变故事的长度或结构；它们只是设定了基调。人工智能在撰写故事其余部分（即晶体结构）时，会始终铭记这一基调。
• “残差”方法（背景嗡嗡声）：这就像有一个背景嗡嗡声在轻轻推动人工智能。如果人工智能开始撰写不符合目标属性的内容，嗡嗡声就会变大，温柔地将其引回正轨。如果人工智能已经走在正确的道路上，嗡嗡声就会保持安静。这种方法非常灵活，允许人工智能优雅地处理缺失的信息。

3. 他们测试了什么？

团队主要通过两种方式测试了这一新系统：

A. 发明新型太阳能材料（发现）
他们要求人工智能设计用于太阳能电池板的高效新材料。

• 结果：人工智能成功生成了数千种它从未见过的全新稳定晶体结构。
• 验证：他们挑选了最佳候选者，并通过超精确的物理模拟（称为 DFT）进行了测试。其中几种由人工智能设计的材料被证明是稳定的，并具有他们所需的高效率。这就像人工智能发明了一种新食谱，而当厨师真正烹饪它时，味道非常美味。

B. 从阴影中解开谜团（复原）
有时，科学家拥有一种晶体，但不知道其确切形状。他们只有 X 射线衍射图案（这就像晶体的阴影或条形码）。

• 结果：研究人员将这些“阴影”输入到 CrystaLLM-𝜋 中。人工智能能够以高精度重建原始的三维晶体结构。
• 验证：即使对于复杂的晶体，它也能发挥作用，并且能够区分同种材料的不同版本（多晶型），例如区分金红石和锐钛矿（二氧化钛的两种不同形态），尽管人工智能在训练期间从未见过这些特定形态。

4. 为什么这很重要？

• 更轻量、更快速：与其他需要海量计算能力（如超级计算机）的人工智能模型不同，该系统在标准显卡上即可高效运行。
• 不会遗忘：人工智能的一个常见问题是，当你教它新技巧时，它会忘记之前所学的一切。CrystaLLM-𝜋 的设计使其能够学习这些新的“旋钮”，而不会忘记如何构建基本晶体。
• 灵活性强：你可以利用它来发明新材料，也可以用它来解决旧谜题，所有这些都基于同一底层系统。

总结

简而言之，CrystaLLM-𝜋 是一种利用人工智能设计晶体的更智能的方法。它不再强迫人工智能“拼写”出所需的属性，而是让人工智能直接“感知”这些属性。这使得科学家能够比以往更快、更准确地发明用于太阳能等领域的新材料，或解析未知晶体的结构。该论文表明，这种方法在实践中行之有效，能够生成真实的、稳定的材料，并通过严格的科学测试。

技术摘要

技术摘要：基于属性条件 Transformer 的晶体材料发现与重构

问题陈述

新功能材料的发现历来受到无机系统巨大成分空间及候选材料表征计算成本高昂的阻碍。虽然生成式人工智能为加速这一过程提供了途径，但现有的基于 Transformer 的方法在尝试逆向设计（生成具有特定目标属性的结构）时面临显著局限。

标准 Transformer 架构通常依赖离散的、基于数字级别的标记化（tokenization）来编码连续的物理属性（如带隙、密度）。这种方法引入了几个关键问题：

1. 表征冲突：使用相同的离散标记来编码局部空间坐标和全局连续属性，导致嵌入空间内产生冲突。
2. 序数关系丢失：数字级别的标记化破坏了连续物理值固有的序数关系，阻碍了平滑的数学插值。
3. 灾难性遗忘：在特定属性数据集上微调预训练模型通常需要架构修改，这会破坏在晶体学信息文件（CIF）的无监督预训练期间获得的基础结构知识。
4. 低效性：序列级条件控制（例如在序列前添加属性标记）增加了序列长度，并扰动了控制 CIF 生成的标记表示，导致不稳定性及结构有效性降低。

方法论：CrystaLLM-$\pi$

为了解决这些局限，作者提出了CrystaLLM-$\pi$（属性注入），这是一种条件自回归框架，将连续属性表示直接集成到 Transformer 的注意力机制中，从而绕过序列级标记化。

核心架构

该模型建立在原始 CrystaLLM 的 GPT-2 架构之上，并在大量未标记 CIF 语料库上进行了预训练。为了实现基于属性的生成，该框架引入了两种新颖的注意力机制，将连续条件向量（$c \in \mathbb{R}^P$）直接注入到多头注意力（MHA）层中：

1. 属性键值（PKV）前缀注意力：

   • 受前缀微调（Prefix Tuning）启发，该方法从条件向量生成“幽灵”键值（KV）对。
   • 这些 KV 对在注意力层中与输入序列的 KV 对进行拼接。
   • 这施加了一种“硬”结构偏置，扩展了有效上下文窗口，同时无需修改前馈神经网络（FFNN）层或输入标记。

2. PKV 残差注意力：

   • 该方法引入了一种“软”条件控制机制。
   • 它在输入查询与条件导出的键/值之间计算并行的“残差”注意力分数。
   • 最终的注意力输出是基础自注意力与残差项的加权和：$A_{out} = A_{base} + \alpha \cdot A_{Residual}$。
   • 权重 $\alpha$ 初始化为零（类似于 LoRA），确保模型最初依赖预训练知识以减轻灾难性遗忘。与改变序列长度从而改变 softmax 归一化的前缀注意力相比，该架构能更优雅地处理缺失或未指定的条件。

训练策略

• 双重优化：采用双重学习率策略。对预训练骨干参数应用保守的学习率以保留基础结构知识，而对新初始化的条件控制层使用较高的学习率以加速适应。
• 损失函数：使用修改后的交叉熵损失函数，纳入对固定 CIF 语法标记的惩罚，以在训练早期加速语法学习。
• 数据处理：该框架利用动态边界标记化方案和条件对齐的批处理来稳定训练。

主要贡献

1. 新颖的条件控制机制：提出了 PKV 前缀和 PKV 残差注意力机制，将连续属性直接集成到注意力层中，避免了序列级标记化的低效性。
2. 结构先验的保留：证明了注意力级条件控制保留了无监督预训练中的丰富结构知识，即使在标签监督稀缺的情况下也能保持高结构有效性。
3. 全面的基准测试：在不同数据集规模（1K 到 653K 样本）和不同的材料设计任务上进行了系统评估，提供了序列级与注意力级条件控制之间的标准化比较。
4. 开源框架：发布了一个轻量级、灵活且可扩展的框架，包含预训练模型、容器化 API 以及用于便捷材料发现的网络界面。

结果

1. 鲁棒性与条件控制有效性

• 带隙条件控制：在 MP 带隙数据集上，预训练模型显著优于从头训练的模型，特别是在目标分布的尾部。前缀架构在不同数据集规模下展现了有效性、校准性和数据效率之间的最佳整体权衡，在完整数据下实现了高 $R^2$ 值（0.97）和低的密度目标平均绝对误差（0.72 g/cm³）。
• 数据稀缺：在低数据 regime（1K 样本）中，残差架构表现出更优越的鲁棒性，在其他方法难以维持结构有效性的情况下仍能保持，这很可能归因于其加性性质最大限度地减少了对预训练先验的干扰。
• 与扩散模型的比较：与基于图的扩散模型 MatterGen 相比，CrystaLLM-$\pi$ 实现了更紧密的目标属性校准，计算成本显著更低（显存更少、训练/推理更快），且在未弛豫输出中具有更高的对称性保持率。

2. 材料发现：光伏候选材料

• 该模型在 5.35K 个标记有光谱限制最大效率（SLME）的无机结构数据集上进行了微调。
• 以 33.2% 的目标 SLME 为条件，模型生成了 16,463 个结构新颖的候选材料。
• DFT 验证：使用密度泛函理论（DFT）验证了候选材料子集。几种材料，如 Cs$_2$NaInAs$_2$（SLME 26.4%）和 NaHfCuS$_3$（SLME 23.3%），被确认为稳定且高效的候选材料。
• 研究强调了从头算（ab initio）验证的重要性，因为一些具有高代理预测 SLME 的候选材料在混合 DFT 表征中由于细微的电子结构特征（例如分裂的导带底）而失败。

3. 从 XRD 重构结构

• 该框架在从 X 射线衍射（XRD）图谱重构晶体结构的任务上进行了测试，该任务需要将高维连续信号与离散 CIF 序列对齐。
• 基准性能：在 MP-20 和 Jarvis-DFT 基准测试中，CrystaLLM-$\pi$ 实现了具有竞争力的结构精度（RMSD 约 0.03–0.04 Å）和匹配率，在特定指标上优于 DiffractGPT 和 Uni3Dar 等基线模型。
• 实验重构：在 Chili-100K 基准测试中，XRD 条件控制模型实现了 49.04% 的结构匹配率（未条件控制基线为 15.89%），并成功重构了每个晶胞多达 83 个原子的结构，而未条件控制模型在超过 40 个原子的系统中则失败。
• 多晶型区分：该模型仅利用成分和 XRD 图谱成功区分了 TiO$_2$ 多晶型（金红石、锐钛矿、板钛矿），甚至重构了训练数据中完全缺失的“板钛矿”相。

意义与主张

本文主张，CrystaLLM-$\pi$ 通过解决连续属性控制与离散结构生成之间的张力，确立了条件自回归晶体生成的新标准。

• 逆向设计能力：该工作表明，自回归生成中的连续控制关键取决于条件信号进入网络的位置。通过将适应局部化到注意力路径中，该框架将生成引导至稀疏的化学空间，而不会侵蚀预训练期间学到的结构先验。
• 效率与可及性：该框架为基于扩散的模型提供了一种轻量级替代方案，显著降低了计算资源需求，同时在发现和重构任务中保持了最先进或接近最先进的性能。
• 通用性：在多样化任务（带隙调整、密度条件控制、SLME 优化和 XRD 结构重构）上的成功表明，该方法可适应材料设计中的各种场景，而无需复杂的架构重新设计。

作者总结道，虽然该框架无法可靠地外推到其训练数据所代表的化学空间之外，但它提供了一个强大且易于使用的工具，用于加速具有特定功能属性的材料发现，以及从实验表征数据中求解结构。