跳转至

WA-DD 分子生成原理

本文档面向没有读过 PocketXMol 原始文献的同事，解释 WA-DD 中分子生成模块用到的模型原理、口袋信息的使用方式、编码与模型结构、输入输出以及完整的生成步骤。阅读后应能理解：系统为什么这样设计、每个参数在控制什么、以及生成结果为什么长这样。

1. 一句话定位

WA-DD 的分子生成不是基于规则拼接片段，而是一个深度生成模型——它学习大量真实蛋白-小分子复合物中的原子相互作用规律，然后在指定口袋区域内从头生成三维分子构象。当前接入的引擎是 PocketXMol。

PocketXMol 来自清华大学团队 2026 年发表于 Cell 的工作：

Peng, X. et al. "Unified modeling of 3D molecular generation via atomic interactions with PocketXMol." Cell (2026). DOI: 10.1016/j.cell.2026.01.003 开源代码：https://github.com/pengxingang/PocketXMol

与传统做法把"分子对接""骨架跃迁""片段生长"等任务拆成不同工具不同，PocketXMol 试图把所有这些任务统一成同一个问题：给定一部分已知原子（蛋白口袋、已有片段、约束），生成剩余未知原子的类型、三维坐标和化学键。

2. 生成模型基础：扩散模型是什么

2.1 类比

扩散模型（Diffusion Model）的核心思想可以类比为：

  1. 正向过程：从一张清晰照片开始，逐步加入噪声，直到完全变成随机雪花点。
  2. 反向过程：训练一个神经网络，学习如何从雪花点一步步恢复出清晰照片。
  3. 生成：从纯噪声出发，执行反向过程，得到一张模型"想象"出来的新照片。

对分子生成来说，"清晰照片"就是一个完整的分子：每个原子的元素类型、三维坐标、原子之间的化学键。"噪声"就是把这些信息打乱：坐标随机偏移、元素类型随机替换、化学键随机改变。

2.2 训练时学什么

训练阶段，PocketXMol 会：

  1. 取一个真实的蛋白-配体复合物结构。
  2. 随机选择一部分原子作为"已知"（比如蛋白口袋原子），另一部分作为"待生成"（配体原子）。
  3. 对"待生成"部分加入不同程度的噪声。
  4. 让模型观察"已知部分"和"当前带噪的待生成部分"，预测原始干净结构。

通过海量数据（论文提到约 1200 万小分子、4 万多蛋白-多肽复合物、8.5 万多蛋白-小分子复合物）的训练，模型学会的不是背下某个分子，而是原子之间合理的空间关系和成键规律。

2.3 采样时如何生成

生成（采样）阶段，模型：

  1. 从随机噪声初始化一个分子。
  2. 把蛋白口袋作为固定条件输入模型。
  3. 迭代执行多步"去噪"：每步预测当前噪声结构对应的更干净结构。
  4. 最后把去噪结果还原成规范的 SDF 分子文件。

步数越多，生成过程越充分，但也越慢。

3. 口袋信息怎么用

3.1 什么是口袋

在结构生物学中，"口袋"（pocket/binding site）是蛋白表面一个适合小分子结合的凹陷区域。WA-DD 用一个中心点 + 半径来定义 PocketXMol 的生成区域，同时用一个PDB 文件保存口袋附近的残基，用于模型输入和可视化。

3.2 口袋资产里的关键字段

WA-DD 的 pocket 资产元数据包含：

字段 含义 用途
center 三维空间中心点 [x, y, z] 对接盒中心和分子生成中心
box_size 对接盒尺寸 [sx, sy, sz] Vina/AutoDock/Uni-Dock 的搜索空间
generation_radius 生成半径（Å） PocketXMol 实际使用的口袋范围
pocket_structure 口袋残基 PDB 模型输入 + 3D 可视化

generation_radius 默认从 box_size 推导：

max_box = max(box_size)
generation_radius = max(10.0, min(24.0, max_box / 2.0 + 4.0))

也就是说，如果 box 是 20 Å 的立方体，生成半径约为 14 Å。用户也可以在提交任务时显式覆盖。

3.3 模型如何利用口袋

PocketXMol 不会把整个蛋白都读入模型，而是只取口袋半径内的残基原子，构建一个蛋白原子图（pocket graph）：

  • 节点：口袋中的原子（包含元素类型、残基信息等特征）。
  • 边：原子之间的空间邻接关系（通常用 KNN 或半径图构建）。

这个口袋图在整个生成过程中是固定不变的条件上下文。模型生成配体时，每一步都要参考口袋原子的位置和化学环境，从而确保新分子与口袋能够形成合理的相互作用（氢键、疏水作用、π-π 堆积等）。

3.4 为什么不用整个蛋白

两个原因：

  1. 计算效率： distant 残基对口袋内结合没有直接影响，去掉可以显著降低显存和计算量。
  2. 聚焦学习：模型在训练时也只关注结合位点附近，因此生成时只需要口袋上下文。

4. 分子和口袋如何编码

PocketXMol 的关键设计是原子级统一表示。它不区分小分子、多肽或蛋白，全部看作"原子 + 键"的图。

4.1 分子图表示

一个分子被表示为图：

  • 节点（node）：每个重原子（氢原子通常被隐式处理）。
  • node_type：元素/原子类型编号（C、N、O、S、P、卤素等）。
  • node_pos：三维坐标 [x, y, z]。
  • 半边（half-edge）：因为图神经网络中边是双向的，但化学键本身无方向，所以用半边存储 i < j 的键。
  • halfedge_type：键类型编号（无键、单键、双键、三键、芳香键、MASK 等）。
  • halfedge_index：键连接的两个原子索引。

4.2 口袋图表示

口袋同样被编码为图：

  • pocket_atom_feature：口袋原子的输入特征（元素、残基类型等 one-hot 或嵌入）。
  • pocket_pos：口袋原子坐标。
  • pocket_knn_edge_index：基于 KNN 构建的口袋原子邻接边。

4.3 任务提示（Task Prompt）

PocketXMol 用一个巧妙的机制统一不同任务：二进制固定掩码（fixed masks）。对图中每个原子/键，模型都知道它是不是"已知、不可改动"的：

  • fixed_node = 1：原子类型已知，不要改。
  • fixed_pos = 1：原子坐标已知，不要动。
  • fixed_halfedge = 1：化学键已知，不要改。

这些掩码作为额外特征拼接进模型输入，成为"任务提示"。不同任务只是掩码设置不同：

任务 固定什么 生成什么
口袋从头生成（pocket_de_novo） 蛋白口袋原子 完整配体所有原子
片段生长（fragment_grow） 蛋白口袋 + 参考 fragment 原子 fragment 之外的扩展部分
骨架跃迁（scaffold_hop） 蛋白口袋 + 部分骨架原子 被替换的骨架部分
连接子设计（linker_design） 蛋白口袋 + 两个固定片段 连接两个片段的 linker

WA-DD 当前开放了前两种，后两种需要用户在 3D 视图中点选保留/替换/锚点原子，页面暂不支持。

5. 模型结构：PMAsymDenoiser

PocketXMol 的核心网络叫 PMAsymDenoiser（Pocket-Molecule Asymmetric Denoiser），源码在 models/maskfill.py。它由四部分组成：

口袋特征编码器（Pocket Encoder）
        ↓
分子嵌入层（Molecule Embedder）
        ↓
去噪主干网络（Denoiser Backbone）
        ↓
输出解码器（Output Decoders）

5.1 口袋编码器（Pocket Encoder）

输入：口袋原子特征 + 口袋原子坐标 + 口袋 KNN 边。

  • 先用一个线性层把 pocket_atom_feature 映射到隐藏维度。
  • 再用一个图神经网络（ContextNodeEdgeNet）在口袋图内部做消息传递，让每个口袋原子获得周围邻居和全局结构信息。
  • 输出 h_pocket：每个口袋原子的上下文向量。

这个编码器让模型"理解"口袋的形状、极性、可成键位点。

5.2 分子嵌入层（Molecule Embedder）

输入：当前带噪的分子（node_in、pos_in、halfedge_in）+ 任务提示掩码。

  • nodetype_embedder：把原子类型编号嵌入成向量。
  • edgetype_embedder：把键类型编号嵌入成向量。
  • 把 fixed_node、fixed_pos 等二进制掩码作为额外通道拼接进去。

这样每个原子和键的初始向量都包含了：当前是什么 + 这个位置是不是被锁定的。

5.3 去噪主干网络（Denoiser Backbone）

这是模型最核心的部分。它也是一个图神经网络（ContextNodeEdgeNet 或 GVP/IPA 变体），但有两个特别之处：

  1. 同时处理节点和边：消息传递不仅在原子之间进行，也在化学键之间进行，因此模型能学习成键约束。
  2. 跨图注意力/上下文融合：分子图的消息传递会读取口袋编码器的输出 h_pocket，让分子原子的更新受口袋原子引导。

经过多层消息传递后，每个原子和键都获得了富含上下文的隐藏表示。

5.4 输出解码器（Output Decoders）

从隐藏表示直接预测三个东西：

  • pred_node：每个原子的元素类型概率分布。
  • pred_pos：每个原子的三维坐标（去噪后的位置）。
  • pred_halfedge：每条半边的键类型概率分布。

在训练时，这三个预测会分别与真实值计算损失；在采样时，它们被用来更新当前分子。

5.5 可选：置信度分数

如果配置里开启 confidence，模型还会对每个原子、坐标、键预测一个置信度，用于后续对多个生成结果排序。

6. 输入与输出

6.1 任务输入

WA-DD 提交分子生成任务时，API 端点 /api/v1/molecule-generation/pocketxmol 接收以下关键参数：

参数 默认值 说明
protein_asset_id 必填 蛋白/已准备蛋白/complex 资产
pocket_asset_id 必填 口袋资产，提供 center/radius/pocket_structure
mode pocket_de_novo pocket_de_novo 或 fragment_grow
reference_ligand_asset_id 可选 fragment_grow 模式下必填
num_mols 20 总生成数量
batch_size 20 每轮同时生成的分子数
num_steps 100 每分子去噪步数
generation_radius 从 box 推导 PocketXMol 半径
mean_atoms / std_atoms / min_atoms 28 / 2 / 5 生成分子大小的正态分布
device cuda:0 GPU 设备

6.2 运行时的真实输入文件

worker 会准备以下文件给 PocketXMol 的 scripts/sample_use.py：

  • protein_path：蛋白 PDB。
  • 可选 input_ligand：参考 fragment SDF（fragment_grow 模式）。
  • task_config：YAML，包含采样参数、口袋参数、任务类型、噪声调度。
  • model_config：YAML，仅指定 checkpoint 路径。

6.3 输出资产

任务完成后，worker 会创建两类资产：

  1. 一个 prepared_ligand_library：
  2. prepared_structure：加氢后的完整 SDF。
  3. generated_structure：原始生成 SDF。
  4. pocket_structure：口袋 PDB。
  5. report：生成报告 JSON。
  6. interaction_report / interaction_table：几何相互作用预览。
  7. molecule_table：每个分子的 SMILES、分子式、分子量、重原子数。

  8. run_config / log：运行配置和日志。

  9. 多个 prepared_ligand：

  10. 每个生成构象一个独立资产。
  11. 可直接用于对接、FEP、相互作用分析。

6.4 加氢处理

PocketXMol 生成的是重原子结构。worker 用 RDKit 的 Chem.AddHs(mol, addCoords=True) 加氢，保持三维坐标不变，并在 SDF 属性中记录：

wa_dd_hydrogen_method: rdkit_addhs_preserve_pose
wa_dd_prepared_for: docking

7. 分子生成的完整步骤

从用户点击提交到资产入库，流程如下：

1. Web 端创建 Job
   - 校验蛋白、口袋、参考配体资产
   - 计算 generation_radius
   - 写入 Postgres，状态 queued
   - 发布 Redis 事件

2. PocketXMol worker 拉取 Job
   - 只处理 job_type=molecule_generation + engine=pocketxmol

3. 准备模型
   - 检查 /modelhub/export/ms/huluxiaohuowa/pocketxmol/current
   - 缺失则通过 modelscope download 下载

4. 准备输入
   - 解析蛋白 PDB
   - 解析口袋资产元数据
   - fragment_grow 模式：读取参考配体第一个分子作为 fragment

5. 写入任务配置 YAML
   - pocket_de_novo：task.name=sbdd，transform.name=ar
   - fragment_grow：task.name=maskfill，保留 fragment 全部重原子

6. 执行采样
   - 调用 scripts/sample_use.py
   - 正向/反向扩散循环
   - 输出 SDF 到工作目录

7. 后处理
   - 收集所有成功 SDF
   - 合并成一个 ligand library
   - 加氢
   - 生成分子信息表
   - 计算几何相互作用预览

8. 写入资产
   - 创建 library 资产
   - 为每个构象创建独立 prepared_ligand 资产
   - 更新 Job 状态为 completed
   - 发布完成事件

8. 两种开放模式的区别

8.1 基于口袋从头生成（pocket_de_novo）

  • 只固定蛋白口袋。
  • 分子从完全噪声开始生成。
  • 适合探索全新化学骨架。
  • 任务配置示例（简化）：
task:
  name: sbdd
  transform:
    name: ar
    part1_pert: small
noise:
  name: maskfill
  num_steps: 100
  ar_config:
    strategy: refine
    r: 3
    threshold_node: 0.98
    threshold_pos: 0.91
    threshold_bond: 0.98
    max_ar_step: 10

这里的 ar 指 autoregressive-like refine，每步先生成/修正一部分原子，再逐步细化。

8.2 片段生长（fragment_grow）

  • 固定蛋白口袋 + 参考 fragment 的所有重原子。
  • 模型只生成 fragment 之外的扩展部分。
  • 适合在已知活性片段基础上做 SAR 扩展。
  • WA-DD 当前采用"整个参考分子的第一个构象作为 fragment"的语义，用户不需要手动选原子。
  • 任务配置示例（简化）：
transforms:
  variable_mol_size:
    not_remove: [0, 1, 2, 3, 4, 5, 6]  # fragment 原子索引
task:
  name: maskfill
  transform:
    name: maskfill
    preset_partition:
      grouped_node_p1: [[0, 1, 2, 3, 4, 5, 6]]
    settings:
      part1_pert:
        fixed: 1

8.3 暂未开放的模式

  • Scaffold hopping：需要用户选择哪些原子属于要被替换的骨架。
  • Linker design：需要用户选择两个或多个片段以及连接锚点。

这两个模式在 PocketXMol 里完全支持，但 WA-DD 页面缺少原子/锚点选择器，所以 API 直接拒绝提交，避免任务语义不清。

9. 关键参数调优指南

参数 影响 建议
num_mols 总候选数 20-100；先小批量验证
batch_size 显存占用、并行度 16-24 GB 显存用 20；小口袋/小半径可尝试 50
num_steps 每分子去噪深度 20 快速验证；100 正式；200+ 更充分但慢
generation_radius 口袋范围、显存 默认足够；盲目增大会显著增加计算量
mean_atoms 目标分子大小 默认 28 重原子适合典型小分子；大环/肽可调
std_atoms 大小波动 默认 2；想多样性大可增大
min_atoms 最小原子数 默认 5；过滤掉过碎的结果

注意：batch_size 和半径对显存的影响通常大于 num_steps。

10. 训练数据与模型能力边界

10.1 训练数据

论文中提到 PocketXMol 在以下数据上联合训练：

  • 约 1200 万个小分子（用于学习一般化学规律）。
  • 约 4 万个蛋白质-多肽复合物。
  • 约 8.5 万个蛋白质-小分子复合物。

联合训练让模型学到的不是某个任务的特解，而是原子间相互作用的通用规律，因此能在 13 个计算基准中的 11 个达到 SOTA。

10.2 能做什么

  • 在已知口袋内生成全新小分子。
  • 在保留 fragment 的前提下扩展分子。
  • 生成与口袋几何互补的构象。
  • 作为下游对接/FEP 的候选输入。

10.3 不能做什么/注意事项

  • 不保证合成可及性：生成的分子在化学上合理，但不一定容易合成。
  • 不替代实验验证：结合亲和力、选择性、ADMET 仍需实验或更精细计算。
  • 几何相互作用预览不是能量计算：worker 输出的 interaction 表只是距离阈值几何预览，用于快速查看匹配，不是 MD 或严格能量模拟。
  • 受限于训练分布：如果目标口袋或分子类型与训练数据差异很大，生成质量可能下降。
  • 多肽/非天然氨基酸：PocketXMol 支持，但 WA-DD 当前小分子流程默认 is_pep: False。

11. 与下游任务的衔接

生成结果已经被 worker 标记为 prepared_for: [docking, wa-dd, fep_md]，可以直接进入下游：

  • 对接：在 Docking 页面选择生成的 prepared_ligand 或 prepared_ligand_library 作为配体。
  • FEP：选择单个 prepared_ligand 构象作为候选配体。
  • 相互作用分析：选择生成时使用的 pocket 资产作为受体上下文，再选 prepared_ligand 查看几何相互作用。

如果下游工具需要 PDBQT 等格式，由对应的 worker 在提交时自动转换，不需要用户手动修改 SDF。

12. 参考资料

  1. PocketXMol 论文：Peng, X. et al. "Unified modeling of 3D molecular generation via atomic interactions with PocketXMol." Cell (2026). DOI: 10.1016/j.cell.2026.01.003
  2. PocketXMol 开源代码：https://github.com/pengxingang/PocketXMol
  3. WA-DD 集成代码：
  4. Worker: src/pocketxmol_worker/worker.py
  5. API: src/wa_dd_web/main.py
  6. Schema: src/wa_dd_web/schemas.py
  7. 用户指南：userguide/molecule-generation.md