Fast Organic Crystal Structure Prediction with Unit Cell Flow Matching¶

会议: NeurIPS 2026
arXiv: 2606.03199
OpenReview: https://openreview.net/forum?id=KD4i8hE0Rc
代码: https://github.com/aspuru-guzik-group/clari
领域: 计算生物
关键词: 晶体结构预测, 流匹配, 晶胞建模, 对偏置注意力, 虚拟筛选

一句话总结¶

Clari 提出了一种直接在单晶胞上联合建模原子坐标与晶格矢量的流匹配模型，摒弃了昂贵的大尺度超晶胞裁剪与三角形更新层，在保持甚至超越生成质量的前提下实现了 15–30 倍采样加速，使数秒级有机晶体结构预测与大规模固态虚拟筛选成为可能。

研究背景与动机¶

有机固体的物理化学性质（如溶解度、稳定性、光电性能等）高度取决于分子在固态下的堆积晶体结构。因此，晶体结构预测（Crystal Structure Prediction, CSP）在药物分子多晶型筛选、农药配方、有机半导体与高能材料等领域中具有决定性作用。然而，由于分子晶体能量超曲面极其复杂且崎岖，传统 CSP 依赖于庞大的计算漏斗：先通过准随机采样或进化算法生成数十万候选结构，再使用高精度但昂贵的密度泛函理论（DFT）进行多轮松弛与能量重排。这种传统范式每预测一个分子往往需要耗费数个 CPU-年，昂贵的算力开销将 CSP 长期局限在极少数重点药物管线中，根本无法支撑现代包含数百万分子的虚拟化合物库的大规模固态高通量筛选。

近期以 OXtal 为代表的生成式晶体模型打破了这一瓶颈，尝试将 CSP 形式化为端到端的数据驱动生成任务。然而，OXtal 继承了 AlphaFold3 类生物大分子模型的一系列沉重架构负担：为了规避显式晶格参数化的几何畸变，它在实空间中构建包含多个对称副本的大尺度超晶胞（bulk crop）进行局部密集建模，并依赖计算复杂度呈立方级增长的三角形更新层（triangle updates）传递成对几何信息。这使得单次分子采样耗时仍在数十秒乃至数分钟量级，对超大化合物库的筛选依然存在显著算力门槛。

面对生成精度与采样吞吐量之间的核心矛盾，本文提出：有机晶体生成的核心瓶颈不在于几何表达能力，而在于对无序堆积冗余的消除与对称性处理。只要将晶格基向量与原子坐标统一表征，并在流匹配框架下引入符合物理统计的先验约束，即可直接在紧凑的单晶胞内完成精确生成。核心 idea：将晶胞晶格矢量显式映射为虚拟原子节点，构建单晶胞流匹配模型 Clari，利用轻量级对偏置注意力彻底取代高开销三角形层，并结合全原子联合建模实现直接能量重排与秒级无冗余晶体预测。

方法详解¶

整体框架¶

Clari 的输入仅需晶胞内分子的二维化学图（原子种类与共价键邻接矩阵，无需任何预先优化的三维构象或 RDKit 标准化处理），输出为单晶胞的三维笛卡尔原子坐标以及描述周期性边界的 $3 \times 3$ 晶格矩阵。

整体推理流程包含三个核心阶段：首先从数据驱动拟合的先验分布中采样初始晶格基底与初始高斯噪声；接着将晶格基向量作为 3 个特殊的虚拟节点拼接至原子坐标序列，在预设时间步长下利用基于 Diffusion Transformer（DiT）的流匹配网络逐步整合几何距离、拓扑距离与自条件端点反馈，预测联合速度场；最后通过 Heun 常微分方程（ODE）求解器积分生成包含全原子（显式氢原子）的三维单晶胞结构，并可选用通用原子模型（UMA）对生成的多候选晶体直接进行无松弛的能量打分重排。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：2D 分子图<br/>(原子特征 F + 键类型 E)"] --> B["先验采样与初始化<br/>数据拟合晶格先验 p0 + 高斯原子噪声"]
    B --> C["虚拟晶格节点拼接与流匹配表征<br/>L/2 与 C 拼接为 (3+n) 统一 Token 流"]
    C --> D["单晶胞流匹配 DiT 主干<br/>对偏置注意力 + 自条件端点递归"]
    D --> E["连续时间速度积分 (Heun ODE)<br/>20-50 步积分生成单晶胞 (L, C)"]
    E --> F["测试时计算与能量重拍<br/>显式全原子免松弛 UMA 能量打分"]
    F --> G["输出：低能量稳定预测晶体结构"]

关键设计¶

1. 虚拟晶格节点与单晶胞联合流匹配：消除超晶胞冗余

针对以往方法（如 OXtal）建模庞大超晶胞所带来的剧烈显存占用与计算开销，Clari 直接在单一晶胞上建立坐标与晶格的统一生成路径。对于包含 $N$ 个原子的晶胞，定义晶格矩阵为 $\bm{L} \in \mathbb{R}^{3 \times 3}$（各行为本原晶格矢量），零中心笛卡尔原子坐标为 $\bm{C} \in \mathbb{R}^{N \times 3}$。Clari 将晶格矢量缩放后视为 3 个特殊的“虚拟原子”节点，将二者行拼接形成增广坐标矩阵： $$\bm{x} = \frac{1}{\sigma}\begin{pmatrix} \frac{1}{2}\bm{L} \\ \bm{C} \end{pmatrix} \in \mathbb{R}^{(3+N) \times 3}$$ 其中 $\sigma$ 为全局方差归一化常数。通过在线性插值路径 $\bm{x}_t = (1-t)\bm{x}_0 + t\bm{x}_1$ 下优化流匹配目标： $$\mathcal{L}_{\mathrm{FM}} = \mathbb{E}_{t, \bm{x}_0, \bm{x}_1} \left[ \|v_\theta(\bm{x}_t, t) - (\bm{x}_1 - \bm{x}_0)\|^2 \right]$$ 网络 $v_\theta$ 能够同时回归晶格变形与原子相对位移，免去了对超晶胞的周期性平移复制，从根本上消除了数据维度的冗余。

2. 对偏置注意力取代三角形层：计算复杂度从三次降至二次

针对 AlphaFold3 与 OXtal 中三角形更新层（Triangle Multiplication / Update）带来的 $O(N^3)$ 巨大显存与计算瓶颈，Clari 采用了基于 Diffusion Transformer（DiT）的对偏置注意力（Pair-Bias Attention）架构。模型将三维成对信息解耦为注意力矩阵的加性偏置，成对特征 $\bm{z}_{ij}$ 由共价键类型、拓扑图最短路径距离、笛卡尔欧氏距离与基于晶格矢量的周期性最小镜像距离（$\min_{\bm{z} \in \{-1,0,1\}^3} \|\bm{p}_i - \bm{p}_j - \bm{L}^\top \bm{z}\|$）多维离散编码拼接而成。在自注意力计算中，成对特征经线性变换后直接注入注意力 Logits，无需在每层网络内部运行高开销的节点三元组闭环更新。搭配 SwiGLU 激活、QK-Norm 与 AdaLN-Zero 全局调制层，Clari 在参数量扩展至 1.7 亿（Clari-L）并处理多达 512 原子的复杂晶胞时，依旧保持了高效线性的吞吐扩展性。

3. 数据拟合晶格先验与周期性几何辅助损失：抑制晶胞畸变与原子碰撞

在连续生成模型中，标准的各向同性高斯先验 $\bm{x}_0 \sim \mathcal{N}(\bm{0}, \bm{I})$ 在晶格矩阵空间会导致严重的几何病态，例如极小晶胞体积、各向异性极端的长宽比或近乎平行的基底夹角。为此，Clari 显式将晶格参数解构为原子数密度 $\rho = N/|\det \bm{L}|$、三个晶胞夹角 $(\alpha, \beta, \gamma)$ 以及归一化轴长 $(a, b, c) = (\|\bm{l}_1\|, \|\bm{l}_2\|, \|\bm{l}_3\|) / V^{1/3}$，并分别在剑桥结构数据库（CSD）训练集上拟合一维与多元高斯分布。采样时，先验分布先生成物理可行的基底形态，再经随机旋转与符号置换形成 $p_0$。此外，为了防止生成中间态的体相重叠与穿模，Clari 在一步外推端点 $\hat{\bm{x}}_1$ 上引入了晶格体积相对误差损失 $\mathcal{L}_{\mathrm{vol}}$ 与周期性成对防碰撞损失 $\mathcal{L}_{\mathrm{pair}}$： $$\mathcal{L}_{\mathrm{pair}} = \sum_{(i,j) \in \Lambda} \left[ |\hat{d}_{ij} - d_{ij}| + 5 \cdot \max(0, \alpha_{ij} - \hat{d}_{ij}) \right]$$ 该设计使模型在去噪轨迹初期就能快速拉开原子间距，彻底抑制非物理的空间位阻冲突。

4. 全原子建模与无松弛测试时能量重排：端到端推理扩展

主流晶体与分子生成模型通常只预测重原子骨架，下游应用需要依赖耗时且复杂的隐式加氢（hydrogen decoration）和力场松弛。Clari 在训练阶段完整保留全部显式氢原子。由于包含精确定位的氢原子坐标与方向，Clari 生成的原始晶胞可直接输入通用原子预训练电位模型（如 UMA, Universal Model for Atoms）进行单点单能计算。在测试阶段，通过轻量级 Best-of-$N$ 策略（例如采样 $N=150$ 篇候选并依据 UMA 能量保留前 30 个），模型在完全跳过昂贵 DFT 局部几何优化的前提下，以仅增加数秒的极小开销显著提升了盲测匹配成功率。

损失函数 / 训练策略¶

Clari 的最终训练损失由基础流匹配速度损失与两个几何辅助目标直接相加构成： $$\mathcal{L}_{\mathrm{Clari}} = \mathcal{L}_{\mathrm{FM}}^{\bm{L}} + \mathcal{L}_{\mathrm{FM}}^{\bm{C}} + \mathcal{L}_{\mathrm{vol}} + \mathcal{L}_{\mathrm{pair}}$$ 时间步分布采用有利于微观构型优化的 $\text{Beta}(1.8, 1)$ 分布。训练过程中引入自条件策略（Self-Conditioning），以 50% 概率将模型前一步估算的外推终点 $\hat{\bm{x}}_1$ 递归作为输入特征。模型使用 Muon 优化器更新 Transformer 核心权重矩阵，并使用 Adam 优化其余参数，在 8 张 H100 GPU 上训练 15 万步即可完全收敛。

实验关键数据¶

主实验¶

在代表有机晶体结构预测最高标准基准的 OXtal 测试集（刚性分子集 Rigid、柔性分子集 Flexible）以及历届国际晶体盲测（CSP Blind Tests 5/6/7）上，评估 COMPACK 15 分子簇匹配率（$\text{Sol}@k$，要求至少匹配 8/15 分子且 $\text{RMSD}_{15} < 2$ Å）。当 $k < n_s$ 时，使用 UMA 能量筛选前 $k$ 个低能候选结构。

数据集	指标	Clari-L (150→30)	OXtal (30)	相对提升 / 差距
Rigid (50)	Sol@30 (↑)	0.772	0.300	+0.472 (+157.3%)
Flexible (50)	Sol@30 (↑)	0.346	0.220	+0.126 (+57.3%)
CSP5 (6)	Sol@30 (↑)	0.789	0.167	+0.622 (+372.5%)
CSP6 (5)	Sol@30 (↑)	0.480	0.200	+0.280 (+140.0%)
CSP7 (8)	Sol@30 (↑)	0.263	0.125	+0.138 (+110.4%)
CSD Teaching (773)	Sol@30 (↑)	0.484	–	新基准强基线

在计算效率方面，单分子采样 150 个晶体候选并在 GPU 上完成 UMA 能量打分选出前 30 个，Clari-L 端到端仅需 6.0 秒（纯生成仅 2.2 秒），相比 OXtal 在相同盲测任务上取得了 15–30× 的纯采样加速 以及 5–8× 的端到端能量重排加速。

消融实验¶

在 CSD 验证集上逐步验证核心架构组件、先验分布与自条件设计的有效性（采用 5000 次 Bootstrap 重抽样）：

配置	碰撞率 Clash Rate (%, ↓)	PoseBusters 构象合格率 (%, ↑)	相对体积误差 Vol. Error (%, ↓)	EMD PDD 距离分布误差 (↓)	说明
A (Base DiT)	25.00 ± 0.50	78.14 ± 0.42	2.07 ± 0.05	11.01 ± 0.05	平均池化预测晶格基础 DiT
B (+ 虚拟晶格节点)	23.25 ± 0.49	79.04 ± 0.41	2.22 ± 0.05	11.09 ± 0.05	架构简化为单流 Token，训练更稳定
C (+ 辅助几何损失)	20.33 ± 0.47	80.11 ± 0.40	1.79 ± 0.04	10.37 ± 0.04	引入体积与成对防碰撞损失
D (+ 自条件自递归)	9.79 ± 0.35	85.23 ± 0.35	1.55 ± 0.04	9.66 ± 0.03	循环输入终点估算，显著降低错误率
E (D + 标准正态先验)	9.32 ± 0.33	84.89 ± 0.36	2.88 ± 0.06	10.71 ± 0.05	替换为各向同性高斯，几何分布显著恶化
Clari-M (最终中型模型)	9.56 ± 0.35	87.34 ± 0.34	1.59 ± 0.04	9.56 ± 0.03	改用 Beta(1.8, 1) 时间步调度
Clari-L (最终大型模型)	7.69 ± 0.32	85.89 ± 0.39	1.50 ± 0.04	9.28 ± 0.03	参数扩展至 173M，全指标全面领先

关键发现¶

自条件机制（Self-Conditioning）贡献最为显著：引入自条件后，晶体碰撞率从 20.33% 腰斩至 9.79%，PoseBusters 构象通过率提升超过 5 个百分点，表明循环回传预测终点有助于模型在离散求解早期建立全局物理约束。
数据拟合晶格先验不可或缺：若将拟合先验退回为各向同性标准高斯先验（模型 E），相对体积误差剧增 85%（从 1.55% 恶化至 2.88%），周期性距离误差明显回升，证实了非晶态各向同性噪声无法为周期性晶胞提供合理的初始拓扑支撑。
化学广泛性与非 RDKit 分子泛化：Clari 能够直接泛化并成功预测富勒烯（如 $\text{C}_{60} \cdot \text{Co}_2(\text{CO})_8$）、多面体低聚倍半硅氧烷笼状分子（POSS cage）以及多核过渡金属配合物（如具有 $\text{Pt}_4(\mu_3\text{-I})_4$ 核心的四聚体）。这类复杂体系在传统依赖 RDKit sanitization 或复杂三维初始构象的方法中往往无法被有效处理。

亮点与洞察¶

物理先验驱动的虚拟晶格节点表示：将周期性晶格基向量视同特殊原子节点并合入注意力流，以极度简化的单晶胞数据结构彻底摆脱了复杂的周期性超晶胞展开，架构极具优雅性与通用性。
注意力与成对偏置的高效平衡：证明了在周期性固态晶体生成任务中，并不需要照搬生物大分子领域昂贵的 $O(N^3)$ 三角形更新层；精心编码的周期性最小镜像成对偏置完全足以引导正确的空间微观堆积。
显式氢原子直接对接测试时计算（TTC）：端到端预测显式氢原子省去了传统两阶段加氢所引入的几何形变，使生成结构与预训练原子能量模型（UMA）能够零缝隙串联，形成极速的“生成-重排”闭环。

局限与展望¶

输入手性与立体化学敏感性欠缺：当前模型输入为二维分子图，对立体异构体和特定手性中心不敏感，在药物化学关键立体异构晶型的精确区分上仍需进一步拓展。
依赖晶胞内分子化学计量数 $Z$：推理时仍需预先指定晶胞内分子的重复单元数 $Z$。虽然常见有机晶体通常集中在 $Z \in \{1, 2, 4\}$ 且容易枚举，但端到端联合预测 $Z$ 仍是值得探索的方向。
非物理空洞与偶发穿模：在超大或高柔性分子中，模型偶尔仍会生成局部密度异常或含有空间空洞的亚稳态结构，未来可探索引入粉末 X 射线衍射（PXRD）或空间群（Space Group）作为条件引导（Steering）来进一步约束生成过程。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙提出虚拟晶格节点单晶胞流匹配与无三角形对偏置架构，开辟了秒级 CSP 新路径。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖消融、OXtal 基准、历届国际盲测及新提出的 CSD Teaching 子集，评估全面严谨。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密，物理图景与机器学习公式表述清晰，动机切入极具说服力。
价值: ⭐⭐⭐⭐⭐ 将耗时数 CPU-年的有机晶体预测降低到几秒，直接打开了百万级分子虚拟筛选在材料与制药工业的大门。