神经常微分方程

简介

神经常微分方程模型是一类新的深度神经网络模型,不同于VGG、ResNet等这些有有限个离散的隐藏层构成的神经网络模型。

例如残差网络、循环神经网络解码器、归一化流等模型建立起复杂的变换,是通过一个变换(相对简单的变换,比如ReLU变换)序列实现的。公式化表示为$$\mathbf h_{l+1} = \mathbf h_l + f_{l}(\mathbf h_l, \theta)$$ 论文中是使用公式$$\mathbf h_{t+1} = \mathbf h_t + f(\mathbf h_t, \theta_t)$$ 使用这种表示方式,便于扩展到常微分方程。

因为常微分方程$$\frac{d\mathbf h(t)}{dt} = f(\mathbf h(t), t, \theta)$$, 给定一个初始解,按照欧拉方法$$\mathbf h_{t+1} = \mathbf h_t + f(\mathbf h_t, t, \theta)$$可以迭代式求出变量\(\mathbf h\)的轨迹。可以看出常微分方程的欧拉迭代公式与离散神经网络模型的迭代过程 是一致的。除了每步迭代更新公式不一样。

类型 连续性 深度
常规神经网络模型 离散 有限个层(N层)
神经常微分方程模型 连续 无限,任意迭代次数

模仿常微分方程,用神经网络来建模方程右侧的函数,这个函数是关于要研究变量、模型参数的函数。

前向计算

采用常微分方程建模变量随时间的变化,那么就可以采用常微分数值解法,比如一阶欧拉法、四阶龙格-库塔法,有很多软件集成了这些方法,统一用常微分方程求解器简写成ODESolver。

反向计算

关于参数\(\mathbb z\)梯度

通过对前向操作进行微分,这样会变成与常规神经网络模型一样,需要计算关于参数和中间状态的导数,这样会增加内存的使用。因此引入一种新的方法,该方法将ODESlover看作黑盒,使用伴随敏感度方法——adjoint sensitivity method 计算梯度。该方法引入一个新的变量\(\mathbf a(t)\),通过直接计算该变量的微分方程解来实现对变量\(\mathbf z(t)\)的梯度求解。

现在考虑优化一个标量损失函数\(\mathbf L(\dots)\), 其输入是神经常微分方程模型最终结果\(\mathbf z_1\),也是ODESolver的结果。现在假设已知\(t_0\)时刻变量的状态值即\(\mathbf z_0 = \mathbf z(t_0)\),常微分方程参数\(\theta\), 损失函数\(L()\)也提前确定了,现在求解损失函数\(L\)关于状态变量\(\mathbf z(t)\)和模型参数\(\theta\)的梯度。

伴随敏感度方法,就是把损失函数\(L\)关于状态变量\(\mathbf z(t)\)的梯度,称为伴随量,这个伴随量命名为\(\mathbf a(t)\), 即\(\mathbf a(t) = \frac{\partial L}{\partial \mathbf z(t)}\)。这个量本身也是随着时间\(t\)变化的。 这个伴随量的动力学方程如下所示

\[\frac{d\mathbf a(t)}{dt} = -\mathbf a(t)^T \frac{\partial f(\mathbf z(t), t, \theta)}{\partial \mathbf z} \]

给定初始值\(\mathbf a(t_1)\),由ODESolver就可以求出\(\mathbf a(t)\) 在各个时间点上的状态值。使用ODESolver的前提条件是\(\frac{\partial f(\mathbf z(t), t, \theta)}{\partial \mathbf z}\) 是已知。

既然\(\mathbf a(t)\) 在各个时间点上的状态值都已经求解出来了,那么就依次获得了\(\frac{\partial L}{\partial z(t)}\) 在各个时间点\(t_1, t_{\sum_{1}^{N-1} \delta_{i}}, t_{\delta_{1}}, t_0\) 的状态值(这里面的\(\delta_{i}\)表示第i个时间步的长度)。这个顺序是逆序的,对应着反向传播的梯度。

现在上公式来推导得出上述微分方程(为避免损耗脑细胞,可以跳过不看):

推导依赖1: 根据链式法则 \(\frac{\partial L}{\partial \mathbf z(t)} = \frac{\partial L}{\partial \mathbf z(t+1)} \frac{\partial z(t+1)}{\partial \mathbf z(t)}\), 因为求关于变量\(\mathbf z(t)\)的梯度,按照时间步,是从最大的时间开始求解再依次求较小的时间。因此这里是这样的链式形式。 此外由于是求解连续梯度,因此将上面改写成 \(\frac{\partial L}{\partial \mathbf z(t)} = \frac{\partial L}{\partial \mathbf z(t+\epsilon)} \frac{\partial \mathbf z(t+\epsilon)}{\partial \mathbf z(t)}\). 注意上面第一个公式t+1更多表达时间步索引,而第二个公式就表示t时间本身。

下面假设向量\(\mathbf z(t), \mathbf z(t+\epsilon)\) 的维度是d; 且向量对向量求导,采用分子布局方式。

可以按照下面来理解

\[L \leftarrow z(t+\epsilon) \leftarrow z(t) \]

\(y=z(t+\epsilon)\),此处\(z(t+\epsilon)\), \(z(t)\)都假设为行向量\(\mathbb {R}^{1 \times d}\),那么

\[\frac{\partial L}{\partial z} = (\frac{dL}{dz_1}, \frac{dL}{dz_2}, \cdots, \frac{dL}{dz_d}) \]

\[\frac{\partial L}{\partial y} = (\frac{dL}{dy_1}, \frac{dL}{dy_2}, \cdots, \frac{dL}{dy_d}) \]

\[\frac{\partial y}{\partial z} = \left[ \begin{array}{} \frac{dy_1}{dz_1} & \frac{dy_1}{dz_2} & \cdots & \frac{dy_1}{dz_d} \\ \frac{dy_2}{dz_1} & \frac{dy_2}{dz_2} & \cdots & \frac{dy_2}{dz_d} \\ \cdots && \\ \frac{dy_d}{dz_1} & \frac{dy_d}{dz_2} & \cdots & \frac{dy_d}{dz_d} \\ \end{array} \right] \]

行对应\(y^T\)的分量,列对应着\(x\)的分量,可以看出

\[\frac{\partial L}{\partial y}\frac{\partial y}{\partial z} = (\frac{dL}{dy_1}, \frac{dL}{dy_2}, \cdots, \frac{dL}{dy_d}) \left[ \begin{array}{} \frac{dy_1}{dz_1} & \frac{dy_1}{dz_2} & \cdots & \frac{dy_1}{dz_d} \\ \frac{dy_2}{dz_1} & \frac{dy_2}{dz_2} & \cdots & \frac{dy_2}{dz_d} \\ \cdots && \\ \frac{dy_d}{dz_1} & \frac{dy_d}{dz_2} & \cdots & \frac{dy_d}{dz_d} \\ \end{array} \right] = (\sum\frac{dL}{dy_i}\frac{dy_i}{dz_1}, \sum\frac{dL}{dy_i}\frac{dy_i}{dz_2}, \cdots, \sum\frac{dL}{dy_i}\frac{dy_i}{dz_d}) \]

因此

\[\frac{\partial L}{\partial z} = \frac{\partial L}{\partial y}\frac{\partial y}{\partial z} \]

推导依赖2: 泰勒公式
既然知道了变量\(\mathbf z(t)\) 的一阶导数,那么很自然由泰勒公式就直接得出 \(\mathbf z(t+\epsilon) = \mathbf z(t) + \epsilon \frac{d\mathbf z(t)}{dt} + o(\epsilon^2)\)

推理依赖3: 导数的定义

\[\begin{aligned} \frac{\mathbf a(t+\epsilon) - \mathbf a(t)}{\epsilon} &= \frac{\frac{\partial L}{\partial \mathbf z(t+\epsilon)} - \frac{\partial L}{\partial \mathbf{z}(t)}}{\epsilon} \\ &= \frac{\frac{\partial L}{\partial \mathbf z(t+\epsilon)} - \frac{\partial L}{\partial \mathbf z(t+\epsilon)} \frac{\partial \mathbf z(t+\epsilon)}{\partial \mathbf z(t)} }{\epsilon} \\ &= \frac{\frac{\partial L}{\partial \mathbf z(t+\epsilon)}}{\epsilon}[I - \frac{\partial \mathbf z(t+\epsilon)}{\partial \mathbf z(t)}] \\ &= \frac{\frac{\partial L}{\partial \mathbf z(t+\epsilon)}}{\epsilon}[I - \frac{\partial }{\partial \mathbf z^(t)}(\mathbf z(t) + f(\mathbf z(t), t, \theta)\epsilon + o(\epsilon^2))] \\ &= \frac{\frac{\partial L}{\partial \mathbf z(t+\epsilon)}}{\epsilon}[I - (I + \frac{\partial f(\mathbf z(t), t, \theta)}{\partial \mathbf z(t)}\epsilon)] \\ &= \frac{\frac{\partial L}{\partial \mathbf z(t+\epsilon)}}{\epsilon}[- \frac{\partial f(\mathbf z(t), t, \theta)}{\partial \mathbf z(t)}\epsilon ] \\ &= -\frac{\partial L}{\partial \mathbf z(t+\epsilon)}\frac{\partial f(\mathbf z(t), t, \theta)}{\partial \mathbf z(t)} \end{aligned} \]

两边同时取关于\(\epsilon\)趋近于0的极限,便可以得到\(\frac{d \mathbf a(t)}{dt} = -\mathbf a(t)\frac{\partial f(\mathbf z(t), t, \theta)}{\partial \mathbf z(t)}\)

设有时间点\(t_1, t_2, \cdots, t_{N-1}, t_{N}\) , 知道初始时间点\(t_N\)上损失值关于变量\(\mathbf z(t)\)的梯度,即\(\mathbf a(t_N) = \frac{\partial L}{d\mathbf z(t_N)}\)。那么由ODESolver可以,依次求解出\(\mathbf a(t_{N-1}), \mathbf a(t_{N-2}), \cdots, \mathbf a(t_1)\).

到这里,关于变量\(\mathbf z(t)\)的神经微分方程网络,其各个时间点前向状态值,及其反向梯度值都完全求解出来了。

关于参数\(\theta\)\(t\)的梯度

将参数\(\theta, t\)关于时间t的微分方程写作为

\[\frac{\partial \theta(t)}{\partial t}=\mathbb 0, \frac{\partial t(t)}{\partial t}= 1 \]

现在将状态变量\(\mathbb z(t)\)进行扩充,添加参数\(\theta, t\),那么其对应的微分方程为

\[\frac{d[\mathbb z, \theta, t]}{dt} = f_{aug}[\mathbb z, \theta, t] = [f(\mathbb z,\theta,t), \mathbb 0, 1] \]

它们对应的伴随量记为

\[a_{aug} = [a, a_{\theta}, a_{t}], a_{\theta}(t) = \frac{dL}{d\theta(t)}, a_{t}(t) = \frac{dL}{d t(t)} \]

对扩充的\(\mathbb z\), 其伴随量a_{aug}的导数可以推导得到为

\[\frac{da_{aug}(t)}{dt} = -a_{aug}(t) \frac{\partial f_{aug}}{\partial [\mathbb z, \theta, t]}(t)=-[a\frac{\partial f}{\partial \mathbb z}, a\frac{\partial f}{\partial \theta}, a\frac{\partial f}{\partial t}] \]

这个可以下面公式推导得出

\[\frac{\partial f_{aug}}{\partial [\mathbb z, \theta, t]}(t)= \left [ \begin{array}{} \frac{\partial f}{\partial \mathbb z}, \frac{\partial f}{\partial \theta}, \frac{\partial f}{\partial t} \\ 0, 0, 0 \\ 0, 0, 0 \\ \end{array} \right ] \]

向量求导布局

矩阵向量求导的本质就是多元函数求导,仅仅是把因变量分量对自变量分量的求导结果排列成了向量矩阵的形式,为未来方便表达与计算而已。但是在矩阵向量求导中,其求导结果会因向量和矩阵的形式而导致结果不唯一,为此引入了求导布局的概念。

比如参考1中,提到m维列向量\(\mathbf y\)对n维列向量\(\mathbf x\)求导。这两个向量求导,一共有m个标量对n个标量分别求导,共计mn个求导。求导的结果就可以排列为一个矩阵。

  • 如果是分子布局,则矩阵的第一个维度以分子的维度为准,那就是m行n列的矩阵即

\[\left[ \begin{array}{} \frac{dy_1}{dx_1} & \frac{dy_1}{dx_2} & \cdots & \frac{dy_1}{dx_n} \\ \frac{dy_2}{dx_1} & \frac{dy_2}{dx_2} & \cdots & \frac{dy_2}{dx_n} \\ \cdots && \\ \frac{dy_m}{dx_1} & \frac{dy_m}{dx_2} & \cdots & \frac{dy_m}{dx_n} \\ \end{array} \right] \]

  • 如果是分母布局,则求导的结果矩阵的第一个维度以分母的维度为准,那就是n行m列的矩阵即

\[\left[ \begin{array}{} \frac{dy_1}{dx_1} & \frac{dy_2}{dx_1} & \cdots & \frac{dy_m}{dx_1} \\ \frac{dy_1}{dx_2} & \frac{dy_2}{dx_2} & \cdots & \frac{dy_m}{dx_2} \\ \cdots && \\ \frac{dy_1}{dx_n} & \frac{dy_2}{dx_n} & \cdots & \frac{dy_m}{dx_n} \\ \end{array} \right] \]

可以看出按分子布局和按分母布局,得到的求导矩阵是互为转置的。

比如参考2中,假设某函数从\(\mathbf f: \mathbb R^n \rightarrow \mathbb R^m\),从n维向量\(\mathbf x \in \mathbb R^n\)映射到m维向量\(\mathbf f(\mathbf x) \in \mathbb R^m\), 求导数为

\[\left[ \begin{array}{} \frac{df_1}{dx_1} & \frac{df_1}{dx_2} & \cdots & \frac{df_1}{dx_n} \\ \frac{df_2}{dx_1} & \frac{df_2}{dx_2} & \cdots & \frac{df_2}{dx_n} \\ \cdots && \\ \frac{df_m}{dx_1} & \frac{df_m}{dx_2} & \cdots & \frac{df_m}{dx_n} \\ \end{array} \right] \]

参考

posted @ 2025-04-06 17:26  星辰大海,绿色星球  阅读(109)  评论(0)    收藏  举报