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

原创苏剑林 PaperWeekly

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

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

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

论文标题：

Parallelizing Non-Linear Sequential Models over the Sequence Length

论文链接：

https://arxiv.org/pdf/2309.12252.pdf

求不动点

原论文对其方法做了非常一般的介绍，而且其侧重点是 PDE 和 ODE，这里我们直接从 RNN 入手。考虑常见的简单非线性 RNN：

由于的存在，它只能串行计算。现在我们在两边都减去：

当然，这改变不了它是非线性 RNN 的实质。然而我们可以发现，假如右端的换成像那样的给定向量，那么这就是一个线性 RNN 了，根据《Google新作试图“复活”RNN：RNN能否再次辉煌？》的结果，它是可以并行计算的。此时，敏捷的读者可能已经猜到后面的步骤了——迭代求解！

首先，将上述 RNN 更改成

从给定出发，反复迭代上式，理想情况下，它会收敛于一个不动点，这就是原来非线性 RNN 的计算结果。当然，理论上通过式（3）迭代的总计算量是比直接通过式（1）递归计算要大的，但由于每一步迭代都是可并行的线性 RNN，并且如果收敛速度比较快时迭代步数不需要太多，那么总的耗时通常都会快于直接非线性 RNN 递归（尤其是序列长度很大时）。

简化形式

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

当我们通过式（3）将非线性 RNN 转为线性 RNN 的迭代之后，同样享受线性 RNN 可对角化的“待遇”，从而提高计算速度。具体来说，在复数域中将 A 对角化为，那么式（3）变为

两端都左乘：

令，那么上式可以简化为

由于 RNN 之后一般都还要接个投影层，所以的 P 原则上可以合并到外接的投影层里边，也就是说，上式理论上具备跟原来的（1）具备同等的表达能力，但由于是对角阵，递归的计算量会明显降低。上式还出现了逆矩阵，不单计算量大，而且不利于优化，所以我们可以干脆将和换成两个不相关的参数矩阵：

只要初始化是就行。

摄动思想

假定，那么式（3）其实就是将原本的非线性 RNN 就分解为一系列线性 RNN：

而假设都是小量，那么对式（1）右端利用得到：

这正好是（8）中的第一个方程，因此如果假设成立，那么或许已经足够接近理想的，后面的每一步迭代都在快速逼近它。从这里我们可以看出，“两边同时减去 ”是关键之处，这使得（3）的第一步迭代就接近于原本非线性 RNN 的一阶线性近似，这可以提高收敛速度，也是数学物理中的经典操作，名曰“摄动”。

加快收敛

根据摄动法的思想，提高收敛速度的关键就是提高近似展开的精度，比如较为简单的改进是只假设是小量，那么根据一阶泰勒展开有（将视为列向量，这里的是 Hadamard 积分）

于是改进的结果就是式（3）变为

其中。更精细的改进是在每一步迭代时，都在前一步迭代结果的基础上进行展开：

于是式（3）变为

其中

。最后的这个迭代格式，实际上就是求方程数值解的“牛顿法”，它具有二次收敛速度。

何必收敛

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

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

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

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

一般情形

前面的“长篇大论”，都只围绕着简单的非线性 RNN 也就是式（1）进行讨论，对于更常用的 LSTM、GRU，结果又如何呢？

以 GRU 为例，它原本的形式为

初始阶段，所有门控都可以近似视为，那么模仿式（9）有

所以可以选取，将 GRU 改写为迭代

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