从Neural ODE到连续Normalizing Flows

Neural ODE通过神经网络建模隐藏状态变换的导数，从而构建隐藏状态连续变换模型。Neural ODE的建模方法应用归一化流模型，就可把离散归一化流变为连续归一化流，显著提升了模型表达力。

Neural ODE

循环神经网络RNN通过组合隐藏状态的变化构建复杂的变换，即

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

式(1)中$t\in\{0\ldots T\}$，且$\mathbf{h}_t\in\mathbb{R}^D$

RNN可以理解为隐藏状态的离散演化，其数学形式为差分方程。若构建一个隐藏状态连续变换的模型，那么数学形式应为微分方程，即

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

该常微分方程可通过神经网络构建。在《Neural Ordinary Differential Equations》中，作者们利用伴随敏感度方法计算网络的参数。

这种通过ODE求解器定义与评估模型，具有内存效率、自适应计算、可扩展和可逆归一化流的产生、以及连续时间序列模型构建的特性。

确切的说，神经网络建模的是离散型和连续型的隐藏状态动力学。由于离散型通过差分方程建模，因此可直接计算时刻$T$的隐藏层状态。与之不同，隐藏层需要通过ODE求解器才能获得隐藏层状态。

Normalizing Flows

在归一化流中，式(1)被用于表达变量的变换，用于构建离散的归一化流。

若$\mathbf{z}(t)$通过双射函数变为$\mathbf{z}(t+1)$，即

$$ \begin{aligned} \mathbf{z}(t+1)=\mathbf{z}(t)+uh(w^T\mathbf{z}(t)+b) \end{aligned}\tag{5} $$

那么，根据change of variable理论，$p(\mathbf{z}(t+1))$概率函数的为

$$ \begin{aligned} logp(\mathbf{z}(t+1))=logp(\mathbf{z}(t))-log\vert1+u^T\frac{\partial h}{\partial\mathbf{z}}\vert \end{aligned}\tag{6} $$

该离散归一化流被称为planar归一化流。离散归一化流的主要缺点是雅可比行列式的计算复杂度为$\mathbf{z}$维度的三次方。

Change of Variables

考虑可逆平滑映射$f:R^d\to R^d$，其逆函数为$f^{-1}=g$. 若利用该映射概率密度函数为$q(\mathbf{z})$的随机变量映射为${\mathbf{z}}'=f(\mathbf{z})$，那么

$$ \begin{aligned} E({z}') =& \int_{-\infty}^{\infty}{z}'q({z}')d{z}' \\ =& \int_{-\infty}^{\infty}f(z)q(z)dz \end{aligned}\tag{3} $$

由于映射$f$为可逆平滑映射，即属于单调函数，那么

$$ \begin{aligned} {z}'q({z}')d{z}' = & f(z)q(z)dz \\ f(z)q({z}')d{z}' = & f(z)df^{-1}({z}') \\ f(z)q({z}')d{z}' = & f(z)q(z)\vert det \frac{\partial{f^{-1}}}{\partial{z}'}\vert d{z}' \end{aligned} $$

即$q({z}')=q(z)\vert det \frac{\partial{f^{-1}}}{\partial{z}'}\vert$

根据逆函数的Jacobians特性，可知，$q({z}')=q(z)\vert det \frac{\partial{f}}{\partial{z}}\vert^{-1}$

其中，由于$dz,d{z}'$表示微体积，则$dz=\vert \frac{d}{d{z}'}f^{-1}({z}')\vert d{z}'$

Continuous Normalizing Flows

通过Neural ODE构建的连续型归一化流显著降低了计算复杂度。这是因为在$f$满足$\mathbf{z}$中均匀Lipschitz连续和$t$中连续的条件下，其变量log概率的变换满足一个常微分方程，其只需要计算矩阵的迹，被称为即时的变量变换，即

$$ \begin{aligned} \frac{\partial logp(\mathbf{z}(t))}{\partial t}=-tr(\frac{df}{d\mathbf{z}(t)}) \end{aligned}\tag{4} $$

对应的，planar归一化流的连续形式为

$$ \begin{aligned} \frac{d\mathbf{z}(t)}{dt}=uh(w^{T}\mathbf{z}(t)+b),~\frac{\partial logp(\mathbf{z}(t))}{\partial t}=-u^{T}\frac{\partial h}{\partial\mathbf{z}(t)} \end{aligned}\tag{5} $$

给定初试条件分布$p(\mathbf{z}(0))$，那么可通过求解ODE(4)可得概率密度$p(\mathbf{z}(t))$

连续归一化流的Instantaneous Change of Variables数学推导可见Neural ODE论文或An introduction to Flow Matching。

最后，为了估计数据分布的似然，流模型的损失函数均为最大化似然。