4.7.1 扩散模型:加噪免费,去噪才要学
4.5 节结清了流模型的账:精确似然的代价是可逆约束,架构被行列式的可解析性捆住手脚;4.6 节的 GraphAF 已在这套约束内做到当年最好。扩散模型(diffusion model)换了交易方式:放弃逐样本似然的精确记账,换回架构的完全自由与训练的平稳(Ho et al., 2020;Song et al., 2021)。它如今是分子生成——尤其是三维生成——的主力范式。本节不推随机微分方程,只把机制讲透:加噪的方向为什么免费,去噪的方向为什么可学,训练为什么只剩一件事。
扩散模型由两条方向相反的马尔可夫链组成。前向过程(forward process)对数据逐步加噪:给 xt−1掺入小量高斯噪声,得到 xt,
调度选得合适,链的末端 xT 便近似各向同性的标准高斯——数据的一切结构被抹平,只剩噪声。直觉是搅拌一滴墨水:清水变浑的规则固定、已知、不可逆;想从浑水里找回那滴墨,才需要本事。训练时前向链不必真跑完,任意中间步可一步合成:
反向过程(reverse process)才是学习对象。倘若拿得到每一步的反向条件 pθ(xt−1 | xt),从纯噪声出发逐步采样,就能走回数据。它的可行性建立在一个数学事实上:每步 β 足够小时,反向转移也近似高斯,其均值指向「噪声略少一点」的版本。于是学习任务化简为一句话——给定 xt 与步号 t,预测当年掺进去的那份噪声,记网络输出为 εθ(xt, t)。训练目标因此极简:
去噪即得分。「预测噪声」看似偏门,实则与得分函数(score function) ∇x log pt(x) 成正比——最优去噪器给出对数密度梯度的方向。DDPM 的离散噪声链(Ho et al., 2020)与随机微分方程框架(Song et al., 2021)是同一机制的两种记法;本教程止步于这层对应,不再展开推导。
与流模型逐项对照,交换的条款看得最清楚。其一,架构。流模型每一层都要可逆、行列式都要可解析(4.5.4 节),扩散的去噪器却是任意网络——分子图上就是一个普通 GNN,三维场景再换等变 GNN。其二,稳定性。均方回归对批内噪声天然平均,训练曲线平稳;GAN 的模式坍塌、策略梯度的高方差都不在议程上。其三,似然。流模型逐样本精确;扩散的似然只有变分上界,实践中干脆不报 NLL,改以样本质量说话。这是明确的折中:质量优先,似然让位。
分子生成为何拥抱扩散?三个理由。第一,架构自由与图结构对味:消息传递网络天生就是去噪器,4.2 节的动作序列里学过的条件建模全部可以复用。第二,原子类型是类别变量,高斯噪声加不上去;离散扩散(discrete diffusion)以类别转移矩阵代替高斯核,把「逐步均匀化」搬进离散空间(Hoogeboom et al., 2021;Austin et al., 2021),D3PM 一类思路即由此而来。第三,Karras et al.(2022)的统一视角(EDM)把扩散设计空间拆成噪声调度、网络参数化、预处理器与采样器等可独立替换的旋钮——各组件的改进从此可以单独归因,工程上从手艺活变成实验设计。
扩散的账单同样明确:采样要迭代数十到上千步,每步一次网络前向。离线批量生成无所谓,交互式设计难耐延迟;少步蒸馏与一致性类方法是眼下的解法,此处按下不表。图 4.7-1 把双向机制画在一处。
4.7.2 三维转向与 SE(3) 等变
4.2 节起,生成的对象一直是二维拓扑图:原子类型、键、邻接。图只回答「谁连谁」,不回答「长成什么形状」。同一张拓扑图对应海量的构象(conformation)——同一分子在三维空间中的不同原子排布;能量、偶极、与蛋白的结合模式这些真正决定性质的量,都长在构象上。三维生成跳过「先出图、再折叠」的两段式,直接生成带坐标的构象:一次回答「连成什么样」与「长什么样」。
直接生成坐标立即招致一个语言问题:坐标没有规范朝向。同一构象整体旋转或平移之后,物理上仍是同一个分子——能量不变、性质不变、与受体的结合不变。数学上,刚体运动(旋转 R 与平移 t 的复合)构成群 SE(3)。网络若不内建这个对称性,就得为同一构象的每个朝向各学一遍表示;浪费参数之外,更糟的是不同朝向给出不同预测,物理自洽无从谈起。
SE(3) 等变性(SE(3) equivariance)。设 g = (R, t) ∈ SE(3) 为刚体运动,作用于坐标 g·x = R x + t。映射 f 称为等变的,若对一切 g 有 f(g·x) = g·f(x):输入转动,输出同步转动。标量函数 h 称为不变的,若 h(g·x) = h(x):能量、性质一类量属此。等变是不变在向量值输出上的推广;两者合起来,意思只有一句——网络的输出只依赖分子的内禀几何,不依赖实验室坐标系怎么摆。
对称性如何写进网络?等变图神经网络(equivariant graph neural network)中的 EGNN 给出目前最简洁的机制(Satorras et al., 2021)。每层维护两套状态:坐标 xi 与标量特征 hi,按三行更新:
三行读下来,等变与不变都不是约束出来的,而是构造出来的:不出现任何坐标的绝对值,一切计算都经由距离与相对向量。对照第 3 章的普通消息传递网络:那边特征不变、坐标根本不参与;EGNN 把坐标升为一等公民,又不必诉诸球谐基函数库一类的重型机械。图 4.7-2 用交换图把这层意思画直。
有了等变去噪器,4.7.1 的框架可以整体搬进三维。等变扩散(equivariant diffusion)(EDM;Hoogeboom et al., 2022)把原子类型(离散)与坐标(连续)放进同一扩散框架:类别变量走转移矩阵,坐标变量走高斯核,同一个 EGNN 充当两者的联合去噪器。两处细节见功力。其一,坐标先做重心归零,消掉平移方向的冗余——平移不改分子,却会稀释训练信号。其二,坐标噪声取各向同性高斯,旋转对称性保持到噪声本身:加噪后的分布在每个朝向上仍等概率,反向链学到的去噪也随之不偏不倚。键的生成通常放在原子与坐标之后按几何规则恢复,或以离散扩散一并处理。
等变到此也有边界。SE(3) 不含镜面反射:手性分子与其镜像不在同一轨道上,等变网络对两者一视同仁——而手性常常正是药效所在,需要显式特征或数据增强补课。此外,扩散一步只保证「几何上像」,不保证「能量上稳」;生成构象常需力场弛豫或经验证才可交付。先验内建得再深,也代替不了物理校验。
习题 4.7-1
推证。记平移算子 Tt(x) = x + t。(i)设平移映射 f(x) = x + b,证明 f 与 Tt 交换:f(Tt x) = Tt f(x) 对一切 t 成立;再说明若换作一般线性映射 f(x) = A x(A 为固定 3×3 矩阵),交换关系何时失效。(ii)设标量函数 g(x) = ‖x‖²。举一个具体的 t 与 x,说明 g(Tt x) ≠ g(x)。(iii)设两点 xi, xj,证明距离 ‖xi−xj‖ 在任意平移与旋转下不变,并据此说明式 (4.7-4) 中以距离平方为输入的消息为何自动获得不变性。
参考解答(i)两头分别展开:f(Tt x) = (x + t) + b = x + t + b;Tt f(x) = f(x) + t = x + b + t。两者恒相等,交换对一切 t 成立。若换作 f(x) = A x:f(Tt x) = A x + A t,而 Tt f(x) = A x + t,相等当且仅当 A t = t——A ≠ I 时,t 不落在 A 的不动子空间,交换即失效。结论写成一句:与一切平移交换的仿射映射必是纯平移;一般线性部分只与旋转部分同床(A 为正交阵时 f 与旋转交换),与平移不交换。(ii)取 x = 0、t = (1, 0, 0):g(Tt x) = ‖t‖² = 1,而 g(x) = 0,不等。绝对坐标不是不变量,任何以绝对坐标为输入的函数都不会平移不变。(iii)Tt 下差向量不变:Ttxi − Ttxj = xi − xj;旋转 R 是正交变换,‖R(xi−xj)‖² = (xi−xj)ᵀRᵀR(xi−xj) = ‖xi−xj‖²。故距离是刚体运动不变量。式 (4.7-4) 的消息 mij 只经由 ‖xi−xj‖² 与标量特征 h 依赖输入,复合不变量仍是不变量——消息的取值与实验室坐标系无关;坐标更新里的相对向量 (xi−xj) 则同步随 R 转动、不随 t 改变,等变因此由构造保证。
4.7.3 条件生成与实用形态
从原型到工具,差的是条件。药物设计里最要紧的条件是蛋白口袋(binding pocket):结合位点的残基给出一份三维的几何与化学上下文——氢键供体的朝向、疏水空腔的形状、金属离子的位置。以口袋为条件的生成把「配体应当长在哪儿、长成什么样」写进输入:口袋原子与配体原子同场进入等变网络,配体在口袋的坐标系内被逐步去噪出来;生成的对象不再是无主的分子,而是「对着这个位点量体裁衣」的候选(Guan et al., 2023)。以靶点结构为条件的三维生成由此把「画分子」变成「对着锁孔配钥匙」。此方向模型众多、迭代极快,本教程点到为止,不列清单。
条件如何注入?分类器无关引导(classifier-free guidance)是当下最通行的做法(Ho & Salimans, 2022):训练时以一定概率把条件遮空,同一个网络于是既会条件去噪、也会无条件去噪;采样时把两者作差外推:
一个减法把条件强度收敛为单个标量旋钮 w:调大则紧贴口袋、牺牲多样性,调小则自由发挥、牺牲结合。无需额外训练分类器,也无需对抗博弈,这是它在三维生成里近乎垄断的原因。
前沿结果的评估口径。三维生成领域论文迭代的节奏远快于基准统一的节奏。不同工作在不同数据集(QM9、GEOM-Drugs、CrossDocked 口袋集)上训练,报不同指标:FCD、SA、QED 之类二维指标之外,又有稳定性(是否经力场弛豫)、结合能(用哪种打分器)、原子价是否合法(判定脚本各异)。口径一换,名次常翻。横向比较前先核对四件事:数据集与划分方式、指标定义与实现版本、后处理(弛豫、去重、过滤)、统计口径(均值还是最优子集)。这正是 4.1 节立下的纪律——模型族可以换,裁判不能换。
习题 4.7-2
论述。为何三维分子生成「必须」等变?从数据效率、物理自洽、采样正确性三个角度展开;并回答:若用一个普通(非等变)GNN 直接以坐标为输入做去噪,分别会在哪三个环节付出什么代价?二维图生成(4.2–4.6 节)为何没有这层要求?
参考解答数据效率:训练集里每个构象只以某一个朝向出现,而非等变网络必须从有限朝向的样本中「顺便」学会所有朝向——同一几何的长尾姿态全靠数据量去填;等变网络则以构造保证一个样本等于它的全部朝向,数据需求下降可达数量级。物理自洽:能量与性质是刚体运动不变量,非等变网络对同一构象的不同朝相给出不同预测,违背定义本身;下游若接能量或力场任务,误差无界。采样正确性:扩散的每一步去噪都应把噪声分布往「等价的朝向族」上对称地拉;去噪器若偏爱某些朝向,反向链会把整体旋转的自由度逐步收缩到一个被偏好的姿态子集,采出的构象系综带人为的各向异性。非等变 GNN 直接吃绝对坐标时:学习环节要为每个朝向重复拟合,泛化环节在未见朝向上崩坏,采样环节引入姿态偏差。二维图生成没有这层要求,因为其对象是拓扑图:邻接与原子类型本身没有朝向,「旋转一张图」不是需要定义的操作;对称性负担是坐标进入生成对象那一刻才背上身的。
4.7.4 基础模型路线
生成之外,性质预测一侧正经历另一场范式迁移。第 3 章 3.7 节的多任务学习给出雏形:共享编码器、多头输出,一个模型同时服务若干任务。把这条路推到极限,就是分子基础模型(foundation model):以大规模图网络为骨架——消息传递与图 Transformer 并用的混合架构——在数百个任务、数百万分子的大混合上预训练。MolGPS 是代表(Sypetkowski et al., 2024):参数量达 1B–3B 量级,预训练混合不同保真度(fidelity)的标签——半经验、DFT、实验测值并存——微调与线性探测在 38 个下游任务中的 26 个上超过此前的单任务最优,且性能随深度、宽度、分子数与标签数的扩展单调上升。
工业逻辑一言可蔽:从「每个任务训一个模型」到「一个底座多任务微调」。旧范式下,数据管道、超参搜索、部署运维逐任务重复;基础模型把成本前置到一次预训练,下游以少量参数与数据适配。少样本任务受益最大——标签稀薄的任务骑在大底座的表示上,才拿得到信号。谱系上,这是 3.7 节硬参数共享的直系后代:共享的范围从「一个批次里的几个任务」扩展到「整个化学空间的所有标签」,共享的载体从几层编码器扩展到数十亿参数的底座。
生成与基础模型正在合流。一条路线把分子写成图序列或文本序列,交给自回归大模型逐 token 生成;另一条以扩散为生成底座,把预训练表示接入去噪器。两条路线都借用「预训练吃下大语料、微调对齐小目标」的同一套工业语法。孰优孰劣未有定论——本节只登记走向,不下判词。
4.7.5 五族谱系:一张表收束全章
四章走过五族生成模型,可以同桌对账了。表 4.7-1 按「似然—采样—稳定性—结构先验」四栏收束:前三栏是机器学习通用的账,第四栏才是本教程的主线——化学的结构性先验注入在了何处。
| 模型族 | 似然 | 采样 | 训练稳定性 | 结构先验的注入方式 |
|---|---|---|---|---|
| VAE(第 2 章) | ELBO 下界,真值不可见 | 一次解码 | 平稳;受后验近似限制 | 编码—解码整体映射,逐步结构无保证 |
| GAN | 无似然,密度未定义 | 一次前向 | 不稳,模式坍塌常见 | 判别器只看整图成品,中间过程无约束 |
| 流模型(4.5–4.6) | 精确可算 | 一次逆变换;自回归型逐维串行 | 平稳 | 可逆约束换解析行列式;GraphAF 以子图掩码写入价键合法性 |
| 扩散模型(本节) | 变分上界,通常不报 NLL | 迭代多步(数十至千步) | 平稳(均方回归) | 任意网络可作去噪器;离散(原子)与连续(坐标)分治;等变先验内建于架构 |
| 强化学习(4.3–4.4) | 不适用,按回报优化 | 逐动作展开 | 方差大,依赖基线与奖励设计 | 价键约束与中间奖励直接写入环境(GCPN) |
读表指南:扩散、流与自回归怎么选。依次三问。一问要不要逐样本似然的监控与比较?要,选流(4.5 节的精确对数似然);不要或可放弃,扩散与自回归皆可。二问采样延迟是否致命?交互式虚拟筛选等不得,自回归与流一次(或一遍)出分子;离线批量生成,扩散的多步迭代无碍,且质量目前最稳。三问训练预算与值守条件如何?低预算、无人值守,优先扩散或流:前者是回归,后者是似然,都不含对抗与高方差回报;GAN 与裸策略梯度慎碰。最后一栏提醒选型者:性能差距的相当部分来自先验注入的位置,而非网络本身的大小。
把 2018 年以来的脉络排成年表,五族的接力看得更清楚。
- 2018图生成起步:GraphVAE、MolGAN 与 GCPN 把「分子图=可微生成对象」立起来,评估靠人工挑选与简单指标。
- 2019–2020基准与似然双线并进:GuacaMol、MOSES 定下评估口径;GraphAF 把自回归与流合流,精确似然可用;同年 DDPM 在图像域立起新范式。
- 2021两块拼图就位:离散扩散(multinomial diffusion、D3PM)处理类别变量,EGNN 给出轻量等变骨架。
- 2022等变扩散 EDM 把原子类型与三维坐标放进同一框架,三维分子生成成为主流范式之一。
- 2023条件三维生成进入口袋场景(TargetDiff 等),生成与亲和力预测、打分函数耦合。
- 2024图基础模型 MolGPS(1B–3B 参数)证实 GNN 的规模收益;扩散侧少步采样与引导技术走向成熟。
- 2025–2026生成与基础模型合流:共享底座上的条件生成渐成工业默认路线(趋势观察,非定论)。
最后回扣 4.1 节。表换到扩散一列,结论一个字不变:无论哪一族模型,有效性、唯一性、新颖性、可合成性都要回到 4.1 节的评估体系去裁决。模型族是选手,评估体系才是裁判;选手越换越快,裁判的权威反而越显。前沿论文的军备竞赛里,最稀缺的从来不是新架构,而是统一口径的基准与统计纪律。
习题 4.7-3
推演。(i)表 4.7-1 中「扩散模型」一行的似然栏写「变分上界,通常不报 NLL」。请说明:为什么流模型能给出精确似然而扩散不能?各用两句话,从生成映射的结构(双射与可逆 vs 分步加噪的随机链)推到密度的可算性。(ii)某团队的需求为:逐样本 NLL 必须可监控;采样可离线批处理;训练预算有限且无人值守。按读表指南给出选型并说明理由。
参考解答(i)流模型:生成方向是确定的双射 z ↦ x,按变量替换公式(式 4.5-4 的多维版),数据密度=基础密度×雅可比行列式绝对值,逐样本精确可算(式 4.5-6 至 4.5-8 的链式分解)。扩散模型:生成是一串随机转移的乘积,x0 要对整条中间链 x1:T 积分(求和)才得到边际密度,该积分不可解,只能以下界(对联合分布的因子化近似)代替——训练优化的恰是这个下界而非真值,故逐样本 NLL 没有精确读数,实践中改以样本质量评估。(ii)选流模型一系(自回归流,如 GraphAF)。理由逐条对需求:NLL 可监控——流是五族中唯一精确逐样本可算的(VAE 只有下界、扩散只有上界、GAN 与 RL 无似然);离线批处理——不必顾虑自回归采样逐维串行的延迟;预算与值守——流与扩散的训练最平稳,而流同时满足第一条,故取流。附带提醒:接受架构上的可逆约束与逐维串行采样,正是这份选型单上写明的价格。
4.7.6 本卷收束
从 4.2 到 4.7,一条主线贯穿:把化学空间的结构性先验注入生成过程。4.2 节把生成分解为动作序列,价键合法性写进动作空间;4.3–4.4 节以回报驯服这个序列决策;4.5 节把链式法则翻译成可逆变换,4.6 节的 GraphAF 拿到精确似然;本节的扩散把先验挪进时间维——噪声调度规定「从结构到噪声」的路径,等变网络让几何先验长在架构里,基础模型则把表示的先验摊到整个化学空间。先验的注入位置在变,注入的必要性从未变:裸的通用生成器造不出分子,能造分子的,都是先验与数据合力之作。
至于 TorchDrug:它把 2018–2021 年那一波图生成范式封存进一套可运行、可读的 API——GCPN 与 GraphAF 的每一步都能亲手拆开看,作为教学标本恰好合身。工具本身已冻结于 v0.2.1(2022 年前后),活跃的后继(Graphium 一系)与整个工具链生态,留给第 5 章盘点。本卷到此收笔。
关键术语
- 扩散模型 (diffusion model)
- 前向固定加噪、反向学习去噪的生成模型族;训练归结为回归所加噪声。
- 前向过程 (forward process)
- 按固定噪声调度把数据逐步抹成近高斯的马尔可夫链,不含可学参数。
- 反向过程 (reverse process)
- 学习的去噪马尔可夫链,从纯噪声出发逐步还原数据分布。
- 得分函数 (score function)
- ∇x log p(x),对数密度的梯度;最优去噪器输出与之成正比。
- 离散扩散 (discrete diffusion)
- 以类别转移矩阵代替高斯核、在离散状态空间上定义的扩散过程。
- 构象 (conformation)
- 同一分子图在三维空间中的特定原子排布;能量与结合性质由它决定。
- SE(3) 等变性 (SE(3) equivariance)
- 输出随输入的刚体运动同步变换的性质;标量输出则退化为不变性。
- 等变图神经网络 (equivariant GNN)
- 坐标等变更新、标量特征不变的图网络;EGNN 以距离消息实现两者。
- 分类器无关引导 (classifier-free guidance)
- 条件与无条件得分之差按权重外推的条件注入与采样技巧。
- 蛋白口袋 (binding pocket)
- 蛋白上与配体结合的空腔位点;三维条件生成的核心条件来源。
- 基础模型 (foundation model)
- 大规模多任务预训练、经微调或探测适配众多下游任务的底座模型。
- 保真度 (fidelity)
- 标签的计算或实验精度等级;多保真预训练混合半经验、DFT 与实验值。
参考文献与延伸阅读
- Ho J, Jain A, Abbeel P. 2020. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems 33 (NeurIPS 2020): 6840–6851.
- Song Y, Sohl-Dickstein J, Kingma DP, Kumar A, Ermon S, Poole B. 2021. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations (ICLR 2021).
- Hoogeboom E, Nielsen D, Jaini P, Forré P, Welling M. 2021. Argmax flows and multinomial diffusion: learning categorical distributions. Advances in Neural Information Processing Systems 34 (NeurIPS 2021).
- Austin J, Johnson DD, Ho J, Tarlow D, van den Berg R. 2021. Structured denoising diffusion models in discrete state spaces. Advances in Neural Information Processing Systems 34 (NeurIPS 2021).
- Satorras VG, Hoogeboom E, Welling M. 2021. E(n) equivariant graph neural networks. Proceedings of the 38th International Conference on Machine Learning (PMLR 139): 732–741.
- Hoogeboom E, Satorras VG, Vignac C, Welling M. 2022. Equivariant diffusion for molecule generation in 3D. Proceedings of the 39th International Conference on Machine Learning (PMLR 162): 8867–8887.
- Karras T, Aittala M, Aila T, Laine S. 2022. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems 35 (NeurIPS 2022).
- Ho J, Salimans T. 2022. Classifier-free diffusion guidance. arXiv:2207.12598.
- Guan J, Zhou W, Huang Y, et al. 2023. 3D equivariant diffusion for target-aware molecule generation and affinity prediction. Advances in Neural Information Processing Systems 36 (NeurIPS 2023).
- Sypetkowski M, Wenkel F, Poursafaei F, et al. 2024. On the scalability of GNNs for molecular graphs. Advances in Neural Information Processing Systems 37 (NeurIPS 2024).