简介：围绕ICML 2024提出的IRDiff检索增强3D分子扩散模型，这份评测文档与可运行代码包面向药物设计、分子生成方向的研究者与工程实践者，重点解决论文复现、代码排错以及基于特定蛋白或口袋开展目标感知分子生成的落地问题。压缩包仅含1个PDF文件，大小约12.23MB，正文把完整评测过程、修正后的项目代码、缺失模块文件、报错位置与修改方法、测试案例及个人分析标注整合在一起，并系统梳理模型背景、PMINet结构以及检索增强与自增强机制，便于按线索对照理解与复用。修正后的代码可在指定蛋白/口袋体系上调用IRDiff生成分子，计算vina_score、vina_docking_score、qvina_score、QED、SA等指标，也支持用自定义数据集微调或再训练，为模型修改与实验设计提供可操作起点。已有121人学习，适合希望从论文复现走向可运行实验、并进一步做微调优化的中高级读者参考。

1. 从头生成的分子扩散模型，为什么在结合亲和力上总是差一口气

分子生成、扩散模型、药物设计这三个词摆在一起，最容易让人产生一个错觉：只要生成的 3D 结构足够“真”，它离成为候选配体就不远了。实际评测里恰恰相反——结构指标最漂亮的模型，结合亲和力往往最难看。在 CrossDocked 2020 这套基准上，自回归方案和早期扩散方案都卡在这个位置：全原子成对距离分布能压得很低，Vina 分数却长期徘徊，而高通量筛选最先看的就是后者。

IRDiff 的思路是把“从零生成”换成“带着参考生成”。它先用一个预测结合亲和力的网络从外部配体库里挑出真正贴合口袋的分子，再把它们的结构信息灌进扩散过程。听起来只是加了个检索模块，实际上改的是扩散模型被条件的那个分布本身——参考分子把“什么样的子结构能在这个口袋里形成强相互作用”这件事，从隐式的模型参数里搬到了显式的输入里。后面几章拆的就是这个“搬”字具体落在了哪一层。

2. PMINet：一个只输出亲和力标量的网络，凭什么反过来指导生成

2.1 双塔等变编码加原子级交叉注意力

PMINet 的结构不复杂，但每一层都有明确分工。蛋白质和配体先各自过一遍浅层的 SE(3)-等变网络，在全连接图上做消息传递，这一步建模的是分子内相互作用——键长、键角、原子局部环境。两路编码互不干扰，蛋白原子坐标保持原样，配体原子坐标随扩散状态变化。

真正把“蛋白-配体”这一对连起来的是原子级交叉注意力：配体原子当 query，蛋白原子当 key/value。每个配体原子去问一遍“我周围有哪些蛋白原子、它们是什么类型”，得到的是一份局部相互作用描述。最后池化到复合物级别，回归出一个亲和力标量。

# PMINet 前向的数据流示意，用于说明张量怎么流动，不是逐行源码 def forward(protein_xyz, protein_type, ligand_xyz, ligand_type): # 1) 两路独立的 SE(3)-等变编码，各自在分子内全连接图上做消息传递 h_p = egnn_stack(protein_xyz, protein_type) # 蛋白原子级特征 [Np, d] h_l = egnn_stack(ligand_xyz, ligand_type) # 配体原子级特征 [Nl, d] # 2) 原子级交叉注意力：配体原子作 query，蛋白原子作 key/value h_cross = cross_attention(q=h_l, kv=h_p) # 每个配体原子一份相互作用上下文 # 3) 池化到复合物级别，回归一个结合亲和力标量 return affinity_head(pool(h_cross))

关键参数是d（隐藏维度）和egnn_stack的层数——原文写的是“两个浅层”，意味着它不追求把几何细节抠到极致，而是要在可接受的算力下把相互作用压成一个够用的向量。坐标输入保证了等变性，原子类型走离散嵌入，这两条决定了后面融合时它输出的到底是什么形态的特征。

2.2 用亲和力预训练换来的到底是什么先验

预训练数据是 PDBbind v2016，3767 个复合物做训练、290 个做测试。训练目标很朴素：让预测的亲和力和真实值靠近。但副产品比目标值本身更有价值——每个蛋白质-配体对都会得到一个“相互作用感知”的嵌入向量，这个向量不是口袋形状的描述子，也不包含任何 3D 体素信息，它压缩的是“这对分子和这个口袋配不配”。

这个性质决定了 PMINet 能同时干两件事。第一件是检索排序：给一个蛋白，扫一遍配体库，按预测亲和力从高到低排。第二件是生成引导：在扩散的每一步，把当前生成状态和蛋白一起喂进去，拿到一份“现在这个分子结合得好不好”的稠密信号，作为特征融合进去噪网络。

注意：原文有一句话容易被跳过——PMINet 只关注蛋白质与配体之间的（对级）相互作用建模，并不显式建模 3D 口袋几何，但 IRDiff 仍然能拿它做 3D 生成。这个“看起来不匹配”的地方恰恰是设计的巧妙处：生成任务真正缺的不是几何精度，而是化学层面的方向感，比如哪些基团能在这个口袋里做氢键供体、哪个位置该放芳环。PMINet 学到的正是这类东西。

2.3 预训练模型和扩散模型是两套参数，各管一段

这里有个工程上很容易混淆的点：PMINet 和扩散主干是分开的两套权重。采样日志里能看到类似 “Restored from ./pretrained_models/PMINet with 0 missing and 0 unexpected” 的行，说明 PMINet 是独立加载的，它不参与扩散主干的梯度更新，也不和主干共享层。

这种“冻结的裁判 + 可训练的选手”结构带来两个直接后果。一是训练阶段自增强分支的梯度不会回传到 PMINet，预训练学到的东西不会被生成任务带偏；二是微调自定义数据集时，如果只换扩散主干而保留同一份 PMINet，那么“亲和力判断标准”是不变的，变的只是生成分布。这一点在做迁移实验时很关键，否则你会分不清性能变化到底来自哪一侧。

3. 检索增强：为什么外部参考配体比模型自己更懂目标口袋

3.1 检索流程：128 个配体怎么变成 1 个参考

检索池不大：从训练集里随机抽 128 个配体，对每个目标蛋白用 PMINet 逐个打分，按预测亲和力取最高的 k 个当参考，默认 k=1。这个池子在数据准备阶段就算好了，落盘成test_prompt_ligand_indices_top3.pt这类索引文件，采样脚本用--test_prompt_indices_path指过去，跑的时候不重新检索。

# 检索阶段的等价逻辑，实际实现是离线批处理 pool = random.sample(train_ligands, n) # n = 128，库大小 scores = [pminet(protein_xyz, lig_xyz) for lig_xyz in pool] # 逐对打分 order = sorted(range(len(pool)), key=lambda i: scores[i]) refs = [pool[i] for i in order[-k:]] # k 默认取 1 save_indices(refs, "test_prompt_ligand_indices.pt")

n和k是这套机制里最需要调的两个数。文件名里的 “top3” 是历史命名，不代表实际取 3 个。消融实验给出的结论很明确：加大n有收益，加大k没收益甚至有害。原因不玄学——库越大，能覆盖到的有效子结构越多样；参考越多，噪声源也越多，多个参考之间的相互作用上下文会互相打架，模型反而不知道听谁的。

3.2 参考特征怎么灌进扩散：交叉注意力的接法

融合发生在特征层，不在坐标层。生成分子的隐状态作为 query，去 attend 参考配体特征和参考复合物里的蛋白特征，得到增强后的表示，再进主干去噪。这里用的是可训练的交叉注意力，不是简单拼接，意味着模型可以自己学“哪一部分参考信息值得抄”。

# 扩散第 t 步的增强融合示意 h_gen_t = denoiser.encode(x_t, atom_types) # 当前生成分子编码 h_prot = protein_encoder(protein_xyz) # 目标蛋白原子特征 h_ref_l = ligand_encoder(ref_ligand_xyz) # 参考配体原子特征 h_ref_p = protein_encoder(ref_protein_xyz) # 参考复合物里的蛋白特征 # 检索增强：生成特征去 attend 参考的配体+蛋白上下文 h_gen_t = cross_attn(q=h_gen_t, kv=concat(h_ref_l, h_ref_p)) # 自增强：再去 attend 当前状态与目标蛋白的相互作用嵌入 h_pmi = pminet.encode(protein_xyz, x_t) h_gen_t = h_gen_t + cross_attn(q=h_gen_t, kv=h_pmi)

两段 attention 的顺序不是随便排的。检索增强提供的是“外部经验”，自增强提供的是“当前状态的自评”，先拿外部经验把方向定下来，再让自评做局部修正，这个先后关系在消融里体现为：单独用任何一个都比基线好，两个一起用达到最佳结合相关指标。

3.3 反向验证：低亲和力参考会把模型带偏

证明检索增强有效的办法，不是看用上参考后涨了多少，而是看故意用错参考会掉多少。实验把参考替换成 PMINet 排名最低的配体，低亲和力组的平均预测亲和力约 3.95，高亲和力组约 7.24，其余设置不变，结合相关指标全面变差。

这个反证比正向增益更有说服力。它说明增强机制本身不是“多给信息总没坏处”的无害操作——扩散网络在每一步都拿参考特征做条件，参考的质量直接改变了生成分布的中心。参考选错了，模型会非常认真地朝着错误方向生成。

还有个更细的观察：生成分子与对应参考分子的指纹相似度，明显高于它与数据库里随机分子的相似度。把生成分子和参考分子摆一起看，能看到相同的关键子片段被保留下来、重新组合到不同位置。这基本坐实了参考的角色——不是模板复制，而是子结构级别的“借用”。

4. 自增强：生成中途把自己喂回 PMINet，等变性为什么没被破坏

4.1 训练与采样的时序

训练阶段，先把蛋白质原子质心平移到零点，从完整分子出发按扩散时间表加噪到任意 t，把 t+1 步的小分子和对应蛋白一起送进 PMINet 拿到双向嵌入，再和扩散模型编码的对应特征融合，最后让网络预测原始特征（去噪目标）。整个过程里蛋白坐标一路保持不变，只有配体在动。

采样阶段反过来：输入给定蛋白的结合位点、外部配体库和参考池大小，先采样配体的原子数目，蛋白质心归零，初始化配体原子坐标与原子类型，提取蛋白原子级嵌入，然后进入迭代更新循环，每一轮同时走检索增强和自增强。

# 训练一步：加噪 + 双增强 + 去噪（示意） t = randint(0, T) x_t = q_sample(x_0, t) # 按扩散表加噪到 t x_t = x_t - x_t.mean(dim=0) # 配体坐标整体平移，配合蛋白质心归零 eps_pred = backbone(x_t, protein_xyz, t) # 主干去噪预测 h_pmi = pminet.encode(protein_xyz, x_t) # 当前状态与蛋白的相互作用嵌入 eps_pred = backbone.fuse(eps_pred, h_pmi) # 特征层融合后才输出 loss = mse(eps_pred, eps_true)

采样耗时可以参考内置案例：batch_size 设 25，显存占用约 19 GB，采 100 个分子、1000 步扩散，整体约 50 分钟。这个量级决定了做消融时不能一把梭，得先把num_samples和num_steps压下来做小规模对比。

4.2 等变性守恒靠的是两条硬约束

第一条是蛋白原子坐标已知且在更新过程中保持不变。蛋白是条件，不是被生成的对象，它的坐标一旦固定，整个系统的参考系就固定了。

第二条是蛋白原子质心平移为零，加上 EGNN 本身的设计，保证逆变换下的 SE(3) 等变性。直白点说，整个网络对“整体旋转和平移”是等变的：你把输入转一下，输出的坐标也跟着转一下，不会变成别的分子。

关键在于增强只作用于 SE(3)-不变的隐藏状态。融合发生在特征层，融合的对象是标量化的相互作用嵌入，坐标更新那条路径完全没被碰。所以加了多少层增强，等变性都不受影响——这是一个结构性的保证，不是靠训练调出来的经验性质。

4.3 一个反直觉的消融：没有预训练的增强会掉点

把增强机制保留、但去掉预训练的 PMINet，生成性能不升反降。这个结果值得多看一眼：说明增强模块本身不是架构魔法，它只是一条通道，通道里流的东西来自预训练。通道建好了但里面是随机初始化，等于每一步都在给去噪网络灌噪声。

反过来，自增强单独用就能显著提升结合相关和属性相关指标，因为此时通道里有真东西——预训练 PMINet 带来的相互作用知识。这也是为什么复现时第一件事是确认 PMINet 权重加载成功，日志里 “0 missing and 0 unexpected” 这行不能跳过。

4.4 训练和采样共享同一套增强路径

两个阶段的差异只在数据来源：训练时 PMINet 看到的是加噪后的真实分子，采样时看到的是模型自己生成到一半的分子。这种“训练-采样一致性”是自增强能奏效的前提。如果训练时喂真实分子、采样时喂生成分子，两者分布差太远，PMINet 给出的嵌入在采样后半段就会失真，生成分子越接近完成、指导信号反而越不可靠。设计上让两阶段走同一条路径，就是在压这个分布差。

5. 复现排查：从 conda 依赖死锁到构象文件丢失的五处坑

5.1 cryptography 和 openssl 的版本死锁

现象：执行conda env create -f irdiff.yml，解析器在LibMambaUnsatisfiableError上直接停住，报cryptography-37.0.1需要openssl <1.1.2a，而 yml 里锁了openssl ==3.1.0，两者不可兼容。

原因：这份 yml 是导出环境时生成的，把一堆间接依赖的精确版本也钉死了。导出的机器上这些版本恰好共存，换一台机器重建时，解析器发现约束集本身自相矛盾。

解决：把 yml 里cryptography的版本号去掉，让它跟着 openssl 走。

# 修改前 - cryptography=37.0.1=py38h9ce1e76_0 # 修改后：只留包名，版本交给解析器 - cryptography

5.2 autodocktools-py3 在 pip 上找不到对应版本

现象：conda 依赖装完之后进入 pip 阶段，报Could not find a version that satisfies the requirement autodocktools-py3==0+unknown，随后CondaEnvException: Pip failed。

原因：0+unknown是本地构建时 setuptools 写进环境快照的伪版本号，公共索引上根本没有这个版本。这是导出环境时最典型的“本地状态泄漏”。

解决：跳过 yml 里的 pip 段，进虚拟环境从上游源码仓库直接构建安装。

conda activate IRDiff # 从源码构建 wheel 并安装，绕开索引上不存在的伪版本号 python -m pip install "git+<该工具上游源码仓库地址>" # 成功时末尾会显示 Successfully installed AutoDockTools-py3-1.5.7.post1+...

剩下几个对接相关的包也手动补齐，避免 yml 里同样的伪版本问题。

pip install meeko==0.1.dev3 pdb2pqr==3.6.1 vina==1.2.2 propka==3.5.0

5.3 einops 缺失和 pyg 版本错配

现象一：跑sample_split.py直接抛ModuleNotFoundError: No module named 'einops'，栈顶指向models/attn.py第 4 行的 import。

原因：yml 漏列了这个依赖，但注意力模块里到处在用rearrange和repeat。属于依赖清单和代码不同步。

现象二：环境装好了、能跑，但训练阶段出奇怪的问题。

原因：pyg版本不对。这个项目对 pyg 的版本敏感，装 2.0.4 之外的版本，训练时会出问题，而且报错位置往往离真正原因很远。

解决：补装 einops，并把 pyg 钉在 2.0.4。

conda install einops conda install pyg=2.0.4 -c pyg -c conda-forge

注意：pyg 这类带编译扩展的包，装完之后建议立刻用python -c "import torch_geometric; print(torch_geometric.__version__)"确认一遍实际生效版本，别只看 conda 的回显。

5.4 写死的绝对路径

现象：数据加载阶段报文件不存在，但明明已经把数据放进./data。

原因：仓库里保留了/path/to/...这种模板占位路径，共两处，一处是评估脚本的输出目录，一处是训练配置的数据路径。

解决：全部改成相对路径。

# configs/training.yml 第 4-5 行 path: ./data/crossdocked_v1.1_rmsd1.0 split: ./data/crossdocked_pocket10_pose_split.pt

# cal_metrics_from_pt.py 第 8 行 eval_path = './eval_results/'

数据侧只需要两个文件：crossdocked_pocket10_pose_split.pt和crossdocked_v1.1_rmsd1.0_pocket10_processed_final.lmdb，都放./data下。预处理流程和 TargetDiff 一致，不需要自己重跑。

5.5 评估脚本不落盘构象，指标算完分子就没了

现象：eval_split.py跑完，./eval_results/里只有metrics_*.pt，找不到任何分子文件，没法做对接复算，也没法目视检查口袋里的姿态。

原因：原脚本在if args.save分支里只torch.save了指标字典（validity、键长分布、全部结果），生成分子的 Mol 对象没有被写出。这是一个纯功能性缺失，不是 bug。

解决：在保存分支末尾追加 SDF 写出，把构象落到./eval_results/Generated_molecules.sdf。

if args.save: torch.save({ 'stability': validity_dict, 'bond_length': all_bond_dist, 'all_results': results }, os.path.join(result_path, f'metrics_{args.eval_step}_{eval_step}.pt')) # 追加：把生成的 RDKit Mol 写成 SDF，供后续对接和目视检查 from rdkit import Chem writer = Chem.SDWriter(os.path.join(result_path, 'Generated_molecules.sdf')) for mol in results['mol_list']: # 字段名按实际 results 结构调整 if mol is not None: writer.write(mol) writer.close()

这一步不是可选项。要算vina_score、vina_docking_score、qvina_score这些对接指标，或者想拿生成分子去微调，没有 SDF 就等于把一轮采样的结果全扔了。

6. 设计取舍：参考给多少、信多少，才是这套增强的边界

这套机制最容易被误读的地方，是把“检索增强”当成一个越大越好的旋钮。消融给的是相反答案：n从 128 往上加，结合指标会涨，因为库变大意味着覆盖到的有效子结构更杂更全；但k往上加，收益很快耗尽甚至转负。这不是调参玄学，是这套方法的定位决定的——它借的是“局部化学经验”，不是“多模板平均”。多个参考一起上，相互作用上下文互相稀释，模型反而拿不到明确方向。

第二个取舍是亲和力和类药性之间的平衡。IRDiff 在 CrossDocked 2020 上拿到平均 Vina 约 -6.03、平均 QED 约 0.53，相比 Pocket2Mol 在平均和中位 Vina 上分别高出约 17.3% 和 46.6%，相比 DecompDiff 分别高出约 6.3% 和 14.1%；高亲和力比例达到 67.4%，也就是三分之二以上的生成分子在测试集上比真实配体结合得更好。但 QED 和 SA 并不是它的强项。原作者的态度很直接：这两个指标在实际药物发现流程里只是粗筛，落在合理区间即可，不值得为它们牺牲亲和力。这个判断在实际项目里是否成立，取决于你处在筛选管线的前段还是后段——前段要多样性，后段要成药性，权重完全不同。

第三个取舍藏在那个反直觉消融里：增强机制脱离预训练 PMINet 会掉点。这意味着整套方法的性能上限被 PMINet 的预训练质量死死绑住。换一个靶点家族、换一个化学空间，先确认 PMINet 在那类复合物上的亲和力预测还靠谱，再谈生成效果；否则增强通道里流的可能是错的知识，而且它以“每一步都在融合”的方式持续放大错误。

想自己验证边界，我一般会这么做：固定k=1，把n从 32、64、128 三档扫一遍，每档采同样数量的分子，看平均 Vina 和 QED 的两条曲线是不是一个涨一个平；再把参考替换成同库里的随机配体做对照，如果指标掉得不多，说明这个靶点上检索增强本来就没多少可借的东西。这套对照不用跑满 1000 步，几百步就够看出趋势。

血泪经验是：复现这类工作，环境跑通只是入场券，真正的坑在“指标算完了但分子没了”和“参考索引和当前数据集不匹配”这两件事上。前者我在 5.5 里补了 SDF 写出，后者没有报错、只会让结果莫名其妙地差。从那以后我每次跑采样前都强制走一遍检查清单——PMINet 权重加载日志、参考索引文件与当前 split 是否对应、评估脚本是否落盘构象，三样缺一样就不开跑。希望帮到你。

本文还有配套的精品资源，点击获取