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

©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能否再次辉煌？》的结果，它是可以并行计算的。此时，敏捷的读者可能已经猜到后面的步骤了——迭代求解！

首先ÿ