Neural Graph Matching Improves Retrieval Augmented Generation in Molecular Machine Learning¶

会议: ICML2025
arXiv: 2502.17874
代码: coleygroup/ms-pred
领域: 分子机器学习 / 图匹配 / 检索增强生成
关键词: 神经图匹配, RAG, 质谱模拟, 分子碎片DAG, GNN, MARASON

一句话总结¶

提出 MARASON，将神经图匹配（Neural Graph Matching）引入分子机器学习的检索增强生成（RAG）框架，通过可微分的碎片级对齐机制，把检索到的参考分子谱图信息有效融入目标分子的质谱预测中，在 NIST 数据集上将 top-1 检索准确率从 19% 提升到 28%。

研究背景与动机¶

检索增强生成（RAG） 在 LLM 领域已是成熟范式，但在分子机器学习中的最优集成方式尚不明确。简单的拼接（concatenation）策略几乎不带来改进。
质谱模拟（MS/MS simulation） 是分子机器学习的重要应用：给定分子结构预测其质谱（m/z 值与峰强度），可加速未知化合物的结构解析，应用于代谢组学、生物标志物发现和环境科学等领域。
化学直觉表明：结构相似的分子具有相似的碎裂模式和谱图。领域专家在比较分子时会进行原子/碎片级的结构对齐，而非仅看全局指纹相似度。
传统图匹配方法（如 Hungarian、RRWM）使用固定亲和度度量（如 Tanimoto + Gaussian kernel），表达能力有限、对噪声不鲁棒。
核心洞察：需要一种端到端可学习的碎片级匹配机制，既能捕捉节点（碎片）亲和度，又能建模边（碎裂层次）信息，从而让 RAG 真正发挥作用。

方法详解¶

整体流程（MARASON）¶

MARASON 基于 ICEBERG 模型（ICML 2024 SOTA），分为三个核心模块：

检索模块：基于 Morgan 指纹 + Tanimoto 相似度从训练集中检索最相似的参考分子及其谱图
神经图匹配模块：对目标和参考分子的碎裂 DAG 进行碎片级对齐
强度预测模块：利用匹配结果融合参考谱图信息，预测目标分子谱图

检索增强处理¶

从训练集数据库中检索 Tanimoto 相似度最高的参考分子 \(\mathcal{M}^r\)
排除加合物类型或仪器类型不匹配的条目
选取最多 3 条碰撞能量最接近的参考谱图，通过插值学习目标碰撞能量下的谱图嵌入
对参考谱图的每个碎片 \(\mathcal{F}_j^r\)，在 ±6 个氢原子质量偏移范围内匹配峰，得到 13 维强度向量
经 Set Transformer + 平均池化得到参考强度嵌入 \(\mathbf{T}^r\)

碎裂 DAG 图匹配¶

ICEBERG-Generate 模型对目标和参考分子各生成一个碎裂有向无环图（DAG），其中节点=碎片，边=碎裂路径。

传统方法（对比基线）：

线性分配问题（Hungarian 算法）：

\[\max_{\mathbf{X}} \text{tr}(\mathbf{M}^\top \mathbf{X}), \quad \text{s.t.} \; \mathbf{X} \in \{0,1\}^{n \times n^r}\]

其中亲和度矩阵 \(m_{i,j} = \text{Tanimoto}(\mathcal{F}_i, \mathcal{F}_j^r)\)。

二次分配问题（RRWM 求解器）进一步建模边亲和度：

\[\max_{\mathbf{X}} \text{vec}(\mathbf{X})^\top \mathbf{K} \text{vec}(\mathbf{X})\]

神经图匹配（核心贡献）：

(1) 碎片级嵌入学习：使用共享 \(\text{GNN}_{\text{frag}}\) 对每个碎片及其母分子编码，拼接碎片嵌入、母分子嵌入、差值嵌入（表示"中性丢失"）、断键数、化学式差异，经 MLP 得到碎片嵌入 \(\mathbf{H}\)（目标）和 \(\mathbf{H}^r\)（参考）。

(2) DAG 层次嵌入学习：构建正向 DAG \(\mathcal{G}\) 和反向 DAG \(\mathcal{G}^{-1}\)，通过双向 GNN 更新嵌入：

\[\bar{\mathbf{H}} \leftarrow \mathbf{H} + \text{GNN}_{\text{fwd}}(\mathbf{H}, \mathcal{G}) + \text{GNN}_{\text{rev}}(\mathbf{H}, \mathcal{G}^{-1})\]

(3) 可微分匹配层：计算余弦相似度矩阵后通过 Softmax 得到连续匹配矩阵：

\[\bar{m}_{i,j} = \text{cosine}(\bar{\mathbf{h}}_i, \bar{\mathbf{h}}_j^r), \quad \bar{\mathbf{X}} = \text{Softmax}(\bar{\mathbf{M}})\]

强度预测¶

将目标碎片嵌入、对齐后的参考碎片嵌入、对齐后的参考强度、匹配分数和全局 Tanimoto 相似度拼接：

\[\text{Input} = [\mathbf{H}, \bar{\mathbf{X}}\mathbf{H}^r, \bar{\mathbf{X}}\mathbf{T}^r, \mathbf{s}, \text{Tanimoto}(\mathcal{M}, \mathcal{M}^r)]\]

其中匹配分数 \(s_i = \sum_{j=1}^{n^r} \bar{x}_{i,j} \bar{m}_{i,j}\)。

输入经 Set Transformer + Attention + MLP 输出最终谱图强度。整个流程端到端可微，匹配层通过梯度反传自动学习最优亲和度度量。

实验关键数据¶

数据集¶

NIST 2020：530,640 条 HCD 谱图，25,541 个唯一分子结构，80/10/10 划分
MassSpecGym：更具挑战性的通用基准，强调对新骨架的泛化

检索准确率（NIST 2020，随机划分，正加合物）¶

方法	Top-1	Top-5	Top-10
3DMolMS	0.055	0.225	0.394
NEIMS (GNN)	0.175	0.515	0.687
MassFormer	0.191	0.550	0.716
ICEBERG	0.189	0.623	0.770
ICEBERG (w/ CE)	0.202	0.639	0.793
MARASON	0.278	0.685	0.827

Top-1 准确率相比无 RAG 的 ICEBERG 提升 47%（0.189→0.278），相比加入碰撞能量的 ICEBERG 提升 38%。

MassSpecGym 检索准确率¶

方法	Top-1	Top-5	Top-20
FraGNNet	0.319	0.632	0.827
MARASON	0.340	0.640	0.854

消融实验（余弦相似度，随机划分）¶

RAG 策略	匹配层	余弦相似度
无 RAG	-	0.739
拼接参考谱图	-	0.737 (−0.3%)
Hungarian	-	0.746 (+0.9%)
RRWM	-	0.742 (+0.4%)
NGM (共享GNN)	Softmax	0.753 (+1.9%)
NGM (独立GNN)	Softmax	0.757 (+2.4%)

关键发现：

简单拼接参考谱图反而略微降低性能（−0.3%），验证了朴素 RAG 方案无效
传统图匹配（Hungarian/RRWM）有一定提升但有限
神经图匹配显著优于传统方法，Softmax 优于 Sinkhorn
独立 GNN（目标/参考/共享各一组）优于共享 GNN

亮点与洞察¶

化学直觉与深度学习的巧妙结合：将领域专家"比较类似碎片"的实践形式化为可微分的神经图匹配，是 RAG 在分子领域的原则性设计范式
嵌套 GNN 架构：碎片级 GNN + DAG 层次 GNN 的嵌套设计，既编码局部子结构信息，又保留全局碎裂层次关系
端到端学习亲和度度量：相比固定的 Tanimoto 度量，可学习的亲和度对噪声和结构歧义更鲁棒（如图 1 所示的等排体识别）
消融实验设计严谨：系统对比了无 RAG、朴素拼接、传统匹配、神经匹配等多种策略，清晰展示了各设计选择的贡献
应用潜力巨大：当前 NIST 仅覆盖约 27K 化合物，MARASON 有望将模拟谱库扩展至 PubChem 的 1.11 亿化合物

局限与展望¶

单参考分子检索：当前仅检索 1 个最相似的参考分子，多参考融合策略可能进一步提升性能
依赖 ICEBERG-Generate：碎裂 DAG 的质量受限于预训练的 ICEBERG-Generate 模型，碎片生成错误会传播到匹配环节
计算开销：嵌套 GNN + 图匹配的计算成本高于简单的前馈模型，论文未详细报告推理延迟
数据集局限：主要在 NIST 2020 上验证，该数据集以小分子（<1500 Da）为主，对大分子和生物大分子的适用性未知
Scaffold 划分性能下降明显：在 Murcko scaffold 划分下所有方法性能均显著下降，说明对分布外骨架的泛化仍是挑战
负加合物实验不充分：论文主要展示正加合物结果，负加合物类型的全面评估缺失

评分¶

新颖性: ⭐⭐⭐⭐ — 神经图匹配 × 分子RAG 的交叉创新，设计范式有通用价值
实验充分度: ⭐⭐⭐⭐ — NIST + MassSpecGym 双数据集验证，消融全面，但缺少推理效率对比
写作质量: ⭐⭐⭐⭐⭐ — 动机清晰、化学直觉与方法论结合紧密、图示精良
价值: ⭐⭐⭐⭐ — 对分子机器学习中的 RAG 设计提供了原则性指导，应用潜力大