Interpreting GFlowNets for Drug Discovery: What probes can and cannot show

技术摘要：解释 GFlowNets 在药物发现中的应用

问题陈述
生成式流网络（GFlowNets），特别是针对分子设计进行训练的 SynFlowNet，提供了一个在遵循合成约束的同时探索化学空间的强大框架。然而，其内部决策策略仍然是不透明的，这限制了它们在药物发现中的应用，因为药物化学家需要具有可解释性的依据来支持所提出的结构。现有的解释性方法通常是为连续潜空间或基于似然的模型设计的，而非针对 GFlowNet 的离散、序列化、基于图的策略。此外，亟需区分策略中真正学习到的化学知识与通过图架构和原子特征化即可获得的平凡信息。

方法论
作者引入了一个多尺度解释性框架，应用于使用特定检查点冻结嵌入（frozen embeddings）的 SynFlowNet（使用 QED 奖励函数训练）。该框架整合了三种互补的方法：

1. 基于梯度的显著性与反事实分析： 应用集成梯度（Integrated Gradients, IG）于“停止（Stop）”动作的对数概率，以生成原子级显著性图。高显著性原子被分组为基团（连通分量和环系统）。随后使用 RDKit 转换规则（如卤素交换、酰胺修饰）对这些基团进行反事实扰动，以测量 QED 的变化（$\Delta$QED）。
2. 稀疏自编码器 (SAEs)：
   • 欠完备型 (MLP-SAE)： 使用压缩自编码器（256 $\to$ 128）来识别主导的、稀疏的潜在因子，并测试其与理化性质（如极性、亲脂性）的相关性。
   • 过完备型 (BatchTopK SAE)： 使用带有 BatchTopK 稀疏机制（$k=64$）的过完备自编码器（256 $\to$ 1024），旨在不使用 $\ell_1$ 正则化的情况下恢复细粒度的特征检测器，包括处于叠加态（superposition）中的特征。
3. 探测与控制实验： 在嵌入上训练浅层前馈分类器和严格的线性探针（岭回归/逻辑回归），以预测化学基团（官能团、卤素）及理化性质的存在。至关重要的是，本研究采用了严谨的控制实验：
   • 描述符基准： 将探针性能与 RDKit 描述符和 ECFP4 指纹进行比较。
   • 随机标签控制： 确保探针不会记忆噪声。
   • 骨架不相交拆分 (Scaffold-Disjoint Splits)： 防止训练/测试拆分中的骨架泄漏。
   • 未训练网络控制： 在具有相同架构的随机初始化网络上运行完全相同的探针，以分离训练与架构的贡献。
   • 特征消融： 通过将单个特征置零来测试下游解码的特异性。

关键结果

• 显著性与反事实分析： 集成梯度突出了具有化学意义的区域（极性取代基、芳香片段）。对这些基团进行反事实编辑会产生微小但可测量的 QED 偏移（例如，+0.015 至 +0.051），尽管相对于 QED 的范围而言效应值较小，且针对随机编辑的受控比较仍是未来的工作。
• 潜在结构： 欠完备型 SAE 识别出与尺寸和极性强相关的因子，但使用这些因子的 MLP 预测器在复合 QED 上的表现较差（$R^2=0.25$），相比于单一性质而言。
• 重构与稀疏性： 过完备型 BatchTopK SAE 以低 NMSE（0.0015）重构嵌入，并保持严格的批均稀疏度（活跃特征数为 64）。在相似的活跃水平下，其表现优于匹配的线性基准（NMSE 为 0.0026）。
• 线性可解码性： 嵌入能够很好地解码理化性质（例如，QED $R^2 \approx 0.59$，尺寸 $R^2 \approx 0.996$）。然而，在随机拆分下，这些得分并不显著优于通过 RDKit 描述符基准实现的得分。
• 未训练网络控制（决定性发现）： 一个具有相同架构的随机初始化网络，其解码每个被探测属性的能力几乎与训练后的策略一样好。对于核心属性（类药性/QED），训练带来的增益微乎其微，仅为 +0.005（未训练 0.581 vs. 训练后 0.586），这处于随机种子的方差范围内。对于分子尺寸，未训练的网络反而略好。
• 特征特异性： 在过完备字典中，单次运行揭示了具有化学富集性的检测器（例如，分别针对碘、氯、溴、氟和硼的特征）。单特征消融显示了属性特异性的影响（特异性比率为 8–23 倍），但特征索引在不同种子间不可复现；只有一小部分核心方向是稳定的。

意义与主张
本文最重大的结论是一个负向控制实验：从 GFlowNet 嵌入中解码出的理化性质反映了图架构和原子特征化，而非学习到的化学知识。

• 未得到支持的内容： 高探针得分（即使接近完美的 AUROC）并非获得化学理解或学习到的结构-活性推理的证据。解码属性的能力主要是输入表示和网络架构的人为产物，这一点已通过未训练网络匹配训练后策略的表现得到了证实。
• 得到支持的内容：
  1. 过完备字典提供了嵌入的有效、精确稀疏分解。
  2. 单次运行包含具有化学富集性的子结构检测器（如每种卤素的特征），这些检测器在多个种子间是稳定的。
  3. 这些特征对探针解码具有属性特异性的干预效应。
  4. 类药性（QED）在稀疏潜在空间中比在稠密嵌入中更具线性可及性。

结论
作者得出结论，解释性方法必须针对架构基准进行严格控制。对于 SynFlowNet，该框架成功识别了细粒度的特征检测器，并验证了稀疏表示的鲁棒性。然而，本研究明确警告不要将高可解码性得分解释为已习得化学知识的证明。论文将其主张限制在字典中存在一个可复现的稳定且具有化学富集性的核心特征，同时断言模型预测理化性质的能力主要是其架构而非训练的功能。未来的工作应探索这些检测器在 SAE 针对未训练网络嵌入进行训练时是否依然存在，并将该框架扩展到多目标设置。