第四章 · 4.6

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

GraphAF: Autoregressive Flows for Molecules
GraphAF 把 4.2 的自回归骨架与 4.5 的流数学焊在一起:每步从高斯噪声出发,经条件可逆变换解码出原子类型与化学键,链式分解给出精确似然。本节拆解单步数据流——节点流、边流与价约束检查;梳理预训练与微调两阶段:ZINC250k 上极大似然立类药先验,PPO 微调追目标性质;核对有效性、penalized LogP 与训练效率的论文量级,末与 GCPN 对照并清点局限。

4.6.1 合流:4.2 的骨架,4.5 的数学

4.4 节的 GCPN 把 4.2 框架里的每个条件分布交给一个策略网络,训练信号是回报。GraphAF 换了一种实现:条件分布不由 softmax 直接给出,而由一次条件可逆变换(conditional invertible transformation)给出——在基础分布的噪声与离散的原子、键之间架一座可以回头的桥。模型名的两个词各认一门课:autoregressive 认 4.2 的链式分解,flow 认 4.5 的变量变换。两者拼起来,就是自回归流(autoregressive flow)(Shi et al., 2020)。

形式化从链式分解开始。生成序确定后,分子图 x 的概率沿式 (4.2-2) 拆开;GraphAF 把「原子—键」两层结构写得更细:

pθ(x) = ∏i=1n p(ai | g<i) · ∏j=1i−1 p(bij | g<i, ai, bi,<j)
(4.6-1)ai 为第 i 个原子的类别,bij 为新原子与旧节点 j 的键类别(含「无键」),g<i 为已生成子图。节点项对应 4.2 分解树的第①层,键项对应第②③层。注意键的条件里含 bi,<j:同一条键序列内部还要按序依赖,理由见 4.6.2。

与 4.5 节的 MAF(masked autoregressive flow,掩码自回归流)逐项对照,可以看清 GraphAF 的位置。MAF 用掩码强制「第 i 维只依赖前 i−1 维」,依赖关系藏在线性掩码里;GraphAF 把这层依赖显式交给编码器:hi = RGCN(g<i)(第 3 章的关系图卷积网络)汇总已生成子图,条件变换 fθ(·; hi) 以 hi 为条件参数。掩码换成编码器,条件从「维度的线性自选」升级为「结构感知的非线性摘要」;维度还随生成步增长——每步处理的向量只有新原子或新键那几维。似然由变量变换精确给出:

log pθ(xi | x<i) = log pZi) + log |det ∂xi/∂εi|, εi ~ N(0, I), xi = fθi; hi)
(4.6-2)pZ基础分布(base distribution)——标准高斯——的密度。变换由 4.5 的耦合层堆叠而成,变量变换公式(change of variables)中的雅可比取三角结构,行列式为对角元之积,代价 O(d)。MAF 靠掩码造出三角雅可比;GraphAF 每个变量单独一个小变换,三角性由构造直接保证(习题 4.6-2)。
定义

自回归流(autoregressive flow)。一串条件可逆变换 xi = fθi; hi):εi 独立采自基础分布,hi 为已生成部分的编码。正向(噪声到数据)逐步采样,反向(数据到噪声)沿同一组变换回程,链式分解与变量变换合起来给出精确似然。它兼得两头:自回归模型的条件表达力,流模型的可逆与可算密度。MAF 是它在固定维欧氏空间的原型,GraphAF 是它在逐步增长的图上的实例——「图」进了条件 hi,「流」管每个变量内部的变换。

这条路线的源头 4.5 节已交代:可逆变换配变量变换,密度即可精确计算(Rezende & Mohamed, 2015),谱系与变体见 Papamakarios 等人的综述(Papamakarios et al., 2021)。放在更大的棋盘上看:变分自编码器只能优化证据下界,密度本身评估不了(Gómez-Bombarelli et al., 2018);策略网络干脆不是密度模型。三条路线里,流是唯一既给精确似然、又允许逐步采样的——GraphAF 把这份「可对账性」搬进了分子图。

还剩离散性这道坎:原子类型与键型都是类别变量,可逆变换却住在连续空间。论文的处理是去量化(dequantization)——训练时给 one-hot 类别向量加 [0, 1)d 上的均匀噪声,得到连续向量再交给流;采样时反变换出连续向量,argmax 译回类别(Dinh et al., 2017;Kingma & Dhariwal, 2018)。

训练侧:xi = one-hot(ai) + u, u ~ U[0, 1)d; 采样侧:ai = argmaxk [ fθi; hi) ]k
(4.6-3)xi 为数据侧的连续向量,d 为类别数(ZINC250k 口径:原子 9 维、键 4 维——3 类键含「无键」)。argmax 译码保证任何连续样本都落回某个类别,采样永不中断;模型分给各类别「单位立方体」的概率质量,则由训练逐步压准。GraphAF 没有用类别概率的 sigmoid 串联或 logit 变换——按论文,离散与连续之间的桥就是均匀去量化。

4.6.2 单步数据流:一次条件变换的完整旅程

把一步拆开慢慢走。设子图 g<i 已生成,含 i−1 个原子与其间的键。

第一步,编码。RGCN 读子图,输出 hi——这一步是「条件」的全部来源,也是 4.2.5 预留接口的落点。

第二步,节点流(node flow)。从 N(0, I) 采 εi,经 L 层耦合变换(参数由 hi 调制)正向变换为 9 维连续向量,argmax 得原子类别 ai

第三步,边流(edge flow)。对新原子与每个旧节点的连键按序生成:先 j = 1,再 j = 2,依此类推;每条键的分布以 g<i、ai 与此前已采的键 bi,<j 为条件。按序而非一次并行,为的是捕捉价态相关——碳的一侧已落双键,另一侧可用的键型随即收紧。这是式 (4.6-1) 中键项内层依赖的意义。

锚点候选数随子图线性增长:第 i 个原子要对 i−1 个旧节点各做一次键判定,整张分子共 n(n−1)/2 次——恰是邻接矩阵上三角的规模。4.2 式 (4.2-4) 的组合动作空间在此没有被一次大 softmax 吞下,而是按「原子一次、键逐次」摊开:每步分支数小、化学语义清晰,代价是步数本身在累积。

第四步,价约束检查。每采一条键,核对两端原子的累计价;违价的键被拒绝,重采基础分布噪声再试。论文的机制是「检查—拒绝—重采」,不是采样前的掩码。

第五步,终止。新原子与旧子图之间一条键都没有生成、或原子数达上限,生成就此停止;否则子图更新为 g<i+1,回到第一步。

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 ∪ 新原子 回到编码(下一轮) 回到 ① 一步 = 一次条件变换:噪声在先、类别在后;键按序补齐,价约束兜底。 与 4.2.6 的采样循环同骨架:分布的实现由「softmax + 掩码」换成「流采样 + 检查拒绝」。 绿框为流的可逆部件;虚线为条件 h 的流动。检查开启时有效性按构造为 100%,关闭时约 68%(4.6.4)。
图 4.6-1 GraphAF 单步生成的完整数据流。左列:RGCN 把已生成子图编码为条件 h;上带:节点流把高斯噪声经 L 层耦合变换映为 9 维连续向量,argmax 译出原子类别;中带:边流对每个旧节点按序采键——每条键以子图、新原子与此前已采的键为条件——随后价约束检查,违价的键被拒绝并重采噪声;下带:无键相连或达上限则终止,否则子图更新、回到编码。整套循环与 4.2.6 的掩码采样循环逐条对应,差别只在条件分布的实现方式。

与 4.2 节式 (4.2-5) 的对照值得单独一说。4.2 的掩码在采样前把非法动作置零、重归一化;GraphAF 在采样后检查、拒绝、重采。落点不同,目标一致:非法结构被从结构上根除,有效性按构造为 100%。分布语义有差别——掩码直接改写候选分布,拒绝重采相当于在基础分布上多次抽签;两者都是硬约束的合法实现,4.2 预告的 68% 对 100% 两组读数,正是这套检查开与关的差别(4.6.4)。

边的按序依赖还有一个训练与采样的不对称注脚。teacher forcing 下全部键由数据轨迹给出,整张图的似然——节点与键、含每条键的内层依赖——一次前向并行算完,内层依赖不增加训练成本;采样时这些依赖变成真实的时间先后:每个新原子要等自己的键逐条落定,下一个原子才能开算。同一个依赖结构,训练里是「免费的图结构」,采样里是「逐次的等待」。式 (4.6-1) 的连乘在两端读出的时间代价完全不同——这是 4.6.4 效率之账与 4.6.7 局限之源。

方法

GraphAF 单步采样循环。① RGCN 编码 g<i 得 hi;② 采 εi ~ N(0, I),条件流正向变换,argmax 得原子类别 ai;③ 对每个旧节点按序:边流采键,条件含 ai 与已采各键;④ 价检查:违价 → 拒绝该键、重采噪声再试;⑤ 新原子无键相连或达上限 → 终止并输出分子;否则更新子图,回到 ①。逐动作的合法性检查换来构造性的 100% 有效;代价是采样侧的分布已被检查改写,不再是原始流分布——用分布偏移换全部样本合法,与 4.2.6 的结论一致。

习题 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) 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)。三个 p(a·) 由节点流给出;四个 p(b·) 由边流给出(b 不存在,首原子无边)。终止项 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,log |det J| = log σ = 0.6931。代入式 (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 就采 b32,等于假设同一原子的各条键互相独立——价化学恰恰不是这样。式 (4.6-1) 把这份相关写进条件,是边流「按序」设计的全部理由。

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

GraphAF 的训练剧本分两幕,正对 4.1 混合范式的两个半边。

第一幕是分布学习半边:在 ZINC250k(约 25 万个类药分子、至多 38 个原子、9 类原子、3 类键)上做极大似然预训练。流的可逆性在此兑现为工程优势:沿数据的生成轨迹做 teacher forcing,整张图全部原子与键的条件似然一次前向并行算完——无须滚动环境,无须判别器。论文配置:预训练 300 轮、batch 32、Adam(学习率 0.001)。

maxθ Σx∈D log pθ(x) = Σx∈D Σi [ log pZi) + log |det Ji| ]
(4.6-4)每个分子取一条生成轨迹(规范序或逐轮随机抽取,4.2.2 的处置照旧);对数似然精确对数似然(exact log-likelihood)可直接计算——这是流对策略网络的结构性优势:策略网络不是密度模型,给不出这个数。

先验质量的差别由此埋下。GCPN 论文同样有预训练,但训练信号自始至终是回报:模仿靠判别器逐动作打分,没有精确似然可言。GraphAF 预训练直接得到一个可评估的概率密度 pθ:可采样、可算似然、可对 4.1 的三比例与 FCD 直接对账。微调的起点因此不同——一边从「判别器的软信号」出发,一边从「精确拟合的分布」出发。「流 MLE + RL 微调」与「纯 RL」两条路线的分水岭就在这里:不是有无强化学习,而是预训练那一步给了什么。

第二幕是目标导向半边,机制经论文核实为 近端策略优化(proximal policy optimization, PPO)(Schulman et al., 2017)。奖励设计与 GCPN 一脉相承:中间奖励(intermediate reward)是对价违反的小惩罚——把 4.2 的硬约束在训练侧再上一道保险;最终奖励(final reward)是目标性质分数(penalized LogP 或 QED)减去空间位阻罚与违反 ZINC 官能团过滤的罚项,并按折扣因子摊回各中间步以稳定训练。

J(θ) = Eτ∼πθ[ Σt=1T γT−t rt ], rt = −λ·1[第 t 步违价](t < T), rT = s(x) − 检罚项
(4.6-5)πθ 即 GraphAF 的逐步采样策略;γ 为折扣因子;s(x) 为目标性质分数;检罚项含位阻罚与官能团过滤罚。策略梯度穿过条件变换与 RGCN 回传——可逆变换处处有定义的雅可比,保证了梯度通路的平滑。

「流的可导」常被误读为「性质打分可以直接反传」。按论文,微调是采样式的策略梯度,不是对 s(x) 做梯度上升。流给微调的实际贡献有三:好初值(预训练分布作为参数初始化)、平滑可导的策略参数化、以及微调全程可随时计算对数似然——策略是否漂离类药流形,流自带仪表盘;策略网络路线没有这个读数。两条训练信号的语义也不同:似然奖励「像数据」,是无偏的分布信号;回报奖励「合目标」,是有偏的任务信号。GraphAF 把两者分阶段使用;GCPN 把模仿回报与任务回报混在同一目标里。分与混,各有效果,也有各自的代价(习题 4.6-3)。

似然信号(预训练)

  • 问的是:「这批分子像不像数据」
  • 来源:数据轨迹的精确对数似然(式 4.6-4)
  • 性质:无偏的分布信号,密度可对账
  • 管住:连接方式与片段组合——「像药」

回报信号(微调)

  • 问的是:「这个方向值不值得走」
  • 来源:价违反惩罚 + 性质分数,按折扣摊回(式 4.6-5)
  • 性质:有偏的任务信号,方差靠 PPO 的裁剪压住
  • 管住:往高分区走的搜索——「合目标」
GraphAF 两阶段训练:先在 ZINC250k 上极大似然预训练得到类药先验,再以 PPO 微调追目标性质;下方对照 GCPN 的纯回报路线 4.1 的混合范式在此落地:先学「像」(分布学习半边),再学「好」(目标导向半边) 阶段一 预训练:极大似然 distribution learning(式 4.6-4) 数据:ZINC250k,约 25 万类药分子 信号:精确对数似然(teacher forcing) 并行:整图条件似然一次前向 预算:300 轮,batch 32,Adam 产出:类药先验 p(可评估的分布) 参数作为初始化 先验即起点 阶段二 微调:策略优化 goal-directed(式 4.6-5,PPO) 算法:PPO 策略梯度 中间奖励:价违反小惩罚 最终奖励:性质分数 − 位阻罚 − 官能团罚 折扣:终局分数摊回各中间步 梯度:穿过条件变换回传至编码器与流 GCPN 路线对照:训练信号自始至终是回报 预训练:判别器逐动作打模仿分(无精确似然);微调:任务回报。两阶段同形,信号不同。 真正的分野不是「有无 RL」,而是预训练给了什么:软的模仿信号,还是硬的精确分布。
图 4.6-2 GraphAF 的两阶段训练范式。阶段一在 ZINC250k 上以精确对数似然预训练,产出可评估的类药先验——4.1 混合范式的分布学习半边;参数随后作为初始化进入阶段二,PPO 以「价违反中间惩罚 + 目标性质最终奖励」微调——目标导向半边。下方对照 GCPN:两阶段结构同形,但预训练信号是判别器的模仿回报而非精确似然,先验质量由此分野。

4.6.4 结果核对:有效性、性质优化与效率

有效性要分两个口径读。不加化学规则约束时,GraphAF 约 68% 的样本化学合法,GCPN 同口径约 20%;开启采样中的价检查后,两家都是 100%(Shi et al., 2020)。无检查口径量的是「学到的分布里有多少质量落在合法分子上」——68% 说明流先验已把价化学学进大半,20% 则说明纯回报训练的先验更弱;100% 量的是工程保证,不是学习本身。跨口径抄数字,两个数都失去含义。

案例

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——近持平,QED 在数据集内有约 0.948 的天然上限,两家都摸到了天花板。
  • 训练时间(单张 V100 + 32 核 CPU):JT-VAE 约 24 小时,GCPN 约 8 小时,GraphAF 约 4 小时。

penalized LogP 领先幅度大,QED 打平。两个读数合起来才有信息:QED 有天花板,比的是谁先到顶;penalized LogP 无上限,比的是谁更敢跑——而「敢跑」的方向未必是药。4.1 的作弊标本对 GraphAF 原样适用:高分分子照样要过五规则、合成可行性与人工检视,分数本身不是终点。分布学习侧的评测论文另在 QM9 与 MOSES 上复算,结论同向——对照不止 ZINC250k 一处成立。

顺带一个读表细节:top-3 指生成集中分数最高的三个分子,量的是寻优的尖峰,不是平均质量;按 4.1 的四象限读法,分布侧指标与目标侧分数应分开报告。两家的 QED 都摸到 0.948 的上限,恰好说明该指标在此已无区分度——表格里没有差距的格子,同样是一种结果。

效率的账要分开算。训练时间减半来自流的并行性:teacher forcing 下整图的条件似然一次前向完成,滚动环境与收集回报的开销省去。采样端则要诚实:原子间的串行依赖没有消失,每个新原子都要一次编码器前向与若干次流变换;批量并行摊薄的是常数,改变不了 O(n) 的步数。训练快、采样慢,这个不对称是 4.6.7 局限的伏笔。

警示

数字可比的前提。GraphAF 的表格可读,是因为对照双方在同一数据划分(ZINC250k)、同一评估脚本、同一检查开关下完成。抄录文献数字前先核对三件事:其一,数据集与划分——ZINC250k、QM9、MOSES 互不可比;其二,有效性是否开启价检查——68% 与 100% 是两个口径,GCPN 的 20% 也是「无检查」口径;其三,训练预算与硬件——「4 小时对 8 小时」的前提是同一张卡。指标没有默认口径,口径写进表格才算数。

4.6.5 GCPN 与 GraphAF:同一骨架的两次实例化

把 4.4 与本节收进一张对照表(表 4.6-1)。

表 4.6-1GCPN 与 GraphAF 的受控对照(Shi et al., 2020;You et al., 2018;TorchDrug 实现统一了两者的流水线)
维度GCPN(4.4)GraphAF(本节)
条件分布实现策略网络 softmax(GNN 打分 + 掩码)条件可逆变换(节点流 + 边流)
预训练信号判别器模仿回报,无精确似然极大似然,精确对数似然(式 4.6-4)
先验来源模仿项混入目标函数预训练分布整体作为初始化
微调机制PPO(任务回报)PPO(性质回报 + 价违反中间惩罚)
有效性机制动作掩码:采样前置零归一化价检查:采样后拒绝、重采
似然可否评估否(策略网络非密度模型)可(变量变换精确计算)
训练效率(同硬件)约 8 小时约 4 小时(整图似然一次前向)
部件清单策略网络、判别器、价值网络、环境编码器、节点流、边流

共同点比差异更根本:两者都是「逐原子自回归 + 图编码器 + 价约束 + 预训练—微调」,是 4.2 框架的两次实例化;连数据(ZINC250k)、评估口径、骨干网络(RGCN)都相同。变量控制到只剩「条件分布的实现与训练信号」,这是 4.1 末指出的教学法价值:换一处、其余不动,差异才有归因。

两个模型还共享同一批失败教训:硬约束抬高合法性的同时压缩探索,追目标分又侵蚀多样性——4.1.6 的张力在两代模型上都可以复测。对照本身也可当作实验设计:拿同一份预训练权重,只换微调目标,观察 top-3 分数与内部多样性的此消彼长,比读十张表格更能说明「分布学习与目标导向为何缺一不可」。

习题 4.6-2

推证:(1) 证明若流变换的雅可比矩阵为下三角,则行列式等于对角元之积,行列式的计算代价为 O(d);(2) 说明 4.5 耦合层与 GraphAF 的逐变量变换如何各自获得三角雅可比;(3) 据此解释:GraphAF 为什么不需要 MAF 式的输入掩码。

参考解答

(1) 行列式按任意一行展开均可,对三角矩阵,非零乘积项只剩「全取对角元」一项,故 det J = ∏k Jkk。乘积只涉及 d 个对角元,代价 O(d),无须 O(d3) 的通用行列式分解。

(2) 耦合层把维度分成两半:一半恒等复制(对应对角元 1),另一半按前一半的函数逐维缩放和平移——每个被变换的维度只依赖自身与未被变换的维度,雅可比块为对角或零,整体三角。GraphAF 对每个变量(9 维原子向量或 4 维键向量)单独堆耦合层,变量之间的自回归依赖不走掩码、走条件 hi:每个小变换内部三角,跨变量的「只看已生成部分」由 RGCN 的输入直接保证。

(3) MAF 的掩码服务于一个目标:在固定维空间里用一个网络同时输出所有维度的变换,还要保住三角雅可比——只好用掩码切断「后维看前维」的通路。GraphAF 把「一个带掩码的大变换」换成「许多个以编码为条件的小变换」,三角性由构造给出,掩码自然没有存在的必要。条件化的代价是每步都要跑一遍编码器;换来的是条件可以非线性、结构感知,且随图的增长自然伸缩。

习题 4.6-3

论述:GCPN 与 GraphAF 各自最可能在什么任务上先失败?从训练信号、先验质量、采样机制三方面展开,并为每种失败模式配一个可观测的诊断指标。

参考解答

GCPN 的软肋在分布学习侧。训练信号是回报:判别器模仿分粗糙且有偏,价约束靠掩码兜底,模型自己学到的分布对价化学的覆盖弱——无检查口径 68% 对 20% 的差距即证。它先失败的任务是无约束的类药分布复现:多样性、FCD、按分子长度分桶的有效性都会露怯。诊断指标:关闭价检查的有效性(量先验质量)与内部多样性。它的高方差回报也使训练不稳,需要价值网络与判别器协同调参。

GraphAF 的软肋在长序列与目标极端化。采样严格串行,teacher forcing 预训练与自由采样推理之间的 exposure bias 沿轨迹累积,分子越长、环系越复杂,后期条件越容易落在训练未见的状态上;均匀去量化又是离散似然的粗糙代理,类别边界附近的质量未必压得准。它先失败的任务是长链与大环类结构:按分子长度分桶的 validity 会随长度下滑。目标导向侧它同样可被奖励作弊——penalized LogP 的高分分子未必成药,诊断要回到五规则与合成可行性的人检,而非分数本身。一句话:GCPN 输在先验弱,GraphAF 输在序列长;两者的目标导向风险同源——打分函数的边界窄。

4.6.6 TorchDrug 中的 GraphAF:可复现的教学遗产

TorchDrug 把 GraphAF 收作标准生成模块,与 GCPN 并列(Zhu et al., 2022)。实现结构与论文一一对应:同一个流类实例化两次,分别作 node_flow 与 edge_flow;RGCN 作骨干编码器;生成任务封装为一个自回归生成任务类,预训练走对数似然目标,微调把准则切到 PPO,任务在 QED 与 penalized LogP 之间选择——4.6.3 的两幕,在代码里就是两个入口。GCPN 与 GraphAF 共用同一条流水线,正合 4.6.5 的受控对照观。评估权仍在框架之外:TorchDrug 造分子,有效性、三比例与打分全部回到 RDKit 口径复核——4.1 的「裁判独立于选手」落到工具分工上。

复现门槛低是它最大的遗产。教程口径的数字一并列出:

  • 数据:ZINC250k,加载与预处理约 3–5 分钟;
  • 预训练:对数似然目标,10 轮,单卡 GPU,batch 128;
  • 微调:PPO 准则,10 轮,目标在 QED 与 penalized LogP 之间选择。

要分清两件事:跑通流水线是一节课的作业,复现论文数字要按论文的 300 轮预算另算——教程的轮数是演示规模,不是论文规模。

环境的账照例要算。TorchDrug 冻结于 v0.2.1(约 2022 年),官方兼容上限为 Python 3.10 与 PyTorch 2.0;新卡、新框架下的安装冲突须回到 5.2 的隔离方案。停更之后,这一代生成思想由 Graphium 等后继延续(5.1),而「预训练学先验、微调达目标」的范式本身并未过时。

4.6.7 局限与去向

四条局限,逐条对账。其一,采样串行:每个新原子一次编码器前向,长分子要 O(n) 步,训练端的并行救不了采样端。其二,误差累积:teacher forcing 预训练、自由采样推理,exposure bias(4.2.3)在长分子上放大——一步走偏,后续条件全在训练未见的状态上外推。其三,去量化粗糙:均匀噪声只是离散似然的一个代理上界,论文附录自承变分去量化是改进方向;桥越粗糙,离散分布与连续模型之间的缝越大。其四,目标可作弊:penalized LogP 的高分不等于成药,4.1 的警示原样成立。

去向也清楚。逐原子的串行与逐键的依赖是自回归的天性;扩散模型把整图一次并行去噪,模式覆盖更稳,向三维构象的扩展也更自然——4.7 接棒。GraphAF 留下的不是最快的采样器,而是一套干净的观念:条件可逆变换给出精确似然,价检查给出构造性合法,两阶段训练缝合「像」与「好」。这三个观念在后续每一代生成模型里都能找到回声。

就本章的线索而言,GraphAF 是终点前的最后一站:4.2 立骨架,4.3 补策略梯度的数学,4.4 与本节用两套训练信号各自点亮骨架,4.5 的流数学在这里完成从理论到分子的最后一跳。带着这套地图去读扩散模型与基础模型,新的只是去噪目标与网络结构,旧问题——合法性、多样性、目标作弊——一个都不会少。


关键术语

自回归流 (autoregressive flow)
条件可逆变换的链式串联:既可逐步采样,又可精确算似然的流模型;MAF 的图上推广。
条件可逆变换 (conditional invertible transformation)
以已生成部分的编码为条件参数的可逆变换,GraphAF 每步生成的执行者。
基础分布 (base distribution)
流模型采样的源头分布;GraphAF 取标准高斯。
变量变换公式 (change of variables)
由可逆变换与雅可比行列式把基础密度换算成数据密度的公式。
去量化 (dequantization)
离散类别加均匀噪声化为连续向量、argmax 译回类别的技术;流进入离散域的桥。
精确对数似然 (exact log-likelihood)
流模型可直接计算的数据密度对数;预训练的优化对象,策略网络给不出的读数。
节点流与边流 (node flow / edge flow)
分别生成原子类别与化学键类别的两套条件流;TorchDrug 中同一流类的两次实例化。
近端策略优化 (proximal policy optimization, PPO)
带裁剪目标的策略梯度算法;GraphAF 与 GCPN 微调阶段的共同选择。
中间奖励与最终奖励 (intermediate / final reward)
逐步的价违反惩罚与终点的性质分数;终局分数按折扣摊回各步。
价约束检查 (valency check)
采样时逐步核对化合价,违价的键被拒绝并重采;100% 有效性的工程来源。

参考文献与延伸阅读

  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 Z, 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):377–385.
  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.