第四章 · 4.6

4.6 GraphAF:自回归流的分子生成

GraphAF: Autoregressive Flows for Molecules
GraphAF 将 4.2 节的自回归分解与 4.5 节的流模型结合:根据已生成子图,把高斯噪声变换为连续向量,再解码为原子或键类别。本节说明节点流、边流和价约束检查的作用,区分连续密度的精确计算与离散分子概率的近似学习,介绍 ZINC250k 上的极大似然预训练和 PPO 性质微调。最后依据原论文解读有效性、性质分数与训练时间,并比较 GraphAF 和 GCPN 的训练机制及适用边界。

4.6.1 自回归分解与条件流

4.4 节的 GCPN 用策略网络输出离散动作概率。GraphAF 则以条件可逆变换(conditional invertible transformation)构造连续条件密度,再把连续样本解码为原子和键。可逆性成立于基础噪声与连续表示之间;最后的类别解码不可逆。模型中的 autoregressive 对应 4.2 节的链式分解,flow 对应 4.5 节的变量变换公式,两者共同构成自回归流(autoregressive flow)(Shi et al., 2020)。

先考虑离散生成过程。固定生成顺序,将分子表示 x 写成原子与键的序列;生成其 n 个原子及相应键的概率可作如下分解:

\[p_{\theta}(x) = \prod_{i=1}^{n} \left[p(a_{i} | g_{<i}) \cdot \prod_{j=1}^{i-1} p(b_{ij} | g_{<i}, a_{i}, b_{i,<j})\right]\]
(4.6-1)ai 为第 i 个原子的类别,bij 为新原子与旧节点 j 的键类别(含“无键”),g<i 为已生成子图。键项还以此前的键 bi,<j 为条件。式中省略了终止事件;若计算完整变长轨迹的概率,还需纳入终止项,见习题 4.6-1。这里是有序表示的链式分解,不能直接等同于对所有生成顺序求和后的无序图概率。

GraphAF 与 MAF(masked autoregressive flow,掩码自回归流)都要求当前变换的参数只依赖此前的变量。MAF 用网络掩码限制这种依赖;GraphAF 用 RGCN(关系图卷积网络)编码当前子图,以图表示预测节点流参数,以图表示和两端节点表示预测边流参数。训练时,GraphAF 同样使用掩码,保证各个条件分布只读取对应的历史子图。以下用 xi 统一表示某一步的连续原子或键向量,用 hi 表示该步的条件编码;其密度由变量变换公式计算:

\[\log p_{\theta}(x_{i} | x_{<i}) = \log p_{Z}(\varepsilon_{i}) + \log \left|\det \frac{\partial \varepsilon_{i}}{\partial x_{i}}\right| \quad \varepsilon_{i} \sim N(0, I) \quad x_{i} = f_{\theta}(\varepsilon_{i}; h_{i})\]
(4.6-2)pZ 是基础分布(base distribution)的密度,GraphAF 采用标准高斯。原论文用条件逐维仿射变换 xi = μ(hi) + σ(hi) ⊙ εi,其中各缩放分量为正。给定历史后,当前向量内部的雅可比为对角矩阵;按生成顺序排列全部变量时,雅可比为块三角矩阵。变量变换公式(change of variables)的行列式因此可由对角元之积计算,乘积成本为 O(d),不包含编码器的计算成本(习题 4.6-2)。
定义

自回归流(autoregressive flow)。由条件可逆变换 xi = fθ(εi; hi) 构成的生成模型,其中 εi 独立采自基础分布,hi 只编码此前的变量。噪声到数据的方向按序采样;数据到噪声的方向在观测历史已知时计算各步逆变换。链式分解与变量变换公式给出连续数据的精确密度。GraphAF 将这一机制用于分子图的连续表示,并通过图编码器表达结构条件。

4.5 节介绍了可逆变换与密度计算的关系(Rezende & Mohamed, 2015);更完整的模型分类见 Papamakarios 等人的综述(Papamakarios et al., 2021)。VAE 通常通过证据下界训练,边缘似然一般需要近似估计;离散自回归策略则可以直接计算给定动作轨迹的概率。流的特点是对连续变量使用可逆映射和显式密度,不能据此断言其他模型无法评估概率,也不能把轨迹概率、连续密度和无序分子图概率混为一谈。

原子类型和键型是类别变量,连续流不能直接作用于这些离散取值。GraphAF 使用去量化(dequantization):训练时给 one-hot 类别向量加上 [0, 1)d 内的均匀噪声,再拟合所得连续向量;采样时通过流生成连续向量,以 argmax 解码为类别(Dinh et al., 2017;Kingma & Dhariwal, 2018)。

\[\text{训练侧:}x_{i} = \operatorname{one\text{-}hot}(a_{i}) + u \quad u \sim U[0, 1)^{d}; \text{采样侧:}a_{i} = \operatorname{argmax}_{k} [ f_{\theta}(\varepsilon_{i}; h_{i}) ]_{k}\]
(4.6-3)xi 为连续向量,d 为类别数:ZINC250k 的原子向量为 9 维,键向量为 4 维,即 3 类化学键加上“无键”。对训练样本,原 one-hot 中的 1 落入 [1, 2),其他分量落入 [0, 1),因此 argmax 能恢复原类别。任意生成向量也能解码为类别,但这不保证其化学有效性。均匀去量化的期望对数密度是离散对数似然的下界;式 (4.6-2) 的精确性只针对连续密度,见 GraphAF 论文附录 A。

4.6.2 单步生成:节点、边与价约束

设已生成的子图 g<i 含 i−1 个原子及其间的键。新增一个原子的过程包含以下五步。

第一步,编码。RGCN 根据当前子图计算节点表示,经读出操作得到图表示 hi,用于确定节点流的条件参数。

第二步,节点流(node flow)。从 N(0, I) 采样 εi,按 hi 决定的平移和缩放做条件仿射变换,得到 9 维连续向量,再以 argmax 得到原子类别 ai。TorchDrug 的实现可堆叠 L 层条件流,具体结构见 4.6.6。

第三步,边流(edge flow)。依次生成新原子与候选旧节点之间的键。每条键以 g<i、ai 和已生成的键 bi,<j 为条件;编码器随当前部分图更新,边流还使用两端节点的表示。例如,碳原子已有一条双键时,剩余可用价会影响后续键的选择。这说明同一新原子的各条键通常不能按条件独立处理。

若考察所有旧节点,第 i 个原子需要 i−1 次键判定,整张分子共需 n(n−1)/2 次。原论文还采用广度优先搜索(BFS)排序,并在所用数据中只建模新原子与最近 12 个原子之间的候选边,以缩短依赖范围。这个窗口依赖数据和排序规则,不能推广为所有分子图的性质。原子与键分步生成缩小了单次动作的类别数,但总成本仍取决于候选边数、编码次数和重采样次数。

第四步,价约束检查。每生成一条候选键,核对两端原子的累计键级是否超出允许价。若违反约束,就拒绝该键,重新采样基础噪声并生成另一键型。这是采样后的检查与重采样;用于并行训练的历史掩码承担不同功能。

第五步,终止与更新。除首原子外,若新原子与旧子图没有形成任何键,就结束生成并保留此前的连通分子;达到原子数上限时也停止。否则将新原子及其键加入子图,继续下一轮。首原子从空图生成,不受“必须连接旧子图”这一条件限制。

GraphAF 单步数据流:编码得条件,节点流采原子类别,边流按序采键,价检查拒绝重采,最后决定终止或更新 ① 编码(条件的来源) 已生成子图 g C C N RGCN 编码器 h = RGCN(g) 条件 h ② 节点流:采样原子类别 ε ~ N(0, I) 基础分布噪声 条件流层 f(·; h) 条件仿射变换 × L 层 连续向量 9 维(ZINC250k) argmax 译码 N ③ 边流 + ④ 价约束检查:按序采键 对每个旧节点 j 按序 条件:g、新原子、已采各键 键 1 键 2 … 价约束检查 违价 → 拒绝该键,重采噪声 合法键写入子图 累计价同步更新 ⑤ 终止与更新 新原子无键相连 / 原子数达上限 → 生成终止,输出分子 子图更新:g ← g ∪ 新原子 回到编码(下一轮) 回到 ① 每轮先生成原子,再按序生成候选键;价检查在每次键采样后执行。 连续向量由条件流生成,argmax 转为类别;训练掩码与采样时的价检查作用不同。 虚线表示条件编码的传递。论文 ZINC250k 结果:开启检查为 100%,关闭约 68%(4.6.4)。
图 4.6-1 GraphAF 的一轮生成。RGCN 将当前子图编码为条件表示;节点流生成连续向量并解码出原子类别;边流根据更新后的部分图和两端节点表示,逐次生成候选键,价检查拒绝违反约束的键并触发重采样。随后判断是否终止,或更新子图继续生成。图中的 L 层条件流对应 TorchDrug 的堆叠实现;原论文的基本变换是条件逐维仿射变换。“无键相连”的终止条件从第二个原子起适用。

4.2 节式 (4.2-5) 讨论的是另一种约束实现:采样前把非法动作的概率置零并重新归一化。若合法动作集合固定,且拒绝采样能够持续到接受为止,两者都会得到原候选分布在合法集合上的条件分布;主要差别在计算方式和重试成本。实际实现还可能设置重试上限。价检查可以排除指定的价违反,却不保证合成可行性或成药性;论文中的有效性数字应连同评估规则理解。

自回归依赖对训练和采样的影响不同。教师强制(teacher forcing)训练时,全部历史原子和键都已由数据给出,模型可用掩码构造各步对应的子图,并行计算条件密度。并行化减少了串行等待,但仍需要处理这些子图,不能视为没有计算成本。采样时,后续条件必须等待前面的原子和键生成,因而保留顺序依赖。

方法

GraphAF 单步采样循环。① RGCN 编码 g<i;② 采样 εi ~ N(0, I),经条件流和 argmax 得到 ai;③ 根据新原子和已生成的键更新条件,逐次采样候选键;④ 若候选键违反价约束,拒绝并重采;⑤ 除首原子外,若新原子没有连接旧子图,则舍弃它并终止;达到原子数上限时也终止,否则继续生成。启用约束后,接受的样本分布已由合法集合重新条件化,应与未经检查的原始采样分布区分。

习题 4.6-1

乙醇的重原子图为 C–C–O(隐氢)。取生成序:a1 = 端碳,a2 = 中碳,a3 = 氧;键:b21 = 单键,b31 = 无键,b32 = 单键。

  1. 按式 (4.6-1) 写出 p(x) 的链式分解,逐项指明哪些来自节点流、哪些来自边流,并说明“终止”以什么事件的概率进入分解。
  2. 设第 3 步(氧)的连续向量某一维由一维仿射条件变换给出:x = μ + σ·ε,ε ~ N(0, 1),编码器输出 μ = 1.2、σ = 2。观测到该维取值 x = 2.2,手算 log p(x | h)。(已知 ln 2π ≈ 1.8379,ln 2 ≈ 0.6931,e−0.125 ≈ 0.8825)
  3. 解释在一般分子图中,b32 的条件为何需要包含 b31。
参考解答

(1) 记 gk 为已完成前 k 个原子的子图。该轨迹的分解为 p(x) = p(a1) · p(a2 | g1) p(b21 | g1, a2) · p(a3 | g2) p(b31 | g2, a3) p(b32 | g2, a3, b31) · p(停 | g3)。其中有三个节点项和三个键项;首原子没有边。尚未达到大小上限时,终止对应下一候选原子与当前子图的所有候选键均为“无键”的事件,其概率还需对该候选原子的类别求和。流先定义连续密度,再通过解码诱导这些类别概率;连续向量上的密度值不能直接代替类别的概率质量。

(2) 反解 ε = (x − μ)/σ = (2.2 − 1.2)/2 = 0.5。标准正态密度 φ(ε) = (2π)−1/2 e−ε²/2,log φ(0.5) = −0.5 × 1.8379 − 0.125 = −1.0440。正向雅可比为 |dx/dε| = σ = 2,逆向雅可比为 |dε/dx| = 1/2。代入式 (4.6-2):log p(x | h) = −1.0440 − 0.6931 = −1.7371,密度约 e−1.737 ≈ 0.176。σ = 2 使分布的尺度扩大,因而对应点的密度等于 φ(0.5) 的一半,对数密度减去 log 2。

(3) b31 与 b32 共享氧原子。若氧与端碳已经成键,其剩余可用价便会减少,进而限制与中碳的键型。乙醇这一条观测轨迹中 b31 为“无键”,但生成模型还必须处理它取其他值的可能状态。将 b31 纳入条件,允许模型表达这种键间相关性,以及已形成子图对后续连接方式的影响。

4.6.3 训练两阶段:先极大似然,后策略微调

GraphAF 先学习训练分子的分布,再用目标性质微调,对应 4.1 节介绍的两类生成目标。

第一阶段在 ZINC250k 上进行极大似然预训练。该数据含约 25 万个类药分子,每个分子至多有 38 个重原子,使用 9 类原子和 3 类化学键。训练沿给定的 BFS 顺序构造历史子图,经掩码并行计算去量化表示的条件密度,不需要为这个似然目标采集强化学习轨迹或训练判别器。论文的密度建模实验采用 10 轮、batch 32、Adam 和 0.001 的学习率;性质优化实验则先预训练 300 轮,再以 PPO 微调。数据中的最大分子大小与生成上限应分别记录,论文的模型生成上限设为 48。

\[\max_{\theta} \sum_{\tilde{x} \in \tilde{D}} \log p_{\theta}(\tilde{x}) = \max_{\theta} \sum_{\tilde{x} \in \tilde{D}} \sum_{i} \left[\log p_{Z}(\varepsilon_{i}) + \log |\det J_{i}|\right]\]
(4.6-4)这里的 D̃ 是一次去量化采样得到的连续训练数据,i 遍历该生成顺序下的节点和键变量,Ji = ∂εi/∂xi 为逆变换雅可比。训练可重新采样去量化噪声,以估计期望目标。精确对数似然(exact log-likelihood)指给定连续表示和生成顺序后的密度对数,不是对解码区域或所有生成顺序积分、求和后的分子概率。

GCPN 也从真实分子中学习先验信息。其原论文用专家状态—动作对最小化 −log πθ(at | st),并在训练中保留这一监督目标,同时优化包含性质和对抗奖励的强化学习目标。因此,两者的区别不能概括为“极大似然与纯回报”。GraphAF 显式拟合去量化连续密度后再做性质微调;GCPN 则结合专家动作监督、对抗学习和策略优化,见 GCPN 原论文第 3.5 节及实验设置。

第二阶段采用近端策略优化(proximal policy optimization, PPO)(Schulman et al., 2017)。奖励包括两部分:中间奖励(intermediate reward)惩罚违反价约束的候选键;最终奖励(final reward)由目标性质奖励与化学过滤罚项组成,并按折扣因子分配给此前各步。以下写出用于说明奖励分配的形式,PPO 实际优化的是基于优势估计和概率比裁剪的代理目标:

\[R(x)=s(x)-\text{过滤罚项},\qquad \widehat{A}_{t}=-\lambda\,\mathbf{1}[\text{第 }t\text{ 步违价}]+\gamma^{T-t}R(x)-c_{t},\quad 1\leq t\leq T\]
(4.6-5)γ 为折扣因子,ct 为用于降低估计方差的逐步移动平均基线;式中省略了节点动作与边动作的分别记号。过滤罚项包含位阻和官能团过滤惩罚。s(x) 是经变换后的性质奖励:论文对 QED 使用系数缩放,对 penalized LogP 使用带温度参数的指数变换。策略梯度通过动作对数密度对模型参数求导,不穿过 argmax 或 RDKit 性质计算。

连续流可导,并不意味着离散分子的性质评分可以直接反向传播。GraphAF 的微调通过采样估计策略梯度,并使用 PPO 约束单次更新幅度。预训练提供参数初值,条件密度提供策略优化所需的可微计算;这些作用仍不能保证目标分数提高时分布质量保持不变。连续对数似然也不足以单独诊断类药性,应同时检查有效性、唯一性、新颖性、内部多样性及与训练分布的距离。似然目标拟合给定数据,奖励目标偏向指定性质,两者都受模型、数据和目标定义的限制。

似然信号(预训练)

  • 目标:拟合训练分子的连续表示
  • 来源:去量化样本的对数密度(式 4.6-4)
  • 口径:连续密度精确可算,离散似然采用下界
  • 评估:连接模式、片段分布和样本多样性

回报信号(微调)

  • 目标:提高指定性质的奖励
  • 来源:价违反惩罚与折扣后的终局奖励(式 4.6-5)
  • 优化:基线降低估计方差,PPO 裁剪限制更新幅度
  • 评估:性质改善及其对分布质量的影响
GraphAF 两阶段训练:先拟合去量化分子表示的密度,再以 PPO 优化目标性质;下方对照 GCPN 的专家监督与强化学习 两类训练目标:预训练拟合分子数据的分布,微调提高指定性质的奖励 阶段一 预训练:极大似然 distribution learning(式 4.6-4) 数据:ZINC250k,约 25 万类药分子 信号:去量化表示的连续对数密度 并行:历史子图掩码与条件密度计算 预算:300 轮,batch 32,Adam 产出:拟合训练数据的生成模型 参数作为初始化 预训练参数 阶段二 微调:策略优化 goal-directed(式 4.6-5,PPO) 算法:PPO 策略梯度 中间奖励:价违反惩罚 最终奖励:性质奖励 − 化学过滤罚项 折扣:终局奖励分配给此前各步 梯度:动作对数密度对参数求导 GCPN 对照:专家动作监督、对抗学习与策略优化 专家目标:动作负对数似然;强化学习:性质、有效性与对抗奖励。论文在训练中同时保留。 比较重点:分布的参数化、监督数据的使用方式,以及预训练与性质优化的衔接。
图 4.6-2 GraphAF 的两阶段训练。第一阶段拟合 ZINC250k 分子的去量化连续表示;第二阶段以这些参数为初值,通过 PPO 优化性质奖励,同时使用价违反和化学过滤惩罚。图中 300 轮是原论文性质优化实验的预训练预算。GCPN 的专家动作负对数似然目标与对抗奖励是不同的训练信号,原论文在强化学习过程中继续使用专家监督,不能将其预训练归为判别器回报。

4.6.4 结果解读:有效性、性质优化与效率

有效性需要区分是否启用价检查。GraphAF 论文报告,在 ZINC250k 实验中,关闭检查时 GraphAF 的有效性约为 68%,作者运行 GCPN 公开代码得到约 20%;启用检查时,两者报告的有效性均为 100%(Shi et al., 2020)。前一口径反映未经规则筛选的生成质量,后一口径同时包含模型与约束机制的作用。这一实验支持 GraphAF 在所测设置下具有较高的无检查有效性,但不足以把差异单独归因于某一种训练信号。

案例

GraphAF 论文报告的结果(ZINC250k;保留各项实验口径)。

  • 有效性(无规则约束):GraphAF 68%,GCPN 20%;(带价检查):两者均 100%。
  • penalized LogP 优化 top-3 分数:GraphAF 12.23 / 11.29 / 11.05;GCPN 7.98 / 7.85 / 7.80。
  • QED 优化 top-3:GraphAF 0.948 / 0.948 / 0.947;GCPN 0.948 / 0.947 / 0.946。按论文保留的小数位,两者结果接近;0.948 是这些实验中的高分水平,不应直接称为已证明的理论上限。
  • 达到论文表 2 所报告生成结果的训练时间(单张 V100 + 32 核 CPU):JT-VAE 约 24 小时,GCPN 约 8 小时,GraphAF 约 4 小时;这不是含 300 轮预训练与 PPO 微调的完整性质优化耗时。

GraphAF 在这组实验中的 penalized LogP 最高分明显高于 GCPN,QED 最高分则接近。penalized LogP 是经验评分,优化可能偏向高疏水性或特定结构模式;提高它不等于提高成药性。应结合五规则、合成可行性和人工结构审查解读高分分子,且这些检查各自也有适用范围。论文另在 QM9 和 MOSES 上评估了 GraphAF 的分布学习表现;这些结果检验了不同数据集上的生成能力,并非在每个数据集上重复了完整的 GCPN 对照。

top-3 是搜索到的三个最高分分子,描述极值表现,不能代替平均质量、多样性或成功率。即使两个模型的最高 QED 接近,它们达到高分的频率、所需采样量和结构分布仍可能不同。按 4.1 节的评估框架,分布指标和目标分数应分别报告,并补充样本量与搜索预算。

训练时间与生成速度也应分开。GraphAF 可在观测历史已知时并行计算条件密度,这有助于提高预训练效率;论文给出的约 4 小时对 8 小时是特定软硬件和训练配置下的实测结果,不能只由模型结构推导。采样仍需逐个生成原子和候选键。若每个新原子最多考察 P 个旧节点,键判定次数为 O(nP);若考察全部旧节点,则为 O(n²)。实际运行时间还包括图编码和拒绝重采样。

警示

数字可比的前提。GraphAF 论文的部分基线来自原始文献,部分来自作者复跑或重新评分,不能将整张表称为完全受控实验。引用时应核对数据集和划分、分子预处理、检查开关、评分脚本、样本量与训练预算。论文性质优化采用 GCPN 的评分脚本,训练时间比较使用同一计算环境;这些信息提高了相应指标的可比性,却不意味着网络宽度、生成上限和全部超参数相同。

4.6.5 GCPN 与 GraphAF 的机制对照

表 4.6-1 按原论文比较两种方法。TorchDrug 提供了统一的任务接口,但接口相近不代表两篇论文的全部实验条件一致。

表 4.6-1GCPN 与 GraphAF 的机制比较(Shi et al., 2020;You et al., 2018)
维度GCPN(4.4)GraphAF(本节)
条件分布实现图编码器与 softmax 动作分布节点、边的条件连续流,再解码为类别
预训练信号专家状态—动作对的负对数似然去量化数据的连续对数密度(式 4.6-4)
先验来源专家监督与对抗学习;训练时继续使用专家目标预训练生成模型提供微调初值
性质优化PPO,结合性质、有效性和对抗奖励PPO,结合性质、价违反与化学过滤奖励
有效性机制环境拒绝非法动作,状态保持不变拒绝违反价约束的候选键并重新采样
可计算的概率量离散动作概率及给定轨迹的概率给定顺序下连续表示的密度;离散似然用去量化下界学习
论文训练耗时比较约 8 小时,达到对应生成实验结果约 4 小时,同一硬件环境下的生成实验
主要训练部件策略网络、判别器、价值估计与分子环境图编码器、节点流、边流;微调时另用奖励、基线与环境

两种方法都基于当前分子图做顺序决策,并结合图表示、化学规则和性质奖励。但动作结构不同:GCPN 每步选择连接动作,可以新增原子,也可以连接已有原子;GraphAF 先生成一个新原子,再依次决定它的候选键。两者的原始网络尺寸和训练目标也不同。因此,表中结果反映整套方法的表现,不构成仅替换条件分布后得到的消融结论。

两种方法都需要处理规则约束、搜索范围和多样性之间的关系。可以在同一模型内部固定预训练权重、采样预算和随机种子,只改变微调奖励或约束强度,联合记录 top-3 分数、平均分、有效性及内部多样性。这样的消融实验有助于判断性质提升来自何种改变;GCPN 与 GraphAF 的参数结构不同,不能直接共用同一份权重。

习题 4.6-2

推证:(1) 证明下三角雅可比矩阵的行列式等于对角元之积,并说明已知对角元后计算行列式的成本为何是 O(d);(2) 比较 4.5 节仿射耦合层与 GraphAF 条件逐维仿射变换的雅可比结构;(3) 解释 GraphAF 为何在并行训练时仍需限制各步对未来节点和边的访问。

参考解答

(1) 下三角矩阵的第一行只有首个元素可能非零。沿第一行展开,得到 det J = J11 det J′,其中 J′ 仍为下三角矩阵。递归展开即得 det J = ∏k Jkk。d 个已知对角元只需 d−1 次乘法,成本为 O(d),不必进行通常成本为 O(d3) 的通用矩阵分解;计算这些对角元所需的网络前向成本另计。

(2) 仿射耦合层把输入分为两组:第一组保持不变,第二组的平移和逐维缩放由第一组决定。按这一顺序排列变量,雅可比为块下三角矩阵,对角块分别是单位矩阵和缩放因子组成的对角矩阵。GraphAF 则在给定历史条件后对当前原子或键向量做逐维仿射变换,当前块内的雅可比是对角矩阵;跨步骤的参数只依赖此前变量,所以整体雅可比为块下三角。两者都便于计算行列式,但条件来源不同。

(3) 顺序采样时,未来变量尚未生成,历史依赖可由当前子图直接保证。并行训练时,完整分子已经作为输入提供;若不限制访问,RGCN 可能读取尚未生成的节点或边,造成未来信息泄漏。GraphAF 原论文第 4.2 节因此使用掩码,让各步只连接到对应的历史子图。这与 MADE、MAF 保持自回归依赖的目的相同,区别在于掩码作用于图结构条件。三角结构是设计要求,掩码是并行实现这一要求的手段。

习题 4.6-3

论述:分别为 GCPN 和 GraphAF 提出一种可能的失败情境,从训练信号、训练分布和采样机制解释原因,并给出可观测的诊断指标。区分论文已报告的结果与需要实验验证的推断。

参考解答

GCPN 的一种风险是性质奖励过强、专家监督和对抗奖励相对过弱时,生成分布偏离训练分子。原论文讨论了减弱这些先验相关目标后出现高 penalized LogP 却不理想的结构。可同时记录性质分数、内部多样性、片段分布距离和关闭价检查后的有效性,并通过改变各目标权重检验原因。GraphAF 论文中 GCPN 的约 20% 无检查有效性是特定实验结果,不能证明其专家预训练不存在,也不能单独判断哪种任务一定先失败。

GraphAF 的一种风险是生成轨迹进入训练数据少见的部分图,尤其在较长分子或依赖窗口之外的连接模式中。这是基于顺序生成机制的推断,需要按分子长度、环大小和生成步数分组,测量无检查有效性、重采样次数和结构分布距离。均匀去量化下界与离散似然之间可能存在间隙,也应避免用连续密度替代生成质量评估。性质微调同样可能过度利用评分函数:若最高分上升而多样性和结构合理性下降,应检查奖励定义、训练分布和化学过滤规则,而非据单个高分判断模型成功。

4.6.6 TorchDrug 中的 GraphAF 实现

TorchDrug 将 GraphAF 封装为可组合的生成模块(Zhu et al., 2022)。官方教程以 RGCN 编码分子图,将同一个 GraphAF 类分别实例化为 node_flow 和 edge_flow,再交给 AutoregressiveGeneration 任务。该实现堆叠条件逐维仿射层;图条件决定各层缩放和平移,并非把向量分成两半的 Real NVP 耦合层。预训练使用负对数似然准则,性质优化使用 PPO;QED 示例还保留加权的负对数似然项。评估时应记录 RDKit 的版本、分子标准化、去重规则和样本分母,以便复核结果。实现细节见 官方分子生成教程与 ConditionalFlow 源码。

官方教程采用以下演示配置:

  • 数据:ZINC250k;教程估计加载与预处理约需 3–5 分钟,实际耗时随设备与缓存变化;
  • 预训练:负对数似然目标,10 轮,单卡 GPU,batch 128;节点流和边流各使用 12 层条件流;
  • 微调:示例训练 10 轮、batch 64;penalized LogP 使用 PPO,QED 使用权重为 0.25 的 PPO 与权重为 1.0 的负对数似然组合。

完成教程示例与复现论文结果是不同的验收目标。论文性质优化采用 300 轮预训练,并有独立的微调预算和超参数搜索;TorchDrug 示例的网络层数、生成上限和奖励配置也需逐项核对。不能只把示例轮数调到 300,就视为完整复现。

本节实现说明对应 TorchDrug 0.2.1 文档。其 官方安装说明列出 Python 3.7–3.10 和 PyTorch ≥ 1.8.0,并未把 PyTorch 2.0 写为兼容上限。实际可安装组合还取决于 CUDA、torch-scatter、torch-cluster 及编译环境,可按 5.2 节的方法建立并记录独立环境。Graphium 等工具的目标与接口需另行比较,不宜将其直接称为 GraphAF 生成方法的继任框架。

4.6.7 局限与后续方向

GraphAF 的局限主要涉及四个方面。其一,采样保留原子和候选键的顺序依赖,训练端的并行计算不能消除这种依赖。其二,训练以真实历史子图为条件,自由采样则使用模型生成的历史,可能出现 4.2.3 节讨论的暴露偏差;影响大小需通过实验判断。其三,均匀去量化优化的是离散对数似然下界,界的松紧会影响连续模型与离散生成目标之间的一致性,论文附录 A 将变分去量化列为改进方向。其四,性质优化仍受评分函数适用范围限制,高 penalized LogP 不等于成药性。

4.7 节将介绍扩散模型与三维生成。部分扩散方法在每个去噪步骤中同时更新多个节点或坐标,但通常仍需多个去噪步骤;采样速度和分布覆盖需要在具体实现与相同预算下比较。GraphAF 提供的比较基准是:条件连续密度如何建模、化学规则在何处施加,以及分布学习与性质优化如何衔接。

至此,4.2 节的顺序生成框架、4.3 节的策略梯度和 4.5 节的变量变换公式已经对应到具体分子生成过程。阅读后续模型时,可以继续用相同问题组织分析:模型表示什么分布,训练目标如何近似,采样怎样满足约束,以及性质提升是否伴随多样性或结构质量的变化。


关键术语

自回归流 (autoregressive flow)
各步参数只依赖此前变量的条件可逆变换序列,可顺序采样并计算连续表示的密度。
条件可逆变换 (conditional invertible transformation)
给定历史编码后,基础噪声与当前连续变量之间的可逆映射;类别解码不属于可逆部分。
基础分布 (base distribution)
流模型生成连续样本时使用的初始噪声分布;GraphAF 采用标准高斯。
变量变换公式 (change of variables)
由可逆变换与雅可比行列式把基础密度换算成数据密度的公式。
去量化 (dequantization)
通过给离散表示加入噪声构造连续训练数据;GraphAF 的均匀去量化对应离散对数似然的下界学习。
精确对数似然 (exact log-likelihood)
本节指给定连续表示和生成顺序后可计算的密度对数,不等同于离散无序分子图的精确概率。
节点流与边流 (node flow / edge flow)
分别建模原子和键的连续表示,经 argmax 解码为类别;边流还使用两端节点的条件表示。
近端策略优化 (proximal policy optimization, PPO)
通过概率比裁剪限制策略更新幅度的算法;GraphAF 与 GCPN 都用它优化性质奖励。
中间奖励与最终奖励 (intermediate / final reward)
生成步骤中的约束相关奖励,以及完整分子的性质和过滤奖励;最终奖励可按折扣分配给此前各步。
价约束检查 (valency check)
逐步检查候选键是否超过两端原子的允许价;GraphAF 拒绝违规键并重新采样。

参考文献与延伸阅读

  1. Shi C, Xu M, Zhu Z, Zhang W, Zhang M, Tang J. 2020. GraphAF: a flow-based autoregressive model for molecular graph generation. International Conference on Learning Representations (ICLR 2020).
  2. You J, Liu B, Ying R, Pande V, Leskovec J. 2018. Graph convolutional policy network for goal-directed molecular graph generation. Advances in Neural Information Processing Systems 31 (NeurIPS 2018):6410–6421.
  3. Rezende DJ, Mohamed S. 2015. Variational inference with normalizing flows. Proceedings of the 32nd International Conference on Machine Learning (PMLR 37):1530–1538.
  4. Papamakarios G, Nalisnick E, Rezende DJ, Mohamed S, Lakshminarayanan B. 2021. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research 22(57):1–64.
  5. Dinh L, Sohl-Dickstein J, Bengio S. 2017. Density estimation using Real NVP. International Conference on Learning Representations (ICLR 2017).
  6. Kingma DP, Dhariwal P. 2018. Glow: generative flow with invertible 1×1 convolutions. Advances in Neural Information Processing Systems 31 (NeurIPS 2018).
  7. Schulman J, Wolski F, Dhariwal P, Radford A, Klimov O. 2017. Proximal policy optimization algorithms. arXiv:1707.06347.
  8. Zhu Z, Shi C, Zhang Z, et al. 2022. TorchDrug: a powerful and flexible machine learning platform for drug discovery. arXiv:2202.08320.
  9. Gómez-Bombarelli R, Wei JN, Duvenaud D, et al. 2018. Automatic chemical design using a data-driven continuous representation of molecules. ACS Central Science 4(3):268–276.