脑洞大开：非线性RNN居然也可以并行计算？

近年来，线性RNN由于其可并行训练以及常数推理成本等特性，吸引了一定研究人员的关注（例如笔者之前写的《Google新作试图“复活”RNN：RNN能否再次辉煌？》），这让RNN在Transformer遍地开花的潮流中仍有“一席之地”。然而，目前看来这“一席之地”只属于线性RNN，因为非线性RNN无法高效地并行训练，所以在架构之争中是“心有余而力不足”。

不过，一篇名为《Parallelizing Non-Linear Sequential Models over the Sequence Length》的论文有不同的看法，它提出了一种迭代算法，宣传可以实现非线性RNN的并行训练！真有如此神奇？接下来我们一探究竟。

求不动点 #

原论文对其方法做了非常一般的介绍，而且其侧重点是PDE和ODE，这里我们直接从RNN入手。考虑常见的简单非线性RNN：
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t)\label{eq:rnn}\end{equation}
由于$\tanh$的存在，它只能串行计算。现在我们在两边都减去$Ax_{t-1}$：
\begin{equation}x_t - Ax_{t-1} = \tanh(Ax_{t-1} + u_t) - Ax_{t-1}\end{equation}
当然，这改变不了它是非线性RNN的实质。然而我们可以发现，假如右端的$x_{t-1}$换成像$u_t$那样的给定向量，那么这就是一个线性RNN了，根据《Google新作试图“复活”RNN：RNN能否再次辉煌？》的结果，它是可以并行计算的。此时，敏捷的读者可能已经猜到后面的步骤了——迭代求解！

首先，将上述RNN更改成
\begin{equation}x_t^{(n)} - Ax_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - Ax_{t-1}^{(n-1)}\label{eq:rnn-iter}\end{equation}
从给定$x_t^{(0)}$出发，反复迭代上式，理想情况下，它会收敛于一个不动点$x_t^*$，这就是原来非线性RNN的计算结果。当然，理论上通过式$\eqref{eq:rnn-iter}$迭代的总计算量是比直接通过式$\eqref{eq:rnn}$递归计算要大的，但由于每一步迭代都是可并行的线性RNN，并且如果收敛速度比较快时迭代步数不需要太多，那么总的耗时通常都会快于直接非线性RNN递归（尤其是序列长度很大时）。

简化形式 #

事实上，非线性RNN之所以慢，无法并行计算还是次要的，最关键是它包含了大量的非element-wise运算，比如式$\eqref{eq:rnn}$的$\tanh$里边的矩阵运算$Ax_{t-1}$；而线性RNN之所以快，除了它允许并行训练之外，更关键的是它能通过对角化来将矩阵乘法变换为element-wise的乘法——对于element-wise乘法来说，即便是串行计算也不会太慢。

当我们通过式$\eqref{eq:rnn-iter}$将非线性RNN转为线性RNN的迭代之后，同样享受线性RNN可对角化的“待遇”，从而提高计算速度。具体来说，在复数域中将$A$对角化为$P\Lambda P^{-1}$，那么式$\eqref{eq:rnn-iter}$变为
\begin{equation}x_t^{(n)} - P\Lambda P^{-1} x_{t-1}^{(n)} = \tanh(P\Lambda P^{-1} x_{t-1}^{(n-1)} + u_t) - P\Lambda P^{-1} x_{t-1}^{(n-1)}\end{equation}
两端都左乘$P^{-1}$：
\begin{equation}P^{-1} x_t^{(n)} - \Lambda P^{-1} x_{t-1}^{(n)} = P^{-1}\tanh(P\Lambda P^{-1} x_{t-1}^{(n-1)} + u_t) - \Lambda P^{-1} x_{t-1}^{(n-1)}\end{equation}
令$y_t = P^{-1} x_t$，那么上式可以简化为
\begin{equation}y_t^{(n)} - \Lambda y_{t-1}^{(n)} = P^{-1}\tanh(P\Lambda y_{t-1}^{(n-1)} + u_t) - \Lambda y_{t-1}^{(n-1)}\end{equation}
由于RNN之后一般都还要接个投影层，所以$x_t = P y_t$的$P$原则上可以合并到外接的投影层里边，也就是说，上式理论上具备跟原来的$\eqref{eq:rnn}$具备同等的表达能力，但由于$\Lambda$是对角阵，递归的计算量会明显降低。上式还出现了逆矩阵$P^{-1}$，不单计算量大，而且不利于优化，所以我们可以干脆将$P^{-1}$和$P\Lambda$换成两个不相关的参数矩阵：
\begin{equation}y_t^{(n)} - \Lambda y_{t-1}^{(n)} = P\tanh(Q y_{t-1}^{(n-1)} + u_t) - \Lambda y_{t-1}^{(n-1)}\end{equation}
只要初始化是$PQ=\Lambda$就行。

摄动思想 #

假定$x_t^{(0)}=0$，那么式$\eqref{eq:rnn-iter}$其实就是将原本的非线性RNN就分解为一系列线性RNN：
\begin{equation}\begin{array}{c}
x_t^{(1)} - Ax_{t-1}^{(1)} = \tanh(u_t)\\
x_t^{(2)} - Ax_{t-1}^{(2)} = \tanh(Ax_{t-1}^{(1)} + u_t) - Ax_{t-1}^{(1)} \\
\vdots \\
x_t^{(n)} - Ax_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - Ax_{t-1}^{(n-1)} \\
\vdots \\
\end{array}\label{eq:rnns}\end{equation}
而假设$x_{t-1},u_t$都是小量，那么对式$\eqref{eq:rnn}$右端利用$\tanh x \approx x$得到：
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t) \approx Ax_{t-1} + u_t \approx Ax_{t-1} + \tanh(u_t)\label{eq:rnn-approx}\end{equation}
这正好是$\eqref{eq:rnns}$中的第一个方程，因此如果假设成立，那么$x_t^{(1)}$或许已经足够接近理想的$x_t^*$，后面的每一步迭代都在快速逼近它。从这里我们可以看出，“两边同时减去$Ax_{t-1}$”是关键之处，这使得$\eqref{eq:rnn-iter}$的第一步迭代就接近于原本非线性RNN的一阶线性近似，这可以提高收敛速度，也是数学物理中的经典操作，名曰“摄动”。

加快收敛 #

根据摄动法的思想，提高收敛速度的关键就是提高近似展开的精度，比如较为简单的改进是只假设$x_{t-1}$是小量，那么根据一阶泰勒展开有（将$u_t$视为列向量，这里的$\circ$是Hadamard积分）
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t) \approx \tanh(u_t) + (\text{sech}^2 u_t\circ A)x_{t-1}\end{equation}
于是改进的结果就是式$\eqref{eq:rnn-iter}$变为
\begin{equation}x_t^{(n)} - A_t x_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - A_t x_{t-1}^{(n-1)}\label{eq:iter-plus1}\end{equation}
其中$A_t = \text{sech}^2 u_t\circ A$。更精细的改进是在每一步迭代时，都在前一步迭代结果的基础上进行展开：
\begin{equation}\begin{aligned}
x_t =&\, \tanh(Ax_{t-1} + u_t) \\
\approx&\, \tanh(Ax_{t-1}^{(n-1)} + u_t) + (\text{sech}^2 (Ax_{t-1}^{(n-1)} + u_t)\circ A)(x_{t-1} - x_{t-1}^{(n-1)})
\end{aligned}\end{equation}
于是式$\eqref{eq:rnn-iter}$变为
\begin{equation}x_t^{(n)} - A_t^{(n)} x_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - A_t^{(n)} x_{t-1}^{(n-1)}\label{eq:iter-plus2}\end{equation}
其中$A_t^{(n)}=\text{sech}^2 (Ax_{t-1}^{(n-1)} + u_t)\circ A$。最后的这个迭代格式，实际上就是求方程数值解的“牛顿法”，它具有二次收敛速度。

何必收敛 #

理论上来说，$\eqref{eq:iter-plus1}$、$\eqref{eq:iter-plus2}$两个改进确实能提高收敛速度，然而它们使得每一步线性递归的矩阵$A$变得跟$t$甚至$n$有关了，这其实会大大增加并行的复杂度，也不能利用“简化形式”一节的对角化技巧来加速。另一方面，如果保持$\eqref{eq:rnn-iter}$这样的迭代格式，虽然有诸多效率上的好处，但收敛方面确实无法得到很好的保障。

难道这两者的矛盾就无法调和了吗？事实上，按照笔者的观点，最直接的做法是“别去管它”——借助非线性RNN导出了$\eqref{eq:rnn-iter}$后，就忘记原本的非线性RNN，将式$\eqref{eq:rnn-iter}$作为基本模型。也就是说，何必忧虑式$\eqref{eq:rnn-iter}$会不会收敛到原来的非线性RNN？直接将它作为新的出发点不好吗？梯度下降学到怎样的结果就是怎样的结果，如果梯度下降学到的结果是不收敛到原来的非线性RNN，那么就意味着不收敛到原来的RNN是更适合的。

抛开这一层思维束缚后，其实很多问题会变得豁然开朗起来。首先，即便是式$\eqref{eq:iter-plus2}$在理论上拥有非常好的收敛速度，但也是有条件的，而且在深度学习的背景下，要保证这些条件会显得很奢侈。换言之，即便是式$\eqref{eq:iter-plus2}$的收敛性也没有绝对保证，所以何必“五十步笑百步”去苛责式$\eqref{eq:rnn-iter}$？其次，将式$\eqref{eq:rnn-iter}$视为新的出发点后，我们可以将它单纯地理解为线性RNN的一种新用法，或者说解决线性RNN缺陷（比如线性RNN不是图灵完备的）的一个思路，这样操作性更强。

总的来说，不去管它的收敛性，似乎更能打破思维僵局，探索更一般的结果。

一般情形 #

前面的“长篇大论”，都只围绕着简单的非线性RNN也就是式$\eqref{eq:rnn}$进行讨论，对于更常用的LSTM、GRU，结果又如何呢？

以GRU为例，它原本的形式为
\begin{equation}\begin{aligned} z_{t} & = \sigma \left( W_{z} x_{t} + U_{z} h_{t - 1} + b_{z} \right) \\
r_{t} & = \sigma \left( W_{r} x_{t} + U_{r} h_{t - 1} + b_{r} \right) \\
\hat{h}_t & = \tanh \left( W_{h} x_{t} + U_{h} (r_t \circ h_{t - 1}) + b_{c} \right)\\
h_{t} & = \left(1 - z_{t}\right) \circ h_{t - 1} + z_{t} \circ \hat{h}_t \end{aligned}\end{equation}
初始阶段，所有门控都可以近似视为$\frac{1}{2}$，那么模仿式$\eqref{eq:rnn-approx}$有
\begin{equation}\begin{aligned}
h_{t} &\, = \left(1 - z_{t}\right) \circ h_{t - 1} + z_{t} \circ \hat{h}_t \\
&\, \approx \frac{1}{2} h_{t - 1} + \frac{1}{2} \hat{h}_t \\
&\, \approx \frac{1}{2} h_{t - 1} + \frac{1}{2} \left(\tanh ( W_{h} x_{t} + b_{c} ) + \frac{1}{2}U_{h} h_{t - 1}\right) \\
&\, = \frac{1}{2} \left(I + \frac{1}{2}U_{h}\right)h_{t - 1} + \frac{1}{2} \tanh ( W_{h} x_{t} + b_{c} ) \\
\end{aligned}\end{equation}
所以可以选取$A=\frac{1}{2} \left(I + \frac{1}{2}U_{h}\right)$，将GRU改写为迭代
\begin{equation}\begin{aligned} z_{t}^{(n)} & = \sigma \left( W_{z} x_{t} + U_{z} h_{t - 1}^{(n-1)} + b_{z} \right) \\
r_{t}^{(n)} & = \sigma \left( W_{r} x_{t} + U_{r} h_{t - 1}^{(n-1)} + b_{r} \right) \\
\hat{h}_t^{(n)} & = \tanh \left( W_{h} x_{t} + U_{h} (r_t^{(n)} \circ h_{t - 1}^{(n-1)}) + b_{c} \right)\\
h_{t}^{(n)} & = Ah_{t-1}^{(n)} - Ah_{t-1}^{(n - 1)} + \left(1 - z_{t}^{(n)}\right) \circ h_{t - 1}^{(n-1)} + z_{t}^{(n)} \circ \hat{h}_t^{(n)} \end{aligned}\end{equation}

总的来说，这种将非线性RNN变为线性RNN迭代的转换，从实践的角度来看，就是以非线性RNN为引，导出一种多层线性RNN的参数共享和组合方法，迭代了几次，那么就有几层线性RNN的计算量。这样自然而言就引发了一个思考：除非可以证明GRU、LSTM等非线性RNN有绝对的优势，否则直接叠加几层“线性RNN+MLP”不好吗？

文章小结 #

本文简单探讨了非线性RNN的并行计算问题——通过数学物理中的“摄动”思想，我们可以将非线性RNN转化为线性RNN的迭代，从而利用线性RNN的可并行性来实现非线性RNN的并行。

转载到请包括本文地址：https://kexue.fm/archives/9783

更详细的转载事宜请参考：《科学空间FAQ》

打赏

微信打赏

支付宝打赏

因为网站后台对打赏并无记录，因此欢迎在打赏时候备注留言。
你还可以点击这里或在下方评论区留言来告知你的建议或需求。