跳转至

Boosting LLM's Molecular Structure Elucidation with Knowledge Enhanced Tree Search Reasoning

会议: ACL 2025
作者: Xiang Zhuang, Bin Wu, Jiyu Cui, Kehua Feng, Xiaotong Li, Huabin Xing, Keyan Ding, Qiang Zhang, Huajun Chen (浙江大学, UCL)
arXiv: 2506.23056
代码: GitHub
领域: LLM/NLP, 化学推理, 分子结构解析
关键词: molecular structure elucidation, MCTS, knowledge base, reward model, spectral data, test-time scaling

一句话总结

提出 K-MSE(Knowledge-enhanced Molecular Structure Elucidation)框架,构建分子子结构知识库扩展 LLM 的化学结构空间覆盖,设计专用分子-光谱打分器替代 LLM 自身评估,结合蒙特卡洛树搜索(MCTS)实现测试时推理缩放,在 MolPuzzle 基准上分别将 GPT-4o-mini 和 GPT-4o 的准确率从 3.7% 和 27.8% 提升至 27.3% 和 39.8%。

研究背景与动机

  • 核心问题: 分子结构解析是化学实验分析的基础任务——从 NMR、IR 等光谱数据推断分子结构。即使专家也需 10-15 分钟处理一个分子。LLM 有潜力自动化这一过程,但面临两大挑战。
  • 现有不足: (1) LLM 缺乏对化学分子结构空间的全面覆盖——对噻吩等非常见结构常误判为苯环(最常见芳香结构);(2) LLM 无法准确评估自身推理结果,缺乏领域知识来判断预测分子与光谱数据的匹配度,导致树搜索推理缺乏有效奖励信号。
  • 研究动机: 通过外部知识增强化学结构覆盖 + 专用打分器提供准确奖励 → 结合 MCTS 实现 LLM 在分子结构解析中的测试时推理缩放。

方法详解

整体框架

K-MSE 由三个组件构成: 1. 分子子结构知识库 \(\mathcal{KB} = \{(s_i, d_i)\}\):包含子结构 SMILES 表示和文本描述,从 MOSES 分子数据库提取环状和链状子结构 2. 分子-光谱打分器:由分子编码器 \(g_m\)(GIN + MLP 处理分子图和指纹)和光谱编码器 \(g_s\)(Transformer 处理 C-NMR/H-NMR 的化学位移、裂分模式、耦合常数)组成 3. MCTS 推理框架:先从知识库检索相关子结构 → 迭代执行选择(UCT)→扩展(Critique+Rewrite)→评估(打分器)→反向传播

关键设计

  1. 知识库构建:从 MOSES 数据库自动提取分子子结构,利用 LLM 结合外部工具生成的结构信息自动生成可靠描述。兼顾多样性和通用性。
  2. 专用打分器:分子编码器使用 GIN 编码分子图 + MLP 编码 Morgan 指纹,光谱编码器将 NMR 化学位移和耦合常数离散化为 token ID 后输入 Transformer。训练采用 NT-Xent 对比学习损失,使匹配的分子-光谱对嵌入相似度最大化。
  3. 打分器双重角色:既作为 MCTS 奖励模型评估候选分子(\(R(a') = \text{sim}(g_m(m_{a'}), g_s(n))\)),又作为知识库检索桥梁——用光谱编码器编码查询光谱,用分子编码器编码子结构,进行 Top-k 检索。

损失函数

打分器训练使用 NT-Xent 对比学习损失:最大化正确分子-光谱对的余弦相似度,最小化批内负样本对的相似度,温度参数 \(\tau\) 控制分布锐度。MCTS 反向传播采用 \(Q(a) = 0.5 \times Q(a') + 0.5 \times Q(a)\) 的加权更新。

实验

主实验——MolPuzzle 基准(216 个分子,zero-shot)

模型 方法 Morgan FTS MACCS FTS ACC
GPT-4o-mini baseline 0.260 0.512 0.037
GPT-4o-mini + Self-Refine 0.287 0.523 0.069
GPT-4o-mini + MCTSr 0.281 0.530 0.069
GPT-4o-mini + K-MSE 0.470 0.651 0.273
GPT-4o baseline 0.493 0.690 0.278
GPT-4o + Self-Consistency 0.551 0.732 0.347
GPT-4o + K-MSE 0.398
Llama-3.2-11B baseline 0.163 0.349 0.014
Llama-3.2-11B + K-MSE 0.298 0.465 0.111

消融实验

消融组件 对 GPT-4o-mini ACC 的影响
完整 K-MSE 0.273
移除知识库 下降明显——LLM 无法识别非常见子结构
用 LLM 替代专用打分器 下降显著——LLM 无法准确评估分子-光谱匹配
移除 Critique 中的分子图像 下降——纯文本 critique 难以发现结构错误
移除 Critique 中的化学式 下降——缺乏化学约束信息

关键发现

  • K-MSE 在所有基座模型上带来大幅提升:GPT-4o-mini ACC +23.6%,GPT-4o ACC +12.0%,Llama-3.2-11B ACC +9.7%
  • 现有通用推理增强方法(Self-Refine, MCTSr, MAD)在分子结构解析上效果有限——缺乏领域知识是核心瓶颈
  • 专用打分器远优于 LLM 自评估——LLM 缺乏判断分子-光谱匹配度的领域知识
  • 知识库的子结构信息对处理非常见分子结构至关重要
  • 作为即插即用框架,K-MSE 可与任何 LLM 组合使用

亮点

  • 首次将测试时推理缩放(test-time scaling)+ 外部知识增强应用于分子结构解析任务
  • 打分器同时充当奖励模型和检索桥梁的双重角色设计精巧
  • 框架的即插即用特性使其具有很强的实用价值
  • MolPuzzle 上 20%+ 的绝对准确率提升非常显著

局限性

  • 仅在 MolPuzzle 基准上评估,该基准规模较小(216 个分子)
  • 打分器的训练数据覆盖范围可能限制其对罕见分子类型的泛化能力
  • MCTS 的迭代次数增加带来显著的推理时间成本(API 调用+打分器推理)
  • 知识库是静态的,未探索在线扩展或自适应更新机制
  • 仅考虑 NMR 和 IR 光谱,未处理质谱(MS)等其他常用分析数据

相关工作

  • LLM 化学推理: ChemCrow (M. Bran et al., 2024) 集成外部工具、ChatDrug (Liu et al., 2024) 分子编辑、STRUCTCHEM (Ouyang et al., 2024) 预定义推理模板
  • 树搜索推理: Tree-of-Thought (Yao et al., 2023)、MCTSr (Zhang et al., 2024a),但缺乏领域特定的准确奖励模型
  • 分子结构解析: MolPuzzle (Guo et al., 2024) 首次提出该任务的 LLM 基准

评分

维度 分数
新颖性 ⭐⭐⭐⭐⭐
技术深度 ⭐⭐⭐⭐⭐
实验完整度 ⭐⭐⭐⭐
写作质量 ⭐⭐⭐⭐
实用性 ⭐⭐⭐⭐