第四章 · 4.2

4.2 逐原子长大:自回归图生成的一般框架

Autoregressive Graph Generation
自回归生成把造一张分子图化约为走一串动作:每步在中间图上添原子、连键或终止,条件分布连乘即得分子概率。本节建立这一框架——状态、动作空间与链式分解;剖析图的线性化不唯一及其对极大似然的影响;辨析与文本生成的三点差异;给出训练目标、条件分布的图编码器接口与带掩码的采样流程,为策略梯度与流模型两条路线铺设共同骨架。

4.2.1 从一张图到一串动作:链式分解

4.1 节把分子生成定义成分布学习:给定训练集,学一个概率 p(x),使采样出的新分子在统计上像训练集、又不与之雷同。难处在图的组合规模:n 个原子之间每一对都可能成键,仅边的有无就有 2n(n−1)/2 种组合,元素类型与键型还要再乘上去。逐点枚举不可行,联合分布必须拆开。拆法中最直接的一条是自回归(autoregressive):不一次造出整张图,让分子逐原子长大,每步只做一个局部决定。生成一张图,由此化约为生成一个动作序列。GraphRNN 在一般图上验证了这条路(You et al., 2018a),Li 等随即把它带回分子(Li et al., 2018);本节的框架,正是 GCPN(4.4 节)与 GraphAF(4.6 节)共用的骨架。

框架由三个对象撑起。状态(state) st:到第 t 步为止长成的中间分子图,s0 为空图。动作空间(action space) A(s):状态 s 上可执行动作的集合,分三类——添加原子(类型 v 取自元素表 Σ,新原子暂为孤立节点);在新原子与每个已有节点之间加键(键型 b 取自 B = {单、双、三、芳},可连多条);终止。生成轨迹(generation trajectory) τ = (a1, …, aT):从空图出发依次执行动作、以终止收尾的动作序列;它诱导出最终分子图 g(τ)。

A(s) = { 加原子:v ∈ Σ }{ 加键:(i, b),i 为既有节点,b ∈ B } ∪ { 终止 }
(4.2-1)三类动作可交替执行;「加原子」与「加键」也可合并为一个复合动作「加入一个带键的原子」(4.2.3 节),两种粒度对应同一框架。
定义

自回归图生成的三件套。状态 st:第 t 步动作执行前的中间分子图。动作:在当前图上的一次局部修改,属于式 (4.2-1) 的三类之一。轨迹:从 s0(空图)出发、经 T 次动作到达终止的序列 τ = (a1, …, aT),其中 aT 为终止动作;轨迹的终点图记 g(τ)。分子与轨迹的关系是多对一:同一张图对应一族轨迹,一族记作 T(x)(4.2.2 节)。

轨迹一旦确定,分子随之确定;概率沿轨迹做链式分解(chain-rule factorization)

pθ(τ) = pθ(a1, …, aT) = ∏t=1T pθ(at | st), st+1 = st ⊕ at
(4.2-2)⊕ 表示「在图上施加动作」。把「词」换成「动作」、「前缀」换成「中间图」,本式与自回归语言模型的 p(wt | w<t) 逐字同构。同构只是起点:动作空间怎么变、合法性谁来管、何时停——三件事文本生成要么没有,要么处理方式完全不同,4.2.3 节逐一展开。

4.2.2 一张图,多条轨迹:线性化不唯一

文本的线性化是现成的:句子怎么写,就怎么生成。图没有「写法」。先加哪个原子、先连哪条键,分子图本身不置一词——同一张图对应一整族合法轨迹。图 4.2-1 给乙醇画出两条:一条从端点碳长起,一条从氧长起,动作串完全不同,终点同构。

乙醇的两种生成轨迹:同一分子,不同动作顺序,终点同构 轨迹 甲:从端点碳长起 s₀ +C C s₁ +C C C s₂ +O C C O s₃ 终止 动作串:+C → +C(连旧 C,单键)→ +O(连第 2 个 C,单键)→ 终止 轨迹 乙:从氧长起 s₀ +O O s₁ +C O C s₂ +C C C O s₃ 终止 动作串:+O → +C(连 O,单键)→ +C(连新 C,单键)→ 终止 两条动作串不同,终点同构——同一张乙醇图(氢原子按惯例隐去)。 顺序不是分子的属性,而是生成过程的属性:同一分子可有 |T(x)| 条合法轨迹,且随原子数迅速增长。 绿描边=本步新加入的原子;绿环=锚点(新原子连向的已有节点)。每帧按最终图的习惯布局重画。
图 4.2-1 同一分子,两条轨迹。乙醇 C2H5OH(隐氢,重原子图 C–C–O)的两条合法生成路径:轨迹甲从端点碳出发,轨迹乙从氧出发。绿描边标记本步新加入的原子,绿环标记锚点——新原子连向的已有节点。两条轨迹的动作序列完全不同,却终止于同构的分子图;这正是图生成与文本生成的第一个分野:线性化不唯一,顺序内生于生成过程。

顺序的两条出路。其一,给每张图指定规范序(canonical order):按宽度优先遍历、按规范化 SMILES 的解析序(第 2 章)等固定唯一顺序,一物一目标,训练最省事;代价是把「顺序」这个本不属于分子的自由度当作先验硬塞给模型。其二,把顺序交给模型:分子概率定义为全部合法轨迹的概率之和,

p(x) = Στ ∈ T(x) pθ(τ), T(x) = { τ : g(τ) ≅ x }
(4.2-3)T(x) 是诱导图与 x 同构的轨迹集合,≅ 表示图同构。|T(x)| 随原子数增长,同一分子在同一模型下有许多条可走的路。

「重复目标」由此进入极大似然。数据集里的一张图,训练时同时有 |T(x)| 个正确答案:梯度按轨迹概率加权,把所有指向 x 的轨迹一起推高(习题 4.2-2 给出加权形式)。概率质量摊在轨迹族上,单条轨迹的似然上限被轨迹数除;模型「最可能走哪条路」变得含混,但分子层面的 p(x) 并不受损。更深一层是图同构(graph isomorphism):数据存的是无标记图,训练器见到的是它的某一种节点编号;编号换了,输入表面就换。条件分布必须对节点重排不敏感,否则式 (4.2-3) 的求和没有意义——4.2.5 节的图编码器从结构上给这道保险。实践常取折中:每轮为每个分子随机抽一条合法轨迹,既不锁死顺序,也不必穷举全族。

注记

与 SMILES 序列生成的对照。SMILES 本身也存在一串多写:同一分子可以写出许多条 SMILES,RDKit 用规范化算法固定其一(第 2 章)。但 SMILES 生成模型生成的是字符串:词表固定、每步选一个字符,合法性由字符串语法兜底。图的自回归生成直接在图上动作,没有字符串这层中介:合法性改由价约束与掩码负责,顺序不再由书写顺序给定,而成为模型的一部分(或由规范序显式固定)。一句话概括:SMILES 模型学「怎么写」,图模型学「怎么搭」——后者把线性化的自由度与结构合法性一并接了过来。

4.2.3 与文本生成的三点差异

说「与语言模型同构」只对了一半。以下三处差异,各自牵动一项建模选择。

差异一:动作空间是组合的、动态的

文本每步在固定词表里挑一个符号,词表量级 104–105,与已生成内容无关。图生成的一步是三个选择的乘积:原子类型 × 锚点(anchor) × 键型。把「加入一个带键的原子」视作一个复合动作,动作总数为

|A(s)| = |Σ| × (1 + n(s)) × |B| + 1
(4.2-4)n(s) 为当前原子数;1+n 中的 1 对应首原子的「无锚点」选项;末尾 +1 是终止动作。

取 |Σ| = 8、|B| = 4:n = 20 时 |A| = 673,n = 40 时 |A| = 1313。词表不随句子变长,动作空间却随分子长大。一步大 softmax 对付不了它,标准做法是沿分解树拆成三个小分布连乘(式 4.2-7、图 4.2-2):先选原子类型,再在旧节点里选锚点,最后选键型。每个子问题的分支数骤减,且各自带着清晰的化学语义。

一步复合动作分解为原子类型、锚点、键型三层选择,外加终止动作 动作 aₜ (给定状态 sₜ) 加一个带键的原子 终止(STOP) ① 原子类型 v ∈ Σ ② 锚点 i:n 个旧节点 ③ 键型 b ∈ {单,双,三,芳} C N O S F 0 1 2 组合维度:|A(sₜ)| = |Σ| × (1 + n(sₜ)) × |B| + 1(式 4.2-4;+1 为终止) 文本语言模型的词表固定;这里的「词表」随状态增长——一步动作是三层选择的乘积。 树的每一层对应条件分布的一次分解:先选原子,再选锚点,最后选键型(式 4.2-7)。
图 4.2-2 动作空间的分解树。一步复合动作拆为三层选择:原子类型(元素表 Σ)、锚点(当前图上 n 个旧节点之一,首原子除外)、键型(单、双、三、芳),另有独立的终止动作。三层之积给出式 (4.2-4) 的组合维度;树的每一层同时对应条件分布的一次因子分解(式 4.2-7),实现时各接一个子分布头。

差异二:中间态必须合法

文本的任何前缀都算合法中间态,最坏不成句。图的中间态受价约束(valence constraint) 管辖——2.1 节的化合价规则在此从解析依据变成生成依据:碳四价、氮三价、氧二价(芳香键按等效价折算),任何使任一原子超出最高价的动作都是非法动作。非法不是「不好」,是「不行」:采样出六价碳,分子直接作废。处理办法是掩码(masking)

p̃(a | s) = pθ(a | s) · m(a, s) / Σa′∈A(s) pθ(a′ | s) · m(a′, s)
(4.2-5)m(a, s) ∈ {0, 1} 为合法性指示函数:动作 a 使任一原子超价则取 0。置零之后必须重归一化,否则概率和小于 1,采样有偏。

硬约束(hard constraint)软约束(soft constraint)之辨在此见分晓。软约束把非法写进损失作惩罚项,模型「倾向」不选,训练不足时非法样本照样冒头,4.1 节的有效率随之下跌;硬约束在结构上根除非法输出,代价是截断后的分布不再是网络的原始输出——训练沿合法数据轨迹时用不着掩码,推理时必须掩码,两侧存在系统性偏移,4.2.6 节回来算这笔账。

警示

非法动作与价约束掩码。掩码必须「置零+归一化」两步齐全。只置零不归一化,分布总和小于 1,缺口被静默丢弃,采样分布系统性有偏,且很难从指标上察觉。另一类错误是掩码过严:把「本步不该加键」误判为「永远不能加键」,模型被捆住手脚,多样性受损。掩码只该反映价约束这类物理规则,不该夹带顺序、风格等软性偏好——后者留给损失函数去说。逐动作掩码换来的是构造性的合法性:4.6 节 GraphAF 的对照数字(68% 对 100%)是这条纪律的实证。

差异三:终止是动作,不是惯例

文本的句号是词表里的一个符号,何时出现由语料惯例决定。图的终止是一个结构决策:模型自己判断「这张图已经完整」。由此生出两种失败模式。其一,早停:停在欠构建状态,产物合法却不是想要的分子——乙醇的轨迹停在两个碳,得到的是乙烷。其二,贪长:不肯终止,分子越长越大,直至价约束把可选动作挤干、或触发最大原子数的工程上限。两种模式都反映在生成长度分布上,4.1 节以训练集长度分布为参照的核对,正是诊断用的那把尺。工程兜底通常双保险:终止进动作空间由模型学,同时设最大原子数上限,到限强制截停。

文本自回归生成

  • 词表固定,与已生成内容无关;
  • 任意前缀都是合法中间态;
  • 句号是词表符号,时机靠语料惯例。

图自回归生成

  • 动作=类型×锚点×键型,随状态组合膨胀;
  • 中间态受价约束,非法动作须掩码归一化;
  • 终止是结构决策,早停与贪长皆有可能。
习题 4.2-1

设元素表 Σ = {C, N, O}(|Σ| = 3),键型 B = {单, 双, 三, 芳}(|B| = 4)。协议:每个复合动作一步完成「加一个原子并指定锚点与键型」(首原子无锚点),另有终止动作。(1) 当前中间图有 8 个原子,求动作总数 |A|。(2) 在「新原子必须与恰好一个旧原子连一条单键」的简化协议下,枚举生成乙醇 C–C–O 的全部轨迹,给出条数。(3) 就这组轨迹讨论「全轨迹求和」与「锁定单一规范轨迹」两种训练目标各自的得失。

参考解答

(1) |A| = 3 × (1+8) × 4 + 1 = 109。(2) 按首原子分三类。首原子为中间碳:第二步加端碳或加氧(各连中碳),第三步补齐另一种,共 2 条;首原子为端碳:第二步必须加中碳,第三步加氧连中碳,1 条;首原子为氧:第二步加中碳连氧,第三步加端碳连中碳,1 条。合计 4 条。(3) 求和目标最大化 p(x) = Στ∈T(x) p(τ):模型无须猜对唯一顺序,对数据增广与图同构更稳健;代价是概率质量摊在 4 条轨迹上,单条轨迹的似然上限被除以轨迹数,且优化涉及 log-sum,梯度按 p(τ)/p(x) 加权(习题 4.2-2)。锁定规范序则一物一目标、训练最简,但把顺序这一自由度硬编码为先验;若规范序与化学直觉不合,模型要额外花容量去记它。实践常用折中:每轮为每个分子随机抽一条轨迹。

4.2.4 训练:沿数据轨迹的极大似然

有了链式分解,训练目标顺理成章:对数据集中分子的合法轨迹做极大似然,

L(θ) = Σx∈D Σt=1Tx log pθ(at(x) | st(x))
(4.2-6)每个分子取一条轨迹(规范序固定,或逐轮随机抽取);若取全部轨迹,内层换成 log Στ∈T(x) pθ(τ),梯度即习题 4.2-2 的加权形式。

这一训练方式就是 teacher forcing:每一步的条件输入取数据轨迹演化出的状态 st(x),标签取数据的下一步动作;模型自己采样出的动作从不进入训练回路。好处直接:每个位置都有密集的监督信号,无须等待整条轨迹的回报,优化是一个干净的多分类问题。代价同样直接:推理时没有数据轨迹可依,模型只能沿自己的采样历史走。训练状态与推理状态的失配即 exposure bias——一步偏差把生成带进训练从未见过的中间图,误差沿轨迹累积;图生成里这一失配尤其严重(习题 4.2-3)。4.3 节的策略梯度从目标层面回应:干脆对采样出的整条轨迹按回报优化。

习题 4.2-2

设每步条件分布均经掩码并重归一化(式 4.2-5),且规定最大原子数上限 N,达到上限即强制终止。证明:(1) 轨迹分布 {p(τ)} 构成合法概率分布,即 Στ p(τ) = 1;(2) 分子概率 p(x) = Στ∈T(x) p(τ) 良定义,且 Σx p(x) = 1;(3) 最大化 log p(x) 的梯度同时提升所有指向 x 的轨迹。

参考解答

(1) 把全部轨迹组织成一棵有限有根树:根为 s0,每个内部节点是一个状态,其子边是掩码后的合法动作,叶节点为终止(或到上限被强制终止)的轨迹末端。在任一节点,子边概率之和为 1——掩码后的条件分布本身归一。对树高归纳:第一层各状态概率和为 1;若第 k 层概率和为 1,则第 k+1 层每个节点把单位质量按归一化条件分布分给子节点,总和仍为 1。上限 N 保证树有限且无无穷路径,叶节点即完整轨迹,故 Στ p(τ) = 1。(2) T(x) 按「诱导图与 x 同构」划分叶节点,不同 T(x) 互不相交且并集为全部叶节点,故 p(x) 良定义且 Σx p(x) = 1。(3) ∇θ log p(x) = ∇ log Στ∈T(x) p(τ) = Στ∈T(x) [p(τ)/p(x)] ∇ log p(τ)。权重非负、和为 1:每条指向 x 的轨迹按其占总概率的份额获得正向梯度,无一例外。这正是 4.2.2 节「重复目标同时推高」的解析形式。

习题 4.2-3

teacher forcing 训练沿数据状态推进,推理沿模型自身采样推进,两者失配即 exposure bias。论述:为何这一失配在图生成中比在文本生成中后果更重?可从状态空间的组合性、掩码重归一化对漂移的放大、错误状态的不可回溯三点展开,并说明 4.3–4.4 节的强化学习如何从目标层面缓解。

参考解答

其一,状态空间组合爆炸:一步偏差把生成带进训练从未覆盖的中间图,后续条件分布在外推,误差沿轨迹累积;文本尚有语言的平滑性兜底,图的邻接结构没有这层缓冲。其二,掩码重归一化会放大漂移:进入少见的价态格局后,剩余合法动作集合可能又小又生僻,被截断的原始分布把质量集中搬到这些动作上,采样进一步偏离数据流形。其三,不可回溯:文本里一个坏词只是句中瑕疵,图里一步坏动作改变整个后续结构空间,没有「教师」把状态拉回,也没有自然的回退机制。缓解之道有二:训练时按课程混入模型自身的采样状态(scheduled sampling 一类);更彻底的是把整条轨迹的回报纳入目标,即以策略梯度对采样分布直接优化——这正是 4.3 节的动机,GCPN(4.4 节)沿此思路在目标层面消解失配。

4.2.5 每步条件分布的接口:图编码器的复用

链式分解把生成压到一处:实现 pθ(at | st)。它要读的输入是一张逐步长大的分子图——第 3 章的图神经网络在此整体复用:若干层消息传递给出节点表示 hi(3.4–3.5 节),读出函数汇总出图表示 g。三个子分布头依次接上,

pθ(a | s) = p(v | s) · p(i | s, v) · p(b | s, v, i)
(4.2-7)一个原子选择分布 × 一个锚点分布 × 一个键型分布的组合分解——图 4.2-2 分解树的概率版本;首原子时锚点分布退化为「无锚点」一类。

消息传递对节点重排天然等变:同一张图换一套编号,节点表示随之换位,三个分布的输出不变。4.2.2 节留下的同构疑虑至此解除——式 (4.2-3) 的求和与编号无关,良定义。接口既定,训练路线分岔成两条:GCPN 把每步决策当作策略、把化学规则与性质奖励当作环境,用强化学习训练(4.3–4.4 节);GraphAF 把条件分布换成可逆变换的推动,得到可精确计算的对数似然,用流模型训练(4.5–4.6 节)。同一骨架、两种目标——本节是两节共用的地基。TorchDrug 已把 GCPN 收作标准生成模块(Zhu et al., 2022),4.4 节直接在其上动手。另有一条不走逐原子决策的路线:先把分子嵌入连续潜空间、再解码(Gómez-Bombarelli et al., 2018),它与本框架互补,留作对照。

4.2.6 采样:温度、贪心与逐动作合法性检查

推理就是沿式 (4.2-2) 逐步执行:每步在掩码归一化后的分布上抽取动作。温度控制分布的尖锐程度,

pT(a | s) = exp(za / T) / Σa′ exp(za′ / T)
(4.2-8)z 为网络输出的原始 logits。T→0 退化为贪心(取众数);T = 1 即原始分布;T > 1 更平坦。

温度采样(temperature sampling)是质量与多样性的旋钮。贪心(或极低温度)反复走同一条众数轨迹,4.1 节的 uniqueness 掉头向下,生成集退化成少数几个模板分子;温度过高则尾部动作(罕见元素、生僻键型)频繁中标,有效率(validity)与平均性质一同下滑。实务做法是扫一组温度、报质量—多样性曲线,而非只报单点。第二个旋钮是合法性:掩码在推理时同样生效,逐动作检查价约束,非法结构从源头断绝——有效率按构造为 100%。GraphAF 的报告给出对照数字:不用化学规则约束采样,约 68% 的样本化学合法;加上逐动作检查,100%(Shi et al., 2020)。这组数字是 4.6 节的卖点伏笔,也是硬约束一辨的实证注脚。代价同样要交代:掩码重归一化移动了采样分布,它不再是训练时那个未截断的条件分布;用一点分布偏移换全部样本合法,在药物设计的场景里几乎总是划算的交易。

方法

带合法性掩码的采样循环。① 置 s = 空图;② 图编码器读 s,输出各子分布的 logits;③ 按价约束生成掩码,置零并重归一化(式 4.2-5);④ 以温度 T 采样(或取众数);⑤ 施加动作更新 s;⑥ 动作为终止、或原子数到达上限,则输出分子(规范化后入库),否则回到 ②。逐动作检查让非法结构无从出现,有效率按构造为 100%;4.4 与 4.6 节沿用同一循环,只换条件分布的实现。

4.2.7 小结与去向

本节立起自回归图生成的框架:状态是中间分子图,动作是添原子、连键与终止,概率沿轨迹链式分解;线性化不唯一带来重复目标,规范序与全轨迹求和是两种处置;与文本的三点差异——组合动作空间、价约束掩码、终止即动作——决定了实现的全部关键选择;极大似然加 teacher forcing 训练接口,图编码器整体复用;采样端以温度调质量与多样性的平衡,以掩码保合法。框架到此闭合,任务未完:极大似然只教模型模仿数据分布,4.1 节立下的另一半目标——造出性质更优的分子——要求把目标从似然换成回报。策略梯度(4.3 节)接棒。


关键术语

自回归 (autoregressive)
以已生成的部分为条件逐项生成剩余部分的建模范式;图生成中逐动作长大分子。
链式分解 (chain-rule factorization)
把联合分布按生成顺序拆成条件分布连乘;与自回归语言模型同构。
状态 (state)
当前中间分子图,条件分布的输入载体。
动作空间 (action space)
每步可执行动作的集合:加原子、加键、终止;随状态变化。
生成轨迹 (generation trajectory)
从空图到终止的动作序列,其诱导图为最终分子。
锚点 (anchor)
新原子所连的已有节点;一步动作的三层选择之一。
规范序 (canonical order)
为图指定的唯一生成顺序,以消除轨迹多义性。
图同构 (graph isomorphism)
两图经节点重排后重合;同一分子多轨迹的根源。
硬约束与软约束 (hard / soft constraint)
前者以掩码从结构上剔除非法动作,后者以损失中的惩罚项劝阻非法动作。
teacher forcing
训练时沿数据轨迹而非模型自身采样推进的条件监督方式。
exposure bias
训练沿数据状态、推理沿自身状态导致的分布失配。
温度采样 (temperature sampling)
以温度参数调节分布尖锐程度的采样方式;调节质量与多样性。

参考文献与延伸阅读

  1. You J, Ying R, Ren X, Hamilton WL, Leskovec J. 2018a. GraphRNN: generating realistic graphs with deep auto-regressive models. Proceedings of the 35th International Conference on Machine Learning (PMLR 80).
  2. You J, Liu B, Ying Z, Pande V, Leskovec J. 2018b. Graph convolutional policy network for goal-directed molecular graph generation. Advances in Neural Information Processing Systems 31 (NeurIPS 2018).
  3. 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).
  4. Li Y, Vinyals O, Dyer C, Pascanu R, Battaglia P. 2018. Learning deep generative models of graphs. arXiv:1803.03324.
  5. 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:268–276.
  6. Zhu Z, Zhu Z, Xu M, et al. 2022. TorchDrug: a powerful and flexible machine learning platform for drug discovery. arXiv:2202.08320.