Neural ODEs

因为想学习Neural ODE视角下的diffusion方法，所以先从入门Neural ODE开始。本篇文章从ODE的定义开始，通过Euler数值法求解ODE与ResNet架构之间的关系引出Neural ODE，并推导其训练所需的reverse-time算法。

一、Ordinary Differential Equation (ODE)

常微分方程（ordinary differential equation, ODE）是未知函数只含有一个自变量的微分方程，例如简单的一阶常微分方程有以下形式： $y^{'} (t) = f (t, y), y (t_{0}) = y_{0}$ 其中 $y (t)$ 表示 $y$ 是以 $t$ 为自变量的函数， $y^{'} (t)$ 是其导数，通常会需要给出某点的初值 $y_{0}$ ，才能解出 $y (t)$ 。当 $f (x, y)$ 取某些特定形式的时候，我们可以得到 $y (t)$ 的解析解，例如 $f (t, y) = y (t)$ ，我们知道 $y (t) = A e^{t}$ 是满足条件的，其中 $A$ 为待定系数，通过代入给定的初值即可得到。但是大多数时候， $f (t, y)$ 比较复杂，无法求得解析解。幸运的是，很多场景中我们的目标只是能够在给定 $t$ 时获得 $y (t)$ 的取值即可，因此只需要数值解。在给定初值时，任意一点 $t$ 处的 $y (t)$ 取值可以由积分得到： $y (t) = y (t_{0}) + \int_{t_{0}}^{t} y^{'} (t) d t = y_{0} + \int_{t_{0}}^{t} f (t, y) d t$

二、Euler法与ResNet

ODE的数值解法是已经被广泛探索的领域，本质就是各种求 $\int_{t_{0}}^{t} f (t, y) d t$ 的数值方法，但之所以想到用ODE来进行数据建模，还是因为最基本的Euler法。Euler法把积分区间 $[t_{0}, t]$ 平均分成若干片段，把积分过程变为在切线方向上逐步前进的过程（类似将区间分割为若干足够小的分区间求和的数值积分方法）。

令片段长度为 $δ$ ，则 $t_{n + 1} = t_{n} + δ$ ，由一阶近似方法可以得到： $y (t_{n + 1}) = y (t_{n}) + (t_{n + 1} - t_{n}) y^{'} (t_{n}) = y (t_{n}) + δ f (t, y (t_{n}))$ 显然， $h$ 越小，即分割的片段越多，则得到的结果越精确。而如果我们仔细看递推表达式，会发现与ResNet的定义形式非常相似： $\begin{aligned} Euler Method : & y (t_{n + 1}) = y (t_{n}) + δ f (t, y (t_{n})) \\ ResNet : & h_{t + 1} = h_{t} + f (θ_{t}, h_{t}) \end{aligned}$ 其中 $h_{t}$ 为第 $t$ 层的隐状态， $f$ 是前向网络， $θ_{t}$ 为第 $t$ 层的网络权重。如此相似的形式告诉我们ResNet一定程度上可以看作是在用Euler法求解ODE初值问题，其中积分域 $t$ 是离散的，在ResNet中代表网络的深度。但有一点不同的是，ResNet每层的权重不一样，对应的ODE中的 $f$ 每层都不同，因此只能说是形式上的相近。

三、Neural ODE

1. 定义

那很自然的，我们会想到用ODE来做更高层次的抽象，用神经网络来表达 $f (t, y (t_{n}))$ ，或者说 $f (t, h (t_{n}), θ)$ ，则整个过程可以表示为： $y = ODESolver (h (t_{0}), f, t_{0}, t_{1}, θ), h (t_{0}) = embed (x)$ 其中 $h (t_{0})$ 为初值条件，是根据输入 $x$ 变换得到的（e.g. embedding过程）， $t_{0}$ 与 $t_{1}$ 是积分上下限的超参（但是也可以做成可训练的）， $θ$ 指模型 $f$ 的参数。而 $ODESolver$ 则可以使用各种各样的求解器，前文所述的Euler法就是最基础的一种。使用ODE的视角做建模有两个非常显著的优势：

Memory efficient：ODE视角只需要用一个网络来建模导数，因此只需要一组参数，而ResNet这样的则是每层都有独立的参数。另外，在训练中，ODE的reverse算法与前向算法一致，都只需求解一次ODE，因此也不需要存储中间状态（下文再展开），可以大幅减少训练的内存需求。
Adaptive computation：很多先进的 $ODESolver$ 都能对递推步数（一定程度上可以理解为模型深度）根据 $f$ 的各种性质进行动态调整，因此模型的计算复杂度能够自适应地根据问题的复杂度调整。

2. 训练

如果仍然使用backpropagation算法计算梯度来进行训练的话，内存占用会非常大，因为需要记录每一步的中间输出，而且传导过程可能与具体采用的 $ODESolver$ 算法有关。而伴随灵敏度方法（adjoint sensitivity method）可以很优雅地解决这个问题，它将梯度计算的流程划归到与前向计算一致，都只需求解一次ODE，因此不仅内存占用小（不用存中间态），同时也与 $ODESolver$ 的算法选择解藕。

Adjoint sensitivity method

整个过程的大致想法是通过伴随状态（adjoint state）将计算梯度的过程也变为求解ODE的过程，再将伴随状态的初值与原问题的初值一起输入到 $ODESolver$ 中，即可在前向求解过程中将梯度与结果一同输出。

对初值的梯度

令最终的loss为 $L$ ，我们先考虑如何求解 $\frac{\partial L}{\partial h (t_{0})}$ （该梯度会用于对输入变换 $g$ 的训练）。首先定义伴随状态（adjoint state）： $a (t) = \frac{\partial L}{\partial h (t)}$ 接下来我们需要得到关于 $a (t)$ 的ODE。当 $t$ 产生 $δ$ 的变化时，我们有： $a (t) = \frac{\partial L}{\partial h (t + δ)} \frac{\partial h (t + δ)}{\partial h (t)} = a (t + δ) \frac{\partial h (t + δ)}{\partial h (t)}$ 根据 $h (t + δ) = h (t) + \int_{t}^{t + δ} f (s, h (s), θ) d s$ ，我们可以进一步得到： $\begin{aligned} a (t) & = a (t + δ) \frac{\partial}{\partial h (t)} (h (t) + \int_{t}^{t + δ} f (s, h (s), θ) d s) \\ = a (t + δ) [1 + \frac{\partial}{\partial h (t)} (\int_{t}^{t + δ} f (s, h (s), θ) d s)] \end{aligned}$ 因此： $\begin{aligned} a^{'} (t) & = lim_{δ \to 0^{+}} \frac{a (t + δ) - a (t)}{δ} \\ = lim_{δ \to 0^{+}} \frac{- a (t + δ) \frac{\partial}{\partial h (t)} (\int_{t}^{t + δ} f (s, h (s), θ) d s)}{δ} \\ = - a (t) \frac{\partial f (t, h (t), θ)}{\partial h (t)} \end{aligned}$ 由于 $a (t_{1})$ 即为 $L$ 关于网络最终输出 $h (t_{1})$ 的梯度，是很容易求的，因此可以作为ODE问题的初值，从而得到： $a^{'} (t) = - a (t) \frac{\partial f (t, h (t), θ)}{\partial h (t)}, a (t_{1}) = \frac{\partial L}{\partial h (t_{1})}$ 则对于 $h (t_{0})$ 的梯度为： $a (t_{0}) = \frac{\partial L}{\partial h (t_{0})} = a (t_{1}) + \int_{t_{1}}^{t_{0}} - a (t) \frac{\partial f (t, h (t), θ)}{\partial h (t)} d t$ 对 $θ$ 以及 $t$ 的梯度

类似的，我们定义对于权重 $θ$ 以及 $t$ 的伴随状态： $a_{θ} (t) = \frac{\partial L}{\partial θ (t)} ， a_{t} (t) = \frac{\partial L}{\partial t}$ 我们认为权重不随时间而变化，因此有： $\frac{d θ (t)}{d t} = 0, \frac{d t}{d t} = 1$ 接下来的想法非常巧妙，我们之前已经分析了如何通过 $a (t)$ 求解 $\frac{\partial L}{\partial h (t)}$ ，而 $f (t, h (t), θ)$ 输出的是 $h^{'} (t)$ ，那如果我们对 $f$ 的输出进行扩增，使其同时也能输出 $θ^{'} (t)$ 和 $(t)^{'}$ ，就可以直接把求解 $\frac{\partial L}{\partial h (t)}$ 的过程中一元微积分的内容替换成多元微积分，从而得到结果。即： $[\frac{d h (t)}{d t}, \frac{d θ (t)}{d t}, \frac{d t}{d t}] = f_{a u g} (t, h (t), θ) = [f (t, h (t), θ), 0, 1]$ 类似的，对伴随状态进行扩增： $a_{a u g} (t) = [a (t), a_{θ} (t), a_{t} (t)]$ 扩增后的向量可以看成只是对之前的向量增加了若干维度，因此仍然可以用之前推导的结果： $\begin{aligned} a_{a u g}^{'} (t) & = - a_{a u g} (t) \frac{\partial f_{a u g}}{\partial [h (t), θ (t), t]} \\ = - [a (t), a_{θ} (t), a_{t} (t)] [\begin{array}{ccc} \frac{\partial f}{\partial h (t)} & \frac{\partial f}{\partial θ (t)} & \frac{\partial f}{\partial t} \\ \frac{\partial 0}{\partial h (t)} & \frac{\partial 0}{\partial θ (t)} & \frac{\partial 0}{\partial t} \\ \frac{\partial 1}{\partial h (t)} & \frac{\partial 1}{\partial θ (t)} & \frac{\partial 1}{\partial t} \end{array}] \\ = - [a (t), a_{θ} (t), a_{t} (t)] [\begin{array}{ccc} \frac{\partial f}{\partial h (t)} & \frac{\partial f}{\partial θ (t)} & \frac{\partial f}{\partial t} \\ 0 & 0 & 0 \\ 0 & 0 & 0 \end{array}] \\ = - [a (t) \frac{\partial f}{\partial h (t)}, a (t) \frac{\partial f}{\partial θ (t)}, a (t) \frac{\partial f}{\partial t}] \end{aligned}$ 而初值是： $a_{a u g} (t_{1}) = [a (t_{1}), a_{θ} (t_{1}), a_{t} (t_{1})] = [\frac{\partial L}{\partial h (t_{1})}, \frac{\partial L}{\partial θ (t_{1})}, \frac{\partial L}{\partial t_{1}}] = [\frac{\partial L}{\partial h (t_{1})}, 0, a (t_{1}) f (t_{1}, h (t_{1}), θ)]$ 其中 $θ$ 项的初值为0是因为 $θ (t_{1})$ 并未参与计算（ $θ (t_{0})$ 用于计算 $h (t_{0} + δ)$ ）。这里 $t$ 的初值在论文中是带负号的，但感觉似乎不应该带负号，目前还没搞明白带负号的原因，github上也有一个issue讨论这个问题，目前没有讨论结果。

综合

将以上综合起来，我们只需要求解augmented后的ODE即可得到梯度： $[\frac{\partial L}{\partial h (t_{0})}, \frac{\partial L}{\partial θ (t_{0})}, \frac{\partial L}{\partial t_{0}}] = a_{a u g} (t_{0}) = ODESolver (a_{a u g} (t_{1}), f_{a u g}, t_{1}, t_{0}, θ)$ 之后就可以使用正常的训练过程了。

参考文献

[1] Neural ODEs

[2] Understanding Neural ODEs

[3] Neural Ordinary Differential Equations (NeurIPS 2018)

[4] Understanding Adjoint Method of Neural ODE