Generative Models II

2026年9月24日约 35 分钟taking-lying-flat

🦨 Variational Inference with Normalizing Flows

变分推断通常需要用一类已知的概率分布去近似难以直接求解的真实后验分布，并在该分布族中寻找最优近似。然而，实际使用的近似分布族往往较为受限，例如均值场近似。这意味着，无论优化过程如何进行，最终得到的解都可能无法真正逼近真实后验分布。这也是变分方法经常受到批评的重要原因之一：不同于 MCMC 等其他推断方法，传统变分推断即使在渐近情形下，也通常无法恢复真实后验分布。已有大量证据表明，更丰富、更忠实于真实后验的近似分布通常能够带来更好的性能

与采用均值场近似的 sigmoid belief networks 相比，深度自回归网络通过引入具有自回归依赖结构的后验近似，获得明显性能提升
与此同时，也有许多研究揭示受限后验近似所带来的负面影响，Turner 和 Sahani 总结两个常见问题

其一，后验分布方差经常被低估，这可能导致预测质量下降，并使基于该后验近似所做出的决策变得不可靠
其二，后验近似能力不足还可能导致模型参数的最大后验估计出现偏差，这一现象在时间序列模型中尤为明显

为解决后验近似表达能力不足的问题，已有一些方法尝试构造更加丰富的后验近似

这些方法通常基于结构化均值场近似，在近似后验内部引入某些基本形式的变量依赖关系
另一种潜在有效的方法是将近似后验设定为混合模型。然而，混合模型方法会限制变分推断的可扩展性，因为在每次参数更新时，它通常需要针对每个混合成分分别计算对数似然及其梯度，这在计算上往往代价较高

作者提出一种新的变分推断后验近似构造方法。首先回顾当前在一般有向图模型中进行推断的主流实践，即基于摊销变分推断和高效 Monte Carlo 梯度估计的方法。随后，论文主要做出如下贡献

提出使用归一化流来刻画近似后验分布。归一化流是一种通过一系列可逆映射将简单概率密度变换为复杂分布的工具。基于归一化流的推断可以得到一个更紧的、经过修正的变分下界，并且新增项仅带来线性时间复杂度的计算开销
表明归一化流可以扩展到无穷小流，从而定义一类在渐近情形下能够恢复真实后验分布的后验近似。这在一定程度上克服变分推断长期以来被批评的一个问题，即近似分布族过于受限，无法逼近真实后验

🐇 Amortized Variational Inference

为进行推断，通常只需要围绕概率模型的边际似然进行分析，而这要求对模型中的缺失变量或潜变量进行边缘化。然而，这一积分在实际中通常是不可解析的，因此转而优化边际似然的一个下界。考虑一个一般的概率模型，其中观测变量为 $x$ ，需要被积分掉的潜变量为 $z$ ，模型参数为 $\theta$ 。为潜变量引入一个近似后验分布 $q_{\phi}(z|x)$ ，并根据变分原理得到边际似然的一个下界

$\begin{aligned} \log p_\theta(\mathbf{x}) &= \log \int p_\theta(\mathbf{x}|\mathbf{z})p(\mathbf{z})d\mathbf{z} \\ &= \log \int \frac{q_\phi(\mathbf{z}|\mathbf{x})}{q_\phi(\mathbf{z}|\mathbf{x})} p_\theta(\mathbf{x}|\mathbf{z})p(\mathbf{z})d\mathbf{z} \\ &\geq -\mathbb{D}_{\mathrm{KL}}[q_\phi(\mathbf{z}|\mathbf{x})\|p(\mathbf{z})] + \mathbb{E}_q[\log p_\theta(\mathbf{x}|\mathbf{z})] = -\mathcal{F}(\mathbf{x}), \end{aligned}$ (1)

其中，最后一个不等式由 Jensen 不等式得到。 $p_{\theta}(x|z)$ 表示似然函数， $p(z)$ 表示潜变量的先验分布。上述下界通常被称为负自由能，或称为证据下界。它由两个部分组成：第一项是近似后验分布与先验分布之间的 KL 散度，可被视为一种正则化项；第二项是期望对数似然，通常对应于重构误差或数据拟合项。该下界为同时优化模型参数 $\theta$ 和变分近似参数 $\phi$ 提供一个统一的目标函数

当前变分推断的主流实践通常使用小批量数据和随机梯度下降来优化这一目标函数。这也是变分推断能够扩展到大规模数据集和复杂概率模型的重要原因。为成功应用变分方法，需要解决两个关键问题

如何高效计算期望对数似然关于变分参数的梯度，即如何计算 $\nabla_{\phi}\mathbb{E}_{q_{\phi}(z)} \big[\log p_{\theta}(x|z)\big]$
如何选择一个尽可能丰富、同时又具有计算可行性的近似后验分布族 $q(\cdot)$

论文旨在构造更具表达能力的近似后验，并结合推断网络与 Monte Carlo 梯度估计进行摊销变分推断。对于连续潜变量，通过重参数化与 Monte Carlo 采样估计期望对数似然的梯度 $\nabla_{\phi}\mathbb{E}_{q_{\phi}(z)}[\log p(x|z)].$ 小批量数据采样与潜变量采样构成双重随机估计。随机反向传播包含以下两个步骤

第一步是重参数化。也就是说，将潜变量重新表示为一个已知基础分布中的随机变量与一个可微变换的组合。这类可微变换可以是位置—尺度变换，也可以是累积分布函数变换。以高斯近似后验为例，若 $q_{\phi}(z)$ 是均值为 $\mu$ 、方差为 $\sigma^2$ 的高斯分布，其中变分参数为 $\phi=\{\mu,\sigma^2\}$ ，则可以利用标准正态分布作为基础分布，将潜变量重写。随机性不再直接来自依赖参数 $\phi$ 的分布 $q_{\phi}(z)$ ，而是来自一个固定的、与参数无关的基础噪声分布。潜变量 $z$ 则通过一个关于 $\phi$ 可微的确定性变换得到

$z \sim \mathcal{N}(z|\mu,\sigma^2) \quad \Longleftrightarrow \quad z = \mu + \sigma \epsilon,\quad \epsilon \sim \mathcal{N}(0,1)$ (2)

第二步是基于 Monte Carlo 的反向传播。在完成重参数化之后，就可以从基础分布中采样，并使用 Monte Carlo 近似来估计期望。同时，由于采样后的变量可以表示为参数 $\phi$ 的可微函数，因此可以直接通过反向传播计算目标函数关于变分参数的梯度。也就是说，原本较难处理的梯度 $\nabla_{\phi}\mathbb{E}_{q_{\phi}(z)}[f_{\theta}(z)]$ 可以转化为在标准正态噪声下对如下形式进行估计 $\mathbb{E}_{\epsilon\sim \mathcal{N}(0,1)} \left[ \nabla_{\phi}f_{\theta}(\mu+\sigma\epsilon) \right]$

📌 第二项重要实践是，近似后验分布 $q_{\phi}(\cdot)$ 通常通过识别模型或推断网络来表示。所谓推断网络，是指一种学习从观测变量到潜在变量之逆向映射的模型。借助推断网络，无需为每一个数据点单独计算一组变分参数，而是可以学习一组全局变分参数 $\phi$ ，使其在训练阶段和测试阶段均可用于推断。这样一来，推断成本便可以通过推断网络参数在不同潜在变量后验估计之间的泛化能力而得到摊销。最简单的一类推断模型是对角高斯分布

$q_{\phi}(z|x)=\mathcal{N}\bigl(z|\mu_{\phi}(x), \mathrm{diag}(\sigma_{\phi}^{2}(x))\bigr),$ (3)

其中，均值函数 $\mu_{\phi}(x)$ 与标准差函数 $\sigma_{\phi}(x)$ 均由深度神经网络进行参数化

📌 DLGM 是一类通用的深层有向图模型，由 $L$ 层高斯潜在变量构成，其中第 $l$ 层的潜在变量记为 $z_l$ 。每一层潜在变量都以非线性方式依赖于其上一层变量；在 DLGM 中，这种非线性依赖关系由深度神经网络进行建模。其联合概率模型可表示为

$p(x,z_1,\ldots,z_L) = p(x|f_0(z_1))\,p(z_L)\prod_{l=1}^{L-1}p(z_l|f_l(z_{l+1})).$ (4)

其中，第 $L$ 层的高斯分布不再依赖于其他随机变量。顶层潜变量的先验通常设定为单位高斯分布 $p(z_L)=\mathcal{N}(0,I),$ 而观测似然 $p_{\theta}(x|z)$ 则可以根据具体数据类型选择适当的分布形式。该似然分布以第一层潜在变量 $z_1$ 为条件，并同样由深度神经网络进行参数化。DLGM 使用连续潜在变量，因此非常适合结合变分下界与随机反向传播方法进行高效的摊销变分推断
然而，真实的后验分布往往比这类假设所能表达的分布更加复杂。因此，如何以可扩展的方式定义多模态的、受约束的后验近似分布，仍然是变分推断领域中的一个重要开放问题

🦨 Normalizing Flows

通过考察变分下界可以看出，使得 $D_{\mathrm{KL}}\bigl[q_{\phi}(z|x)\|p_{\theta}(z|x)\bigr]=0$ 成立的最优变分分布，应当满足 $q_{\phi}(z|x)=p_{\theta}(z|x),$ 也就是说，近似后验分布 $q_{\phi}(z|x)$ 必须与真实后验分布完全一致

在通常采用的 $q(\cdot)$ 分布族下，例如独立高斯分布或其他均值场近似，这一目标显然难以实现。变分方法的一个重要局限就在于：由于可选近似分布族的表达能力有限，即使在渐近情形下也未必能够得到真实后验分布
理想的变分分布族 $q_{\phi}(z|x)$ 应当具有高度灵活的表达能力，最好能够足够灵活，以至于真实后验分布本身也包含在该分布族之中。通向这一理想目标的一条路径，便是基于归一化流的思想

归一化流描述的是：一个概率密度如何通过一系列可逆映射逐步变换。通过反复应用变量变换公式，初始概率密度会沿着这一系列可逆映射不断流动。在该序列的终点仍可得到一个合法的概率分布，因此这种流被称为归一化流

有限流（Finite Flows）：密度变换的基本规则考虑一个可逆且光滑的映射 $f:\mathbb{R}^d \to \mathbb{R}^d$ ，其逆映射记为 $f^{-1}=g$ ，即满足复合关系 $g \circ f(z)=z$ 。若使用该映射对服从分布 $q(z)$ 的随机变量 $z$ 进行变换，则所得随机变量 $z' = f(z)$ 的分布为

$q(\mathbf{z}') = q(\mathbf{z}) \left| \det \frac{\partial f^{-1}}{\partial \mathbf{z}'} \right| = q(\mathbf{z}) \left| \det \frac{\partial f}{\partial \mathbf{z}} \right|^{-1}$ (5)

其中，最后一个等式可由链式法则，即反函数定理推出，并且是可逆函数雅可比矩阵的一项基本性质。通过组合多个简单映射，并反复应用上述变量替换规则，可以构造出任意复杂的概率密度。设初始随机变量 $z_0$ 服从分布 $q_0$ ，并依次经过 $K$ 个变换 $f_k$ ，则经过连续变换后得到的密度 $q_K(z_K)$ 可表示为

$\begin{gathered} \mathbf{z}_K = f_K \circ \dots \circ f_2 \circ f_1(\mathbf{z}_0) \\ \ln q_K(\mathbf{z}_K) = \ln q_0(\mathbf{z}_0) - \sum_{k=1}^K \ln \left| \det \frac{\partial f_k}{\partial \mathbf{z}_{k-1}} \right|, \end{gathered}$ (6)

由随机变量序列 $z_k = f_k(z_{k-1})$ 所形成的轨迹称为流；由相应分布序列 $q_k$ 所形成的演化路径，则称为归一化流。这类变换具有一个重要性质，通常被称为无意识统计学家定律（LOTUS）：关于变换后密度 $q_K$ 的期望，可在不显式求得 $q_K$ 的情况下计算。对于任意函数 $h(z)$

$\mathbb{E}_{q_K}[h(\mathbf{z})] = \mathbb{E}_{q_0}[h(f_K \circ f_{K-1} \circ \dots \circ f_1(\mathbf{z}_0))],$ (7)

无穷小流（Infinitesimal Flows）：归一化流的长度趋于无穷。在这种情况下得到的是一种无穷小流。它不再由有限个变换组成的序列来描述；相反被表述为一个偏微分方程，用于刻画初始密度 $q_0(z)$ 如何随时间演化，算子 $T$ 描述连续时间下的动力学过程

$\frac{\partial}{\partial t} q_t(z) = T_t[q_t(z)]$ (8)

Langevin 流：一类重要的流由 Langevin 随机微分方程（stochastic differential equation, SDE）给出

$dz(t) = F(z(t),t)dt + G(z(t),t)d\xi(t),$ (9)

其中 $d\xi(t)$ 表示标准 Wiener 过程的增量；从零出发的该过程满足 $\mathbb{E}[\xi_i(t)] = 0,$ 并且 $\mathbb{E}[\xi_i(t)\xi_j(t\prime)] = \delta_{ij}\min(t,t\prime).$ 这里 $F$ 是漂移向量， $D = GG^\top$ 是扩散矩阵。若将一个初始密度为 $q_0(z)$ 的随机变量 $z$ 通过 Langevin 流进行变换，则密度的变换规律由 Fokker–Planck 方程 给出；在概率论中，它也被称为 Kolmogorov 方程。经过变换后，时刻 $t$ 的样本密度 $q_t(z)$ 将按照如下方程演化

$\frac{\partial}{\partial t}q_t(z) = -\sum_i \frac{\partial}{\partial z_i} \left[ F_i(z,t)q_t \right] + \frac{1}{2} \sum_{i,j} \frac{\partial^2}{\partial z_i \partial z_j} \left[ D_{ij}(z,t)q_t \right].$ (10)

在机器学习中，最常使用的一种 Langevin 流形式为 $F(z,t) = -\nabla_z L(z)$ 以及 $G(z,t)=\sqrt{2}I,$ 其中 $L(z)$ 是能量函数，即目标未归一化密度的负对数。在这种设定下， $q_t(z)$ 的平稳解由 Boltzmann 分布给出 $q_\infty(z) \propto e^{-L(z)}.$
也就是说，如果从某一初始密度 $q_0(z)$ 出发，并让其中的样本 $z_0$ 按照 Langevin 随机微分方程不断演化，那么最终得到的点 $z_\infty$ 将服从如下分布 $q_\infty(z) \propto e^{-L(z)}.$ 在理想条件下，这一极限分布可以对应于真实后验分布。因此，Langevin 流为从复杂密度中进行采样提供一种重要途径

Hamiltonian 流：Hamiltonian Monte Carlo，简称 HMC，也可以从归一化流的角度加以描述

具体而言，它可以被看作是在一个扩展空间上的归一化流。该扩展空间记为 $\tilde{z} = (z,\omega),$ 其中 $z$ 是原始变量， $\omega$ 通常表示辅助动量变量。其动力学由 Hamiltonian 函数决定， $M$ 表示正定质量矩阵，动量服从协方差为该矩阵的零均值高斯分布

$H(z,\omega) = L(z) + \frac{1}{2}\omega^\top M^{-1}\omega.$ (11)

🦫 Inference with Normalizing Flows

为能够使用有限归一化流进行可扩展的推断，必须指定一类可逆变换，并且需要一种能够高效计算雅可比行列式的机制

虽然构造可逆的参数化函数是相对直接的，例如使用可逆神经网络，但这类方法在计算雅可比行列式时通常具有较高复杂度。其复杂度通常随 $O(LD^3)$ 增长，其中 $D$ 表示隐藏层的维度， $L$ 表示所使用的隐藏层数量
计算雅可比行列式的梯度还需要若干额外操作，这些操作同样具有 $O(LD^3)$ 的复杂度，并且往往涉及矩阵求逆，而矩阵求逆在数值上可能是不稳定的。因此需要设计这样一类归一化流：它们要么能够以较低代价计算行列式，要么根本不需要显式计算雅可比矩阵

考虑如下形式的一族变换 $f(z)=z+u h(w^\top z+b),$ 其中 $\lambda=\{w\in \mathbb{R}^D,\ u\in \mathbb{R}^D,\ b\in \mathbb{R}\}$ 是自由参数， $h(\cdot)$ 是一个光滑的逐元素非线性函数。对于该映射，可以利用矩阵行列式引理，在 $O(D)$ 时间内计算对数雅可比行列式项。定义

$\begin{gathered} \psi(\mathbf{z}) = h'(\mathbf{w}^\top \mathbf{z} + b)\mathbf{w} \\ \left| \det \frac{\partial f}{\partial \mathbf{z}} \right| = |\det(\mathbf{I} + \mathbf{u}\psi(\mathbf{z})^\top)| = |1 + \mathbf{u}^\top \psi(\mathbf{z})|. \end{gathered}$ (12)

若将任意初始密度 $q_0(z)$ 依次通过一系列形如映射 $f_k$ 进行变换，则所得密度 $q_K(z_K)$ 可以隐式表示为

$\begin{gathered} \mathbf{z}_K = f_K \circ f_{K-1} \circ \dots \circ f_1(\mathbf{z}) \\ \ln q_K(\mathbf{z}_K) = \ln q_0(\mathbf{z}) - \sum_{k=1}^K \ln |1 + \mathbf{u}_k^\top \psi_k(\mathbf{z}_{k-1})|. \end{gathered}$ (13)

由此所定义的流，对超平面 $w^\top z+b=0$ 上各点施加相同位移，位移方向由向量 u 决定，幅度随超平面的位置变化，从而改变初始密度 $q_0$ 。因此，这类映射被称为平面流（planar flows）。采用 tanh 时，可施加 $w^\top u>-1$ 作为保证可逆且雅可比非奇异的充分条件

作为另一种选择，还可以考虑这样一族变换：它们围绕某一参考点 $z_0$ 对初始密度 $q_0$ 进行局部修改。该变换族定义为

$\begin{gathered} f(\mathbf z)=\mathbf z+\beta h(\alpha,r)(\mathbf z-\mathbf z_0),\\ \det\frac{\partial f}{\partial\mathbf z} =[1+\beta h(\alpha,r)]^{D-1}[1+\beta h(\alpha,r)+\beta h'(\alpha,r)r],\\ r=\|\mathbf z-\mathbf z_0\|_2,\quad h(\alpha,r)=\frac{1}{\alpha+r},\quad h'=\frac{\partial h}{\partial r} \end{gathered}$ (14)

该映射的参数为 $\lambda=\{z_0\in\mathbb R^D,\ \alpha>0,\ \beta>-\alpha\}.$ 这一变换族同样允许以线性时间计算行列式。它围绕参考点 $z_0$ 施加径向的收缩与扩张，因此被称为径向流（radial flows）

基于流的自由能界：若使用长度为 $K$ 的归一化流来参数化近似后验分布，即 $q_\phi(z|x):=q_K(z_K),$ 则自由能写成关于初始分布 $q_0(z_0)$ 的期望

$\begin{aligned} \mathcal{F}(\mathbf{x}) &= \mathbb{E}_{q_\phi(z|x)}[\log q_\phi(\mathbf{z}|\mathbf{x}) - \log p(\mathbf{x}, \mathbf{z})] = \mathbb{E}_{q_0(z_0)} [\ln q_K(\mathbf{z}_K) - \log p(\mathbf{x}, \mathbf{z}_K)] \\ &= \mathbb{E}_{q_0(z_0)} [\ln q_0(\mathbf{z}_0)] - \mathbb{E}_{q_0(z_0)} [\log p(\mathbf{x}, \mathbf{z}_K)] - \mathbb{E}_{q_0(z_0)} \left[ \sum_{k=1}^K \ln |1 + \mathbf{u}_k^\top \psi_k(\mathbf{z}_{k-1})| \right]. \end{aligned}$ (15)

归一化流及其对应的自由能界可以与任意变分优化方案结合使用，包括广义变分 EM。对于摊销变分推断，作者使用深度神经网络构造推断模型，将观测变量 $x$ 映射为初始密度 $q_0=\mathcal N(\mu,\operatorname{diag}(\sigma^2))$ 的参数，其中 $\mu\in\mathbb{R}^D,$ $\sigma\in\mathbb{R}_{>0}^D$ ，同时也映射得到归一化流的参数 $\lambda$
在实验中，作者采用自由能的一个简单退火版本，其中 $\beta_t\in[0,1]$ 是逆温度参数，其调度策略为 $\beta_t=\min(1,0.01+t/10000).$

$\mathcal{F}_{\beta_t}(x) = \mathbb{E}_{q_0(z_0)}[\ln q_0(z_0)] - \beta_t \mathbb{E}_{q_0(z_0)}[\log p(x, z_K)] - \mathbb{E}_{q_0(z_0)} \left[ \sum_{k=1}^K \ln |1 + u_k^\top \psi_k(z_{k-1})| \right],$ (16)

Algorithm 1 Variational Inference With Normalizing Flows

🕊️ Neural Ordinary Differential Equations

提出一类新的深度神经网络模型。不同于显式指定一个离散的隐藏层序列，使用神经网络来参数化隐藏状态的导数

模型的输出由黑箱微分方程求解器计算得到。此类连续深度模型在理论上具备恒定的内存开销，能够根据不同输入的复杂性自适应地调整数值求解策略，并允许在计算精度与推理效率之间进行显式权衡
进一步构造连续归一化流模型。作为一类生成模型，该方法能够在无需对数据维度进行划分或重排的情况下，通过最大似然目标进行直接训练。在优化过程中，模型可以在不依赖 ODE 求解器内部实现细节的前提下，对任意黑箱求解过程进行可扩展的反向传播，从而使 ODE 模块能够被嵌入更大规模的神经网络系统中，并支持端到端训练

诸如残差网络、循环神经网络解码器以及归一化流等模型，通常通过对隐藏状态施加一系列变换，来构建复杂的状态变换过程

$\mathbf{h}_{t+1}=\mathbf{h}_t+f(\mathbf{h}_t,\theta_t)$ (1)

其中 $t\in \{0,\ldots,T\}$ 且 $\mathbf{h}_t\in \mathbb{R}^D$ 。这些迭代式更新可以被视为对某一连续变换的欧拉离散化

当不断增加网络层数，并同时减小每一步的更新幅度时，会发生什么？在极限情形下，使用由神经网络指定的常微分方程来参数化隐藏单元的连续动力学

$\frac{d\mathbf{h}(t)}{dt}=f(\mathbf{h}(t),t,\theta)$ (2)

从输入层 $\mathbf{h}(0)$ 出发，可将输出层 $\mathbf{h}(T)$ 定义为该 ODE 初值问题在某一时间 $T$ 处的解。该值可以通过黑箱微分方程求解器计算得到；求解器会在必要的位置评估隐藏单元动力学函数 $f$ ，以便以期望精度确定最终解
残差网络定义一组有限的离散变换序列，ODE 网络定义一个向量场，该向量场以连续方式变换状态

使用 ODE 求解器定义并评估模型具有以下若干优势

伴随方法通过计算求解器输入的梯度，减少中间状态存储。通过计算一个标量损失函数相对于任意 ODE 求解器输入的梯度，而无需通过求解器的内部运算过程进行反向传播。由于不需要存储前向传播中的任何中间量，模型训练时的内存开销可以作为深度的函数保持恒定
ODE 求解器依据误差要求自适应分配计算量。欧拉方法可能是求解 ODE 的最简单方法。现代 ODE 求解器能够对近似误差的增长进行控制，监测误差水平，并在计算过程中动态调整评估策略，以达到所要求的精度水平。这使得模型的评估成本能够随问题复杂度而变化。训练完成后，还可以降低求解精度，以适应实时应用或低功耗应用场景
将隐藏状态动力学参数化为时间的连续函数，可使相邻层共享参数
连续变量替换公式可用于构造可逆的密度模型。连续变换带来的一个意外优势是，变量替换公式变得更容易计算。利用它构造一类新的可逆密度模型。该模型避免归一化流中常见的单隐藏单元瓶颈，并且可以直接通过最大似然进行训练
连续时间动力学能够处理不规则时刻到达的观测。不同于循环神经网络需要将观测区间和发射区间离散化，连续定义的动力学能够自然整合在任意时间到达的数据

🦨 Reverse-mode automatic differentiation of ODE solutions

训练连续深度网络时面临的主要技术困难，是如何通过 ODE 求解器执行反向模式微分，也即通常所说的反向传播。直接对前向传播过程中的每一步运算进行微分虽然较为直接，但会带来较高的内存开销，并引入额外的数值误差

论文将 ODE 求解器视为一个黑箱，并采用伴随敏感度方法来计算梯度。该方法通过在时间上反向求解一个增广的 ODE 系统来获得梯度，适用于任意 ODE 求解器。该方法的计算复杂度随问题规模线性增长，具有较低的内存开销，并且能够显式控制数值误差

考虑优化一个标量损失函数 $L(\cdot)$ ，其输入为 ODE 求解器的输出

$L(\mathbf{z}(t_1)) = L \left( \mathbf{z}(t_0) + \int_{t_0}^{t_1} f(\mathbf{z}(t), t, \theta)dt \right) = L(\mathrm{ODESolve}(\mathbf{z}(t_0), f, t_0, t_1, \theta))$ (3)

为优化 $L$ ，需要计算其关于参数 $\theta$ 的梯度。第一步是确定损失函数的梯度如何在每一个时刻依赖于隐藏状态 $z(t)$ 。该量被称为伴随变量 $a(t)=\frac{\partial L}{\partial z(t)}.$ 其动态演化由另一个 ODE 给出，可以将其理解为链式法则在连续时间情形下的瞬时对应形式

$\frac{da(t)}{dt} = -a(t) \frac{\partial f(z(t), t, \theta)}{\partial z}$ (4)

可以通过再次调用 ODE 求解器来计算 $\partial L/\partial z(t_0)$ 。这一求解器需要沿时间反向运行，其初始值为 $\partial L/\partial z(t_1)$ 。复杂之处在于求解该 ODE 需要知道整条轨迹上的 $z(t)$ 取值。然而，可以在反向求解伴随变量的同时，从终点 $z(t_1)$ 出发，重新计算 $z(t)$ 的轨迹

ODE 解的反向模式微分过程。伴随敏感度方法通过沿时间反向求解一个增广 ODE 系统来完成梯度计算。该增广系统同时包含原始状态，以及损失函数相对于状态的敏感度。如果损失函数直接依赖于多个观测时刻处的状态，则在每一个观测时刻，伴随状态都必须按照损失函数相对于该观测状态的偏导方向进行更新

关于参数 $\theta$ 的梯度计算需要进一步求解第三个积分，该积分同时依赖于状态 $z(t)$ 与伴随变量 $a(t)$

$\frac{dL}{d\theta}=-\int_{t_1}^{t_0}a(t)\frac{\partial f(z(t),t,\theta)}{\partial\theta}\,dt$ (5)

式中出现的向量-雅可比积 $a(t)\frac{\partial f}{\partial z}$ 和 $a(t)\frac{\partial f}{\partial \theta}$ 可以通过自动微分高效计算，其时间成本与一次函数 $f$ 的评估成本相近。用于求解的所有积分，都可以在一次 ODE 求解器调用中同时完成：具体而言，可以将原始状态、伴随变量以及其他偏导项拼接为一个统一的向量。大多数 ODE 求解器都支持在多个时间点输出状态 $z(t)$ 。当损失函数依赖于这些中间状态时，反向模式微分必须被拆分为一系列独立的求解过程，即在每一对相邻输出时间点之间分别进行求解。在每一个观测时刻，伴随变量都需要根据相应的偏导数 $\frac{\partial L}{\partial z(t_i)}$ 进行修正

Algorithm 1 Reverse-Mode Derivative Of An ODE Initial Value Problem

🦥 Continuous Normalizing Flows

归一化流（如 NICE）通过一系列可逆变换构造复杂分布。对于单步变换 $f$ ，根据变量替换公式，变换后的对数密度等于变换前的对数密度减去雅可比行列式绝对值的对数

$\mathbf{z}_1 = f(\mathbf{z}_0) \implies \log p(\mathbf{z}_1) = \log p(\mathbf{z}_0) - \log \left| \det \frac{\partial f}{\partial \mathbf{z}_0} \right|$ (6)

一个典型例子是平面归一化流和其对应的对数密度变化为

$\mathbf{z}(t+1) = \mathbf{z}(t) + u h(w^\top \mathbf{z}(t) + b), \quad \log p(\mathbf{z}(t+1)) = \log p(\mathbf{z}(t)) - \log \left| 1 + u^\top \frac{\partial h}{\partial \mathbf{z}} \right|$ (7)

一般而言，使用变量替换公式时的主要瓶颈在于，需要计算雅可比矩阵 $\partial f/\partial z$ 的行列式。该计算的复杂度通常会随着状态维度 $z$ 或隐藏单元数量呈三次增长。因此，已有研究通常需要在归一化流层数与计算代价之间进行权衡。将模型从离散层序列推广到连续变换后，归一化常数变化的计算反而得到简化

瞬时变量替换公式：设 $z(t)$ 是一个随时间变化的有限维连续随机变量，其概率密度为 $p(z(t))$ 。设 $\frac{dz}{dt}=f(z(t),t)$ 为描述 $z(t)$ 连续时间变换的微分方程。若 $f$ 关于 $z$ 连续可微且一致 Lipschitz 连续，并且关于 $t$ 连续，则对数概率密度的变化同样满足一个微分方程

$\frac{d \log p(\mathbf{z}(t))}{dt} = -\mathrm{tr}\left( \frac{df}{d\mathbf{z}(t)} \right)$ (8)

在连续时间设定下，只需要计算雅可比矩阵的迹。此外，与标准的有限层归一化流不同，微分方程中的动力学函数 $f$ 本身并不需要显式满足双射条件；只要 ODE 解的唯一性得到保证，整个连续变换就会自动构成双射

作为瞬时变量替换公式的一个应用例子，可以考察平面流的连续时间对应形式，以及其归一化常数的变化。给定初始分布 $p(z(0))$ ，可以通过求解这一联合 ODE，从而从 $p(z(t))$ 中采样，并同时评估样本在该时刻的概率密度

$\frac{d\mathbf{z}(t)}{dt} = u h(w^\top \mathbf{z}(t) + b), \quad \frac{d \log p(\mathbf{z}(t))}{dt} = -u^\top \frac{\partial h}{\partial \mathbf{z}(t)}$ (9)

使用多个隐藏单元且仅具有线性计算成本：虽然行列式函数并不是线性的，但迹运算是线性的。即 $\operatorname{tr}\left(\sum_n J_n\right) = \sum_n \operatorname{tr}(J_n).$ 这意味着，如果动力学函数由多个函数之和给出，那么对数密度的微分方程同样可以写成求和形式

$\frac{d\mathbf{z}(t)}{dt} = \sum_{n=1}^M f_n(\mathbf{z}(t)), \quad \frac{d\log p(\mathbf{z}(t))}{dt} = -\sum_{n=1}^M \mathrm{tr} \left( \frac{\partial f_n}{\partial \mathbf{z}} \right)$ (10)

可以以较低成本评估包含大量隐藏单元的流模型，其计算复杂度仅随隐藏单元数量 $M$ 线性增长。相比之下，在标准归一化流中评估这类宽流层通常需要 $O(M^3)$ 的计算代价，这也是传统归一化流架构往往采用多层、但每层仅包含较少隐藏单元的重要原因

时间依赖动力学：可以将流模型的参数指定为时间 $t$ 的函数，使微分方程中的动力学函数 $f(z(t),t)$ 随时间发生变化。这种参数化方式可以看作一种超网络形式。此外，作者还为每个隐藏单元引入门控机制 $\frac{dz}{dt} = \sum_n \sigma_n(t) f_n(z).$ 其中 $\sigma_n(t)\in(0,1)$ 是由神经网络学习得到的门控函数，用于决定动态函数 $f_n(z)$ 应当在何时发挥作用。作者将这类模型称为连续归一化流（Continuous Normalizing Flows）

🦫 Proof of the Instantaneous Change of Variables Theorem

瞬时变量替换公式描述连续变换下对数密度的演化。设 $z(t)$ 是一个有限维连续随机变量，其概率密度 $p(z(t))$ 随时间变化。设 $\frac{dz}{dt}=f(z(t),t)$ 为描述 $z(t)$ 连续时间变换的微分方程。若 $f$ 关于 $z$ 连续可微且一致 Lipschitz 连续，并且关于 $t$ 连续，则对数概率密度的变化同样满足如下微分方程

$\frac{d \log p(\mathbf{z}(t))}{dt} = -\mathrm{tr}\left( \frac{df}{d\mathbf{z}(t)} \right)$ (11)

𝑷𝒓𝒐𝒐𝒇：考虑 $\log p(z(t))$ 在时间上的有限变化，并取其无穷小极限。首先，将时间发生一个无穷小变化 $\varepsilon$ 时 $z$ 的变换记为

$z(t+\varepsilon)=T_{\varepsilon}(z(t))$ (12)

假设 $f$ 关于 $z(t)$ 满足 Lipschitz 连续性，并且关于时间 $t$ 连续；为使用下文的经典 Taylor 展开，再假设向量场足够光滑，所需导数在考察的轨迹邻域有界。因此，根据 Picard 存在唯一性定理，每一个初值问题都存在唯一解。同时还假设 $z(t)$ 是有界的。这些条件意味着 $f,T_{\varepsilon}$ 以及变换 $T_{\varepsilon}$ 关于 $z$ 的雅可比矩阵 $\frac{\partial T_\epsilon}{\partial z}$ 都是有界的

1️⃣ 利用离散变量替换公式以及导数的定义，写出 $\log p(z(t))$ 关于时间的变化率

$\begin{aligned} \frac{d \log p(\mathbf{z}(t))}{dt} &= \lim_{\varepsilon \to 0^+} \frac{\log p(\mathbf{z}(t)) - \log \left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right| - \log p(\mathbf{z}(t))}{\varepsilon} \\ &= - \lim_{\varepsilon \to 0^+} \frac{\log \left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right|}{\varepsilon} \\ &= - \lim_{\varepsilon \to 0^+} \frac{\frac{\partial}{\partial \varepsilon} \log \left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right|}{\frac{\partial}{\partial \varepsilon} \varepsilon} \qquad \text{(by L'Hôpital's rule)} \\ &= - \lim_{\varepsilon \to 0^+} \frac{\frac{\partial}{\partial \varepsilon} \left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right|}{\left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right|} \qquad \left( \left. \frac{\partial \log(\mathbf{z})}{\partial \mathbf{z}} \right|_{\mathbf{z}=1} = 1 \right) \\ &= - \underbrace{ \left( \lim_{\varepsilon \to 0^+} \frac{\partial}{\partial \varepsilon} \left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right| \right) }_{\text{bounded}} \underbrace{ \left( \lim_{\varepsilon \to 0^+} \frac{1}{\left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right|} \right) }_{=1} \\ &= - \lim_{\varepsilon \to 0^+} \frac{\partial}{\partial \varepsilon} \left| \det \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right| \end{aligned}$ (13)

2️⃣ 利用 Jacobi 公式来表示行列式的导数

$\begin{aligned} \frac{d \log p(\mathbf{z}(t))}{dt} &= - \lim_{\varepsilon \to 0^+} \mathrm{tr} \left( \mathrm{adj} \left( \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right) \frac{\partial}{\partial \varepsilon} \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right) \\ &= -\mathrm{tr} \left( \underbrace{ \left( \lim_{\varepsilon \to 0^+} \mathrm{adj} \left( \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right) \right) }_{=I} \left( \lim_{\varepsilon \to 0^+} \frac{\partial}{\partial \varepsilon} \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right) \right) \\ &= -\mathrm{tr} \left( \lim_{\varepsilon \to 0^+} \frac{\partial}{\partial \varepsilon} \frac{\partial}{\partial \mathbf{z}} T_\varepsilon(\mathbf{z}(t)) \right) \end{aligned}$ (14)

3️⃣ 将 $T_{\varepsilon}$ 用其 Taylor 展开代入。由于 $z(t)$ 在无穷小时间 $\varepsilon$ 内按照动力学函数 $f$ 演化，因此有

$\begin{aligned} \frac{d \log p(\mathbf{z}(t))}{dt} &= -\mathrm{tr} \left( \lim_{\varepsilon \to 0^+} \frac{\partial}{\partial \varepsilon} \frac{\partial}{\partial \mathbf{z}} (\mathbf{z} + \varepsilon f(\mathbf{z}(t), t) + \mathcal{O}(\varepsilon^2) + \mathcal{O}(\varepsilon^3) + \dots) \right) \\ &= -\mathrm{tr} \left( \lim_{\varepsilon \to 0^+} \frac{\partial}{\partial \varepsilon} \left( I + \frac{\partial}{\partial \mathbf{z}} \varepsilon f(\mathbf{z}(t), t) + \mathcal{O}(\varepsilon^2) + \mathcal{O}(\varepsilon^3) + \dots \right) \right) \\ &= -\mathrm{tr} \left( \lim_{\varepsilon \to 0^+} \left( \frac{\partial}{\partial \mathbf{z}} f(\mathbf{z}(t), t) + \mathcal{O}(\varepsilon) + \mathcal{O}(\varepsilon^2) + \dots \right) \right) \\ &= -\mathrm{tr} \left( \frac{\partial}{\partial \mathbf{z}} f(\mathbf{z}(t), t) \right) \end{aligned}$ (15)

🦨 Connection to Fokker-Planck and Liouville PDEs

Fokker–Planck 方程是一类著名的偏微分方程，用于描述随机微分方程所对应的概率密度函数如何随时间演化

本文将瞬时变量替换公式与 Fokker–Planck 方程的一个特殊情形联系起来：当扩散项为零时，Fokker–Planck 方程退化为 Liouville 方程。与瞬时变量替换公式中的设定类似，令 $z(t)\in \mathbb{R}^D$ 随时间按照 $\frac{dz(t)}{dt}=f(z(t),t)$ 动力学演化。则 Liouville 方程以偏微分方程的形式描述密度 $p(z,t)$ 在空间中某一固定位置处的变化

$\frac{\partial p(\mathbf{z}, t)}{\partial t} = -\sum_{i=1}^D \frac{\partial}{\partial \mathbf{z}_i} [f_i(\mathbf{z}, t) p(\mathbf{z}, t)]$ (16)

上式并不容易直接使用，因为它需要计算 $\frac{\partial p(z,t)}{\partial z}$ 这类关于空间变量的偏导数，而这些偏导通常需要通过有限差分方法近似。因此，不再在空间中的固定点上评估 $p(\cdot,t)$ ，而是沿着某个粒子的运动轨迹 $z(t)$ 来考察密度的变化

$\begin{aligned} \frac{d}{dt}p(\mathbf z(t),t) &=\nabla_{\mathbf z}p(\mathbf z(t),t)^\top f(\mathbf z(t),t) +\left.\frac{\partial p(\mathbf z,t)}{\partial t}\right|_{\mathbf z=\mathbf z(t)}\\ &=\sum_i f_i\frac{\partial p}{\partial z_i} -\sum_i\left(p\frac{\partial f_i}{\partial z_i}+f_i\frac{\partial p}{\partial z_i}\right)\\ &=-p(\mathbf z(t),t)\sum_i\frac{\partial f_i(\mathbf z(t),t)}{\partial z_i}. \end{aligned}$ (17)

在密度为正处应用对数的链式法则，即可得到瞬时变量替换公式

$\frac{d \log p(\mathbf{z}(t), t)}{dt} = \frac{1}{p(\mathbf{z}(t), t)} \frac{d p(\mathbf{z}(t), t)}{dt} = - \sum_{i=1}^D \frac{\partial f_i(\mathbf{z}(t), t)}{\partial \mathbf{z}_i}$ (18)

可以与状态轨迹 $z(t)$ 结合，形成一个维度为 $D+1$ 的常微分方程系统

$\frac{d}{dt} \begin{bmatrix} \mathbf{z}(t) \\ \log p(\mathbf{z}(t), t) \end{bmatrix} = \begin{bmatrix} f(\mathbf{z}(t), t) \\ - \sum_{i=1}^D \frac{\partial f_i(\mathbf{z}(t), t)}{\partial \mathbf z_i} \end{bmatrix}$ (19)

与 Fokker–Planck 方程和 Liouville 方程相比，瞬时变量替换公式在实际计算中更具实用性。原因在于，它可以被更容易地数值求解：只需要在跟踪轨迹 $z(t)$ 的基础上额外维护一个标量状态，即对数密度 $\log p(z(t),t)$ 。相比之下，基于有限差分近似求解 Liouville 方程的方法通常需要在整个状态空间上构造网格，而网格规模会随维度 $D$ 呈指数级增长

🐇 A Modern Proof of the Adjoint Method

令 $z(t)$ 满足微分方程 $\frac{dz(t)}{dt}=f(z(t),t,\theta),$ 其中 $\theta$ 表示模型参数。证明：若定义伴随状态为 $a(t)=\frac{dL}{dz(t)},$ 则该伴随状态满足如下微分方程

$\frac{d\mathbf{a}(t)}{dt} = -\mathbf{a}(t) \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}(t)}$ (20)

伴随状态可以理解为损失函数 $L$ 关于指定时刻 $t$ 处隐藏状态的梯度。在标准神经网络中，隐藏层 $h_t$ 的梯度可由下一层 $h_{t+1}$ 的梯度通过链式法则得到： $\frac{dL}{d\mathbf{h}_t} = \frac{dL}{d\mathbf{h}_{t+1}} \frac{d\mathbf{h}_{t+1}}{d\mathbf{h}_t}.$ 对于连续隐藏状态而言，经过一个无穷小时间增量 $\varepsilon$ 后，状态的变化可以表示为

$\mathbf{z}(t+\varepsilon)=\mathbf{z}(t)+\int_t^{t+\varepsilon} f(\mathbf{z}(\tau),\tau,\theta)\,d\tau=T_\varepsilon(\mathbf{z}(t),t)$ (21)

$\frac{dL}{d\mathbf{z}(t)} = \frac{dL}{d\mathbf{z}(t+\varepsilon)} \frac{d\mathbf{z}(t+\varepsilon)}{d\mathbf{z}(t)} \qquad \text{or} \qquad \mathbf{a}(t) = \mathbf{a}(t+\varepsilon) \frac{\partial T_\varepsilon(\mathbf{z}(t), t)}{\partial \mathbf{z}(t)}$ (22)

从导数定义出发证明伴随方程

$\begin{aligned} \frac{d\mathbf{a}(t)}{dt} &= \lim_{\varepsilon \to 0^+} \frac{\mathbf{a}(t+\varepsilon) - \mathbf{a}(t)}{\varepsilon} \\ &= \lim_{\varepsilon \to 0^+} \frac{\mathbf{a}(t+\varepsilon) - \mathbf{a}(t+\varepsilon) \frac{\partial}{\partial \mathbf{z}(t)} T_\varepsilon(\mathbf{z}(t))}{\varepsilon} \qquad \text{(chain rule)} \\ &= \lim_{\varepsilon \to 0^+} \frac{\mathbf{a}(t+\varepsilon) - \mathbf{a}(t+\varepsilon) \frac{\partial}{\partial \mathbf{z}(t)} (\mathbf{z}(t) + \varepsilon f(\mathbf{z}(t), t, \theta) + \mathcal{O}(\varepsilon^2))}{\varepsilon} \qquad \text{(Taylor series around } \mathbf{z}(t)\text{)} \\ &= \lim_{\varepsilon \to 0^+} \frac{\mathbf{a}(t+\varepsilon) - \mathbf{a}(t+\varepsilon) \left( I + \varepsilon \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}(t)} + \mathcal{O}(\varepsilon^2) \right)}{\varepsilon} \\ &= \lim_{\varepsilon \to 0^+} \frac{-\varepsilon \mathbf{a}(t+\varepsilon) \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}(t)} + \mathcal{O}(\varepsilon^2)}{\varepsilon} \\ &= \lim_{\varepsilon \to 0^+} -\mathbf{a}(t+\varepsilon) \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}(t)} + \mathcal{O}(\varepsilon) \\ &= -\mathbf{a}(t) \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}(t)} \end{aligned}$ (23)

在离散神经网络中，梯度从最后一层向前逐层传播；类似地，在 ODE 系统中，伴随状态也需要沿时间反方向求解。若损失函数仅依赖于最后时刻 $t_N$ 的状态，则可在终点处设定伴随状态的初始条件 $a(t_N) = \frac{dL}{dz(t_N)}.$ 随后，通过沿时间反向积分伴随方程，可以得到任意时刻的状态梯度，包括初始时刻 $t_0$ 处的梯度

$\underbrace{\mathbf{a}(t_N) = \frac{dL}{d\mathbf{z}(t_N)}}_{\text{initial condition of adjoint diffeq.}},\qquad \underbrace{\mathbf{a}(t_0) = \mathbf{a}(t_N) + \int_{t_N}^{t_0} \frac{d\mathbf{a}(t)}{dt} dt = \mathbf{a}(t_N) - \int_{t_N}^{t_0} \mathbf{a}(t) \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}(t)}\,dt}_{\text{gradient wrt. initial value}}$ (24)

这里假设损失函数 $L$ 仅依赖于最后一个时间点 $t_N$ 。如果 $L$ 还依赖于中间时刻 $t_1,t_2,\ldots,t_{N-1}$ 的状态，则可以在每个相邻时间区间上重复执行伴随反向传播步骤，即依次在区间 $[t_{N-1},t_N]$ 、 $[t_{N-2},t_{N-1}]$ 等上反向求解，并将各区间得到的梯度进行累加。接下来，可以将上述伴随方程推广到参数 $\theta$ 以及积分起止时间 $t_0$ 和 $t_N$ 的梯度计算。将 $\theta$ 视为关于时间不变的状态变量，同时将时间 $t$ 本身也作为一个状态变量。于是有

$\frac{\partial \theta(t)}{\partial t} = \mathbf{0} \qquad \frac{dt(t)}{dt} = 1$ (25)

将它们与 $z(t)$ 组合起来，可以构造增广状态及其对应的微分方程

$\frac{d}{dt} \begin{bmatrix} \mathbf{z} \\ \theta \\ t \end{bmatrix} (t) = f_{aug}([\mathbf{z}, \theta, t]) := \begin{bmatrix} f([\mathbf{z}, \theta, t]) \\ \mathbf{0} \\ 1 \end{bmatrix} , \mathbf{a}_{aug} := \begin{bmatrix} \mathbf{a} & \mathbf{a}_\theta & \mathbf{a}_t \end{bmatrix} , \mathbf{a}_\theta(t) := \frac{dL}{d\theta(t)} , \mathbf{a}_t(t) := \frac{dL}{dt(t)}$ (26)

此时，增广动力学函数的雅可比矩阵具有如下形式

$\frac{\partial f_{aug}}{\partial [\mathbf{z}, \theta, t]} = \begin{bmatrix} \frac{\partial f}{\partial \mathbf{z}} & \frac{\partial f}{\partial \theta} & \frac{\partial f}{\partial t} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & \mathbf{0} \end{bmatrix} (t)$ (27)

其中每一个 0 都表示维度相匹配的零矩阵。将该雅可比矩阵代入伴随方程，可以得到

$\frac{d\mathbf{a}_{aug}(t)}{dt} = - [\mathbf{a}(t) \quad \mathbf{a}_\theta(t) \quad \mathbf{a}_t(t)] \frac{\partial f_{aug}}{\partial [\mathbf{z}, \theta, t]} (t) = - [\mathbf{a} \frac{\partial f}{\partial \mathbf{z}} \quad \mathbf{a} \frac{\partial f}{\partial \theta} \quad \mathbf{a} \frac{\partial f}{\partial t}] (t)$ (28)

其中，第一项正是前面已经得到的伴随微分方程；第二项可以用于计算损失函数关于参数的总梯度。这里固定初始状态，且损失不直接依赖参数。通过在整个时间区间上积分该项，并设定终点条件 $a_{\theta}(t_N)=0$

$\frac{dL}{d\theta} = \mathbf{a}_\theta(t_0) = - \int_{t_N}^{t_0} \mathbf{a}(t) \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \theta} dt$ (29)

最后，还可以得到损失函数关于积分终止时间 $t_N$ 和起始时间 $t_0$ 的梯度。这里固定初始状态，且损失不显式依赖时间；时间伴随的终端值须取终止时间梯度的负值

$\begin{aligned} \frac{dL}{dt_N}&=\mathbf a(t_N)f(\mathbf z(t_N),t_N,\theta),\qquad \mathbf a_t(t_N)=-\frac{dL}{dt_N},\\ \frac{dL}{dt_0}&=\mathbf a_t(t_0) =\mathbf a_t(t_N)-\int_{t_N}^{t_0}\mathbf a(t)\frac{\partial f(\mathbf z(t),t,\theta)}{\partial t}\,dt\\ &=-\mathbf a(t_0)f(\mathbf z(t_0),t_0,\theta). \end{aligned}$ (30)

因此，结合状态伴随方程、初始状态梯度、参数梯度以及时间边界梯度的表达式，就可以得到初值问题 ODE 求解器所有可能输入变量的梯度。伴随敏感度方法不仅能够反向传播到初始状态，还能够同时处理模型参数以及积分区间端点的梯度，从而为连续深度模型的端到端训练提供完整的微分机制

Algorithm 2 Complete Reverse-Mode Derivative Of An ODE Initial Value Problem

🐼 FFJORD

可逆生成模型通过一个易于求逆的神经网络，将来自简单分布的样本点映射到复杂分布

典型代表包括 NICE、Real NVP 和 Glow。这类模型易于采样，并且可以通过变量替换公式进行最大似然训练。然而，为避免计算雅可比行列式时出现 $O(D^3)$ 的高昂代价，模型架构往往必须受到一些不自然的限制，例如对维度进行划分，或者使用秩一权重矩阵
近期，Chen et al. 提出连续归一化流，使用常微分方程来定义从潜变量到数据的映射。在该模型中，似然可以通过迹运算进行计算，其计算复杂度为 $O(D^2)$ 。这使得模型可以采用更加灵活的网络结构，但这种灵活性仍然是受限的

作者利用 Hutchinson 迹估计器，对对数密度给出一个可扩展的无偏估计。由此得到的模型是一类连续时间可逆生成模型，它能够估计对数密度，并通过一次 ODE 求解完成采样，同时允许使用不受特殊结构限制的神经网络架构。这里的无偏性针对精确积分下的对数密度；指数化后的密度估计一般不再无偏。将该方法称为 Free-form Jacobian of Reversible Dynamics，简称 FFJORD

🐿️ Background: Generative Models and Change of Variables

不同于直接参数化一个已经归一化的分布，利用变量替换公式，可以通过可逆函数 $f: \mathbb{R}^D \to \mathbb{R}^D$ ，将一个归一化的基础分布 $p_z(z)$ 进行变换，从而隐式地定义一个复杂的归一化分布 $p_x(x)$ 。给定随机变量 $z \sim p_z(z)$ ，令 $x = f(z)$ ，则 $x$ 的对数密度满足

$\log p_x(\mathbf{x}) = \log p_z(\mathbf{z}) - \log \left| \det \frac{\partial f(\mathbf{z})}{\partial \mathbf{z}} \right|$ (1)

其中 $\frac{\partial f(z)}{\partial z}$ 是函数 $f$ 的雅可比矩阵。一般而言，计算该雅可比矩阵的对数行列式需要 $O(D^3)$ 的时间复杂度。大量研究致力于设计受限的神经网络结构，使雅可比行列式的计算更加可处理。这些方法大体可以分为以下三类

早期平面流等变分归一化流：通过限制函数 $f$ 的具体形式，可以利用多种行列式恒等式来降低计算复杂度。这类模型通常缺少易于计算的解析逆映射，因而难以同时高效地直接最大似然训练和采样，其逆映射记为 $f^{-1}$ 。然而，它们在变分推断中可用于构造近似后验分布
自回归变换：通过采用自回归模型，并为各个维度指定一个顺序，可以使函数 $f$ 的雅可比矩阵被约束为下三角矩阵。逆变换需要对 $f$ 进行 $D$ 次顺序计算；当维度 $D$ 较大时，这一代价往往难以承受
分块变换：通过对维度进行划分，并在不同部分之间使用仿射变换，可以使雅可比行列式的计算变得非常高效，同时逆映射 $f^{-1}$ 的计算成本也可以与正向映射 $f$ 保持一致。这类方法允许使用卷积网络结构，在图像数据的密度估计任务中表现优异

🪺 Continuous Normalizing Flows

Chen et al.（2018）定义一类针对数据 $\mathbf{x}\in \mathbb{R}^D$ 的生成模型

该模型不再直接参数化一个显式的可逆变换映射，而是通过求解连续时间动力系统来隐式定义从基础分布到数据分布的变换
生成过程首先从一个基础分布中采样： $\mathbf{z}_0 \sim p_{\mathbf{z}_0}(\mathbf{z}_0).$ 随后，给定一个由参数化函数定义的常微分方程 $\frac{\partial \mathbf{z}(t)}{\partial t} = f(\mathbf{z}(t), t; \theta),$ 并以 $\mathbf{z}(t_0)=\mathbf{z}_0$ 作为初始条件求解该初值问题，从而得到数据样本 $\mathbf{x}=\mathbf{z}(t_1).$ 这类模型被称为连续归一化流。在该模型下，对数密度的变化满足另一个微分方程，该方程被称为瞬时变量替换公式

$\frac{d \log p(\mathbf{z}(t))}{dt} = -\mathrm{Tr}\left(\frac{\partial f}{\partial \mathbf{z}(t)}\right).$ (2)

在连续时间动力学的演化过程中，样本点对数密度的瞬时变化率由向量场 $f$ 关于状态变量 $\mathbf{z}(t)$ 的雅可比矩阵的迹决定。因此，通过沿时间积分，可以得到对数密度的总变化量

$\log p(\mathbf{z}(t_1)) = \log p(\mathbf{z}(t_0)) - \int_{t_0}^{t_1} \mathrm{Tr}\left(\frac{\partial f}{\partial \mathbf{z}(t)}\right) dt.$ (3)

给定一个数据点，通过求解一个联合初值问题，同时计算出生成该数据点的初始潜变量，以及模型赋予该数据点的对数似然

$\underbrace{\begin{bmatrix} \mathbf{z}_0 \\ \log p(\mathbf{x}) - \log p_{\mathbf{z}_0}(\mathbf{z}_0) \end{bmatrix}}_{\text{solutions}} = \begin{bmatrix} \mathbf{x} \\ 0 \end{bmatrix} + \underbrace{\int_{t_1}^{t_0} \begin{bmatrix} f(\mathbf{z}(t), t; \theta) \\ \mathrm{Tr}\left(\frac{\partial f}{\partial \mathbf{z}(t)}\right) \end{bmatrix} dt}_{\text{dynamics}}, \qquad \underbrace{\begin{bmatrix} \mathbf{z}(t_1) \\ \log p(\mathbf{x}) - \log p(\mathbf{z}(t_1)) \end{bmatrix} = \begin{bmatrix} \mathbf{x} \\ 0 \end{bmatrix}}_{\text{initial values}}$ (4)

该系统联合积分状态变量 $\mathbf{z}(t)$ 及其对数密度的变化。给定数据 $\mathbf{x}=\mathbf{z}(t_1)$ ，从 $t_1$ 反向积分至 $t_0$ ，得到潜变量及相应的对数密度差。将该差值与基础分布对数密度 $\log p_{\mathbf{z}_0}(\mathbf{z}_0)$ 相加，即得数据 $\mathbf{x}$ 的对数似然 $\log p(\mathbf{x})$
对状态 ODE，函数 $f$ 关于状态局部 Lipschitz、关于时间连续即可保证局部解存在且唯一；在整个积分区间存在解还需排除有限时间爆炸。为应用经典密度公式并计算梯度，通常进一步要求状态可微。采用有限权重、平滑 Lipschitz 激活函数（如 softplus 或 tanh）的常见前馈网络可满足相应条件

CNF 的训练目标是最大化对数似然。该目标依赖于一个初值问题的解，而该初值问题的动力学函数由参数 $\theta$ 所刻画。对于任意一个作用在初值问题解上的标量损失函数，设其形式为

$L(\mathbf{z}(t_1)) = L\left(\mathbf z(t_0)+\int_{t_0}^{t_1} f(\mathbf{z}(t), t; \theta)dt\right)$ (5)

其中 $\mathbf{z}(t_1)$ 是由动力系统从初始状态 $\mathbf{z}(t_0)$ 演化得到的终端状态。Pontryagin 表明，该损失函数关于参数 $\theta$ 的导数同样可以表示为一个初值问题的解，其形式为

$\frac{dL}{d\theta} = -\int_{t_1}^{t_0} \frac{\partial L}{\partial \mathbf{z}(t)} \frac{\partial f(\mathbf{z}(t), t; \theta)}{\partial \theta} dt.$ (6)

$\frac{\partial L}{\partial \mathbf{z}(t)}$ 被称为该 ODE 的伴随状态。Chen 首先使用一个黑箱 ODE 求解器计算终端状态 $\mathbf{z}(t_1)$ ，随后再调用一次 ODE 求解器，从终端处的梯度 $\frac{\partial L}{\partial \mathbf{z}(t_1)}$ 作为初始条件，反向求解所对应的伴随系统，从而计算损失函数关于参数 $\theta$ 的梯度

🦨 Scalable Density Evaluation with Unrestricted Architectures

从离散时间动力学切换到连续时间动力学，可以将归一化流的主要计算瓶颈从 $O(D^3)$ 降低到 $O(D^2)$ ，代价是需要引入数值 ODE 求解器。这一变化使模型能够采用更具表达能力的网络结构

在 Rezende 与 Mohamed（2015）提出的原始归一化流模型中，每一层变换只是一个仅包含单个隐藏单元的一层神经网络
在平面连续归一化流中，Chen et al.（2018）所使用的瞬时变换是一个包含多个隐藏单元的一层神经网络

当单次网络求值成本随维度线性增长时，精确计算雅可比迹 $\operatorname{Tr}\left(\frac{\partial f}{\partial \mathbf{z}(t)}\right)$ 需要 $O(D^2)$ 的计算量。逐行反向自动微分的总成本约相当于对 $f$ 进行 $D$ 次求值：每次反向计算得到 $f$ 的一行雅可比，并从中取出一个对角元素。若单次网络求值成本为 $C_f$ ，则精确求迹的总成本为 $O(D C_f)$ 。FFJORD 利用以下两个性质降低计算成本

反向模式自动微分可直接计算向量—雅可比积 $\mathbf{v}^T \frac{\partial f}{\partial \mathbf{z}}$ ，其计算成本与一次 $f$ 前向求值同阶
Hutchinson 恒等式将矩阵迹写为随机二次型的期望 $\operatorname{Tr}(A) = \mathbb{E}_{p(\epsilon)} \left[ \epsilon^T A \epsilon \right]$ ，其中噪声分布 $p(\epsilon)$ 满足 $\mathbb{E}[\epsilon]=0, \operatorname{Cov}(\epsilon)=I.$ 采样一个噪声向量并计算对应的二次型，即得到矩阵迹的单样本无偏估计

为保证每次调用 ODE 求解器时动力学是确定性的，可以在一次 ODE 求解过程中固定同一个噪声向量 $ϵ$ 。在精确积分且可交换期望与积分的条件下，固定噪声不改变对数密度估计的无偏性；实际数值求解仍有积分误差。常用的噪声分布包括标准高斯分布或 Rademacher 分布

$\begin{aligned} \log p(\mathbf{z}(t_1)) &= \log p(\mathbf{z}(t_0)) - \int_{t_0}^{t_1} \mathrm{Tr} \left( \frac{\partial f}{\partial \mathbf{z}(t)} \right) dt \\ &= \log p(\mathbf{z}(t_0)) - \int_{t_0}^{t_1} \mathbb{E}_{p(\boldsymbol{\epsilon})} \left[ \boldsymbol{\epsilon}^T \frac{\partial f}{\partial \mathbf{z}(t)} \boldsymbol{\epsilon} \right] dt \\ &= \log p(\mathbf{z}(t_0)) - \mathbb{E}_{p(\boldsymbol{\epsilon})} \left[ \int_{t_0}^{t_1} \boldsymbol{\epsilon}^T \frac{\partial f}{\partial \mathbf{z}(t)} \boldsymbol{\epsilon} dt \right] \end{aligned}$ (7)

在动力学网络的结构中，常常会存在瓶颈层，即某些隐藏层的宽度小于输入维度。可以利用迹的循环不变性，将 Hutchinson 估计转移到瓶颈空间。用于估计 $\operatorname{Tr}(A)$ 的估计器，其方差与矩阵的对称部分及噪声分布有关，可用 $\|A\|_F^2$ 给出上界。转移到低维空间可能降低方差，但并无普遍保证；瓶颈空间须使用相应维度、零均值且单位协方差的噪声。若将动力学函数看作两个函数的复合 $f = g \circ h(\mathbf{z}),$ 则有

$\underbrace{\mathrm{Tr} \left( \frac{\partial f}{\partial \mathbf{z}} \right)}_{D \times D} = \underbrace{\mathrm{Tr} \left( \frac{\partial g}{\partial h} \frac{\partial h}{\partial \mathbf{z}} \right)}_{D \times D} = \underbrace{\mathrm{Tr} \left( \frac{\partial h}{\partial \mathbf{z}} \frac{\partial g}{\partial h} \right)}_{H \times H} = \mathbb{E}_{p(\boldsymbol{\epsilon})} \left[ \boldsymbol{\epsilon}^T \frac{\partial h}{\partial \mathbf{z}} \frac{\partial g}{\partial h} \boldsymbol{\epsilon} \right].$ (8)

作者的完整方法结合连续时间动力学和基于 Hutchinson 迹估计器的高效对数似然估计方法，从而构造出一种可扩展的可逆生成模型。与以往方法不同，该模型不要求雅可比矩阵具有特殊结构约束，因此可以使用更加自由的神经网络架构。将该方法命名为 FFJORD

设一次计算动力学函数 $f$ 的代价为 $O(DH)$ ，其中 D 表示数据维度，H 表示 $f$ 中隐藏层宽度；这里采用原论文固定深度、简化网络成本的比较，深网络的内部层成本还应另计。对未利用特殊雅可比结构的有限流，重复 L 个可逆变换并显式计算稠密行列式的简化代价为 $O((DH + D^3)L).$
沿用该简化计数，在连续归一化流 CNF 中，主要密度计算由行列式改为迹，从 $O(D^3)$ 降低到 $O(D^2)$ 。由于 ODE 求解器在积分过程中需要多次调用动力学函数 $f$ ，若调用次数记为 $\hat{L}$ ，则 CNF 的似然计算代价为 $O((DH + D^2)\hat{L})$
FFJORD 使用 Hutchinson 估计器，避免逐项计算雅可比矩阵的对角元素，将迹估计的维度相关成本从 $O(D^2)$ 降至 $O(D)$ 。计入一次向量—雅可比积后，整体对数似然估计成本为 $O((DH + D)\hat{L}).$

Algorithm 1 Unbiased Stochastic Log-Density Estimation Using The FFJORD Model