AtomMOF: All-Atom Flow Matching for MOF-Adsorbate Structure Prediction¶
会议: NeurIPS 2026
arXiv: 2602.07351
代码: https://github.com/nayoung10/AtomMOF
领域: 物理/科学计算
关键词: 金属有机框架(MOF)、全原子流匹配、晶体结构预测、主客体吸附复合物、Feynman-Kac 引导
一句话总结¶
本文提出首个面向金属有机框架(MOF)及主客体吸附复合物的全原子变分流匹配生成模型 AtomMOF,直接从构建单元与吸附物的 2D 分子图无约束预测全原子 3D 晶体坐标与晶胞参数,结合机器学习原子间势(MLIP)驱动的 Feynman-Kac 采样引导,在 BW 基准上将结构匹配率提升 35.00%、RMSD 降低 32.68%,并在 ODAC25 吸附构型采样中比传统巨正则蒙特卡洛(GCMC)快数个数量级。
研究背景与动机¶
金属有机框架(MOF)是一类由金属节点与有机配体组装而成的多孔晶体材料,因比表面积巨大且化学性质高度可调,在气体分离、催化以及直接空气碳捕集(DAC,即从含水蒸气竞争的环境空气中选择性吸附极低浓度的 \(\text{CO}_2\))等前沿科学问题中占据核心地位。要评估一种候选 MOF 的碳捕集性能,关键在于精确预测框架主体与吸附物客体分子之间的三维相互作用构型。传统的计算材料学工作流依赖巨正则蒙特卡洛(GCMC)配合密度泛函理论(DFT)进行力场采样与能量弛豫,但这类模拟在单胞往往包含数百个原子的 MOF 体系上计算开销极其高昂;更致命的是,传统模拟必须以已知的主体 MOF 三维晶体结构为前提,无法直接对由海量构建单元组合而成的假想(hypothetical)MOF 空间进行从头筛选。
近年来深度生成模型开始被引入 MOF 结构预测,但现有方法(如 MOFFlow、MOF-BFN 以及引入可旋转键扭转角的 MOFFlow-2)普遍依赖粗粒度(coarse-grained)表征或刚体假设,即强行固定由化学信息学工具生成的局部键长与键角,仅预测构建单元整体的刚体旋转平移或局部二面角。这种刚性约束无法刻画真实晶体中由局部化学环境与主客体相互作用引发的连续键长形变、键角弯曲与配位几何畸变,导致在金属节点附近频繁出现空间位阻冲突(steric clashes)。此外,现有生成模型全部孤立地预测纯主体框架结构,完全忽略了决定材料实际服役性能的吸附物分子(如 \(\text{CO}_2\)、\(\text{H}_2\text{O}\))与框架的联合构型建模。
为了同时打破刚体几何近似与孤立主体建模的双重局限,本文跳出刚性片段拼接范式,直接在全原子无约束空间中统一建模框架与客体分子的协同几何演化。核心 idea:提出基于全原子扩散 Transformer(DiT)的变分流匹配框架 AtomMOF,在无刚体约束下端到端地将构建单元与吸附物的 2D 分子图映射为 3D 全原子平衡坐标与晶胞参数,并引入由机器学习原子间势(MLIP)驱动的 Feynman-Kac 粒子引导机制以保障全原子自由度下的几何有效性与热力学稳定性。
方法详解¶
整体框架¶
AtomMOF 的目标是从给定的构建单元集合 \(\mathcal{B} = \mathcal{B}_{\text{MOF}} \cup \mathcal{B}_{\text{ads}}\)(包含 \(M\) 个有机配体与金属节点的分子图 \(\{\mathcal{G}_i\}_{i=1}^M\),以及 \(P\) 个吸附物分子的分子图 \(\{\mathcal{H}_j\}_{j=1}^P\),纯框架预测时 \(P=0\))出发,联合预测包含 \(N\) 个原子的晶胞平衡结构 \(\mathcal{S} = (\mathbf{A}, \mathbf{X}, \bm{\ell})\),其中 \(\mathbf{A} \in \mathcal{A}^N\) 为原子元素类型,\(\mathbf{X} \in \mathbb{R}^{N \times 3}\) 为全原子笛卡尔坐标,\(\bm{\ell} = (\bm{d}, \bm{\phi}) \in \mathbb{R}_+^3 \times [0^\circ, 180^\circ]^3\) 为晶胞三边长 \((a,b,c)\) 与三夹角 \((\alpha,\beta,\gamma)\)。模型整体由三个核心部分协同工作:首先在无约束坐标与晶胞空间建立基于拉普拉斯后验的变分流匹配(VFM)概率路径;接着通过一个三阶段全原子 DiT 去噪网络(由构建单元特征编码器、成对偏置交互主干网络与双头解码器组成)从含噪状态预测干净的全原子坐标与晶胞参数;最后在推理采样后期将常微分方程(ODE)转化为随机微分方程(SDE),利用预训练机器学习原子间势(MLIP)对粒子群执行 Feynman-Kac 重采样引导,过滤空间碰撞并向低能稳定态偏置。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入:构建单元与吸附物2D分子图 B<br/>及先验采样状态 (X0, l0)"] --> B["全原子变分流匹配框架<br/>零质心高斯坐标先验 + 对数正态/均匀晶胞先验<br/>沿线性插值路径构造含噪状态 (Xt, lt)"]
B --> C["成对几何偏置的全原子DiT架构<br/>Encoder编码局部RDKit坐标/键级/同块掩码<br/>24层DiT Trunk融合单原子与成对注意力偏置<br/>Decoder双头并行输出预测干净态 (X̂1, l̂1)"]
C --> D["MLIP驱动的Feynman-Kac推理引导<br/>后期(t≥0.8)将流ODE转为含得分项的SDE<br/>按MLIP能量势函数对K个粒子重采样"]
D --> E["输出:物理有效且低能量的<br/>MOF-吸附物3D全原子晶体结构 S"]
关键设计¶
1. 全原子变分流匹配框架:在无刚体约束下联合生成全原子坐标与晶胞参数
针对以往模型固定键长与键角导致无法刻画框架柔性形变及主客体微扰的痛点,AtomMOF 彻底放弃刚体与扭转角参数化,直接在全原子笛卡尔坐标 \(\mathbf{X}\) 与晶胞参数 \(\bm{\ell}\) 构成的联合连续空间 \(\mathbf{z} = (\mathbf{X}, \bm{\ell}) \in \mathbb{R}^{N \times 3} \times \mathbb{R}^6\) 上采用变分流匹配(Variational Flow Matching, VFM)。在先验分布 \(p_0(\mathbf{z}_0)\) 的设计上,原子坐标采用零质心高斯分布 \(\mathbf{X}_0 = \mathbf{P}\bm{\varepsilon}\)(其中 \(\bm{\varepsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\),投影矩阵 \(\mathbf{P} = \mathbf{I} - \frac{1}{N}\mathbf{1}\mathbf{1}^\top\) 消除平移自由度以保证质心位于原点),晶胞边长 \(\bm{d}\) 从训练集最大似然估计拟合的对数正态分布 \(\operatorname{LogNormal}(\bm{\mu}, \bm{\sigma})\) 中独立采样,晶胞夹角 \(\bm{\phi}\) 则从 \([60^\circ, 120^\circ]\) 均匀采样。给定真实平衡结构 \(\mathbf{z}_1 \sim q\),时间步 \(t \sim \mathcal{U}[0, 1]\) 下的条件插值路径定义为线性插值 \(\mathbf{z}_t = (1-t)\mathbf{z}_0 + t\mathbf{z}_1\)。
不同于标准流匹配直接回归速度场并用 \(L^2\) 损失训练,VFM 通过变分后验分布 \(q_t^\theta(\mathbf{z}_1 \mid \mathbf{z}_t)\) 逼近真实干净数据后验 \(p_{1|t}(\mathbf{z}_1 \mid \mathbf{z}_t)\)。本文将该变分后验参数化为以神经网络输出 \(\bm{\mu}_t^\theta(\mathbf{X}_t, \bm{\ell}_t \mid \mathcal{B}) = (\hat{\mathbf{X}}_1, \hat{\bm{\ell}}_1)\) 为均值的拉普拉斯分布(Laplace distribution),从而将负对数似然最小化自然转化为对坐标异常值更鲁棒、收敛更快的 \(L^1\) 重建目标。推理时,模型直接由预测的干净态导出速度场 \(\mathbf{v}_t(\mathbf{z}_t) = (\bm{\mu}_t^\theta(\mathbf{z}_t) - \mathbf{z}_t)/(1-t)\),仅需通过 ODE 求解器积分 10 步即可高效生成三维晶体结构。
2. 成对几何偏置的全原子DiT架构:融合构建单元局部化学先验与全局原子交互
放开全原子自由度后,模型必须在缺乏刚性模板约束的情况下准确维持共价键拓扑并捕捉长程周期性相互作用。为此,去噪网络 \(\bm{\mu}_t^\theta(\mathbf{X}_t, \bm{\ell}_t \mid \mathcal{B})\) 设计了“编码器(Encoder)- 主干(Trunk)- 解码器(Decoder)”三阶段全原子扩散 Transformer(DiT)架构。在 Encoder 阶段,模型首先提取每个构建单元的局部化学先验:对有机配体与吸附物使用 RDKit 生成局部参考坐标 \(\mathbf{X}_{\text{local}}\),对金属节点则通过 Kabsch 算法对齐训练集同类节点模板(BW 数据集 9 种节点取平均几何,ODAC25 数据集 887 种节点按 SMILES 索引随机采样参考结构),将局部参考坐标、形式电荷与原子类型拼接后线性映射为初始单原子特征 \(\mathbf{C} \in \mathbb{R}^{N \times D_Q}\)。同时,Encoder 计算局部参考位移向量 \(\mathbf{D}_{ij} = \mathbf{X}_{\text{local}, i} - \mathbf{X}_{\text{local}, j}\) 及其欧氏距离 \(\|\mathbf{D}_{ij}\|\),与化学键类型 \(\mathcal{B}_{\text{bond}}\) 一起嵌入,并乘上指示原子 \(i, j\) 是否属于同一构建单元的二值分块邻接掩码 \(\bm{v}_{ij}\),从而构造仅保留块内局部几何关系的成对表征 \(\mathbf{P} \in \mathbb{R}^{N \times N \times D_P}\)。随后,单原子特征加上当前含噪全局坐标 \(\mathbf{X}_t\) 与含噪晶胞 \(\bm{\ell}_t\) 的线性投影得到 \(\mathbf{Q}\),成对表征融入单原子 MLP 上下文后作为自注意力 logits 的加性偏置(additive bias),送入 3 层浅层 \(\text{DiT}_{\text{encoder}}\)。
在 Trunk 阶段,编码器输出的 \(\mathbf{Q}, \mathbf{P}\) 被线性投影至更高维的潜空间表征 \(\mathbf{S}\) 与 \(\mathbf{Z}\),并在进入主干变换前通过外积和形式将单原子上下文动态注入成对偏置 \(\mathbf{Z}_{ij} \leftarrow \mathbf{Z}_{ij} + \operatorname{Linear}(\mathbf{S}_i) + \operatorname{Linear}(\mathbf{S}_j)\)。随后,24 层深度 \(\text{DiT}_{\text{trunk}}\) 模块利用时间步正弦嵌入的自适应层归一化(AdaLN)调制每一层的自注意力与前馈网络,并在自注意力计算中持续以 \(\mathbf{Z}\) 作为结构偏置:
这种设计既利用分块成对偏置锚定了构建单元内部的局部化学键几何,又允许全原子自注意力在无刚体束缚下自由建模跨构建单元及框架-吸附物之间的非共价作用。Trunk 输出经全局残差连接更新 \(\mathbf{Q} \leftarrow \mathbf{Q} + \operatorname{Linear}(\mathbf{S}')\) 后送入 3 层 \(\text{DiT}_{\text{decoder}}\),最终由两个并行线性头分别直接投影出单原子去噪坐标 \(\hat{\mathbf{X}}_1 \in \mathbb{R}^{N \times 3}\),以及经全局平均池化(Global Average Pooling)后投影出的晶胞参数 \(\hat{\bm{\ell}}_1 \in \mathbb{R}^6\)。
3. MLIP驱动的Feynman-Kac推理引导:无需重训练即可提升结构有效性与能量稳定性
由于全原子生成完全解除了键长与刚体约束,纯生成模型在复杂多孔体系采样时仍可能产生原子重叠或高能亚稳态构型。为了在推理阶段注入物理力场先验,本文引入由预训练机器学习原子间势(MLIP,具体采用 eSEN-ODAC25-Full)引导的 Feynman-Kac(FK)粒子滤波机制,从玻尔兹曼指数倾斜的目标分布 \(p_{\text{target}}(\mathbf{X}, \bm{\ell} \mid \mathcal{B}) \propto p_\theta(\mathbf{X}, \bm{\ell} \mid \mathcal{B}) \exp(-\lambda E(\mathcal{S}))\) 中进行采样(其中 \(E(\cdot)\) 为 MLIP 评估的势能或原子重叠惩罚,\(\lambda\) 为引导强度)。为了给粒子群赋予随机探索能力以配合重采样,模型将确定性的坐标流 ODE 转化为等价的随机微分方程(SDE):
其中扩散系数取 \(g(t) = \sqrt{(1-t)/2}\),得分函数 \(s(\mathbf{X}_t, t)\) 直接由零质心高斯先验下的速度场 \(\mathbf{v}(\mathbf{X}_t, t)\) 解析给出,无需额外训练 score 网络(晶胞参数 \(\bm{\ell}_t\) 仍保持确定性步进)。在采样过程中,算法并行维护 \(K\) 个粒子轨迹;每当时间步进入去噪后期(\(t \ge t_{\text{start}} = 0.8\))且步数满足重采样间隔 \(\Delta\tau\) 时,利用去噪器一步预测的干净态 \(\hat{\mathcal{S}}_1^{(k)}\) 计算即时势能 \(\mathcal{E}_t^{(k)} = E(\hat{\mathcal{S}}_1^{(k)})\),按即时势函数权重 \(w_k = \exp(-\lambda \mathcal{E}_t^{(k)})\) 归一化后进行多项式分布重采样(Multinomial Resampling),复制低能高合理性粒子并淘汰高能冲突粒子。之所以将引导推迟到 \(t_{\text{start}} = 0.8\) 之后才开启,是因为预训练 MLIP 仅在干净晶体结构上训练,在早期高噪声阶段(\(t\) 较小时)一步预测的 \(\hat{\mathcal{S}}_1\) 误差较大、能量估计不可靠,而在轨迹后期的精细几何调整阶段引入引导则能精准消除空间位阻并大幅降低形成能误差。
损失函数 / 训练策略¶
模型的训练目标为式 (4) 所示的拉普拉斯变分流匹配 \(L^1\) 重建损失,将晶胞参数拆分为边长 \(\bm{d}\) 与夹角 \(\bm{\phi}\) 分别赋权:
其中权重设为 \(\lambda_{\text{coord}} = 1.0\)、\(\lambda_{\text{lengths}} = 1.0\)、\(\lambda_{\text{angle}} = 0.01\)。为了应对不同 MOF 单胞原子数 \(N\) 的剧烈变化以及自注意力 \(O(N^2)\) 的显存开销,训练采用基于原子数平方上限(\(N_{\max}^2\))的动态批处理(dynamic batching)策略,并配合三阶段课程学习(Curriculum Learning):在预训练阶段(500 epochs)、微调阶段 1(100 epochs)与微调阶段 2(500 epochs)中,BW 数据集的最大原子数上限依次从 200 提升至 300 再到无限制(ODAC25 数据集依次从 300 提升至 400 再到无限制)。模型共训练三个尺寸变体:AtomMOF-S(38.4M 参数,Trunk 维度 256/128)、AtomMOF-M(129M 参数,Trunk 维度 512/256)和 AtomMOF-L(491M 参数,Trunk 维度 1024/512),均在 8 张 NVIDIA B200 GPU 上使用 AdamW 优化器(学习率 \(1 \times 10^{-4}\),零权重衰减,前 10,000 步线性 warmup)完成训练。
实验关键数据¶
主实验¶
实验分为两部分:第一部分在 BW 数据集上评估纯 MOF 框架从头结构预测精度(表 1)及 Zeo++ 孔道物理性质保持度(表 2);第二部分在 ODAC25 数据集(153,249 个训练样本、3,582 个测试样本,含 \(\text{CO}_2\)、\(\text{H}_2\text{O}\)、\(\text{N}_2\)、\(\text{O}_2\) 等吸附物)上评估 MOF-吸附物复合物的构型覆盖率与吸附能优化成功率(表 3)。
表 1:BW 数据集纯 MOF 结构预测精度对比(10 步 ODE 采样,匹配容差 \(\text{ltol}=0.3, \text{angle\_tol}=10^\circ\))
| 方法 | 采样数 (# Samples) | \(\text{stol}=0.3\) MR (%) \(\uparrow\) | \(\text{stol}=0.3\) RMSD (Å) \(\downarrow\) | \(\text{stol}=0.5\) MR (%) \(\uparrow\) | \(\text{stol}=0.5\) RMSD (Å) \(\downarrow\) |
|---|---|---|---|---|---|
| RS / EA (CrySPY) | 20 | 0.00 | – | 0.00 | – |
| DiffCSP | 1 / 5 | 0.01 / 0.08 | 0.1554 / 0.1299 | 0.23 / 0.87 | 0.3896 / 0.3982 |
| MOFFlow | 1 / 5 | 5.28 / 8.68 | 0.2036 / 0.2039 | 21.93 / 32.71 | 0.3329 / 0.3290 |
| MOF-BFN | 1 / 5 | 6.96 / 12.11 | 0.1979 / 0.1950 | 27.66 / 40.15 | 0.3234 / 0.3131 |
| MOFFlow-2(之前 SOTA) | 1 / 5 | 8.20 / 15.98 | 0.1894 / 0.1842 | 28.71 / 43.95 | 0.3094 / 0.2925 |
| AtomMOF-M(本文) | 1 / 5 | 10.05 / 19.85 | 0.1328 / 0.1249 | 28.63 / 46.08 | 0.2628 / 0.2396 |
| AtomMOF-L(本文) | 1 / 5 | 11.07 / 21.17 | 0.1275 / 0.1199 | 29.96 / 47.47 | 0.2558 / 0.2319 |
表 2:BW 数据集预测结构的 Zeo++ 几何与孔道性质 RMSE 对比(AtomMOF 仅用 10 步 ODE,MOFFlow-2 用默认 50 步)
| 模型 | VSA (\(\text{m}^2/\text{cm}^3\)) \(\downarrow\) | GSA (\(\text{m}^2/\text{g}\)) \(\downarrow\) | AV (\(\text{Å}^3\)) \(\downarrow\) | UCV (\(\text{Å}^3\)) \(\downarrow\) | VF \(\downarrow\) | PLD (Å) \(\downarrow\) | LCD (Å) \(\downarrow\) | DST (\(\text{g}/\text{cm}^3\)) \(\downarrow\) |
|---|---|---|---|---|---|---|---|---|
| MOFFlow-2 | 312.8 | 443.6 | 485.0 | 427.1 | 0.0445 | 1.3751 | 1.5361 | 0.0312 |
| AtomMOF-M | 226.9 | 282.2 | 210.0 | 153.8 | 0.0257 | 1.0147 | 0.9616 | 0.0177 |
| AtomMOF-L | 215.2 | 275.4 | 241.4 | 192.4 | 0.0247 | 1.0131 | 0.9435 | 0.0165 |
表 3:ODAC25 数据集 MOF-吸附物结构预测与 MLIP 弛豫结果(AtomMOF-L)
| 评估维度 / 指标 | 对比基线 / 1 Sample | 10 Samples / AtomMOF-L | 核心对比结论 |
|---|---|---|---|
| 构型覆盖率 COV (%) vs 耗时 | GCMC: 21.06% (\(\sim 10^4\text{ s}\)) | AtomMOF-L: 23.11% (\(<1\text{ s}\)) | 覆盖率更高且推理速度比 GCMC 快约 4 个数量级 |
| 吸附能弛豫成功率 SR (%) \(\uparrow\) | 12.30% (1 Sample) | 18.59% (10 Samples) | 300 步内力收敛至 \(<0.05\text{ eV/Å}\) 且能量差 \(\le 0.1\text{ eV}\) |
| 发现更低吸附能比例 \(\varepsilon_{\min} < 0\) (%) \(\uparrow\) | 6.30% (1 Sample) | 14.20% (10 Samples) | 超越参考数据集中已知最低能量构型的系统占比 |
| 平均最小吸附能差 \(\varepsilon_{\min}\) (eV) \(\downarrow\) | +0.0841 (1 Sample) | -0.0462 (10 Samples) | 负值表明 10 次采样最优构型的平均吸附能比参考真值更稳定 |
消融实验¶
本文围绕 Feynman-Kac(FK)引导机制及各超参数(粒子数 \(K\)、引导强度 \(\lambda\)、重采样间隔 \(\Delta\tau\)、开启时间 \(t_{\text{start}}\) 与势函数形式 \(G_t\))在 AtomMOF-M(200 步 SDE 积分,1000 个随机测试样本)上进行了详尽消融(表 4)。
表 4:MLIP 驱动的 Feynman-Kac 引导机制与超参数敏感性消融实验(BW 数据集)
| 实验组别 | 配置参数 | MOFChecker 有效率 (%) \(\uparrow\) | 结构有效率 Struct. val. (%) \(\uparrow\) | 形成能误差 \(\Delta E_f\) (eV/atom) \(\downarrow\) | 说明 |
|---|---|---|---|---|---|
| 基准对比 | No steering(无引导) | 34.16 | 81.70 | 0.12700 | 纯流匹配生成基线 |
| 完整引导(表 4 最优) | FK Steering (\(K=16, \lambda=2.0, t_{\text{start}}=0.8\)) | 37.87 (+10.86%) | 91.70 (+12.24%) | 0.02005 (-84.19%) | 形成能误差大幅降低 84.19% |
| 粒子数 \(K\) 消融 | \(K=4\) / \(K=8\) / \(K=16\) | 35.79 / 35.92 / 37.87 | 87.00 / 88.90 / 91.70 | 0.07235 / 0.04828 / 0.02005 | 粒子数越多,探索与低能筛选能力越强 |
| 引导强度 \(\lambda\) 消融 (\(K=8\)) | \(\lambda=1.0\) / \(2.0\) / \(5.0\) / \(10.0\) | 36.61 / 35.92 / 36.55 / 36.68 | 88.10 / 88.90 / 90.10 / 90.30 | 0.05987 / 0.04828 / 0.03508 / 0.03148 | 增大 \(\lambda\) 显著压低形成能误差并减少原子碰撞 |
| 开启时间 \(t_{\text{start}}\) 消融 (\(K=8\)) | \(t_{\text{start}}=0.0\) / \(0.2\) / \(0.4\) / \(0.6\) / \(0.8\) | 34.58 / 35.23 / 35.13 / 35.03 / 35.92 | 82.50 / 85.90 / 86.50 / 88.10 / 88.90 | 0.10370 / 0.07469 / 0.06201 / 0.05314 / 0.04828 | 越晚开启引导效果越好,因早期含噪结构上 MLIP 不准 |
| 势函数类型消融 (\(K=8\)) | Immediate / Difference / Max / Sum | 35.92 / 35.09 / 35.71 / 35.39 | 88.90 / 86.10 / 89.00 / 89.00 | 0.04828 / 0.08443 / 0.04828 / 0.04838 | Immediate 势函数 \(\exp(-\lambda \mathcal{E}_t)\) 综合表现最优 |
关键发现¶
- 全原子无约束建模在严格几何容差下优势显著:在更严格的站点匹配容差 \(\text{stol}=0.3\)(单样本)下,AtomMOF-L 将匹配率从 MOFFlow-2 的 8.20% 提升至 11.07%(相对提升 35.00%),RMSD 从 0.1894 Å 降至 0.1275 Å(相对降低 32.68%);同时在比表面积(VSA/GSA)、孔径(PLD/LCD)与晶胞体积(UCV)等全部 8 项 Zeo++ 孔道物性指标上取得最低误差,证明放开键长与键角约束能有效避免金属节点处的空间碰撞。
- 多孔晶体生成首次展现出清晰的 Scaling Law:当模型参数量从 38.4M(AtomMOF-S)扩展至 129M(AtomMOF-M)再到 491M(AtomMOF-L)时,匹配率与 RMSD 均随参数量的对数呈高度线性改善趋势(\(\text{stol}=0.3\) 下 Pearson 相关系数达 \(\pm 0.93\),Spearman 秩相关系数达 \(\pm 1.00\))。在推理效率上,AtomMOF 仅需 3 步 ODE 积分即可超越 MOFFlow-2 在 50 步下的峰值匹配率(\(\sim 7.6\%\)),并在 10 步左右趋于饱和。
- 引导时机 \(t_{\text{start}}\) 是 MLIP 测试时引导成败的关键:消融实验表明,若从 \(t_{\text{start}}=0.0\) 全程开启 FK 引导,形成能误差高达 0.10370 eV/atom,结构有效率仅 82.50%;而延迟至 \(t_{\text{start}}=0.8\) 开启时,形成能误差骤降至 0.04828 eV/atom(\(K=16\) 时进一步降至 0.02005 eV/atom)。这是因为 MLIP 仅在平衡态附近具备可靠的势能面估计,在去噪末期介入才能发挥正向修正作用。
- 生成模型可作为局部能量极小值搜索的高效初始化器:在 ODAC25 主客体预测中,AtomMOF-L 采样 10 个候选构型并用 MLIP 弛豫后,在 14.20% 的测试体系中找到了比数据集参考轨迹最低能态还要稳定的吸附构型(平均 \(\varepsilon_{\min} = -0.0462\text{ eV}\)),展示了替代耗时 GCMC 预采样的潜力。
亮点与洞察¶
- 用局部化学先验偏置替代硬性刚体约束:本文没有将 RDKit 生成的构建单元局部坐标作为不可变的刚体框架,而是将其编码为 DiT 注意力层中的块内成对距离与键级加性偏置(pairwise attention bias)。这种软约束设计既让模型在初始化时获知分子图拓扑与近似几何,又允许全原子坐标在深度 Transformer 交互中根据晶体环境发生连续弹性形变,优雅地解决了金属配位节点处的拼接冲突。
- 零重训练成本的 ODE-to-SDE 物理势场测试时扩展(Test-Time Scaling):本文最令人眼前一亮的技巧是在推理最后 20% 的时间步(\(t \ge 0.8\))利用零质心高斯先验的解析关系将确定性流 ODE 无缝转换为 SDE,并引入预训练 MLIP(或解析碰撞惩罚 \(E_{\text{overlap}}\))执行 Feynman-Kac 粒子重采样。这在不修改训练流程、无需对力场求梯度(gradient-free)的前提下,将形成能误差削减了 84% 以上。
- 可迁移的设计思路:这种“分块内局部参考几何作 Attention Bias + 全原子无约束流匹配 + 末期无梯度物理势能粒子滤波”的范式,可以直接迁移到共价有机框架(COF)、分子晶体多晶型预测、蛋白质-配体柔性诱导契合对接(induced-fit docking)以及非均相催化剂表面吸附构型采样等一切兼具“模块化拓扑块”与“局部柔性形变”的科学计算生成任务中。
局限与展望¶
- 金属节点模板对训练集先验的依赖:在构建单原子初始特征 \(\mathbf{C}\) 时,有机配体可由 RDKit 直接生成局部参考坐标,但无机金属节点(尤其是包含 887 种复杂节点的 ODAC25 数据集)仍需通过 SMILES 索引从训练集中检索并随机采样参考几何结构。当测试集出现全新的、未见过的金属簇拓扑时,这种查表式初始化将难以直接泛化。
- 全原子平方级注意力复杂度对超大单胞的限制:DiT 主干中的成对表征 \(\mathbf{Z} \in \mathbb{R}^{N \times N \times D_P}\) 与全原子自注意力具有 \(O(N^2)\) 的计算与显存开销(训练中不得不通过 \(N_{\max}^2\) 动态批处理与课程学习加以缓解)。对于单胞包含数千乃至上万原子的超大介孔 MOF 或高负载多分子协同吸附体系,全连接成对偏置将成为明显的扩展瓶颈。
- 改进方向:未来可引入基于局部图神经网络的从头金属簇构象预测器以彻底摆脱节点模板检索;在主干架构中引入稀疏周期性注意力或距离截断成对偏置以支持万级原子体系;同时将吸附温度、分压等宏观热力学状态变量作为条件注入流匹配模型,实现完整的吸附等温线端到端预测。
相关工作与启发¶
- vs MOFFlow / MOFFlow-2 / MOF-BFN:这些专门针对 MOF 的生成基线均建立在粗粒度或半柔性假设之上(MOFFlow 与 MOF-BFN 将构建单元视为刚体预测旋转与平移,MOFFlow-2 虽引入可旋转键扭转角但仍固定局部键长与键角),且仅能生成无客体的空框架;本文 AtomMOF 则完全放开全原子自由度并联合建模吸附物分子,在严格容差(\(\text{stol}=0.3\))下取得了压倒性的匹配率与孔道物性还原优势。
- vs DiffCSP / FlowMM(通用无机晶体结构预测):DiffCSP 与 FlowMM 针对小单胞无机晶体设计,缺乏对模块化大分子构建单元内部化学键拓扑的显式利用,因此在数百原子的 MOF 体系上几乎完全失效(BW 数据集上 DiffCSP 匹配率不足 1%);AtomMOF 通过块内成对几何注意力偏置有效桥接了分子图先验与周期性晶体生成。
- vs AlphaFold 3 / All-Atom DiT(Joshi et al., 2025):AtomMOF 在架构上汲取了生物大分子复合物预测(单特征 + 成对偏置 Transformer)与全原子扩散 Transformer 的思想,并将其成功拓展至具有周期性边界条件(PBC)、可变晶胞参数 \(\bm{\ell}\) 以及多分子气相吸附的结晶多孔材料体系。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次实现全原子无约束 MOF 及主客体吸附复合物联合生成,并结合 MLIP 驱动的 Feynman-Kac 引导与多孔晶体 Scaling Law 分析。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 BW 纯框架精度、Zeo++ 8 项孔道物性、38.4M 至 491M 参数与步数扩展律、FK 引导全超参消融,以及 ODAC25 与 GCMC/MLIP 弛豫的对比。
- 写作质量: ⭐⭐⭐⭐⭐ 问题形式化、变分流匹配推导、三阶段 DiT 伪代码及附录实验细节层次分明,图表信息密度高且逻辑严密。
- 价值: ⭐⭐⭐⭐⭐ 为碳捕集(DAC)等能源环境材料的高通量筛选提供了比传统 GCMC 快数个数量级的全原子结构生成与低能构型初始化工具,代码已开源。