Packora: 生成式分子晶体结构预测的系统设计

# Packora：生成式分子晶体结构预测的系统设计  
来源：https://arxiv.org/html/2608.26962  
###### 摘要  
分子晶体结构预测在制药、农用化学品及有机电子学领域具有重要意义，其中分子构象与堆积的微小差异会显著影响材料性质。我们提出Packora——一种基于流匹配的生成式模型，用于分子晶体结构预测，能够从分子图中联合预测原子坐标与晶格参数。该模型支持多组分与有机金属晶体，并可在单一模型中依据分子构象、立体化学标记和空间群信息的任意子集进行条件生成。受CCDC CSP盲测启发，我们将生成与排序分开评估：生成阶段用于衡量生成器质量，排序阶段则在统一的弛豫与排序流程下评估端到端性能。我们还系统研究了架构、训练、条件机制、推理和缩放等关键设计选择，提出基于可缓存对偶推理、训练目标与数值求解器选择、条件随机失衡及对偶/单一表示平衡缩放的有效设计方案。Packora在结构生成与排序基准测试中均优于基线模型，在六个生成基准中实现最佳预算匹配覆盖率，同时具有更高的实验构型恢复率、更低的实验构型排名及更快的排序收敛速度。  
图1：每个目标晶体30个候选结构的覆盖情况。  
我们在每个目标生成30个候选结构的预算下，报告晶体级别求解率 \(\mathrm{Sol}_{C}\)。括号内数字表示各基准测试的目标数量。对于OXtal，我们直接引用每个目标生成30个候选结构时发表的结果；对于CLARI和Packora，我们为每个目标生成1000个候选结构，并计算5000次30个候选结构自助重抽样的平均值；误差棒表示重抽样间的一个标准差。Packora在所有六个基准测试中均达到最高求解率。  

## 1 引言  
分子晶体广泛应用于制药、农用化学品和有机电子学领域。其性质强烈依赖于分子构象与堆积方式，同一化学体系的不同多晶型可能表现出截然不同的溶解度、稳定性、生物利用度、电荷传输和机械性能。分子晶体结构预测旨在根据分子组分预测可能的晶体结构，以加速药物开发与功能材料设计。分子晶体结构预测通常包含两个互补阶段：结构生成与结构排序。传统工作流通过随机、准随机或进化搜索生成候选堆积结构，随后使用基于密度泛函理论等计算成本高昂的能量计算进行弛豫与排序。FastCSP主要通过采用机器学习原子间势进行结构弛豫与能量计算来加速排序阶段，而初始候选结构仍来自随机结构生成器。  

生成模型通过学习实验观测的晶体结构，并将有限的候选预算集中于合理堆积构型，为结构生成提供了更快捷的方法。尽管近期取得进展，但现有方法在所需分子信息、支持的化学体系等方面存在差异，且评估方案在候选预算与后处理流程上不尽统一，使得生成与排序性能难以直接比较。此外，架构、训练、条件机制、推理和缩放等关键设计选择的影响尚未完全明晰。  

本文提出Packora——一种全原子生成式分子晶体模型，能够从包含原子类型、键类型和形式电荷的分子图中联合预测原子坐标与晶胞参数。该模型支持多组分与有机金属晶体，并显式建模氢原子。单一模型还可利用通常可获取或易于推导的额外结构信息（如分子模板、立体化学标记和空间群信息）进行条件生成，从而适应不同的先验知识水平。  

我们还引入与CCDC CSP盲测对齐的双轨评估体系，将结构生成与结构排序分离。生成轨道衡量固定候选集是否包含实验结构（无需弛豫或基于能量的排序），从而隔离生成器的覆盖能力。排序轨道则对每个生成器的候选结构进行统一的弛豫、筛选、去重与晶格能排序流程，衡量实验构型是否被排在靠前位置。这种分离实现了生成器质量与其生成候选结构在实际晶体结构预测工作流中下游实用性的直接比较。  

最后，我们系统研究了架构、训练策略、条件机制、推理和缩放等关键设计选择。我们直接比较了DiT与配对偏置注意力、Pairformer和Pairixer，并对时间采样、平移增强和损失函数设计等训练选择进行消融实验。这些研究表明：可缓存的对偶推理能显著提升生成质量；精心设计的训练与采样策略至关重要；对偶与单一表示的平衡缩放在模型规模增大时能保持覆盖能力。  

## 2 相关工作  
**基于搜索的分子晶体结构预测**通常包括探索分子构象与晶体堆积，随后进行几何弛豫与能量排序。结构生成与排序长期以来一直是分子晶体结构预测的核心组成部分，第七届CCDC盲测明确将其分为两个独立阶段进行评估。候选结构通常通过在选定空间群中采样晶胞参数、分子位置和取向来生成，柔性分子还需额外采样分子内自由度。生成后的结构经过去重或聚类，并采用分层评估方式，通常先使用力场进行初始弛豫和筛选，再采用色散校正密度泛函理论进行精修与排序。Genarris采用受约束随机生成结合基于多样性的筛选，而GAtor使用第一性原理遗传算法配合分子晶体育种算子与结构生态位划分。FastCSP采用机器学习原子间势加速弛豫与排序阶段，同时保留随机结构生成作为初始候选。  

**生成式分子晶体结构预测**则学习实验晶体结构的分布。OXtal是一种在无晶格分子裁剪数据上训练的全原子扩散模型，将AlphaFold3的Pairformer主干适应于原子级表示。MolCrystalFlow将每个分子表示为刚体，并在分子质心、取向和晶格参数上学习测地线流。PackFlow应用由机器学习原子间势能量与力引导的强化学习后训练，提高物理有效性并将提案集中于低能盆地。CLARI生成显式晶胞并用配对偏置注意力替代三角形层，据报比OXtal快15–30倍。然而，这些方法均无法在单一模型中根据推理时需求选择性地提供或省略分子构象、立体化学标记与空间群信息。  

**生成式晶体结构预测**领域还有一条相关研究线专注于无机材料的晶体生成与结构预测。CDVAE结合变分自编码器与基于扩散的坐标生成，而DiffCSP则在给定组分条件下对分数坐标和晶格进行联合扩散。后续方法采用流匹配构建周期晶体生成模型（如FlowMM和CrystalFlow）、随机插值法（如OMatG）或基于周期变量的贝叶斯流。生成式晶体结构预测还扩展至金属有机框架，其中金属节点与有机连接体可作为模块化构件处理。MOFFlow和MOF-BFN分别使用黎曼流和贝叶斯流生成已知构件的晶格与刚体构型，而MOFFlow-2还额外生成构件并建模连接体扭转角。AtomMOF则通过全原子分辨率直接生成MOF及MOF-吸附质结构，消除了刚体假设。  

## 3 方法  
Packora基于分子规格、化学式单位数 \(Z\) 及可选辅助信息，使用条件坐标-晶格流预测原子坐标与晶胞参数。下文描述所选的Packora配置，其设计选择由第5节的受控研究验证。  

### 3.1 问题表述  
**晶体表示**：我们使用 \(N\) 个原子表示晶体：原子类型 \(\mathbf{A}=(A_{1},\ldots,A_{N})\in\mathcal{A}^{N}\)，笛卡尔坐标 \(\mathbf{X}=(\mathbf{x}_{1},\ldots,\mathbf{x}_{N})^{\top}\in\mathbb{R}^{N\times 3}\)，以及晶胞矩阵 \(\mathbf{L}\in\mathbb{R}^{3\times 3}\)（其行向量为格矢）。沿这些格矢的整数平移定义对应的无限周期晶体。由于同一晶格可由不同基矢和全局旋转表示，我们通过两步规范晶胞：首先对晶体进行Niggli约化以获得约化基矢 \(\widetilde{\mathbf{L}}\)，减少等效晶格基矢间的歧义；其次通过构建旋转不变的Gram矩阵 \(\widetilde{\mathbf{L}}\widetilde{\mathbf{L}}^{\top}\) 并取其唯一下三角Cholesky因子，消除剩余的全局旋转自由度。