4.5 流模型：变量变换与雅可比行列式

4.5.1 训练准则：下界、对抗与显式密度

4.2 节介绍了自回归生成，4.3–4.4 节讨论了回报优化。本节转向另一项建模要求：给定连续样本时，直接计算模型赋予它的密度。为说明这一要求，先比较几类常见训练准则。变分自编码器（Kingma & Welling, 2014）引入潜变量 z，边际似然 p_θ(x) 需要对 z 积分；对常用神经网络解码器，这个积分通常难以直接计算，因此训练中使用证据下界(evidence lower bound, ELBO)：

ELBO 给出边际对数似然的下界，但仅凭该曲线通常无法确定界的松紧；VAE 的似然也可用重要性采样等方法估计。生成对抗网络采用无似然训练(likelihood-free)，通过生成器与判别器的对抗目标拟合数据。生成器仍定义了一个概率分布，只是通常没有可直接求值的数据密度。策略梯度则优化期望回报；所用策略仍可具有显式动作概率，并可与似然训练结合。训练准则与概率表示因而是两个不同的分类维度。

流模型(normalizing flow)通过可逆映射与基础密度构造数据密度。本节讨论由有限个、可计算行列式的变换组成的流，其精确对数似然(exact log-likelihood)指给定连续样本后可直接计算的对数密度，不必用变分下界代替该密度。这个术语不表示训练无采样误差，也不表示离散数据的概率自动精确可得（Rezende & Mohamed, 2015；Papamakarios et al., 2021）。表 4.5-1 概括了常见情形。

显式密度可以提供逐样本训练信号，并在数据表示、坐标尺度和预处理一致时支持负对数似然（NLL）比较。它与 4.1 节的有效性、多样性和性质指标互补，密度较高本身并不保证分子更合理。分布学习还可与性质奖励衔接：GraphAF 先拟合去量化后的连续表示，再用回报微调。其离散对数似然下界见 4.5.7 节；此前先推导连续变量的密度公式。

4.5.2 一维：概率质量守恒与变量变换

设随机变量 X 的密度为 pₓ(x)，令 Z=f(X)。先假定 f 在两个区间之间构成双射(bijection)，连续可微，且 f′ 处处非零。这组条件保证逆映射也连续可微。严格单调能保证映射到其像集的一一对应，但单独的严格单调并不排除零导数，例如 f(x)=x³ 在原点的导数为零。密度换算的依据是概率质量在变换下守恒：对应区间描述同一个随机事件。

取小区间 [x₀, x₀+Δ]，并在密度连续的点 x₀ 附近令 Δ 趋于零。事件 {X ∈ [x₀, x₀+Δ]} 等价于 Z 落在该区间经 f 映射所得的像区间中；若 f 递减，像区间的左右端点顺序相反。在数据空间中，这一概率近似为：

在潜空间中，由一阶展开 f(x₀+Δ)−f(x₀)≈f′(x₀)Δ，像区间的长度约为 |f′(x₀)|Δ。因此，同一概率也近似为：

将两式相等，约去 Δ 并取极限，即得变量变换(change of variables)公式，也称变量替换公式：

公式说明，守恒的是概率质量，密度则随坐标变换而改变。例如取 f(x)=2x，即生成方向 x=z/2。若 Z 为标准正态，X 的标准差减半，峰值加倍：p_x(0)=2p_z(0)。同一概率质量分布在更短的区间内，密度相应升高。对于全实数轴上的 f(x)=x²，正负两个分支会映到同一位置，需要按分支求和处理正向密度，不能直接使用单一逆映射的公式。若生成映射为 x=z³，且基础密度在原点连续并为正，则生成密度会在原点附近发散；这来自生成映射在该点的零导数。图 4.5-1 示意了正逆方向与质量守恒。

4.5.3 多维推广：雅可比行列式作为体积缩放因子

对向量 x ∈ Rⁿ，可微变换 f: Rⁿ → Rⁿ 的导数由雅可比矩阵(Jacobian matrix) J_f(x) 表示，其第 (i,j) 个元素为 ∂f_i/∂x_j。在 x₀ 附近，一阶近似为 f(x₀+δ)≈f(x₀)+J_f(x₀)δ。以下假定 f 为连续可微双射，且雅可比行列式处处非零。

线性映射 J 将单位立方体映为平行多面体，其 n 维体积为 |det J|。因此，x₀ 附近一个体积为 δV 的小区域，经 f 映射后的体积近似为 |det J_f(x₀)|δV。比较两个对应区域的概率，约去体积并取极限，得到多维公式：

训练通常使用对数密度。若基础分布 p_z 的各维相互独立，例如多维标准正态，其对数密度可逐维相加：

多个可逆层可以复合为 f=f_L∘…∘f₂∘f₁。令 h₀=x，h_ℓ=f_ℓ(h_ℓ−1)，最终 h_L=z。复合映射的雅可比按链式法则相乘，行列式的乘积在对数域中变为求和：

对连续训练样本，可最大化 Σ_x log p_θ(x)，并用自动微分计算参数梯度。接下来需要设计既可逆、又能高效计算行列式的变换；这正是流结构与普通神经网络结构的主要区别。

4.5.4 行列式计算与结构约束

若不利用特殊结构，计算 n 维深网络的行列式通常需要先形成 n×n 雅可比。反向模式自动微分可逐个输出计算雅可比的各行，前向模式可逐个输入计算各列；显式存储需要 O(n²) 空间，通用 LU 分解还需要 O(n³) 运算。这些成本不包含构造雅可比所需的网络计算。维数较大且每个训练样本都需评估时，开销会明显增加。此外，普通网络不自动满足全局可逆性。

常见设计从易求逆、易计算行列式的结构出发，再复合多层。本节介绍两类：仿射耦合层将坐标分组，保留一组并变换另一组；自回归流按次序构造条件变换，使雅可比具有三角结构。它们并非所有流模型的完整分类，而是理解 GraphAF 所需的两个基础例子。

4.5.5 仿射耦合层：分组与条件变换

耦合层(coupling layer)将坐标分组，使一组变量为另一组的可逆变换提供条件。NICE 使用加性形式（Dinh et al., 2015），Real NVP 使用仿射形式（Dinh et al., 2017）。以下令 x=(x_A,x_B)，保留 A 组，将其作为 B 组缩放与平移的条件。两组常取相同大小，但不要求严格各占一半；也可通过棋盘或通道掩码划分。正向变换为：

给定 y，先由恒等分支恢复 x_A=y_A，再计算 s(y_A) 与 t(y_A)。此时 B 组的仿射系数已知，逆变换可逐元素求出：

按 y=(y_A,y_B) 与 x=(x_A,x_B) 分块，雅可比的左上块为单位阵 I，右上块为零，因为 y_A 与 x_B 无关。右下块是逐元素缩放形成的对角阵 diag(exp s_i)。左下块 ∂y_B/∂x_A 可以稠密，但块三角矩阵的行列式只取决于对角块：

单个耦合层保留 x_A 不变，因而不能独立实现任意联合变换。常用做法是在相邻层交换分组、改变掩码或置换坐标，使不同维度依次作为条件输入和被变换变量。Glow 还在层间使用可逆 1×1 卷积混合通道（Kingma & Dhariwal, 2018）。这类复合既保留各层易于求逆和计算行列式的性质，也增加了维度之间的相互作用。

4.5.6 自回归流：条件密度与三角雅可比

自回归结构也能使雅可比行列式便于计算。按照 4.2 节的链式法则，联合密度可沿固定次序分解为 p(x) = ∏_i p(x_i | x_<i)。自回归流(autoregressive flow)用依赖前缀的可逆标量变换(scalar transformation)构造这些条件密度。以条件仿射形式为例，网络读取 x_<i，输出位置 μ_i 与正尺度 σ_i，并采用各维独立的基础分布：

从联合变换看，z_i 只依赖 x_≤i，对 x_j（j > i）的偏导为零，所以雅可比为下三角矩阵。其对角元为 ∂z_i/∂x_i = 1/σ_i，因而 log |det J| = −Σ log σ_i，与式 (4.5-11) 的逐条件计算一致。两种推导描述的是同一个密度：一种从联合变量变换出发，另一种对每个条件密度应用一维公式后相乘。

GraphAF 将条件变换用于图生成，条件网络读取历史子图，当前变量则是原子或键类别的去量化向量。这里需要区分条件仿射变换与前节的仿射耦合层：GraphAF 的条件由当前向量之外的历史提供，可对当前向量的全部分量同时缩放和平移；RealNVP 型耦合层则在同一输入内分组，以 A 组决定 B 组的变换，并保留 A 组不变。两者都能得到易计算的行列式，但结构不同；仿射变换也只是一般耦合层可采用的一类可逆变换。可对照 GraphAF 原论文第 4.1 节与 TorchDrug 的 ConditionalFlow 实现阅读。

计算方向取决于条件网络的实现。MAF 使用掩码网络约束前缀依赖：密度求值时，完整数据 x 已知，可在一次网络前向中并行计算各维条件参数，再得到全部 z_i 与对数密度。采样时则按 x_i = μ_i(x_<i) + σ_i(x_<i)·z_i 逐维恢复，必须先得到所需前缀。若采用 IAF 那样依赖基础变量前缀的掩码结构，通常可快速采样，而任意观测点的密度求值较慢。这种并行能力来自掩码网络的具体设计；仅仅使用真实前缀训练，并不保证循环式条件网络也能一次并行求值。相关比较见 MAF 原论文。

GraphAF 的历史掩码承担类似作用：训练时虽已知完整分子图，各步骤的图编码仍只能访问当时已生成的原子和键。掩码阻断未来信息，使多个步骤的条件密度能在一次前向计算中并行求值。采样时历史尚未形成，仍需按生成次序推进。这种防止信息泄漏的历史掩码，应与检查候选键是否违反价约束的采样规则区分开来；后者涉及无效候选的拒绝与重采样，不能替代历史依赖约束。

4.5.7 结构、数值与离散数据

流模型的密度可计算性依赖结构设计。实际使用时，还需考虑以下四个方面。

其一，架构约束。本节讨论的可微双射保持连续空间的维数，变换层必须满足可逆性，并使求逆与行列式计算的成本可控。耦合层和自回归结构各自限制了单层的依赖关系，多层复合可以增加表达能力。具体拟合效果仍取决于变换族、深度、条件网络与数据分布，不能仅凭可逆性断言同等参数量下的性能高低。

其二，数值稳定。exp(s) 过大或过小可能导致溢出、下溢或较差的数值条件。对数密度和对数行列式按式 (4.5-6)、(4.5-7) 累加，能避免直接连乘带来的部分问题，但不能消除前向变换中的溢出。因此实现中常限制对数尺度的幅值，并检查逆变换误差。Glow 使用的 ActNorm(activation normalization)是逐通道仿射变换，以首批数据初始化平移和尺度；其行列式易于计算，尺度非零时可逆。这样的初始化有助于控制初始激活尺度，但不保证后续训练始终稳定（Kingma & Dhariwal, 2018）。

其三，离散数据。变量变换公式直接描述连续密度。将离散数据用于连续流，通常需要去量化(dequantization)。例如，给整数向量 a 加入 u ∼ U([0,1)^d)，令 x = a + u；若以单位网格作为离散化区域，则类别质量为 \(P_\theta(a)=\int_{[0,1)^d}p_\theta(a+u)\,du\)。由 Jensen 不等式，\(\mathbb{E}_u[\log p_\theta(a+u)]\leq\log P_\theta(a)\)。因此，连续流可精确计算每个去量化样本的密度，但对噪声取期望后的训练目标是离散对数似然的下界；单次加噪仅提供这一期望的随机估计。若还做数值缩放，也应计入相应的密度修正。

GraphAF 对原子和键的 one-hot 向量逐分量加入 U[0,1) 噪声，学习所得连续向量的条件密度；生成时对连续向量取 argmax 得到类别，并非将一个标量按区间映射为类别。固定历史后，某一类别的概率应对其 argmax 区域积分。训练用的 one-hot 单位盒子位于相应区域内，盒子上的去量化目标仍给出下界，不能直接视为精确类别概率。进一步计算无序分子图的概率，还需处理生成顺序以及约束采样规则的影响。GraphAF 原论文附录 A说明了均匀去量化与离散似然下界的关系；4.6 节将展开这一训练与采样流程。

其四，方向成本。密度求值和采样分别使用变换的两个方向，效率不一定相同。仿射耦合层的两个方向成本接近；MAF 与 IAF 则分别偏向快速密度求值和快速采样。选型应结合训练目标、采样量与时延要求，同时考虑条件网络和流层数，而非要求所有结构在两个方向都具有相同效率。

本节的推导以概率质量守恒为起点：雅可比行列式修正局部体积变化，多层复合的对数行列式按层相加。耦合层和自回归流通过限制依赖结构，使这项计算更容易。应用到分子生成时，连续密度、离散类别概率和分子图概率需要分别定义。下一节将说明 GraphAF 如何结合历史子图条件、去量化训练与价约束采样，并在似然预训练后进行奖励微调。